From 952fe096e00362691daddface7761be602bc34ee Mon Sep 17 00:00:00 2001 From: ashwinnathan-openai Date: Mon, 6 Jul 2026 17:17:14 +0000 Subject: [PATCH] [codex] Support concurrent reasoning summary delivery --- codex-rs/codex-api/src/common.rs | 26 +++- .../src/endpoint/responses_websocket.rs | 1 + codex-rs/codex-api/src/lib.rs | 2 + codex-rs/codex-api/src/sse/responses.rs | 81 +++++++++-- codex-rs/codex-api/tests/clients.rs | 3 + codex-rs/codex-api/tests/sse_end_to_end.rs | 10 +- codex-rs/core/src/client.rs | 112 ++++++++++++++- codex-rs/core/src/client_common_tests.rs | 4 + codex-rs/core/src/client_tests.rs | 136 +++++++++++++++++- codex-rs/core/src/compact.rs | 2 +- codex-rs/core/src/compact_remote_v2.rs | 20 ++- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/response_event_buffer.rs | 46 ++++++ codex-rs/core/src/session/turn.rs | 59 +++++++- codex-rs/core/src/turn_timing.rs | 3 +- codex-rs/core/tests/suite/client.rs | 8 ++ .../core/tests/suite/client_websockets.rs | 12 +- codex-rs/core/tests/suite/items.rs | 90 +++++++++--- codex-rs/memories/write/src/runtime.rs | 2 +- codex-rs/otel/src/events/session_telemetry.rs | 5 +- 20 files changed, 560 insertions(+), 63 deletions(-) create mode 100644 codex-rs/core/src/response_event_buffer.rs diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index 254a5cb38f..cab1e46eeb 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -74,7 +74,10 @@ pub struct MemorySummarizeOutput { pub enum ResponseEvent { Created, SafetyBuffering(SafetyBuffering), - OutputItemDone(ResponseItem), + OutputItemDone { + item: ResponseItem, + output_index: Option, + }, OutputItemAdded(ResponseItem), /// Emitted when the server includes `OpenAI-Model` on the stream response. /// This can differ from the requested model when backend safety routing applies. @@ -104,6 +107,11 @@ pub enum ResponseEvent { delta: String, summary_index: i64, }, + ReasoningSummaryDone { + item_id: String, + text: String, + summary_index: i64, + }, ReasoningContentDelta { delta: String, content_index: i64, @@ -148,6 +156,17 @@ pub struct Reasoning { pub context: Option, } +#[derive(Debug, Serialize, Clone, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum ReasoningSummaryDelivery { + ConcurrentCutoff, +} + +#[derive(Debug, Serialize, Clone, PartialEq)] +pub struct StreamOptions { + pub reasoning_summary_delivery: ReasoningSummaryDelivery, +} + #[derive(Debug, Serialize, Default, Clone, PartialEq)] #[serde(rename_all = "snake_case")] pub enum TextFormatType { @@ -209,6 +228,8 @@ pub struct ResponsesApiRequest { pub reasoning: Option, pub store: bool, pub stream: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub stream_options: Option, pub include: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub service_tier: Option, @@ -233,6 +254,7 @@ impl From<&ResponsesApiRequest> for ResponseCreateWsRequest { reasoning: request.reasoning.clone(), store: request.store, stream: request.stream, + stream_options: request.stream_options.clone(), include: request.include.clone(), service_tier: request.service_tier.clone(), prompt_cache_key: request.prompt_cache_key.clone(), @@ -258,6 +280,8 @@ pub struct ResponseCreateWsRequest { pub reasoning: Option, pub store: bool, pub stream: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub stream_options: Option, pub include: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub service_tier: Option, diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 7f261bba36..50daf59958 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -851,6 +851,7 @@ mod tests { reasoning: None, store: false, stream: true, + stream_options: None, include: vec!["reasoning.encrypted_content".to_string()], service_tier: Some("priority".to_string()), prompt_cache_key: Some("cache-key".to_string()), diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index 72de4d0585..1d1a4c8c6e 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -34,11 +34,13 @@ pub use crate::common::RawMemory; pub use crate::common::RawMemoryMetadata; pub use crate::common::Reasoning; pub use crate::common::ReasoningContext; +pub use crate::common::ReasoningSummaryDelivery; pub use crate::common::ResponseCreateWsRequest; pub use crate::common::ResponseEvent; pub use crate::common::ResponseStream; pub use crate::common::ResponsesApiRequest; pub use crate::common::ResponsesWsRequest; +pub use crate::common::StreamOptions; pub use crate::common::TextControls; pub use crate::common::WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY; pub use crate::common::WS_REQUEST_HEADER_TRACESTATE_CLIENT_METADATA_KEY; diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index 17d31a02ee..033e08225b 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -165,8 +165,10 @@ pub struct ResponsesStreamEvent { response: Option, item: Option, item_id: Option, + output_index: Option, call_id: Option, delta: Option, + text: Option, summary_index: Option, content_index: Option, safety_buffering: Option, @@ -329,7 +331,10 @@ pub fn process_responses_event( "response.output_item.done" => { if let Some(item_val) = event.item { if let Ok(item) = serde_json::from_value::(item_val) { - return Ok(Some(ResponseEvent::OutputItemDone(item))); + return Ok(Some(ResponseEvent::OutputItemDone { + item, + output_index: event.output_index, + })); } debug!("failed to parse ResponseItem from output_item.done"); } @@ -358,6 +363,17 @@ pub fn process_responses_event( })); } } + "response.reasoning_summary_text.done" => { + if let (Some(item_id), Some(text), Some(summary_index)) = + (event.item_id, event.text, event.summary_index) + { + return Ok(Some(ResponseEvent::ReasoningSummaryDone { + item_id, + text, + summary_index, + })); + } + } "response.reasoning_text.delta" => { if let (Some(delta), Some(content_index)) = (event.delta, event.content_index) { return Ok(Some(ResponseEvent::ReasoningContentDelta { @@ -734,6 +750,7 @@ mod tests { async fn parses_items_and_completed() { let item1 = json!({ "type": "response.output_item.done", + "output_index": 0, "item": { "type": "message", "role": "assistant", @@ -745,6 +762,7 @@ mod tests { let item2 = json!({ "type": "response.output_item.done", + "output_index": 1, "item": { "type": "message", "role": "assistant", @@ -769,16 +787,22 @@ mod tests { assert_matches!( &events[0], - Ok(ResponseEvent::OutputItemDone(ResponseItem::Message { - role, - phase: Some(MessagePhase::Commentary), - .. - })) if role == "assistant" + Ok(ResponseEvent::OutputItemDone { + item: ResponseItem::Message { + role, + phase: Some(MessagePhase::Commentary), + .. + }, + output_index: Some(0), + }) if role == "assistant" ); assert_matches!( &events[1], - Ok(ResponseEvent::OutputItemDone(ResponseItem::Message { role, .. })) + Ok(ResponseEvent::OutputItemDone { + item: ResponseItem::Message { role, .. }, + output_index: Some(1), + }) if role == "assistant" ); @@ -796,6 +820,32 @@ mod tests { } } + #[tokio::test] + async fn parses_reasoning_summary_done() { + let events = run_sse(vec![ + json!({ + "type": "response.reasoning_summary_text.done", + "item_id": "reasoning-1", + "summary_index": 0, + "text": "Checking" + }), + json!({ + "type": "response.completed", + "response": { "id": "resp1" } + }), + ]) + .await; + + assert_matches!( + &events[0], + ResponseEvent::ReasoningSummaryDone { + item_id, + text, + summary_index: 0, + } if item_id == "reasoning-1" && text == "Checking" + ); + } + #[tokio::test] async fn error_when_missing_completed() { let item1 = json!({ @@ -814,7 +864,7 @@ mod tests { assert_eq!(events.len(), 2); - assert_matches!(events[0], Ok(ResponseEvent::OutputItemDone(_))); + assert_matches!(events[0], Ok(ResponseEvent::OutputItemDone { .. })); match &events[1] { Err(ApiError::Stream(msg)) => { @@ -849,12 +899,15 @@ mod tests { assert_eq!(events.len(), 2); assert_matches!( &events[0], - ResponseEvent::OutputItemDone(ResponseItem::ToolSearchCall { - call_id, - execution, - arguments, + ResponseEvent::OutputItemDone { + item: ResponseItem::ToolSearchCall { + call_id, + execution, + arguments, + .. + }, .. - }) if call_id.as_deref() == Some("search-1") + } if call_id.as_deref() == Some("search-1") && execution == "client" && arguments == &json!({"query": "calendar create", "limit": 1}) ); @@ -1071,7 +1124,7 @@ mod tests { matches!(ev, ResponseEvent::Created) } fn is_output(ev: &ResponseEvent) -> bool { - matches!(ev, ResponseEvent::OutputItemDone(_)) + matches!(ev, ResponseEvent::OutputItemDone { .. }) } fn is_completed(ev: &ResponseEvent) -> bool { matches!(ev, ResponseEvent::Completed { .. }) diff --git a/codex-rs/codex-api/tests/clients.rs b/codex-rs/codex-api/tests/clients.rs index 60a9852cfa..ab4ab1d08f 100644 --- a/codex-rs/codex-api/tests/clients.rs +++ b/codex-rs/codex-api/tests/clients.rs @@ -321,6 +321,7 @@ async fn responses_client_stream_request_preserves_item_ids() -> Result<()> { reasoning: None, store: false, stream: true, + stream_options: None, include: Vec::new(), service_tier: None, prompt_cache_key: None, @@ -407,6 +408,7 @@ async fn streaming_client_retries_on_transport_error() -> Result<()> { reasoning: None, store: false, stream: true, + stream_options: None, include: Vec::new(), service_tier: None, prompt_cache_key: None, @@ -526,6 +528,7 @@ async fn azure_store_sends_ids_and_headers() -> Result<()> { reasoning: None, store: true, stream: true, + stream_options: None, include: Vec::new(), service_tier: None, prompt_cache_key: None, diff --git a/codex-rs/codex-api/tests/sse_end_to_end.rs b/codex-rs/codex-api/tests/sse_end_to_end.rs index 2526de281d..611193e077 100644 --- a/codex-rs/codex-api/tests/sse_end_to_end.rs +++ b/codex-rs/codex-api/tests/sse_end_to_end.rs @@ -140,14 +140,20 @@ async fn responses_stream_parses_items_and_completed_end_to_end() -> Result<()> assert_eq!(events.len(), 3); match &events[0] { - ResponseEvent::OutputItemDone(ResponseItem::Message { role, .. }) => { + ResponseEvent::OutputItemDone { + item: ResponseItem::Message { role, .. }, + .. + } => { assert_eq!(role, "assistant"); } other => panic!("unexpected first event: {other:?}"), } match &events[1] { - ResponseEvent::OutputItemDone(ResponseItem::Message { role, .. }) => { + ResponseEvent::OutputItemDone { + item: ResponseItem::Message { role, .. }, + .. + } => { assert_eq!(role, "assistant"); } other => panic!("unexpected second event: {other:?}"), diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 43a758b0e8..63883c8d7c 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -56,6 +56,7 @@ use codex_api::ResponsesWebsocketConnection as ApiWebSocketConnection; use codex_api::ResponsesWsRequest; use codex_api::SharedAuthProvider; use codex_api::SseTelemetry; +use codex_api::StreamOptions; use codex_api::TransportError; use codex_api::WebsocketTelemetry; use codex_api::auth_header_telemetry; @@ -111,7 +112,9 @@ use crate::client_common::Prompt; use crate::client_common::ResponseEvent; use crate::client_common::ResponseStream; use crate::feedback_tags; +use crate::response_event_buffer::OutputItemDoneBuffer; use crate::responses_metadata::CodexResponsesMetadata; +use crate::responses_metadata::CodexResponsesRequestKind; use crate::responses_metadata::subagent_header_value; use crate::util::emit_feedback_auth_recovery_tags; use codex_feedback::FeedbackRequestTags; @@ -311,6 +314,7 @@ fn responses_request_properties_match( reasoning: previous_reasoning, store: previous_store, stream: previous_stream, + stream_options: _, include: previous_include, service_tier: previous_service_tier, prompt_cache_key: previous_prompt_cache_key, @@ -327,6 +331,7 @@ fn responses_request_properties_match( reasoning: current_reasoning, store: current_store, stream: current_stream, + stream_options: _, include: current_include, service_tier: current_service_tier, prompt_cache_key: current_prompt_cache_key, @@ -342,6 +347,8 @@ fn responses_request_properties_match( && previous_reasoning == current_reasoning && previous_store == current_store && previous_stream == current_stream + // Stream options control delivery for the current response, not the context + // referenced by `previous_response_id`. && previous_include == current_include && previous_service_tier == current_service_tier && previous_prompt_cache_key == current_prompt_cache_key @@ -392,6 +399,16 @@ fn sideband_websocket_auth_headers(api_auth: &dyn AuthProvider) -> ApiHeaderMap } impl ModelClient { + pub(crate) fn should_request_concurrent_reasoning_summaries( + provider_info: &ModelProviderInfo, + model_info: &ModelInfo, + summary: ReasoningSummaryConfig, + ) -> bool { + provider_info.is_openai() + && model_info.supports_reasoning_summaries + && summary != ReasoningSummaryConfig::None + } + #[allow(clippy::too_many_arguments)] /// Creates a new session-scoped `ModelClient`. /// @@ -848,6 +865,19 @@ impl ModelClient { } else { (prompt.base_instructions.text.clone(), Some(tools)) }; + // Keep setup and background requests byte-for-byte compatible; only normal turns + // consume concurrent summary delivery in the turn event loop. + let stream_options = (matches!( + responses_metadata.request_kind, + Some(CodexResponsesRequestKind::Turn) + ) && Self::should_request_concurrent_reasoning_summaries( + self.state.provider.info(), + model_info, + summary, + )) + .then_some(StreamOptions { + reasoning_summary_delivery: codex_api::ReasoningSummaryDelivery::ConcurrentCutoff, + }); let reasoning = Self::build_reasoning(model_info, effort, summary); let include = if reasoning.is_some() { vec!["reasoning.encrypted_content".to_string()] @@ -882,6 +912,7 @@ impl ModelClient { reasoning, store: provider.is_azure_responses_endpoint(), stream: true, + stream_options, include, service_tier, prompt_cache_key, @@ -1411,6 +1442,7 @@ impl ModelClientSession { let inference_trace_attempt = inference_trace.start_attempt(); inference_trace_attempt.add_request_headers(&mut options.extra_headers); inference_trace_attempt.record_started(&request); + let reorder_output_items = request.stream_options.is_some(); let client = ApiResponsesClient::new( transport, client_setup.api_provider, @@ -1426,6 +1458,7 @@ impl ModelClientSession { request_session_telemetry, inference_trace_attempt, Arc::clone(&self.client.state.provider), + reorder_output_items, ); return Ok(stream); } @@ -1522,6 +1555,7 @@ impl ModelClientSession { } else { session_telemetry_for_request(session_telemetry, &request) }; + let reorder_output_items = request.stream_options.is_some(); let mut client_metadata = self .client .build_ws_client_metadata(responses_metadata, model_info.use_responses_lite); @@ -1537,6 +1571,7 @@ impl ModelClientSession { }; if warmup { ws_payload.generate = Some(false); + ws_payload.stream_options = None; } match self @@ -1628,6 +1663,7 @@ impl ModelClientSession { request_session_telemetry, inference_trace_attempt, Arc::clone(&self.client.state.provider), + reorder_output_items, ); self.websocket_session.last_response_rx = Some(last_request_rx); return Ok(WebsocketStreamOutcome::Stream(stream)); @@ -1872,11 +1908,30 @@ fn add_responses_lite_header(headers: &mut ApiHeaderMap, use_responses_lite: boo const RESPONSE_STREAM_CHANNEL_CAPACITY: usize = 1600; const STREAM_DROPPED_REASON: &str = "response stream dropped before provider terminal event"; +async fn send_output_item_done_events( + tx_event: &mpsc::Sender>, + items_added: &mut Vec, + ready: Vec<(ResponseItem, Option)>, +) -> bool { + for (item, output_index) in ready { + items_added.push(item.clone()); + if tx_event + .send(Ok(ResponseEvent::OutputItemDone { item, output_index })) + .await + .is_err() + { + return false; + } + } + true +} + fn map_response_stream( api_stream: codex_api::ResponseStream, session_telemetry: SessionTelemetry, inference_trace_attempt: InferenceTraceAttempt, provider: SharedModelProvider, + reorder_output_items: bool, ) -> (ResponseStream, oneshot::Receiver) { let codex_api::ResponseStream { rx_event, @@ -1892,6 +1947,7 @@ fn map_response_stream( session_telemetry, inference_trace_attempt, provider, + reorder_output_items, ) } @@ -1901,6 +1957,7 @@ fn map_response_events( session_telemetry: SessionTelemetry, inference_trace_attempt: InferenceTraceAttempt, provider: SharedModelProvider, + reorder_output_items: bool, ) -> (ResponseStream, oneshot::Receiver) where S: futures::Stream> @@ -1918,6 +1975,7 @@ where let mut logged_error = false; let mut tx_last_response = Some(tx_last_response); let mut items_added: Vec = Vec::new(); + let mut output_item_done_buffer = reorder_output_items.then(OutputItemDoneBuffer::default); let (request_start, mut ttft_ms) = (Instant::now(), None); let mut api_stream = api_stream; let upstream_request_id = upstream_request_id.as_deref(); @@ -1940,13 +1998,12 @@ where break; }; match event { - Ok(ResponseEvent::OutputItemDone(item)) => { - items_added.push(item.clone()); - if tx_event - .send(Ok(ResponseEvent::OutputItemDone(item))) - .await - .is_err() - { + Ok(ResponseEvent::OutputItemDone { item, output_index }) => { + let ready = match output_item_done_buffer.as_mut() { + Some(buffer) => buffer.push(item, output_index), + None => vec![(item, output_index)], + }; + if !send_output_item_done_events(&tx_event, &mut items_added, ready).await { inference_trace_attempt.record_cancelled( STREAM_DROPPED_REASON, upstream_request_id, @@ -1960,6 +2017,22 @@ where token_usage, end_turn, }) => { + if let Some(buffer) = output_item_done_buffer.as_mut() { + if !send_output_item_done_events( + &tx_event, + &mut items_added, + buffer.finish(), + ) + .await + { + inference_trace_attempt.record_cancelled( + STREAM_DROPPED_REASON, + upstream_request_id, + &items_added, + ); + return; + } + } feedback_tags!(last_model_response_id = &response_id); if let Some(usage) = &token_usage { session_telemetry.sse_event_completed( @@ -2011,6 +2084,21 @@ where } } Err(err) => { + if let Some(buffer) = output_item_done_buffer.as_mut() + && !send_output_item_done_events( + &tx_event, + &mut items_added, + buffer.finish(), + ) + .await + { + inference_trace_attempt.record_cancelled( + STREAM_DROPPED_REASON, + upstream_request_id, + &items_added, + ); + return; + } let response_debug_context = extract_response_debug_context_from_api_error(&err); let upstream_request_id = @@ -2034,6 +2122,16 @@ where } } } + if let Some(buffer) = output_item_done_buffer.as_mut() + && !send_output_item_done_events(&tx_event, &mut items_added, buffer.finish()).await + { + inference_trace_attempt.record_cancelled( + STREAM_DROPPED_REASON, + upstream_request_id, + &items_added, + ); + return; + } inference_trace_attempt.record_failed( "stream closed before response.completed", upstream_request_id, diff --git a/codex-rs/core/src/client_common_tests.rs b/codex-rs/core/src/client_common_tests.rs index 38c4aaecb5..b8d0a364a3 100644 --- a/codex-rs/core/src/client_common_tests.rs +++ b/codex-rs/core/src/client_common_tests.rs @@ -116,6 +116,7 @@ fn serializes_text_verbosity_when_set() { reasoning: None, store: false, stream: true, + stream_options: None, include: vec![], prompt_cache_key: None, service_tier: None, @@ -163,6 +164,7 @@ fn serializes_text_schema_with_strict_format() { reasoning: None, store: false, stream: true, + stream_options: None, include: vec![], prompt_cache_key: None, service_tier: None, @@ -224,6 +226,7 @@ fn omits_text_when_not_set() { reasoning: None, store: false, stream: true, + stream_options: None, include: vec![], prompt_cache_key: None, service_tier: None, @@ -247,6 +250,7 @@ fn serializes_flex_service_tier_when_set() { reasoning: None, store: false, stream: true, + stream_options: None, include: vec![], prompt_cache_key: None, service_tier: Some(ServiceTier::Flex.to_string()), diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 82c464625d..6217f27400 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -577,21 +577,25 @@ async fn dropped_response_stream_traces_cancelled_partial_output() -> anyhow::Re // item in history, so the trace should preserve it when the stream is // abandoned. let item = output_message("msg-1", "partial answer"); - let api_stream = futures::stream::iter([Ok(ResponseEvent::OutputItemDone(item))]) - .chain(futures::stream::pending()); + let api_stream = futures::stream::iter([Ok(ResponseEvent::OutputItemDone { + item, + output_index: Some(0), + })]) + .chain(futures::stream::pending()); let (mut stream, _) = super::map_response_events( /*upstream_request_id*/ None, api_stream, test_session_telemetry(), attempt, test_model_provider(), + /*reorder_output_items*/ false, ); let observed = stream .next() .await .expect("mapped stream should yield output item")?; - assert!(matches!(observed, ResponseEvent::OutputItemDone(_))); + assert!(matches!(observed, ResponseEvent::OutputItemDone { .. })); // Dropping the consumer is how turn interruption/preemption stops polling // the provider stream. The mapper task observes that drop asynchronously @@ -635,6 +639,7 @@ async fn response_stream_records_last_model_feedback_ids() { test_session_telemetry(), InferenceTraceAttempt::disabled(), test_model_provider(), + /*reorder_output_items*/ false, ); while stream.next().await.is_some() {} @@ -650,6 +655,122 @@ async fn response_stream_records_last_model_feedback_ids() { ); } +#[tokio::test] +async fn response_stream_reorders_completed_items_for_concurrent_cutoff() { + let api_stream = futures::stream::iter([ + Ok(ResponseEvent::OutputItemDone { + item: output_message("msg-1", "second"), + output_index: Some(1), + }), + Ok(ResponseEvent::OutputItemDone { + item: output_message("msg-0", "first"), + output_index: Some(0), + }), + Ok(ResponseEvent::Completed { + response_id: "resp-123".to_string(), + token_usage: None, + end_turn: Some(true), + }), + ]); + let (mut stream, last_response_rx) = super::map_response_events( + /*upstream_request_id*/ None, + api_stream, + test_session_telemetry(), + InferenceTraceAttempt::disabled(), + test_model_provider(), + /*reorder_output_items*/ true, + ); + + let mut completed_item_ids = Vec::new(); + while let Some(event) = stream.next().await { + if let ResponseEvent::OutputItemDone { item, .. } = + event.expect("mapped stream should succeed") + { + completed_item_ids.push(item.id().expect("output item id").to_string()); + } + } + + assert_eq!(completed_item_ids, ["msg-0", "msg-1"]); + let last_response = last_response_rx.await.expect("last response"); + assert_eq!( + last_response + .items_added + .iter() + .map(|item| item.id().expect("output item id")) + .collect::>(), + ["msg-0", "msg-1"] + ); +} + +#[tokio::test] +async fn response_stream_flushes_buffered_items_before_error() { + let api_stream = futures::stream::iter([ + Ok(ResponseEvent::OutputItemDone { + item: output_message("msg-1", "second"), + output_index: Some(1), + }), + Err(ApiError::Stream("boom".to_string())), + ]); + let (mut stream, _) = super::map_response_events( + /*upstream_request_id*/ None, + api_stream, + test_session_telemetry(), + InferenceTraceAttempt::disabled(), + test_model_provider(), + /*reorder_output_items*/ true, + ); + + let first = stream + .next() + .await + .expect("buffered item should be flushed") + .expect("buffered item should succeed"); + assert!(matches!( + first, + ResponseEvent::OutputItemDone { + output_index: Some(1), + .. + } + )); + assert!( + stream + .next() + .await + .expect("mapped error should be emitted") + .is_err() + ); +} + +#[tokio::test] +async fn response_stream_flushes_buffered_items_at_eof() { + let api_stream = futures::stream::iter([Ok(ResponseEvent::OutputItemDone { + item: output_message("msg-1", "second"), + output_index: Some(1), + })]); + let (mut stream, _) = super::map_response_events( + /*upstream_request_id*/ None, + api_stream, + test_session_telemetry(), + InferenceTraceAttempt::disabled(), + test_model_provider(), + /*reorder_output_items*/ true, + ); + + let first = stream + .next() + .await + .expect("buffered item should be flushed") + .expect("buffered item should succeed"); + assert!(matches!( + first, + ResponseEvent::OutputItemDone { + output_index: Some(1), + .. + } + )); + assert!(stream.next().await.is_none()); +} + #[tokio::test] async fn bedrock_unauthorized_error_uses_provider_mapping() { let provider = create_model_provider( @@ -693,10 +814,10 @@ async fn dropped_backpressured_response_stream_traces_cancelled_partial_output() for _ in 0..super::RESPONSE_STREAM_CHANNEL_CAPACITY { events.push_back(ResponseEvent::Created); } - events.push_back(ResponseEvent::OutputItemDone(output_message( - "msg-1", - "partial answer", - ))); + events.push_back(ResponseEvent::OutputItemDone { + item: output_message("msg-1", "partial answer"), + output_index: Some(0), + }); let api_stream = NotifyAfterEventStream { events, yielded: 0, @@ -710,6 +831,7 @@ async fn dropped_backpressured_response_stream_traces_cancelled_partial_output() test_session_telemetry(), attempt, test_model_provider(), + /*reorder_output_items*/ false, ); // Fill the mapper channel with non-terminal events, then yield one output diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index d2a9fd5222..a0368fb949 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -688,7 +688,7 @@ async fn drain_to_completed( )); }; match event { - Ok(ResponseEvent::OutputItemDone(item)) => { + Ok(ResponseEvent::OutputItemDone { item, .. }) => { sess.record_conversation_items(turn_context, std::slice::from_ref(&item)) .await; } diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 6e110045bb..be2f4155b1 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -404,7 +404,7 @@ async fn collect_compaction_output( let mut completed_token_usage = None; while let Some(event) = stream.next().await { match event? { - ResponseEvent::OutputItemDone(item) => { + ResponseEvent::OutputItemDone { item, .. } => { output_item_count += 1; if let ResponseItem::Compaction { .. } = item { compaction_count += 1; @@ -829,12 +829,18 @@ mod tests { internal_chat_message_metadata_passthrough: None, }; let stream = response_stream(vec![ - Ok(ResponseEvent::OutputItemDone(message( - "assistant", - "IGNORED_COMPACT_REPLY", - Some(MessagePhase::FinalAnswer), - ))), - Ok(ResponseEvent::OutputItemDone(compaction.clone())), + Ok(ResponseEvent::OutputItemDone { + item: message( + "assistant", + "IGNORED_COMPACT_REPLY", + Some(MessagePhase::FinalAnswer), + ), + output_index: Some(0), + }), + Ok(ResponseEvent::OutputItemDone { + item: compaction.clone(), + output_index: Some(1), + }), Ok(ResponseEvent::Completed { response_id: "resp-compact".to_string(), token_usage: Some(TokenUsage { diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index df8753a386..29445f4fd1 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -12,6 +12,7 @@ mod client_common; mod realtime_context; mod realtime_conversation; mod realtime_prompt; +mod response_event_buffer; mod responses_metadata; mod responses_retry; pub(crate) mod session; diff --git a/codex-rs/core/src/response_event_buffer.rs b/codex-rs/core/src/response_event_buffer.rs new file mode 100644 index 0000000000..981ddb4dab --- /dev/null +++ b/codex-rs/core/src/response_event_buffer.rs @@ -0,0 +1,46 @@ +use std::collections::BTreeMap; + +use codex_protocol::models::ResponseItem; + +/// Restores response order for completed items without delaying incremental events. +#[derive(Default)] +pub(crate) struct OutputItemDoneBuffer { + next_output_index: usize, + pending: BTreeMap, +} + +impl OutputItemDoneBuffer { + pub(crate) fn push( + &mut self, + item: ResponseItem, + output_index: Option, + ) -> Vec<(ResponseItem, Option)> { + let Some(output_index) = output_index else { + return vec![(item, None)]; + }; + + if output_index < self.next_output_index { + return vec![(item, Some(output_index))]; + } + + self.pending.insert(output_index, item); + self.take_ready() + } + + pub(crate) fn finish(&mut self) -> Vec<(ResponseItem, Option)> { + let pending = std::mem::take(&mut self.pending); + pending + .into_iter() + .map(|(output_index, item)| (item, Some(output_index))) + .collect() + } + + fn take_ready(&mut self) -> Vec<(ResponseItem, Option)> { + let mut ready = Vec::new(); + while let Some(item) = self.pending.remove(&self.next_output_index) { + ready.push((item, Some(self.next_output_index))); + self.next_output_index += 1; + } + ready + } +} diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index f16b525002..5faf5b1b9c 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -6,6 +6,7 @@ use std::sync::atomic::Ordering; use crate::SkillInjections; use crate::build_skill_injections; +use crate::client::ModelClient; use crate::client::ModelClientSession; use crate::client_common::Prompt; use crate::client_common::ResponseEvent; @@ -1909,6 +1910,15 @@ async fn try_run_sampling_request( turn_context.provider.info().name.as_str(), ); let sampling_timing_guard = turn_context.turn_timing_state.begin_sampling(); + let uses_concurrent_reasoning_summaries = matches!( + responses_metadata.request_kind, + Some(CodexResponsesRequestKind::Turn) + ) + && ModelClient::should_request_concurrent_reasoning_summaries( + turn_context.provider.info(), + &turn_context.model_info, + turn_context.reasoning_summary, + ); let mut stream = client_session .stream( prompt, @@ -1928,6 +1938,7 @@ async fn try_run_sampling_request( let mut needs_follow_up = false; let mut last_agent_message: Option = None; let mut active_item: Option = None; + let mut current_response_item_id: Option = None; let mut active_tool_argument_diff_consumer: Option<( String, Box, @@ -1985,7 +1996,7 @@ async fn try_run_sampling_request( match event { ResponseEvent::Created => {} - ResponseEvent::OutputItemDone(item) => { + ResponseEvent::OutputItemDone { item, .. } => { if let Some((_, mut consumer)) = active_tool_argument_diff_consumer.take() && let Ok(Some(event)) = consumer.finish() { @@ -2080,6 +2091,7 @@ async fn try_run_sampling_request( } } ResponseEvent::OutputItemAdded(item) => { + current_response_item_id = item.id().map(str::to_owned); if let ResponseItem::CustomToolCall { call_id, name, @@ -2291,6 +2303,9 @@ async fn try_run_sampling_request( delta, summary_index, } => { + if uses_concurrent_reasoning_summaries { + continue; + } if let Some(active) = active_item.as_ref() { if !active_item_is_streaming_to_client { continue; @@ -2309,6 +2324,9 @@ async fn try_run_sampling_request( } } ResponseEvent::ReasoningSummaryPartAdded { summary_index } => { + if uses_concurrent_reasoning_summaries { + continue; + } if let Some(active) = active_item.as_ref() { if !active_item_is_streaming_to_client { continue; @@ -2323,6 +2341,45 @@ async fn try_run_sampling_request( error_or_panic("ReasoningSummaryPartAdded without active item".to_string()); } } + ResponseEvent::ReasoningSummaryDone { + item_id, + text, + summary_index, + } => { + // Concurrent cutoff can cancel a partial summary at the next output-item + // boundary. Present only atomic parts that finish while their item is current. + if !uses_concurrent_reasoning_summaries + || defer_streamed_turn_items_for_contributors + || current_response_item_id.as_deref() != Some(item_id.as_str()) + { + continue; + } + let Some(active) = active_item.as_ref() else { + continue; + }; + if !active_item_is_streaming_to_client || active.id() != item_id { + continue; + } + if summary_index > 0 { + sess.send_event( + &turn_context, + EventMsg::AgentReasoningSectionBreak(AgentReasoningSectionBreakEvent { + item_id: item_id.clone(), + summary_index, + }), + ) + .await; + } + let event = ReasoningContentDeltaEvent { + thread_id: sess.thread_id.to_string(), + turn_id: turn_context.sub_id.clone(), + item_id, + delta: text, + summary_index, + }; + sess.send_event(&turn_context, EventMsg::ReasoningContentDelta(event)) + .await; + } ResponseEvent::ReasoningContentDelta { delta, content_index, diff --git a/codex-rs/core/src/turn_timing.rs b/codex-rs/core/src/turn_timing.rs index 04f7aeb69b..e291ce45cb 100644 --- a/codex-rs/core/src/turn_timing.rs +++ b/codex-rs/core/src/turn_timing.rs @@ -329,11 +329,12 @@ impl TurnTimingStateInner { fn response_event_records_turn_ttft(event: &ResponseEvent) -> bool { match event { - ResponseEvent::OutputItemDone(item) | ResponseEvent::OutputItemAdded(item) => { + ResponseEvent::OutputItemDone { item, .. } | ResponseEvent::OutputItemAdded(item) => { response_item_records_turn_ttft(item) } ResponseEvent::OutputTextDelta(_) | ResponseEvent::ReasoningSummaryDelta { .. } + | ResponseEvent::ReasoningSummaryDone { .. } | ResponseEvent::ReasoningContentDelta { .. } => true, ResponseEvent::Created | ResponseEvent::ServerModel(_) diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index c87539f940..ef16a83cf2 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -2412,6 +2412,13 @@ async fn configured_reasoning_summary_is_sent() -> anyhow::Result<()> { .and_then(|reasoning| reasoning.get("context")), None ); + pretty_assertions::assert_eq!( + request_body + .get("stream_options") + .and_then(|stream_options| stream_options.get("reasoning_summary_delivery")) + .and_then(|value| value.as_str()), + Some("concurrent_cutoff") + ); Ok(()) } @@ -2583,6 +2590,7 @@ async fn reasoning_summary_is_omitted_when_disabled() -> anyhow::Result<()> { .and_then(|reasoning| reasoning.get("summary")), None ); + pretty_assertions::assert_eq!(request_body.get("stream_options"), None); Ok(()) } diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index d5218cdcab..4c1313bb8f 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -386,7 +386,12 @@ async fn responses_websocket_request_prewarm_reuses_connection() { ]]) .await; - let harness = websocket_harness_with_options(&server, /*runtime_metrics_enabled*/ true).await; + let mut provider = websocket_provider(&server); + provider.name = "OpenAI".to_string(); + let mut harness = + websocket_harness_with_provider_options(provider, /*runtime_metrics_enabled*/ true).await; + harness.model_info.supports_reasoning_summaries = true; + harness.summary = ReasoningSummary::Concise; let mut client_session = harness.client.new_session(); let prompt = prompt_with_input(vec![message_item("hello")]); let responses_metadata = prewarm_metadata(&harness, /*turn_id*/ None); @@ -422,6 +427,7 @@ async fn responses_websocket_request_prewarm_reuses_connection() { assert_eq!(warmup["type"].as_str(), Some("response.create")); assert_eq!(warmup["generate"].as_bool(), Some(false)); + assert_eq!(warmup.get("stream_options"), None); assert_eq!(warmup["tools"], serde_json::json!([])); let warmup_turn_metadata: serde_json::Value = serde_json::from_str( warmup["client_metadata"]["x-codex-turn-metadata"] @@ -436,6 +442,10 @@ async fn responses_websocket_request_prewarm_reuses_connection() { assert_eq!(follow_up["type"].as_str(), Some("response.create")); assert_eq!(follow_up["previous_response_id"].as_str(), Some("warm-1")); assert_eq!(follow_up["input"], serde_json::json!([])); + assert_eq!( + follow_up["stream_options"]["reasoning_summary_delivery"].as_str(), + Some("concurrent_cutoff") + ); server.shutdown().await; } diff --git a/codex-rs/core/tests/suite/items.rs b/codex-rs/core/tests/suite/items.rs index 515d496b8b..dfd4337026 100644 --- a/codex-rs/core/tests/suite/items.rs +++ b/codex-rs/core/tests/suite/items.rs @@ -3,6 +3,7 @@ use anyhow::Ok; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; +use codex_protocol::config_types::ReasoningSummary; use codex_protocol::config_types::Settings; use codex_protocol::items::AgentMessageContent; use codex_protocol::items::TurnItem; @@ -1085,21 +1086,56 @@ async fn plan_mode_handles_missing_plan_close_tag() -> anyhow::Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn reasoning_content_delta_has_item_metadata() -> anyhow::Result<()> { +async fn concurrent_cutoff_emits_only_current_reasoning_summary_done() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let TestCodex { codex, .. } = test_codex().build(&server).await?; - - let stream = sse(vec![ - ev_response_created("resp-1"), - ev_reasoning_item_added("reasoning-1", &[""]), - ev_reasoning_summary_text_delta("step one"), - ev_reasoning_item("reasoning-1", &["step one"], &[]), - ev_completed("resp-1"), - ]); - mount_sse_once(&server, stream).await; + let TestCodex { codex, .. } = test_codex() + .with_model_info_override("gpt-5.4", |model_info| { + model_info.supports_reasoning_summaries = true; + }) + .with_config(|config| { + config.model_reasoning_summary = Some(ReasoningSummary::Auto); + }) + .build(&server) + .await?; + let summary_done = |item_id: &str, summary_index: i64, text: &str| { + serde_json::json!({ + "type": "response.reasoning_summary_text.done", + "item_id": item_id, + "summary_index": summary_index, + "text": text, + }) + }; + let with_output_index = |mut event: serde_json::Value, output_index: usize| { + event["output_index"] = serde_json::json!(output_index); + event + }; + let response_mock = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_reasoning_item_added("reasoning-1", &[""]), + // Concurrent cutoff renders atomic done events, not partial legacy events. + ev_reasoning_summary_text_delta("partial"), + summary_done("reasoning-1", 0, "step one"), + ev_message_item_added("message-1", ""), + // Once a later output item starts, summaries for the earlier item are stale. + summary_done("reasoning-1", 1, "late step"), + ev_output_text_delta("Done"), + with_output_index( + ev_reasoning_item("reasoning-1", &["step one", "late step"], &[]), + /*output_index*/ 0, + ), + with_output_index( + ev_assistant_message("message-1", "Done"), + /*output_index*/ 1, + ), + ev_completed("resp-1"), + ]), + ) + .await; codex .submit(Op::UserInput { @@ -1123,13 +1159,31 @@ async fn reasoning_content_delta_has_item_metadata() -> anyhow::Result<()> { }) .await; - let delta_event = wait_for_event_match(&codex, |ev| match ev { - EventMsg::ReasoningContentDelta(event) => Some(event.clone()), - _ => None, - }) - .await; - assert_eq!(delta_event.item_id, reasoning_item.id); - assert_eq!(delta_event.delta, "step one"); + let mut summary_deltas = Vec::new(); + let mut summary_sections = Vec::new(); + loop { + match wait_for_event(&codex, |_| true).await { + EventMsg::ReasoningContentDelta(event) => { + summary_deltas.push((event.item_id, event.delta, event.summary_index)); + } + EventMsg::AgentReasoningSectionBreak(event) => { + summary_sections.push((event.item_id, event.summary_index)); + } + EventMsg::TurnComplete(_) => break, + _ => {} + } + } + + assert_eq!( + summary_deltas, + vec![(reasoning_item.id, "step one".to_string(), 0)] + ); + assert_eq!(summary_sections, Vec::new()); + assert_eq!( + response_mock.single_request().body_json()["stream_options"]["reasoning_summary_delivery"] + .as_str(), + Some("concurrent_cutoff") + ); Ok(()) } diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 6c7ad3978f..cdb399f0e1 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -294,7 +294,7 @@ impl MemoryStartupContext { while let Some(message) = stream.next().await.transpose()? { match message { ResponseEvent::OutputTextDelta(delta) => result.push_str(&delta), - ResponseEvent::OutputItemDone(item) => { + ResponseEvent::OutputItemDone { item, .. } => { if result.is_empty() && let codex_protocol::models::ResponseItem::Message { content, .. } = item && let Some(text) = content_items_to_text(&content) diff --git a/codex-rs/otel/src/events/session_telemetry.rs b/codex-rs/otel/src/events/session_telemetry.rs index 440f8f04a5..a3711253c3 100644 --- a/codex-rs/otel/src/events/session_telemetry.rs +++ b/codex-rs/otel/src/events/session_telemetry.rs @@ -417,7 +417,7 @@ impl SessionTelemetry { handle_responses_span.record("otel.name", SessionTelemetry::responses_type(event)); match event { - ResponseEvent::OutputItemDone(item) => { + ResponseEvent::OutputItemDone { item, .. } => { handle_responses_span.record("from", "output_item_done"); if let ResponseItem::FunctionCall { name, .. } = item { handle_responses_span.record("tool_name", name.as_str()); @@ -1188,13 +1188,14 @@ impl SessionTelemetry { fn responses_type(event: &ResponseEvent) -> String { match event { ResponseEvent::Created => "created".into(), - ResponseEvent::OutputItemDone(item) | ResponseEvent::OutputItemAdded(item) => { + ResponseEvent::OutputItemDone { item, .. } | ResponseEvent::OutputItemAdded(item) => { SessionTelemetry::responses_item_type(item) } ResponseEvent::Completed { .. } => "completed".into(), ResponseEvent::OutputTextDelta(_) => "text_delta".into(), ResponseEvent::ToolCallInputDelta { .. } => "tool_input_delta".into(), ResponseEvent::ReasoningSummaryDelta { .. } => "reasoning_summary_delta".into(), + ResponseEvent::ReasoningSummaryDone { .. } => "reasoning_summary_done".into(), ResponseEvent::ReasoningContentDelta { .. } => "reasoning_content_delta".into(), ResponseEvent::ReasoningSummaryPartAdded { .. } => { "reasoning_summary_part_added".into()