From b70a351b2e467175731947c5b08a4527e1c6f171 Mon Sep 17 00:00:00 2001 From: Chris Liu <11517021+chris7iu@users.noreply.github.com> Date: Thu, 12 Nov 2026 07:44:46 +0000 Subject: [PATCH] fix(stream): retry pre-output TLS close_notify EOF and abort thinking-only streams on transient drops --- src/providers/codex/client.rs | 22 ++ src/providers/codex/events.rs | 22 +- src/providers/codex/mod.rs | 266 ++++++++++++++++++- src/providers/codex/translate/live_stream.rs | 17 ++ src/providers/cursor/mod.rs | 25 +- src/providers/cursor/sse.rs | 9 + src/providers/cursor/tool_bridge.rs | 8 + src/providers/grok/mod.rs | 106 +++++++- src/providers/grok/translate/stream.rs | 17 ++ src/providers/opencode/chat.rs | 51 +++- src/providers/opencode/messages.rs | 84 +++++- src/providers/opencode/responses.rs | 60 ++++- 12 files changed, 614 insertions(+), 73 deletions(-) diff --git a/src/providers/codex/client.rs b/src/providers/codex/client.rs index a6c98d90..c9d95b77 100644 --- a/src/providers/codex/client.rs +++ b/src/providers/codex/client.rs @@ -3046,6 +3046,10 @@ fn is_retryable_transport_error(err: &CodexError) -> bool { || message.contains("transport error") || message.contains("connection reset") || message.contains("connection closed") + || message.contains("peer closed connection") + || message.contains("close_notify") + || message.contains("unexpected eof") + || message.contains("unexpected-eof") || message.contains("timed out") || message.contains("econnreset") || message.contains("etimedout") @@ -3060,6 +3064,10 @@ fn is_retryable_reqwest_error(err: &reqwest::Error) -> bool { let msg = err.to_string().to_lowercase(); msg.contains("connection reset") || msg.contains("connection closed") + || msg.contains("peer closed connection") + || msg.contains("close_notify") + || msg.contains("unexpected eof") + || msg.contains("unexpected-eof") || msg.contains("econnreset") || msg.contains("etimedout") || msg.contains("epipe") @@ -5692,6 +5700,20 @@ mod tests { assert!(!should_retry_without_continuation(&timeout, None)); } + #[test] + fn statusless_tls_close_notify_unexpected_eof_is_retryable() { + let err = CodexError { + status: 0, + message: "WebSocket stream error: IO error: peer closed connection without sending TLS close_notify: https://docs.rs/rustls/latest/rustls/manual/_03_howto/index.html#unexpected-eof".to_string(), + detail: None, + retry_after: None, + usage_limit: None, + origin: CodexErrorOrigin::WebSocket, + }; + assert!(is_retryable_transport_error(&err)); + assert!(super::super::events::retryable_message(&err.message)); + } + #[tokio::test] async fn http_request_compression_gzip_encodes_body_and_preserves_traffic_capture() { use std::io::Read; diff --git a/src/providers/codex/events.rs b/src/providers/codex/events.rs index 1c63740a..1f35cbc0 100644 --- a/src/providers/codex/events.rs +++ b/src/providers/codex/events.rs @@ -189,17 +189,16 @@ pub(crate) fn classify_event_failure(payload: &Value) -> Option Option 400, + _ if policy_violation => 400, Some("insufficient_quota") => 429, Some("usage_not_included") => 403, _ => match kind { @@ -549,7 +548,7 @@ fn scalar_string_value(value: &Value) -> Option { } } -fn retryable_message(message: &str) -> bool { +pub(super) fn retryable_message(message: &str) -> bool { [ "server error", "internal server error", @@ -570,7 +569,10 @@ fn retryable_message(message: &str) -> bool { "etimedout", "und_err_socket", "fetch failed", + "peer closed connection", + "close_notify", "unexpected eof", + "unexpected-eof", ] .iter() .any(|needle| message.contains(needle)) diff --git a/src/providers/codex/mod.rs b/src/providers/codex/mod.rs index 4e0f9d12..a9f19303 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -1331,6 +1331,22 @@ fn remaining_live_stream_response( let error_message = failure .as_ref() .map_or(message.as_str(), |failure| failure.message.as_str()); + if !translator.has_non_thinking_output() + && retryable_mid_thinking_event_failure(&payload, failure.as_ref()) + { + log_mid_thinking_stream_abort( + &ctx.req_id, + &request_body.model, + error_message, + ); + let _ = tx + .send(Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + error_message.to_string(), + ))) + .await; + return; + } let chunk = translator.error_chunk( error_message, error_type, @@ -1382,6 +1398,19 @@ fn remaining_live_stream_response( return; } } + if !translator.has_non_thinking_output() + && retryable_mid_thinking_codex_error(&err) + { + let reason = codex_error_message(&err); + log_mid_thinking_stream_abort(&ctx.req_id, &request_body.model, reason); + let _ = tx + .send(Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + reason.to_string(), + ))) + .await; + return; + } let error_type = codex_stream_error_type(&err); let chunk = translator.error_chunk( codex_error_message(&err), @@ -1408,11 +1437,18 @@ fn remaining_live_stream_response( let _ = tx.send(Ok(Bytes::from(chunk))).await; return; } - let chunk = translator.error_chunk( - "Upstream event stream closed before terminal Codex response event", - "api_error", - ctx.traffic.as_deref(), - ); + let reason = "Upstream event stream closed before terminal Codex response event"; + if !translator.has_non_thinking_output() { + log_mid_thinking_stream_abort(&ctx.req_id, &request_body.model, reason); + let _ = tx + .send(Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + reason, + ))) + .await; + return; + } + let chunk = translator.error_chunk(reason, "api_error", ctx.traffic.as_deref()); if !chunk.is_empty() { record_live_stream_progress(&ctx, &chunk); let _ = tx.send(Ok(Bytes::from(chunk))).await; @@ -1561,6 +1597,45 @@ fn drop_live_continuation_for_retry(continuation: &mut Option bool { + if err.usage_limit.is_some() { + return false; + } + if matches!( + err.detail.as_deref(), + Some(websocket::WEBSOCKET_KEEPALIVE_FAILURE_DETAIL) + | Some(websocket::WEBSOCKET_MISSING_TERMINAL_DETAIL) + ) { + return true; + } + matches!(err.status, 500 | 502 | 503 | 504 | 529) + || (err.status == 0 && retryable_live_message(codex_error_message(err))) +} + +fn retryable_mid_thinking_event_failure( + payload: &serde_json::Value, + failure: Option<&events::CodexEventFailure>, +) -> bool { + if events::usage_limit_from_event(payload).is_some() { + return false; + } + matches!( + failure.map(|failure| failure.kind), + Some(events::CodexFailureKind::Transient | events::CodexFailureKind::Overloaded) + ) +} + fn retryable_live_message(message: &str) -> bool { let lower = message.to_ascii_lowercase(); [ @@ -1569,8 +1644,17 @@ fn retryable_live_message(message: &str) -> bool { "you can retry your request", "temporarily unavailable", "timed out", + "idle timeout", + "keepalive error", + "transport error", "connection closed", + "peer closed connection", + "close_notify", + "unexpected eof", + "unexpected-eof", "connection reset", + "econnreset", + "etimedout", "broken pipe", "epipe", ] @@ -3355,4 +3439,176 @@ mod tests { )); websocket::invalidate_codex_websocket_pool_owner(&owner); } + + #[test] + fn live_start_tls_close_notify_unexpected_eof_is_retryable() { + let err = client::CodexError { + status: 0, + message: "WebSocket stream error: IO error: peer closed connection without sending TLS close_notify: https://docs.rs/rustls/latest/rustls/manual/_03_howto/index.html#unexpected-eof".to_string(), + detail: None, + retry_after: None, + usage_limit: None, + origin: client::CodexErrorOrigin::WebSocket, + }; + assert!(retryable_live_start_codex_error(&err)); + assert!(retryable_mid_thinking_codex_error(&err)); + } + + #[tokio::test] + async fn live_stream_aborts_body_on_mid_thinking_transport_drop_but_preserves_sse_error_after_text_or_policy() + { + use http_body_util::BodyExt as _; + + // 1. Mid-thinking TLS close_notify drop -> downstream HTTP body aborts with Err(io::Error). + { + let (tx, rx, ctx, request_body) = empty_terminal_fixture(); + tx.send(Ok(serde_json::json!({ + "type": "response.reasoning_summary_text.delta", + "output_index": 0, + "delta": "Inspecting files..." + }))) + .await + .unwrap(); + tx.send(Err(client::CodexError { + status: 0, + message: "WebSocket stream error: IO error: peer closed connection without sending TLS close_notify: https://docs.rs/rustls/latest/rustls/manual/_03_howto/index.html#unexpected-eof".to_string(), + detail: None, + retry_after: None, + usage_limit: None, + origin: client::CodexErrorOrigin::WebSocket, + })) + .await + .unwrap(); + drop(tx); + + let continuation = ContinuationReservation::for_owner_turn(None, None); + let response = match live_stream_response_once( + rx, + "msg_thinking_abort".to_string(), + "claude-opus-4-8", + ctx, + continuation, + request_body, + LiveStreamCompaction { + compact_boundary: false, + attempt: None, + }, + config::CodexTransport::WebSocket, + false, + ) + .await + { + LiveStreamStart::Response(response) => response, + LiveStreamStart::Retry { error, .. } => { + panic!("expected Response after thinking delta, got Retry: {error}") + } + }; + let result = response.into_body().collect().await; + assert!( + result.is_err(), + "expected mid-thinking transport drop to abort body with io::Error" + ); + } + + // 2. Post-text TLS close_notify drop -> preserves event: error + clean EOF so visible text isn't duplicated. + { + let (tx, rx, ctx, request_body) = empty_terminal_fixture(); + tx.send(Ok(serde_json::json!({ + "type": "response.output_text.delta", + "output_index": 0, + "delta": "Hello world" + }))) + .await + .unwrap(); + tx.send(Err(client::CodexError { + status: 0, + message: "WebSocket stream error: IO error: peer closed connection without sending TLS close_notify".to_string(), + detail: None, + retry_after: None, + usage_limit: None, + origin: client::CodexErrorOrigin::WebSocket, + })) + .await + .unwrap(); + drop(tx); + + let continuation = ContinuationReservation::for_owner_turn(None, None); + let response = match live_stream_response_once( + rx, + "msg_text_error".to_string(), + "claude-opus-4-8", + ctx, + continuation, + request_body, + LiveStreamCompaction { + compact_boundary: false, + attempt: None, + }, + config::CodexTransport::WebSocket, + false, + ) + .await + { + LiveStreamStart::Response(response) => response, + LiveStreamStart::Retry { error, .. } => { + panic!("expected Response after text delta, got Retry: {error}") + } + }; + let bytes = response.into_body().collect().await.unwrap().to_bytes(); + let sse = String::from_utf8_lossy(&bytes); + assert!(sse.contains("Hello world")); + assert!(sse.contains("event: error")); + } + + // 3. Permanent policy block during thinking -> preserves event: error + clean EOF (do NOT retry deterministic error). + { + let (tx, rx, ctx, request_body) = empty_terminal_fixture(); + tx.send(Ok(serde_json::json!({ + "type": "response.reasoning_summary_text.delta", + "output_index": 0, + "delta": "Thinking..." + }))) + .await + .unwrap(); + tx.send(Ok(serde_json::json!({ + "type": "response.failed", + "response": { + "status": "failed", + "error": { + "code": "cyber_policy_violation", + "message": "blocked by policy" + } + } + }))) + .await + .unwrap(); + drop(tx); + + let continuation = ContinuationReservation::for_owner_turn(None, None); + let response = match live_stream_response_once( + rx, + "msg_policy_error".to_string(), + "claude-opus-4-8", + ctx, + continuation, + request_body, + LiveStreamCompaction { + compact_boundary: false, + attempt: None, + }, + config::CodexTransport::WebSocket, + false, + ) + .await + { + LiveStreamStart::Response(response) => response, + LiveStreamStart::Retry { error, .. } => { + panic!("expected Response after thinking delta, got Retry: {error}") + } + }; + let bytes = response.into_body().collect().await.unwrap().to_bytes(); + let sse = String::from_utf8_lossy(&bytes); + assert!(sse.contains("event: error")); + } + } } diff --git a/src/providers/codex/translate/live_stream.rs b/src/providers/codex/translate/live_stream.rs index 74f95081..c882417f 100644 --- a/src/providers/codex/translate/live_stream.rs +++ b/src/providers/codex/translate/live_stream.rs @@ -68,6 +68,7 @@ pub struct LiveStreamTranslator { web_search_results: Vec, deferred_text: Vec<(usize, String)>, semantic_output_started: bool, + non_thinking_output_emitted: bool, // Seeds Claude Code's live subagent counter until the provider returns // authoritative usage in the terminal message_delta. estimated_input_tokens: u64, @@ -103,6 +104,7 @@ impl LiveStreamTranslator { web_search_results: Vec::new(), deferred_text: Vec::new(), semantic_output_started: false, + non_thinking_output_emitted: false, estimated_input_tokens, incomplete_response_policy: IncompleteResponsePolicy::Error, finished: false, @@ -204,6 +206,10 @@ impl LiveStreamTranslator { self.semantic_output_started } + pub fn has_non_thinking_output(&self) -> bool { + self.non_thinking_output_emitted + } + pub fn ping_chunk(&mut self, traffic: Option<&TrafficCapture>) -> Vec { let mut out = Vec::new(); if !self.finished { @@ -343,6 +349,7 @@ impl LiveStreamTranslator { }, ); if !deferred { + self.non_thinking_output_emitted = true; self.ensure_message_start(traffic, out); self.emit( traffic, @@ -359,6 +366,7 @@ impl LiveStreamTranslator { "function_call" => { self.close_thinking(traffic, out); self.saw_tool_use = true; + self.non_thinking_output_emitted = true; let index = self.anthropic_index; self.anthropic_index += 1; let call_id = item @@ -493,6 +501,7 @@ impl LiveStreamTranslator { }, ); if !deferred { + self.non_thinking_output_emitted = true; self.ensure_message_start(traffic, out); self.emit( traffic, @@ -520,6 +529,7 @@ impl LiveStreamTranslator { return; } let index = *index; + self.non_thinking_output_emitted = true; self.emit( traffic, out, @@ -583,6 +593,7 @@ impl LiveStreamTranslator { *emitted_args = true; let index = *index; self.semantic_output_started = true; + self.non_thinking_output_emitted = true; self.emit( traffic, out, @@ -601,6 +612,7 @@ impl LiveStreamTranslator { if let Some((index, repaired)) = repaired_read { self.blocks_by_output_index.remove(&output_index); self.semantic_output_started = true; + self.non_thinking_output_emitted = true; self.emit( traffic, out, @@ -714,6 +726,7 @@ impl LiveStreamTranslator { if *deferred { self.deferred_text.push((*index, std::mem::take(text))); } else { + self.non_thinking_output_emitted = true; self.emit( traffic, out, @@ -736,6 +749,7 @@ impl LiveStreamTranslator { .. } => { self.semantic_output_started = true; + self.non_thinking_output_emitted = true; if let Some(final_args) = payload .get("item") .and_then(|item| item.get("arguments")) @@ -808,6 +822,7 @@ impl LiveStreamTranslator { fn emit_web_searches(&mut self, traffic: Option<&TrafficCapture>, out: &mut Vec) { let searches = std::mem::take(&mut self.web_searches); for search in searches { + self.non_thinking_output_emitted = true; self.ensure_message_start(traffic, out); self.emit( traffic, @@ -877,6 +892,7 @@ impl LiveStreamTranslator { } for (index, text) in std::mem::take(&mut self.deferred_text) { + self.non_thinking_output_emitted = true; self.emit( traffic, out, @@ -942,6 +958,7 @@ impl LiveStreamTranslator { traffic: Option<&TrafficCapture>, out: &mut Vec, ) { + self.non_thinking_output_emitted = true; let mapped = map_codex_usage_to_anthropic(&usage, Some(self.web_search_requests)); self.emit( traffic, diff --git a/src/providers/cursor/mod.rs b/src/providers/cursor/mod.rs index e8b4e208..a7703f9b 100644 --- a/src/providers/cursor/mod.rs +++ b/src/providers/cursor/mod.rs @@ -383,7 +383,7 @@ fn live_sse_response( model: &str, ctx: &RequestContext, ) -> Response { - let (tx, rx) = tokio::sync::mpsc::unbounded_channel::(); + let (tx, rx) = tokio::sync::mpsc::unbounded_channel::>(); let message_id = message_id.to_string(); let model = model.to_string(); let monitor = ctx.monitor.clone(); @@ -400,11 +400,11 @@ fn live_sse_response( } let outcome = match send_failed { - true => PumpOutcome::Failed("Cursor run already closed".to_string()), + true => PumpOutcome::Aborted("Cursor run already closed".to_string()), false => { pump_live_bridge(&mut bridge, &message_id, &model, &mut |chunk| { sse.extend_from_slice(chunk); - let _ = tx.send(Bytes::copy_from_slice(chunk)); + let _ = tx.send(Ok(Bytes::copy_from_slice(chunk))); }) .await } @@ -415,6 +415,13 @@ fn live_sse_response( bridge.pending = Some(pending); BridgeRegistry::insert(bridge); } + PumpOutcome::Aborted(error) => { + drop(bridge); + let _ = tx.send(Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + error, + ))); + } // Nothing left to resume: the run is dropped, which closes its // stream upstream. PumpOutcome::Replay | PumpOutcome::Ended | PumpOutcome::Failed(_) => drop(bridge), @@ -434,9 +441,7 @@ fn live_sse_response( }); let body = Body::from_stream(futures_util::stream::unfold(rx, |mut rx| async move { - rx.recv() - .await - .map(|chunk| (Ok::(chunk), rx)) + rx.recv().await.map(|item| (item, rx)) })); let headers = [ (http::header::CONTENT_TYPE, "text/event-stream"), @@ -454,7 +459,10 @@ fn finish_buffered_run(bridge: LiveBridge, outcome: PumpOutcome) { bridge.pending = Some(pending); BridgeRegistry::insert(bridge); } - PumpOutcome::Replay | PumpOutcome::Ended | PumpOutcome::Failed(_) => drop(bridge), + PumpOutcome::Replay + | PumpOutcome::Ended + | PumpOutcome::Aborted(_) + | PumpOutcome::Failed(_) => drop(bridge), } } @@ -478,6 +486,9 @@ async fn pump_buffered( sse.extend_from_slice(chunk) }) .await; + if let PumpOutcome::Aborted(ref error) = outcome { + sse.extend_from_slice(&self::sse::format_sse_error(error)); + } (sse, outcome) } diff --git a/src/providers/cursor/sse.rs b/src/providers/cursor/sse.rs index df3e409f..1f16ca90 100644 --- a/src/providers/cursor/sse.rs +++ b/src/providers/cursor/sse.rs @@ -107,6 +107,7 @@ pub struct CursorSseFramer<'a> { started: bool, thinking_open: bool, text_open: bool, + non_thinking_emitted: bool, next_index: i32, thinking_index: i32, text_index: i32, @@ -129,6 +130,7 @@ impl<'a> CursorSseFramer<'a> { started: false, thinking_open: false, text_open: false, + non_thinking_emitted: false, next_index: 0, thinking_index: -1, text_index: -1, @@ -146,6 +148,10 @@ impl<'a> CursorSseFramer<'a> { std::mem::take(&mut self.output) } + pub fn has_non_thinking_output(&self) -> bool { + self.non_thinking_emitted + } + pub fn ensure_start(&mut self) { if self.started { return; @@ -202,6 +208,7 @@ impl<'a> CursorSseFramer<'a> { } self.ensure_start(); self.text_open = true; + self.non_thinking_emitted = true; self.text_index = self.next_index; self.next_index += 1; @@ -320,6 +327,7 @@ impl<'a> CursorSseFramer<'a> { /// Emit a tool_use pause: content block + message_delta with /// stop_reason="tool_use" + message_stop. pub fn emit_tool_pause(&mut self, tool_use_id: &str, tool_name: &str, partial_json: &str) { + self.non_thinking_emitted = true; self.close_open_blocks(); self.ensure_start(); let index = self.next_content_block_index(); @@ -362,6 +370,7 @@ impl<'a> CursorSseFramer<'a> { if self.finalized { return; } + self.non_thinking_emitted = true; self.ensure_start(); self.close_open_blocks(); diff --git a/src/providers/cursor/tool_bridge.rs b/src/providers/cursor/tool_bridge.rs index a74775f7..e62dcd8d 100644 --- a/src/providers/cursor/tool_bridge.rs +++ b/src/providers/cursor/tool_bridge.rs @@ -72,6 +72,8 @@ pub enum PumpOutcome { Replay, /// The turn ended. Ended, + /// Upstream transport/idle stream failed before any non-thinking block was emitted. + Aborted(String), Failed(String), } @@ -367,6 +369,12 @@ pub async fn pump_live_bridge( return PumpOutcome::Ended; } LiveEvent::Error(error) => { + if !framer.has_non_thinking_output() + && (error.starts_with("read body:") + || error == "Cursor upstream went quiet with the run open") + { + return PumpOutcome::Aborted(error); + } framer.finalize(); sink(&framer.take_output()); sink(&format_sse_error(&error)); diff --git a/src/providers/grok/mod.rs b/src/providers/grok/mod.rs index 8f63b2c4..0c6b4a6d 100644 --- a/src/providers/grok/mod.rs +++ b/src/providers/grok/mod.rs @@ -3,7 +3,7 @@ pub mod client; pub mod count_tokens; pub mod translate; -use std::convert::Infallible; +use std::io; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; @@ -288,9 +288,9 @@ where }; let stream = futures_util::stream::unfold(state, |mut state| async move { state - .next_output() + .next_item() .await - .map(|bytes| (Ok::(Bytes::from(bytes)), state)) + .map(|item| (item.map(Bytes::from), state)) }); ( [ @@ -321,7 +321,7 @@ impl GrokStreamState where S: Stream> + Unpin, { - async fn next_output(&mut self) -> Option> { + async fn next_item(&mut self) -> Option, io::Error>> { if self.terminal { return None; } @@ -332,10 +332,15 @@ where loop { let chunk = match self.upstream.next().await { Some(Ok(chunk)) => chunk, - Some(Err(_)) => return Some(self.fail_at("transport", "upstream_stream")), + Some(Err(err)) => { + if !self.translator.has_non_thinking_output() { + return Some(Err(self.abort_transport(&err.message))); + } + return Some(Ok(self.fail_at("transport", "upstream_stream"))); + } None => { if self.decoder.finish().is_err() || !self.reducer.finished() { - return Some(self.fail_at("decoder", "incomplete_stream")); + return Some(Ok(self.fail_at("decoder", "incomplete_stream"))); } self.terminal = true; self.finish_capture(true); @@ -354,7 +359,7 @@ where } let events = match self.decoder.push(&chunk) { Ok(events) => events, - Err(_) => return Some(self.fail_at("decoder", "malformed_sse")), + Err(_) => return Some(Ok(self.fail_at("decoder", "malformed_sse"))), }; let mut out = Vec::new(); for event in events { @@ -364,7 +369,7 @@ where if let Some(capture) = self.stream_capture.as_mut() { capture.malformed("json", "malformed_event"); } - return Some(self.fail_at("json", "malformed_event")); + return Some(Ok(self.fail_at("json", "malformed_event"))); } }; if let Some(capture) = self.stream_capture.as_mut() { @@ -372,7 +377,7 @@ where } let reduced = match self.reducer.push(value) { Ok(events) => events, - Err(_) => return Some(self.fail_at("reducer", "invalid_event")), + Err(_) => return Some(Ok(self.fail_at("reducer", "invalid_event"))), }; let usage = reduced.iter().find_map(|event| match event { translate::reducer::ReducerEvent::Finish { @@ -384,7 +389,7 @@ where }); match self.translator.render(reduced) { Ok(bytes) => out.extend(bytes), - Err(_) => return Some(self.fail_at("render", "invalid_event")), + Err(_) => return Some(Ok(self.fail_at("render", "invalid_event"))), } if let Some((input_tokens, output_tokens)) = usage && let Some(monitor) = self.monitor.as_ref() @@ -395,16 +400,43 @@ where self.terminal = true; self.capture_downstream(&out); self.finish_capture(true); - return if out.is_empty() { None } else { Some(out) }; + return if out.is_empty() { None } else { Some(Ok(out)) }; } } if !out.is_empty() { self.capture_downstream(&out); - return Some(out); + return Some(Ok(out)); } } } + fn abort_transport(&mut self, message: &str) -> io::Error { + self.terminal = true; + let mut fields = serde_json::Map::new(); + fields.insert("reqId".into(), serde_json::json!(self.req_id)); + fields.insert("stage".into(), serde_json::json!("transport")); + fields.insert("kind".into(), serde_json::json!("upstream_stream")); + fields.insert("bytes".into(), serde_json::json!(self.bytes)); + fields.insert("chunks".into(), serde_json::json!(self.chunks)); + crate::logging::create_logger("grok").warn("grok_stream_failed", Some(fields)); + if let Some(capture) = self.stream_capture.as_mut() { + capture.malformed("transport", "upstream_stream"); + } + if let Some(traffic) = self.traffic.as_ref() { + traffic.write_json( + "060-grok-stream-error", + &serde_json::json!({ + "stage": "transport", + "kind": "upstream_stream", + "bytes": self.bytes, + "chunks": self.chunks, + }), + ); + } + self.finish_capture(false); + io::Error::new(io::ErrorKind::ConnectionReset, message.to_string()) + } + fn fail_at(&mut self, stage: &str, kind: &str) -> Vec { self.error_sent = true; let mut fields = serde_json::Map::new(); @@ -1035,4 +1067,54 @@ mod tests { assert!(captured.contains("[redacted len=6]")); assert!(!captured.contains("secret")); } + + #[tokio::test] + async fn streaming_transport_error_aborts_before_text_and_emits_sse_error_after_text() { + // Thinking-only followed by upstream transport error -> body aborts with io::Error. + let thinking_then_err = futures_util::stream::iter(vec![ + Ok(Bytes::from_static( + b"data: {\"type\":\"response.reasoning_summary_text.delta\",\"delta\":\"thinking...\"}\n\n", + )), + Err(client::GrokError { + status: StatusCode::BAD_GATEWAY, + retry_after: None, + message: "Grok upstream stream failed".into(), + }), + ]); + let response = stream_body( + thinking_then_err, + "msg_think".into(), + "grok-4.5".into(), + 0, + None, + "req_think".into(), + None, + ); + assert!(response.into_body().collect().await.is_err()); + + // Text delta followed by upstream transport error -> preserves event: error + clean EOF. + let text_then_err = futures_util::stream::iter(vec![ + Ok(Bytes::from_static( + b"data: {\"type\":\"response.output_text.delta\",\"delta\":\"visible\"}\n\n", + )), + Err(client::GrokError { + status: StatusCode::BAD_GATEWAY, + retry_after: None, + message: "Grok upstream stream failed".into(), + }), + ]); + let response = stream_body( + text_then_err, + "msg_text".into(), + "grok-4.5".into(), + 0, + None, + "req_text".into(), + None, + ); + let bytes = response.into_body().collect().await.unwrap().to_bytes(); + let sse = String::from_utf8_lossy(&bytes); + assert!(sse.contains("visible")); + assert!(sse.contains("event: error")); + } } diff --git a/src/providers/grok/translate/stream.rs b/src/providers/grok/translate/stream.rs index b344484f..4b700ca5 100644 --- a/src/providers/grok/translate/stream.rs +++ b/src/providers/grok/translate/stream.rs @@ -98,6 +98,7 @@ pub struct StreamTranslator { model: String, started: bool, finished: bool, + non_thinking_emitted: bool, search_blocks: GrokSearchBlocks, // Seeds Claude Code's live usage counter until the provider returns // authoritative usage in the terminal message_delta. @@ -163,6 +164,7 @@ impl StreamTranslator { model, started: false, finished: false, + non_thinking_emitted: false, search_blocks, estimated_input_tokens: 0, } @@ -173,6 +175,10 @@ impl StreamTranslator { self } + pub fn has_non_thinking_output(&self) -> bool { + self.non_thinking_emitted + } + pub fn render(&mut self, events: Vec) -> anyhow::Result> { if self.finished && !events.is_empty() { anyhow::bail!("event after terminal completion"); @@ -196,6 +202,17 @@ impl StreamTranslator { serde_json::json!({"type":"message_start","message":{"id":self.message_id,"type":"message","role":"assistant","model":self.model,"content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":self.estimated_input_tokens,"output_tokens":0}}}), ); } + if matches!( + event, + ReducerEvent::TextStart(_) + | ReducerEvent::TextDelta(_, _) + | ReducerEvent::ToolStart(_, _, _) + | ReducerEvent::ToolDelta(_, _) + | ReducerEvent::HostedSearch { .. } + | ReducerEvent::Finish { .. } + ) { + self.non_thinking_emitted = true; + } if matches!(event, ReducerEvent::Finish { .. }) { self.finished = true; } diff --git a/src/providers/opencode/chat.rs b/src/providers/opencode/chat.rs index e120c4da..99eca164 100644 --- a/src/providers/opencode/chat.rs +++ b/src/providers/opencode/chat.rs @@ -1,4 +1,4 @@ -use std::convert::Infallible; +use std::io; use std::sync::Arc; use axum::body::Body; @@ -1017,6 +1017,16 @@ impl LiveStreamTranslator { pub fn is_finished(&self) -> bool { self.state.finished } + + pub fn has_non_thinking_output(&self) -> bool { + self.state.finished + || self.state.pending_stop.is_some() + || self + .state + .blocks + .iter() + .any(|block| !matches!(block.kind, BlockKind::Thinking { .. })) + } } fn validate_post_done_metadata(data: &str) -> anyhow::Result<()> { @@ -1081,9 +1091,9 @@ pub fn stream_body( }; let stream = futures_util::stream::unfold(state, |mut state| async move { state - .next_output() + .next_item() .await - .map(|bytes| (Ok::(Bytes::from(bytes)), state)) + .map(|item| (item.map(Bytes::from), state)) }); Body::from_stream(stream) } @@ -1106,7 +1116,12 @@ impl OpenCodeChatStreamState where S: futures_util::Stream> + Unpin, { + #[cfg(test)] async fn next_output(&mut self) -> Option> { + self.next_item().await.and_then(Result::ok) + } + + async fn next_item(&mut self) -> Option, io::Error>> { if self.terminal { return None; } @@ -1117,19 +1132,24 @@ where loop { let chunk = match self.upstream.next().await { Some(Ok(chunk)) => chunk, - Some(Err(_)) => return Some(self.fail_at("transport", "upstream_stream")), + Some(Err(err)) => { + if !self.translator.has_non_thinking_output() { + return Some(Err(self.abort_transport(&err.message))); + } + return Some(Ok(self.fail_at("transport", "upstream_stream"))); + } None => { let output = match self.translator.finish() { Ok(output) => output, - Err(_) => return Some(self.fail_at("decoder", "incomplete_stream")), + Err(_) => return Some(Ok(self.fail_at("decoder", "incomplete_stream"))), }; if self.capture_decoder.finish().is_err() { - return Some(self.fail_at("capture", "incomplete_stream")); + return Some(Ok(self.fail_at("capture", "incomplete_stream"))); } self.terminal = true; self.capture_downstream(&output); self.finish_capture(true); - return (!output.is_empty()).then_some(output); + return (!output.is_empty()).then_some(Ok(output)); } }; if self.bytes == 0 @@ -1143,7 +1163,7 @@ where let output = match self.translator.push(&chunk) { Ok(output) => output, - Err(_) => return Some(self.fail_at("translation", "invalid_event")), + Err(_) => return Some(Ok(self.fail_at("translation", "invalid_event"))), }; if !output.is_empty() { let (input_tokens, output_tokens) = usage_from_anthropic_sse(&output); @@ -1160,18 +1180,27 @@ where } if self.translator.is_finished() { if self.translator.finish().is_err() || self.capture_decoder.finish().is_err() { - return Some(self.fail_at("decoder", "trailing_incomplete_frame")); + return Some(Ok(self.fail_at("decoder", "trailing_incomplete_frame"))); } self.terminal = true; self.finish_capture(true); - return (!output.is_empty()).then_some(output); + return (!output.is_empty()).then_some(Ok(output)); } if !output.is_empty() { - return Some(output); + return Some(Ok(output)); } } } + fn abort_transport(&mut self, message: &str) -> io::Error { + self.terminal = true; + if let Some(capture) = self.stream_capture.as_mut() { + capture.malformed("transport", "upstream_stream"); + } + self.finish_capture(false); + io::Error::new(io::ErrorKind::ConnectionReset, message.to_string()) + } + fn capture_upstream(&mut self, bytes: &[u8]) { let events = match self.capture_decoder.push(bytes) { Ok(events) => events, diff --git a/src/providers/opencode/messages.rs b/src/providers/opencode/messages.rs index b0d6a075..f00463fa 100644 --- a/src/providers/opencode/messages.rs +++ b/src/providers/opencode/messages.rs @@ -1,4 +1,4 @@ -use std::convert::Infallible; +use std::io; use std::sync::Arc; use axum::body::Body; @@ -41,6 +41,7 @@ pub fn stream_body( let state = MessagesStreamState { upstream: upstream.into_stream(), decoder: SseDecoder::default(), + non_thinking_emitted: false, terminal: false, error_sent: false, monitor, @@ -51,10 +52,7 @@ pub fn stream_body( traffic, }; let stream = futures_util::stream::unfold(state, |mut state| async move { - state - .next_output() - .await - .map(|bytes| (Ok::(bytes), state)) + state.next_item().await.map(|item| (item, state)) }); Body::from_stream(stream) } @@ -62,6 +60,7 @@ pub fn stream_body( struct MessagesStreamState { upstream: S, decoder: SseDecoder, + non_thinking_emitted: bool, terminal: bool, error_sent: bool, monitor: Option, @@ -76,7 +75,12 @@ impl MessagesStreamState where S: futures_util::Stream> + Unpin, { + #[cfg(test)] async fn next_output(&mut self) -> Option { + self.next_item().await.and_then(Result::ok) + } + + async fn next_item(&mut self) -> Option> { if self.terminal { return None; } @@ -87,12 +91,17 @@ where let chunk = match self.upstream.next().await { Some(Ok(chunk)) => chunk, - Some(Err(_)) => return Some(self.fail_at("transport", "upstream_stream")), + Some(Err(err)) => { + if !self.non_thinking_emitted { + return Some(Err(self.abort_transport(&err.message))); + } + return Some(Ok(self.fail_at("transport", "upstream_stream"))); + } None => { if self.decoder.finish().is_err() { - return Some(self.fail_at("decoder", "incomplete_stream")); + return Some(Ok(self.fail_at("decoder", "incomplete_stream"))); } - return Some(self.fail_at("protocol", "missing_message_stop")); + return Some(Ok(self.fail_at("protocol", "missing_message_stop"))); } }; if self.bytes == 0 @@ -105,18 +114,18 @@ where let events = match self.decoder.push(&chunk) { Ok(events) => events, - Err(_) => return Some(self.fail_at("decoder", "malformed_sse")), + Err(_) => return Some(Ok(self.fail_at("decoder", "malformed_sse"))), }; let mut input_tokens = None; let mut output_tokens = None; let mut terminal = false; for event in events { if terminal { - return Some(self.fail_at("protocol", "event_after_message_stop")); + return Some(Ok(self.fail_at("protocol", "event_after_message_stop"))); } let value: serde_json::Value = match serde_json::from_str(&event.data) { Ok(value) => value, - Err(_) => return Some(self.fail_at("json", "malformed_event")), + Err(_) => return Some(Ok(self.fail_at("json", "malformed_event"))), }; if let Some(capture) = self.stream_capture.as_mut() { capture.upstream_event(event.event.as_deref(), &value); @@ -135,7 +144,32 @@ where .or(output_tokens); let kind = value.get("type").and_then(serde_json::Value::as_str); if event.event.as_deref() == Some("error") || kind == Some("error") { - return Some(self.fail_at("upstream", "error_event")); + return Some(Ok(self.fail_at("upstream", "error_event"))); + } + if kind == Some("content_block_start") + && value + .pointer("/content_block/type") + .and_then(serde_json::Value::as_str) + .is_some_and(|block_type| { + !matches!(block_type, "thinking" | "redacted_thinking") + }) + { + self.non_thinking_emitted = true; + } + if kind == Some("content_block_delta") + && value + .pointer("/delta/type") + .and_then(serde_json::Value::as_str) + .is_some_and(|delta_type| { + !matches!(delta_type, "thinking_delta" | "signature_delta") + }) + { + self.non_thinking_emitted = true; + } + if matches!(kind, Some("message_delta" | "message_stop")) + || event.event.as_deref() == Some("message_stop") + { + self.non_thinking_emitted = true; } if event.event.as_deref() == Some("message_stop") || kind == Some("message_stop") { terminal = true; @@ -152,12 +186,32 @@ where } if terminal { if self.decoder.finish().is_err() { - return Some(self.fail_at("decoder", "trailing_incomplete_frame")); + return Some(Ok(self.fail_at("decoder", "trailing_incomplete_frame"))); } self.terminal = true; self.finish_capture(true); } - Some(chunk) + Some(Ok(chunk)) + } + + fn abort_transport(&mut self, message: &str) -> io::Error { + self.terminal = true; + if let Some(capture) = self.stream_capture.as_mut() { + capture.malformed("transport", "upstream_stream"); + } + if let Some(traffic) = self.traffic.as_ref() { + traffic.write_json( + "060-opencode-messages-stream-error", + &serde_json::json!({ + "stage": "transport", + "kind": "upstream_stream", + "bytes": self.bytes, + "chunks": self.chunks, + }), + ); + } + self.finish_capture(false); + io::Error::new(io::ErrorKind::ConnectionReset, message.to_string()) } fn fail_at(&mut self, stage: &str, kind: &str) -> Bytes { @@ -267,6 +321,7 @@ mod tests { let mut state = MessagesStreamState { upstream, decoder: SseDecoder::default(), + non_thinking_emitted: false, terminal: false, error_sent: false, monitor: None, @@ -297,6 +352,7 @@ mod tests { let mut state = MessagesStreamState { upstream, decoder: SseDecoder::default(), + non_thinking_emitted: false, terminal: false, error_sent: false, monitor: None, diff --git a/src/providers/opencode/responses.rs b/src/providers/opencode/responses.rs index 147e0dc7..2930e871 100644 --- a/src/providers/opencode/responses.rs +++ b/src/providers/opencode/responses.rs @@ -1,4 +1,4 @@ -use std::convert::Infallible; +use std::io; use std::sync::Arc; use axum::body::Body; @@ -68,9 +68,9 @@ pub fn stream_body( }; let stream = futures_util::stream::unfold(state, |mut state| async move { state - .next_output() + .next_item() .await - .map(|bytes| (Ok::(Bytes::from(bytes)), state)) + .map(|item| (item.map(Bytes::from), state)) }); Body::from_stream(stream) } @@ -95,7 +95,12 @@ impl ResponsesStreamState where S: futures_util::Stream> + Unpin, { + #[cfg(test)] async fn next_output(&mut self) -> Option> { + self.next_item().await.and_then(Result::ok) + } + + async fn next_item(&mut self) -> Option, io::Error>> { if self.terminal { return None; } @@ -107,24 +112,31 @@ where loop { let chunk = match self.upstream.next().await { Some(Ok(chunk)) => chunk, - Some(Err(_)) => return Some(self.fail_at("transport", "upstream_stream")), + Some(Err(err)) => { + if !self.translator.has_non_thinking_output() + && self.pending_completion.is_none() + { + return Some(Err(self.abort_transport(&err.message))); + } + return Some(Ok(self.fail_at("transport", "upstream_stream"))); + } None => { if self.decoder.finish().is_err() { - return Some(self.fail_at("decoder", "incomplete_stream")); + return Some(Ok(self.fail_at("decoder", "incomplete_stream"))); } let Some(completion) = self.pending_completion.take() else { - return Some(self.fail_at("decoder", "incomplete_stream")); + return Some(Ok(self.fail_at("decoder", "incomplete_stream"))); }; let output = match self.translator.accept(&completion, self.traffic.as_deref()) { Ok(output) => output, - Err(_) => return Some(self.fail_at("translation", "invalid_event")), + Err(_) => return Some(Ok(self.fail_at("translation", "invalid_event"))), }; self.terminal = true; self.record_progress(&output); self.capture_downstream(&output); self.finish_capture(true); - return (!output.is_empty()).then_some(output); + return (!output.is_empty()).then_some(Ok(output)); } }; @@ -138,14 +150,14 @@ where let events = match self.decoder.push(&chunk) { Ok(events) => events, - Err(_) => return Some(self.fail_at("decoder", "malformed_sse")), + Err(_) => return Some(Ok(self.fail_at("decoder", "malformed_sse"))), }; let mut output = Vec::new(); for event in events { let data = event.data.trim(); if data == "[DONE]" { if self.pending_completion.is_none() || self.done_seen { - return Some(self.fail_at("protocol", "premature_done")); + return Some(Ok(self.fail_at("protocol", "premature_done"))); } self.done_seen = true; if let Some(capture) = self.stream_capture.as_mut() { @@ -156,7 +168,7 @@ where } let value: serde_json::Value = match serde_json::from_str(data) { Ok(value) => value, - Err(_) => return Some(self.fail_at("json", "malformed_event")), + Err(_) => return Some(Ok(self.fail_at("json", "malformed_event"))), }; if let Some(capture) = self.stream_capture.as_mut() { capture.upstream_event(event.event.as_deref(), &value); @@ -166,7 +178,7 @@ where if event_type == Some("ping") { continue; } - return Some(self.fail_at("protocol", "event_after_completion")); + return Some(Ok(self.fail_at("protocol", "event_after_completion"))); } if matches!( event_type, @@ -178,18 +190,38 @@ where } let translated = match self.translator.accept(&value, self.traffic.as_deref()) { Ok(translated) => translated, - Err(_) => return Some(self.fail_at("translation", "invalid_event")), + Err(_) => return Some(Ok(self.fail_at("translation", "invalid_event"))), }; output.extend(translated); } if !output.is_empty() { self.record_progress(&output); self.capture_downstream(&output); - return Some(output); + return Some(Ok(output)); } } } + fn abort_transport(&mut self, message: &str) -> io::Error { + self.terminal = true; + if let Some(capture) = self.stream_capture.as_mut() { + capture.malformed("transport", "upstream_stream"); + } + if let Some(traffic) = self.traffic.as_ref() { + traffic.write_json( + "060-opencode-responses-stream-error", + &serde_json::json!({ + "stage": "transport", + "kind": "upstream_stream", + "bytes": self.bytes, + "chunks": self.chunks, + }), + ); + } + self.finish_capture(false); + io::Error::new(io::ErrorKind::ConnectionReset, message.to_string()) + } + fn record_progress(&self, output: &[u8]) { let Some(monitor) = self.monitor.as_ref() else { return;