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 {