openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-13 21:30:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
This commit is contained in:
cassirer-openai
2026-11-13 04:36:03 +00:00
committed by copyberry
parent aa5f492031
commit b8dde1eb40
26 changed files with 599 additions and 442 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/client.rs
+32 -18
View File
@@ -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
+11 -7
View File
@@ -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
+2 -2
View File
@@ -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
+1 -1
View File
@@ -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
+29 -27
View File
@@ -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
+1 -1
View File
@@ -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
+2 -1
View File
@@ -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
+8
View File
@@ -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!(
codex-rs/core/tests/responses_headers.rs
+3 -3
View File
@@ -167,7 +167,7 @@ async fn responses_stream_includes_subagent_header_on_review() {
let mut stream = client_session
.stream(
&prompt,
&Arc::new(prompt),
&model_info,
&session_telemetry,
effort,
@@ -317,7 +317,7 @@ async fn responses_stream_includes_subagent_header_on_other() {
let mut stream = client_session
.stream(
&prompt,
&Arc::new(prompt),
&model_info,
&session_telemetry,
effort,
@@ -452,7 +452,7 @@ async fn responses_respects_model_info_overrides_from_config() {
let mut stream = client_session
.stream(
&prompt,
&Arc::new(prompt),
&model_info,
&session_telemetry,
effort,
codex-rs/core/tests/suite/client.rs
+2 -2
View File
@@ -1880,7 +1880,7 @@ async fn send_request_with_provider(provider: ModelProviderInfo) {
let mut stream = client_session
.stream(
&prompt,
&Arc::new(prompt),
&model_info,
&session_telemetry,
effort,
@@ -3660,7 +3660,7 @@ async fn azure_responses_request_does_not_store_and_preserves_prefixed_item_ids(
let mut stream = client_session
.stream(
&prompt,
&Arc::new(prompt),
&model_info,
&session_telemetry,
effort,
codex-rs/core/tests/suite/client_websockets.rs
+14 -15
View File
@@ -962,7 +962,7 @@ async fn responses_websocket_request_prewarm_traces_logical_request() {
.expect("turn started");
let inference_trace = InferenceTraceContext::enabled(
writer,
Arc::clone(&writer),
harness.thread_id.to_string(),
"turn-1".to_string(),
harness.model_info.slug.clone(),
@@ -991,7 +991,7 @@ async fn responses_websocket_request_prewarm_traces_logical_request() {
}
let attribution = serde_json::json!([{"provenance": {"type": "user"}}]);
prompt.input.push(serde_json::from_value(serde_json::json!({
Arc::make_mut(&mut prompt).input.push(serde_json::from_value(serde_json::json!({
"type": "message", "role": "user", "content": [{"type": "input_text", "text": "follow up"}],
"internal_chat_message_metadata_passthrough": {"content_item_metadata": attribution}
})).unwrap());
@@ -1021,6 +1021,7 @@ async fn responses_websocket_request_prewarm_traces_logical_request() {
assert_eq!(follow_up["previous_response_id"].as_str(), Some("warm-1"));
assert_eq!(follow_up["input"], serde_json::json!([]));
writer.flush().expect("publish the live trace");
let rollout = replay_bundle(trace_dir.path()).expect("replay trace");
let mut reader = codex_rollout_trace::PayloadReader::new(trace_dir.path());
assert_eq!(rollout.inference_calls.len(), 2);
@@ -2555,7 +2556,7 @@ async fn responses_websocket_v2_creates_without_previous_response_id_when_non_in
let mut session = harness.client.new_session();
let prompt_one = prompt_with_input(vec![message_item("hello")]);
let mut prompt_two = prompt_with_input(vec![message_item("hello"), message_item("second")]);
prompt_two.output_schema = output_schema;
Arc::make_mut(&mut prompt_two).output_schema = output_schema;
let mut model_info = harness.model_info.clone();
if let Some(model) = model {
model_info.slug = model.to_string();
@@ -2834,20 +2835,18 @@ fn assistant_message_item(id: &str, text: &str) -> ResponseItem {
}
}
fn prompt_with_input(input: Vec<ResponseItem>) -> Prompt {
let mut prompt = Prompt::default();
prompt.input = input;
prompt.base_instructions.text.clear();
prompt
fn prompt_with_input(input: Vec<ResponseItem>) -> Arc<Prompt> {
prompt_with_input_and_instructions(input, "")
}
fn prompt_with_input_and_instructions(input: Vec<ResponseItem>, instructions: &str) -> Prompt {
let mut prompt = prompt_with_input(input);
fn prompt_with_input_and_instructions(input: Vec<ResponseItem>, instructions: &str) -> Arc<Prompt> {
let mut prompt = Prompt::default();
prompt.input = input;
prompt.base_instructions = BaseInstructions {
text: instructions.to_string(),
provenance: None,
};
prompt
Arc::new(prompt)
}
fn websocket_provider(server: &WebSocketTestServer) -> ModelProviderInfo {
@@ -3271,7 +3270,7 @@ fn assert_counter<const N: usize>(
async fn stream_until_complete(
client_session: &mut ModelClientSession,
harness: &WebsocketTestHarness,
prompt: &Prompt,
prompt: &Arc<Prompt>,
) {
stream_until_complete_with_service_tier(
client_session,
@@ -3285,7 +3284,7 @@ async fn stream_until_complete(
async fn stream_until_complete_with_model_info(
client_session: &mut ModelClientSession,
harness: &WebsocketTestHarness,
prompt: &Prompt,
prompt: &Arc<Prompt>,
model_info: &ModelInfo,
expected_response_id: &str,
) {
@@ -3318,7 +3317,7 @@ async fn stream_until_complete_with_model_info(
async fn stream_until_complete_with_service_tier(
client_session: &mut ModelClientSession,
harness: &WebsocketTestHarness,
prompt: &Prompt,
prompt: &Arc<Prompt>,
service_tier: Option<ServiceTier>,
) {
let responses_metadata = turn_metadata(harness, /*turn_id*/ None);
@@ -3335,7 +3334,7 @@ async fn stream_until_complete_with_service_tier(
async fn stream_until_complete_with_metadata(
client_session: &mut ModelClientSession,
harness: &WebsocketTestHarness,
prompt: &Prompt,
prompt: &Arc<Prompt>,
service_tier: Option<ServiceTier>,
responses_metadata: &CodexResponsesMetadata,
) {
codex-rs/memories/write/src/phase1.rs
+1 -1
View File
@@ -311,7 +311,7 @@ mod job {
prompt.output_schema_strict = true;
let (result, token_usage) = context
.stream_stage_one_prompt(config, &prompt, stage_one_context)
.stream_stage_one_prompt(config, prompt, stage_one_context)
.await?;
let output = StageOneOutput::parse(&result, config.memories.version)?;
codex-rs/memories/write/src/runtime.rs
+2 -2
View File
@@ -319,7 +319,7 @@ impl MemoryStartupContext {
pub(crate) async fn stream_stage_one_prompt(
&self,
config: &Config,
prompt: &Prompt,
prompt: Prompt,
context: &StageOneRequestContext,
) -> anyhow::Result<(String, Option<TokenUsage>)> {
let installation_id = resolve_installation_id(&config.codex_home).await?;
@@ -364,7 +364,7 @@ impl MemoryStartupContext {
.await;
let mut stream = client_session
.stream(
prompt,
&Arc::new(prompt),
&context.model_info,
&context.session_telemetry,
context.reasoning_effort.clone(),
codex-rs/memories/write/src/startup_tests.rs
+1 -1
View File
@@ -725,7 +725,7 @@ async fn memories_startup_phase1_uses_live_thread_service_tier_and_detached_meta
context
.stream_stage_one_prompt(
&test.config,
&codex_core::Prompt::default(),
codex_core::Prompt::default(),
&request_context,
)
.await?;
codex-rs/rollout-trace/README.md
+9
View File
@@ -139,6 +139,15 @@ output settings. For a continuation from an unrecorded WebSocket warmup, the
start event also links the full formatted input so the reducer can resolve the
missing history.
One worker per bundle records events in submission order. For inference starts,
the transport shares its encoded request and Core shares the prepared prompt.
The worker serializes the prompt, compresses both bodies, and writes them before
publishing the start event. Remote compaction uses the same worker to snapshot
the model, instructions, input history, and parallel tool-call setting. Up to
eight jobs can wait for the worker. Use `TraceWriter::flush` to read a live
snapshot. Shutdown waits for the worker to finish accepted jobs, seal the shard,
and write the closing marker. A recording failure prevents the marker.
Each shard is one Zstandard frame (level 3, 8 MiB window, checksum enabled).
Each record contains a little-endian 16-bit ID length, a little-endian 64-bit
body length, the UTF-8 ID, then the JSON bytes. The writer flushes a record
codex-rs/rollout-trace/src/code_cell.rs
+25 -54
View File
@@ -16,7 +16,6 @@ use crate::model::CodeCellRuntimeStatus;
use crate::model::CodexTurnId;
use crate::model::ModelVisibleCallId;
use crate::payload::RawPayloadKind;
use crate::payload::RawPayloadRef;
use crate::raw_event::RawTraceEventContext;
use crate::raw_event::RawTraceEventPayload;
use crate::writer::TraceWriter;
@@ -86,14 +85,16 @@ impl CodeCellTraceContext {
let CodeCellTraceContextState::Enabled(context) = &self.state else {
return;
};
append_with_context_best_effort(
context,
if let Err(err) = context.writer.enqueue_event(
RawTraceEventContext::turn(&context.thread_id, &context.codex_turn_id),
RawTraceEventPayload::CodeCellStarted {
runtime_cell_id: context.runtime_cell_id.clone(),
model_visible_call_id: model_visible_call_id.into(),
source_js: source_js.into(),
},
);
) {
warn!("failed to append rollout trace event: {err:#}");
}
}
/// Records the first response returned by the public code-mode `exec` tool.
@@ -106,12 +107,16 @@ impl CodeCellTraceContext {
let CodeCellTraceContextState::Enabled(context) = &self.state else {
return;
};
append_with_context_best_effort(
context,
RawTraceEventPayload::CodeCellInitialResponse {
runtime_cell_id: context.runtime_cell_id.clone(),
status: code_cell_status_for_runtime_response(response),
response_payload: code_cell_response_payload(context, response),
let runtime_cell_id = context.runtime_cell_id.clone();
let status = code_cell_status_for_runtime_response(response);
let _ = context.writer.event_with_payload(
RawTraceEventContext::turn(&context.thread_id, &context.codex_turn_id),
RawPayloadKind::ToolResult,
&CodeCellResponseTracePayload { response },
move |reference| RawTraceEventPayload::CodeCellInitialResponse {
runtime_cell_id,
status,
response_payload: Some(reference),
},
);
}
@@ -121,12 +126,16 @@ impl CodeCellTraceContext {
let CodeCellTraceContextState::Enabled(context) = &self.state else {
return;
};
append_with_context_best_effort(
context,
RawTraceEventPayload::CodeCellEnded {
runtime_cell_id: context.runtime_cell_id.clone(),
status: code_cell_status_for_runtime_response(response),
response_payload: code_cell_response_payload(context, response),
let runtime_cell_id = context.runtime_cell_id.clone();
let status = code_cell_status_for_runtime_response(response);
let _ = context.writer.event_with_payload(
RawTraceEventContext::turn(&context.thread_id, &context.codex_turn_id),
RawPayloadKind::ToolResult,
&CodeCellResponseTracePayload { response },
move |reference| RawTraceEventPayload::CodeCellEnded {
runtime_cell_id,
status,
response_payload: Some(reference),
},
);
}
@@ -145,41 +154,3 @@ fn code_cell_status_for_runtime_response(response: &RuntimeResponse) -> CodeCell
}
}
}
fn code_cell_response_payload(
context: &EnabledCodeCellTraceContext,
response: &RuntimeResponse,
) -> Option<RawPayloadRef> {
write_json_payload_best_effort(
&context.writer,
RawPayloadKind::ToolResult,
&CodeCellResponseTracePayload { response },
)
}
fn write_json_payload_best_effort(
writer: &TraceWriter,
kind: RawPayloadKind,
payload: &impl Serialize,
) -> Option<RawPayloadRef> {
match writer.write_json_payload(kind, payload) {
Ok(payload_ref) => Some(payload_ref),
Err(err) => {
warn!("failed to write rollout trace payload: {err:#}");
None
}
}
}
fn append_with_context_best_effort(
context: &EnabledCodeCellTraceContext,
payload: RawTraceEventPayload,
) {
let event_context = RawTraceEventContext {
thread_id: Some(context.thread_id.clone()),
codex_turn_id: Some(context.codex_turn_id.clone()),
};
if let Err(err) = context.writer.append_with_context(event_context, payload) {
warn!("failed to append rollout trace event: {err:#}");
}
}
codex-rs/rollout-trace/src/compaction.rs
+51 -93
View File
@@ -12,7 +12,6 @@ use std::sync::atomic::Ordering;
use codex_protocol::models::ResponseItem;
use serde::Serialize;
use tracing::warn;
use crate::inference::serialize_trace_items;
use crate::model::AgentThreadId;
@@ -120,19 +119,47 @@ impl CompactionTraceContext {
matches!(self.state, CompactionTraceContextState::Enabled(_))
}
/// Starts a new upstream attempt and records the exact compact endpoint request.
pub fn start_attempt(&self, request: &impl Serialize) -> CompactionTraceAttempt {
/// Starts an upstream attempt. The recording worker builds the snapshot if tracing is enabled.
pub fn start_attempt<T: Serialize>(
&self,
request: impl FnOnce() -> T + Send + 'static,
) -> CompactionTraceAttempt {
let CompactionTraceContextState::Enabled(context) = &self.state else {
return CompactionTraceAttempt::disabled();
return CompactionTraceAttempt {
state: CompactionTraceAttemptState::Disabled,
};
};
let compaction_request_id = next_compaction_request_id();
let attempt = CompactionTraceAttempt {
state: CompactionTraceAttemptState::Enabled(EnabledCompactionTraceAttempt {
context: context.clone(),
compaction_request_id: next_compaction_request_id(),
compaction_request_id: compaction_request_id.clone(),
}),
};
attempt.record_started(request);
let compaction_id = context.compaction_id.clone();
let thread_id = context.thread_id.clone();
let codex_turn_id = context.codex_turn_id.clone();
let model = context.model.clone();
let provider_name = context.provider_name.clone();
let _ = context.writer.defer_event(
RawTraceEventContext::turn(&thread_id, &codex_turn_id),
move |bundle| {
let request_payload = bundle.write_body(
RawPayloadKind::CompactionRequest,
&serde_json::to_vec(&request())?,
)?;
Ok(RawTraceEventPayload::CompactionRequestStarted {
compaction_id,
compaction_request_id,
thread_id,
codex_turn_id,
model,
provider_name,
request_payload,
})
},
);
attempt
}
@@ -144,67 +171,20 @@ impl CompactionTraceContext {
let CompactionTraceContextState::Enabled(context) = &self.state else {
return;
};
let checkpoint_payload = match context
.writer
.write_json_payload(RawPayloadKind::CompactionCheckpoint, checkpoint)
{
Ok(payload_ref) => payload_ref,
Err(err) => {
warn!("failed to write rollout trace payload: {err:#}");
return;
}
};
let event_context = RawTraceEventContext {
thread_id: Some(context.thread_id.clone()),
codex_turn_id: Some(context.codex_turn_id.clone()),
};
if let Err(err) = context.writer.append_with_context(
event_context,
RawTraceEventPayload::CompactionInstalled {
compaction_id: context.compaction_id.clone(),
let compaction_id = context.compaction_id.clone();
let _ = context.writer.event_with_payload(
RawTraceEventContext::turn(&context.thread_id, &context.codex_turn_id),
RawPayloadKind::CompactionCheckpoint,
checkpoint,
move |checkpoint_payload| RawTraceEventPayload::CompactionInstalled {
compaction_id,
checkpoint_payload,
},
) {
warn!("failed to append rollout trace event: {err:#}");
}
);
}
}
impl CompactionTraceAttempt {
/// Builds an attempt that records nothing.
fn disabled() -> Self {
Self {
state: CompactionTraceAttemptState::Disabled,
}
}
fn record_started(&self, request: &impl Serialize) {
let CompactionTraceAttemptState::Enabled(attempt) = &self.state else {
return;
};
let Some(request_payload) = write_json_payload_best_effort(
&attempt.context.writer,
RawPayloadKind::CompactionRequest,
request,
) else {
return;
};
append_with_context_best_effort(
&attempt.context,
RawTraceEventPayload::CompactionRequestStarted {
compaction_id: attempt.context.compaction_id.clone(),
compaction_request_id: attempt.compaction_request_id.clone(),
thread_id: attempt.context.thread_id.clone(),
codex_turn_id: attempt.context.codex_turn_id.clone(),
model: attempt.context.model.clone(),
provider_name: attempt.context.provider_name.clone(),
request_payload,
},
);
}
/// Records the non-streaming compact endpoint response payload.
///
/// Compaction responses use the same response-item preservation rules as
@@ -215,19 +195,16 @@ impl CompactionTraceAttempt {
return;
};
let response_payload = TracedCompactionCompleted { output_items };
let Some(response_payload) = write_json_payload_best_effort(
&attempt.context.writer,
let ctx = &attempt.context;
let compaction_id = ctx.compaction_id.clone();
let compaction_request_id = attempt.compaction_request_id.clone();
let _ = ctx.writer.event_with_payload(
RawTraceEventContext::turn(&ctx.thread_id, &ctx.codex_turn_id),
RawPayloadKind::CompactionResponse,
&response_payload,
) else {
return;
};
append_with_context_best_effort(
&attempt.context,
RawTraceEventPayload::CompactionRequestCompleted {
compaction_id: attempt.context.compaction_id.clone(),
compaction_request_id: attempt.compaction_request_id.clone(),
move |response_payload| RawTraceEventPayload::CompactionRequestCompleted {
compaction_id,
compaction_request_id,
response_payload,
},
);
@@ -246,8 +223,8 @@ impl CompactionTraceAttempt {
let CompactionTraceAttemptState::Enabled(attempt) = &self.state else {
return;
};
append_with_context_best_effort(
&attempt.context,
let _ = attempt.context.writer.enqueue_event(
RawTraceEventContext::turn(&attempt.context.thread_id, &attempt.context.codex_turn_id),
RawTraceEventPayload::CompactionRequestFailed {
compaction_id: attempt.context.compaction_id.clone(),
compaction_request_id: attempt.compaction_request_id.clone(),
@@ -261,22 +238,3 @@ fn next_compaction_request_id() -> CompactionRequestId {
let ordinal = NEXT_COMPACTION_REQUEST.fetch_add(1, Ordering::Relaxed);
format!("compaction_request:{ordinal}")
}
fn write_json_payload_best_effort(
writer: &TraceWriter,
kind: RawPayloadKind,
payload: &impl Serialize,
) -> Option<crate::RawPayloadRef> {
writer.write_json_payload(kind, payload).ok()
}
fn append_with_context_best_effort(
context: &EnabledCompactionTraceContext,
payload: RawTraceEventPayload,
) {
let event_context = RawTraceEventContext {
thread_id: Some(context.thread_id.clone()),
codex_turn_id: Some(context.codex_turn_id.clone()),
};
let _ = context.writer.append_with_context(event_context, payload);
}
codex-rs/rollout-trace/src/inference.rs
+15 -8
View File
@@ -177,11 +177,11 @@ impl InferenceTraceAttempt {
/// Captures the transport's encoded JSON and the immutable Core prompt that produced it.
/// A continuation from an unrecorded warmup also supplies its full formatted context.
pub fn record_started(
pub fn record_started<P: Serialize + Send + Sync + 'static>(
&self,
outgoing: impl AsRef<[u8]>,
prepared_prompt: &impl Serialize,
warmup_context: impl FnOnce() -> serde_json::Result<Option<Vec<u8>>>,
outgoing: impl AsRef<[u8]> + Send + 'static,
prepared_prompt: Arc<P>,
warmup_context: impl FnOnce() -> serde_json::Result<Option<Vec<u8>>> + Send + 'static,
) {
let InferenceTraceAttemptState::Enabled(attempt) = &self.state else {
return;
@@ -192,14 +192,14 @@ impl InferenceTraceAttempt {
let inference_call_id = attempt.inference_call_id.clone();
let model = ctx.model.clone();
let provider_name = ctx.provider_name.clone();
let start = ctx.writer.record_event(
let start = ctx.writer.defer_event(
RawTraceEventContext::turn(&thread_id, &codex_turn_id),
move |bundle| {
let request_payload =
bundle.write_body(RawPayloadKind::InferenceRequest, outgoing.as_ref())?;
let prepared_prompt_payload = bundle.write_body(
RawPayloadKind::PreparedPrompt,
&serde_json::to_vec(prepared_prompt)?,
&serde_json::to_vec(prepared_prompt.as_ref())?,
)?;
let warmup_context_payload = warmup_context()?
.map(|bytes| bundle.write_body(RawPayloadKind::WarmupContext, &bytes))
@@ -364,7 +364,7 @@ fn record_partial_response(
let ctx = &attempt.context;
let envelope = RawTraceEventContext::turn(&ctx.thread_id, &ctx.codex_turn_id);
if output_items.is_empty() {
let _ = ctx.writer.append_with_context(envelope, event(None));
let _ = ctx.writer.enqueue_event(envelope, event(None));
} else {
let payload = TracedResponseStreamOutput {
response_id: None,
@@ -469,9 +469,12 @@ mod tests {
"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)
});
attempt.record_completed("resp-1", Some("req-1"), &None, &[]);
writer.flush()?;
let rollout = replay_bundle(temp.path())?;
let inference = rollout
.inference_calls
@@ -521,3 +524,7 @@ mod tests {
);
}
}
#[cfg(test)]
#[path = "inference_recording_tests.rs"]
mod recording_tests;
codex-rs/rollout-trace/src/inference_recording_tests.rs
+144
View File
@@ -0,0 +1,144 @@
#![cfg(test)]
use std::collections::BTreeMap;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::mpsc;
use std::thread;
use pretty_assertions::assert_eq;
use serde::Serialize;
use serde::Serializer;
use tempfile::TempDir;
use crate::InferenceTraceAttempt;
use crate::InferenceTraceContext;
use crate::PayloadReader;
use crate::RawPayloadKind;
use crate::RawTraceEvent;
use crate::RawTraceEventContext;
use crate::RawTraceEventPayload;
use crate::TraceWriter;
fn recording(dir: &TempDir) -> anyhow::Result<(Arc<TraceWriter>, InferenceTraceAttempt)> {
let writer = Arc::new(TraceWriter::create(
dir.path(),
"trace".into(),
"rollout".into(),
"root".into(),
)?);
let attempt = InferenceTraceContext::enabled(
Arc::clone(&writer),
"root".into(),
"turn".into(),
"model".into(),
"provider".into(),
)
.start_attempt();
Ok((writer, attempt))
}
#[test]
fn inference_dispatch_can_finish_before_prepared_prompt_serialization() -> anyhow::Result<()> {
struct Prompt {
caller: thread::ThreadId,
release: Mutex<mpsc::Receiver<()>>,
}
impl Serialize for Prompt {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
assert_ne!(thread::current().id(), self.caller);
self.release.lock().unwrap().recv().unwrap();
serializer.serialize_str("prepared")
}
}
let dir = TempDir::new()?;
let (writer, attempt) = recording(&dir)?;
let (release, gate) = mpsc::channel();
let wire = br#"{ "input": [] }"#;
attempt.record_started(
wire,
Arc::new(Prompt {
caller: thread::current().id(),
release: Mutex::new(gate),
}),
|| Ok(None),
);
attempt.record_completed("response", /*upstream_request_id*/ None, &None, &[]);
release.send(())?;
drop(attempt);
drop(writer);
let events = std::fs::read_to_string(dir.path().join("trace.jsonl"))?
.lines()
.map(serde_json::from_str)
.collect::<Result<Vec<RawTraceEvent>, _>>()?;
let [start, end, closed] = &events[..] else {
panic!("shutdown must publish both accepted events before recording closure")
};
assert!(matches!(
closed.payload,
RawTraceEventPayload::RecordingCompleted
));
let RawTraceEventPayload::InferenceStarted {
request_payload,
prepared_prompt_payload,
..
} = &start.payload
else {
panic!("start must precede completion")
};
assert!(matches!(
end.payload,
RawTraceEventPayload::InferenceCompleted { .. }
));
let mut reader = PayloadReader::new(dir.path());
assert_eq!(reader.read(request_payload)?, wire);
assert_eq!(reader.read(prepared_prompt_payload)?, br#""prepared""#);
Ok(())
}
#[test]
fn deferred_serialization_failure_suppresses_the_start_and_closing_marker() -> anyhow::Result<()> {
let dir = TempDir::new()?;
let (writer, attempt) = recording(&dir)?;
// JSON cannot encode an array as an object key.
attempt.record_started(b"{}", Arc::new(BTreeMap::from([([1, 2], true)])), || {
Ok(None)
});
assert!(writer.flush().is_err());
drop(attempt);
drop(writer);
assert_eq!(std::fs::read_to_string(dir.path().join("trace.jsonl"))?, "");
Ok(())
}
#[test]
fn borrowed_response_encoding_failure_preserves_the_queued_start() -> anyhow::Result<()> {
let dir = TempDir::new()?;
let (writer, attempt) = recording(&dir)?;
attempt.record_started(b"{}", Arc::new("prepared"), || Ok(None));
// Handoff succeeds even though JSON cannot represent this response's map key.
writer.event_with_payload(
RawTraceEventContext::turn("root", "turn"),
RawPayloadKind::InferenceResponse,
&BTreeMap::from([([1, 2], true)]),
|_| unreachable!("a failed payload has no reference to publish"),
)?;
assert!(
writer
.flush()
.unwrap_err()
.to_string()
.contains("key must be a string")
);
drop(attempt);
drop(writer);
let log = std::fs::read_to_string(dir.path().join("trace.jsonl"))?;
// Exactly one event survived. Parsing the whole log rejects a spurious closing marker.
let event: RawTraceEvent = serde_json::from_str(&log)?;
assert!(matches!(
event.payload,
RawTraceEventPayload::InferenceStarted { .. }
));
Ok(())
}
codex-rs/rollout-trace/src/protocol_event.rs
-12
View File
@@ -11,7 +11,6 @@
//! variants a compile-time prompt to decide whether the trace should capture
//! them.
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ExecCommandBeginEvent;
use codex_protocol::protocol::ExecCommandEndEvent;
@@ -42,7 +41,6 @@ pub(crate) fn codex_turn_trace_event(
thread_id: AgentThreadId,
default_turn_id: &str,
event: &EventMsg,
store_agent_message: impl FnOnce(&ResponseItem) -> Option<crate::RawPayloadRef>,
) -> Option<CodexTurnTraceEvent> {
match event {
EventMsg::TurnStarted(event) => {
@@ -82,16 +80,6 @@ pub(crate) fn codex_turn_trace_event(
},
})
}
EventMsg::RawResponseItem(event) => Some(CodexTurnTraceEvent {
context_turn_id: default_turn_id.to_string(),
payload: RawTraceEventPayload::HistoryItemAccepted {
item_id: event.item.id()?.to_string(),
agent_message_payload: match &event.item {
item @ ResponseItem::AgentMessage { .. } => Some(store_agent_message(item)?),
_ => None,
},
},
}),
_ => None,
}
}
codex-rs/rollout-trace/src/reducer/tool/mcp_tests.rs
+1
View File
@@ -72,6 +72,7 @@ fn mcp_calls_keep_callable_identity_and_classify_the_backend() -> anyhow::Result
},
},
);
writer.flush()?;
let before = replay_bundle(temp.path())?;
// Namespaced write_stdin calls must neither require terminal arguments
// nor create terminal sessions when those arguments happen to be valid.
codex-rs/rollout-trace/src/reducer/tool/terminal_tests.rs
+1
View File
@@ -538,6 +538,7 @@ fn invalid_write_stdin_arguments_remain_failed_tool_calls() -> anyhow::Result<()
);
dispatch.record_failed("invalid arguments");
writer.flush()?;
let rollout = replay_bundle(temp.path())?;
let call = &rollout.tool_calls["call-stdin"];
assert_eq!(call.execution.status, ExecutionStatus::Failed);
codex-rs/rollout-trace/src/thread.rs
+118 -107
View File
@@ -6,6 +6,7 @@
//! through session code.
use codex_protocol::mcp::CallToolResult;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::AgentStatus;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::SessionSource;
@@ -26,7 +27,6 @@ use crate::ExecutionStatus;
use crate::InferenceTraceContext;
use crate::McpCallTraceContext;
use crate::RawPayloadKind;
use crate::RawPayloadRef;
use crate::RawTraceEventContext;
use crate::RawTraceEventPayload;
use crate::RolloutStatus;
@@ -130,13 +130,28 @@ impl ThreadTraceContext {
root_thread_id,
thread_id: metadata.thread_id.clone(),
};
record_thread_started(&context, metadata);
context.writer.check()?;
// Thread creation reports storage failures before returning the new handle.
let metadata_payload = context
.writer
.write_json_payload(RawPayloadKind::SessionMetadata, &metadata)?;
context.writer.append(RawTraceEventPayload::ThreadStarted {
thread_id: metadata.thread_id,
agent_path: metadata.agent_path,
metadata_payload: Some(metadata_payload),
})?;
Ok(Self {
state: ThreadTraceContextState::Enabled(context),
})
}
/// Publishes queued evidence before reading a snapshot of a running bundle.
pub fn flush(&self) -> anyhow::Result<()> {
match &self.state {
ThreadTraceContextState::Disabled => Ok(()),
ThreadTraceContextState::Enabled(context) => context.writer.flush(),
}
}
/// Returns whether this handle will write trace events.
///
/// Most methods have their own disabled fast path. Callers should branch on
@@ -209,15 +224,16 @@ impl ThreadTraceContext {
let Some(event_type) = wrapped_protocol_event_type(event) else {
return;
};
let Some(event_payload) =
context.write_json_payload_best_effort(RawPayloadKind::ProtocolEvent, event)
else {
return;
};
context.append_best_effort(RawTraceEventPayload::ProtocolEventObserved {
event_type: event_type.to_string(),
event_payload,
});
let event_type = event_type.to_owned();
let _ = context.writer.event_with_payload(
RawTraceEventContext::default(),
RawPayloadKind::ProtocolEvent,
event,
move |event_payload| RawTraceEventPayload::ProtocolEventObserved {
event_type,
event_payload,
},
);
}
/// Records turn lifecycle and history acceptance from existing protocol events.
@@ -225,10 +241,30 @@ impl ThreadTraceContext {
let ThreadTraceContextState::Enabled(context) = &self.state else {
return;
};
if let EventMsg::RawResponseItem(event) = event {
let Some(item_id) = event.item.id().map(ToString::to_string) else {
return;
};
let accepted = move |agent_message_payload| RawTraceEventPayload::HistoryItemAccepted {
item_id,
agent_message_payload,
};
let turn = RawTraceEventContext::turn(&context.thread_id, default_turn_id);
if matches!(event.item, ResponseItem::AgentMessage { .. }) {
// Publish the receipt and its history event together, in worker order.
let _ = context.writer.event_with_payload(
turn,
RawPayloadKind::ProtocolEvent,
&event.item,
move |payload| accepted(Some(payload)),
);
} else {
let _ = context.writer.enqueue_event(turn, accepted(None));
}
return;
}
let Some(trace_event) =
codex_turn_trace_event(context.thread_id.clone(), default_turn_id, event, |item| {
context.write_json_payload_best_effort(RawPayloadKind::ProtocolEvent, item)
})
codex_turn_trace_event(context.thread_id.clone(), default_turn_id, event)
else {
return;
};
@@ -250,10 +286,43 @@ impl ThreadTraceContext {
let Some(trace_event) = tool_runtime_trace_event(event) else {
return;
};
let Some(payload) = context.raw_tool_runtime_payload(trace_event) else {
return;
};
context.append_with_context_best_effort(codex_turn_id.into(), payload);
let envelope = RawTraceEventContext::turn(&context.thread_id, &codex_turn_id.into());
match trace_event {
crate::protocol_event::ToolRuntimeTraceEvent::Started {
tool_call_id,
kind,
payload,
} => {
let tool_call_id = tool_call_id.to_owned();
let _ = context.writer.event_with_payload(
envelope,
RawPayloadKind::ToolRuntimeEvent,
&payload,
move |runtime_payload| RawTraceEventPayload::ToolCallRuntimeStarted {
tool_call_id,
kind,
runtime_payload,
},
);
}
crate::protocol_event::ToolRuntimeTraceEvent::Ended {
tool_call_id,
status,
payload,
} => {
let tool_call_id = tool_call_id.to_owned();
let _ = context.writer.event_with_payload(
envelope,
RawPayloadKind::ToolRuntimeEvent,
&payload,
move |runtime_payload| RawTraceEventPayload::ToolCallRuntimeEnded {
tool_call_id,
status,
runtime_payload,
},
);
}
}
}
/// Records the full MCP backend result before callbacks, sanitization, or UI truncation.
@@ -271,20 +340,19 @@ impl ThreadTraceContext {
struct Payload<'a> {
result: &'a Result<CallToolResult, String>,
}
let Some(runtime_payload) = context
.write_json_payload_best_effort(RawPayloadKind::ToolRuntimeEvent, &Payload { result })
else {
return;
let tool_call_id = tool_call_id.to_owned();
let status = if result.is_ok() {
ExecutionStatus::Completed
} else {
ExecutionStatus::Failed
};
context.append_with_context_best_effort(
codex_turn_id.to_owned(),
RawTraceEventPayload::ToolCallRuntimeEnded {
tool_call_id: tool_call_id.to_owned(),
status: if result.is_ok() {
ExecutionStatus::Completed
} else {
ExecutionStatus::Failed
},
let _ = context.writer.event_with_payload(
RawTraceEventContext::turn(&context.thread_id, codex_turn_id),
RawPayloadKind::ToolRuntimeEvent,
&Payload { result },
move |runtime_payload| RawTraceEventPayload::ToolCallRuntimeEnded {
tool_call_id,
status,
runtime_payload,
},
);
@@ -326,20 +394,21 @@ impl ThreadTraceContext {
};
let child_codex_turn_id = child_codex_turn_id.into();
let parent_thread_id = parent_thread_id.into();
let carried_payload =
context.write_json_payload_best_effort(RawPayloadKind::AgentResult, payload);
context.append_with_context_best_effort(
child_codex_turn_id.clone(),
RawTraceEventPayload::AgentResultObserved {
edge_id: format!(
"edge:agent_result:{}:{child_codex_turn_id}:{parent_thread_id}",
context.thread_id
),
child_thread_id: context.thread_id.clone(),
let child_thread_id = context.thread_id.clone();
let message = payload.message.to_owned();
let edge_id =
format!("edge:agent_result:{child_thread_id}:{child_codex_turn_id}:{parent_thread_id}");
let _ = context.writer.event_with_payload(
RawTraceEventContext::turn(&child_thread_id, &child_codex_turn_id),
RawPayloadKind::AgentResult,
payload,
move |reference| RawTraceEventPayload::AgentResultObserved {
edge_id,
child_thread_id,
child_codex_turn_id,
parent_thread_id,
message: payload.message.to_string(),
carried_payload,
message,
carried_payload: Some(reference),
},
);
}
@@ -503,70 +572,12 @@ fn start_root_in_root(
ThreadTraceContext::start(writer, thread_id, metadata)
}
fn record_thread_started(
context: &EnabledThreadTraceContext,
metadata: ThreadStartedTraceMetadata,
) {
let metadata_payload =
context.write_json_payload_best_effort(RawPayloadKind::SessionMetadata, &metadata);
context.append_best_effort(RawTraceEventPayload::ThreadStarted {
thread_id: metadata.thread_id,
agent_path: metadata.agent_path,
metadata_payload,
});
}
impl EnabledThreadTraceContext {
fn write_json_payload_best_effort(
&self,
kind: RawPayloadKind,
payload: &impl Serialize,
) -> Option<RawPayloadRef> {
match self.writer.write_json_payload(kind, payload) {
Ok(payload_ref) => Some(payload_ref),
Err(err) => {
warn!("failed to write rollout trace payload: {err:#}");
None
}
}
}
fn raw_tool_runtime_payload(
&self,
trace_event: crate::protocol_event::ToolRuntimeTraceEvent<'_>,
) -> Option<RawTraceEventPayload> {
match trace_event {
crate::protocol_event::ToolRuntimeTraceEvent::Started {
tool_call_id,
kind,
payload,
} => {
let runtime_payload = self
.write_json_payload_best_effort(RawPayloadKind::ToolRuntimeEvent, &payload)?;
Some(RawTraceEventPayload::ToolCallRuntimeStarted {
tool_call_id: tool_call_id.to_string(),
kind,
runtime_payload,
})
}
crate::protocol_event::ToolRuntimeTraceEvent::Ended {
tool_call_id,
status,
payload,
} => {
let runtime_payload = self
.write_json_payload_best_effort(RawPayloadKind::ToolRuntimeEvent, &payload)?;
Some(RawTraceEventPayload::ToolCallRuntimeEnded {
tool_call_id: tool_call_id.to_string(),
status,
runtime_payload,
})
}
}
}
fn append_best_effort(&self, payload: RawTraceEventPayload) {
if let Err(err) = self.writer.append(payload) {
if let Err(err) = self
.writer
.enqueue_event(RawTraceEventContext::default(), payload)
{
warn!("failed to append rollout trace event: {err:#}");
}
}
@@ -580,7 +591,7 @@ impl EnabledThreadTraceContext {
thread_id: Some(self.thread_id.clone()),
codex_turn_id: Some(codex_turn_id),
};
if let Err(err) = self.writer.append_with_context(event_context, payload) {
if let Err(err) = self.writer.enqueue_event(event_context, payload) {
warn!("failed to append rollout trace event: {err:#}");
}
}
codex-rs/rollout-trace/src/thread_tests.rs
+17 -3
View File
@@ -56,6 +56,7 @@ fn create_in_root_writes_replayable_lifecycle_events() -> anyhow::Result<()> {
thread_trace.record_ended(RolloutStatus::Completed);
let bundle_dir = single_bundle_dir(temp.path())?;
thread_trace.flush()?;
let replayed = replay_bundle(&bundle_dir)?;
assert_eq!(replayed.status, RolloutStatus::Completed);
@@ -165,6 +166,7 @@ fn spawned_thread_start_appends_to_root_bundle() -> anyhow::Result<()> {
})?;
child_trace.record_ended(RolloutStatus::Completed);
let bundle_dir = single_bundle_dir(temp.path())?;
root_trace.flush()?;
let replayed = replay_bundle(&bundle_dir)?;
assert_eq!(fs::read_dir(temp.path())?.count(), 1);
@@ -262,7 +264,11 @@ fn disabled_thread_context_accepts_trace_calls_without_writing() -> anyhow::Resu
thread_trace.inference_trace_context("turn-1", "gpt-test", "test-provider");
let inference_attempt = inference_trace.start_attempt();
let input = serde_json::json!({ "kind": "inference" });
inference_attempt.record_started(serde_json::to_vec(&input)?, &input, || Ok(None));
inference_attempt.record_started(
serde_json::to_vec(&input)?,
std::sync::Arc::new(input),
|| Ok(None),
);
let token_usage: Option<codex_protocol::protocol::TokenUsage> = None;
inference_attempt.record_completed("response-1", Some("req-1"), &token_usage, &[]);
inference_attempt.record_failed("inference failed", /*upstream_request_id*/ None, &[]);
@@ -275,7 +281,7 @@ fn disabled_thread_context_accepts_trace_calls_without_writing() -> anyhow::Resu
);
assert!(!compaction_trace.is_enabled());
let compaction_attempt =
compaction_trace.start_attempt(&serde_json::json!({ "kind": "compaction" }));
compaction_trace.start_attempt::<()>(|| panic!("disabled tracing built the request"));
compaction_attempt.record_completed(&[]);
compaction_attempt.record_failed("compaction failed");
compaction_trace.record_installed(&CompactionCheckpointTracePayload {
@@ -308,11 +314,16 @@ fn compaction_contexts_share_identity_across_models() -> anyhow::Result<()> {
let compaction_trace =
thread_trace.compaction_trace_context("turn-1", "compaction-1", model, "test-provider");
assert!(compaction_trace.is_enabled());
let caller = std::thread::current().id();
compaction_trace
.start_attempt(&serde_json::json!({ "model": model }))
.start_attempt(move || {
assert_ne!(std::thread::current().id(), caller);
serde_json::json!({ "model": model })
})
.record_failed("test failure");
}
thread_trace.flush()?;
let replayed = replay_bundle(&single_bundle_dir(temp.path())?)?;
let mut attempts = replayed
.compaction_requests
@@ -340,6 +351,7 @@ fn protocol_wrapper_records_selected_events_as_raw_payloads() -> anyhow::Result<
thread_trace.record_protocol_event(&EventMsg::ShutdownComplete);
thread_trace.flush()?;
let event_log = fs::read_to_string(single_bundle_dir(temp.path())?.join("trace.jsonl"))?;
let protocol_event_seen = event_log.lines().any(|line| {
let event: crate::RawTraceEvent = serde_json::from_str(line).expect("raw trace event");
@@ -377,6 +389,7 @@ fn intercepted_exec_records_patch_without_a_terminal() -> anyhow::Result<()> {
},
}));
let bundle = single_bundle_dir(temp.path())?;
trace.flush()?;
let before = replay_bundle(&bundle)?.tool_calls["patch-call"].clone();
let changes = HashMap::from([(
PathBuf::from("hello.txt"),
@@ -415,6 +428,7 @@ fn intercepted_exec_records_patch_without_a_terminal() -> anyhow::Result<()> {
},
);
trace.flush()?;
let rollout = replay_bundle(&bundle)?;
let call = &rollout.tool_calls["patch-call"];
assert_eq!(call.kind, ToolCallKind::ApplyPatch);
codex-rs/rollout-trace/src/tool_dispatch.rs
+16 -43
View File
@@ -14,7 +14,6 @@ use codex_protocol::models::SearchToolCallParams;
use serde::Serialize;
use serde_json::Value as JsonValue;
use serde_json::json;
use tracing::warn;
use crate::model::AgentThreadId;
use crate::model::CodeModeRuntimeToolId;
@@ -25,7 +24,6 @@ use crate::model::ToolCallId;
use crate::model::ToolCallKind;
use crate::model::ToolCallSummary;
use crate::payload::RawPayloadKind;
use crate::payload::RawPayloadRef;
use crate::raw_event::RawToolCallRequester;
use crate::raw_event::RawTraceEventContext;
use crate::raw_event::RawTraceEventPayload;
@@ -201,15 +199,16 @@ fn record_started(context: &EnabledToolDispatchTraceContext, invocation: ToolDis
tool_namespace: tool_namespace.as_deref(),
payload: &payload,
};
let request_payload =
write_json_payload_best_effort(&context.writer, RawPayloadKind::ToolInvocation, &request);
let (model_visible_call_id, code_mode_runtime_tool_id, requester) =
requester_fields(invocation.requester);
append_with_context_best_effort(
context,
RawTraceEventPayload::ToolCallStarted {
tool_call_id: context.tool_call_id.clone(),
let tool_call_id = context.tool_call_id.clone();
let _ = context.writer.event_with_payload(
RawTraceEventContext::turn(&context.thread_id, &context.codex_turn_id),
RawPayloadKind::ToolInvocation,
&request,
move |reference| RawTraceEventPayload::ToolCallStarted {
tool_call_id,
model_visible_call_id,
code_mode_runtime_tool_id,
requester,
@@ -219,7 +218,7 @@ fn record_started(context: &EnabledToolDispatchTraceContext, invocation: ToolDis
input_preview,
output_preview: None,
},
invocation_payload: request_payload,
invocation_payload: Some(reference),
},
);
}
@@ -340,45 +339,19 @@ fn append_tool_call_ended(
status: ExecutionStatus,
response: &DispatchedToolTraceResponse<'_>,
) {
let response_payload =
write_json_payload_best_effort(&context.writer, RawPayloadKind::ToolResult, response);
append_with_context_best_effort(
context,
RawTraceEventPayload::ToolCallEnded {
tool_call_id: context.tool_call_id.clone(),
let tool_call_id = context.tool_call_id.clone();
let _ = context.writer.event_with_payload(
RawTraceEventContext::turn(&context.thread_id, &context.codex_turn_id),
RawPayloadKind::ToolResult,
response,
move |reference| RawTraceEventPayload::ToolCallEnded {
tool_call_id,
status,
result_payload: response_payload,
result_payload: Some(reference),
},
);
}
fn write_json_payload_best_effort(
writer: &TraceWriter,
kind: RawPayloadKind,
payload: &impl Serialize,
) -> Option<RawPayloadRef> {
match writer.write_json_payload(kind, payload) {
Ok(payload_ref) => Some(payload_ref),
Err(err) => {
warn!("failed to write rollout trace payload: {err:#}");
None
}
}
}
fn append_with_context_best_effort(
context: &EnabledToolDispatchTraceContext,
payload: RawTraceEventPayload,
) {
let event_context = RawTraceEventContext {
thread_id: Some(context.thread_id.clone()),
codex_turn_id: Some(context.codex_turn_id.clone()),
};
if let Err(err) = context.writer.append_with_context(event_context, payload) {
warn!("failed to append rollout trace event: {err:#}");
}
}
#[cfg(test)]
mod tests {
use super::*;
codex-rs/rollout-trace/src/writer.rs
+93 -41
View File
@@ -1,21 +1,22 @@
//! Hot-path trace bundle writer.
//!
//! Writers serialize whole payloads before publishing the events that reference them.
//! Closing the bundle seals the final shard before recording completion.
//! A single worker publishes live events in order while model and tool work proceeds.
//! Shutdown waits for the worker to drain accepted events and seal the bundle.
use std::fmt;
use std::fs::File;
use std::io::BufWriter;
use std::io::Write;
use std::path::Path;
use std::sync::Mutex;
use std::sync::Arc;
use std::sync::OnceLock;
use std::sync::PoisonError;
use std::sync::mpsc;
use std::thread;
use std::time::SystemTime;
use std::time::UNIX_EPOCH;
use anyhow::Context;
use anyhow::Result;
use anyhow::anyhow;
use anyhow::bail;
use serde::Serialize;
@@ -32,15 +33,16 @@ use crate::raw_event::RawTraceEventContext;
use crate::raw_event::RawTraceEventPayload;
use crate::shard::ShardWriter;
/// Serializes access to the event log and its ordered payload shards.
pub struct TraceWriter {
bundle: Mutex<BundleWriter>,
}
// Bound the number of retained prompts before their serialized sizes are known.
const MAX_PENDING_JOBS: usize = 8;
type WriteJob = Box<dyn FnOnce(&mut BundleWriter) + Send>;
impl fmt::Debug for TraceWriter {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("TraceWriter").finish_non_exhaustive()
}
/// Serializes trace writes in submission order. Explicit writes wait until their evidence is readable.
#[derive(Debug)]
pub struct TraceWriter {
jobs: Option<mpsc::SyncSender<WriteJob>>,
worker: Option<thread::JoinHandle<()>>,
failure: Arc<OnceLock<String>>,
}
/// Owns the bundle files until shutdown. Each body is flushed before its event is published.
@@ -51,7 +53,7 @@ pub(crate) struct BundleWriter {
next_seq: u64,
next_payload_ordinal: u64,
// Keep the first error so explicit reads and shutdown can detect an incomplete recording.
failure: OnceLock<String>,
failure: Arc<OnceLock<String>>,
}
impl BundleWriter {
@@ -90,7 +92,7 @@ impl BundleWriter {
) -> Result<RawTraceEvent> {
if matches!(payload, RawTraceEventPayload::RecordingCompleted) {
self.shards.finish()?;
// Timestamp the closing marker after the shard is sealed.
// Closure is observed here, after draining earlier jobs.
wall_time_unix_ms = unix_time_ms();
}
let event = RawTraceEvent {
@@ -133,28 +135,50 @@ impl TraceWriter {
std::fs::write(&manifest_path, serde_json::to_vec_pretty(&manifest)?)
.with_context(|| format!("write trace manifest {}", manifest_path.display()))?;
let bundle = BundleWriter {
let failure = Arc::new(OnceLock::new());
let mut bundle = BundleWriter {
rollout_id: manifest.rollout_id,
events: BufWriter::new(event_log),
shards: ShardWriter::new(payloads_dir),
next_seq: 1,
next_payload_ordinal: 1,
failure: OnceLock::new(),
failure: Arc::clone(&failure),
};
let (jobs, receiver) = mpsc::sync_channel::<WriteJob>(MAX_PENDING_JOBS);
let worker = thread::Builder::new()
.name("rollout-trace".into())
.spawn(move || {
for job in receiver {
job(&mut bundle);
}
})
.context("start rollout trace writer")?;
Ok(Self {
bundle: Mutex::new(bundle),
jobs: Some(jobs),
worker: Some(worker),
failure,
})
}
fn call<T>(&self, write: impl FnOnce(&mut BundleWriter) -> Result<T>) -> Result<T> {
self.bundle
.lock()
.unwrap_or_else(PoisonError::into_inner)
.record(write)
fn submit(&self, job: impl FnOnce(&mut BundleWriter) + Send + 'static) -> Result<()> {
check_recording(&self.failure)?;
self.jobs
.as_ref()
.context("closed rollout trace writer")?
.send(Box::new(job))
.map_err(|_| anyhow!("rollout trace worker stopped"))
}
pub(crate) fn check(&self) -> Result<()> {
self.call(|_| Ok(()))
/// A caller that needs a published reference or a live read barrier waits in FIFO order.
fn call<T: Send + 'static>(
&self,
write: impl FnOnce(&mut BundleWriter) -> Result<T> + Send + 'static,
) -> Result<T> {
let (reply, result) = mpsc::channel();
self.submit(move |bundle| {
let _ = reply.send(bundle.record(write));
})?;
result.recv().context("rollout trace worker stopped")?
}
/// Publishes a complete JSON body and returns its location in the bundle.
@@ -163,7 +187,9 @@ impl TraceWriter {
kind: RawPayloadKind,
value: &impl Serialize,
) -> Result<RawPayloadRef> {
self.call(|bundle| bundle.write_body(kind, &serde_json::to_vec(value)?))
check_recording(&self.failure)?;
let body = serde_json::to_vec(value);
self.call(move |bundle| bundle.write_body(kind, &body?))
}
/// Appends one raw event with no extra envelope context.
@@ -181,38 +207,64 @@ impl TraceWriter {
self.call(move |bundle| bundle.append(context, payload, now))
}
/// Writes related payloads and their event under the same lock.
pub(crate) fn record_event(
/// The worker writes this event's bodies before publishing its references.
pub(crate) fn defer_event(
&self,
context: RawTraceEventContext,
build: impl FnOnce(&mut BundleWriter) -> Result<RawTraceEventPayload>,
build: impl FnOnce(&mut BundleWriter) -> Result<RawTraceEventPayload> + Send + 'static,
) -> Result<()> {
self.call(|bundle| {
let payload = build(bundle)?;
bundle.append(context, payload, unix_time_ms())?;
Ok(())
let now = unix_time_ms();
self.submit(move |bundle| {
let _ = bundle.record(|bundle| {
let payload = build(bundle)?;
bundle.append(context, payload, now)
});
})
}
/// Publishes a JSON body before the event that points to it.
pub(crate) fn enqueue_event(
&self,
context: RawTraceEventContext,
payload: RawTraceEventPayload,
) -> Result<()> {
self.defer_event(context, move |_| Ok(payload))
}
/// Encodes a borrowed observation before returning. The worker handles its
/// bytes or serialization error in order, preserving previously accepted evidence.
pub(crate) fn event_with_payload(
&self,
context: RawTraceEventContext,
kind: RawPayloadKind,
value: &impl Serialize,
event: impl FnOnce(RawPayloadRef) -> RawTraceEventPayload,
event: impl FnOnce(RawPayloadRef) -> RawTraceEventPayload + Send + 'static,
) -> Result<()> {
self.record_event(context, |bundle| {
let reference = bundle.write_body(kind, &serde_json::to_vec(value)?)?;
check_recording(&self.failure)?;
let body = serde_json::to_vec(value);
self.defer_event(context, move |bundle| {
let reference = bundle.write_body(kind, &body?)?;
Ok(event(reference))
})
}
/// Waits for earlier records to become readable, or returns the first recording failure.
pub fn flush(&self) -> Result<()> {
self.call(|_| Ok(()))
}
}
impl Drop for TraceWriter {
fn drop(&mut self) {
if !self.bundle.is_poisoned() && !std::thread::panicking() {
let _ = self.append(RawTraceEventPayload::RecordingCompleted);
// The closing event seals the shard after earlier jobs. Joining makes it visible on return.
if !thread::panicking() {
let _ = self.enqueue_event(
RawTraceEventContext::default(),
RawTraceEventPayload::RecordingCompleted,
);
}
drop(self.jobs.take());
if let Some(worker) = self.worker.take() {
let _ = worker.join();
}
}
}
@@ -355,7 +407,7 @@ mod tests {
.is_err()
);
std::fs::create_dir(temp.path().join("payloads"))?;
assert!(writer.check().is_err());
assert!(writer.flush().is_err());
assert!(
writer
.write_json_payload(RawPayloadKind::ProtocolEvent, &json!({}))
@@ -387,7 +439,7 @@ mod tests {
})
.is_err()
);
assert!(writer.check().is_err());
assert!(writer.flush().is_err());
drop(writer);
assert_eq!(
std::fs::read_to_string(temp.path().join("trace.jsonl"))?,