Merge pull request #2211 from bobleer/bob/unify-live-agent-error-retries

fix(agent): unify live model retry budget

Author: bobleer

src/apps/cli/src/dispatch/runner.rs

@@ -220,27 +220,55 @@ pub(crate) fn process_alive(pid: u32) -> bool {
         return false;
     };
     // SAFETY: signal 0 performs liveness/permission checking only.
-    if unsafe { libc::kill(pid, 0) } == 0 {
-        #[cfg(target_os = "linux")]
-        {
-            // A zombie still answers to kill(0), but it has already exited and
-            // must not be treated as an authenticated leader for escalation.
-            if let Ok(stat) = std::fs::read_to_string(format!("/proc/{pid}/stat")) {
-                if stat
-                    .rsplit_once(") ")
-                    .and_then(|(_, fields)| fields.split_whitespace().next())
-                    == Some("Z")
-                {
-                    return false;
-                }
+    if unsafe { libc::kill(pid, 0) } != 0
+        && !matches!(
+            std::io::Error::last_os_error().raw_os_error(),
+            Some(libc::EPERM)
+        )
+    {
+        return false;
+    }
+
+    #[cfg(target_os = "linux")]
+    {
+        // A zombie still answers to kill(0), but it has already exited and
+        // must not be treated as an authenticated leader for escalation.
+        if let Ok(stat) = std::fs::read_to_string(format!("/proc/{pid}/stat")) {
+            if stat
+                .rsplit_once(") ")
+                .and_then(|(_, fields)| fields.split_whitespace().next())
+                == Some("Z")
+            {
+                return false;
             }
         }
-        return true;
     }
-    matches!(
-        std::io::Error::last_os_error().raw_os_error(),
-        Some(libc::EPERM)
-    )
+
+    #[cfg(target_os = "macos")]
+    {
+        // macOS also reports zombies as present to kill(0). Query the process
+        // state before using a leader PID to authenticate SIGKILL escalation;
+        // a failed/empty query means the process disappeared during the check.
+        let output = Command::new("ps")
+            .args(["-p", &pid.to_string(), "-o", "stat="])
+            .output();
+        let Ok(output) = output else {
+            return false;
+        };
+        if !output.status.success() {
+            return false;
+        }
+        return String::from_utf8_lossy(&output.stdout)
+            .trim_start()
+            .chars()
+            .next()
+            .is_some_and(|state| state != 'Z');
+    }
+
+    #[cfg(not(target_os = "macos"))]
+    {
+        true
+    }
 }
 
 #[cfg(not(unix))]

src/apps/cli/tests/cli_command_contracts/exec_cli_contracts.rs

@@ -367,6 +367,48 @@ fn stream_json_patch_success_emits_one_success_terminal() {
     );
 }
 
+#[test]
+fn stream_json_malformed_sse_retries_then_completes() {
+    let server = MockOpenAiServer::malformed_sse_then_immediate();
+    let environment = CliTestEnvironment::new();
+    environment.configure_mock_model(server.base_url());
+    let mut command = environment.std_command();
+    command.args([
+        "exec",
+        "exercise malformed provider stream retry",
+        "--output-format",
+        "stream-json",
+    ]);
+    let output = command_output_with_timeout(&mut command, std::time::Duration::from_secs(30));
+    server.assert_chat_completion_requests(2);
+
+    let stdout = stdout(&output);
+    assert!(output.status.success(), "{}\n{stdout}", stderr(&output));
+    let events = jsonl_events(&stdout);
+    assert!(
+        events.iter().any(|value| {
+            value["event"]["type"] == "TextChunk"
+                && value["event"]["text"]
+                    .as_str()
+                    .is_some_and(|text| text.contains(STREAM_COMPLETED_MARKER))
+        }),
+        "retried model stream did not complete: {stdout}"
+    );
+    assert_eq!(
+        events
+            .iter()
+            .filter(|value| is_terminal_event(value))
+            .count(),
+        1,
+        "retried stream must emit exactly one terminal envelope: {stdout}"
+    );
+    assert_eq!(
+        events.last().expect("retried stream terminal event")["event"]["type"],
+        "DialogTurnCompleted",
+        "retried stream terminal must be last: {stdout}"
+    );
+}
+
 #[test]
 fn stream_json_provider_http_403_emits_one_error_terminal() {
     let server = MockOpenAiServer::http_403("provider authorization denied");
@@ -493,7 +535,7 @@ fn stream_json_disconnect_then_exhausted_retry_failure_emits_one_error_terminal(
         "stream-json",
     ]);
     let output = command_output_with_timeout(&mut command, std::time::Duration::from_secs(30));
-    server.assert_chat_completion_requests(11);
+    server.assert_chat_completion_requests(10);
 
     let stdout = stdout(&output);
     assert!(!output.status.success(), "{stdout}");

src/apps/cli/tests/support/mod.rs

@@ -303,6 +303,7 @@ enum MockModelResponse {
     Gated,
     Http403 { reason: String },
     DisconnectThenHttp403,
+    MalformedSseThenImmediate,
 }
 
 impl MockOpenAiServer {
@@ -324,6 +325,10 @@ impl MockOpenAiServer {
         Self::spawn(MockModelResponse::DisconnectThenHttp403)
     }
 
+    pub(crate) fn malformed_sse_then_immediate() -> Self {
+        Self::spawn(MockModelResponse::MalformedSseThenImmediate)
+    }
+
     pub(crate) fn base_url(&self) -> &str {
         &self.base_url
     }
@@ -405,11 +410,14 @@ impl MockOpenAiServer {
                             &disconnect_tx,
                         );
                         attempt += 1;
-                        if matches!(
-                            response,
-                            MockModelResponse::Http403 { .. }
-                                | MockModelResponse::DisconnectThenHttp403
-                        ) {
+                        let accepts_more_requests =
+                            matches!(
+                                response,
+                                MockModelResponse::Http403 { .. }
+                                    | MockModelResponse::DisconnectThenHttp403
+                            ) || (matches!(response, MockModelResponse::MalformedSseThenImmediate)
+                                && attempt < 2);
+                        if accepts_more_requests {
                             continue;
                         }
                         break;
@@ -473,6 +481,13 @@ fn serve_model_response(
         )
         .expect("write mock response headers");
 
+    if matches!(response, MockModelResponse::MalformedSseThenImmediate) && attempt == 0 {
+        write_chunk(stream, b"data: not-json\n\n").expect("write malformed SSE frame");
+        let _ = stream.write_all(b"0\r\n\r\n");
+        let _ = stream.flush();
+        return;
+    }
+
     write_sse_chunk(
         stream,
         &json!({

src/crates/adapters/ai-adapters/src/client.rs

@@ -19,6 +19,7 @@ use crate::trace::{
 use crate::types::ProxyConfig;
 use crate::types::*;
 use anyhow::Result;
+use bitfun_core_types::errors::{AiProviderError, ErrorCategory};
 use format::ApiFormat;
 use log::warn;
 use reqwest::Client;
@@ -223,30 +224,36 @@ impl AIClient {
         extra_body: Option<serde_json::Value>,
         trace: Option<ModelExchangeTraceConfig>,
     ) -> Result<StreamResponse> {
-        let max_tries = SEND_MESSAGE_STREAM_ATTEMPTS;
-        match ApiFormat::parse(&self.config.format)? {
-            ApiFormat::OpenAIChat => {
-                openai::chat::send_stream(self, messages, tools, extra_body, max_tries, trace).await
-            }
-            ApiFormat::OpenAIResponses => {
-                openai::responses::send_stream(self, messages, tools, extra_body, max_tries, trace)
-                    .await
-            }
-            ApiFormat::Anthropic => {
-                anthropic::request::send_stream(self, messages, tools, extra_body, max_tries, trace)
-                    .await
-            }
-            ApiFormat::Gemini => {
-                gemini::request::send_stream(self, messages, tools, extra_body, max_tries, trace)
-                    .await
-            }
-            ApiFormat::GeminiCodeAssist => {
-                gemini::code_assist::send_stream(
-                    self, messages, tools, extra_body, max_tries, trace,
-                )
-                .await
-            }
-        }
+        self.send_message_stream_with_extra_body_and_max_attempts(
+            messages,
+            tools,
+            extra_body,
+            SEND_MESSAGE_STREAM_ATTEMPTS,
+            trace,
+        )
+        .await
+    }
+
+    /// Open one model stream without an adapter-owned retry loop.
+    ///
+    /// Runtime owners with a broader attempt lifecycle use this entry point so
+    /// connection, HTTP, TTFT, parsing, and in-stream failures all consume one
+    /// shared retry budget instead of multiplying nested retry loops.
+    pub async fn send_message_stream_once(
+        &self,
+        messages: Vec<Message>,
+        tools: Option<Vec<ToolDefinition>>,
+        trace: Option<ModelExchangeTraceConfig>,
+    ) -> Result<StreamResponse> {
+        let custom_body = self.config.custom_request_body.clone();
+        self.send_message_stream_with_extra_body_and_max_attempts(
+            messages,
+            tools,
+            custom_body,
+            1,
+            trace,
+        )
+        .await
     }
 
     pub async fn send_message(
@@ -306,15 +313,33 @@ impl AIClient {
         max_attempts: usize,
     ) -> Result<GeminiResponse> {
         for attempt in 0..max_attempts {
-            let stream_response = self
+            let stream_response = match self
                 .send_message_stream_with_extra_body_and_max_attempts(
                     messages.clone(),
                     tools.clone(),
                     extra_body.clone(),
-                    max_attempts,
+                    1,
                     trace.clone(),
                 )
-                .await?;
+                .await
+            {
+                Ok(response) => response,
+                Err(error) => {
+                    if attempt == max_attempts - 1 {
+                        return Err(error);
+                    }
+                    let delay_ms = send_message_retry_delay_ms_for_error(attempt, &error);
+                    warn!(
+                        "Retrying AI stream request after error: attempt={}/{}, delay_ms={}, error={}",
+                        attempt + 1,
+                        max_attempts,
+                        delay_ms,
+                        error
+                    );
+                    tokio::time::sleep(Duration::from_millis(delay_ms)).await;
+                    continue;
+                }
+            };
             let trace_handle = stream_response.trace_handle.clone();
 
             match response_aggregator::aggregate_stream_response(stream_response).await {
@@ -333,7 +358,7 @@ impl AIClient {
                     if attempt == max_attempts - 1 {
                         return Err(error);
                     }
-                    let delay_ms = send_message_retry_delay_ms(attempt, &error.to_string());
+                    let delay_ms = send_message_retry_delay_ms_for_error(attempt, &error);
                     warn!(
                         "Retrying aggregated AI stream after error: attempt={}/{}, delay_ms={}, error={}",
                         attempt + 1,
@@ -419,22 +444,52 @@ impl AIClient {
     }
 }
 
+#[cfg(test)]
 fn send_message_retry_delay_ms(attempt_index: usize, error_message: &str) -> u64 {
+    send_message_retry_delay_ms_with_provider(attempt_index, error_message, None)
+}
+
+fn send_message_retry_delay_ms_for_error(attempt_index: usize, error: &anyhow::Error) -> u64 {
+    send_message_retry_delay_ms_with_provider(
+        attempt_index,
+        &error.to_string(),
+        error.downcast_ref::<AiProviderError>(),
+    )
+}
+
+fn send_message_retry_delay_ms_with_provider(
+    attempt_index: usize,
+    error_message: &str,
+    provider_error: Option<&AiProviderError>,
+) -> u64 {
     let shift = u32::try_from(attempt_index)
         .unwrap_or(u32::MAX)
         .min(SEND_MESSAGE_MAX_RETRY_EXPONENT_SHIFT);
     let msg = error_message.to_lowercase();
-    let is_rate_limit =
-        msg.contains("429") || msg.contains("rate limit") || msg.contains("too many requests");
+    let is_rate_limit = provider_error
+        .is_some_and(|error| error.category == ErrorCategory::RateLimit)
+        || msg.contains("429")
+        || msg.contains("rate limit")
+        || msg.contains("too many requests");
 
-    if is_rate_limit {
+    let fallback = if is_rate_limit {
         SEND_MESSAGE_RATE_LIMIT_RETRY_BASE_DELAY_MS
             .saturating_mul(1u64 << shift)
             .min(SEND_MESSAGE_MAX_RATE_LIMIT_DELAY_MS)
     } else {
         SEND_MESSAGE_RETRY_BASE_DELAY_MS
             .saturating_mul(1u64 << shift)
             .min(SEND_MESSAGE_MAX_EXPONENTIAL_DELAY_MS)
+    };
+
+    match provider_error.and_then(|error| error.retry_after_ms) {
+        Some(retry_after_ms) if is_rate_limit => retry_after_ms
+            .max(fallback)
+            .min(SEND_MESSAGE_MAX_RATE_LIMIT_DELAY_MS),
+        Some(retry_after_ms) if retry_after_ms > 0 => {
+            retry_after_ms.min(SEND_MESSAGE_MAX_RATE_LIMIT_DELAY_MS)
+        }
+        Some(_) | None => fallback,
     }
 }