diff --git a/codex-rs/api-client/src/client/sse.rs b/codex-rs/api-client/src/client/sse.rs index 8c9354a711..5e944ec50e 100644 --- a/codex-rs/api-client/src/client/sse.rs +++ b/codex-rs/api-client/src/client/sse.rs @@ -1,4 +1,5 @@ use std::time::Duration; +use std::time::Instant; use bytes::Bytes; use codex_otel::otel_event_manager::OtelEventManager; @@ -30,9 +31,16 @@ pub async fn process_sse_wire( let mut buffer = String::new(); loop { + let start = Instant::now(); let result = timeout(max_idle_duration, stream.next()).await; + let duration = start.elapsed(); match result { Err(_) => { + otel_event_manager.sse_event_failed( + None, + duration, + &"idle timeout waiting for SSE", + ); let _ = tx_event .send(Err(Error::Stream( "stream idle timeout fired before Completed event".to_string(), @@ -42,6 +50,7 @@ pub async fn process_sse_wire( return; } Ok(Some(Err(err))) => { + otel_event_manager.sse_event_failed(None, duration, &err); let _ = tx_event.send(Err(err)).await; return; } @@ -49,6 +58,11 @@ pub async fn process_sse_wire( let chunk_str = match std::str::from_utf8(&chunk) { Ok(s) => s, Err(err) => { + otel_event_manager.sse_event_failed( + None, + duration, + &format!("UTF8 error: {err}"), + ); let _ = tx_event .send(Err(Error::Other(format!( "Invalid UTF-8 in SSE chunk: {err}" @@ -62,13 +76,29 @@ pub async fn process_sse_wire( buffer.push_str(&chunk_str); while let Some(frame) = next_frame(&mut buffer) { - if !handle_frame(frame, &mut decoder, &tx_event, &otel_event_manager).await { + if !handle_frame( + frame, + &mut decoder, + &tx_event, + &otel_event_manager, + duration, + ) + .await + { return; } } } Ok(None) => { - if !drain_buffer(&mut buffer, &mut decoder, &tx_event, &otel_event_manager).await { + if !drain_buffer( + &mut buffer, + &mut decoder, + &tx_event, + &otel_event_manager, + duration, + ) + .await + { return; } return; @@ -99,12 +129,13 @@ async fn drain_buffer( decoder: &mut D, tx_event: &mpsc::Sender>, otel_event_manager: &OtelEventManager, + duration: Duration, ) -> bool where D: crate::client::WireResponseDecoder + Send, { while let Some(frame) = next_frame(buffer) { - if !handle_frame(frame, decoder, tx_event, otel_event_manager).await { + if !handle_frame(frame, decoder, tx_event, otel_event_manager, duration).await { return false; } } @@ -114,7 +145,7 @@ where } let remainder = std::mem::take(buffer); - handle_frame(remainder, decoder, tx_event, otel_event_manager).await + handle_frame(remainder, decoder, tx_event, otel_event_manager, duration).await } async fn handle_frame( @@ -122,23 +153,38 @@ async fn handle_frame( decoder: &mut D, tx_event: &mpsc::Sender>, otel_event_manager: &OtelEventManager, + duration: Duration, ) -> bool where D: crate::client::WireResponseDecoder + Send, { - if let Some(data) = parse_sse_frame(&frame) { - if let Err(e) = decoder.on_frame(&data, tx_event, otel_event_manager).await { - let _ = tx_event.send(Err(e)).await; - return false; + if let Some(frame) = parse_sse_frame(&frame) { + if frame.data.trim() == "[DONE]" { + otel_event_manager.sse_event_kind(&frame.event); + return true; } + + match decoder + .on_frame(&frame.data, tx_event, otel_event_manager) + .await + { + Ok(_) => { + otel_event_manager.sse_event_kind(&frame.event); + } + Err(e) => { + otel_event_manager.sse_event_failed(Some(&frame.event), duration, &e); + let _ = tx_event.send(Err(e)).await; + return false; + } + }; } true } -fn parse_sse_frame(frame: &str) -> Option { +fn parse_sse_frame(frame: &str) -> Option { let mut data = String::new(); - let mut saw_event = false; + let mut event: Option = None; let mut saw_data_line = false; for raw_line in frame.split('\n') { @@ -148,10 +194,10 @@ fn parse_sse_frame(frame: &str) -> Option { } if let Some(rest) = line.strip_prefix("event:") { - if rest.trim_start().is_empty() && !saw_data_line { - continue; + let trimmed = rest.trim_start(); + if !trimmed.is_empty() { + event = Some(trimmed.to_string()); } - saw_event = true; continue; } @@ -171,11 +217,19 @@ fn parse_sse_frame(frame: &str) -> Option { } } - if data.is_empty() && !saw_event && !saw_data_line { + if data.is_empty() && event.is_none() && !saw_data_line { return None; } - Some(data) + Some(SseFrame { + event: event.unwrap_or_else(|| "message".to_string()), + data, + }) +} + +struct SseFrame { + event: String, + data: String, } #[cfg(test)] diff --git a/codex-rs/api-client/src/decode_wire/chat.rs b/codex-rs/api-client/src/decode_wire/chat.rs index eb9cc6471f..85633479b8 100644 --- a/codex-rs/api-client/src/decode_wire/chat.rs +++ b/codex-rs/api-client/src/decode_wire/chat.rs @@ -1,10 +1,14 @@ use async_trait::async_trait; use codex_otel::otel_event_manager::OtelEventManager; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ReasoningItemContent; +use codex_protocol::models::ResponseItem; use serde_json::Value; use tokio::sync::mpsc; use tracing::debug; use crate::client::WireResponseDecoder; +use crate::error::Error; use crate::error::Result; use crate::stream::WireEvent; @@ -16,15 +20,19 @@ struct FunctionCallState { arguments: String, } +#[derive(Default)] pub struct WireChatSseDecoder { fn_call_state: FunctionCallState, + created_emitted: bool, + assistant_started: bool, + assistant_text: String, + reasoning_started: bool, + reasoning_text: String, } impl WireChatSseDecoder { pub fn new() -> Self { - Self { - fn_call_state: FunctionCallState::default(), - } + Self::default() } } @@ -36,11 +44,11 @@ impl WireResponseDecoder for WireChatSseDecoder { tx: &mpsc::Sender>, _otel: &OtelEventManager, ) -> Result<()> { - // Chat sends a terminal "[DONE]" frame; ignore it. - let Ok(parsed_chunk) = serde_json::from_str::(json) else { + // Chat sends a terminal "[DONE]" frame; ignore it. Treat other parse errors as failures. + let parsed_chunk = serde_json::from_str::(json).map_err(|err| { debug!("failed to parse Chat SSE JSON: {}", json); - return Ok(()); - }; + Error::Other(format!("failed to parse Chat SSE JSON: {err}")) + })?; let choices = parsed_chunk .get("choices") @@ -49,10 +57,29 @@ impl WireResponseDecoder for WireChatSseDecoder { .unwrap_or_default(); for choice in choices { + if !self.created_emitted { + let _ = tx.send(Ok(WireEvent::Created)).await; + self.created_emitted = true; + } + if let Some(delta) = choice.get("delta") { if let Some(content) = delta.get("content").and_then(|c| c.as_array()) { for piece in content { if let Some(text) = piece.get("text").and_then(|t| t.as_str()) { + if !self.assistant_started { + self.assistant_started = true; + let message = ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: String::new(), + }], + }; + let value = serde_json::to_value(message) + .unwrap_or_else(|_| Value::String(String::new())); + let _ = tx.send(Ok(WireEvent::OutputItemAdded(value))).await; + } + self.assistant_text.push_str(text); let _ = tx .send(Ok(WireEvent::OutputTextDelta(text.to_string()))) .await; @@ -80,6 +107,19 @@ impl WireResponseDecoder for WireChatSseDecoder { if let Some(reasoning) = delta.get("reasoning_content").and_then(|c| c.as_array()) { for entry in reasoning { if let Some(text) = entry.get("text").and_then(|t| t.as_str()) { + if !self.reasoning_started { + self.reasoning_started = true; + let reasoning_item = ResponseItem::Reasoning { + id: String::new(), + summary: vec![], + content: None, + encrypted_content: None, + }; + let value = serde_json::to_value(reasoning_item) + .unwrap_or_else(|_| Value::String(String::new())); + let _ = tx.send(Ok(WireEvent::OutputItemAdded(value))).await; + } + self.reasoning_text.push_str(text); let _ = tx .send(Ok(WireEvent::ReasoningContentDelta(text.to_string()))) .await; @@ -88,23 +128,69 @@ impl WireResponseDecoder for WireChatSseDecoder { } } - if let Some(finish_reason) = choice.get("finish_reason").and_then(|f| f.as_str()) - && finish_reason == "tool_calls" - && self.fn_call_state.active - { - let function_name = self.fn_call_state.name.take().unwrap_or_default(); - let call_id = self.fn_call_state.call_id.take().unwrap_or_default(); - let arguments = self.fn_call_state.arguments.clone(); - self.fn_call_state = FunctionCallState::default(); + if let Some(finish_reason) = choice.get("finish_reason").and_then(|f| f.as_str()) { + match finish_reason { + "tool_calls" if self.fn_call_state.active => { + let function_name = self.fn_call_state.name.take().unwrap_or_default(); + let call_id = self.fn_call_state.call_id.take().unwrap_or_default(); + let arguments = self.fn_call_state.arguments.clone(); + self.fn_call_state = FunctionCallState::default(); - let item = serde_json::json!({ - "type": "function_call", - "id": call_id, - "call_id": call_id, - "name": function_name, - "arguments": arguments, - }); - let _ = tx.send(Ok(WireEvent::OutputItemDone(item))).await; + let item = serde_json::json!({ + "type": "function_call", + "id": call_id, + "call_id": call_id, + "name": function_name, + "arguments": arguments, + }); + let _ = tx.send(Ok(WireEvent::OutputItemDone(item))).await; + } + "stop" | "length" => { + if self.reasoning_started { + let mut content = Vec::new(); + if !self.reasoning_text.is_empty() { + content.push(ReasoningItemContent::ReasoningText { + text: self.reasoning_text.clone(), + }); + } + let reasoning_item = ResponseItem::Reasoning { + id: String::new(), + summary: vec![], + content: Some(content), + encrypted_content: None, + }; + let value = serde_json::to_value(reasoning_item) + .unwrap_or_else(|_| Value::String(String::new())); + let _ = tx.send(Ok(WireEvent::OutputItemDone(value))).await; + } + + if self.assistant_started { + let message = ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: self.assistant_text.clone(), + }], + }; + let value = serde_json::to_value(message) + .unwrap_or_else(|_| Value::String(String::new())); + let _ = tx.send(Ok(WireEvent::OutputItemDone(value))).await; + } + + let _ = tx + .send(Ok(WireEvent::Completed { + response_id: String::new(), + token_usage: None, + })) + .await; + + self.assistant_started = false; + self.assistant_text.clear(); + self.reasoning_started = false; + self.reasoning_text.clear(); + } + _ => {} + } } } diff --git a/codex-rs/api-client/src/decode_wire/responses.rs b/codex-rs/api-client/src/decode_wire/responses.rs index 98af5a86f1..0a9c29f93e 100644 --- a/codex-rs/api-client/src/decode_wire/responses.rs +++ b/codex-rs/api-client/src/decode_wire/responses.rs @@ -35,10 +35,10 @@ impl WireResponseDecoder for WireResponsesSseDecoder { tx: &mpsc::Sender>, otel: &OtelEventManager, ) -> Result<()> { - let Ok(event) = serde_json::from_str::(json) else { + let event = serde_json::from_str::(json).map_err(|err| { debug!("failed to parse Responses SSE JSON: {}", json); - return Ok(()); - }; + Error::Other(format!("failed to parse Responses SSE JSON: {err}")) + })?; match event.event_type.as_str() { "response.created" => { diff --git a/codex-rs/provider-config/src/lib.rs b/codex-rs/provider-config/src/lib.rs index beed624829..c8417dc1d3 100644 --- a/codex-rs/provider-config/src/lib.rs +++ b/codex-rs/provider-config/src/lib.rs @@ -302,13 +302,20 @@ impl ModelProviderInfo { account_id: None, }) } else { - match self.api_key()? { - Some(key) => Some(AuthContext { + match self.api_key() { + Ok(Some(key)) => Some(AuthContext { mode: AuthMode::ApiKey, bearer_token: Some(key), account_id: None, }), - None => auth.clone(), + Ok(None) => auth.clone(), + Err(err) => { + if auth.is_some() { + auth.clone() + } else { + return Err(err); + } + } } };