openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add core rollout trace recorder

Browse Source
This commit is contained in:
Albin Cassirer
2026-04-15 16:07:59 -07:00
parent eff93282cf
commit 856ef510c2
5 changed files with 1195 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/client.rs
+104 -16
View File
@@ -77,6 +77,9 @@ use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::W3cTraceContext;
use codex_rollout_trace::CompactionTraceContext;
use codex_rollout_trace::InferenceTraceAttempt;
use codex_rollout_trace::InferenceTraceContext;
use codex_tools::create_tools_json_for_responses_api;
use eventsource_stream::Event;
use eventsource_stream::EventStreamError;
@@ -408,6 +411,7 @@ impl ModelClient {
effort: Option<ReasoningEffortConfig>,
summary: ReasoningSummaryConfig,
session_telemetry: &SessionTelemetry,
compaction_trace: &CompactionTraceContext,
) -> Result<Vec<ResponseItem>> {
if prompt.input.is_empty() {
return Ok(Vec::new());
@@ -462,10 +466,18 @@ impl ModelClient {
extra_headers.extend(build_conversation_headers(Some(
self.state.conversation_id.to_string(),
)));
client
.compact_input(&payload, extra_headers)
.await
.map_err(map_api_error)
let trace_attempt = compaction_trace.start_attempt(&payload);
match client.compact_input(&payload, extra_headers).await {
Ok(output) => {
trace_attempt.record_completed(&output);
Ok(output)
}
Err(err) => {
let err = map_api_error(err);
trace_attempt.record_failed(&err);
Err(err)
}
}
}
pub(crate) async fn create_realtime_call_with_headers(
@@ -1140,6 +1152,7 @@ impl ModelClientSession {
summary: ReasoningSummaryConfig,
service_tier: Option<ServiceTier>,
turn_metadata_header: Option<&str>,
inference_trace: &InferenceTraceContext,
) -> Result<ResponseStream> {
if let Some(path) = &*CODEX_RS_SSE_FIXTURE {
warn!(path, "Streaming from fixture");
@@ -1148,7 +1161,11 @@ impl ModelClientSession {
self.client.state.provider.stream_idle_timeout(),
)
.map_err(map_api_error)?;
let (stream, _last_request_rx) = map_response_stream(stream, session_telemetry.clone());
let (stream, _last_request_rx) = map_response_stream(
stream,
session_telemetry.clone(),
InferenceTraceAttempt::disabled(),
);
return Ok(stream);
}
@@ -1182,6 +1199,8 @@ impl ModelClientSession {
summary,
service_tier,
)?;
let inference_trace_attempt = inference_trace.start_attempt();
inference_trace_attempt.record_started(&request);
let client = ApiResponsesClient::new(
transport,
client_setup.api_provider,
@@ -1192,12 +1211,17 @@ impl ModelClientSession {
match stream_result {
Ok(stream) => {
let (stream, _) = map_response_stream(stream, session_telemetry.clone());
let (stream, _) = map_response_stream(
stream,
session_telemetry.clone(),
inference_trace_attempt,
);
return Ok(stream);
}
Err(ApiError::Transport(
unauthorized_transport @ TransportError::Http { status, .. },
)) if status == StatusCode::UNAUTHORIZED => {
inference_trace_attempt.record_failed(&unauthorized_transport);
pending_retry = PendingUnauthorizedRetry::from_recovery(
handle_unauthorized(
unauthorized_transport,
@@ -1208,7 +1232,11 @@ impl ModelClientSession {
);
continue;
}
Err(err) => return Err(map_api_error(err)),
Err(err) => {
let err = map_api_error(err);
inference_trace_attempt.record_failed(&err);
return Err(err);
}
}
}
}
@@ -1239,6 +1267,7 @@ impl ModelClientSession {
turn_metadata_header: Option<&str>,
warmup: bool,
request_trace: Option<W3cTraceContext>,
inference_trace: &InferenceTraceContext,
) -> Result<WebsocketStreamOutcome> {
let auth_manager = self.client.state.auth_manager.clone();
@@ -1313,17 +1342,35 @@ impl ModelClientSession {
let ws_request = self.prepare_websocket_request(ws_payload, &request);
self.websocket_session.last_request = Some(request);
let inference_trace_attempt = if warmup {
// Prewarm sends `generate=false`; it is connection setup, not a
// model inference attempt that should appear in rollout traces.
InferenceTraceAttempt::disabled()
} else {
inference_trace.start_attempt()
};
inference_trace_attempt.record_started(&ws_request);
let stream_result = self.websocket_session.connection.as_ref().ok_or_else(|| {
map_api_error(ApiError::Stream(
"websocket connection is unavailable".to_string(),
))
})?;
let stream_result = stream_result
let stream_result = match stream_result
.stream_request(ws_request, self.websocket_session.connection_reused())
.await
.map_err(map_api_error)?;
let (stream, last_request_rx) =
map_response_stream(stream_result, session_telemetry.clone());
{
Ok(stream_result) => stream_result,
Err(err) => {
let err = map_api_error(err);
inference_trace_attempt.record_failed(&err);
return Err(err);
}
};
let (stream, last_request_rx) = map_response_stream(
stream_result,
session_telemetry.clone(),
inference_trace_attempt,
);
self.websocket_session.last_response_rx = Some(last_request_rx);
return Ok(WebsocketStreamOutcome::Stream(stream));
}
@@ -1382,6 +1429,7 @@ impl ModelClientSession {
return Ok(());
}
let disabled_trace = InferenceTraceContext::disabled();
match self
.stream_responses_websocket(
prompt,
@@ -1393,6 +1441,7 @@ impl ModelClientSession {
turn_metadata_header,
/*warmup*/ true,
current_span_w3c_trace_context(),
&disabled_trace,
)
.await
{
@@ -1416,12 +1465,11 @@ impl ModelClientSession {
}
#[allow(clippy::too_many_arguments)]
/// Streams a single model request within the current turn.
/// Streams a single model request without rollout tracing.
///
/// The caller is responsible for passing per-turn settings explicitly (model selection,
/// reasoning settings, telemetry context, and turn metadata). This method will prefer the
/// Responses WebSocket transport when the provider supports it and it remains healthy, and will
/// fall back to the HTTP Responses API transport otherwise.
/// This is the public client API. It routes through the same transport code
/// as traced Codex turns, but supplies a disabled trace context so tracing
/// does not leak into callers that only need model streaming.
pub async fn stream(
&mut self,
prompt: &Prompt,
@@ -1431,6 +1479,37 @@ impl ModelClientSession {
summary: ReasoningSummaryConfig,
service_tier: Option<ServiceTier>,
turn_metadata_header: Option<&str>,
) -> Result<ResponseStream> {
let disabled_trace = InferenceTraceContext::disabled();
self.stream_with_trace(
prompt,
model_info,
session_telemetry,
effort,
summary,
service_tier,
turn_metadata_header,
&disabled_trace,
)
.await
}
#[allow(clippy::too_many_arguments)]
/// Streams a model request with an explicit rollout trace context.
///
/// The context may be enabled or disabled. Transport code records against it
/// unconditionally so HTTP, WebSocket, retry, and fallback paths do not need
/// separate trace/no-trace branches.
pub(crate) async fn stream_with_trace(
&mut self,
prompt: &Prompt,
model_info: &ModelInfo,
session_telemetry: &SessionTelemetry,
effort: Option<ReasoningEffortConfig>,
summary: ReasoningSummaryConfig,
service_tier: Option<ServiceTier>,
turn_metadata_header: Option<&str>,
inference_trace: &InferenceTraceContext,
) -> Result<ResponseStream> {
let wire_api = self.client.state.provider.wire_api;
match wire_api {
@@ -1448,6 +1527,7 @@ impl ModelClientSession {
turn_metadata_header,
/*warmup*/ false,
request_trace,
inference_trace,
)
.await?
{
@@ -1466,6 +1546,7 @@ impl ModelClientSession {
summary,
service_tier,
turn_metadata_header,
inference_trace,
)
.await
}
@@ -1561,6 +1642,7 @@ fn parent_thread_id_header_value(session_source: &SessionSource) -> Option<Strin
fn map_response_stream<S>(
api_stream: S,
session_telemetry: SessionTelemetry,
inference_trace_attempt: InferenceTraceAttempt,
) -> (ResponseStream, oneshot::Receiver<LastResponse>)
where
S: futures::Stream<Item = std::result::Result<ResponseEvent, ApiError>>
@@ -1601,6 +1683,11 @@ where
usage.total_tokens,
);
}
inference_trace_attempt.record_completed(
&response_id,
&token_usage,
&items_added,
);
if let Some(sender) = tx_last_response.take() {
let _ = sender.send(LastResponse {
response_id: response_id.clone(),
@@ -1625,6 +1712,7 @@ where
}
Err(err) => {
let mapped = map_api_error(err);
inference_trace_attempt.record_failed(&mapped);
if !logged_error {
session_telemetry.see_event_completed_failed(&mapped);
logged_error = true;
codex-rs/core/src/compact_remote.rs
+41 -1
View File
@@ -13,6 +13,7 @@ use crate::context_manager::ContextManager;
use crate::context_manager::TotalTokenUsageBreakdown;
use crate::context_manager::estimate_response_item_model_visible_bytes;
use crate::context_manager::is_codex_generated_item;
use crate::rollout_trace::CompactionCheckpointTracePayload;
use codex_analytics::CompactionImplementation;
use codex_analytics::CompactionPhase;
use codex_analytics::CompactionReason;
@@ -26,6 +27,7 @@ use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::CompactedItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::TurnStartedEvent;
use codex_rollout_trace::CompactionTraceContext;
use futures::TryFutureExt;
use tokio_util::sync::CancellationToken;
use tracing::error;
@@ -114,7 +116,11 @@ async fn run_remote_compact_task_inner_impl(
turn_context: &Arc<TurnContext>,
initial_context_injection: InitialContextInjection,
) -> CodexResult<()> {
let compaction_item = TurnItem::ContextCompaction(ContextCompactionItem::new());
let context_compaction_item = ContextCompactionItem::new();
// Use the UI compaction item ID as the trace compaction ID so protocol lifecycle events,
// endpoint attempts, and the installed history checkpoint all have one join key.
let compaction_id = context_compaction_item.id.clone();
let compaction_item = TurnItem::ContextCompaction(context_compaction_item);
sess.emit_turn_item_started(turn_context, &compaction_item)
.await;
let mut history = sess.clone_history().await;
@@ -131,6 +137,10 @@ async fn run_remote_compact_task_inner_impl(
"trimmed history items before remote compaction"
);
}
// This is the history selected for remote compaction, after any trimming required to fit the
// compact endpoint. The checkpoint below records it separately from the next sampling request,
// whose prompt will repeat current developer/context prefix items.
let trace_input_history = history.raw_items().to_vec();
// Required to keep `/undo` available after compaction
let ghost_snapshots: Vec<ResponseItem> = history
.raw_items()
@@ -157,6 +167,21 @@ async fn run_remote_compact_task_inner_impl(
personality: turn_context.personality,
output_schema: None,
};
// Remote compaction is the only compaction shape rollout tracing supports. The trace context
// records the exact `/responses/compact` request and response; normal sampling requests remain
// traced through the inference path.
let compaction_trace = sess.services.rollout_trace.as_ref().map_or_else(
CompactionTraceContext::disabled,
|trace| {
trace.compaction_trace_context(
sess.conversation_id.to_string(),
turn_context.sub_id.clone(),
compaction_id.clone(),
turn_context.model_info.slug.clone(),
turn_context.provider.name.clone(),
)
},
);
let mut new_history = sess
.services
@@ -167,6 +192,7 @@ async fn run_remote_compact_task_inner_impl(
turn_context.reasoning_effort,
turn_context.reasoning_summary,
&turn_context.session_telemetry,
&compaction_trace,
)
.or_else(|err| async {
let total_usage_breakdown = sess.get_total_token_usage_breakdown().await;
@@ -200,6 +226,20 @@ async fn run_remote_compact_task_inner_impl(
message: String::new(),
replacement_history: Some(new_history.clone()),
};
if let Some(trace) = sess.services.rollout_trace.as_ref() {
// Install is the semantic boundary where the compact endpoint's output becomes live
// thread history. Keep it distinct from the later inference request so the reducer can
// still represent repeated developer/context prefix items exactly as the model saw them.
trace.record_compaction_installed(
sess.conversation_id.to_string(),
turn_context.sub_id.clone(),
compaction_id,
&CompactionCheckpointTracePayload {
input_history: &trace_input_history,
replacement_history: &new_history,
},
);
}
sess.replace_compacted_history(new_history, reference_context_item, compacted_item)
.await;
sess.recompute_token_usage(turn_context).await;
codex-rs/core/src/lib.rs
+1
View File
@@ -137,6 +137,7 @@ pub use project_doc::LOCAL_PROJECT_DOC_FILENAME;
pub use project_doc::discover_project_doc_paths;
pub use project_doc::read_project_docs;
mod rollout;
mod rollout_trace;
pub(crate) mod safety;
mod session_rollout_init_error;
pub mod shell;
codex-rs/core/src/rollout_trace.rs
+886
View File
@@ -0,0 +1,886 @@
//! Opt-in producer for the rollout trace bundle.
//!
//! This module is the deliberately thin bridge from `codex-core` into
//! `codex-rollout-trace`. Core emits raw observations; the trace crate's
//! offline reducer owns the semantic graph.
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use crate::agent::AgentStatus;
use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolOutput;
use crate::tools::context::ToolPayload;
use codex_protocol::ThreadId;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ExecCommandSource;
use codex_protocol::protocol::ExecCommandStatus;
use codex_protocol::protocol::PatchApplyStatus;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::TurnAbortReason;
use codex_rollout_trace::AgentThreadId;
use codex_rollout_trace::CodeCellRuntimeStatus;
use codex_rollout_trace::CodeModeRuntimeToolId;
use codex_rollout_trace::CompactionTraceContext;
use codex_rollout_trace::ExecutionStatus;
use codex_rollout_trace::InferenceTraceContext;
use codex_rollout_trace::ModelVisibleCallId;
use codex_rollout_trace::RawPayloadKind;
use codex_rollout_trace::RawPayloadRef;
use codex_rollout_trace::RawToolCallRequester;
use codex_rollout_trace::RawTraceEventContext;
use codex_rollout_trace::RawTraceEventPayload;
use codex_rollout_trace::RolloutStatus;
use codex_rollout_trace::ToolCallKind;
use codex_rollout_trace::ToolCallSummary;
use codex_rollout_trace::TraceWriter;
use serde::Serialize;
use tracing::debug;
use tracing::warn;
use uuid::Uuid;
/// Environment variable that enables local trace-bundle recording.
///
/// The value is a root directory. Each independent root session gets one child
/// bundle directory. Spawned child threads share their root session's bundle so
/// one reduced `state.json` describes the whole multi-agent rollout tree.
pub(crate) const CODEX_ROLLOUT_TRACE_ROOT_ENV: &str = "CODEX_ROLLOUT_TRACE_ROOT";
/// Lightweight handle stored in `SessionServices`.
///
/// Cloning the handle is cheap; all sequencing and file ownership remains
/// inside `TraceWriter`.
#[derive(Clone, Debug)]
pub(crate) struct RolloutTraceRecorder {
writer: Arc<TraceWriter>,
root_thread_id: AgentThreadId,
}
/// Metadata captured once at thread/session start.
///
/// This payload is intentionally operational rather than reduced: it is a raw
/// payload that later reducers can mine as the reduced thread model evolves.
#[derive(Serialize)]
pub(crate) struct ThreadStartedTraceMetadata {
pub(crate) thread_id: String,
pub(crate) agent_path: String,
pub(crate) task_name: Option<String>,
pub(crate) nickname: Option<String>,
pub(crate) agent_role: Option<String>,
pub(crate) session_source: SessionSource,
pub(crate) cwd: PathBuf,
pub(crate) rollout_path: Option<PathBuf>,
pub(crate) model: String,
pub(crate) provider_name: String,
pub(crate) approval_policy: String,
pub(crate) sandbox_policy: String,
}
/// History replacement checkpoint persisted when compaction installs new live history.
///
/// The checkpoint keeps compaction separate from ordinary sampling snapshots:
/// `input_history` is the live thread history selected for compaction, while
/// `replacement_history` is what future prompts may carry after the checkpoint.
#[derive(Serialize)]
pub(crate) struct CompactionCheckpointTracePayload<'a> {
pub(crate) input_history: &'a [ResponseItem],
pub(crate) replacement_history: &'a [ResponseItem],
}
/// Raw invocation payload for the canonical Codex tool boundary.
///
/// Protocol events may add runtime detail later, but this envelope preserves
/// the caller-facing request for both direct model calls and code-mode nested
/// calls.
#[derive(Serialize)]
struct DispatchedToolTraceRequest<'a> {
tool_name: &'a str,
tool_namespace: Option<&'a str>,
payload: serde_json::Value,
}
/// Raw response payload for dispatch-level tool trace events.
#[derive(Serialize)]
#[serde(rename_all = "snake_case", tag = "type")]
enum DispatchedToolTraceResponse<'a> {
DirectResponse {
response_item: &'a ResponseInputItem,
},
CodeModeResponse {
value: serde_json::Value,
},
Error {
error: &'a str,
},
}
/// Raw code-mode response captured at the runtime boundary.
///
/// The reducer keeps the graph small and uses this payload as evidence for
/// future viewers that need exact content items or stored-value details.
#[derive(Serialize)]
struct CodeCellResponseTracePayload<'a> {
response: &'a codex_code_mode::RuntimeResponse,
}
/// Trace-only payload for the notification a finished child sends back to its parent.
#[derive(Serialize)]
struct AgentResultTracePayload<'a> {
child_agent_path: &'a str,
message: &'a str,
status: &'a AgentStatus,
}
impl RolloutTraceRecorder {
/// Creates and starts a trace bundle if `CODEX_ROLLOUT_TRACE_ROOT` is set.
///
/// Trace startup is best-effort. A tracing failure must not make the Codex
/// session unusable, because traces are diagnostic and can be enabled while
/// debugging unrelated production failures.
pub(crate) fn maybe_create(
thread_id: ThreadId,
metadata: ThreadStartedTraceMetadata,
) -> Option<Self> {
let root = std::env::var_os(CODEX_ROLLOUT_TRACE_ROOT_ENV)?;
let root = PathBuf::from(root);
match Self::create_in_root(root.as_path(), thread_id, metadata) {
Ok(recorder) => Some(recorder),
Err(err) => {
warn!("failed to initialize rollout trace recorder: {err:#}");
None
}
}
}
fn create_in_root(
root: &Path,
thread_id: ThreadId,
metadata: ThreadStartedTraceMetadata,
) -> anyhow::Result<Self> {
let trace_id = Uuid::new_v4().to_string();
let thread_id = thread_id.to_string();
let bundle_dir = root.join(format!("trace-{trace_id}-{thread_id}"));
let writer = TraceWriter::create(
&bundle_dir,
trace_id.clone(),
thread_id.clone(),
thread_id.clone(),
)?;
let recorder = Self {
writer: Arc::new(writer),
root_thread_id: thread_id.clone(),
};
recorder.append_best_effort(RawTraceEventPayload::RolloutStarted {
trace_id,
root_thread_id: thread_id,
});
recorder.record_thread_started(metadata);
debug!("recording rollout trace at {}", bundle_dir.display());
Ok(recorder)
}
#[cfg(test)]
pub(crate) fn create_in_root_for_test(
root: &Path,
thread_id: ThreadId,
metadata: ThreadStartedTraceMetadata,
) -> anyhow::Result<Self> {
Self::create_in_root(root, thread_id, metadata)
}
/// Wraps selected UI/protocol events in the trace bundle.
///
/// We intentionally skip high-volume stream deltas here. Inference/tool
/// hooks emit typed raw events; protocol wrappers are debug breadcrumbs, not
/// the canonical transcript.
pub(crate) fn record_protocol_event(&self, event: &EventMsg) {
let Some(event_type) = wrapped_protocol_event_type(event) else {
return;
};
let event_payload =
match self.write_json_payload_best_effort(RawPayloadKind::ProtocolEvent, event) {
Some(event_payload) => event_payload,
None => return,
};
self.append_best_effort(RawTraceEventPayload::ProtocolEventObserved {
event_type: event_type.to_string(),
event_payload,
});
}
/// Emits the lifecycle event and metadata for one thread in this rollout tree.
///
/// Root sessions call this immediately after `RolloutStarted`; spawned
/// child sessions call it on the inherited recorder. Keeping children in
/// the root bundle preserves one raw payload namespace and one reduced
/// `RolloutTrace` for the whole multi-agent task.
pub(crate) fn record_thread_started(&self, metadata: ThreadStartedTraceMetadata) {
let metadata_payload =
self.write_json_payload_best_effort(RawPayloadKind::SessionMetadata, &metadata);
self.append_best_effort(RawTraceEventPayload::ThreadStarted {
thread_id: metadata.thread_id,
agent_path: metadata.agent_path,
metadata_payload,
});
}
/// Emits typed turn lifecycle events from the UI/protocol lifecycle.
pub(crate) fn record_codex_turn_event(
&self,
thread_id: AgentThreadId,
default_turn_id: &str,
event: &EventMsg,
) {
match event {
EventMsg::TurnStarted(event) => {
self.append_with_context_best_effort(
thread_id.clone(),
event.turn_id.clone(),
RawTraceEventPayload::CodexTurnStarted {
codex_turn_id: event.turn_id.clone(),
thread_id,
},
);
}
EventMsg::TurnComplete(event) => {
self.append_with_context_best_effort(
thread_id,
event.turn_id.clone(),
RawTraceEventPayload::CodexTurnEnded {
codex_turn_id: event.turn_id.clone(),
status: ExecutionStatus::Completed,
},
);
}
EventMsg::TurnAborted(event) => {
let turn_id = event
.turn_id
.clone()
.unwrap_or_else(|| default_turn_id.to_string());
self.append_with_context_best_effort(
thread_id,
turn_id.clone(),
RawTraceEventPayload::CodexTurnEnded {
codex_turn_id: turn_id,
status: execution_status_for_abort_reason(&event.reason),
},
);
}
_ => {}
}
}
/// Emits typed runtime tool events from existing protocol lifecycle events.
///
/// The protocol event stays separate from the caller-facing invocation and
/// result payloads. Reducers attach it to `ToolCall.raw_runtime_payload_ids`
/// and can also use it to build richer objects such as terminal operations.
pub(crate) fn record_tool_call_event(
&self,
thread_id: AgentThreadId,
codex_turn_id: String,
event: &EventMsg,
) {
let Some(payload) = self.tool_call_trace_payload(event) else {
return;
};
self.append_with_context_best_effort(thread_id, codex_turn_id, payload);
}
/// Emits the parent runtime object for one model-authored code-mode cell.
///
/// This must run before JavaScript starts because the runtime can request
/// nested tools before the initial custom-tool response is available.
pub(crate) fn record_code_cell_started(
&self,
thread_id: AgentThreadId,
codex_turn_id: String,
runtime_cell_id: &str,
model_visible_call_id: &str,
source_js: &str,
) {
self.append_with_context_best_effort(
thread_id,
codex_turn_id,
RawTraceEventPayload::CodeCellStarted {
runtime_cell_id: runtime_cell_id.to_string(),
model_visible_call_id: model_visible_call_id.to_string(),
source_js: source_js.to_string(),
},
);
}
/// Emits the first response returned by the public code-mode `exec` tool.
///
/// A yielded response returns control to the model while the cell keeps
/// running. A terminal response is followed by `CodeCellEnded` so the
/// reducer can distinguish "first model-visible output" from runtime end.
pub(crate) fn record_code_cell_initial_response(
&self,
thread_id: AgentThreadId,
codex_turn_id: String,
response: &codex_code_mode::RuntimeResponse,
) {
let response_payload = self.code_cell_response_payload(response);
self.append_with_context_best_effort(
thread_id,
codex_turn_id,
RawTraceEventPayload::CodeCellInitialResponse {
runtime_cell_id: code_cell_runtime_id(response).to_string(),
status: code_cell_status_for_runtime_response(response),
response_payload,
},
);
}
/// Emits the terminal lifecycle point for a code-mode cell.
pub(crate) fn record_code_cell_ended(
&self,
thread_id: AgentThreadId,
codex_turn_id: String,
response: &codex_code_mode::RuntimeResponse,
) {
let response_payload = self.code_cell_response_payload(response);
self.append_with_context_best_effort(
thread_id,
codex_turn_id,
RawTraceEventPayload::CodeCellEnded {
runtime_cell_id: code_cell_runtime_id(response).to_string(),
status: code_cell_status_for_runtime_response(response),
response_payload,
},
);
}
/// Emits a generic lifecycle start for direct/code-mode tools without a
/// richer protocol-backed lifecycle.
///
/// The registry calls this after it has resolved a concrete handler. At that
/// point we know the tool call is valid, but we are still before
/// approval/pre-use hooks, so blocked tools are represented as failed tool
/// executions instead of disappearing from the trace.
pub(crate) fn record_dispatched_tool_call_started(&self, invocation: &ToolInvocation) {
let request = DispatchedToolTraceRequest {
tool_name: invocation.tool_name.name.as_str(),
tool_namespace: invocation.tool_name.namespace.as_deref(),
payload: dispatched_tool_payload(&invocation.payload),
};
let request_payload =
self.write_json_payload_best_effort(RawPayloadKind::ToolInvocation, &request);
let (model_visible_call_id, code_mode_runtime_tool_id, requester) =
dispatched_tool_requester_fields(invocation);
self.append_with_context_best_effort(
invocation.session.conversation_id.to_string(),
invocation.turn.sub_id.clone(),
RawTraceEventPayload::ToolCallStarted {
tool_call_id: invocation.call_id.clone(),
model_visible_call_id,
code_mode_runtime_tool_id,
requester,
kind: dispatched_tool_kind(invocation),
summary: ToolCallSummary::Generic {
label: dispatched_tool_label(invocation),
input_preview: Some(truncate_preview(&invocation.payload.log_payload())),
output_preview: None,
},
invocation_payload: request_payload,
},
);
}
/// Emits the caller-facing result for a dispatch-level tool lifecycle.
pub(crate) fn record_dispatched_tool_call_ended(
&self,
invocation: &ToolInvocation,
status: ExecutionStatus,
result: &dyn ToolOutput,
response_call_id: &str,
tool_payload: &ToolPayload,
) {
let direct_response_item;
let response = match invocation.source {
ToolCallSource::Direct => {
direct_response_item = result.to_response_item(response_call_id, tool_payload);
DispatchedToolTraceResponse::DirectResponse {
response_item: &direct_response_item,
}
}
ToolCallSource::CodeMode { .. } => DispatchedToolTraceResponse::CodeModeResponse {
value: result.code_mode_result(tool_payload),
},
ToolCallSource::JsRepl => return,
};
self.append_dispatched_tool_call_ended(invocation, status, &response);
}
/// Emits a failed end event for a dispatch-level tool lifecycle.
pub(crate) fn record_dispatched_tool_call_failed(
&self,
invocation: &ToolInvocation,
error: &str,
) {
let response = DispatchedToolTraceResponse::Error { error };
self.append_dispatched_tool_call_ended(invocation, ExecutionStatus::Failed, &response);
}
fn append_dispatched_tool_call_ended(
&self,
invocation: &ToolInvocation,
status: ExecutionStatus,
response: &DispatchedToolTraceResponse<'_>,
) {
let response_payload =
self.write_json_payload_best_effort(RawPayloadKind::ToolResult, response);
self.append_with_context_best_effort(
invocation.session.conversation_id.to_string(),
invocation.turn.sub_id.clone(),
RawTraceEventPayload::ToolCallEnded {
tool_call_id: invocation.call_id.clone(),
status,
result_payload: response_payload,
},
);
}
/// Builds reusable inference trace context for one Codex turn.
///
/// The returned context is intentionally not "an inference call" yet.
/// Transport code owns retry/fallback attempts and calls `start_attempt`
/// only after it has built the concrete request payload for that attempt.
pub(crate) fn inference_trace_context(
&self,
thread_id: AgentThreadId,
codex_turn_id: String,
model: String,
provider_name: String,
) -> InferenceTraceContext {
InferenceTraceContext::enabled(
Arc::clone(&self.writer),
thread_id,
codex_turn_id,
model,
provider_name,
)
}
/// Builds remote-compaction trace context for one checkpoint.
///
/// Rollout tracing currently has a first-class checkpoint model only for remote compaction.
/// The compact endpoint is a model-facing request whose output replaces live history, so it
/// needs both request/response attempt events and a later checkpoint event when processed
/// replacement history is installed.
pub(crate) fn compaction_trace_context(
&self,
thread_id: AgentThreadId,
codex_turn_id: String,
compaction_id: String,
model: String,
provider_name: String,
) -> CompactionTraceContext {
CompactionTraceContext::enabled(
Arc::clone(&self.writer),
thread_id,
codex_turn_id,
compaction_id,
model,
provider_name,
)
}
/// Emits the checkpoint where remote-compacted history replaces live thread history.
///
/// This checkpoint is deliberately separate from the compact endpoint response: Codex filters
/// and reinjects context before replacement history becomes live. The reducer uses this event
/// to connect the pre-compaction history to the processed replacement items without treating
/// repeated developer/context prefix items as part of the replacement itself.
pub(crate) fn record_compaction_installed(
&self,
thread_id: AgentThreadId,
codex_turn_id: String,
compaction_id: String,
checkpoint: &CompactionCheckpointTracePayload<'_>,
) {
let Some(checkpoint_payload) =
self.write_json_payload_best_effort(RawPayloadKind::CompactionCheckpoint, checkpoint)
else {
return;
};
self.append_with_context_best_effort(
thread_id,
codex_turn_id,
RawTraceEventPayload::CompactionInstalled {
compaction_id,
checkpoint_payload,
},
);
}
/// Emits the v2 child-to-parent completion message as an explicit graph edge.
///
/// This notification is not a tool call in the child: it is runtime delivery
/// from the completed child turn into the parent's mailbox. Without a
/// trace-owned edge the viewer would have to infer the relationship from a
/// later parent prompt snapshot, which loses the runtime timing and source.
pub(crate) fn record_agent_result_interaction(
&self,
child_thread_id: AgentThreadId,
child_codex_turn_id: String,
parent_thread_id: AgentThreadId,
child_agent_path: &str,
message: &str,
status: &AgentStatus,
) {
let carried_payload = self.write_json_payload_best_effort(
RawPayloadKind::AgentResult,
&AgentResultTracePayload {
child_agent_path,
message,
status,
},
);
self.append_with_context_best_effort(
child_thread_id.clone(),
child_codex_turn_id.clone(),
RawTraceEventPayload::AgentResultObserved {
edge_id: format!(
"edge:agent_result:{child_thread_id}:{child_codex_turn_id}:{parent_thread_id}"
),
child_thread_id,
child_codex_turn_id,
parent_thread_id,
message: message.to_string(),
carried_payload,
},
);
}
/// Emits terminal trace events for graceful session shutdown.
///
/// Child agent sessions share their root recorder, so ending a child thread
/// must not close the whole rollout. Only the root thread's shutdown emits
/// `RolloutEnded`.
pub(crate) fn record_thread_ended(&self, thread_id: AgentThreadId, status: RolloutStatus) {
self.append_best_effort(RawTraceEventPayload::ThreadEnded {
thread_id: thread_id.clone(),
status: status.clone(),
});
if thread_id == self.root_thread_id {
self.append_best_effort(RawTraceEventPayload::RolloutEnded { status });
}
}
fn write_json_payload_best_effort(
&self,
kind: RawPayloadKind,
payload: &impl Serialize,
) -> Option<codex_rollout_trace::RawPayloadRef> {
match self.writer.write_json_payload(kind, payload) {
Ok(payload_ref) => Some(payload_ref),
Err(err) => {
warn!("failed to write rollout trace payload: {err:#}");
None
}
}
}
fn append_best_effort(&self, payload: RawTraceEventPayload) {
if let Err(err) = self.writer.append(payload) {
warn!("failed to append rollout trace event: {err:#}");
}
}
fn append_with_context_best_effort(
&self,
thread_id: AgentThreadId,
codex_turn_id: String,
payload: RawTraceEventPayload,
) {
let context = RawTraceEventContext {
thread_id: Some(thread_id),
codex_turn_id: Some(codex_turn_id),
};
if let Err(err) = self.writer.append_with_context(context, payload) {
warn!("failed to append rollout trace event: {err:#}");
}
}
fn tool_call_trace_payload(&self, event: &EventMsg) -> Option<RawTraceEventPayload> {
match event {
EventMsg::ExecCommandBegin(event) if event.source != ExecCommandSource::UserShell => {
self.tool_runtime_started_payload(&event.call_id, event)
}
EventMsg::ExecCommandEnd(event) if event.source != ExecCommandSource::UserShell => self
.tool_runtime_ended_payload(
&event.call_id,
execution_status_for_exec_status(&event.status),
event,
),
EventMsg::PatchApplyBegin(event) => {
self.tool_runtime_started_payload(&event.call_id, event)
}
EventMsg::PatchApplyEnd(event) => self.tool_runtime_ended_payload(
&event.call_id,
execution_status_for_patch_status(&event.status),
event,
),
EventMsg::McpToolCallBegin(event) => {
self.tool_runtime_started_payload(&event.call_id, event)
}
EventMsg::McpToolCallEnd(event) => self.tool_runtime_ended_payload(
&event.call_id,
if event.result.is_ok() {
ExecutionStatus::Completed
} else {
ExecutionStatus::Failed
},
event,
),
EventMsg::CollabAgentSpawnBegin(event) => {
self.tool_runtime_started_payload(&event.call_id, event)
}
EventMsg::CollabAgentSpawnEnd(event) => self.tool_runtime_ended_payload(
&event.call_id,
if event.new_thread_id.is_some() {
ExecutionStatus::Completed
} else {
ExecutionStatus::Failed
},
event,
),
EventMsg::CollabAgentInteractionBegin(event) => {
self.tool_runtime_started_payload(&event.call_id, event)
}
EventMsg::CollabAgentInteractionEnd(event) => {
self.tool_runtime_ended_payload(&event.call_id, ExecutionStatus::Completed, event)
}
EventMsg::CollabWaitingBegin(event) => {
self.tool_runtime_started_payload(&event.call_id, event)
}
EventMsg::CollabWaitingEnd(event) => {
self.tool_runtime_ended_payload(&event.call_id, ExecutionStatus::Completed, event)
}
EventMsg::CollabCloseBegin(event) => {
self.tool_runtime_started_payload(&event.call_id, event)
}
EventMsg::CollabCloseEnd(event) => {
self.tool_runtime_ended_payload(&event.call_id, ExecutionStatus::Completed, event)
}
_ => None,
}
}
fn tool_runtime_started_payload(
&self,
tool_call_id: &str,
event: &impl Serialize,
) -> Option<RawTraceEventPayload> {
let runtime_payload =
self.write_json_payload_best_effort(RawPayloadKind::ToolRuntimeEvent, event)?;
Some(RawTraceEventPayload::ToolCallRuntimeStarted {
tool_call_id: tool_call_id.to_string(),
runtime_payload,
})
}
fn tool_runtime_ended_payload(
&self,
tool_call_id: &str,
status: ExecutionStatus,
event: &impl Serialize,
) -> Option<RawTraceEventPayload> {
let runtime_payload =
self.write_json_payload_best_effort(RawPayloadKind::ToolRuntimeEvent, event)?;
Some(RawTraceEventPayload::ToolCallRuntimeEnded {
tool_call_id: tool_call_id.to_string(),
status,
runtime_payload,
})
}
fn code_cell_response_payload(
&self,
response: &codex_code_mode::RuntimeResponse,
) -> Option<RawPayloadRef> {
self.write_json_payload_best_effort(
RawPayloadKind::ToolResult,
&CodeCellResponseTracePayload { response },
)
}
}
fn execution_status_for_abort_reason(reason: &TurnAbortReason) -> ExecutionStatus {
match reason {
TurnAbortReason::Interrupted | TurnAbortReason::Replaced | TurnAbortReason::ReviewEnded => {
ExecutionStatus::Cancelled
}
}
}
fn execution_status_for_exec_status(status: &ExecCommandStatus) -> ExecutionStatus {
match status {
ExecCommandStatus::Completed => ExecutionStatus::Completed,
ExecCommandStatus::Failed => ExecutionStatus::Failed,
ExecCommandStatus::Declined => ExecutionStatus::Cancelled,
}
}
fn execution_status_for_patch_status(status: &PatchApplyStatus) -> ExecutionStatus {
match status {
PatchApplyStatus::Completed => ExecutionStatus::Completed,
PatchApplyStatus::Failed => ExecutionStatus::Failed,
PatchApplyStatus::Declined => ExecutionStatus::Cancelled,
}
}
fn code_cell_runtime_id(response: &codex_code_mode::RuntimeResponse) -> &str {
match response {
codex_code_mode::RuntimeResponse::Yielded { cell_id, .. }
| codex_code_mode::RuntimeResponse::Terminated { cell_id, .. }
| codex_code_mode::RuntimeResponse::Result { cell_id, .. } => cell_id,
}
}
fn code_cell_status_for_runtime_response(
response: &codex_code_mode::RuntimeResponse,
) -> CodeCellRuntimeStatus {
match response {
codex_code_mode::RuntimeResponse::Yielded { .. } => CodeCellRuntimeStatus::Yielded,
codex_code_mode::RuntimeResponse::Terminated { .. } => CodeCellRuntimeStatus::Terminated,
codex_code_mode::RuntimeResponse::Result { error_text, .. } => {
if error_text.is_some() {
CodeCellRuntimeStatus::Failed
} else {
CodeCellRuntimeStatus::Completed
}
}
}
}
fn dispatched_tool_requester_fields(
invocation: &ToolInvocation,
) -> (
Option<ModelVisibleCallId>,
Option<CodeModeRuntimeToolId>,
RawToolCallRequester,
) {
match &invocation.source {
ToolCallSource::Direct => (
Some(invocation.call_id.clone()),
None,
RawToolCallRequester::Model,
),
ToolCallSource::CodeMode {
cell_id,
runtime_tool_call_id,
} => (
None,
Some(runtime_tool_call_id.clone()),
RawToolCallRequester::CodeCell {
runtime_cell_id: cell_id.clone(),
},
),
ToolCallSource::JsRepl => (None, None, RawToolCallRequester::Model),
}
}
fn dispatched_tool_kind(invocation: &ToolInvocation) -> ToolCallKind {
if let ToolPayload::Mcp { server, tool, .. } = &invocation.payload {
return ToolCallKind::Mcp {
server: server.clone(),
tool: tool.clone(),
};
}
match invocation.tool_name.name.as_str() {
"exec_command" | "local_shell" | "shell" | "shell_command" => ToolCallKind::ExecCommand,
"write_stdin" => ToolCallKind::WriteStdin,
"apply_patch" => ToolCallKind::ApplyPatch,
"web_search" | "web_search_preview" => ToolCallKind::Web,
"image_generation" | "image_query" => ToolCallKind::ImageGeneration,
"spawn_agent" => ToolCallKind::SpawnAgent,
"send_message" => ToolCallKind::SendMessage,
"followup_task" => ToolCallKind::AssignAgentTask,
"wait_agent" => ToolCallKind::WaitAgent,
"close_agent" => ToolCallKind::CloseAgent,
other => ToolCallKind::Other {
name: other.to_string(),
},
}
}
fn dispatched_tool_label(invocation: &ToolInvocation) -> String {
if let ToolPayload::Mcp { server, tool, .. } = &invocation.payload {
return format!("mcp:{server}:{tool}");
}
invocation.tool_name.to_string()
}
fn dispatched_tool_payload(payload: &ToolPayload) -> serde_json::Value {
match payload {
ToolPayload::Function { arguments } => serde_json::json!({
"type": "function",
"arguments": arguments,
}),
ToolPayload::ToolSearch { arguments } => serde_json::json!({
"type": "tool_search",
"arguments": arguments,
}),
ToolPayload::Custom { input } => serde_json::json!({
"type": "custom",
"input": input,
}),
ToolPayload::LocalShell { params } => serde_json::json!({
"type": "local_shell",
"command": params.command,
"workdir": params.workdir,
"timeout_ms": params.timeout_ms,
}),
ToolPayload::Mcp {
server,
tool,
raw_arguments,
} => serde_json::json!({
"type": "mcp",
"server": server,
"tool": tool,
"raw_arguments": raw_arguments,
}),
}
}
fn truncate_preview(value: &str) -> String {
const MAX_PREVIEW_CHARS: usize = 160;
let mut preview = value.chars().take(MAX_PREVIEW_CHARS).collect::<String>();
if value.chars().count() > MAX_PREVIEW_CHARS {
preview.push_str("...");
}
preview
}
fn wrapped_protocol_event_type(event: &EventMsg) -> Option<&'static str> {
match event {
EventMsg::SessionConfigured(_) => Some("session_configured"),
EventMsg::TurnStarted(_) => Some("turn_started"),
EventMsg::TurnComplete(_) => Some("turn_complete"),
EventMsg::TurnAborted(_) => Some("turn_aborted"),
EventMsg::ThreadNameUpdated(_) => Some("thread_name_updated"),
EventMsg::ThreadRolledBack(_) => Some("thread_rolled_back"),
EventMsg::Error(_) => Some("error"),
EventMsg::Warning(_) => Some("warning"),
EventMsg::ShutdownComplete => Some("shutdown_complete"),
_ => None,
}
}
#[cfg(test)]
#[path = "rollout_trace_tests.rs"]
mod tests;
codex-rs/core/src/rollout_trace_tests.rs
+163
View File
@@ -0,0 +1,163 @@
use std::fs;
use std::path::Path;
use std::path::PathBuf;
use codex_protocol::AgentPath;
use codex_protocol::ThreadId;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_rollout_trace::ExecutionStatus;
use codex_rollout_trace::RawTraceEventPayload;
use codex_rollout_trace::RolloutStatus;
use tempfile::TempDir;
use super::*;
#[test]
fn create_in_root_writes_replayable_lifecycle_events() -> anyhow::Result<()> {
let temp = TempDir::new()?;
let thread_id = ThreadId::new();
let recorder = RolloutTraceRecorder::create_in_root(
temp.path(),
thread_id,
ThreadStartedTraceMetadata {
thread_id: thread_id.to_string(),
agent_path: "/root".to_string(),
task_name: None,
nickname: None,
agent_role: None,
session_source: SessionSource::Exec,
cwd: PathBuf::from("/workspace"),
rollout_path: Some(PathBuf::from("/tmp/rollout.jsonl")),
model: "gpt-test".to_string(),
provider_name: "test-provider".to_string(),
approval_policy: "never".to_string(),
sandbox_policy: format!("{:?}", SandboxPolicy::DangerFullAccess),
},
)
.expect("trace recorder");
recorder.record_thread_ended(thread_id.to_string(), RolloutStatus::Completed);
let bundle_dir = single_bundle_dir(temp.path())?;
let replayed = codex_rollout_trace::replay_bundle(&bundle_dir)?;
assert_eq!(replayed.status, RolloutStatus::Completed);
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);
Ok(())
}
#[test]
fn spawned_thread_start_appends_to_root_bundle() -> anyhow::Result<()> {
let temp = TempDir::new()?;
let root_thread_id = ThreadId::new();
let child_thread_id = ThreadId::new();
let recorder = RolloutTraceRecorder::create_in_root(
temp.path(),
root_thread_id,
minimal_metadata(root_thread_id),
)
.expect("trace recorder");
recorder.record_thread_started(ThreadStartedTraceMetadata {
thread_id: child_thread_id.to_string(),
agent_path: "/root/repo_file_counter".to_string(),
task_name: Some("repo_file_counter".to_string()),
nickname: Some("Kepler".to_string()),
agent_role: Some("worker".to_string()),
session_source: SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id: root_thread_id,
depth: 1,
agent_path: Some(
AgentPath::try_from("/root/repo_file_counter").map_err(anyhow::Error::msg)?,
),
agent_nickname: Some("Kepler".to_string()),
agent_role: Some("worker".to_string()),
}),
cwd: PathBuf::from("/workspace"),
rollout_path: Some(PathBuf::from("/tmp/child-rollout.jsonl")),
model: "gpt-test".to_string(),
provider_name: "test-provider".to_string(),
approval_policy: "never".to_string(),
sandbox_policy: format!("{:?}", SandboxPolicy::DangerFullAccess),
});
recorder.record_thread_ended(child_thread_id.to_string(), RolloutStatus::Completed);
let bundle_dir = single_bundle_dir(temp.path())?;
let replayed = codex_rollout_trace::replay_bundle(&bundle_dir)?;
assert_eq!(fs::read_dir(temp.path())?.count(), 1);
assert_eq!(replayed.threads.len(), 2);
assert_eq!(
replayed.threads[&child_thread_id.to_string()].agent_path,
"/root/repo_file_counter"
);
assert_eq!(replayed.status, RolloutStatus::Running);
assert_eq!(
replayed.threads[&child_thread_id.to_string()]
.execution
.status,
ExecutionStatus::Completed
);
assert_eq!(replayed.raw_payloads.len(), 2);
Ok(())
}
#[test]
fn protocol_wrapper_records_selected_events_as_raw_payloads() -> anyhow::Result<()> {
let temp = TempDir::new()?;
let thread_id = ThreadId::new();
let recorder =
RolloutTraceRecorder::create_in_root(temp.path(), thread_id, minimal_metadata(thread_id))
.expect("trace recorder");
recorder.record_protocol_event(&EventMsg::ShutdownComplete);
let event_log = fs::read_to_string(single_bundle_dir(temp.path())?.join("trace.jsonl"))?;
let protocol_event_seen = event_log.lines().any(|line| {
let event: codex_rollout_trace::RawTraceEvent =
serde_json::from_str(line).expect("raw trace event");
matches!(
event.payload,
RawTraceEventPayload::ProtocolEventObserved {
event_type,
..
} if event_type == "shutdown_complete"
)
});
assert!(protocol_event_seen);
Ok(())
}
fn minimal_metadata(thread_id: ThreadId) -> ThreadStartedTraceMetadata {
ThreadStartedTraceMetadata {
thread_id: thread_id.to_string(),
agent_path: "/root".to_string(),
task_name: None,
nickname: None,
agent_role: None,
session_source: SessionSource::Exec,
cwd: PathBuf::from("/workspace"),
rollout_path: None,
model: "gpt-test".to_string(),
provider_name: "test-provider".to_string(),
approval_policy: "never".to_string(),
sandbox_policy: "danger-full-access".to_string(),
}
}
fn single_bundle_dir(root: &Path) -> anyhow::Result<PathBuf> {
let mut entries = fs::read_dir(root)?
.map(|entry| entry.map(|entry| entry.path()))
.collect::<Result<Vec<_>, _>>()?;
entries.sort();
assert_eq!(entries.len(), 1);
Ok(entries.remove(0))
}