Merge pull request #135 from TransformerOptimus/feat/bench-runner-llm-policy

adithyn7 · web-flow · commit 6e1048e8edcb · 2026-06-12T17:13:44.000+05:30
Bench-runner LLM transport fix: HTTP/1.1 + no pool + 15s header timeout (LlmPolicy)
diff --git a/apps/desktop/src-tauri/src/agent_bridge/commands.rs b/apps/desktop/src-tauri/src/agent_bridge/commands.rs
@@ -592,6 +592,7 @@ pub fn provider_to_llm_config(p: &ProviderConfig, model: &str) -> LlmClientConfi
         extra_headers: vec![],
         thinking: None,
         disable_cache_control: false,
+        policy: Default::default(),
     }
 }
 
diff --git a/crates/agent/src/agent/loop_.rs b/crates/agent/src/agent/loop_.rs
@@ -1555,6 +1555,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             std::path::PathBuf::from("/tmp"),
         );
@@ -1950,6 +1951,52 @@ mod tests {
         assert_eq!(result, "ok");
     }
 
+    #[tokio::test]
+    async fn test_retry_on_header_timeout() {
+        // HeaderTimeout is the new non-API error class added for bench-runner's
+        // 15s header-arrival timeout. Like ChunkTimeout it must take the
+        // retryable fall-through arm — otherwise a stalled-headers re-send never
+        // fires and the run dies on the first attempt.
+        let mock = MockLlm::new(vec![
+            Err(AgentError::HeaderTimeout(15_000)),
+            Ok(text_response("recovered")),
+        ]);
+        let (mut agent, _rx) = make_agent(mock, ToolRegistry::new(), None);
+        agent.config.retry_config.max_retries = 3;
+
+        let result = agent.run(ChatMessage::user("test")).await.unwrap().unwrap_done();
+        assert_eq!(result, "recovered", "HeaderTimeout must be retryable");
+    }
+
+    #[test]
+    fn test_llm_policy_values() {
+        let app = crate::llm::LlmPolicy::default();
+        assert!(!app.http1_only);
+        assert!(!app.no_pool);
+        assert_eq!(app.header_timeout_ms, None);
+
+        let bench = crate::llm::LlmPolicy::bench();
+        assert!(bench.http1_only);
+        assert!(bench.no_pool);
+        assert_eq!(bench.header_timeout_ms, Some(15_000));
+    }
+
+    #[tokio::test]
+    async fn test_retry_on_chunk_timeout_nonapi_error() {
+        // ChunkTimeout takes the SAME fall-through arm in call_llm_with_retry as
+        // HttpError (mid-stream decode). If this retries, the non-API error path
+        // re-sends; if it doesn't, the decode-bypass bug is in the classification.
+        let mock = MockLlm::new(vec![
+            Err(AgentError::ChunkTimeout(60)),
+            Ok(text_response("recovered")),
+        ]);
+        let (mut agent, _rx) = make_agent(mock, ToolRegistry::new(), None);
+        agent.config.retry_config.max_retries = 3;
+
+        let result = agent.run(ChatMessage::user("test")).await.unwrap().unwrap_done();
+        assert_eq!(result, "recovered", "non-API error (ChunkTimeout) must re-send");
+    }
+
     #[tokio::test]
     async fn test_no_retry_on_401() {
         let mock = MockLlm::new(vec![Err(AgentError::LlmApiError {
@@ -2297,6 +2344,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             std::path::PathBuf::from("/tmp"),
         );
@@ -2362,6 +2410,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             std::path::PathBuf::from("/tmp"),
         );
@@ -2427,6 +2476,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             std::path::PathBuf::from("/tmp"),
         );
@@ -2521,6 +2571,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             std::path::PathBuf::from("/tmp"),
         );
@@ -2617,6 +2668,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             std::path::PathBuf::from("/tmp"),
         );
@@ -2699,6 +2751,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             std::path::PathBuf::from("/tmp"),
         );
@@ -2773,6 +2826,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             work_dir.path().to_path_buf(),
         );
@@ -3022,6 +3076,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             tmp.path().to_path_buf(),
         );
@@ -3089,6 +3144,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             tmp.path().to_path_buf(),
         );
@@ -3164,6 +3220,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             tmp.path().to_path_buf(),
         );
@@ -3230,6 +3287,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             tmp.path().to_path_buf(),
         );
@@ -3298,6 +3356,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             tmp.path().to_path_buf(),
         );
@@ -3372,6 +3431,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             tmp.path().to_path_buf(),
         );
@@ -3590,6 +3650,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             tmp.path().to_path_buf(),
         );
