perf(agent): share rpc response wait across result types

.github/workflows/ci.yml

@@ -653,6 +653,46 @@ jobs:
           if-no-files-found: error
           retention-days: 7
 
+      - name: Measure Agent RPC monomorphization
+        if: github.head_ref == 'perf/win7-agent-rpc-monomorphization'
+        shell: pwsh
+        run: |
+          $outputDir = Join-Path $env:RUNNER_TEMP "dbx-agent-rpc-mono"
+          New-Item -ItemType Directory -Force -Path $outputDir | Out-Null
+          $env:TAURI_CONFIG = Get-Content src-tauri/tauri.webview2-win7-fixed.conf.json -Raw
+          $env:RUSTC_WRAPPER = ""
+          $timer = [System.Diagnostics.Stopwatch]::StartNew()
+          cargo -Z host-config -Z target-applies-to-host -Z build-std=std,panic_abort `
+            --config .github/fixtures/win7-host-repro-config.toml `
+            rustc --locked --package dbx --lib --release --features custom-protocol `
+            --target x86_64-win7-windows-msvc -- `
+            "-Zdump-mono-stats=$outputDir" "-Zdump-mono-stats-format=json"
+          if ($LASTEXITCODE -ne 0) {
+            exit $LASTEXITCODE
+          }
+          $timer.Stop()
+          $monoFile = Get-ChildItem $outputDir -File -Filter "*.mono_items.json" | Select-Object -First 1
+          if (!$monoFile) {
+            throw "dbx_lib monomorphization data is missing."
+          }
+          $items = Get-Content $monoFile.FullName -Raw | ConvertFrom-Json
+          "dbx_lib profile: $([math]::Round($timer.Elapsed.TotalSeconds, 2)) seconds" >> $env:GITHUB_STEP_SUMMARY
+          foreach ($symbol in @("AgentDriverClient::call_typed_with_timeout_and_cancel::{closure#0}", "AgentDriverClient::call_value_with_timeout_and_cancel::{closure#0}", "AgentRuntimeClient::call_typed::{closure#0}")) {
+            $item = $items | Where-Object { $_.name.EndsWith($symbol) } | Select-Object -First 1
+            if ($item) {
+              "$symbol`: $($item.instantiation_count) instances, estimate $($item.total_estimate)" >> $env:GITHUB_STEP_SUMMARY
+            }
+          }
+
+      - name: Upload Agent RPC monomorphization data
+        if: github.head_ref == 'perf/win7-agent-rpc-monomorphization'
+        uses: actions/upload-artifact@v6
+        with:
+          name: DBX-win7-agent-rpc-mono
+          path: ${{ runner.temp }}/dbx-agent-rpc-mono/*.mono_items.json
+          if-no-files-found: error
+          retention-days: 7
+
       - name: Measure dbx library command registry expansion
         if: vars.WIN7_REGISTRY_AB == 'true'
         shell: pwsh

crates/dbx-driver-agent/src/agent_driver.rs

@@ -1,3 +1,4 @@
+use std::any::Any;
 use std::collections::{HashMap, VecDeque};
 use std::fmt;
 use std::io::{BufRead, BufReader, BufWriter, Write};
@@ -17,6 +18,14 @@ use crate::agent_recovery::{RecoveryDecision, RecoveryPolicy, RecoveryScope};
 use crate::models::connection::DatabaseConnectionInfo;
 
 type PendingAgentResponse = tokio::sync::oneshot::Sender<Result<Value, String>>;
+type AgentDecodedResult = Box<dyn Any + Send>;
+type AgentResultDecoder = fn(Value) -> Result<AgentDecodedResult, String>;
+
+fn deserialize_agent_result<T: DeserializeOwned + Send + 'static>(value: Value) -> Result<AgentDecodedResult, String> {
+    serde_json::from_value::<T>(value)
+        .map(|result| Box::new(result) as AgentDecodedResult)
+        .map_err(|error| format!("Failed to deserialize agent result: {error}"))
+}
 
 #[derive(Clone)]
 struct CachedAgentQuery {
@@ -1809,6 +1818,26 @@ impl AgentDriverClient {
         timeout_duration: Option<Duration>,
         cancel_token: Option<CancellationToken>,
     ) -> Result<T, AgentCallError> {
+        let result = self
+            .call_decoded_with_timeout_and_cancel(
+                method,
+                params,
+                timeout_duration,
+                cancel_token,
+                deserialize_agent_result::<T>,
+            )
+            .await?;
+        Ok(*result.downcast::<T>().expect("Agent result decoder returned a different type"))
+    }
+
+    async fn call_decoded_with_timeout_and_cancel(
+        &mut self,
+        method: &str,
+        params: Value,
+        timeout_duration: Option<Duration>,
+        cancel_token: Option<CancellationToken>,
+        decode: AgentResultDecoder,
+    ) -> Result<AgentDecodedResult, AgentCallError> {
         if let Some(runtime) = &self.shared_runtime {
             let agent_session_id = self.agent_session_id.clone();
             let mut params = params;
@@ -1822,10 +1851,15 @@ impl AgentDriverClient {
                     .ok_or_else(|| "Agent RPC parameters must be an object".to_string())?
                     .insert("agentSessionId".to_string(), Value::String(session_id.clone()));
             }
-            return runtime
+            let result: Value = runtime
                 .call_typed(method, params, timeout_duration, cancel_token)
                 .await
-                .map_err(|error| error.with_legacy_session_id(agent_session_id.as_deref()));
+                .map_err(|error| error.with_legacy_session_id(agent_session_id.as_deref()))?;
+            return decode(result).map_err(|message| AgentCallError::ContractViolation {
+                rpc_code: None,
+                message,
+                reason: ContractViolationReason::InvalidResult,
+            });
         }
         self.next_id += 1;
         let id = self.next_id;
@@ -1924,8 +1958,7 @@ impl AgentDriverClient {
             let result = if let Some(err) = resp.get("error") {
                 Err(format_agent_rpc_error(err))
             } else if let Some(result_val) = resp.get("result") {
-                serde_json::from_value::<T>(result_val.clone())
-                    .map_err(|e| format!("Failed to deserialize agent result: {e}"))
+                decode(result_val.clone())
             } else {
                 Err(format!("Agent response missing both 'result' and 'error': {line}"))
             };