fix(llm-client): fall back on retryable pre-answer stream errors

crates/libsy-llm-client/src/client.rs

@@ -6,7 +6,6 @@
 
 use std::borrow::Cow;
 use std::collections::{BTreeMap, BTreeSet, HashMap};
-use std::future::ready;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{Duration, SystemTime};
 
@@ -563,7 +562,8 @@ impl TranslatingLlmClient {
         if status.is_success() {
             if streaming {
                 let upstream_headers = response.headers().clone();
-                let chunks = match prepare_response_stream(response, backend, model).await {
+                let chunks = match prepare_response_stream(response, backend, metadata, model).await
+                {
                     Ok(chunks) => chunks,
                     Err(error) => {
                         metrics::record_upstream_attempt(None);
@@ -913,9 +913,11 @@ fn stream_with_deadline(
     .boxed()
 }
 
+/// Check initial stream events before returning them, while fallback is still possible.
 async fn prepare_response_stream(
     response: reqwest::Response,
     backend: &Backend,
+    metadata: Option<&Metadata>,
     model: &ModelId,
 ) -> Result<LlmResponseStream> {
     let bytes = response.bytes_stream().map(|chunk| {
@@ -924,42 +926,73 @@ async fn prepare_response_stream(
             .map_err(convert_reqwest_error)
     });
     let mut chunks = decode_stream(bytes, backend.wire_format())?;
-    match chunks.next().await {
-        None => Ok(stream::empty().boxed()),
-        // Nothing has reached the caller, so transport failures can still be retried.
-        Some(Err(error @ (LlmClientError::Transport { .. } | LlmClientError::Timeout { .. }))) => {
-            Err(error)
-        }
-        Some(first) => {
-            // An in-band context overflow skips retries and advances to another candidate.
-            if let Some(message) = first_event_overflow(&first, backend) {
-                return Err(LlmClientError::ContextWindowExceeded {
-                    model: model.clone(),
-                    message,
-                });
+    let mut prefix = Vec::with_capacity(2);
+    while let Some(next) = chunks.next().await {
+        match next {
+            // Nothing has reached the caller, so transport failures can still be retried.
+            Err(error @ (LlmClientError::Transport { .. } | LlmClientError::Timeout { .. })) => {
+                return Err(error);
+            }
+            Ok(event) => {
+                if let Some(error) = pre_answer_error(&event, backend, metadata, model) {
+                    return Err(error);
+                }
+                // OpenAI streams often send a role-only prelude before the answer.
+                let hold_role = prefix.is_empty()
+                    && matches!(backend, Backend::OpenAiChat(_))
+                    && !event.normalized().is_empty()
+                    && event
+                        .normalized()
+                        .iter()
+                        .all(|chunk| matches!(chunk, LlmResponseChunk::MessageStart { .. }));
+                prefix.push(Ok(event));
+                if !hold_role {
+                    break;
+                }
+            }
+            Err(error) => {
+                prefix.push(Err(error));
+                break;
             }
-            Ok(stream::once(ready(first)).chain(chunks).boxed())
         }
     }
+    Ok(stream::iter(prefix).chain(chunks).boxed())
 }
 
-// The overflow message when a stream's first event is an in-band provider rejection
-// of the whole request, rather than the start of a response.
-fn first_event_overflow(
-    first: &Result<LlmResponseStreamEvent>,
+/// Classify a provider rejection before any stream event reaches the caller.
+fn pre_answer_error(
+    event: &LlmResponseStreamEvent,
     backend: &Backend,
-) -> Option<String> {
-    first
-        .as_ref()
-        .ok()?
-        .normalized()
-        .iter()
-        .find_map(|chunk| match chunk {
-            LlmResponseChunk::StreamError { message } if backend.is_context_overflow(message) => {
-                Some(message.clone())
-            }
-            _ => None,
-        })
+    metadata: Option<&Metadata>,
+    model: &ModelId,
+) -> Option<LlmClientError> {
+    let message = event.normalized().iter().find_map(|chunk| match chunk {
+        LlmResponseChunk::StreamError { message } => Some(message),
+        _ => None,
+    })?;
+    if backend.is_context_overflow(message) {
+        return Some(LlmClientError::ContextWindowExceeded {
+            model: model.clone(),
+            message: message.clone(),
+        });
+    }
+    if !matches!(backend, Backend::OpenAiChat(_)) {
+        return None;
+    }
+    let code = event
+        .preservation()?
+        .raw()
+        .get("error")?
+        .get("code")?
+        .as_u64()
+        .and_then(|code| u16::try_from(code).ok())?;
+    if !metrics::is_retryable_http_status(code) {
+        return None;
+    }
+    Some(LlmClientError::UpstreamHttp {
+        status: StatusCode::from_u16(code).ok()?,
+        body: redact_forwarded_headers(message.clone(), metadata, backend.is_forwarding_auth()),
+    })
 }
 
 // Uses Retry-After when supplied, capped so an upstream cannot stall a request indefinitely.
@@ -2862,6 +2895,48 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn retry_exhaustion_preserves_stream_error_message_without_forwarded_credentials()
+    -> Result<()> {
+        let server = MockServer::start().await;
+        let sse = "data: {\"error\":{\"code\":503,\"message\":\"provider overloaded for Bearer client-key\"}}\n\n";
+        Mock::given(method("POST"))
+            .and(path("/v1/chat/completions"))
+            .and(wiremock::matchers::header(
+                "authorization",
+                "Bearer client-key",
+            ))
+            .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
+            .expect(2)
+            .mount(&server)
+            .await;
+
+        let mut backend_config = forwarding_config(&format!("{}/v1", server.uri()));
+        backend_config.max_retries = 1;
+        let client = TranslatingLlmClient::new(&[ModelConfig::new(
+            "gpt",
+            Backend::OpenAiChat(backend_config),
+            None,
+        )])?;
+        let mut headers = HeaderMap::new();
+        headers.insert(
+            "authorization",
+            http::HeaderValue::from_static("Bearer client-key"),
+        );
+        let mut request = request_with_headers("gpt", headers);
+        request.llm_request.stream = true;
+
+        let Err(error) = client.call_rewrite_model(request, None).await else {
+            panic!("expected an upstream stream error");
+        };
+        let LlmClientError::UpstreamHttp { status, body } = error else {
+            panic!("expected the upstream stream error, got {error:?}");
+        };
+        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
+        assert_eq!(body, "provider overloaded for [REDACTED]");
+        Ok(())
+    }
+
     #[tokio::test]
     async fn streaming_openai_chat_preserves_usage_opt_out()
     -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {

crates/libsy-llm-client/src/run.rs

@@ -2928,10 +2928,9 @@ mod tests {
             .map_err(|source| LibsyError::client_call("strong", source))?;
         assert_eq!(completion_text(&aggregate), "ok");
 
-        // Only a classified overflow re-routes; any other in-band error event still
-        // streams through to the consumer without trying another candidate.
+        // A non-retryable in-band error still reaches the caller.
         let (_server, calls, result) = run_streaming_candidates(
-            "data: {\"error\":{\"code\":500,\"message\":\"engine exploded\",\"object\":\"error\"},\"model\":\"weak\"}\n\ndata: [DONE]\n\n",
+            "data: {\"error\":{\"code\":401,\"message\":\"invalid key\",\"object\":\"error\"},\"model\":\"weak\"}\n\ndata: [DONE]\n\n",
         )
         .await;
         let (_, response) = result?;
@@ -2941,7 +2940,7 @@ mod tests {
             .into_agg()
             .await
             .expect_err("the in-band error should surface to the consumer");
-        assert!(error.to_string().contains("engine exploded"));
+        assert!(error.to_string().contains("invalid key"));
 
         // Once content has streamed, a retry would repeat delivered output, so a later
         // overflow event keeps today's semantics and surfaces during aggregation.
@@ -2975,4 +2974,78 @@ mod tests {
         assert_eq!(completion_text(&aggregate), "");
         Ok(())
     }
+
+    /// A role-only prelude leaves fallback available for a retryable provider error.
+    #[tokio::test]
+    async fn pre_answer_retryable_stream_errors_fall_back_to_the_next_candidate() -> Result<()> {
+        for weak_sse in [
+            "data: {\"error\":{\"code\":503,\"message\":\"provider overloaded\"}}\n\n",
+            "data: {\"error\":{\"code\":429,\"message\":\"rate limited\"}}\n\n",
+            "data: {\"id\":\"turn\",\"model\":\"weak\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\ndata: {\"error\":{\"code\":503,\"message\":\"provider overloaded\"}}\n\n",
+        ] {
+            let (_server, calls, result) = run_streaming_candidates(weak_sse).await;
+            let (_, response) = result?;
+            assert_eq!(&*calls.lock(), &["weak", "strong"]);
+            assert_eq!(response.served_model().map(ModelId::as_str), Some("strong"));
+            let aggregate = response
+                .llm_response
+                .into_agg()
+                .await
+                .map_err(|source| LibsyError::client_call("strong", source))?;
+            assert_eq!(completion_text(&aggregate), "ok");
+        }
+        Ok(())
+    }
+
+    /// Once answer content arrives, later errors must not replay the request on another model.
+    #[tokio::test]
+    async fn stream_fallback_stops_at_answer_content_and_non_retryable_errors() -> Result<()> {
+        for weak_sse in [
+            "data: {\"error\":{\"code\":400,\"message\":\"invalid request\"}}\n\n",
+            "data: {\"error\":{\"code\":\"busy\",\"message\":\"not an HTTP status\"}}\n\n",
+            "data: {\"id\":\"turn\",\"model\":\"weak\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"partial\"},\"finish_reason\":null}]}\n\ndata: {\"error\":{\"code\":503,\"message\":\"provider overloaded\"}}\n\n",
+        ] {
+            let (_server, calls, result) = run_streaming_candidates(weak_sse).await;
+            let (_, response) = result?;
+            assert_eq!(&*calls.lock(), &["weak"]);
+            response
+                .llm_response
+                .into_agg()
+                .await
+                .expect_err("the stream error should reach the caller");
+        }
+
+        let (_server, calls, result) = run_streaming_candidates(STRONG_SSE).await;
+        let (_, response) = result?;
+        assert_eq!(&*calls.lock(), &["weak"]);
+        let aggregate = response
+            .llm_response
+            .into_agg()
+            .await
+            .map_err(|source| LibsyError::client_call("weak", source))?;
+        assert_eq!(completion_text(&aggregate), "ok");
+
+        let (_server, calls, result) = run_streaming_candidates(
+            "data: {\"id\":\"turn\",\"model\":\"weak\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\ndata: {\"id\":\"turn\",\"model\":\"weak\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n",
+        )
+        .await;
+        let (_, response) = result?;
+        assert_eq!(&*calls.lock(), &["weak"]);
+        let LlmResponse::Stream(stream) = response.llm_response else {
+            panic!("expected a stream");
+        };
+        let events = stream.collect::<Vec<_>>().await;
+        assert_eq!(events.len(), 2);
+        assert!(matches!(
+            events[0].as_ref().expect("role event").normalized(),
+            [LlmResponseChunk::MessageStart { .. }]
+        ));
+        assert!(events[1]
+            .as_ref()
+            .expect("answer event")
+            .normalized()
+            .iter()
+            .any(|chunk| matches!(chunk, LlmResponseChunk::TextDelta { text, .. } if text == "ok")));
+        Ok(())
+    }
 }