Move rollout trace recording to a background worker (#61253)

## Why

Recording inference starts synchronously serializes prompts and writes trace payloads before model work can proceed. Defer that work so inference dispatch can continue while the trace is recorded.

## What changed

- Use one worker per trace bundle to write events in submission order, with up to eight pending jobs.
- Share prepared prompts and encoded requests with the worker, and build remote compaction snapshots there. Write payloads before publishing events that reference them.
- Add `TraceWriter::flush` and `ThreadTraceContext::flush` for live reads. Drain accepted jobs and seal the bundle on shutdown; suppress the closing marker after a recording failure.

## Testing

Add tests for deferred prompt serialization, ordered start and completion events, shutdown draining, and serialization failures that preserve earlier evidence without publishing a closing marker. Verify compaction snapshots are built on the worker and skipped when tracing is disabled.

GitOrigin-RevId: 32fd787bb1d281f1814df2f08c73b667628f8add

codex-rs/core/src/client.rs

@@ -373,7 +373,7 @@ struct WebsocketSession {
     /// Owner of the cached state, including before a connection is opened.
     auth_owner: Option<AuthOwner>,
     connection_key: Option<ResponsesConnectionKey>,
-    last_request: Option<ResponsesApiRequest>,
+    last_request: Option<Arc<ResponsesApiRequest>>,
     last_response_rx: Option<oneshot::Receiver<LastResponse>>,
     last_response_from_untraced_warmup: bool,
     connection_reused: StdMutex<bool>,
@@ -1971,7 +1971,7 @@ impl ResponsesHttpExecutor<'_> {
         turn_state: &Arc<OnceLock<String>>,
         expected_auth_owner: Option<AuthOwner>,
         event_sender: Option<&Sender<ProtocolEvent>>,
-        prompt: &Prompt,
+        prompt: &Arc<Prompt>,
         model_info: &ModelInfo,
         session_telemetry: &SessionTelemetry,
         effort: Option<ReasoningEffortConfig>,
@@ -2104,7 +2104,11 @@ impl ResponsesHttpExecutor<'_> {
             .with_telemetry(Some(request_telemetry), Some(sse_telemetry));
             let stream_result = client
                 .stream_request_with_observer(request, options, |wire| {
-                    inference_trace_attempt.record_started(wire, prompt, || Ok(None));
+                    inference_trace_attempt.record_started(
+                        wire.clone(),
+                        Arc::clone(prompt),
+                        || Ok(None),
+                    );
                 })
                 .await;
 
@@ -2181,7 +2185,7 @@ impl ModelClientSession {
     )]
     async fn stream_responses_websocket(
         &mut self,
-        prompt: &Prompt,
+        prompt: &Arc<Prompt>,
         model_info: &ModelInfo,
         session_telemetry: &SessionTelemetry,
         effort: Option<ReasoningEffortConfig>,
@@ -2388,6 +2392,7 @@ impl ModelClientSession {
                     .prepare_response_items_for_request(&mut request.input);
                 Some(original_item_ids)
             };
+            let mut request = Arc::new(request);
             let mut ws_payload = ResponseCreateWsRequest {
                 previous_response_id,
                 input: incremental_items.as_deref().unwrap_or(&request.input),
@@ -2396,7 +2401,7 @@ impl ModelClientSession {
                     Some(client_metadata),
                     request_trace.as_ref(),
                 ),
-                ..ResponseCreateWsRequest::from(&request)
+                ..ResponseCreateWsRequest::from(request.as_ref())
             };
             self.client.set_guardian_metadata(
                 &mut ws_payload.client_metadata,
@@ -2426,7 +2431,8 @@ impl ModelClientSession {
             // only [B] (or []) with previous_response_id pointing to the warmup.
             // Replay has no inference row for that ID. Keep the delta as sent and
             // supply the full provider-formatted input so replay can recover A.
-            let warmup_request = previous_response_id_from_untraced_warmup.then_some(&request);
+            let warmup_request =
+                previous_response_id_from_untraced_warmup.then(|| Arc::clone(&request));
 
             let websocket_connection =
                 self.websocket_session.connection.as_ref().ok_or_else(|| {
@@ -2455,20 +2461,28 @@ impl ModelClientSession {
                     self.websocket_session.connection_reused(),
                     Some(Arc::clone(&self.turn_state)),
                     |wire| {
-                        inference_trace_attempt.record_started(wire, prompt, move || {
-                            #[derive(serde::Serialize)]
-                            struct WarmupContext<'a> {
-                                input: &'a [ResponseItem],
-                            }
-                            warmup_request
-                                .map(|r| serde_json::to_vec(&WarmupContext { input: &r.input }))
-                                .transpose()
-                        })
+                        inference_trace_attempt.record_started(
+                            wire.clone(),
+                            Arc::clone(prompt),
+                            move || {
+                                #[derive(serde::Serialize)]
+                                struct WarmupContext<'a> {
+                                    input: &'a [ResponseItem],
+                                }
+                                warmup_request
+                                    .map(|r| serde_json::to_vec(&WarmupContext { input: &r.input }))
+                                    .transpose()
+                            },
+                        )
                     },
                 )
                 .await;
             if let Some(original_item_ids) = original_item_ids {
-                for (item, original_item_id) in request.input.iter_mut().zip(original_item_ids) {
+                for (item, original_item_id) in Arc::make_mut(&mut request)
+                    .input
+                    .iter_mut()
+                    .zip(original_item_ids)
+                {
                     item.set_id(original_item_id);
                 }
             }
@@ -2548,7 +2562,7 @@ impl ModelClientSession {
     #[allow(clippy::too_many_arguments)]
     pub async fn prewarm_websocket(
         &mut self,
-        prompt: &Prompt,
+        prompt: &Arc<Prompt>,
         model_info: &ModelInfo,
         session_telemetry: &SessionTelemetry,
         effort: Option<ReasoningEffortConfig>,
@@ -2613,7 +2627,7 @@ impl ModelClientSession {
     /// branches.
     pub async fn stream(
         &mut self,
-        prompt: &Prompt,
+        prompt: &Arc<Prompt>,
         model_info: &ModelInfo,
         session_telemetry: &SessionTelemetry,
         effort: Option<ReasoningEffortConfig>,

codex-rs/core/src/client_tests.rs

@@ -977,7 +977,7 @@ fn websocket_incremental_reuse_tracks_raw_result_metadata() -> anyhow::Result<()
         )?;
 
         let mut session = client.new_session();
-        session.websocket_session.last_request = Some(previous);
+        session.websocket_session.last_request = Some(Arc::new(previous));
         let (sender, receiver) = tokio::sync::oneshot::channel();
         sender
             .send(super::LastResponse {
@@ -1028,7 +1028,7 @@ fn websocket_continuation_reports_unavailable_response_state() {
         )
         .expect("build continuation request");
     let mut session = client.new_session();
-    session.websocket_session.last_request = Some(request.clone());
+    session.websocket_session.last_request = Some(Arc::new(request.clone()));
     assert_eq!(
         session.prepare_websocket_request(&request).err(),
         Some(NonIncrementalReason::NoPreviousResponse),
@@ -1092,14 +1092,14 @@ async fn responses_http_preserves_raw_tool_metadata_for_openai_custom_endpoint()
     let output = output_with_tool_result_metadata(ToolResultMetadata::new(&json!({
         "private": "raw-result-metadata",
     })));
-    let prompt = Prompt {
+    let prompt = Arc::new(Prompt {
         input: vec![output.clone()],
         base_instructions: BaseInstructions {
             text: String::new(),
             ..Default::default()
         },
         ..Default::default()
-    };
+    });
     let responses_metadata = test_responses_metadata_for_client(
         &client,
         /*turn_id*/ None,
@@ -1260,7 +1260,7 @@ fn websocket_continuation_reset_reason_survives_failed_reconnect_and_turn_bounda
             )
             .expect("build continuation request");
         let mut session = client.new_session();
-        session.websocket_session.last_request = Some(request);
+        session.websocket_session.last_request = Some(Arc::new(request));
         session.websocket_session.reset(Some(reason));
         session.websocket_session.reset(/*reason*/ None);
         session.websocket_session.reset(Some(later_reason));
@@ -1462,7 +1462,7 @@ fn started_inference_attempt(temp: &TempDir) -> anyhow::Result<InferenceTraceAtt
     })?;
 
     let inference_trace = InferenceTraceContext::enabled(
-        writer,
+        Arc::clone(&writer),
         "thread-root".to_string(),
         "turn-1".to_string(),
         "gpt-test".to_string(),
@@ -1477,7 +1477,11 @@ fn started_inference_attempt(temp: &TempDir) -> anyhow::Result<InferenceTraceAtt
             "content": [{"type": "input_text", "text": "hello"}]
         }],
     });
-    attempt.record_started(serde_json::to_vec(&request)?, &request, || Ok(None));
+    attempt.record_started(serde_json::to_vec(&request)?, Arc::new(request), || {
+        Ok(None)
+    });
+    // Cancellation tests can start replaying as soon as this helper returns.
+    writer.flush()?;
     Ok(attempt)
 }
 

codex-rs/core/src/compact.rs

@@ -335,7 +335,7 @@ async fn run_compact_task_inner_impl(
             turn_context.as_ref(),
             &mut client_session,
             &responses_metadata,
-            &prompt,
+            &Arc::new(prompt),
             compaction_metadata.phase(),
         )
         .await;
@@ -842,7 +842,7 @@ async fn drain_to_completed(
     turn_context: &TurnContext,
     client_session: &mut ModelClientSession,
     responses_metadata: &CodexResponsesMetadata,
-    prompt: &Prompt,
+    prompt: &Arc<Prompt>,
     phase: CompactionPhase,
 ) -> CodexResult<CompactionResponse> {
     let mut includes_reasoning = turn_context.model_info().use_responses_lite;

codex-rs/core/src/compact_remote_v2.rs

@@ -291,7 +291,7 @@ struct RemoteCompactionV2Output {
 async fn run_remote_compaction_request_v2(
     sess: &Session,
     step_context: &StepContext,
-    prompt: &Prompt,
+    prompt: &Arc<Prompt>,
     responses_metadata: &CodexResponsesMetadata,
     transport: &mut CompactionRequestTransport<'_>,
     submitted_audio: Option<&crate::context_manager::rejected_audio::SubmittedFileAudio>,

codex-rs/core/src/compact_remote_v2_attempt.rs

@@ -177,7 +177,7 @@ pub(super) async fn run_remote_compact_v2_attempt(
             .iter()
             .any(|item| matches!(item, ResponseItem::AdditionalTools { .. }));
     let tool_router = &step_context.tool_router;
-    let mut prompt = Prompt {
+    let mut prompt = Arc::new(Prompt {
         input,
         tools: if tools_in_history {
             tool_router
@@ -195,7 +195,7 @@ pub(super) async fn run_remote_compact_v2_attempt(
         output_schema_strict: true,
         prompt_cache_ttl: turn_context.prompt_cache_ttl,
         cyber_access_program: turn_context.cyber_access_program,
-    };
+    });
 
     let mut responses_metadata = match captured_window {
         Some(window) => {
@@ -225,13 +225,20 @@ pub(super) async fn run_remote_compact_v2_attempt(
         response_id,
         token_usage,
     } = loop {
-        prompt.input.push(ResponseItem::CompactionTrigger {});
-        let trace_attempt = compaction_trace.start_attempt(&serde_json::json!({
-            "model": model_info.slug.as_str(),
-            "instructions": prompt.base_instructions.text.as_str(),
-            "input": &prompt.input,
-            "parallel_tool_calls": prompt.parallel_tool_calls,
-        }));
+        Arc::make_mut(&mut prompt)
+            .input
+            .push(ResponseItem::CompactionTrigger {});
+        // A retry can edit its input while the worker still holds this attempt's snapshot.
+        let trace_prompt = Arc::clone(&prompt);
+        let trace_model = model_info.slug.clone();
+        let trace_attempt = compaction_trace.start_attempt(move || {
+            serde_json::json!({
+                "model": trace_model,
+                "instructions": trace_prompt.base_instructions.text.as_str(),
+                "input": &trace_prompt.input,
+                "parallel_tool_calls": trace_prompt.parallel_tool_calls,
+            })
+        });
         let compaction_output_result = run_remote_compaction_request_v2(
             sess,
             step_context,
@@ -246,30 +253,25 @@ pub(super) async fn run_remote_compact_v2_attempt(
                 .as_ref()
                 .map(|output| std::slice::from_ref(&output.compaction_output)),
         );
-        prompt.input.pop();
-        match compaction_output_result {
+        let rejected_id = match compaction_output_result {
             Ok(output) => break output,
-            Err(error) => {
-                let Some(id) = submitted_audio
-                    .take()
-                    .and_then(|audio| audio.rejected_file_id(error.details()))
-                else {
-                    return Err(error);
-                };
-                unavailable.insert(id);
-                rejected_audio::project_compaction_audio(
-                    &mut prompt.input,
-                    &mut prompt_input_metadata,
-                    &unavailable,
-                );
-            }
-        }
+            Err(error) => submitted_audio
+                .take()
+                .and_then(|audio| audio.rejected_file_id(error.details()))
+                .ok_or(error)?,
+        };
+        unavailable.insert(rejected_id);
+        let input = &mut Arc::make_mut(&mut prompt).input;
+        input.pop();
+        rejected_audio::project_compaction_audio(input, &mut prompt_input_metadata, &unavailable);
     };
+    let mut prompt_input = Arc::unwrap_or_clone(prompt).input;
+    prompt_input.pop(); // Exclude the request-only trigger from checkpoint history.
     Ok(RemoteCompactV2Attempt {
         rejected_audio_file_ids: unavailable,
         input_goal_ids,
         trace_input_history,
-        prompt_input: prompt.input,
+        prompt_input,
         prompt_input_metadata,
         compaction_output,
         compaction_response_id: response_id,

codex-rs/core/src/session/startup_prewarm.rs

@@ -441,7 +441,7 @@ async fn schedule_startup_prewarm_inner(
     // Prewarm establishes the request baseline before the first turn can change effort.
     client_session
         .prewarm_websocket(
-            &startup_prompt,
+            &Arc::new(startup_prompt),
             &step_context.settings.model_info,
             &step_context.session_telemetry,
             session

codex-rs/core/src/session/turn.rs

@@ -3000,7 +3000,8 @@ async fn try_run_sampling_request(
         None
     };
     let mut resampling = crate::guardian::ResponseResampling::new(&sess);
-    let mut stream_request = async |prompt: &Prompt| {
+    let prompt = Arc::new(prompt);
+    let mut stream_request = async |prompt: &Arc<Prompt>| {
         client_session
             .stream(
                 prompt,

codex-rs/core/src/tools/tool_dispatch_trace_tests.rs

@@ -121,6 +121,7 @@ async fn dispatch_lifecycle_trace_records_direct_and_code_mode_requesters() -> a
     let temp = TempDir::new()?;
     let (mut session, turn) = make_session_and_context().await;
     attach_test_trace(&mut session, &turn, temp.path())?;
+    let trace = session.services.rollout_thread_trace.clone();
     session.services.rollout_thread_trace.start_code_cell_trace(
         turn.sub_id.as_str(),
         "cell-1",
@@ -164,6 +165,7 @@ async fn dispatch_lifecycle_trace_records_direct_and_code_mode_requesters() -> a
         )
         .await?;
 
+    trace.flush()?;
     let replayed = codex_rollout_trace::replay_bundle(single_bundle_dir(temp.path())?)?;
     assert_eq!(
         replayed.tool_calls["direct-call"].model_visible_call_id,
@@ -214,6 +216,7 @@ async fn dispatch_lifecycle_trace_records_unsupported_tool_failures() -> anyhow:
     let temp = TempDir::new()?;
     let (mut session, turn) = make_session_and_context().await;
     attach_test_trace(&mut session, &turn, temp.path())?;
+    let trace = session.services.rollout_thread_trace.clone();
 
     let registry = ToolRegistry::empty_for_test();
     let session = Arc::new(session);
@@ -234,6 +237,7 @@ async fn dispatch_lifecycle_trace_records_unsupported_tool_failures() -> anyhow:
         .await;
 
     assert!(matches!(result, Err(FunctionCallError::RespondToModel(_))));
+    trace.flush()?;
     let replayed = codex_rollout_trace::replay_bundle(single_bundle_dir(temp.path())?)?;
     let tool_call = &replayed.tool_calls["unsupported-call"];
     assert_eq!(tool_call.execution.status, ExecutionStatus::Failed);
@@ -247,6 +251,7 @@ async fn dispatch_lifecycle_trace_records_incompatible_payload_failures() -> any
     let temp = TempDir::new()?;
     let (mut session, turn) = make_session_and_context().await;
     attach_test_trace(&mut session, &turn, temp.path())?;
+    let trace = session.services.rollout_thread_trace.clone();
 
     let registry = ToolRegistry::with_handler_for_test(Arc::new(TestHandler {
         tool_name: codex_tools::ToolName::plain(codex_code_mode::PUBLIC_TOOL_NAME),
@@ -271,6 +276,7 @@ async fn dispatch_lifecycle_trace_records_incompatible_payload_failures() -> any
         .await;
 
     assert!(matches!(result, Err(FunctionCallError::Fatal(_))));
+    trace.flush()?;
     let replayed = codex_rollout_trace::replay_bundle(single_bundle_dir(temp.path())?)?;
     let tool_call = &replayed.tool_calls["incompatible-call"];
     assert_eq!(tool_call.execution.status, ExecutionStatus::Failed);
@@ -291,6 +297,7 @@ async fn missing_code_mode_wait_traces_only_the_wait_tool_call() -> anyhow::Resu
         session.services.executed_tool_calls.clone(),
     );
     attach_test_trace(&mut session, &turn, temp.path())?;
+    let trace = session.services.rollout_thread_trace.clone();
 
     let registry = ToolRegistry::with_handler_for_test(Arc::new(CodeModeWaitHandler::new(
         /*description_override*/ None, /*parameters_override*/ None,
@@ -318,6 +325,7 @@ async fn missing_code_mode_wait_traces_only_the_wait_tool_call() -> anyhow::Resu
         .dispatch_any_with_state(invocation, /*call_state*/ None)
         .await?;
 
+    trace.flush()?;
     let replayed = codex_rollout_trace::replay_bundle(single_bundle_dir(temp.path())?)?;
     assert_eq!(replayed.code_cells.len(), 0);
     assert!(