codex_rollout_trace/
thread.rs

1//! Thread-scoped rollout trace helpers.
2//!
3//! A rollout bundle can contain a root thread plus spawned child threads. This
4//! context owns the stable identity for one thread inside that bundle. Keeping
5//! thread-local event methods here avoids repeatedly plumbing `thread_id`
6//! through session code.
7
8use codex_protocol::protocol::AgentStatus;
9use codex_protocol::protocol::EventMsg;
10use codex_protocol::protocol::SessionSource;
11use serde::Serialize;
12use std::path::Path;
13use std::path::PathBuf;
14use std::sync::Arc;
15use tracing::debug;
16use tracing::warn;
17use uuid::Uuid;
18
19use crate::AgentThreadId;
20use crate::CodeCellTraceContext;
21use crate::CodexTurnId;
22use crate::CompactionId;
23use crate::CompactionTraceContext;
24use crate::InferenceTraceContext;
25use crate::McpCallTraceContext;
26use crate::RawPayloadKind;
27use crate::RawPayloadRef;
28use crate::RawTraceEventContext;
29use crate::RawTraceEventPayload;
30use crate::RolloutStatus;
31use crate::ToolCallId;
32use crate::ToolDispatchInvocation;
33use crate::ToolDispatchTraceContext;
34use crate::TraceWriter;
35use crate::protocol_event::codex_turn_trace_event;
36use crate::protocol_event::tool_runtime_trace_event;
37use crate::protocol_event::wrapped_protocol_event_type;
38
39/// Environment variable that enables local trace-bundle recording.
40///
41/// The value is a root directory. Each independent root session gets one child
42/// bundle directory. Spawned child threads share their root session's bundle so
43/// one reduced `state.json` describes the whole multi-agent rollout tree.
44pub const CODEX_ROLLOUT_TRACE_ROOT_ENV: &str = "CODEX_ROLLOUT_TRACE_ROOT";
45
46/// Metadata captured once at thread/session start.
47///
48/// This payload is intentionally operational rather than reduced: it is a raw
49/// payload that later reducers can mine as the reduced thread model evolves.
50#[derive(Serialize)]
51pub struct ThreadStartedTraceMetadata {
52    pub thread_id: String,
53    pub agent_path: String,
54    pub task_name: Option<String>,
55    pub nickname: Option<String>,
56    pub agent_role: Option<String>,
57    pub session_source: SessionSource,
58    pub cwd: std::path::PathBuf,
59    pub rollout_path: Option<std::path::PathBuf>,
60    pub model: String,
61    pub provider_name: String,
62    pub approval_policy: String,
63    pub sandbox_policy: String,
64}
65
66/// Trace-only payload for a child completion notification delivered to its parent.
67#[derive(Serialize)]
68pub struct AgentResultTracePayload<'a> {
69    pub child_agent_path: &'a str,
70    pub message: &'a str,
71    pub status: &'a AgentStatus,
72}
73
74/// No-op capable trace handle for one thread in a rollout bundle.
75#[derive(Clone, Debug)]
76pub struct ThreadTraceContext {
77    state: ThreadTraceContextState,
78}
79
80#[derive(Clone, Debug)]
81enum ThreadTraceContextState {
82    Disabled,
83    Enabled(EnabledThreadTraceContext),
84}
85
86#[derive(Clone, Debug)]
87struct EnabledThreadTraceContext {
88    writer: Arc<TraceWriter>,
89    root_thread_id: AgentThreadId,
90    thread_id: AgentThreadId,
91}
92
93impl ThreadTraceContext {
94    /// Builds a context that accepts trace calls and records nothing.
95    pub fn disabled() -> Self {
96        Self {
97            state: ThreadTraceContextState::Disabled,
98        }
99    }
100
101    /// Starts a root thread trace from `CODEX_ROLLOUT_TRACE_ROOT`, or disables tracing.
102    ///
103    /// Trace startup is best-effort. A tracing failure must not make the Codex
104    /// session unusable, because traces are diagnostic and can be enabled while
105    /// debugging unrelated production failures.
106    pub fn start_root_or_disabled(metadata: ThreadStartedTraceMetadata) -> Self {
107        let Some(root) = std::env::var_os(CODEX_ROLLOUT_TRACE_ROOT_ENV) else {
108            return Self::disabled();
109        };
110        let root = PathBuf::from(root);
111        match start_root_in_root(root.as_path(), metadata) {
112            Ok(context) => context,
113            Err(err) => {
114                warn!("failed to initialize rollout trace bundle: {err:#}");
115                Self::disabled()
116            }
117        }
118    }
119
120    /// Starts a root trace in a known directory.
121    ///
122    /// This is public for tests that need replayable trace bundles without
123    /// mutating process environment.
124    pub fn start_root_in_root_for_test(
125        root: &Path,
126        metadata: ThreadStartedTraceMetadata,
127    ) -> anyhow::Result<Self> {
128        start_root_in_root(root, metadata)
129    }
130
131    /// Starts one thread lifecycle inside an existing rollout bundle.
132    pub(crate) fn start(
133        writer: Arc<TraceWriter>,
134        root_thread_id: AgentThreadId,
135        metadata: ThreadStartedTraceMetadata,
136    ) -> Self {
137        let context = EnabledThreadTraceContext {
138            writer,
139            root_thread_id,
140            thread_id: metadata.thread_id.clone(),
141        };
142        record_thread_started(&context, metadata);
143        Self {
144            state: ThreadTraceContextState::Enabled(context),
145        }
146    }
147
148    /// Returns whether this handle will write trace events.
149    ///
150    /// Most methods have their own disabled fast path. Callers should branch on
151    /// this only when preparing trace payloads would otherwise clone data the
152    /// production path needs to move elsewhere.
153    pub fn is_enabled(&self) -> bool {
154        matches!(self.state, ThreadTraceContextState::Enabled(_))
155    }
156
157    /// Starts a fresh child thread in this context's rollout tree.
158    ///
159    /// Callers should use [`ThreadTraceContext::disabled`] for resumed children:
160    /// reusing the parent trace would emit a duplicate `ThreadStarted` event
161    /// for an existing thread id and make the bundle unreplayable.
162    pub fn start_child_thread_trace_or_disabled(
163        &self,
164        metadata: ThreadStartedTraceMetadata,
165    ) -> Self {
166        match &self.state {
167            ThreadTraceContextState::Disabled => Self::disabled(),
168            ThreadTraceContextState::Enabled(context) => Self::start(
169                Arc::clone(&context.writer),
170                context.root_thread_id.clone(),
171                metadata,
172            ),
173        }
174    }
175
176    /// Emits terminal trace events for graceful thread shutdown.
177    ///
178    /// Spawned child sessions share their root bundle, so only the root
179    /// thread end closes the rollout. Child thread ends update the child thread
180    /// execution state without marking the whole bundle complete.
181    pub fn record_ended(&self, status: RolloutStatus) {
182        let ThreadTraceContextState::Enabled(context) = &self.state else {
183            return;
184        };
185        context.append_best_effort(RawTraceEventPayload::ThreadEnded {
186            thread_id: context.thread_id.clone(),
187            status: status.clone(),
188        });
189        if context.thread_id == context.root_thread_id {
190            context.append_best_effort(RawTraceEventPayload::RolloutEnded { status });
191        }
192    }
193
194    /// Wraps selected protocol events as raw trace breadcrumbs.
195    ///
196    /// High-volume stream deltas stay out of this wrapper; typed inference,
197    /// tool, terminal, and code-mode hooks provide the canonical runtime data.
198    pub fn record_protocol_event(&self, event: &EventMsg) {
199        let ThreadTraceContextState::Enabled(context) = &self.state else {
200            return;
201        };
202        let Some(event_type) = wrapped_protocol_event_type(event) else {
203            return;
204        };
205        let Some(event_payload) =
206            context.write_json_payload_best_effort(RawPayloadKind::ProtocolEvent, event)
207        else {
208            return;
209        };
210        context.append_best_effort(RawTraceEventPayload::ProtocolEventObserved {
211            event_type: event_type.to_string(),
212            event_payload,
213        });
214    }
215
216    /// Emits typed Codex turn lifecycle events from protocol lifecycle events.
217    pub fn record_codex_turn_event(&self, default_turn_id: &str, event: &EventMsg) {
218        let ThreadTraceContextState::Enabled(context) = &self.state else {
219            return;
220        };
221        let Some(trace_event) =
222            codex_turn_trace_event(context.thread_id.clone(), default_turn_id, event)
223        else {
224            return;
225        };
226        context.append_with_context_best_effort(
227            trace_event.context_turn_id.clone(),
228            trace_event.payload,
229        );
230    }
231
232    /// Emits typed runtime tool events from existing protocol lifecycle events.
233    ///
234    /// These events are runtime observations on an already-dispatched tool. The
235    /// dispatch trace records the caller-facing boundary; these payloads explain
236    /// what Codex did while executing that boundary.
237    pub fn record_tool_call_event(&self, codex_turn_id: impl Into<CodexTurnId>, event: &EventMsg) {
238        let ThreadTraceContextState::Enabled(context) = &self.state else {
239            return;
240        };
241        let Some(trace_event) = tool_runtime_trace_event(event) else {
242            return;
243        };
244        let Some(payload) = context.raw_tool_runtime_payload(trace_event) else {
245            return;
246        };
247        context.append_with_context_best_effort(codex_turn_id.into(), payload);
248    }
249
250    /// Emits the v2 child-to-parent completion message as an explicit graph edge.
251    ///
252    /// The notification is runtime delivery from a completed child turn into
253    /// the parent's mailbox, not a tool call executed by the child. Recording it
254    /// directly preserves timing and source without making the reducer infer
255    /// the edge from a later parent prompt snapshot.
256    pub fn record_agent_result_interaction(
257        &self,
258        child_codex_turn_id: impl Into<CodexTurnId>,
259        parent_thread_id: impl Into<AgentThreadId>,
260        payload: &AgentResultTracePayload<'_>,
261    ) {
262        let ThreadTraceContextState::Enabled(context) = &self.state else {
263            return;
264        };
265        let child_codex_turn_id = child_codex_turn_id.into();
266        let parent_thread_id = parent_thread_id.into();
267        let carried_payload =
268            context.write_json_payload_best_effort(RawPayloadKind::AgentResult, payload);
269        context.append_with_context_best_effort(
270            child_codex_turn_id.clone(),
271            RawTraceEventPayload::AgentResultObserved {
272                edge_id: format!(
273                    "edge:agent_result:{}:{child_codex_turn_id}:{parent_thread_id}",
274                    context.thread_id
275                ),
276                child_thread_id: context.thread_id.clone(),
277                child_codex_turn_id,
278                parent_thread_id,
279                message: payload.message.to_string(),
280                carried_payload,
281            },
282        );
283    }
284
285    /// Emits a turn-start lifecycle event.
286    ///
287    /// Most production turn lifecycle wiring lives outside this PR layer, but
288    /// trace-focused integration tests need a small explicit hook so reducer
289    /// inputs remain valid without exercising the full session loop.
290    pub fn record_codex_turn_started(&self, codex_turn_id: impl Into<CodexTurnId>) {
291        let ThreadTraceContextState::Enabled(context) = &self.state else {
292            return;
293        };
294        let codex_turn_id = codex_turn_id.into();
295        context.append_with_context_best_effort(
296            codex_turn_id.clone(),
297            RawTraceEventPayload::CodexTurnStarted {
298                codex_turn_id,
299                thread_id: context.thread_id.clone(),
300            },
301        );
302    }
303
304    /// Starts a first-class code-mode cell lifecycle and returns its trace handle.
305    pub fn start_code_cell_trace(
306        &self,
307        codex_turn_id: impl Into<CodexTurnId>,
308        runtime_cell_id: impl Into<String>,
309        model_visible_call_id: impl Into<String>,
310        source_js: impl Into<String>,
311    ) -> CodeCellTraceContext {
312        let context = self.code_cell_trace_context(codex_turn_id, runtime_cell_id);
313        context.record_started(model_visible_call_id, source_js);
314        context
315    }
316
317    /// Builds a trace handle for an already-started code-mode runtime cell.
318    pub fn code_cell_trace_context(
319        &self,
320        codex_turn_id: impl Into<CodexTurnId>,
321        runtime_cell_id: impl Into<String>,
322    ) -> CodeCellTraceContext {
323        let ThreadTraceContextState::Enabled(context) = &self.state else {
324            return CodeCellTraceContext::disabled();
325        };
326        CodeCellTraceContext::enabled(
327            Arc::clone(&context.writer),
328            context.thread_id.clone(),
329            codex_turn_id,
330            runtime_cell_id,
331        )
332    }
333
334    /// Starts one dispatch-level tool lifecycle and returns its trace handle.
335    ///
336    /// `invocation` is lazy because adapting core tool objects into trace-owned
337    /// payloads can clone large arguments. Disabled tracing should not pay that
338    /// cost on the hot tool-dispatch path.
339    pub fn start_tool_dispatch_trace(
340        &self,
341        invocation: impl FnOnce() -> Option<ToolDispatchInvocation>,
342    ) -> ToolDispatchTraceContext {
343        let ThreadTraceContextState::Enabled(context) = &self.state else {
344            return ToolDispatchTraceContext::disabled();
345        };
346        let Some(invocation) = invocation() else {
347            return ToolDispatchTraceContext::disabled();
348        };
349        ToolDispatchTraceContext::start(Arc::clone(&context.writer), invocation)
350    }
351
352    /// Builds reusable inference trace context for one Codex turn.
353    ///
354    /// The returned context is intentionally not "an inference call" yet.
355    /// Transport code owns retry/fallback attempts and calls `start_attempt`
356    /// only after it has built the concrete request payload for that attempt.
357    pub fn inference_trace_context(
358        &self,
359        codex_turn_id: impl Into<CodexTurnId>,
360        model: impl Into<String>,
361        provider_name: impl Into<String>,
362    ) -> InferenceTraceContext {
363        let ThreadTraceContextState::Enabled(context) = &self.state else {
364            return InferenceTraceContext::disabled();
365        };
366        InferenceTraceContext::enabled(
367            Arc::clone(&context.writer),
368            context.thread_id.clone(),
369            codex_turn_id.into(),
370            model.into(),
371            provider_name.into(),
372        )
373    }
374
375    /// Builds remote-compaction trace context for one checkpoint.
376    ///
377    /// Rollout tracing currently has a first-class checkpoint model only for remote compaction.
378    /// The compact endpoint is a model-facing request whose output replaces live history, so it
379    /// needs both request/response attempt events and a later checkpoint event when processed
380    /// replacement history is installed.
381    pub fn compaction_trace_context(
382        &self,
383        codex_turn_id: impl Into<CodexTurnId>,
384        compaction_id: impl Into<CompactionId>,
385        model: impl Into<String>,
386        provider_name: impl Into<String>,
387    ) -> CompactionTraceContext {
388        let ThreadTraceContextState::Enabled(context) = &self.state else {
389            return CompactionTraceContext::disabled();
390        };
391        CompactionTraceContext::enabled(
392            Arc::clone(&context.writer),
393            context.thread_id.clone(),
394            codex_turn_id.into(),
395            compaction_id.into(),
396            model.into(),
397            provider_name.into(),
398        )
399    }
400
401    /// Starts bridge correlation for one concrete MCP backend request.
402    ///
403    /// Dispatch-level tool IDs remain compact and UI-friendly. This UUID is
404    /// deliberately separate: it is only for cross-process log joins where a
405    /// rollout-local counter would collide across samples.
406    pub fn start_mcp_call_trace(&self, tool_call_id: impl Into<ToolCallId>) -> McpCallTraceContext {
407        let ThreadTraceContextState::Enabled(context) = &self.state else {
408            return McpCallTraceContext::disabled();
409        };
410        let mcp_call_id = Uuid::new_v4().to_string();
411        let trace = McpCallTraceContext::enabled(mcp_call_id.clone());
412        context.append_best_effort(RawTraceEventPayload::McpToolCallCorrelationAssigned {
413            tool_call_id: tool_call_id.into(),
414            mcp_call_id,
415        });
416        trace
417    }
418}
419
420fn start_root_in_root(
421    root: &Path,
422    metadata: ThreadStartedTraceMetadata,
423) -> anyhow::Result<ThreadTraceContext> {
424    let trace_id = Uuid::new_v4().to_string();
425    let thread_id = metadata.thread_id.clone();
426    let bundle_dir = root.join(format!("trace-{trace_id}-{thread_id}"));
427    let writer = TraceWriter::create(
428        &bundle_dir,
429        trace_id.clone(),
430        thread_id.clone(),
431        thread_id.clone(),
432    )?;
433    let writer = Arc::new(writer);
434
435    if let Err(err) = writer.append(RawTraceEventPayload::RolloutStarted {
436        trace_id,
437        root_thread_id: thread_id.clone(),
438    }) {
439        warn!("failed to append rollout trace event: {err:#}");
440    }
441
442    debug!("recording rollout trace at {}", bundle_dir.display());
443    Ok(ThreadTraceContext::start(writer, thread_id, metadata))
444}
445
446fn record_thread_started(
447    context: &EnabledThreadTraceContext,
448    metadata: ThreadStartedTraceMetadata,
449) {
450    let metadata_payload =
451        context.write_json_payload_best_effort(RawPayloadKind::SessionMetadata, &metadata);
452    context.append_best_effort(RawTraceEventPayload::ThreadStarted {
453        thread_id: metadata.thread_id,
454        agent_path: metadata.agent_path,
455        metadata_payload,
456    });
457}
458
459impl EnabledThreadTraceContext {
460    fn write_json_payload_best_effort(
461        &self,
462        kind: RawPayloadKind,
463        payload: &impl Serialize,
464    ) -> Option<RawPayloadRef> {
465        match self.writer.write_json_payload(kind, payload) {
466            Ok(payload_ref) => Some(payload_ref),
467            Err(err) => {
468                warn!("failed to write rollout trace payload: {err:#}");
469                None
470            }
471        }
472    }
473
474    fn raw_tool_runtime_payload(
475        &self,
476        trace_event: crate::protocol_event::ToolRuntimeTraceEvent<'_>,
477    ) -> Option<RawTraceEventPayload> {
478        match trace_event {
479            crate::protocol_event::ToolRuntimeTraceEvent::Started {
480                tool_call_id,
481                payload,
482            } => {
483                let runtime_payload = self
484                    .write_json_payload_best_effort(RawPayloadKind::ToolRuntimeEvent, &payload)?;
485                Some(RawTraceEventPayload::ToolCallRuntimeStarted {
486                    tool_call_id: tool_call_id.to_string(),
487                    runtime_payload,
488                })
489            }
490            crate::protocol_event::ToolRuntimeTraceEvent::Ended {
491                tool_call_id,
492                status,
493                payload,
494            } => {
495                let runtime_payload = self
496                    .write_json_payload_best_effort(RawPayloadKind::ToolRuntimeEvent, &payload)?;
497                Some(RawTraceEventPayload::ToolCallRuntimeEnded {
498                    tool_call_id: tool_call_id.to_string(),
499                    status,
500                    runtime_payload,
501                })
502            }
503        }
504    }
505
506    fn append_best_effort(&self, payload: RawTraceEventPayload) {
507        if let Err(err) = self.writer.append(payload) {
508            warn!("failed to append rollout trace event: {err:#}");
509        }
510    }
511
512    fn append_with_context_best_effort(
513        &self,
514        codex_turn_id: CodexTurnId,
515        payload: RawTraceEventPayload,
516    ) {
517        let event_context = RawTraceEventContext {
518            thread_id: Some(self.thread_id.clone()),
519            codex_turn_id: Some(codex_turn_id),
520        };
521        if let Err(err) = self.writer.append_with_context(event_context, payload) {
522            warn!("failed to append rollout trace event: {err:#}");
523        }
524    }
525}
526
527#[cfg(test)]
528#[path = "thread_tests.rs"]
529mod tests;
codex_rollout_trace/thread.rs

codex_rollout_trace/
thread.rs
