From d93071fae87f73ba1d251ebb3808773f18f0b77a Mon Sep 17 00:00:00 2001 From: Brent Traut Date: Tue, 27 Oct 2026 20:34:51 +0000 Subject: [PATCH] Avoid retaining unneeded output items in HTTP response streams (#57563) ## Why HTTP response streams retained copies of completed output items even when inference tracing was disabled, keeping their payloads alive until the response completed. ## What changed Collect completed output items only when inference tracing or WebSocket last-response bookkeeping needs them. Make the last-response sender optional and create its channel only for WebSocket requests. ## Testing Add a regression test that verifies an untraced HTTP stream releases a completed item's shared payload after delivery, while the response remains open. Update existing tracing and WebSocket bookkeeping tests for the optional sender. GitOrigin-RevId: 7695534c6d00f3002005324c43c23bf733f01804 --- codex-rs/core/src/client.rs | 49 ++++++++++++++----------- codex-rs/core/src/client_tests.rs | 43 +++++++++++++++++++--- codex-rs/rollout-trace/src/inference.rs | 10 ++++- 3 files changed, 73 insertions(+), 29 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 3ce7006b82..df64477120 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -2014,14 +2014,14 @@ impl ResponsesHttpExecutor<'_> { match stream_result { Ok(stream) => { - let (stream, _) = map_response_stream( + return Ok(map_response_stream( stream, request_session_telemetry, inference_trace_attempt, Arc::clone(&self.client.state.provider), interceptors, - ); - return Ok(stream); + /*last_response_tx*/ None, + )); } Err(ApiError::Transport(unauthorized_transport)) if self @@ -2357,14 +2357,16 @@ impl ModelClientSession { ); err })?; - let (stream, last_request_rx) = map_response_stream( + let (last_response_tx, last_response_rx) = oneshot::channel(); + let stream = map_response_stream( stream_result, request_session_telemetry, inference_trace_attempt, Arc::clone(&self.client.state.provider), interceptors, + Some(last_response_tx), ); - self.websocket_session.last_response_rx = Some(last_request_rx); + self.websocket_session.last_response_rx = Some(last_response_rx); return Ok(WebsocketStreamOutcome::Stream(stream)); } } @@ -2619,18 +2621,20 @@ fn map_response_stream( inference_trace_attempt: InferenceTraceAttempt, provider: SharedModelProvider, interceptors: Vec>, -) -> (ResponseStream, oneshot::Receiver) { + last_response_tx: Option>, +) -> ResponseStream { let upstream_request_id = api_stream.upstream_request_id.take(); let interrupt = api_stream.interrupt.take(); - let (mut stream, last_response) = map_response_events( + let mut stream = map_response_events( upstream_request_id, crate::model_request::intercept_stream(Box::pin(api_stream), interceptors), session_telemetry, inference_trace_attempt, provider, + last_response_tx, ); stream.interrupt = interrupt; - (stream, last_response) + stream } fn map_response_events( @@ -2639,7 +2643,8 @@ fn map_response_events( session_telemetry: SessionTelemetry, inference_trace_attempt: InferenceTraceAttempt, provider: SharedModelProvider, -) -> (ResponseStream, oneshot::Receiver) + mut last_response_tx: Option>, +) -> ResponseStream where S: futures::Stream> + Unpin @@ -2648,13 +2653,13 @@ where { let (tx_event, rx_event) = mpsc::channel::>(RESPONSE_STREAM_CHANNEL_CAPACITY); - let (tx_last_response, rx_last_response) = oneshot::channel::(); let consumer_dropped = CancellationToken::new(); let consumer_dropped_for_stream = consumer_dropped.clone(); tokio::spawn(async move { let mut logged_error = false; - let mut tx_last_response = Some(tx_last_response); + let collect_output_items = + last_response_tx.is_some() || inference_trace_attempt.needs_output_items(); let mut items_added: Vec = Vec::new(); let mut active_item_id: Option = None; let (request_start, mut ttft_ms) = (Instant::now(), None); @@ -2690,7 +2695,9 @@ where ); } active_item_id = None; - items_added.push(item.clone()); + if collect_output_items { + items_added.push(item.clone()); + } if tx_event .send(Ok(ResponseEvent::OutputItemDone(item))) .await @@ -2720,12 +2727,13 @@ where &token_usage, &items_added, ); - if let Some(sender) = tx_last_response.take() { + if let Some(sender) = last_response_tx.take() { let _ = sender.send(LastResponse { response_id: response_id.clone(), - items_added: std::mem::take(&mut items_added), + items_added, }); } + items_added = Vec::new(); if tx_event .send(Ok(ResponseEvent::Completed { response_id, @@ -2791,14 +2799,11 @@ where ); }); - ( - ResponseStream { - rx_event, - interrupt: None, - consumer_dropped: consumer_dropped_for_stream, - }, - rx_last_response, - ) + ResponseStream { + rx_event, + interrupt: None, + consumer_dropped: consumer_dropped_for_stream, + } } /// Handles a 401 response by optionally refreshing ChatGPT tokens once. diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index a08faf802c..e5dd102416 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -1657,6 +1657,34 @@ async fn summarize_memories_returns_empty_for_empty_input() { assert_eq!(output.len(), 0); } +#[tokio::test] +async fn untraced_http_stream_releases_completed_output_before_response_ends() { + // This item's shared payload lets us observe whether the mapper kept a copy. + let tools = Arc::from([json!({"type": "function", "name": "lookup"})]); + let payload = Arc::downgrade(&tools); + let item = ResponseItem::AdditionalTools { + id: None, + role: "developer".to_string(), + tools, + }; + let events = futures::stream::iter([Ok(ResponseEvent::OutputItemDone(item))]) + .chain(futures::stream::pending()); + let mut stream = super::map_response_events( + /*upstream_request_id*/ None, + events, + test_session_telemetry(), + InferenceTraceAttempt::disabled(), + test_model_provider(), + /*last_response_tx*/ None, + ); + + assert!(matches!( + stream.next().await, + Some(Ok(ResponseEvent::OutputItemDone(_))) + )); + assert_eq!(payload.strong_count(), 0); +} + #[tokio::test] async fn dropped_response_stream_traces_cancelled_partial_output() -> anyhow::Result<()> { let temp = TempDir::new()?; @@ -1669,12 +1697,13 @@ async fn dropped_response_stream_traces_cancelled_partial_output() -> anyhow::Re let item = output_message("1", "partial answer"); let api_stream = futures::stream::iter([Ok(ResponseEvent::OutputItemDone(item))]) .chain(futures::stream::pending()); - let (mut stream, _) = super::map_response_events( + let mut stream = super::map_response_events( /*upstream_request_id*/ None, api_stream, test_session_telemetry(), attempt, test_model_provider(), + /*last_response_tx*/ None, ); let observed = stream @@ -1725,12 +1754,13 @@ async fn response_stream_records_last_model_feedback_ids() { end_turn: Some(true), }), ]); - let (mut stream, _) = super::map_response_events( + let mut stream = super::map_response_events( Some("req-123".to_string()), api_stream, test_session_telemetry(), InferenceTraceAttempt::disabled(), test_model_provider(), + /*last_response_tx*/ None, ); while stream.next().await.is_some() {} @@ -1948,12 +1978,13 @@ async fn dropped_backpressured_response_stream_traces_cancelled_partial_output() notify: Arc::clone(&backpressured_item_yielded), }; - let (stream, _) = super::map_response_events( + let stream = super::map_response_events( /*upstream_request_id*/ None, api_stream, test_session_telemetry(), attempt, test_model_provider(), + /*last_response_tx*/ None, ); // Fill the mapper channel with non-terminal events, then yield one output @@ -2244,7 +2275,8 @@ async fn intercepted_output_reaches_trace_and_websocket_bookkeeping() -> anyhow: })) .await?; drop(tx_event); - let (mut stream, last_response) = super::map_response_stream( + let (last_response_tx, last_response_rx) = tokio::sync::oneshot::channel(); + let mut stream = super::map_response_stream( codex_api::ResponseStream { rx_event, upstream_request_id: None, @@ -2254,6 +2286,7 @@ async fn intercepted_output_reaches_trace_and_websocket_bookkeeping() -> anyhow: attempt, test_model_provider(), vec![Box::new(ReplaceOutput)], + Some(last_response_tx), ); let mut delivered = Vec::new(); while let Some(event) = stream.next().await { @@ -2262,7 +2295,7 @@ async fn intercepted_output_reaches_trace_and_websocket_bookkeeping() -> anyhow: } } assert_eq!(delivered, vec![output_message("1", "transformed")]); - assert_eq!(last_response.await?.items_added, delivered); + assert_eq!(last_response_rx.await?.items_added, delivered); let rollout = replay_bundle(temp.path())?; let payload = rollout .raw_payloads diff --git a/codex-rs/rollout-trace/src/inference.rs b/codex-rs/rollout-trace/src/inference.rs index 7cea4d0f97..aec924605f 100644 --- a/codex-rs/rollout-trace/src/inference.rs +++ b/codex-rs/rollout-trace/src/inference.rs @@ -1,8 +1,9 @@ //! Hot-path helpers for recording upstream inference attempts. //! -//! The model client should not need to know whether rollout tracing is enabled. //! A disabled context records nothing, which keeps one-shot HTTP calls, -//! WebSocket reuse, and retry/fallback attempts on the same code path. +//! WebSocket reuse, and retry/fallback attempts on the same lifecycle path. +//! Transport callers can check whether terminal records need output items +//! before retaining them. use std::fmt::Display; use std::sync::Arc; @@ -142,6 +143,11 @@ impl InferenceTraceAttempt { } } + /// Whether terminal trace records need the completed output items. + pub fn needs_output_items(&self) -> bool { + matches!(&self.state, InferenceTraceAttemptState::Enabled(_)) + } + fn inference_call_id(&self) -> Option<&str> { match &self.state { InferenceTraceAttemptState::Disabled => None,