openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-30 19:18:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Avoid retaining unneeded output items in HTTP response streams (#57563)

Browse Source 
## 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
This commit is contained in:
Brent Traut
2026-10-27 20:42:47 +00:00
committed by copyberry
parent 33d306d457
commit d93071fae8
3 changed files with 73 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/client.rs
+27 -22
View File
@@ -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<Box<dyn codex_extension_api::ModelResponseInterceptor>>,
) -> (ResponseStream, oneshot::Receiver<LastResponse>) {
last_response_tx: Option<oneshot::Sender<LastResponse>>,
) -> 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<S>(
@@ -2639,7 +2643,8 @@ fn map_response_events<S>(
session_telemetry: SessionTelemetry,
inference_trace_attempt: InferenceTraceAttempt,
provider: SharedModelProvider,
) -> (ResponseStream, oneshot::Receiver<LastResponse>)
mut last_response_tx: Option<oneshot::Sender<LastResponse>>,
) -> ResponseStream
where
S: futures::Stream<Item = std::result::Result<ResponseEvent, ApiError>>
+ Unpin
@@ -2648,13 +2653,13 @@ where
{
let (tx_event, rx_event) =
mpsc::channel::<Result<ResponseEvent>>(RESPONSE_STREAM_CHANNEL_CAPACITY);
let (tx_last_response, rx_last_response) = oneshot::channel::<LastResponse>();
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<ResponseItem> = Vec::new();
let mut active_item_id: Option<ResponseItemId> = 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.
codex-rs/core/src/client_tests.rs
+38 -5
View File
@@ -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