From e73744b3f46dcc6786ab08405c2715c898511281 Mon Sep 17 00:00:00 2001 From: cassirer-openai Date: Fri, 13 Nov 2026 01:07:02 +0000 Subject: [PATCH] Store rollout trace payloads in compressed shards (#61222) ## What changed - Replace individual JSON payload files with numbered Zstandard shards in trace schema v3. Add `PayloadReader` to resolve payload references, cache decoded shards, and validate payloads during replay. - Capture exact outgoing HTTP bodies and WebSocket messages alongside prepared prompts. Preserve full formatted context for continuations from unrecorded WebSocket warmups while retaining the actual incremental request. - Flush payload records before publishing references and seal the final shard before recording completion. Replace `recording_complete` with `recording_status` to distinguish unclosed recordings, closed recordings with partial payloads, and complete recordings. Reduce trace bundles with the binary that recorded them. ## Testing Add coverage for shard rotation, oversized records, live reads, cache eviction, and malformed or corrupt payloads. Extend HTTP and WebSocket integration assertions to verify captured requests and prepared prompts, including warmup continuations. GitOrigin-RevId: 40ce3c11383f48e08f52e7120b706c61a449df9c --- codex-rs/Cargo.lock | 2 + codex-rs/app-server/Cargo.toml | 1 + .../app-server/tests/suite/v2/guardian_v2.rs | 9 +- .../app-server/tests/suite/v2/mcp_tool.rs | 7 +- codex-rs/cli/tests/rollout_trace.rs | 79 +++-- codex-rs/codex-api/src/endpoint/responses.rs | 14 +- .../src/endpoint/responses_websocket.rs | 35 +- codex-rs/core/src/client.rs | 35 +- codex-rs/core/src/client_common.rs | 6 +- codex-rs/core/src/client_tests.rs | 14 +- codex-rs/core/src/session/tests.rs | 5 +- .../core/tests/suite/client_websockets.rs | 84 ++++- codex-rs/http-client/src/request.rs | 5 + codex-rs/rollout-trace/Cargo.toml | 1 + codex-rs/rollout-trace/README.md | 33 +- codex-rs/rollout-trace/src/bundle.rs | 5 +- codex-rs/rollout-trace/src/compaction.rs | 9 +- codex-rs/rollout-trace/src/inference.rs | 279 +++++++-------- codex-rs/rollout-trace/src/lib.rs | 4 + .../rollout-trace/src/model/conversation.rs | 3 + codex-rs/rollout-trace/src/model/mod.rs | 23 +- codex-rs/rollout-trace/src/model/session.rs | 2 + codex-rs/rollout-trace/src/payload.rs | 60 +++- codex-rs/rollout-trace/src/raw_event.rs | 30 +- .../rollout-trace/src/reducer/code_cell.rs | 2 +- .../src/reducer/code_cell_tests.rs | 38 +- .../rollout-trace/src/reducer/conversation.rs | 9 +- .../rollout-trace/src/reducer/inference.rs | 24 +- codex-rs/rollout-trace/src/reducer/mod.rs | 77 +++-- .../rollout-trace/src/reducer/test_support.rs | 3 + codex-rs/rollout-trace/src/reducer/thread.rs | 4 +- .../rollout-trace/src/reducer/tool/agents.rs | 12 +- .../src/reducer/tool/terminal.rs | 6 +- .../src/reducer/tool/terminal_tests.rs | 24 +- codex-rs/rollout-trace/src/shard.rs | 324 ++++++++++++++++++ codex-rs/rollout-trace/src/shard_tests.rs | 132 +++++++ codex-rs/rollout-trace/src/thread_tests.rs | 79 ++++- codex-rs/rollout-trace/src/writer.rs | 259 ++++++++------ 38 files changed, 1238 insertions(+), 500 deletions(-) create mode 100644 codex-rs/rollout-trace/src/shard.rs create mode 100644 codex-rs/rollout-trace/src/shard_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 96a5f49128..d50587d71d 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2144,6 +2144,7 @@ dependencies = [ "codex-queue-extension", "codex-rmcp-client", "codex-rollout", + "codex-rollout-trace", "codex-runtime-agent-instance", "codex-sandboxing", "codex-shell-command", @@ -4899,6 +4900,7 @@ dependencies = [ "tempfile", "tracing", "uuid", + "zstd", ] [[package]] diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 959373d773..5e6ecb808a 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -144,6 +144,7 @@ axum = { workspace = true, default-features = false, features = [ "tokio", ] } base64 = { workspace = true } +codex-rollout-trace = { workspace = true } codex-utils-cargo-bin = { workspace = true } core_test_support = { workspace = true } flate2 = { workspace = true } diff --git a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs index 76dee71d43..8393bb17d3 100644 --- a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs +++ b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs @@ -2078,6 +2078,7 @@ async fn guardian_v2_routes_scoped_tool_approvals( let mut denied_results = Vec::new(); for entry in std::fs::read_dir(traces.path())? { let bundle = entry?.path(); + let mut reader = codex_rollout_trace::PayloadReader::new(&bundle); for line in std::fs::read_to_string(bundle.join("trace.jsonl"))?.lines() { let event: Value = serde_json::from_str(line)?; let payload = &event["payload"]; @@ -2086,12 +2087,8 @@ async fn guardian_v2_routes_scoped_tool_approvals( && payload["tool_call_id"] == "guardian-action-0" { assert_eq!(payload["status"], "failed"); - let path = payload["runtime_payload"]["path"] - .as_str() - .expect("MCP runtime payload path"); - let result: Value = - serde_json::from_reader(std::fs::File::open(bundle.join(path))?)?; - denied_results.push(result); + let reference = serde_json::from_value(payload["runtime_payload"].clone())?; + denied_results.push(reader.read_json(&reference)?); } } } diff --git a/codex-rs/app-server/tests/suite/v2/mcp_tool.rs b/codex-rs/app-server/tests/suite/v2/mcp_tool.rs index ca2e06e467..dfc2368555 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_tool.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_tool.rs @@ -1129,11 +1129,8 @@ async fn mcp_tool_call_completion_notification_contains_truncated_large_result() }) .collect(); assert_eq!(completions.len(), 1); - let payload_path = completions[0]["runtime_payload"]["path"] - .as_str() - .expect("MCP runtime payload path"); - let payload: serde_json::Value = - serde_json::from_reader(std::fs::File::open(bundle.join(payload_path))?)?; + let reference = serde_json::from_value(completions[0]["runtime_payload"].clone())?; + let payload = codex_rollout_trace::PayloadReader::new(&bundle).read_json(&reference)?; let expected = json!({ "content": [ {"type": "text", "text": "large-mcp-content-".repeat(DEFAULT_OUTPUT_BYTES_CAP / 8)}, diff --git a/codex-rs/cli/tests/rollout_trace.rs b/codex-rs/cli/tests/rollout_trace.rs index 8f8e106d16..78d5802965 100644 --- a/codex-rs/cli/tests/rollout_trace.rs +++ b/codex-rs/cli/tests/rollout_trace.rs @@ -25,6 +25,7 @@ use core_test_support::streaming_sse::StreamingSseChunk; use core_test_support::streaming_sse::start_streaming_sse_server; use pretty_assertions::assert_eq; use serde_json::json; +use std::path::Path; use tempfile::TempDir; use tokio::sync::oneshot; @@ -64,7 +65,12 @@ async fn enabled_recording_cannot_silently_fail_to_start() -> Result<()> { let mut bundles = std::fs::read_dir(trace_root)?; let bundle = bundles.next().unwrap()?.path(); assert!(bundles.next().is_none()); - assert!(codex_rollout_trace::replay_bundle(bundle)?.recording_complete); + assert_eq!( + reduce_with_recording_binary(&codex, home.path(), &bundle) + .await? + .recording_status, + codex_rollout_trace::RecordingStatus::Complete + ); Ok(()) } @@ -141,23 +147,13 @@ async fn accepted_output_survives_success_and_stream_failure() -> Result<()> { ); assert!(app.shutdown_gracefully().await?.success()); - // Exercise the exported artifact, not just an in-process reducer call. // There is no following request to prove acceptance of this final item. let bundle = std::fs::read_dir(traces.path())?.next().unwrap()?.path(); - let reduced = tokio::process::Command::new(&codex) - .env("CODEX_HOME", home.path()) - .args(["debug", "trace-reduce"]) - .arg(&bundle) - .output() - .await?; - assert!( - reduced.status.success(), - "{}", - String::from_utf8_lossy(&reduced.stderr) + let rollout = reduce_with_recording_binary(&codex, home.path(), &bundle).await?; + assert_eq!( + rollout.recording_status, + codex_rollout_trace::RecordingStatus::Complete ); - let rollout: RolloutTrace = - serde_json::from_reader(std::fs::File::open(bundle.join("state.json"))?)?; - assert!(rollout.recording_complete); let call = rollout.inference_calls.values().next().unwrap(); let output_id = &call.response_item_ids[0]; let output = &rollout.conversation_items[output_id]; @@ -186,6 +182,20 @@ async fn accepted_output_survives_success_and_stream_failure() -> Result<()> { .iter() .find(|request| request.url.path().ends_with("/responses")) .unwrap(); + let mut reader = codex_rollout_trace::PayloadReader::new(&bundle); + assert_eq!( + reader.read(&rollout.raw_payloads[&call.raw_request_payload_id])?, + request.body.as_slice() + ); + let prepared = reader.read_json(&rollout.raw_payloads[&call.prepared_prompt_payload_id])?; + assert!(prepared["base_instructions"]["text"].is_string()); + assert!( + prepared["input"] + .as_array() + .unwrap() + .iter() + .any(|item| item["role"] == "user") + ); assert_eq!( request.headers["x-codex-inference-call-id"].to_str()?, call.inference_call_id @@ -263,20 +273,11 @@ async fn accepted_output_survives_interruption_and_shutdown() -> Result<()> { // Reduce immediately after shutdown: the background stream recorder must // have preserved the accepted output even without response.completed. let bundle = std::fs::read_dir(traces.path())?.next().unwrap()?.path(); - let reduced = tokio::process::Command::new(&codex) - .env("CODEX_HOME", home.path()) - .args(["debug", "trace-reduce"]) - .arg(&bundle) - .output() - .await?; - assert!( - reduced.status.success(), - "{}", - String::from_utf8_lossy(&reduced.stderr) + let rollout = reduce_with_recording_binary(&codex, home.path(), &bundle).await?; + assert_eq!( + rollout.recording_status, + codex_rollout_trace::RecordingStatus::Complete ); - let rollout: RolloutTrace = - serde_json::from_reader(std::fs::File::open(bundle.join("state.json"))?)?; - assert!(rollout.recording_complete); assert_eq!(rollout.inference_calls.len(), 1); let call = rollout.inference_calls.values().next().unwrap(); assert_eq!(call.execution.status, ExecutionStatus::Cancelled); @@ -302,3 +303,25 @@ async fn accepted_output_survives_interruption_and_shutdown() -> Result<()> { assert_eq!(accepted_messages, vec![output_id]); Ok(()) } + +/// Reduce with the executable that recorded the bundle, then read its exported graph. +async fn reduce_with_recording_binary( + codex: &Path, + home: &Path, + bundle: &Path, +) -> Result { + let reduced = tokio::process::Command::new(codex) + .env("CODEX_HOME", home) + .args(["debug", "trace-reduce"]) + .arg(bundle) + .output() + .await?; + assert!( + reduced.status.success(), + "{}", + String::from_utf8_lossy(&reduced.stderr) + ); + Ok(serde_json::from_reader(std::fs::File::open( + bundle.join("state.json"), + )?)?) +} diff --git a/codex-rs/codex-api/src/endpoint/responses.rs b/codex-rs/codex-api/src/endpoint/responses.rs index 9c41a802ea..e60876005d 100644 --- a/codex-rs/codex-api/src/endpoint/responses.rs +++ b/codex-rs/codex-api/src/endpoint/responses.rs @@ -57,6 +57,16 @@ impl ResponsesClient { } } + pub async fn stream_request( + &self, + request: ResponsesApiRequest, + options: ResponsesOptions, + ) -> Result { + self.stream_request_with_observer(request, options, |_| {}) + .await + } + + /// Gives the observer a shared copy of the exact JSON before transport dispatch. #[instrument( name = "responses.stream_request", level = "info", @@ -67,10 +77,11 @@ impl ResponsesClient { api.path = "/responses" ) )] - pub async fn stream_request( + pub async fn stream_request_with_observer( &self, request: ResponsesApiRequest, options: ResponsesOptions, + on_encoded: impl FnOnce(&bytes::Bytes), ) -> Result { let ResponsesOptions { session_id, @@ -83,6 +94,7 @@ impl ResponsesClient { let body = EncodedJsonBody::encode(&request) .map_err(|e| ApiError::Stream(format!("failed to encode responses request: {e}")))?; + on_encoded(body.shared_bytes()); let mut headers = extra_headers; if let Some(ref thread_id) = thread_id { insert_header(&mut headers, "x-client-request-id", thread_id); diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index bced460a6b..20493b6704 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -42,6 +42,7 @@ use tokio::sync::oneshot; use tokio::time::Instant; use tokio_tungstenite::tungstenite::Error as WsError; use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::Utf8Bytes; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::handshake::client::Response; use tokio_tungstenite::tungstenite::protocol::CloseFrame; @@ -283,17 +284,29 @@ impl ResponsesWebsocketConnection { .is_none_or(|stream| stream.pump_task.is_finished()) } + pub async fn stream_request( + &self, + request: ResponsesWsRequest<'_>, + connection_reused: bool, + turn_state: Option>>, + ) -> Result { + self.stream_request_with_observer(request, connection_reused, turn_state, |_| {}) + .await + } + + /// Gives the observer a shared copy of the exact JSON before transport dispatch. #[instrument( name = "responses_websocket.stream_request", level = "info", skip_all, fields(transport = "responses_websocket", api.path = "/responses") )] - pub async fn stream_request( + pub async fn stream_request_with_observer( &self, request: ResponsesWsRequest<'_>, connection_reused: bool, turn_state: Option>>, + on_encoded: impl FnOnce(&bytes::Bytes), ) -> Result { let (tx_event, rx_event) = mpsc::channel::>(1600); @@ -328,7 +341,8 @@ impl ResponsesWebsocketConnection { warmup: ws_request.generate == Some(false), connection_reused, }; - let request_text = serialize_websocket_request(&request)?; + let request_text = Utf8Bytes::from(serialize_websocket_request(&request)?); + on_encoded(request_text.as_ref()); let current_span = Span::current(); tokio::spawn( @@ -827,7 +841,7 @@ fn is_websocket_receive_failure( async fn run_websocket_response_stream( ws_stream: &mut WsStream, tx_event: mpsc::Sender>, - request_text: String, + request_text: Utf8Bytes, idle_timeout: Duration, telemetry: Option>, turn_state: Option<&OnceLock>, @@ -873,7 +887,7 @@ async fn run_websocket_response_stream( "type": "response.interrupt", "response_id": response_id, "mode": "discard_partial_items", - }).to_string(), + }).to_string().into(), idle_timeout, /*telemetry*/ None, timing_log_context.connection_reused, @@ -1104,7 +1118,7 @@ fn safety_buffering_for_event( async fn send_websocket_request( ws_stream: &mut WsStream, - request_text: String, + request_text: Utf8Bytes, idle_timeout: Duration, telemetry: Option<&Arc>, connection_reused: bool, @@ -1116,13 +1130,10 @@ async fn send_websocket_request( return Err(ApiError::WebsocketMessageTooLarge); } let request_start = Instant::now(); - let result = tokio::time::timeout( - idle_timeout, - ws_stream.send(Message::Text(request_text.into())), - ) - .await - .map_err(|_| ApiError::Stream("idle timeout sending websocket request".into())) - .and_then(|result| result.map_err(map_ws_stream_error)); + let result = tokio::time::timeout(idle_timeout, ws_stream.send(Message::Text(request_text))) + .await + .map_err(|_| ApiError::Stream("idle timeout sending websocket request".into())) + .and_then(|result| result.map_err(map_ws_stream_error)); if let Some(t) = telemetry.as_ref() { t.on_ws_request( diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index afe9865f05..86958e5f62 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -2096,14 +2096,17 @@ impl ResponsesHttpExecutor<'_> { } request.input = input; } - inference_trace_attempt.record_started(&request); let client = ApiResponsesClient::new( transport, client_setup.api_provider, client_setup.api_auth, ) .with_telemetry(Some(request_telemetry), Some(sse_telemetry)); - let stream_result = client.stream_request(request, options).await; + let stream_result = client + .stream_request_with_observer(request, options, |wire| { + inference_trace_attempt.record_started(wire, prompt, || Ok(None)); + }) + .await; match stream_result { Ok(stream) => { @@ -2367,13 +2370,6 @@ impl ModelClientSession { } else { inference_trace.start_attempt() }; - if previous_response_id_from_untraced_warmup { - // The transport can reuse an untraced warmup response id and omit the - // already-sent input, but rollout replay needs the logical model-visible - // request rather than the compressed websocket delta. - inference_trace_attempt.record_started(&request); - } - let (previous_response_id, mut incremental_items) = match continuation { Some(continuation) => (Some(continuation.response_id), Some(continuation.items)), None => (None, None), @@ -2426,9 +2422,11 @@ impl ModelClientSession { let ResponsesWsRequest::ResponseCreate(payload) = &mut ws_request; payload.input = input; } - if !previous_response_id_from_untraced_warmup { - inference_trace_attempt.record_started(&ws_request); - } + // A generate=false warmup can send [A]. An inference that reuses it sends + // 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 websocket_connection = self.websocket_session.connection.as_ref().ok_or_else(|| { @@ -2452,10 +2450,21 @@ impl ModelClientSession { ], ); let stream_result = websocket_connection - .stream_request( + .stream_request_with_observer( ws_request, 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() + }) + }, ) .await; if let Some(original_item_ids) = original_item_ids { diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index df995f71a9..6f2fbe5b64 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -18,10 +18,12 @@ use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; -/// API request payload for a single model turn -#[derive(Debug, Clone)] +/// Core inputs to provider request formatting for a single model turn. +/// The rollout trace serializes this shape before transport-specific rewrites. +#[derive(Debug, Clone, serde::Serialize)] pub struct Prompt { /// Conversation context input items. + #[serde(serialize_with = "codex_rollout_trace::serialize_trace_items")] pub input: Vec, /// Tool definitions to inject into this request, empty when supplied by history. diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 5eabca60fa..707231b02b 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -1469,14 +1469,15 @@ fn started_inference_attempt(temp: &TempDir) -> anyhow::Result anyhow::Re rollout.threads["thread-root"].conversation_item_ids, inference.request_item_ids ); - assert_eq!(rollout.raw_payloads.len(), 2); + assert_eq!(rollout.raw_payloads.len(), 3); Ok(()) } @@ -2073,7 +2074,7 @@ async fn dropped_backpressured_response_stream_traces_cancelled_partial_output() assert_eq!(inference.execution.status, ExecutionStatus::Cancelled); assert_eq!(inference.response_item_ids.len(), 1); - assert_eq!(rollout.raw_payloads.len(), 2); + assert_eq!(rollout.raw_payloads.len(), 3); Ok(()) } @@ -2898,8 +2899,9 @@ async fn intercepted_output_reaches_trace_and_websocket_bookkeeping() -> anyhow: .values() .find(|payload| payload.kind == codex_rollout_trace::RawPayloadKind::InferenceResponse) .expect("response trace payload"); - let recorded: serde_json::Value = - serde_json::from_slice(&std::fs::read(temp.path().join(&payload.path))?)?; + let recorded: serde_json::Value = serde_json::from_slice( + codex_rollout_trace::PayloadReader::new(temp.path()).read(payload)?, + )?; assert_eq!(recorded["output_items"], serde_json::to_value(&delivered)?); Ok(()) } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 8e6bdfb309..743b949807 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -9016,7 +9016,10 @@ async fn submission_loop_channel_close_runs_full_thread_teardown() -> anyhow::Re let bundle = std::fs::read_dir(traces.path())?.next().unwrap()?.path(); let rollout = replay_bundle(bundle)?; assert_eq!(rollout.status, RolloutStatus::Aborted); - assert!(rollout.recording_complete); + assert_eq!( + rollout.recording_status, + codex_rollout_trace::RecordingStatus::Complete + ); assert_eq!(1, calls.load(std::sync::atomic::Ordering::SeqCst)); assert_eq!( diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index ae69634c9a..7a6830d46c 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -915,12 +915,13 @@ async fn responses_websocket_request_prewarm_traces_logical_request() { let server = start_websocket_server(vec![vec![ vec![ev_response_created("warm-1"), ev_completed("warm-1")], vec![ev_response_created("resp-1"), ev_completed("resp-1")], + vec![ev_response_created("resp-2"), ev_completed("resp-2")], ]]) .await; let harness = websocket_harness_with_options(&server, /*runtime_metrics_enabled*/ true).await; let mut client_session = harness.client.new_session(); - let prompt = prompt_with_input(vec![message_item("hello")]); + let mut prompt = prompt_with_input_and_instructions(vec![message_item("hello")], "Base rules"); let prewarm_responses_metadata = prewarm_metadata(&harness, /*turn_id*/ None); client_session @@ -989,6 +990,29 @@ 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!({ + "type": "message", "role": "user", "content": [{"type": "input_text", "text": "follow up"}], + "internal_chat_message_metadata_passthrough": {"content_item_metadata": attribution} + })).unwrap()); + let mut next = client_session + .stream( + &prompt, + &harness.model_info, + &harness.session_telemetry, + harness.effort.clone(), + harness.summary, + /*service_tier*/ None, + &responses_metadata, + &inference_trace, + ) + .await + .expect("second websocket stream"); + while let Some(event) = next.next().await { + if matches!(event, Ok(ResponseEvent::Completed { .. })) { + break; + } + } let connection = server.single_connection(); let follow_up = connection .get(1) @@ -998,20 +1022,50 @@ async fn responses_websocket_request_prewarm_traces_logical_request() { assert_eq!(follow_up["input"], serde_json::json!([])); let rollout = replay_bundle(trace_dir.path()).expect("replay trace"); - let inference = rollout - .inference_calls - .values() - .next() - .expect("inference should be present"); - assert_eq!(inference.request_item_ids.len(), 1); - assert_eq!( - rollout.conversation_items[&inference.request_item_ids[0]] - .body - .parts, - vec![ConversationPart::Text { - text: "hello".to_string(), - }], - ); + let mut reader = codex_rollout_trace::PayloadReader::new(trace_dir.path()); + assert_eq!(rollout.inference_calls.len(), 2); + for (index, response_id, previous, wire_items, prepared_items) in + [(1, "resp-1", "warm-1", 0, 1), (2, "resp-2", "resp-1", 1, 2)] + { + let call = rollout + .inference_calls + .values() + .find(|call| call.response_id.as_deref() == Some(response_id)) + .unwrap(); + let outgoing = connection[index].body_json(); + assert_eq!( + reader + .read_json(&rollout.raw_payloads[&call.raw_request_payload_id]) + .unwrap(), + outgoing + ); + assert_eq!(outgoing["previous_response_id"], previous); + assert_eq!(outgoing["input"].as_array().unwrap().len(), wire_items); + let prepared = reader + .read_json(&rollout.raw_payloads[&call.prepared_prompt_payload_id]) + .unwrap(); + assert_eq!(prepared["base_instructions"]["text"], "Base rules"); + assert_eq!( + prepared["input"], + serde_json::to_value(&prompt.input[..prepared_items]).unwrap() + ); + if index == 2 { + assert_eq!( + prepared["input"][1]["internal_chat_message_metadata_passthrough"]["content_item_metadata"], + attribution + ); + assert!(outgoing["input"][0]["internal_chat_message_metadata_passthrough"]["content_item_metadata"].is_null()); + } + assert_eq!(call.request_item_ids.len(), prepared_items + 1); + assert_eq!( + rollout.conversation_items[&call.request_item_ids[1]] + .body + .parts, + vec![ConversationPart::Text { + text: "hello".to_string() + }] + ); + } server.shutdown().await; } diff --git a/codex-rs/http-client/src/request.rs b/codex-rs/http-client/src/request.rs index b21ce6280e..6a96bbcf3d 100644 --- a/codex-rs/http-client/src/request.rs +++ b/codex-rs/http-client/src/request.rs @@ -33,6 +33,11 @@ impl EncodedJsonBody { &self.bytes } + /// Observers can clone this handle to retain the transport JSON without copying it. + pub fn shared_bytes(&self) -> &Bytes { + &self.bytes + } + pub(crate) fn trace_bytes(&self) -> &[u8] { self.trace_bytes.as_ref().unwrap_or(&self.bytes) } diff --git a/codex-rs/rollout-trace/Cargo.toml b/codex-rs/rollout-trace/Cargo.toml index 048026f60a..98793fd99b 100644 --- a/codex-rs/rollout-trace/Cargo.toml +++ b/codex-rs/rollout-trace/Cargo.toml @@ -21,6 +21,7 @@ serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } sha2 = { workspace = true } tracing = { workspace = true } +zstd = { workspace = true } uuid = { workspace = true, features = ["v4"] } [dev-dependencies] diff --git a/codex-rs/rollout-trace/README.md b/codex-rs/rollout-trace/README.md index 010eaddaf6..efcc30ebdf 100644 --- a/codex-rs/rollout-trace/README.md +++ b/codex-rs/rollout-trace/README.md @@ -67,7 +67,7 @@ flowchart TD subgraph Bundle["trace bundle"] Manifest["manifest.json\ntrace_id, rollout_id, root_thread_id"] Events["trace.jsonl\nordered raw event spine"] - Payloads["payloads/*.json\nlarge raw evidence"] + Payloads["payloads/*.zst\nlarge raw evidence"] end Reducer["replay_bundle\ndeterministic offline reducer"] @@ -110,7 +110,7 @@ writer. Disabled contexts accept the same calls and record nothing. When tracing is enabled, a thread can start or resume only if its recorder initializes successfully. A later write failure stops recording. Partial -bundles remain available for replay. The reducer sets `recording_complete` in +bundles remain available for replay. The reducer sets `recording_status` to `complete` in `state.json` when it reads the closing record written by the last trace handle, including when the recorded task was interrupted. @@ -120,10 +120,33 @@ A trace bundle contains: - `manifest.json`: trace identity and bundle metadata. - `trace.jsonl`: append-only raw events ordered by writer-assigned `seq`. -- `payloads/*.json`: raw requests, responses, tool inputs/results, runtime +- `payloads/*.zst`: raw requests, responses, tool inputs/results, runtime events, terminal output, compaction data, and protocol snapshots. - `state.json`: reduced trace for consumers, written by `codex debug trace-reduce`. +Schema v3 stores whole JSON bodies in numbered Zstandard shards. Every reference in +`trace.jsonl` carries the payload ID and shard number. The reducer collects these +references in `state.json` under `raw_payloads`, so the same index serves every +payload kind. `PayloadReader::read` resolves a reference to JSON bytes and caches +the two most recently read shards. An oversized shard occupies the cache alone. +Reduce a bundle with the binary that recorded it. + +For an inference call, `raw_request_payload_id` resolves to the outgoing HTTP +body or WebSocket message, including its `previous_response_id` and incremental +`input`. `prepared_prompt_payload_id` resolves to Core's input to formatting, +including base instructions and their provenance, ordered history, tools and +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. + +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 +before publishing its reference, and rotates before the next record would +exceed 32 MiB uncompressed. An oversized record occupies one shard. The final +frame is sealed before the writer publishes `RecordingCompleted`. A live bundle +can be reduced through its published records with the last frame still open. + `trace_id` identifies this diagnostic artifact. `rollout_id` identifies the Codex rollout/session being observed. Keeping those separate lets us reason about the stored trace without confusing it with the product-level session identity. @@ -143,7 +166,7 @@ By default this writes `/state.json`. Within the recording binary, flowchart LR Model["model-visible payloads\nrequests and response output items"] Runtime["runtime observations\ntool dispatch, terminal output, code-mode JSON"] - RawPayloads["payloads/*.json\nexact evidence"] + RawPayloads["payloads/*.zst\nexact evidence"] Reducer["reducer"] Conversation["ConversationItem\nwhat the model saw"] ToolCall["ToolCall\nruntime tool boundary"] @@ -214,7 +237,7 @@ the same raw event log, payload directory, and reduced `state.json`. The reducer is strict where the raw evidence should be self-consistent: - raw events are replayed in `seq` order; -- payload files must exist before events refer to them; +- payload records must be readable before events refer to them; - reduced object IDs are stable within one replay; - runtime events may be queued until the model-visible source or delivery target has been observed; diff --git a/codex-rs/rollout-trace/src/bundle.rs b/codex-rs/rollout-trace/src/bundle.rs index 2a72a5d6c2..c66ee5ee22 100644 --- a/codex-rs/rollout-trace/src/bundle.rs +++ b/codex-rs/rollout-trace/src/bundle.rs @@ -10,9 +10,10 @@ pub(crate) const RAW_EVENT_LOG_FILE_NAME: &str = "trace.jsonl"; pub(crate) const PAYLOADS_DIR_NAME: &str = "payloads"; /// Conventional file name for the reducer's serialized `RolloutTrace`. pub const REDUCED_STATE_FILE_NAME: &str = "state.json"; -pub(crate) const TRACE_MANIFEST_SCHEMA_VERSION: u32 = 2; +/// On-disk bundle layout described by manifest.json. +pub(crate) const TRACE_MANIFEST_SCHEMA_VERSION: u32 = 3; /// Bump for every change to the consumer-facing `state.json` schema. -pub(crate) const REDUCED_TRACE_SCHEMA_VERSION: u32 = 2; +pub(crate) const REDUCED_TRACE_SCHEMA_VERSION: u32 = 3; /// Manifest stored at the root of a trace bundle. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] diff --git a/codex-rs/rollout-trace/src/compaction.rs b/codex-rs/rollout-trace/src/compaction.rs index a5c18186f8..2377b6738c 100644 --- a/codex-rs/rollout-trace/src/compaction.rs +++ b/codex-rs/rollout-trace/src/compaction.rs @@ -14,7 +14,7 @@ use codex_protocol::models::ResponseItem; use serde::Serialize; use tracing::warn; -use crate::inference::TraceResponseItems; +use crate::inference::serialize_trace_items; use crate::model::AgentThreadId; use crate::model::CodexTurnId; use crate::model::CompactionId; @@ -71,7 +71,8 @@ struct EnabledCompactionTraceAttempt { #[derive(Serialize)] struct TracedCompactionCompleted<'a> { - output_items: TraceResponseItems<'a>, + #[serde(serialize_with = "serialize_trace_items")] + output_items: &'a [ResponseItem], } /// History replacement checkpoint persisted when compaction installs new live history. @@ -213,9 +214,7 @@ impl CompactionTraceAttempt { let CompactionTraceAttemptState::Enabled(attempt) = &self.state else { return; }; - let response_payload = TracedCompactionCompleted { - output_items: TraceResponseItems(output_items), - }; + let response_payload = TracedCompactionCompleted { output_items }; let Some(response_payload) = write_json_payload_best_effort( &attempt.context.writer, RawPayloadKind::CompactionResponse, diff --git a/codex-rs/rollout-trace/src/inference.rs b/codex-rs/rollout-trace/src/inference.rs index 5411420f61..a0d9cbd231 100644 --- a/codex-rs/rollout-trace/src/inference.rs +++ b/codex-rs/rollout-trace/src/inference.rs @@ -18,7 +18,6 @@ use serde::Serialize; use serde::Serializer; use serde::ser::Error; use serde::ser::SerializeSeq; -use serde_json::Value as JsonValue; use uuid::Uuid; use crate::model::AgentThreadId; @@ -78,7 +77,7 @@ enum InferenceTraceAttemptState { struct EnabledInferenceTraceAttempt { context: EnabledInferenceTraceContext, inference_call_id: InferenceCallId, - terminal_recorded: AtomicBool, + awaiting_terminal: AtomicBool, } /// Non-delta response payload saved for completed or interrupted inference streams. @@ -92,7 +91,8 @@ struct TracedResponseStreamOutput<'a> { response_id: Option<&'a str>, upstream_request_id: Option<&'a str>, token_usage: Option<&'a TokenUsage>, - output_items: TraceResponseItems<'a>, + #[serde(serialize_with = "serialize_trace_items")] + output_items: &'a [ResponseItem], } impl InferenceTraceContext { @@ -132,7 +132,7 @@ impl InferenceTraceContext { state: InferenceTraceAttemptState::Enabled(EnabledInferenceTraceAttempt { context: context.clone(), inference_call_id: next_inference_call_id(), - terminal_recorded: AtomicBool::new(false), + awaiting_terminal: AtomicBool::new(false), }), } } @@ -175,34 +175,51 @@ impl InferenceTraceAttempt { headers.insert(INFERENCE_CALL_ID_HEADER, inference_call_id); } - /// Records the request payload replay should treat as the model-visible inference input. - /// - /// This is usually the exact provider request. Callers may instead pass a - /// logical request when the transport omits already-sent input, such as - /// websocket reuse after an untraced warmup response. - pub fn record_started(&self, request: &impl Serialize) { + /// 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( + &self, + outgoing: impl AsRef<[u8]>, + prepared_prompt: &impl Serialize, + warmup_context: impl FnOnce() -> serde_json::Result>>, + ) { let InferenceTraceAttemptState::Enabled(attempt) = &self.state else { return; }; - let Some(request_payload) = write_json_payload_best_effort( - &attempt.context.writer, - RawPayloadKind::InferenceRequest, - request, - ) else { - return; - }; - - append_with_context_best_effort( - &attempt.context, - RawTraceEventPayload::InferenceStarted { - inference_call_id: attempt.inference_call_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, + let ctx = &attempt.context; + let thread_id = ctx.thread_id.clone(); + let codex_turn_id = ctx.codex_turn_id.clone(); + 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( + 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)?, + )?; + let warmup_context_payload = warmup_context()? + .map(|bytes| bundle.write_body(RawPayloadKind::WarmupContext, &bytes)) + .transpose()?; + Ok(RawTraceEventPayload::InferenceStarted { + inference_call_id, + thread_id, + codex_turn_id, + model, + provider_name, + request_payload, + prepared_prompt_payload, + warmup_context_payload, + }) }, ); + // A transport encoding error produces no outgoing request or start row. + attempt + .awaiting_terminal + .store(start.is_ok(), Ordering::Release); } /// Records successful provider completion and serializes the observed output items. @@ -221,22 +238,23 @@ impl InferenceTraceAttempt { let Some(attempt) = self.take_terminal_attempt() else { return; }; - let Some(response_payload) = write_response_payload_best_effort( - attempt, - Some(response_id), + let payload = TracedResponseStreamOutput { + response_id: Some(response_id), upstream_request_id, - token_usage.as_ref(), + token_usage: token_usage.as_ref(), output_items, - ) else { - return; }; - - append_with_context_best_effort( - &attempt.context, - RawTraceEventPayload::InferenceCompleted { - inference_call_id: attempt.inference_call_id.clone(), - response_id: Some(response_id.to_string()), - upstream_request_id: upstream_request_id.map(str::to_string), + let inference_call_id = attempt.inference_call_id.clone(); + let response_id = Some(response_id.to_owned()); + let upstream_request_id = upstream_request_id.map(str::to_owned); + let _ = attempt.context.writer.event_with_payload( + RawTraceEventContext::turn(&attempt.context.thread_id, &attempt.context.codex_turn_id), + RawPayloadKind::InferenceResponse, + &payload, + move |response_payload| RawTraceEventPayload::InferenceCompleted { + inference_call_id, + response_id, + upstream_request_id, response_payload, }, ); @@ -252,23 +270,17 @@ impl InferenceTraceAttempt { let Some(attempt) = self.take_terminal_attempt() else { return; }; - let partial_response_payload = if output_items.is_empty() { - None - } else { - write_response_payload_best_effort( - attempt, - /*response_id*/ None, - upstream_request_id, - /*token_usage*/ None, - output_items, - ) - }; - append_with_context_best_effort( - &attempt.context, - RawTraceEventPayload::InferenceFailed { - inference_call_id: attempt.inference_call_id.clone(), - upstream_request_id: upstream_request_id.map(str::to_string), - error: error.to_string(), + let inference_call_id = attempt.inference_call_id.clone(); + let error = error.to_string(); + let upstream = upstream_request_id.map(str::to_owned); + record_partial_response( + attempt, + upstream_request_id, + output_items, + move |partial_response_payload| RawTraceEventPayload::InferenceFailed { + inference_call_id, + upstream_request_id: upstream, + error, partial_response_payload, }, ); @@ -288,23 +300,17 @@ impl InferenceTraceAttempt { let Some(attempt) = self.take_terminal_attempt() else { return; }; - let partial_response_payload = if output_items.is_empty() { - None - } else { - write_response_payload_best_effort( - attempt, - /*response_id*/ None, - upstream_request_id, - /*token_usage*/ None, - output_items, - ) - }; - append_with_context_best_effort( - &attempt.context, - RawTraceEventPayload::InferenceCancelled { - inference_call_id: attempt.inference_call_id.clone(), - upstream_request_id: upstream_request_id.map(str::to_string), - reason: reason.to_string(), + let inference_call_id = attempt.inference_call_id.clone(); + let reason = reason.to_string(); + let upstream = upstream_request_id.map(str::to_owned); + record_partial_response( + attempt, + upstream_request_id, + output_items, + move |partial_response_payload| RawTraceEventPayload::InferenceCancelled { + inference_call_id, + upstream_request_id: upstream, + reason, partial_response_payload, }, ); @@ -315,88 +321,64 @@ impl InferenceTraceAttempt { InferenceTraceAttemptState::Disabled => return None, InferenceTraceAttemptState::Enabled(attempt) => attempt, }; - if attempt.terminal_recorded.swap(true, Ordering::AcqRel) { - return None; - } - Some(attempt) + attempt + .awaiting_terminal + .swap(false, Ordering::AcqRel) + .then_some(attempt) } } -/// Serializes a response item for trace evidence rather than future request construction. -/// -/// The protocol serializer intentionally omits some readable reasoning content -/// when shaping items for later model requests. Rollout traces need the item as -/// Codex received it, so this helper restores that content in the raw payload. -fn trace_response_item_json(item: &ResponseItem) -> serde_json::Result { - let mut value = serde_json::to_value(item)?; - - if let ResponseItem::Reasoning { - content: Some(content), - .. - } = item - && let JsonValue::Object(object) = &mut value - { - object.insert("content".to_string(), serde_json::to_value(content)?); - } - - Ok(value) -} - -/// Serializes response evidence one item at a time, propagating failures to the writer. -pub(crate) struct TraceResponseItems<'a>(pub &'a [ResponseItem]); - -impl Serialize for TraceResponseItems<'_> { - fn serialize(&self, serializer: S) -> Result { - let mut sequence = serializer.serialize_seq(Some(self.0.len()))?; - for item in self.0 { - let value = trace_response_item_json(item).map_err(S::Error::custom)?; +/// Serializes Core history or received response items for replay, preserving +/// readable reasoning that the outbound request serializer can omit. +pub fn serialize_trace_items( + items: &[ResponseItem], + serializer: S, +) -> Result { + let mut sequence = serializer.serialize_seq(Some(items.len()))?; + for item in items { + if let ResponseItem::Reasoning { + content: Some(content), + .. + } = item + { + let mut value = serde_json::to_value(item).map_err(S::Error::custom)?; + value["content"] = serde_json::to_value(content).map_err(S::Error::custom)?; sequence.serialize_element(&value)?; + } else { + sequence.serialize_element(item)?; } - sequence.end() } + sequence.end() } fn next_inference_call_id() -> InferenceCallId { Uuid::new_v4().to_string() } -fn write_json_payload_best_effort( - writer: &TraceWriter, - kind: RawPayloadKind, - payload: &impl Serialize, -) -> Option { - writer.write_json_payload(kind, payload).ok() -} - -fn write_response_payload_best_effort( +fn record_partial_response( attempt: &EnabledInferenceTraceAttempt, - response_id: Option<&str>, upstream_request_id: Option<&str>, - token_usage: Option<&TokenUsage>, output_items: &[ResponseItem], -) -> Option { - let response_payload = TracedResponseStreamOutput { - response_id, - upstream_request_id, - token_usage, - output_items: TraceResponseItems(output_items), - }; - write_json_payload_best_effort( - &attempt.context.writer, - RawPayloadKind::InferenceResponse, - &response_payload, - ) -} - -fn append_with_context_best_effort( - context: &EnabledInferenceTraceContext, - payload: RawTraceEventPayload, + event: impl FnOnce(Option) -> RawTraceEventPayload + Send + 'static, ) { - 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); + 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)); + } else { + let payload = TracedResponseStreamOutput { + response_id: None, + upstream_request_id, + token_usage: None, + output_items, + }; + let _ = ctx.writer.event_with_payload( + envelope, + RawPayloadKind::InferenceResponse, + &payload, + move |reference| event(Some(reference)), + ); + } } #[cfg(test)] @@ -471,7 +453,7 @@ mod tests { thread_id: "thread-root".to_string(), })?; let context = InferenceTraceContext::enabled( - writer, + Arc::clone(&writer), "thread-root".to_string(), "turn-1".to_string(), "gpt-test".to_string(), @@ -479,14 +461,15 @@ mod tests { ); let attempt = context.start_attempt(); - attempt.record_started(&json!({ + let request = json!({ "model": "gpt-test", "input": [{ "type": "message", "role": "user", "content": [{"type": "input_text", "text": "hello"}] }], - })); + }); + attempt.record_started(serde_json::to_vec(&request)?, &request, || Ok(None)); attempt.record_completed("resp-1", Some("req-1"), &None, &[]); let rollout = replay_bundle(temp.path())?; @@ -501,7 +484,7 @@ mod tests { assert_eq!(inference.codex_turn_id, "turn-1"); assert_eq!(inference.execution.status, ExecutionStatus::Completed); assert_eq!(inference.upstream_request_id, Some("req-1".to_string())); - assert_eq!(rollout.raw_payloads.len(), 2); + assert_eq!(rollout.raw_payloads.len(), 3); Ok(()) } @@ -521,18 +504,20 @@ mod tests { }; let normal = serde_json::to_value(&item).expect("response item serializes"); - let traced = trace_response_item_json(&item).expect("trace item serializes"); + let traced = + serialize_trace_items(std::slice::from_ref(&item), serde_json::value::Serializer) + .expect("trace item serializes"); assert_eq!(normal.get("content"), None); assert_eq!( traced, - json!({ + json!([{ "type": "reasoning", "id": "rs_1", "summary": [{"type": "summary_text", "text": "summary"}], "content": [{"type": "text", "text": "raw reasoning"}], "encrypted_content": "encoded", - }), + }]), ); } } diff --git a/codex-rs/rollout-trace/src/lib.rs b/codex-rs/rollout-trace/src/lib.rs index 1318f70058..eb2c6e5335 100644 --- a/codex-rs/rollout-trace/src/lib.rs +++ b/codex-rs/rollout-trace/src/lib.rs @@ -17,6 +17,7 @@ mod payload; mod protocol_event; mod raw_event; mod reducer; +mod shard; mod thread; mod tool_dispatch; mod writer; @@ -35,10 +36,12 @@ pub use compaction::CompactionTraceContext; pub use inference::InferenceTraceAttempt; /// Shared recorder context for inference attempts within one Codex turn. pub use inference::InferenceTraceContext; +pub use inference::serialize_trace_items; /// Trace-owned MCP execution correlation propagated to bridge request metadata. pub use mcp::McpCallTraceContext; /// Public reduced trace model returned by replay. pub use model::*; +pub use payload::PayloadCapture; /// Stable identifier for one raw payload inside a rollout bundle. pub use payload::RawPayloadId; /// Coarse role labels for raw payload files. @@ -57,6 +60,7 @@ pub use raw_event::RawTraceEventContext; pub use raw_event::RawTraceEventPayload; /// Replay a bundle recorded by this executable into a `RolloutTrace`. pub use reducer::replay_bundle; +pub use shard::PayloadReader; /// Raw payload captured when a child agent reports completion to its parent. pub use thread::AgentResultTracePayload; /// Environment variable that enables local trace-bundle recording. diff --git a/codex-rs/rollout-trace/src/model/conversation.rs b/codex-rs/rollout-trace/src/model/conversation.rs index 1dda12eeae..14511ff001 100644 --- a/codex-rs/rollout-trace/src/model/conversation.rs +++ b/codex-rs/rollout-trace/src/model/conversation.rs @@ -191,7 +191,10 @@ pub struct InferenceCall { /// Runtime tool calls whose model-visible call item came from this response. pub tool_call_ids_started_by_response: Vec, pub usage: Option, + /// Whole outgoing HTTP body or WebSocket message. pub raw_request_payload_id: RawPayloadId, + /// Ordered Core input and tool/output settings, before provider formatting. + pub prepared_prompt_payload_id: RawPayloadId, /// Full upstream response payload. `None` while running or after pre-stream failures. pub raw_response_payload_id: Option, } diff --git a/codex-rs/rollout-trace/src/model/mod.rs b/codex-rs/rollout-trace/src/model/mod.rs index 438b0aca64..0ab80477b2 100644 --- a/codex-rs/rollout-trace/src/model/mod.rs +++ b/codex-rs/rollout-trace/src/model/mod.rs @@ -51,6 +51,18 @@ pub type EdgeId = String; /// Reducer-owned ID for request/log correlation metadata. pub type CorrelationId = String; +/// Whether the recorder sealed the bundle and captured all referenced payload bytes. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RecordingStatus { + /// No closing event was observed. A live or interrupted bundle can have this state. + Unclosed, + /// The event log closed cleanly with one or more intentional capture gaps. + ClosedWithPartialPayloads, + /// The log closed and every referenced payload was captured in full. + Complete, +} + /// Canonical reduced graph for one Codex rollout. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct RolloutTrace { @@ -67,9 +79,7 @@ pub struct RolloutTrace { /// Wall-clock timestamp for terminal rollout status. `None` means running or partial trace. pub ended_at_unix_ms: Option, pub status: RolloutStatus, - /// All recorder handles closed without a recorded write failure; not task success. - #[serde(default)] - pub recording_complete: bool, + pub recording_status: RecordingStatus, pub root_thread_id: AgentThreadId, pub threads: BTreeMap, pub codex_turns: BTreeMap, @@ -88,27 +98,26 @@ pub struct RolloutTrace { pub compaction_requests: BTreeMap, /// Information-flow edges between threads, cells, tools, and runtime resources. pub interaction_edges: BTreeMap, - /// Raw JSON payloads keyed by raw-payload ID. Most point at files outside this object. + /// Index from raw payload IDs to their shard and role. pub raw_payloads: BTreeMap, } impl RolloutTrace { /// Builds an empty reduced trace that a reducer can populate. pub(crate) fn new( - schema_version: u32, trace_id: String, rollout_id: String, root_thread_id: AgentThreadId, started_at_unix_ms: i64, ) -> Self { Self { - schema_version, + schema_version: crate::bundle::REDUCED_TRACE_SCHEMA_VERSION, trace_id, rollout_id, started_at_unix_ms, ended_at_unix_ms: None, status: RolloutStatus::Running, - recording_complete: false, + recording_status: RecordingStatus::Unclosed, root_thread_id, threads: BTreeMap::new(), codex_turns: BTreeMap::new(), diff --git a/codex-rs/rollout-trace/src/model/session.rs b/codex-rs/rollout-trace/src/model/session.rs index 03f336582a..acbf223a06 100644 --- a/codex-rs/rollout-trace/src/model/session.rs +++ b/codex-rs/rollout-trace/src/model/session.rs @@ -55,6 +55,8 @@ pub struct AgentThread { #[serde(rename_all = "snake_case", tag = "type")] pub enum AgentOrigin { Root, + /// Child session metadata was omitted by a capture limit. + Unknown, Spawned { parent_thread_id: AgentThreadId, /// Interaction edge that carried the spawn task. diff --git a/codex-rs/rollout-trace/src/payload.rs b/codex-rs/rollout-trace/src/payload.rs index 5efc7dc3ee..3b883fcbf4 100644 --- a/codex-rs/rollout-trace/src/payload.rs +++ b/codex-rs/rollout-trace/src/payload.rs @@ -15,20 +15,66 @@ pub type RawPayloadId = String; pub struct RawPayloadRef { pub raw_payload_id: RawPayloadId, /// Payload role. This lets details UI choose syntax highlighting and labels - /// without opening the payload file first. + /// without decoding the payload. pub kind: RawPayloadKind, - /// Path relative to the trace bundle root. - /// - /// The writer always materializes payloads as bundle-local files. Keeping - /// this as a plain path avoids exposing storage modes we do not produce. - pub path: String, + /// Shard containing the retained bytes. Absent when the bundle budget omitted the body. + pub shard: Option, + #[serde(default, skip_serializing_if = "PayloadCapture::is_complete")] + pub capture: PayloadCapture, + /// Interleave these values with the stored JSON fragments to recover the original bytes. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub parts: Vec, +} + +/// Describes a configured size limit. An I/O or encoding error still fails the writer. +#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "status", rename_all = "snake_case")] +pub enum PayloadCapture { + #[default] + Complete, + Truncated { + original_bytes: usize, + retained_bytes: usize, + }, + Omitted { + original_bytes: usize, + }, +} + +impl PayloadCapture { + pub fn is_complete(&self) -> bool { + matches!(self, Self::Complete) + } +} + +impl RawPayloadRef { + pub(crate) fn references(&self) -> impl Iterator { + let mut pending = vec![self]; + std::iter::from_fn(move || { + let next = pending.pop()?; + pending.extend(&next.parts); + Some(next) + }) + } + + pub fn is_complete(&self) -> bool { + self.capture.is_complete() && self.parts.iter().all(RawPayloadRef::is_complete) + } } /// Coarse role of a raw payload. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] #[serde(rename_all = "snake_case", tag = "type", content = "value")] pub enum RawPayloadKind { InferenceRequest, + /// Core input and tool/output settings before provider-specific formatting. + PreparedPrompt, + /// One exact item in a prepared prompt or oversized provider input/output array. + ModelItem, + /// An inline image or oversized text part. Losing its bytes preserves its message and position. + ContentValue, + /// Provider-formatted context for replay of a continuation from an untraced warmup. + WarmupContext, /// Full upstream inference response or non-delta response stream summary. InferenceResponse, CompactionRequest, diff --git a/codex-rs/rollout-trace/src/raw_event.rs b/codex-rs/rollout-trace/src/raw_event.rs index 044a8a122f..3fa06eb26f 100644 --- a/codex-rs/rollout-trace/src/raw_event.rs +++ b/codex-rs/rollout-trace/src/raw_event.rs @@ -23,7 +23,7 @@ use serde_json::Value; pub type RawEventSeq = u64; /// Current raw event envelope schema version. -pub(crate) const RAW_TRACE_EVENT_SCHEMA_VERSION: u32 = 2; +pub(crate) const RAW_TRACE_EVENT_SCHEMA_VERSION: u32 = 3; /// One append-only raw trace event. /// @@ -49,6 +49,16 @@ pub struct RawTraceEventContext { pub codex_turn_id: Option, } +impl RawTraceEventContext { + /// Scopes an observation to a model or tool turn. + pub(crate) fn turn(thread_id: &str, codex_turn_id: &str) -> Self { + Self { + thread_id: Some(thread_id.to_owned()), + codex_turn_id: Some(codex_turn_id.to_owned()), + } + } +} + /// Runtime requester as observed at the raw tool boundary. /// /// This intentionally uses runtime-local identifiers. The reducer is the only @@ -112,7 +122,13 @@ pub enum RawTraceEventPayload { codex_turn_id: CodexTurnId, model: String, provider_name: String, + /// Actual outgoing HTTP body or WebSocket message. request_payload: RawPayloadRef, + /// Core prompt before provider formatting. + prepared_prompt_payload: RawPayloadRef, + /// Full provider-formatted context only when the previous response was an untraced warmup. + #[serde(default, skip_serializing_if = "Option::is_none")] + warmup_context_payload: Option, }, InferenceCompleted { inference_call_id: InferenceCallId, @@ -279,9 +295,15 @@ impl RawTraceEventPayload { metadata_payload, .. } => metadata_payload.iter().collect(), RawTraceEventPayload::InferenceStarted { - request_payload, .. - } - | RawTraceEventPayload::InferenceCompleted { + request_payload, + prepared_prompt_payload, + warmup_context_payload, + .. + } => [request_payload, prepared_prompt_payload] + .into_iter() + .chain(warmup_context_payload) + .collect(), + RawTraceEventPayload::InferenceCompleted { response_payload: request_payload, .. } diff --git a/codex-rs/rollout-trace/src/reducer/code_cell.rs b/codex-rs/rollout-trace/src/reducer/code_cell.rs index 742dc7a6b9..f61279a5d2 100644 --- a/codex-rs/rollout-trace/src/reducer/code_cell.rs +++ b/codex-rs/rollout-trace/src/reducer/code_cell.rs @@ -492,7 +492,7 @@ impl TraceReducer { let Some(request_payload) = request_payload else { return Ok(()); }; - let payload = self.read_payload_json(request_payload)?; + let payload = self.payload_reader.read_json(request_payload)?; if payload.get("tool_name").and_then(Value::as_str) != Some("wait") { return Ok(()); } diff --git a/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs b/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs index ecfb263c89..317c647658 100644 --- a/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs @@ -13,6 +13,8 @@ use crate::payload::RawPayloadKind; use crate::raw_event::RawToolCallRequester; use crate::raw_event::RawTraceEventPayload; use crate::reducer::test_support::append_inference_request; +use crate::reducer::test_support::append_inference_start; +use crate::reducer::test_support::append_inference_start_for_thread; use crate::reducer::test_support::create_started_writer; use crate::reducer::test_support::generic_summary; use crate::reducer::test_support::message; @@ -34,14 +36,7 @@ fn code_cell_lifecycle_links_nested_tools_waits_and_outputs() -> anyhow::Result< "input": [message("user", "count files")] }), )?; - writer.append(RawTraceEventPayload::InferenceStarted { - inference_call_id: "inference-1".to_string(), - thread_id: "thread-root".to_string(), - codex_turn_id: "turn-1".to_string(), - model: "gpt-test".to_string(), - provider_name: "test-provider".to_string(), - request_payload: request, - })?; + append_inference_start(&writer, "inference-1", "turn-1", request)?; let response = writer.write_json_payload( RawPayloadKind::InferenceResponse, &json!({ @@ -125,14 +120,7 @@ fn code_cell_lifecycle_links_nested_tools_waits_and_outputs() -> anyhow::Result< }] }), )?; - writer.append(RawTraceEventPayload::InferenceStarted { - inference_call_id: "inference-2".to_string(), - thread_id: "thread-root".to_string(), - codex_turn_id: "turn-2".to_string(), - model: "gpt-test".to_string(), - provider_name: "test-provider".to_string(), - request_payload: followup, - })?; + append_inference_start(&writer, "inference-2", "turn-2", followup)?; let wait_request = writer.write_json_payload( RawPayloadKind::ToolInvocation, &json!({ @@ -368,14 +356,7 @@ fn cancelled_turn_terminates_unfinished_code_cell() -> anyhow::Result<()> { "input": [message("user", "count files")] }), )?; - writer.append(RawTraceEventPayload::InferenceStarted { - inference_call_id: "inference-1".to_string(), - thread_id: "thread-root".to_string(), - codex_turn_id: "turn-1".to_string(), - model: "gpt-test".to_string(), - provider_name: "test-provider".to_string(), - request_payload: request, - })?; + append_inference_start(&writer, "inference-1", "turn-1", request)?; let response = writer.write_json_payload( RawPayloadKind::InferenceResponse, &json!({ @@ -461,14 +442,7 @@ fn runtime_code_cell_ids_can_repeat_across_threads_and_after_resume() -> anyhow: "input": [message("user", "run code")] }), )?; - writer.append(RawTraceEventPayload::InferenceStarted { - inference_call_id: inference_call_id.to_string(), - thread_id: thread_id.to_string(), - codex_turn_id: turn_id.to_string(), - model: "gpt-test".to_string(), - provider_name: "test-provider".to_string(), - request_payload: request, - })?; + append_inference_start_for_thread(&writer, thread_id, turn_id, inference_call_id, request)?; writer.append_with_context( trace_context_for_thread(thread_id, turn_id), RawTraceEventPayload::CodeCellStarted { diff --git a/codex-rs/rollout-trace/src/reducer/conversation.rs b/codex-rs/rollout-trace/src/reducer/conversation.rs index 6f00a5b53d..beb8546e64 100644 --- a/codex-rs/rollout-trace/src/reducer/conversation.rs +++ b/codex-rs/rollout-trace/src/reducer/conversation.rs @@ -49,7 +49,8 @@ impl TraceReducer { { return Ok(()); } - let item = normalize::normalize_model_item(&self.read_payload_json(payload)?, payload)?; + let item = + normalize::normalize_model_item(&self.payload_reader.observed_json(payload)?, payload)?; if item.agent_message.is_none() || item.responses_item_id.as_deref() != Some(wire_id) { bail!("accepted agent message does not match receipt {wire_id}"); } @@ -73,7 +74,7 @@ impl TraceReducer { codex_turn_id: &str, request_payload: &RawPayloadRef, ) -> Result> { - let payload = self.read_payload_json(request_payload)?; + let payload = self.payload_reader.read_json(request_payload)?; let Some(input) = payload.get("input") else { bail!( "inference request payload {} did not contain input", @@ -169,7 +170,7 @@ impl TraceReducer { inference_call_id: &InferenceCallId, response_payload: &RawPayloadRef, ) -> Result> { - let payload = self.read_payload_json(response_payload)?; + let payload = self.payload_reader.read_json(response_payload)?; let Some(output_items) = payload.get("output_items").and_then(Value::as_array) else { bail!( "inference response payload {} did not contain output_items", @@ -345,7 +346,7 @@ impl TraceReducer { compaction_id: &CompactionId, checkpoint_payload: &RawPayloadRef, ) -> Result { - let payload = self.read_payload_json(checkpoint_payload)?; + let payload = self.payload_reader.read_json(checkpoint_payload)?; let input_history = required_array(&payload, "input_history", checkpoint_payload)?; let replacement_history = required_array(&payload, "replacement_history", checkpoint_payload)?; diff --git a/codex-rs/rollout-trace/src/reducer/inference.rs b/codex-rs/rollout-trace/src/reducer/inference.rs index 5becb59ca0..55fa426a47 100644 --- a/codex-rs/rollout-trace/src/reducer/inference.rs +++ b/codex-rs/rollout-trace/src/reducer/inference.rs @@ -5,12 +5,14 @@ use anyhow::Result; use anyhow::bail; +use anyhow::ensure; use super::TraceReducer; use crate::model::ExecutionStatus; use crate::model::ExecutionWindow; use crate::model::InferenceCall; use crate::model::InferenceCallId; +use crate::payload::RawPayloadKind; use crate::payload::RawPayloadRef; use crate::raw_event::RawEventSeq; use crate::raw_event::RawTraceEventPayload; @@ -26,6 +28,8 @@ pub(super) struct StartedInferenceCall { pub(super) model: String, pub(super) provider_name: String, pub(super) request_payload: RawPayloadRef, + pub(super) prepared_prompt_payload: RawPayloadRef, + pub(super) warmup_context_payload: Option, } impl TraceReducer { @@ -50,10 +54,25 @@ impl TraceReducer { ); } + ensure!( + started.request_payload.kind == RawPayloadKind::InferenceRequest + && started.prepared_prompt_payload.kind == RawPayloadKind::PreparedPrompt + && started + .warmup_context_payload + .as_ref() + .is_none_or(|p| p.kind == RawPayloadKind::WarmupContext), + "inference {} has mislabeled request evidence", + started.inference_call_id + ); let inference_call_id = started.inference_call_id.clone(); let thread_id = started.thread_id.clone(); let codex_turn_id = started.codex_turn_id.clone(); - let request_payload = started.request_payload.clone(); + // A warmup has no inference row to follow. Preserve its formatted context for + // the transcript while keeping the actual (often empty) delta available for mocks. + let context = started + .warmup_context_payload + .as_ref() + .unwrap_or(&started.request_payload); let Some(turn) = self.rollout.codex_turns.get(&codex_turn_id) else { bail!( "inference start {inference_call_id} referenced unknown codex turn {codex_turn_id}" @@ -72,7 +91,7 @@ impl TraceReducer { &inference_call_id, &thread_id, &codex_turn_id, - &request_payload, + context, )?; self.thread_mut(&thread_id)?; @@ -99,6 +118,7 @@ impl TraceReducer { tool_call_ids_started_by_response: Vec::new(), usage: None, raw_request_payload_id: started.request_payload.raw_payload_id, + prepared_prompt_payload_id: started.prepared_prompt_payload.raw_payload_id, raw_response_payload_id: None, }, ); diff --git a/codex-rs/rollout-trace/src/reducer/mod.rs b/codex-rs/rollout-trace/src/reducer/mod.rs index d57608ae6d..4710b3452a 100644 --- a/codex-rs/rollout-trace/src/reducer/mod.rs +++ b/codex-rs/rollout-trace/src/reducer/mod.rs @@ -6,23 +6,20 @@ use std::fs::File; use std::io::BufRead; use std::io::BufReader; use std::path::Path; -use std::path::PathBuf; use anyhow::Context; use anyhow::Result; use anyhow::bail; use anyhow::ensure; -use serde::de::IgnoredAny; -use serde_json::Value; use crate::bundle::MANIFEST_FILE_NAME; use crate::bundle::RAW_EVENT_LOG_FILE_NAME; -use crate::bundle::REDUCED_TRACE_SCHEMA_VERSION; use crate::bundle::TRACE_MANIFEST_SCHEMA_VERSION; use crate::bundle::TraceBundleManifest; use crate::model::ExecutionStatus; use crate::model::RolloutTrace; -use crate::payload::RawPayloadRef; +use crate::shard::PayloadReader; + use crate::raw_event::RAW_TRACE_EVENT_SCHEMA_VERSION; use crate::raw_event::RawTraceEvent; use crate::raw_event::RawTraceEventPayload; @@ -48,9 +45,10 @@ use self::inference::StartedInferenceCall; use self::tool::ObservedAgentResultEdge; use self::tool::PendingAgentInteractionEdge; use self::tool::ToolCallStarted; +use crate::RecordingStatus; /// Replays a trace bundle written by this executable into a reduced rollout graph. -/// A valid prefix has `recording_complete = false`; malformed events or missing +/// A valid prefix has `RecordingStatus::Unclosed`. Malformed events or missing /// or invalid referenced payloads return an error. pub fn replay_bundle(bundle_dir: impl AsRef) -> Result { let bundle_dir = bundle_dir.as_ref(); @@ -64,13 +62,12 @@ pub fn replay_bundle(bundle_dir: impl AsRef) -> Result { ); let mut reducer = TraceReducer { rollout: RolloutTrace::new( - REDUCED_TRACE_SCHEMA_VERSION, manifest.trace_id, manifest.rollout_id, manifest.root_thread_id, manifest.started_at_unix_ms, ), - bundle_dir: bundle_dir.to_path_buf(), + payload_reader: PayloadReader::new(bundle_dir), next_conversation_item_ordinal: 1, next_terminal_operation_ordinal: 1, terminal_ids_by_process: BTreeMap::new(), @@ -104,7 +101,7 @@ pub fn replay_bundle(bundle_dir: impl AsRef) -> Result { continue; } ensure!( - !reducer.rollout.recording_complete, + reducer.rollout.recording_status == RecordingStatus::Unclosed, "trace event after recording completed" ); let event: RawTraceEvent = serde_json::from_str(&line) @@ -132,14 +129,10 @@ pub fn replay_bundle(bundle_dir: impl AsRef) -> Result { } reducer.apply_event(event)?; } - for payload in reducer.rollout.raw_payloads.values() { - // IgnoredAny checks JSON structure but skips UTF-8 validation inside strings. - // Read one body as text so opaque payloads get both checks without a JSON tree. - let body = std::fs::read_to_string(reducer.bundle_dir.join(&payload.path)) - .with_context(|| format!("read payload {}", payload.raw_payload_id))?; - serde_json::from_str::(&body) - .with_context(|| format!("invalid payload {}", payload.raw_payload_id))?; - } + reducer.payload_reader.validate( + &reducer.rollout.raw_payloads, + reducer.rollout.recording_status != RecordingStatus::Unclosed, + )?; // Spawn edges prefer the child task message as their target, but a child can // fail before that message is ever reduced. Only after replaying the whole // bundle do we know which spawn deliveries need the child-thread fallback. @@ -209,7 +202,7 @@ pub fn replay_bundle(bundle_dir: impl AsRef) -> Result { struct TraceReducer { rollout: RolloutTrace, - bundle_dir: PathBuf, + payload_reader: PayloadReader, /// First (wall time, sequence) by receiving thread and Responses item ID. first_acceptances: BTreeMap>, /// Delivered items awaiting their first model-visible sighting, keyed by thread and wire ID. @@ -261,27 +254,39 @@ struct TraceReducer { } impl TraceReducer { - fn read_payload_json(&self, payload: &RawPayloadRef) -> Result { - // Reducers keep raw bodies out of the graph, but typed replay sometimes - // needs a small subset of fields to build semantic objects. - let payload_path = self.bundle_dir.join(&payload.path); - let file = File::open(&payload_path) - .with_context(|| format!("open payload {}", payload.raw_payload_id))?; - serde_json::from_reader(BufReader::new(file)) - .with_context(|| format!("parse payload {}", payload.raw_payload_id)) - } - fn apply_event(&mut self, event: RawTraceEvent) -> Result<()> { // Raw payload refs are reducer-wide evidence, not owned by a single // semantic arm. Keep this bookkeeping separate so typed reduction can // stay strict without duplicating payload insertion in every case. - for payload in event.payload.raw_payload_refs() { - self.insert_raw_payload(payload); + for payload in event + .payload + .raw_payload_refs() + .into_iter() + .flat_map(crate::RawPayloadRef::references) + { + let old = self + .rollout + .raw_payloads + .insert(payload.raw_payload_id.clone(), payload.clone()); + ensure!( + old.as_ref().is_none_or(|old| old == payload), + "conflicting references for payload {}", + payload.raw_payload_id + ); } match event.payload { RawTraceEventPayload::RecordingCompleted => { - self.rollout.recording_complete = true; + self.rollout.recording_status = if self + .rollout + .raw_payloads + .values() + .all(|p| p.capture.is_complete()) + { + RecordingStatus::Complete + } else { + RecordingStatus::ClosedWithPartialPayloads + }; } RawTraceEventPayload::RolloutStarted { trace_id, @@ -366,6 +371,8 @@ impl TraceReducer { model, provider_name, request_payload, + prepared_prompt_payload, + warmup_context_payload, } => { self.start_inference_call( event.seq, @@ -377,6 +384,8 @@ impl TraceReducer { model, provider_name, request_payload, + prepared_prompt_payload, + warmup_context_payload, }, )?; } @@ -637,10 +646,4 @@ impl TraceReducer { // must see both before finding their source and outputs in the ordered history. self.flush_pending_code_cell_starts() } - - fn insert_raw_payload(&mut self, payload: &RawPayloadRef) { - self.rollout - .raw_payloads - .insert(payload.raw_payload_id.clone(), payload.clone()); - } } diff --git a/codex-rs/rollout-trace/src/reducer/test_support.rs b/codex-rs/rollout-trace/src/reducer/test_support.rs index bd12e9d652..bf938b04d4 100644 --- a/codex-rs/rollout-trace/src/reducer/test_support.rs +++ b/codex-rs/rollout-trace/src/reducer/test_support.rs @@ -127,6 +127,9 @@ pub(crate) fn append_inference_start_for_thread( request_payload: RawPayloadRef, ) -> anyhow::Result<()> { writer.append(RawTraceEventPayload::InferenceStarted { + prepared_prompt_payload: writer + .write_json_payload(RawPayloadKind::PreparedPrompt, &json!({"input": []}))?, + warmup_context_payload: None, inference_call_id: inference_call_id.to_string(), thread_id: thread_id.to_string(), codex_turn_id: codex_turn_id.to_string(), diff --git a/codex-rs/rollout-trace/src/reducer/thread.rs b/codex-rs/rollout-trace/src/reducer/thread.rs index b36c018fa0..c9b8f2b647 100644 --- a/codex-rs/rollout-trace/src/reducer/thread.rs +++ b/codex-rs/rollout-trace/src/reducer/thread.rs @@ -229,10 +229,10 @@ impl TraceReducer { } fn thread_started_metadata( - &self, + &mut self, metadata_payload: &RawPayloadRef, ) -> Result { - let value = self.read_payload_json(metadata_payload)?; + let value = self.payload_reader.read_json(metadata_payload)?; serde_json::from_value(value) .with_context(|| format!("parse thread metadata {}", metadata_payload.raw_payload_id)) } diff --git a/codex-rs/rollout-trace/src/reducer/tool/agents.rs b/codex-rs/rollout-trace/src/reducer/tool/agents.rs index b395eb15d2..cc3c397bd4 100644 --- a/codex-rs/rollout-trace/src/reducer/tool/agents.rs +++ b/codex-rs/rollout-trace/src/reducer/tool/agents.rs @@ -82,7 +82,7 @@ impl TraceReducer { match kind { ToolCallKind::AssignAgentTask => { let payload: CollabAgentInteractionBeginEvent = - serde_json::from_value(self.read_payload_json(runtime_payload)?)?; + serde_json::from_value(self.payload_reader.read_json(runtime_payload)?)?; self.queue_message_agent_interaction( tool_call_id, InteractionEdgeKind::AssignAgentTask, @@ -93,7 +93,7 @@ impl TraceReducer { } ToolCallKind::SendMessage => { let payload: CollabAgentInteractionBeginEvent = - serde_json::from_value(self.read_payload_json(runtime_payload)?)?; + serde_json::from_value(self.payload_reader.read_json(runtime_payload)?)?; self.queue_message_agent_interaction( tool_call_id, InteractionEdgeKind::SendMessage, @@ -104,7 +104,7 @@ impl TraceReducer { } ToolCallKind::CloseAgent => { let payload: CollabCloseBeginEvent = - serde_json::from_value(self.read_payload_json(runtime_payload)?)?; + serde_json::from_value(self.payload_reader.read_json(runtime_payload)?)?; self.upsert_close_agent_interaction( tool_call_id, payload.receiver_thread_id.to_string(), @@ -135,7 +135,7 @@ impl TraceReducer { if matches!(kind, ToolCallKind::Mcp { .. }) { return Ok(()); } - let runtime_payload_json = self.read_payload_json(runtime_payload)?; + let runtime_payload_json = self.payload_reader.read_json(runtime_payload)?; if runtime_payload_json.get("agent_thread_id").is_some() { let payload: SubAgentActivityEvent = serde_json::from_value(runtime_payload_json)?; return self.end_sub_agent_activity(wall_time_unix_ms, tool_call_id, &kind, &payload); @@ -277,7 +277,7 @@ impl TraceReducer { }) } - fn agent_message_content_from_invocation(&self, tool_call_id: &str) -> Result { + fn agent_message_content_from_invocation(&mut self, tool_call_id: &str) -> Result { let tool_call = self.rollout.tool_calls.get(tool_call_id).with_context(|| { format!("agent activity referenced unknown tool call {tool_call_id}") })?; @@ -296,7 +296,7 @@ impl TraceReducer { "agent activity tool call {tool_call_id} referenced missing invocation payload {invocation_payload_id}" ) })?; - let invocation = self.read_payload_json(invocation_payload)?; + let invocation = self.payload_reader.read_json(invocation_payload)?; let arguments = invocation .get("payload") .and_then(|payload| payload.get("arguments")) diff --git a/codex-rs/rollout-trace/src/reducer/tool/terminal.rs b/codex-rs/rollout-trace/src/reducer/tool/terminal.rs index 536438fc5a..1e0fbfda03 100644 --- a/codex-rs/rollout-trace/src/reducer/tool/terminal.rs +++ b/codex-rs/rollout-trace/src/reducer/tool/terminal.rs @@ -57,7 +57,7 @@ impl TraceReducer { return Ok(None); }; - let payload = self.read_payload_json(invocation_payload)?; + let payload = self.payload_reader.read_json(invocation_payload)?; let Some(request) = parse_dispatch_terminal_request(payload).with_context(|| { format!( "parse terminal invocation payload {} as dispatch payload", @@ -92,7 +92,7 @@ impl TraceReducer { return Ok(None); }; - let payload = self.read_payload_json(runtime_payload)?; + let payload = self.payload_reader.read_json(runtime_payload)?; let payload: ExecCommandBeginPayload = serde_json::from_value(payload).with_context(|| { format!( @@ -173,7 +173,7 @@ impl TraceReducer { }; let response = response_payload .map(|payload| { - let value = self.read_payload_json(payload)?; + let value = self.payload_reader.read_json(payload)?; let response = parse_terminal_response_payload( value, &operation_kind, diff --git a/codex-rs/rollout-trace/src/reducer/tool/terminal_tests.rs b/codex-rs/rollout-trace/src/reducer/tool/terminal_tests.rs index e66773510f..8a7e95f744 100644 --- a/codex-rs/rollout-trace/src/reducer/tool/terminal_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/tool/terminal_tests.rs @@ -22,6 +22,7 @@ use crate::model::ToolCallSummary; use crate::payload::RawPayloadKind; use crate::raw_event::RawTraceEventPayload; use crate::reducer::test_support::ROOT_THREAD_ID; +use crate::reducer::test_support::append_inference_start; use crate::reducer::test_support::create_started_writer; use crate::reducer::test_support::generic_summary; use crate::reducer::test_support::message; @@ -555,9 +556,10 @@ fn invalid_write_stdin_arguments_remain_failed_tool_calls() -> anyhow::Result<() json!({"type": "error", "error": "invalid arguments"}), ), ] { - let payload = std::fs::read(temp.path().join(&rollout.raw_payloads[id].path))?; + let mut reader = crate::PayloadReader::new(temp.path()); + let payload = reader.read(&rollout.raw_payloads[id])?; assert_eq!( - serde_json::from_slice::(&payload)?, + serde_json::from_slice::(payload)?, expected ); } @@ -654,14 +656,7 @@ fn append_inference_with_tool_call(writer: &TraceWriter) -> anyhow::Result<()> { "input": [message("user", "run tests")] }), )?; - writer.append(RawTraceEventPayload::InferenceStarted { - inference_call_id: "inference-1".to_string(), - thread_id: "thread-root".to_string(), - codex_turn_id: "turn-1".to_string(), - model: "gpt-test".to_string(), - provider_name: "test-provider".to_string(), - request_payload: request, - })?; + append_inference_start(writer, "inference-1", "turn-1", request)?; let response = writer.write_json_payload( RawPayloadKind::InferenceResponse, @@ -696,13 +691,6 @@ fn append_followup_with_tool_output(writer: &TraceWriter) -> anyhow::Result<()> }] }), )?; - writer.append(RawTraceEventPayload::InferenceStarted { - inference_call_id: "inference-2".to_string(), - thread_id: "thread-root".to_string(), - codex_turn_id: "turn-2".to_string(), - model: "gpt-test".to_string(), - provider_name: "test-provider".to_string(), - request_payload: request, - })?; + append_inference_start(writer, "inference-2", "turn-2", request)?; Ok(()) } diff --git a/codex-rs/rollout-trace/src/shard.rs b/codex-rs/rollout-trace/src/shard.rs new file mode 100644 index 0000000000..e47a0e40b1 --- /dev/null +++ b/codex-rs/rollout-trace/src/shard.rs @@ -0,0 +1,324 @@ +//! Whole JSON bodies in ordered Zstandard shards. Each record is a little-endian +//! u16 ID length and u64 body length, followed by the UTF-8 ID and the body bytes. +//! A flush publishes one record without resetting the compression window. + +use std::collections::BTreeMap; +use std::collections::HashMap; +use std::fs::File; +use std::io::BufWriter; +use std::io::ErrorKind; +use std::io::Read; +use std::io::Write; +use std::ops::Range; +use std::path::Path; +use std::path::PathBuf; + +use anyhow::Context; +use anyhow::Result; +use anyhow::ensure; +use serde::de::IgnoredAny; +use zstd::stream::write::Encoder; + +use crate::bundle::PAYLOADS_DIR_NAME; +use crate::payload::PayloadCapture; +use crate::payload::RawPayloadId; +use crate::payload::RawPayloadRef; + +const TARGET_BYTES: usize = 32 * 1024 * 1024; +const WINDOW_LOG: u32 = 23; // 8 MiB of match history across payloads. +const HEADER_BYTES: usize = 10; + +fn shard_path(directory: &Path, number: u64) -> PathBuf { + directory.join(format!("{number:06}.zst")) +} + +pub(crate) struct ShardWriter { + directory: PathBuf, + active: Option>>, + number: u64, + bytes: usize, +} + +impl ShardWriter { + pub(crate) fn new(directory: PathBuf) -> Self { + Self { + directory, + active: None, + number: 0, + bytes: 0, + } + } + + /// Returns the shard containing this whole record once it is readable. + pub(crate) fn append(&mut self, id: &str, body: &[u8]) -> Result { + let id_len = u16::try_from(id.len())?; + let size = HEADER_BYTES + id.len() + body.len(); + if self.bytes.saturating_add(size) > TARGET_BYTES { + self.finish()?; + } + let encoder = match self.active.as_mut() { + Some(active) => active, + None => { + self.number += 1; + self.bytes = 0; + let path = shard_path(&self.directory, self.number); + let file = File::create_new(&path) + .with_context(|| format!("create shard {}", path.display()))?; + let mut encoder = Encoder::new(BufWriter::new(file), 3)?; + encoder.window_log(WINDOW_LOG)?; + encoder.include_checksum(true)?; + self.active.insert(encoder) + } + }; + encoder.write_all(&id_len.to_le_bytes())?; + encoder.write_all(&(body.len() as u64).to_le_bytes())?; + encoder.write_all(id.as_bytes())?; + encoder.write_all(body)?; + encoder.flush()?; + self.bytes += size; + Ok(self.number) + } + + /// Flushes the frame footer before the event log can publish RecordingCompleted. + pub(crate) fn finish(&mut self) -> Result<()> { + if let Some(active) = self.active.take() { + active.finish()?.flush()?; + } + Ok(()) + } +} + +struct DecodedShard { + number: u64, + sealed: bool, + bytes: Box<[u8]>, + records: HashMap>, +} + +impl DecodedShard { + fn open(directory: &Path, number: u64) -> Result { + let path = shard_path(directory, number); + let file = File::open(&path).with_context(|| format!("open shard {}", path.display()))?; + let mut decoder = zstd::stream::read::Decoder::new(file)?; + decoder.window_log_max(WINDOW_LOG)?; + let mut bytes = Vec::new(); + let sealed = match decoder.read_to_end(&mut bytes) { + Ok(_) => true, + // A live writer may have flushed complete records without sealing the frame. + Err(error) if error.kind() == ErrorKind::UnexpectedEof => false, + Err(error) => { + return Err(error).with_context(|| format!("decompress shard {}", path.display())); + } + }; + + let mut records = HashMap::new(); + let mut offset = 0; + while bytes.len() - offset >= HEADER_BYTES { + let header = &bytes[offset..offset + HEADER_BYTES]; + let id_len = usize::from(u16::from_le_bytes(header[..2].try_into()?)); + let body_len = usize::try_from(u64::from_le_bytes(header[2..].try_into()?))?; + let id_start = offset + HEADER_BYTES; + let Some(body_start) = id_start.checked_add(id_len) else { + break; + }; + let Some(end) = body_start.checked_add(body_len) else { + break; + }; + if end > bytes.len() { + break; + } + let id = std::str::from_utf8(&bytes[id_start..body_start])?; + ensure!( + records.insert(id.to_owned(), body_start..end).is_none(), + "duplicate payload {id} in shard {number}" + ); + offset = end; + } + ensure!( + !sealed || offset == bytes.len(), + "partial record in sealed shard {number}" + ); + Ok(Self { + number, + sealed, + bytes: bytes.into_boxed_slice(), + records, + }) + } +} + +/// Keeps the two most recently read shards so replay can cross a shard boundary +/// in either direction. A shard larger than the target occupies the cache alone. +pub struct PayloadReader { + directory: PathBuf, + shards: Vec, + assembled: Vec, +} + +impl PayloadReader { + /// Opens payloads on demand relative to this bundle root. + pub fn new(bundle: impl AsRef) -> Self { + Self { + directory: bundle.as_ref().join(PAYLOADS_DIR_NAME), + shards: Vec::new(), + assembled: Vec::new(), + } + } + + /// Reads the captured body verbatim, including from a flushed, unfinished shard. + pub fn read(&mut self, payload: &RawPayloadRef) -> Result<&[u8]> { + self.assembled = Vec::new(); + ensure!( + payload.is_complete(), + "partial payload {}", + payload.raw_payload_id + ); + if payload.parts.is_empty() { + return self.read_inner(payload, /*require_sealed*/ false); + } + let mut bytes = Vec::new(); + self.assemble(payload, &mut bytes)?; + self.assembled = bytes; + Ok(&self.assembled) + } + + /// Read the surviving structure for reduction. Known gaps occupy their original slots. + /// Exact replay uses `read`, which rejects a body with any missing descendant. + pub(crate) fn observed_json(&mut self, payload: &RawPayloadRef) -> Result { + if payload.is_complete() { + return self.read_json(payload); + } + let mut bytes = Vec::new(); + self.assemble(payload, &mut bytes)?; + Ok(serde_json::from_slice(&bytes)?) + } + + fn assemble(&mut self, payload: &RawPayloadRef, bytes: &mut Vec) -> Result<()> { + if !payload.capture.is_complete() { + serde_json::to_writer( + bytes, + &serde_json::json!({"$trace_unavailable": payload.raw_payload_id}), + )?; + } else if payload.parts.is_empty() { + bytes.extend_from_slice(self.read_inner(payload, /*require_sealed*/ false)?); + } else { + let fragments = decode_fragments( + self.read_inner(payload, /*require_sealed*/ false)?, + payload.parts.len(), + )?; + bytes.extend_from_slice(fragments[0].as_bytes()); + for (part, suffix) in payload.parts.iter().zip(&fragments[1..]) { + self.assemble(part, bytes)?; + bytes.extend_from_slice(suffix.as_bytes()); + } + } + Ok(()) + } + + /// Returns the captured prefix of a truncated body, which can end inside a UTF-8 character. + pub fn read_prefix(&mut self, payload: &RawPayloadRef) -> Result<&[u8]> { + ensure!( + matches!(payload.capture, PayloadCapture::Truncated { .. }), + "expected a truncated payload" + ); + self.read_inner(payload, /*require_sealed*/ false) + } + + /// Decodes the complete JSON at this reference. + pub fn read_json(&mut self, payload: &RawPayloadRef) -> Result { + serde_json::from_slice(self.read(payload)?) + .with_context(|| format!("parse payload {}", payload.raw_payload_id)) + } + + /// Checks every published body, including payloads the semantic reducer does not inspect. + pub(crate) fn validate( + &mut self, + index: &BTreeMap, + recording_closed: bool, + ) -> Result<()> { + // Group reads to decompress each shard just once during validation. + // Only the last shard can legitimately lack its footer in a live recording. + let mut refs: Vec<_> = index.values().collect(); + refs.sort_unstable_by_key(|p| p.shard); + let last = refs.last().and_then(|p| p.shard); + for payload in refs { + if matches!(payload.capture, PayloadCapture::Omitted { .. }) { + ensure!(payload.shard.is_none(), "omitted payload has a shard"); + continue; + } + let sealed = recording_closed || payload.shard != last; + let bytes = self.read_inner(payload, sealed)?; + if let PayloadCapture::Truncated { + original_bytes, + retained_bytes, + } = payload.capture + { + ensure!( + bytes.len() == retained_bytes && retained_bytes < original_bytes, + "invalid retained length" + ); + continue; + } + if payload.parts.is_empty() { + // IgnoredAny skips UTF-8 checks inside strings; validate the encoding separately. + serde_json::from_str::(std::str::from_utf8(bytes)?) + .with_context(|| format!("invalid payload {}", payload.raw_payload_id))?; + } else { + // Filling the slots with null verifies that this layout is a JSON document. + let fragments = decode_fragments(bytes, payload.parts.len())?; + serde_json::from_str::(&fragments.join("null"))?; + } + } + Ok(()) + } + + fn read_inner(&mut self, payload: &RawPayloadRef, require_sealed: bool) -> Result<&[u8]> { + let shard = payload.shard.context("payload omitted by size limit")?; + let hit = self.shards.iter().position(|cached| { + cached.number == shard + && (cached.sealed + || (!require_sealed && cached.records.contains_key(&payload.raw_payload_id))) + }); + if let Some(index) = hit { + self.shards.swap(0, index); + } else { + // An open shard can have gained a record or footer since the previous read. + self.shards + .retain(|cached| cached.number != shard && cached.bytes.len() <= TARGET_BYTES); + self.shards.truncate(1); + let decoded = DecodedShard::open(&self.directory, shard)?; + if decoded.bytes.len() > TARGET_BYTES { + self.shards.clear(); + } + self.shards.insert(0, decoded); + } + let cached = &self.shards[0]; + ensure!( + !require_sealed || cached.sealed, + "missing footer in payload shard {shard}" + ); + let range = cached + .records + .get(&payload.raw_payload_id) + .with_context(|| { + format!( + "missing payload {} in shard {shard}", + payload.raw_payload_id + ) + })?; + Ok(&cached.bytes[range.clone()]) + } +} + +fn decode_fragments(bytes: &[u8], parts: usize) -> Result> { + let fragments: Vec = serde_json::from_slice(bytes)?; + ensure!( + fragments.len() == parts + 1, + "invalid payload fragment count" + ); + Ok(fragments) +} + +#[cfg(test)] +#[path = "shard_tests.rs"] +mod tests; diff --git a/codex-rs/rollout-trace/src/shard_tests.rs b/codex-rs/rollout-trace/src/shard_tests.rs new file mode 100644 index 0000000000..47a62c62a0 --- /dev/null +++ b/codex-rs/rollout-trace/src/shard_tests.rs @@ -0,0 +1,132 @@ +#![cfg(test)] + +use super::*; +use crate::RawPayloadKind; +use pretty_assertions::assert_eq; +use tempfile::TempDir; + +#[test] +fn published_records_survive_rotation_and_remain_readable_while_open() -> Result<()> { + let bundle = TempDir::new()?; + let directory = bundle.path().join(PAYLOADS_DIR_NAME); + std::fs::create_dir(&directory)?; + let mut writer = ShardWriter::new(directory); + let mut reader = PayloadReader::new(bundle.path()); + let mut append = |id: &str, body: &[u8]| -> Result { + let shard = writer.append(id, body)?; + Ok(RawPayloadRef { + raw_payload_id: id.to_owned(), + kind: RawPayloadKind::InferenceRequest, + shard: Some(shard), + capture: Default::default(), + parts: Vec::new(), + }) + }; + + let first = append("first", b"{}")?; + assert_eq!(reader.read(&first)?, b"{}"); + let body = br#"{"text":"second\nline"}"#; + let second = append("second", body)?; + assert_eq!(second.shard, first.shard); + assert_eq!(reader.read(&second)?, body); + + // A record larger than the target gets its own shard and remains whole. + let mut large = vec![b'x'; TARGET_BYTES + 1]; + large[0] = b'"'; + *large.last_mut().unwrap() = b'"'; + let oversized = append("large", &large)?; + let last = append("last", b"null")?; + assert_ne!(first.shard, oversized.shard); + assert_ne!(oversized.shard, last.shard); + assert_eq!(reader.read(&oversized)?, large.as_slice()); + assert_eq!(reader.shards.len(), 1); + assert_eq!(reader.read(&first)?, b"{}"); + assert_eq!(reader.read(&last)?, b"null"); + assert!(reader.read_inner(&last, /*require_sealed*/ true).is_err()); + writer.finish()?; + assert_eq!(reader.read_inner(&last, /*require_sealed*/ true)?, b"null"); + Ok(()) +} + +#[test] +fn rotation_counts_the_record_header_and_id() -> Result<()> { + let directory = TempDir::new()?; + let mut writer = ShardWriter::new(directory.path().to_owned()); + let first = writer.append("a", b"x")?; + let remaining = TARGET_BYTES - (HEADER_BYTES + 1 + 1); + let exact_fit = writer.append( + "longer", + &vec![0; remaining - HEADER_BYTES - "longer".len()], + )?; + let next = writer.append("c", b"x")?; + assert_eq!([first, exact_fit, next], [1, 1, 2]); + Ok(()) +} + +#[test] +fn live_prefix_survives_a_partial_record_but_a_sealed_frame_rejects_it() -> Result<()> { + let directory = TempDir::new()?; + let mut writer = ShardWriter::new(directory.path().to_owned()); + let shard = writer.append("published", b"true")?; + let encoder = writer.active.as_mut().expect("open shard"); + // A valid header and ID promise four body bytes, but the writer stops after two. + encoder.write_all(&1u16.to_le_bytes())?; + encoder.write_all(&4u64.to_le_bytes())?; + encoder.write_all(b"xnu")?; + encoder.flush()?; + let live = DecodedShard::open(directory.path(), shard)?; + assert!(!live.sealed); + assert_eq!(live.records.len(), 1); + assert_eq!(&live.bytes[live.records["published"].clone()], b"true"); + writer.finish()?; + assert!( + DecodedShard::open(directory.path(), shard) + .is_err_and(|error| error.to_string().contains("partial record")) + ); + Ok(()) +} + +#[test] +fn duplicate_ids_cannot_resolve_to_different_bodies() -> Result<()> { + let directory = TempDir::new()?; + let mut writer = ShardWriter::new(directory.path().to_owned()); + let shard = writer.append("same", b"true")?; + writer.append("same", b"false")?; + writer.finish()?; + assert!( + DecodedShard::open(directory.path(), shard) + .is_err_and(|error| error.to_string().contains("duplicate payload same")) + ); + Ok(()) +} + +#[test] +fn reads_across_a_boundary_keep_the_two_most_recent_shards() -> Result<()> { + let bundle = TempDir::new()?; + let directory = bundle.path().join(PAYLOADS_DIR_NAME); + std::fs::create_dir(&directory)?; + let mut writer = ShardWriter::new(directory.clone()); + let mut refs = Vec::new(); + for shard in 1..=3 { + let id = format!("id-{shard}"); + let shard = writer.append(&id, b"null")?; + writer.finish()?; + refs.push(RawPayloadRef { + raw_payload_id: id, + kind: RawPayloadKind::InferenceRequest, + shard: Some(shard), + capture: Default::default(), + parts: Vec::new(), + }); + } + let mut reader = PayloadReader::new(bundle.path()); + reader.read(&refs[0])?; + reader.read(&refs[1])?; + // Removing the backing file makes a redundant decompression observable. + std::fs::remove_file(shard_path(&directory, refs[0].shard.unwrap()))?; + assert_eq!(reader.read(&refs[0])?, b"null"); + reader.read(&refs[2])?; + std::fs::remove_file(shard_path(&directory, refs[1].shard.unwrap()))?; + assert!(reader.read(&refs[1]).is_err()); + Ok(()) +} diff --git a/codex-rs/rollout-trace/src/thread_tests.rs b/codex-rs/rollout-trace/src/thread_tests.rs index 28902eef9e..e4fe9db9c9 100644 --- a/codex-rs/rollout-trace/src/thread_tests.rs +++ b/codex-rs/rollout-trace/src/thread_tests.rs @@ -62,22 +62,28 @@ fn create_in_root_writes_replayable_lifecycle_events() -> anyhow::Result<()> { assert_eq!(replayed.root_thread_id, thread_id.to_string()); assert_eq!(replayed.threads[&thread_id.to_string()].agent_path, "/root"); assert_eq!(replayed.raw_payloads.len(), 1); - assert!(!replayed.recording_complete); + assert_eq!(replayed.recording_status, crate::RecordingStatus::Unclosed); thread_trace.record_protocol_event(&EventMsg::ShutdownComplete); drop(thread_trace); let complete = replay_bundle(&bundle_dir)?; assert_eq!(complete.status, RolloutStatus::Completed); - assert!(complete.recording_complete); + assert_eq!(complete.recording_status, crate::RecordingStatus::Complete); // A syntactically valid prefix is still an incomplete export. let path = bundle_dir.join("trace.jsonl"); let events = fs::read_to_string(&path)?; fs::write(&path, events.strip_suffix('\n').unwrap())?; - assert!(!replay_bundle(&bundle_dir)?.recording_complete); + assert_eq!( + replay_bundle(&bundle_dir)?.recording_status, + crate::RecordingStatus::Unclosed + ); let (prefix, _) = events.trim_end().rsplit_once('\n').unwrap(); fs::write(&path, prefix)?; - assert!(!replay_bundle(&bundle_dir)?.recording_complete); + assert_eq!( + replay_bundle(&bundle_dir)?.recording_status, + crate::RecordingStatus::Unclosed + ); fs::write(&path, &events)?; // Completion closes the log; even a correctly numbered later event is invalid. @@ -93,11 +99,33 @@ fn create_in_root_writes_replayable_lifecycle_events() -> anyhow::Result<()> { .values() .find(|payload| payload.kind == RawPayloadKind::ProtocolEvent) .unwrap(); - fs::write(bundle_dir.join(&opaque.path), "{\"truncated\":")?; + let shard = opaque.shard.unwrap(); + let shard_path = bundle_dir.join(format!("payloads/{shard:06}.zst")); + let compressed = fs::read(&shard_path)?; + // Even when all records can be decoded, a completed bundle requires the frame checksum. + fs::write(&shard_path, &compressed[..compressed.len() - 1])?; assert!(replay_bundle(&bundle_dir).is_err()); - fs::write(bundle_dir.join(&opaque.path), b"{\"value\":\"\xff\"}")?; + let mut corrupt_checksum = compressed.clone(); + *corrupt_checksum.last_mut().unwrap() ^= 1; + fs::write(&shard_path, corrupt_checksum)?; assert!(replay_bundle(&bundle_dir).is_err()); - fs::remove_file(bundle_dir.join(&opaque.path))?; + fs::write(&shard_path, &compressed)?; + + let mut reader = crate::PayloadReader::new(&bundle_dir); + let original = reader.read(opaque)?; + let mut decoded = zstd::decode_all(compressed.as_slice())?; + let start = decoded + .windows(original.len()) + .position(|body| body == original) + .unwrap(); + let string_byte = start + original.iter().position(|&b| b == b'"').unwrap() + 1; + for (index, bad_byte) in [(start, b'?'), (string_byte, 0xff)] { + let old_byte = std::mem::replace(&mut decoded[index], bad_byte); + fs::write(&shard_path, zstd::encode_all(decoded.as_slice(), 3)?)?; + assert!(replay_bundle(&bundle_dir).is_err()); + decoded[index] = old_byte; + } + fs::remove_file(&shard_path)?; assert!(replay_bundle(&bundle_dir).is_err()); Ok(()) @@ -157,9 +185,15 @@ fn spawned_thread_start_appends_to_root_bundle() -> anyhow::Result<()> { root_trace.record_ended(RolloutStatus::Completed); drop(root_trace); // Root shutdown does not close the recording while a child still owns it. - assert!(!replay_bundle(&bundle_dir)?.recording_complete); + assert_eq!( + replay_bundle(&bundle_dir)?.recording_status, + crate::RecordingStatus::Unclosed + ); drop(child_trace); - assert!(replay_bundle(&bundle_dir)?.recording_complete); + assert_eq!( + replay_bundle(&bundle_dir)?.recording_status, + crate::RecordingStatus::Complete + ); Ok(()) } @@ -178,18 +212,30 @@ fn resumed_child_rejects_a_failed_shared_writer() -> anyhow::Result<()> { // A parent payload write fails while the child is evicted. let bundle = single_bundle_dir(temp.path())?; - let payloads = bundle.join("payloads"); - let backup = temp.path().join("saved-payloads"); - fs::rename(&payloads, &backup)?; - root.record_protocol_event(&EventMsg::ShutdownComplete); - fs::rename(backup, payloads)?; + // Occupy the next shard path so rotation fails without moving an open file. + let next_shard = bundle.join("payloads/000002.zst"); + fs::write(&next_shard, [])?; + let ThreadTraceContextState::Enabled(context) = &root.state else { + unreachable!() + }; + let oversized = "x".repeat(32 * 1024 * 1024); + assert!( + context + .writer + .write_json_payload(RawPayloadKind::ProtocolEvent, &oversized) + .is_err() + ); + fs::remove_file(next_shard)?; assert!( root.resume_child_thread_trace(child_id.to_string()) .is_err() ); drop(root); - assert!(!replay_bundle(&bundle)?.recording_complete); + assert_eq!( + replay_bundle(&bundle)?.recording_status, + crate::RecordingStatus::Unclosed + ); Ok(()) } @@ -215,7 +261,8 @@ fn disabled_thread_context_accepts_trace_calls_without_writing() -> anyhow::Resu let inference_trace = thread_trace.inference_trace_context("turn-1", "gpt-test", "test-provider"); let inference_attempt = inference_trace.start_attempt(); - inference_attempt.record_started(&serde_json::json!({ "kind": "inference" })); + let input = serde_json::json!({ "kind": "inference" }); + inference_attempt.record_started(serde_json::to_vec(&input)?, &input, || Ok(None)); let token_usage: Option = None; inference_attempt.record_completed("response-1", Some("req-1"), &token_usage, &[]); inference_attempt.record_failed("inference failed", /*upstream_request_id*/ None, &[]); diff --git a/codex-rs/rollout-trace/src/writer.rs b/codex-rs/rollout-trace/src/writer.rs index f3cae91c78..0ce8c754e3 100644 --- a/codex-rs/rollout-trace/src/writer.rs +++ b/codex-rs/rollout-trace/src/writer.rs @@ -1,16 +1,15 @@ //! Hot-path trace bundle writer. //! -//! All producers share one writer. Write failures prevent further recording and -//! suppress the final closing event. +//! Writers serialize whole payloads before publishing the events that reference them. +//! Closing the bundle seals the final shard before recording completion. +use std::fmt; use std::fs::File; -use std::fs::OpenOptions; use std::io::BufWriter; use std::io::Write; use std::path::Path; -use std::path::PathBuf; use std::sync::Mutex; -use std::sync::MutexGuard; +use std::sync::OnceLock; use std::sync::PoisonError; use std::time::SystemTime; use std::time::UNIX_EPOCH; @@ -31,25 +30,84 @@ use crate::raw_event::RAW_TRACE_EVENT_SCHEMA_VERSION; use crate::raw_event::RawTraceEvent; use crate::raw_event::RawTraceEventContext; use crate::raw_event::RawTraceEventPayload; +use crate::shard::ShardWriter; -/// Local trace bundle writer. -/// -/// The writer appends raw events and writes payload files. It does not keep a -/// reduced `RolloutTrace` in memory; replay is owned by the reducer. -#[derive(Debug)] +/// Serializes access to the event log and its ordered payload shards. pub struct TraceWriter { - inner: Mutex, + bundle: Mutex, } -#[derive(Debug)] -struct TraceWriterInner { - manifest: TraceBundleManifest, - payloads_dir: PathBuf, - event_log: BufWriter, +impl fmt::Debug for TraceWriter { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TraceWriter").finish_non_exhaustive() + } +} + +/// Owns the bundle files until shutdown. Each body is flushed before its event is published. +pub(crate) struct BundleWriter { + rollout_id: String, + events: BufWriter, + shards: ShardWriter, next_seq: u64, next_payload_ordinal: u64, - // Preserve the first failure even when a producer cannot return it to its caller. - error: Option, + // Keep the first error so explicit reads and shutdown can detect an incomplete recording. + failure: OnceLock, +} + +impl BundleWriter { + fn record(&mut self, write: impl FnOnce(&mut Self) -> Result) -> Result { + check_recording(&self.failure)?; + write(self).inspect_err(|error| { + let error = format!("{error:#}"); + tracing::error!("rollout recording failed: {error}"); + let _ = self.failure.set(error); + }) + } + + pub(crate) fn write_body( + &mut self, + kind: RawPayloadKind, + body: &[u8], + ) -> Result { + let ordinal = self.next_payload_ordinal; + let raw_payload_id = format!("raw_payload:{ordinal}"); + let shard = self.shards.append(&raw_payload_id, body)?; + self.next_payload_ordinal += 1; + Ok(RawPayloadRef { + raw_payload_id, + kind, + shard: Some(shard), + capture: crate::PayloadCapture::Complete, + parts: Vec::new(), + }) + } + + fn append( + &mut self, + context: RawTraceEventContext, + payload: RawTraceEventPayload, + mut wall_time_unix_ms: i64, + ) -> Result { + if matches!(payload, RawTraceEventPayload::RecordingCompleted) { + self.shards.finish()?; + // Timestamp the closing marker after the shard is sealed. + wall_time_unix_ms = unix_time_ms(); + } + let event = RawTraceEvent { + schema_version: RAW_TRACE_EVENT_SCHEMA_VERSION, + seq: self.next_seq, + wall_time_unix_ms, + rollout_id: self.rollout_id.clone(), + thread_id: context.thread_id, + codex_turn_id: context.codex_turn_id, + payload, + }; + serde_json::to_writer(&mut self.events, &event)?; + self.events.write_all(b"\n")?; + self.events.flush()?; + self.next_seq += 1; + Ok(event) + } } impl TraceWriter { @@ -60,7 +118,7 @@ impl TraceWriter { rollout_id: String, root_thread_id: AgentThreadId, ) -> Result { - let bundle_dir = bundle_dir.as_ref().to_path_buf(); + let bundle_dir = bundle_dir.as_ref(); let payloads_dir = bundle_dir.join(PAYLOADS_DIR_NAME); std::fs::create_dir_all(&payloads_dir) .with_context(|| format!("create trace payload dir {}", payloads_dir.display()))?; @@ -69,48 +127,43 @@ impl TraceWriter { let manifest = TraceBundleManifest::new(trace_id, rollout_id, root_thread_id, started_at_unix_ms); let event_log_path = bundle_dir.join(RAW_EVENT_LOG_FILE_NAME); - let event_log = OpenOptions::new() - .create_new(true) - .write(true) - .open(&event_log_path) + let event_log = File::create_new(&event_log_path) .with_context(|| format!("open trace event log {}", event_log_path.display()))?; - write_json_file(&bundle_dir.join(MANIFEST_FILE_NAME), &manifest)?; + let manifest_path = bundle_dir.join(MANIFEST_FILE_NAME); + std::fs::write(&manifest_path, serde_json::to_vec_pretty(&manifest)?) + .with_context(|| format!("write trace manifest {}", manifest_path.display()))?; + let 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(), + }; Ok(Self { - inner: Mutex::new(TraceWriterInner { - manifest, - payloads_dir, - event_log: BufWriter::new(event_log), - next_seq: 1, - next_payload_ordinal: 1, - error: None, - }), + bundle: Mutex::new(bundle), }) } - /// Writes a JSON payload file and returns its reduced-state reference. + fn call(&self, write: impl FnOnce(&mut BundleWriter) -> Result) -> Result { + self.bundle + .lock() + .unwrap_or_else(PoisonError::into_inner) + .record(write) + } + + pub(crate) fn check(&self) -> Result<()> { + self.call(|_| Ok(())) + } + + /// Publishes a complete JSON body and returns its location in the bundle. pub fn write_json_payload( &self, kind: RawPayloadKind, value: &impl Serialize, ) -> Result { - let mut inner = self.lock_inner(); - inner.check()?; - let ordinal = inner.next_payload_ordinal; - inner.next_payload_ordinal += 1; - let raw_payload_id = format!("raw_payload:{ordinal}"); - let relative_path = format!("{PAYLOADS_DIR_NAME}/{ordinal}.json"); - let absolute_path = inner.payloads_dir.join(format!("{ordinal}.json")); - // Payload files are created before the event that references them. A - // replay interrupted after an event is appended should never point at a - // payload file that the writer planned but had not written yet. - let result = write_json_file(&absolute_path, value); - inner.record_result(result)?; - Ok(RawPayloadRef { - raw_payload_id, - kind, - path: relative_path, - }) + self.call(|bundle| bundle.write_body(kind, &serde_json::to_vec(value)?)) } /// Appends one raw event with no extra envelope context. @@ -118,83 +171,57 @@ impl TraceWriter { self.append_with_context(RawTraceEventContext::default(), payload) } - /// Appends one raw event with explicit thread/turn context. + /// Appends an event and waits for it to become readable. pub fn append_with_context( &self, context: RawTraceEventContext, payload: RawTraceEventPayload, ) -> Result { - let mut inner = self.lock_inner(); - inner.check()?; - let event = RawTraceEvent { - schema_version: RAW_TRACE_EVENT_SCHEMA_VERSION, - seq: inner.next_seq, - wall_time_unix_ms: unix_time_ms(), - rollout_id: inner.manifest.rollout_id.clone(), - thread_id: context.thread_id, - codex_turn_id: context.codex_turn_id, - payload, - }; - let result = (|| { - serde_json::to_writer(&mut inner.event_log, &event)?; - inner.event_log.write_all(b"\n")?; - inner.event_log.flush()?; + let now = unix_time_ms(); + self.call(move |bundle| bundle.append(context, payload, now)) + } + + /// Writes related payloads and their event under the same lock. + pub(crate) fn record_event( + &self, + context: RawTraceEventContext, + build: impl FnOnce(&mut BundleWriter) -> Result, + ) -> Result<()> { + self.call(|bundle| { + let payload = build(bundle)?; + bundle.append(context, payload, unix_time_ms())?; Ok(()) - })(); - inner.record_result(result)?; - inner.next_seq += 1; - Ok(event) + }) } - /// Reports any earlier recording failure, including failures in detached producers. - pub(crate) fn check(&self) -> Result<()> { - self.lock_inner().check() - } - - fn lock_inner(&self) -> MutexGuard<'_, TraceWriterInner> { - // Preserve the event log after a panic in tracing code. Dropping the - // writer would lose subsequent diagnostic events in exactly the session - // we are trying to debug. - self.inner.lock().unwrap_or_else(PoisonError::into_inner) - } -} - -impl TraceWriterInner { - fn check(&self) -> Result<()> { - if let Some(error) = &self.error { - bail!("rollout recording failed: {error}"); - } - Ok(()) - } - - fn record_result(&mut self, result: Result) -> Result { - if let Err(error) = &result { - let error = format!("{error:#}"); - tracing::error!("rollout recording failed: {error}"); - self.error = Some(error); - } - result + /// Publishes a JSON body before the event that points to it. + pub(crate) fn event_with_payload( + &self, + context: RawTraceEventContext, + kind: RawPayloadKind, + value: &impl Serialize, + event: impl FnOnce(RawPayloadRef) -> RawTraceEventPayload, + ) -> Result<()> { + self.record_event(context, |bundle| { + let reference = bundle.write_body(kind, &serde_json::to_vec(value)?)?; + Ok(event(reference)) + }) } } impl Drop for TraceWriter { fn drop(&mut self) { - // The last producer releases the writer, even for interrupted turns. - // This marks recorder closure, not successful task completion. - if !self.inner.is_poisoned() && !std::thread::panicking() { + if !self.bundle.is_poisoned() && !std::thread::panicking() { let _ = self.append(RawTraceEventPayload::RecordingCompleted); } } } -fn write_json_file(path: &Path, value: &impl Serialize) -> Result<()> { - let file = File::create(path).with_context(|| format!("create {}", path.display()))?; - let mut writer = BufWriter::new(file); - serde_json::to_writer_pretty(&mut writer, value) - .with_context(|| format!("write JSON {}", path.display()))?; - writer - .flush() - .with_context(|| format!("flush JSON {}", path.display())) +fn check_recording(failure: &OnceLock) -> Result<()> { + if let Some(error) = failure.get() { + bail!("rollout recording failed: {error}"); + } + Ok(()) } pub(crate) fn unix_time_ms() -> i64 { @@ -259,6 +286,9 @@ mod tests { }), )?; writer.append(RawTraceEventPayload::InferenceStarted { + prepared_prompt_payload: writer + .write_json_payload(RawPayloadKind::PreparedPrompt, &json!({"input": []}))?, + warmup_context_payload: None, inference_call_id: "inference-1".to_string(), thread_id: "thread-root".to_string(), codex_turn_id: "turn-1".to_string(), @@ -306,8 +336,8 @@ mod tests { Some(inference_response.raw_payload_id), ); assert_eq!( - rollout.raw_payloads[&metadata_payload.raw_payload_id].path, - "payloads/1.json" + rollout.raw_payloads[&metadata_payload.raw_payload_id].shard, + Some(1) ); Ok(()) @@ -345,8 +375,11 @@ mod tests { let writer = TraceWriter::create(temp.path(), "trace".into(), "rollout".into(), "root".into())?; // A read-only file makes the real buffered write fail on every platform. - writer.lock_inner().event_log = - std::io::BufWriter::new(std::fs::File::open(temp.path().join("trace.jsonl"))?); + let file = std::fs::File::open(temp.path().join("trace.jsonl"))?; + writer.call(move |bundle| { + bundle.events = std::io::BufWriter::new(file); + Ok(()) + })?; assert!( writer .append(RawTraceEventPayload::RolloutEnded {