[rollout_trace] Record core session rollout traces (#18877)

## Summary Wires rollout trace recording into `codex-core` session and turn execution. This records the core model request/response, compaction, and session lifecycle boundaries needed for replay without yet tracing every nested runtime/tool boundary. ## Stack This is PR 2/5 in the rollout trace stack. - [#18876](https://github.com/openai/codex/pull/18876): Add rollout trace crate - [#18877](https://github.com/openai/codex/pull/18877): Record core session rollout traces - [#18878](https://github.com/openai/codex/pull/18878): Trace tool and code-mode boundaries - [#18879](https://github.com/openai/codex/pull/18879): Trace sessions and multi-agent edges - [#18880](https://github.com/openai/codex/pull/18880): Add debug trace reduction command ## Review Notes This layer is the first live integration point. The important review question is whether trace recording is isolated from normal session behavior: trace failures should not become user-visible execution failures, and recording should preserve the existing turn/session lifecycle semantics. The PR depends on the reducer/data model from the first stack entry and only introduces the core recorder surface that later PRs use for richer runtime and relationship events.
2026-07-01 00:31:56 +08:00 · 2026-04-22 17:00:48 +00:00
parent 79ea577156
commit f67383bcba
23 changed files with 627 additions and 17 deletions
@@ -13,6 +13,7 @@ use std::sync::atomic::Ordering;
 use codex_protocol::models::ResponseItem;
 use serde::Serialize;
 use serde_json::Value as JsonValue;
+use tracing::warn;

 use crate::inference::trace_response_item_json;
 use crate::model::AgentThreadId;
@@ -74,6 +75,17 @@ struct TracedCompactionCompleted {
    output_items: Vec<JsonValue>,
 }

+/// 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 struct CompactionCheckpointTracePayload<'a> {
+    pub input_history: &'a [ResponseItem],
+    pub replacement_history: &'a [ResponseItem],
+}
+
 impl CompactionTraceContext {
    /// Builds a context that accepts trace calls and records nothing.
    pub fn disabled() -> Self {
@@ -118,6 +130,40 @@ impl CompactionTraceContext {
        attempt.record_started(request);
        attempt
    }
+
+    /// Records the point where compacted history becomes the live thread history.
+    ///
+    /// The checkpoint belongs to the same semantic compaction lifecycle as the
+    /// compact endpoint attempts, so the context reuses its stable compaction ID.
+    pub fn record_installed(&self, checkpoint: &CompactionCheckpointTracePayload<'_>) {
+        let CompactionTraceContextState::Enabled(context) = &self.state else {
+            return;
+        };
+        let checkpoint_payload = match context
+            .writer
+            .write_json_payload(RawPayloadKind::CompactionCheckpoint, checkpoint)
+        {
+            Ok(payload_ref) => payload_ref,
+            Err(err) => {
+                warn!("failed to write rollout trace payload: {err:#}");
+                return;
+            }
+        };
+
+        let event_context = RawTraceEventContext {
+            thread_id: Some(context.thread_id.clone()),
+            codex_turn_id: Some(context.codex_turn_id.clone()),
+        };
+        if let Err(err) = context.writer.append_with_context(
+            event_context,
+            RawTraceEventPayload::CompactionInstalled {
+                compaction_id: context.compaction_id.clone(),
+                checkpoint_payload,
+            },
+        ) {
+            warn!("failed to append rollout trace event: {err:#}");
+        }
+    }
 }

 impl CompactionTraceAttempt {
@@ -184,6 +230,14 @@ impl CompactionTraceAttempt {
        );
    }

+    /// Records the compact endpoint result without forcing callers to branch on trace events.
+    pub fn record_result<E: Display>(&self, result: Result<&[ResponseItem], E>) {
+        match result {
+            Ok(output_items) => self.record_completed(output_items),
+            Err(err) => self.record_failed(err),
+        }
+    }
+
    /// Records pre-response failures from the compact endpoint.
    pub fn record_failed(&self, error: impl Display) {
        let CompactionTraceAttemptState::Enabled(attempt) = &self.state else {
@@ -12,11 +12,14 @@ mod inference;
 mod model;
 mod payload;
 mod raw_event;
+mod recorder;
 mod reducer;
 mod writer;

 /// Conventional reduced-state cache name written next to a raw trace bundle.
 pub use bundle::REDUCED_STATE_FILE_NAME;
+/// Raw checkpoint payload for a remote compaction install event.
+pub use compaction::CompactionCheckpointTracePayload;
 /// No-op-capable handle for recording remote-compaction requests.
 pub use compaction::CompactionTraceAttempt;
 /// Shared recorder context for a compaction checkpoint.
@@ -43,6 +46,12 @@ pub use raw_event::RawTraceEvent;
 pub use raw_event::RawTraceEventContext;
 /// Typed payload for one raw trace event.
 pub use raw_event::RawTraceEventPayload;
+/// Environment variable that enables local trace-bundle recording.
+pub use recorder::CODEX_ROLLOUT_TRACE_ROOT_ENV;
+/// Best-effort hot-path recorder for one rollout trace bundle.
+pub use recorder::RolloutTraceRecorder;
+/// Raw metadata captured when a thread starts.
+pub use recorder::ThreadStartedTraceMetadata;
 /// Replay a raw trace bundle and write/read its reduced `RolloutTrace`.
 pub use reducer::replay_bundle;
 /// Append-only writer used by hot-path Codex instrumentation.
@@ -0,0 +1,228 @@
+//! Opt-in hot-path producer for rollout trace bundles.
+
+use std::path::Path;
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use codex_protocol::ThreadId;
+use codex_protocol::protocol::SessionSource;
+use serde::Serialize;
+use tracing::debug;
+use tracing::warn;
+use uuid::Uuid;
+
+use crate::AgentThreadId;
+use crate::CodexTurnId;
+use crate::CompactionId;
+use crate::CompactionTraceContext;
+use crate::InferenceTraceContext;
+use crate::RawPayloadKind;
+use crate::RawPayloadRef;
+use crate::RawTraceEventPayload;
+use crate::TraceWriter;
+
+/// 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 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`. Disabled handles intentionally accept the same calls
+/// as enabled handles so hot-path session code can describe traceable events
+/// without repeatedly branching on whether diagnostic recording is enabled.
+#[derive(Clone, Debug)]
+pub struct RolloutTraceRecorder {
+    state: RolloutTraceRecorderState,
+}
+
+#[derive(Clone, Debug)]
+enum RolloutTraceRecorderState {
+    Disabled,
+    Enabled(EnabledRolloutTraceRecorder),
+}
+
+#[derive(Clone, Debug)]
+struct EnabledRolloutTraceRecorder {
+    writer: Arc<TraceWriter>,
+}
+
+/// 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 struct ThreadStartedTraceMetadata {
+    pub thread_id: String,
+    pub agent_path: String,
+    pub task_name: Option<String>,
+    pub nickname: Option<String>,
+    pub agent_role: Option<String>,
+    pub session_source: SessionSource,
+    pub cwd: PathBuf,
+    pub rollout_path: Option<PathBuf>,
+    pub model: String,
+    pub provider_name: String,
+    pub approval_policy: String,
+    pub sandbox_policy: String,
+}
+
+impl RolloutTraceRecorder {
+    /// Builds a recorder handle that accepts trace calls and records nothing.
+    pub fn disabled() -> Self {
+        Self {
+            state: RolloutTraceRecorderState::Disabled,
+        }
+    }
+
+    /// Creates and starts a root trace bundle, or returns a disabled recorder.
+    ///
+    /// 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. The returned recorder has not
+    /// emitted `ThreadStarted`; session setup records that event uniformly for
+    /// root and inherited child recorders.
+    pub fn create_root_or_disabled(thread_id: ThreadId) -> Self {
+        let Some(root) = std::env::var_os(CODEX_ROLLOUT_TRACE_ROOT_ENV) else {
+            return Self::disabled();
+        };
+        let root = PathBuf::from(root);
+        match Self::create_in_root(root.as_path(), thread_id) {
+            Ok(recorder) => recorder,
+            Err(err) => {
+                warn!("failed to initialize rollout trace recorder: {err:#}");
+                Self::disabled()
+            }
+        }
+    }
+
+    fn create_in_root(root: &Path, thread_id: ThreadId) -> 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 = EnabledRolloutTraceRecorder {
+            writer: Arc::new(writer),
+        };
+
+        recorder.append_best_effort(RawTraceEventPayload::RolloutStarted {
+            trace_id,
+            root_thread_id: thread_id,
+        });
+
+        debug!("recording rollout trace at {}", bundle_dir.display());
+        Ok(Self::enabled(recorder))
+    }
+
+    fn enabled(inner: EnabledRolloutTraceRecorder) -> Self {
+        Self {
+            state: RolloutTraceRecorderState::Enabled(inner),
+        }
+    }
+
+    /// 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 fn record_thread_started(&self, metadata: ThreadStartedTraceMetadata) {
+        let RolloutTraceRecorderState::Enabled(recorder) = &self.state else {
+            return;
+        };
+        let metadata_payload =
+            recorder.write_json_payload_best_effort(RawPayloadKind::SessionMetadata, &metadata);
+        recorder.append_best_effort(RawTraceEventPayload::ThreadStarted {
+            thread_id: metadata.thread_id,
+            agent_path: metadata.agent_path,
+            metadata_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 fn inference_trace_context(
+        &self,
+        thread_id: impl Into<AgentThreadId>,
+        codex_turn_id: impl Into<CodexTurnId>,
+        model: impl Into<String>,
+        provider_name: impl Into<String>,
+    ) -> InferenceTraceContext {
+        let RolloutTraceRecorderState::Enabled(recorder) = &self.state else {
+            return InferenceTraceContext::disabled();
+        };
+
+        InferenceTraceContext::enabled(
+            Arc::clone(&recorder.writer),
+            thread_id.into(),
+            codex_turn_id.into(),
+            model.into(),
+            provider_name.into(),
+        )
+    }
+
+    /// 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 fn compaction_trace_context(
+        &self,
+        thread_id: impl Into<AgentThreadId>,
+        codex_turn_id: impl Into<CodexTurnId>,
+        compaction_id: impl Into<CompactionId>,
+        model: impl Into<String>,
+        provider_name: impl Into<String>,
+    ) -> CompactionTraceContext {
+        let RolloutTraceRecorderState::Enabled(recorder) = &self.state else {
+            return CompactionTraceContext::disabled();
+        };
+
+        CompactionTraceContext::enabled(
+            Arc::clone(&recorder.writer),
+            thread_id.into(),
+            codex_turn_id.into(),
+            compaction_id.into(),
+            model.into(),
+            provider_name.into(),
+        )
+    }
+}
+
+impl EnabledRolloutTraceRecorder {
+    fn write_json_payload_best_effort(
+        &self,
+        kind: RawPayloadKind,
+        payload: &impl Serialize,
+    ) -> Option<RawPayloadRef> {
+        match self.writer.write_json_payload(kind, payload) {
+            Ok(payload_ref) => Some(payload_ref),
+            Err(err) => {
+                warn!("failed to write rollout trace payload: {err:#}");
+                None
+            }
+        }
+    }
+
+    fn append_best_effort(&self, payload: RawTraceEventPayload) {
+        if let Err(err) = self.writer.append(payload) {
+            warn!("failed to append rollout trace event: {err:#}");
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "recorder_tests.rs"]
+mod tests;
@@ -0,0 +1,162 @@
+use std::fs;
+use std::path::Path;
+use std::path::PathBuf;
+
+use codex_protocol::AgentPath;
+use codex_protocol::ThreadId;
+use codex_protocol::protocol::SandboxPolicy;
+use codex_protocol::protocol::SessionSource;
+use codex_protocol::protocol::SubAgentSource;
+use tempfile::TempDir;
+
+use super::*;
+use crate::CompactionCheckpointTracePayload;
+use crate::RolloutStatus;
+use crate::replay_bundle;
+
+#[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).expect("trace recorder");
+    recorder.record_thread_started(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),
+    });
+
+    let bundle_dir = single_bundle_dir(temp.path())?;
+    let replayed = replay_bundle(&bundle_dir)?;
+
+    assert_eq!(replayed.status, RolloutStatus::Running);
+    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).expect("trace recorder");
+    recorder.record_thread_started(minimal_metadata(root_thread_id));
+
+    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),
+    });
+    let bundle_dir = single_bundle_dir(temp.path())?;
+    let replayed = 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,
+        crate::ExecutionStatus::Running
+    );
+    assert_eq!(replayed.raw_payloads.len(), 2);
+
+    Ok(())
+}
+
+#[test]
+fn disabled_recorder_accepts_trace_calls_without_writing() -> anyhow::Result<()> {
+    let temp = TempDir::new()?;
+    let thread_id = ThreadId::new();
+    let recorder = RolloutTraceRecorder::disabled();
+
+    recorder.record_thread_started(minimal_metadata(thread_id));
+
+    let inference_trace =
+        recorder.inference_trace_context(thread_id, "turn-1", "gpt-test", "test-provider");
+    let inference_attempt = inference_trace.start_attempt();
+    inference_attempt.record_started(&serde_json::json!({ "kind": "inference" }));
+    let token_usage: Option<codex_protocol::protocol::TokenUsage> = None;
+    inference_attempt.record_completed("response-1", &token_usage, &[]);
+    inference_attempt.record_failed("inference failed");
+
+    let compaction_trace = recorder.compaction_trace_context(
+        thread_id,
+        "turn-1",
+        "compaction-1",
+        "gpt-test",
+        "test-provider",
+    );
+    let compaction_attempt =
+        compaction_trace.start_attempt(&serde_json::json!({ "kind": "compaction" }));
+    compaction_attempt.record_completed(&[]);
+    compaction_attempt.record_failed("compaction failed");
+    compaction_trace.record_installed(&CompactionCheckpointTracePayload {
+        input_history: &[],
+        replacement_history: &[],
+    });
+
+    assert_eq!(fs::read_dir(temp.path())?.count(), 0);
+
+    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))
+}