diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index a3f65b2355..494927277c 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -299,6 +299,8 @@ use crate::rollout::RolloutRecorderParams; use crate::rollout::map_session_init_error; use crate::rollout::metadata; use crate::rollout::policy::EventPersistenceMode; +use crate::rollout_trace::RolloutTraceRecorder; +use crate::rollout_trace::ThreadStartedTraceMetadata; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; use crate::shell; use crate::shell_snapshot::ShellSnapshot; @@ -440,6 +442,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) metrics_service_name: Option, pub(crate) inherited_shell_snapshot: Option>, pub(crate) inherited_exec_policy: Option>, + pub(crate) inherited_rollout_trace: Option, pub(crate) user_shell_override: Option, pub(crate) parent_trace: Option, pub(crate) analytics_events_client: Option, @@ -494,6 +497,7 @@ impl Codex { inherited_shell_snapshot, user_shell_override, inherited_exec_policy, + inherited_rollout_trace, parent_trace: _, analytics_events_client, } = args; @@ -690,6 +694,7 @@ impl Codex { agent_control, environment, analytics_events_client, + inherited_rollout_trace, ) .await .map_err(|e| { @@ -1697,6 +1702,7 @@ impl Session { agent_control: AgentControl, environment: Option>, analytics_events_client: Option, + inherited_rollout_trace: Option, ) -> anyhow::Result> { debug!( "Configuring session: model={}; provider={:?}", @@ -1833,6 +1839,40 @@ impl Session { let rollout_path = rollout_recorder .as_ref() .map(|rec| rec.rollout_path().to_path_buf()); + let trace_agent_path = session_configuration + .session_source + .get_agent_path() + .unwrap_or_else(codex_protocol::AgentPath::root); + let trace_task_name = + (!trace_agent_path.is_root()).then(|| trace_agent_path.name().to_string()); + let trace_metadata = ThreadStartedTraceMetadata { + thread_id: conversation_id.to_string(), + agent_path: trace_agent_path.to_string(), + task_name: trace_task_name, + nickname: session_configuration.session_source.get_nickname(), + agent_role: session_configuration.session_source.get_agent_role(), + session_source: session_configuration.session_source.clone(), + cwd: session_configuration.cwd.to_path_buf(), + rollout_path: rollout_path.clone(), + model: session_configuration.collaboration_mode.model().to_string(), + provider_name: config.model_provider_id.clone(), + approval_policy: session_configuration.approval_policy.value().to_string(), + sandbox_policy: format!("{:?}", session_configuration.sandbox_policy.get()), + }; + let rollout_trace = if let Some(rollout_trace) = inherited_rollout_trace { + rollout_trace.record_thread_started(trace_metadata); + Some(rollout_trace) + } else if matches!( + session_configuration.session_source, + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) + ) { + // Spawned child threads are part of their root rollout tree. If + // the parent had no trace recorder, do not create an orphan child + // bundle that looks like an independent rollout. + None + } else { + RolloutTraceRecorder::maybe_create(conversation_id, trace_metadata) + }; let mut post_session_configured_events = Vec::::new(); @@ -2107,6 +2147,7 @@ impl Session { analytics_events_client, hooks, rollout: Mutex::new(rollout_recorder), + rollout_trace, user_shell: Arc::new(default_shell), agent_identity_manager: Arc::new(AgentIdentityManager::new( config.as_ref(), @@ -2873,6 +2914,18 @@ impl Session { /// Persist the event to rollout and send it to clients. pub(crate) async fn send_event(&self, turn_context: &TurnContext, msg: EventMsg) { let legacy_source = msg.clone(); + if let Some(rollout_trace) = &self.services.rollout_trace { + rollout_trace.record_codex_turn_event( + self.conversation_id.to_string(), + &turn_context.sub_id, + &legacy_source, + ); + rollout_trace.record_tool_call_event( + self.conversation_id.to_string(), + turn_context.sub_id.clone(), + &legacy_source, + ); + } let event = Event { id: turn_context.sub_id.clone(), msg, @@ -2925,13 +2978,19 @@ impl Session { return; } - self.forward_child_completion_to_parent(*parent_thread_id, child_agent_path, status) - .await; + self.forward_child_completion_to_parent( + turn_context, + *parent_thread_id, + child_agent_path, + status, + ) + .await; } /// Sends the standard completion envelope from a spawned MultiAgentV2 child to its parent. async fn forward_child_completion_to_parent( &self, + turn_context: &TurnContext, parent_thread_id: ThreadId, child_agent_path: &codex_protocol::AgentPath, status: AgentStatus, @@ -2949,9 +3008,19 @@ impl Session { child_agent_path.clone(), parent_agent_path, Vec::new(), - message, + message.clone(), /*trigger_turn*/ false, ); + if let Some(rollout_trace) = &self.services.rollout_trace { + rollout_trace.record_agent_result_interaction( + self.conversation_id.to_string(), + turn_context.sub_id.clone(), + parent_thread_id.to_string(), + child_agent_path.as_str(), + &message, + &status, + ); + } if let Err(err) = self .services .agent_control @@ -2990,6 +3059,9 @@ impl Session { // Persist the event into rollout (recorder filters as needed) let rollout_items = vec![RolloutItem::EventMsg(event.msg.clone())]; self.persist_rollout_items(&rollout_items).await; + if let Some(rollout_trace) = &self.services.rollout_trace { + rollout_trace.record_protocol_event(&event.msg); + } self.deliver_event_raw(event).await; } @@ -5887,6 +5959,12 @@ mod handlers { msg: EventMsg::ShutdownComplete, }; sess.send_event_raw(event).await; + if let Some(rollout_trace) = &sess.services.rollout_trace { + rollout_trace.record_thread_ended( + sess.conversation_id.to_string(), + codex_rollout_trace::RolloutStatus::Completed, + ); + } true } @@ -7843,8 +7921,19 @@ async fn try_run_sampling_request( auth_mode = sess.services.auth_manager.auth_mode(), features = sess.features.enabled_features(), ); + let inference_trace = sess.services.rollout_trace.as_ref().map_or_else( + codex_rollout_trace::InferenceTraceContext::disabled, + |trace| { + trace.inference_trace_context( + sess.conversation_id.to_string(), + turn_context.sub_id.clone(), + turn_context.model_info.slug.clone(), + turn_context.provider.name.clone(), + ) + }, + ); let mut stream = client_session - .stream( + .stream_with_trace( prompt, &turn_context.model_info, &turn_context.session_telemetry, @@ -7852,6 +7941,7 @@ async fn try_run_sampling_request( turn_context.reasoning_summary, turn_context.config.service_tier, turn_metadata_header, + &inference_trace, ) .instrument(trace_span!("stream_request")) .or_cancel(&cancellation_token) diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 9cd58e044f..1afb0f09e9 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -95,6 +95,7 @@ pub(crate) async fn run_codex_thread_interactive( inherited_shell_snapshot: None, user_shell_override: None, inherited_exec_policy: Some(Arc::clone(&parent_session.services.exec_policy)), + inherited_rollout_trace: None, parent_trace: None, analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), })) diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 5db38f7b72..31c040c148 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -11,6 +11,7 @@ use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejection; use crate::mcp::McpManager; use crate::plugins::PluginsManager; +use crate::rollout_trace::RolloutTraceRecorder; use crate::skills_watcher::SkillsWatcher; use crate::tools::code_mode::CodeModeService; use crate::tools::network_approval::NetworkApprovalService; @@ -42,6 +43,7 @@ pub(crate) struct SessionServices { pub(crate) analytics_events_client: AnalyticsEventsClient, pub(crate) hooks: Hooks, pub(crate) rollout: Mutex>, + pub(crate) rollout_trace: Option, pub(crate) user_shell: Arc, pub(crate) agent_identity_manager: Arc, pub(crate) shell_snapshot_tx: watch::Sender>>, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index b8ad75b49e..fc3fe9fb29 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -41,6 +41,7 @@ use codex_protocol::protocol::Op; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionConfiguredEvent; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::W3cTraceContext; @@ -925,6 +926,9 @@ impl ThreadManagerState { } Some(_) | None => crate::file_watcher::WatchRegistration::default(), }; + let inherited_rollout_trace = self + .inherited_rollout_trace_for_source(&session_source) + .await; let CodexSpawnOk { codex, thread_id, .. } = Codex::spawn(CodexSpawnArgs { @@ -944,6 +948,7 @@ impl ThreadManagerState { metrics_service_name, inherited_shell_snapshot, inherited_exec_policy, + inherited_rollout_trace, user_shell_override, parent_trace, analytics_events_client: self.analytics_events_client.clone(), @@ -988,6 +993,24 @@ impl ThreadManagerState { pub(crate) fn notify_thread_created(&self, thread_id: ThreadId) { let _ = self.thread_created_tx.send(thread_id); } + + async fn inherited_rollout_trace_for_source( + &self, + session_source: &SessionSource, + ) -> Option { + // Only v2 thread-spawn children inherit a recorder. Independent + // top-level threads still create their own rollout bundles. + let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + }) = session_source + else { + return None; + }; + self.get_thread(*parent_thread_id) + .await + .ok() + .and_then(|thread| thread.codex.session.services.rollout_trace.clone()) + } } /// Return a fork snapshot cut strictly before the nth user message (0-based).