@@ -3675,6 +3736,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             tmp.path().to_path_buf(),
         );
@@ -3730,6 +3792,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             tmp.path().to_path_buf(),
         );
@@ -3792,6 +3855,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             std::path::PathBuf::from("/tmp"),
         );
diff --git a/crates/agent/src/error.rs b/crates/agent/src/error.rs
@@ -12,6 +12,9 @@ pub enum AgentError {
     #[error("SSE chunk timeout: no data received for {0}s")]
     ChunkTimeout(u64),
 
+    #[error("LLM header timeout: response headers did not arrive within {0}ms")]
+    HeaderTimeout(u64),
+
     #[error("HTTP error: {0}")]
     HttpError(#[from] reqwest::Error),
 
diff --git a/crates/agent/src/llm/anthropic.rs b/crates/agent/src/llm/anthropic.rs
@@ -581,6 +581,7 @@ mod tests {
             extra_headers: vec![],
             thinking: None,
             disable_cache_control: false,
+            policy: Default::default(),
         }
     }
 
diff --git a/crates/agent/src/llm/client.rs b/crates/agent/src/llm/client.rs
@@ -50,6 +50,49 @@ static SHARED_HTTP_CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
         .expect("Failed to create HTTP client")
 });
 
+/// Per-binary behavior knobs for the LLM HTTP client. The default is **permissive**
+/// (byte-identical to the desktop app's historical behavior — pooled HTTP/2 via the
+/// shared client); `LlmPolicy::bench()` is the strict variant the headless
+/// `bench-runner` opts into so the eval harness matches opencode's working transport
+/// shape (HTTP/1.1, fresh connection per request, 15s header-arrival timeout). Gating
+/// these behind a policy keeps the shipped app client unchanged.
+#[derive(Debug, Clone)]
+pub struct LlmPolicy {
+    /// Pin the LLM client to HTTP/1.1. `false` = let reqwest negotiate (defaults to h2
+    /// when ALPN offers it).
+    pub http1_only: bool,
+    /// Disable connection pooling — every LLM request opens a fresh TCP+TLS connection.
+    /// `false` = use the shared pooled client.
+    pub no_pool: bool,
+    /// Abort the request if response headers do not arrive within this many ms. `None`
+    /// = no header-arrival timeout (only the per-chunk idle timeout applies once the
+    /// body starts).
+    pub header_timeout_ms: Option<u64>,
+}
+
+impl Default for LlmPolicy {
+    /// Permissive — preserves the historical desktop-app behavior exactly.
+    fn default() -> Self {
+        Self {
+            http1_only: false,
+            no_pool: false,
+            header_timeout_ms: None,
+        }
+    }
+}
+
+impl LlmPolicy {
+    /// Strict policy for the eval harness: matches opencode's working transport shape
+    /// to neutralize router-side mid-stream resets on long Fireworks/Kimi runs.
+    pub fn bench() -> Self {
+        Self {
+            http1_only: true,
+            no_pool: true,
+            header_timeout_ms: Some(15_000),
+        }
+    }
+}
+
 /// Configuration for the LLM HTTP client.
 #[derive(Debug, Clone)]
 pub struct LlmClientConfig {
@@ -75,6 +118,9 @@ pub struct LlmClientConfig {
     pub thinking: Option<serde_json::Value>,
     /// Skip cache_control / prompt_cache_key — set for one-shot calls like compaction where writes never read back.
     pub disable_cache_control: bool,
+    /// Per-binary LLM transport behavior. Default = pooled HTTP/2 (app); `bench()` =
+    /// HTTP/1.1, no pool, 15s header timeout (matches opencode).
+    pub policy: LlmPolicy,
 }
 
 /// HTTP client that speaks either OpenAI chat-completions or the Anthropic
@@ -86,10 +132,24 @@ pub struct LlmClient {
 
 impl LlmClient {
     pub fn new(config: LlmClientConfig) -> Self {
-        Self {
-            config,
-            http: SHARED_HTTP_CLIENT.clone(),
-        }
+        // Build a dedicated client when the policy deviates from default; otherwise
+        // share the pooled process-wide client. Bench-runner takes this path to match
+        // opencode's transport shape (h1.1, no pool) which has zero decode deaths on
+        // the same router that gave us 26.
+        let http = if config.policy.http1_only || config.policy.no_pool {
+            let mut b = reqwest::Client::builder()
+                .connect_timeout(std::time::Duration::from_secs(30));
+            if config.policy.http1_only {
+                b = b.http1_only();
+            }
+            if config.policy.no_pool {
+                b = b.pool_max_idle_per_host(0);
+            }
+            b.build().expect("Failed to build dedicated LLM HTTP client")
+        } else {
+            SHARED_HTTP_CLIENT.clone()
+        };
+        Self { config, http }
     }
 
     /// Send a streaming chat completion request and return the assembled response.
@@ -216,18 +276,12 @@ impl LlmClient {
             }
         }
 
-        let response = match self
-            .http
-            .post(&url)
-            .headers(headers)
-            .json(&request_body)
-            .send()
-            .await
-        {
+        let send_future = self.http.post(&url).headers(headers).json(&request_body).send();
+        let response = match send_with_header_timeout(send_future, self.config.policy.header_timeout_ms).await {
             Ok(resp) => resp,
             Err(e) => {
                 log::error!("[LLM] HTTP request failed: {e}");
-                return Err(e.into());
+                return Err(e);
             }
         };
 
@@ -287,18 +341,12 @@ impl LlmClient {
             }
         }
 
-        let response = match self
-            .http
-            .post(&url)
-            .headers(headers)
-            .json(&request_body)
-            .send()
-            .await
-        {
+        let send_future = self.http.post(&url).headers(headers).json(&request_body).send();
+        let response = match send_with_header_timeout(send_future, self.config.policy.header_timeout_ms).await {
             Ok(resp) => resp,
             Err(e) => {
                 log::error!("[LLM] HTTP request failed: {e}");
-                return Err(e.into());
+                return Err(e);
             }
         };
 
@@ -318,6 +366,23 @@ impl LlmClient {
     }
 }
 
+/// Race the `.send()` future against an optional header-arrival timeout. Bounds the
+/// request-issue + header-receive phase only; once headers arrive the body stream is
+/// returned and bounded separately by `CHUNK_IDLE_TIMEOUT` in the SSE reader. A
+/// `HeaderTimeout` is retryable by `call_llm_with_retry` (same arm as `HttpError`).
+async fn send_with_header_timeout<F>(fut: F, timeout_ms: Option<u64>) -> Result<reqwest::Response, AgentError>
+where
+    F: std::future::Future<Output = Result<reqwest::Response, reqwest::Error>>,
+{
+    match timeout_ms {
+        Some(ms) => match tokio::time::timeout(std::time::Duration::from_millis(ms), fut).await {
+            Ok(res) => res.map_err(AgentError::HttpError),
+            Err(_) => Err(AgentError::HeaderTimeout(ms)),
+        },
+        None => fut.await.map_err(AgentError::HttpError),
+    }
+}
+
 /// Strip orphaned tool_result messages whose originating tool_use is no longer
 /// present (e.g. lost to compaction). Shared by both provider paths.
 fn sanitize_orphaned_tool_results(messages: &[ChatMessage]) -> Vec<ChatMessage> {
diff --git a/crates/agent/src/llm/mod.rs b/crates/agent/src/llm/mod.rs
@@ -3,5 +3,5 @@ pub mod sse;
 pub mod client;
 pub mod anthropic;
 
-pub use client::{LlmClient, LlmClientConfig, LlmProvider};
+pub use client::{LlmClient, LlmClientConfig, LlmPolicy, LlmProvider};
 pub use types::{ChatMessage, ToolDefinition, ToolCall, LlmResponse, Usage, MessageContent, ContentBlock, ImageUrlContent, Provider};
diff --git a/crates/agent/src/session.rs b/crates/agent/src/session.rs
@@ -425,6 +425,7 @@ mod tests {
                 extra_headers: vec![],
                 thinking: None,
                 disable_cache_control: false,
+                policy: Default::default(),
             },
             PathBuf::from("/tmp"),
         );
diff --git a/crates/agent/tests/integration.rs b/crates/agent/tests/integration.rs
diff --git a/crates/bench-runner/src/main.rs b/crates/bench-runner/src/main.rs

Original file line number	Diff line number	Diff line change
`@@ -592,6 +592,7 @@ pub fn provider_to_llm_config(p: &ProviderConfig, model: &str) -> LlmClientConfi`
`592`	`592`	`extra_headers: vec![],`
`593`	`593`	`thinking: None,`
`594`	`594`	`disable_cache_control: false,`
	`595`	`+ policy: Default::default(),`
`595`	`596`	`}`
`596`	`597`	`}`
`597`	`598`
Original file line number	Diff line number	Diff line change
`@@ -581,6 +581,7 @@ mod tests {`
`581`	`581`	`extra_headers: vec![],`
`582`	`582`	`thinking: None,`
`583`	`583`	`disable_cache_control: false,`
	`584`	`+ policy: Default::default(),`
`584`	`585`	`}`
`585`	`586`	`}`
`586`	`587`