mirror of
https://github.com/openai/codex.git
synced 2026-11-13 21:30:47 +00:00
Store rollout trace payloads in compressed shards (#61222)
## What changed - Replace individual JSON payload files with numbered Zstandard shards in trace schema v3. Add `PayloadReader` to resolve payload references, cache decoded shards, and validate payloads during replay. - Capture exact outgoing HTTP bodies and WebSocket messages alongside prepared prompts. Preserve full formatted context for continuations from unrecorded WebSocket warmups while retaining the actual incremental request. - Flush payload records before publishing references and seal the final shard before recording completion. Replace `recording_complete` with `recording_status` to distinguish unclosed recordings, closed recordings with partial payloads, and complete recordings. Reduce trace bundles with the binary that recorded them. ## Testing Add coverage for shard rotation, oversized records, live reads, cache eviction, and malformed or corrupt payloads. Extend HTTP and WebSocket integration assertions to verify captured requests and prepared prompts, including warmup continuations. GitOrigin-RevId: 40ce3c11383f48e08f52e7120b706c61a449df9c
This commit is contained in:
committed by
copyberry
parent
662daabad2
commit
e73744b3f4
Generated
+2
@@ -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]]
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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)?);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)},
|
||||
|
||||
@@ -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<RolloutTrace> {
|
||||
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"),
|
||||
)?)?)
|
||||
}
|
||||
|
||||
@@ -57,6 +57,16 @@ impl<T: HttpTransport> ResponsesClient<T> {
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn stream_request(
|
||||
&self,
|
||||
request: ResponsesApiRequest,
|
||||
options: ResponsesOptions,
|
||||
) -> Result<ResponseStream, ApiError> {
|
||||
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<T: HttpTransport> ResponsesClient<T> {
|
||||
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<ResponseStream, ApiError> {
|
||||
let ResponsesOptions {
|
||||
session_id,
|
||||
@@ -83,6 +94,7 @@ impl<T: HttpTransport> ResponsesClient<T> {
|
||||
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);
|
||||
|
||||
@@ -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<Arc<OnceLock<String>>>,
|
||||
) -> Result<ResponseStream, ApiError> {
|
||||
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<Arc<OnceLock<String>>>,
|
||||
on_encoded: impl FnOnce(&bytes::Bytes),
|
||||
) -> Result<ResponseStream, ApiError> {
|
||||
let (tx_event, rx_event) =
|
||||
mpsc::channel::<std::result::Result<ResponseEvent, ApiError>>(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<std::result::Result<ResponseEvent, ApiError>>,
|
||||
request_text: String,
|
||||
request_text: Utf8Bytes,
|
||||
idle_timeout: Duration,
|
||||
telemetry: Option<Arc<dyn WebsocketTelemetry>>,
|
||||
turn_state: Option<&OnceLock<String>>,
|
||||
@@ -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<dyn WebsocketTelemetry>>,
|
||||
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(
|
||||
|
||||
+22
-13
@@ -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 {
|
||||
|
||||
@@ -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<ResponseItem>,
|
||||
|
||||
/// Tool definitions to inject into this request, empty when supplied by history.
|
||||
|
||||
@@ -1469,14 +1469,15 @@ fn started_inference_attempt(temp: &TempDir) -> anyhow::Result<InferenceTraceAtt
|
||||
"test-provider".to_string(),
|
||||
);
|
||||
let attempt = inference_trace.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));
|
||||
Ok(attempt)
|
||||
}
|
||||
|
||||
@@ -1802,7 +1803,7 @@ async fn dropped_response_stream_traces_cancelled_partial_output() -> 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(())
|
||||
}
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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 `<trace-bundle>/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;
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<Option<Vec<u8>>>,
|
||||
) {
|
||||
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<JsonValue> {
|
||||
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<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
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<S: Serializer>(
|
||||
items: &[ResponseItem],
|
||||
serializer: S,
|
||||
) -> Result<S::Ok, S::Error> {
|
||||
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<crate::RawPayloadRef> {
|
||||
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<crate::RawPayloadRef> {
|
||||
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<crate::RawPayloadRef>) -> 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",
|
||||
}),
|
||||
}]),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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<ToolCallId>,
|
||||
pub usage: Option<TokenUsage>,
|
||||
/// 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<RawPayloadId>,
|
||||
}
|
||||
|
||||
@@ -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<i64>,
|
||||
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<AgentThreadId, AgentThread>,
|
||||
pub codex_turns: BTreeMap<CodexTurnId, CodexTurn>,
|
||||
@@ -88,27 +98,26 @@ pub struct RolloutTrace {
|
||||
pub compaction_requests: BTreeMap<CompactionRequestId, CompactionRequest>,
|
||||
/// Information-flow edges between threads, cells, tools, and runtime resources.
|
||||
pub interaction_edges: BTreeMap<EdgeId, InteractionEdge>,
|
||||
/// 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<RawPayloadId, RawPayloadRef>,
|
||||
}
|
||||
|
||||
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(),
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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<u64>,
|
||||
#[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<RawPayloadRef>,
|
||||
}
|
||||
|
||||
/// 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<Item = &Self> {
|
||||
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,
|
||||
|
||||
@@ -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<CodexTurnId>,
|
||||
}
|
||||
|
||||
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<RawPayloadRef>,
|
||||
},
|
||||
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,
|
||||
..
|
||||
}
|
||||
|
||||
@@ -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(());
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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<Vec<String>> {
|
||||
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<Vec<String>> {
|
||||
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<ReducedCompactionCheckpoint> {
|
||||
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)?;
|
||||
|
||||
@@ -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<RawPayloadRef>,
|
||||
}
|
||||
|
||||
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,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -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<Path>) -> Result<RolloutTrace> {
|
||||
let bundle_dir = bundle_dir.as_ref();
|
||||
@@ -64,13 +62,12 @@ pub fn replay_bundle(bundle_dir: impl AsRef<Path>) -> Result<RolloutTrace> {
|
||||
);
|
||||
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<Path>) -> Result<RolloutTrace> {
|
||||
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<Path>) -> Result<RolloutTrace> {
|
||||
}
|
||||
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::<IgnoredAny>(&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<Path>) -> Result<RolloutTrace> {
|
||||
|
||||
struct TraceReducer {
|
||||
rollout: RolloutTrace,
|
||||
bundle_dir: PathBuf,
|
||||
payload_reader: PayloadReader,
|
||||
/// First (wall time, sequence) by receiving thread and Responses item ID.
|
||||
first_acceptances: BTreeMap<String, BTreeMap<String, (i64, crate::RawEventSeq)>>,
|
||||
/// 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<Value> {
|
||||
// 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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -229,10 +229,10 @@ impl TraceReducer {
|
||||
}
|
||||
|
||||
fn thread_started_metadata(
|
||||
&self,
|
||||
&mut self,
|
||||
metadata_payload: &RawPayloadRef,
|
||||
) -> Result<ThreadStartedMetadata> {
|
||||
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))
|
||||
}
|
||||
|
||||
@@ -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<String> {
|
||||
fn agent_message_content_from_invocation(&mut self, tool_call_id: &str) -> Result<String> {
|
||||
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"))
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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::<serde_json::Value>(&payload)?,
|
||||
serde_json::from_slice::<serde_json::Value>(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(())
|
||||
}
|
||||
|
||||
@@ -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<Encoder<'static, BufWriter<File>>>,
|
||||
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<u64> {
|
||||
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<String, Range<usize>>,
|
||||
}
|
||||
|
||||
impl DecodedShard {
|
||||
fn open(directory: &Path, number: u64) -> Result<Self> {
|
||||
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<DecodedShard>,
|
||||
assembled: Vec<u8>,
|
||||
}
|
||||
|
||||
impl PayloadReader {
|
||||
/// Opens payloads on demand relative to this bundle root.
|
||||
pub fn new(bundle: impl AsRef<Path>) -> 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<serde_json::Value> {
|
||||
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<u8>) -> 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::Value> {
|
||||
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<RawPayloadId, RawPayloadRef>,
|
||||
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::<IgnoredAny>(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::<IgnoredAny>(&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<Vec<String>> {
|
||||
let fragments: Vec<String> = serde_json::from_slice(bytes)?;
|
||||
ensure!(
|
||||
fragments.len() == parts + 1,
|
||||
"invalid payload fragment count"
|
||||
);
|
||||
Ok(fragments)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "shard_tests.rs"]
|
||||
mod tests;
|
||||
@@ -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<RawPayloadRef> {
|
||||
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(())
|
||||
}
|
||||
@@ -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<codex_protocol::protocol::TokenUsage> = None;
|
||||
inference_attempt.record_completed("response-1", Some("req-1"), &token_usage, &[]);
|
||||
inference_attempt.record_failed("inference failed", /*upstream_request_id*/ None, &[]);
|
||||
|
||||
@@ -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<TraceWriterInner>,
|
||||
bundle: Mutex<BundleWriter>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct TraceWriterInner {
|
||||
manifest: TraceBundleManifest,
|
||||
payloads_dir: PathBuf,
|
||||
event_log: BufWriter<File>,
|
||||
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<File>,
|
||||
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<String>,
|
||||
// Keep the first error so explicit reads and shutdown can detect an incomplete recording.
|
||||
failure: OnceLock<String>,
|
||||
}
|
||||
|
||||
impl BundleWriter {
|
||||
fn record<T>(&mut self, write: impl FnOnce(&mut Self) -> Result<T>) -> Result<T> {
|
||||
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<RawPayloadRef> {
|
||||
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<RawTraceEvent> {
|
||||
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<Self> {
|
||||
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<T>(&self, write: impl FnOnce(&mut BundleWriter) -> Result<T>) -> Result<T> {
|
||||
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<RawPayloadRef> {
|
||||
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<RawTraceEvent> {
|
||||
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<RawTraceEventPayload>,
|
||||
) -> 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<T>(&mut self, result: Result<T>) -> Result<T> {
|
||||
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<String>) -> 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 {
|
||||
|
||||
Reference in New Issue
Block a user