From 52521a5e408c86b22dfe37edab7cec2376854961 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Tue, 3 Mar 2026 17:03:45 -0800 Subject: [PATCH] feat(app-server): propagate app-server trace context into core (#13368) ### Summary Propagate trace context originating at app-server RPC method handlers -> codex core submission loop (so this includes spans such as `run_turn`!). This implements PR 2 of the app-server tracing rollout. This also removes the old lower-level env-based reparenting in core so explicit request/submission ancestry wins instead of being overridden by ambient `TRACEPARENT` state. ### What changed - Added `trace: Option` to codex_protocol::Submission - Taught `Codex::submit()` / `submit_with_id()` to automatically capture the current span context when constructing or forwarding a submission - Wrapped the core submission loop in a submission_dispatch span parented from Submission.trace - Warn on invalid submission trace carriers and ignore them cleanly - Removed the old env-based downstream reparenting path in core task execution - Stopped OTEL provider init from implicitly attaching env trace context process-wide - Updated mcp-server Submission call sites for the new field Added focused unit tests for: - capturing trace context into Submission - preferring `Submission.trace` when building the core dispatch span ### Why PR 1 gave us consistent inbound request spans in app-server, but that only covered the transport boundary. For long-running work like turns and reviews, the important missing piece was preserving ancestry after the request handler returns and core continues work on a different async path. This change makes that handoff explicit and keeps the parentage rules simple: - app-server request span sets the current context - `Submission.trace` snapshots that context - core restores it once, at the submission boundary - deeper core spans inherit naturally That also lets us stop relying on env-based reparenting for this path, which was too ambient and could override explicit ancestry. --- codex-rs/Cargo.lock | 6 + codex-rs/core/Cargo.toml | 2 + codex-rs/core/src/codex.rs | 475 ++++++++++++------- codex-rs/core/src/codex_delegate.rs | 50 +- codex-rs/core/src/tasks/regular.rs | 3 - codex-rs/exec/Cargo.toml | 4 + codex-rs/exec/src/lib.rs | 110 ++++- codex-rs/mcp-server/src/codex_tool_runner.rs | 1 + codex-rs/mcp-server/src/message_processor.rs | 1 + codex-rs/otel/src/otel_provider.rs | 22 - codex-rs/otel/src/traces/otel_manager.rs | 8 - codex-rs/protocol/src/protocol.rs | 3 + 12 files changed, 489 insertions(+), 196 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 48e8c598d..1974883dd 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1851,6 +1851,7 @@ dependencies = [ "notify", "once_cell", "openssl-sys", + "opentelemetry", "opentelemetry_sdk", "os_info", "predicates", @@ -1880,6 +1881,7 @@ dependencies = [ "toml 0.9.11+spec-1.1.0", "toml_edit 0.24.0+spec-1.1.0", "tracing", + "tracing-opentelemetry", "tracing-subscriber", "tracing-test", "url", @@ -1916,6 +1918,7 @@ dependencies = [ "codex-arg0", "codex-cloud-requirements", "codex-core", + "codex-otel", "codex-protocol", "codex-utils-absolute-path", "codex-utils-cargo-bin", @@ -1925,6 +1928,8 @@ dependencies = [ "codex-utils-sandbox-summary", "core_test_support", "libc", + "opentelemetry", + "opentelemetry_sdk", "owo-colors", "predicates", "pretty_assertions", @@ -1936,6 +1941,7 @@ dependencies = [ "tempfile", "tokio", "tracing", + "tracing-opentelemetry", "tracing-subscriber", "ts-rs", "uuid", diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index eafedc9d7..ebb01edc4 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -161,6 +161,7 @@ core_test_support = { workspace = true } ctor = { workspace = true } insta = { workspace = true } maplit = { workspace = true } +opentelemetry = { workspace = true } predicates = { workspace = true } pretty_assertions = { workspace = true } test-case = "3.3.1" @@ -170,6 +171,7 @@ opentelemetry_sdk = { workspace = true, features = [ ] } serial_test = { workspace = true } tempfile = { workspace = true } +tracing-opentelemetry = { workspace = true } tracing-subscriber = { workspace = true } tracing-test = { workspace = true, features = ["no-env-filter"] } walkdir = { workspace = true } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 1ead1871f..87aac7318 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -70,6 +70,8 @@ use codex_hooks::HooksConfig; use codex_network_proxy::NetworkProxy; use codex_network_proxy::NetworkProxyAuditMetadata; use codex_network_proxy::normalize_host; +use codex_otel::current_span_w3c_trace_context; +use codex_otel::set_parent_from_w3c_trace_context; use codex_protocol::ThreadId; use codex_protocol::approvals::ExecPolicyAmendment; use codex_protocol::approvals::NetworkPolicyAmendment; @@ -539,14 +541,21 @@ impl Codex { /// Submit the `op` wrapped in a `Submission` with a unique ID. pub async fn submit(&self, op: Op) -> CodexResult { let id = Uuid::now_v7().to_string(); - let sub = Submission { id: id.clone(), op }; + let sub = Submission { + id: id.clone(), + op, + trace: None, + }; self.submit_with_id(sub).await?; Ok(id) } /// Use sparingly: prefer `submit()` so Codex is responsible for generating /// unique IDs for each submission. - pub async fn submit_with_id(&self, sub: Submission) -> CodexResult<()> { + pub async fn submit_with_id(&self, mut sub: Submission) -> CodexResult<()> { + if sub.trace.is_none() { + sub.trace = current_span_w3c_trace_context(); + } self.tx_sub .send(sub) .await @@ -3686,176 +3695,230 @@ async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiv // To break out of this loop, send Op::Shutdown. while let Ok(sub) = rx_sub.recv().await { debug!(?sub, "Submission"); - match sub.op.clone() { - Op::Interrupt => { - handlers::interrupt(&sess).await; - } - Op::CleanBackgroundTerminals => { - handlers::clean_background_terminals(&sess).await; - } - Op::RealtimeConversationStart(params) => { - if let Err(err) = - handle_realtime_conversation_start(&sess, sub.id.clone(), params).await - { - sess.send_event_raw(Event { - id: sub.id.clone(), - msg: EventMsg::Error(ErrorEvent { - message: err.to_string(), - codex_error_info: Some(CodexErrorInfo::Other), - }), - }) - .await; + let dispatch_span = submission_dispatch_span(&sub); + let should_exit = async { + match sub.op.clone() { + Op::Interrupt => { + handlers::interrupt(&sess).await; + false } - } - Op::RealtimeConversationAudio(params) => { - handle_realtime_conversation_audio(&sess, sub.id.clone(), params).await; - } - Op::RealtimeConversationText(params) => { - handle_realtime_conversation_text(&sess, sub.id.clone(), params).await; - } - Op::RealtimeConversationClose => { - handle_realtime_conversation_close(&sess, sub.id.clone()).await; - } - Op::OverrideTurnContext { - cwd, - approval_policy, - sandbox_policy, - windows_sandbox_level, - model, - effort, - summary, - service_tier, - collaboration_mode, - personality, - } => { - let collaboration_mode = if let Some(collab_mode) = collaboration_mode { - collab_mode - } else { - let state = sess.state.lock().await; - state.session_configuration.collaboration_mode.with_updates( - model.clone(), - effort, - None, + Op::CleanBackgroundTerminals => { + handlers::clean_background_terminals(&sess).await; + false + } + Op::RealtimeConversationStart(params) => { + if let Err(err) = + handle_realtime_conversation_start(&sess, sub.id.clone(), params).await + { + sess.send_event_raw(Event { + id: sub.id.clone(), + msg: EventMsg::Error(ErrorEvent { + message: err.to_string(), + codex_error_info: Some(CodexErrorInfo::Other), + }), + }) + .await; + } + false + } + Op::RealtimeConversationAudio(params) => { + handle_realtime_conversation_audio(&sess, sub.id.clone(), params).await; + false + } + Op::RealtimeConversationText(params) => { + handle_realtime_conversation_text(&sess, sub.id.clone(), params).await; + false + } + Op::RealtimeConversationClose => { + handle_realtime_conversation_close(&sess, sub.id.clone()).await; + false + } + Op::OverrideTurnContext { + cwd, + approval_policy, + sandbox_policy, + windows_sandbox_level, + model, + effort, + summary, + service_tier, + collaboration_mode, + personality, + } => { + let collaboration_mode = if let Some(collab_mode) = collaboration_mode { + collab_mode + } else { + let state = sess.state.lock().await; + state.session_configuration.collaboration_mode.with_updates( + model.clone(), + effort, + None, + ) + }; + handlers::override_turn_context( + &sess, + sub.id.clone(), + SessionSettingsUpdate { + cwd, + approval_policy, + sandbox_policy, + windows_sandbox_level, + collaboration_mode: Some(collaboration_mode), + reasoning_summary: summary, + service_tier, + personality, + ..Default::default() + }, ) - }; - handlers::override_turn_context( - &sess, - sub.id.clone(), - SessionSettingsUpdate { - cwd, - approval_policy, - sandbox_policy, - windows_sandbox_level, - collaboration_mode: Some(collaboration_mode), - reasoning_summary: summary, - service_tier, - personality, - ..Default::default() - }, - ) - .await; - } - Op::UserInput { .. } | Op::UserTurn { .. } => { - handlers::user_input_or_turn(&sess, sub.id.clone(), sub.op).await; - } - Op::ExecApproval { - id: approval_id, - turn_id, - decision, - } => { - handlers::exec_approval(&sess, approval_id, turn_id, decision).await; - } - Op::PatchApproval { id, decision } => { - handlers::patch_approval(&sess, id, decision).await; - } - Op::UserInputAnswer { id, response } => { - handlers::request_user_input_response(&sess, id, response).await; - } - Op::DynamicToolResponse { id, response } => { - handlers::dynamic_tool_response(&sess, id, response).await; - } - Op::AddToHistory { text } => { - handlers::add_to_history(&sess, &config, text).await; - } - Op::GetHistoryEntryRequest { offset, log_id } => { - handlers::get_history_entry_request(&sess, &config, sub.id.clone(), offset, log_id) .await; - } - Op::ListMcpTools => { - handlers::list_mcp_tools(&sess, &config, sub.id.clone()).await; - } - Op::RefreshMcpServers { config } => { - handlers::refresh_mcp_servers(&sess, config).await; - } - Op::ReloadUserConfig => { - handlers::reload_user_config(&sess).await; - } - Op::ListCustomPrompts => { - handlers::list_custom_prompts(&sess, sub.id.clone()).await; - } - Op::ListSkills { cwds, force_reload } => { - handlers::list_skills(&sess, sub.id.clone(), cwds, force_reload).await; - } - Op::ListRemoteSkills { - hazelnut_scope, - product_surface, - enabled, - } => { - handlers::list_remote_skills( - &sess, - &config, - sub.id.clone(), + false + } + Op::UserInput { .. } | Op::UserTurn { .. } => { + handlers::user_input_or_turn(&sess, sub.id.clone(), sub.op).await; + false + } + Op::ExecApproval { + id: approval_id, + turn_id, + decision, + } => { + handlers::exec_approval(&sess, approval_id, turn_id, decision).await; + false + } + Op::PatchApproval { id, decision } => { + handlers::patch_approval(&sess, id, decision).await; + false + } + Op::UserInputAnswer { id, response } => { + handlers::request_user_input_response(&sess, id, response).await; + false + } + Op::DynamicToolResponse { id, response } => { + handlers::dynamic_tool_response(&sess, id, response).await; + false + } + Op::AddToHistory { text } => { + handlers::add_to_history(&sess, &config, text).await; + false + } + Op::GetHistoryEntryRequest { offset, log_id } => { + handlers::get_history_entry_request( + &sess, + &config, + sub.id.clone(), + offset, + log_id, + ) + .await; + false + } + Op::ListMcpTools => { + handlers::list_mcp_tools(&sess, &config, sub.id.clone()).await; + false + } + Op::RefreshMcpServers { config } => { + handlers::refresh_mcp_servers(&sess, config).await; + false + } + Op::ReloadUserConfig => { + handlers::reload_user_config(&sess).await; + false + } + Op::ListCustomPrompts => { + handlers::list_custom_prompts(&sess, sub.id.clone()).await; + false + } + Op::ListSkills { cwds, force_reload } => { + handlers::list_skills(&sess, sub.id.clone(), cwds, force_reload).await; + false + } + Op::ListRemoteSkills { hazelnut_scope, product_surface, enabled, - ) - .await; - } - Op::DownloadRemoteSkill { hazelnut_id } => { - handlers::export_remote_skill(&sess, &config, sub.id.clone(), hazelnut_id).await; - } - Op::Undo => { - handlers::undo(&sess, sub.id.clone()).await; - } - Op::Compact => { - handlers::compact(&sess, sub.id.clone()).await; - } - Op::DropMemories => { - handlers::drop_memories(&sess, &config, sub.id.clone()).await; - } - Op::UpdateMemories => { - handlers::update_memories(&sess, &config, sub.id.clone()).await; - } - Op::ThreadRollback { num_turns } => { - handlers::thread_rollback(&sess, sub.id.clone(), num_turns).await; - } - Op::SetThreadName { name } => { - handlers::set_thread_name(&sess, sub.id.clone(), name).await; - } - Op::RunUserShellCommand { command } => { - handlers::run_user_shell_command(&sess, sub.id.clone(), command).await; - } - Op::ResolveElicitation { - server_name, - request_id, - decision, - } => { - handlers::resolve_elicitation(&sess, server_name, request_id, decision).await; - } - Op::Shutdown => { - if handlers::shutdown(&sess, sub.id.clone()).await { - break; + } => { + handlers::list_remote_skills( + &sess, + &config, + sub.id.clone(), + hazelnut_scope, + product_surface, + enabled, + ) + .await; + false } + Op::DownloadRemoteSkill { hazelnut_id } => { + handlers::export_remote_skill(&sess, &config, sub.id.clone(), hazelnut_id) + .await; + false + } + Op::Undo => { + handlers::undo(&sess, sub.id.clone()).await; + false + } + Op::Compact => { + handlers::compact(&sess, sub.id.clone()).await; + false + } + Op::DropMemories => { + handlers::drop_memories(&sess, &config, sub.id.clone()).await; + false + } + Op::UpdateMemories => { + handlers::update_memories(&sess, &config, sub.id.clone()).await; + false + } + Op::ThreadRollback { num_turns } => { + handlers::thread_rollback(&sess, sub.id.clone(), num_turns).await; + false + } + Op::SetThreadName { name } => { + handlers::set_thread_name(&sess, sub.id.clone(), name).await; + false + } + Op::RunUserShellCommand { command } => { + handlers::run_user_shell_command(&sess, sub.id.clone(), command).await; + false + } + Op::ResolveElicitation { + server_name, + request_id, + decision, + } => { + handlers::resolve_elicitation(&sess, server_name, request_id, decision).await; + false + } + Op::Shutdown => handlers::shutdown(&sess, sub.id.clone()).await, + Op::Review { review_request } => { + handlers::review(&sess, &config, sub.id.clone(), review_request).await; + false + } + _ => false, // Ignore unknown ops; enum is non_exhaustive to allow extensions. } - Op::Review { review_request } => { - handlers::review(&sess, &config, sub.id.clone(), review_request).await; - } - _ => {} // Ignore unknown ops; enum is non_exhaustive to allow extensions. + } + .instrument(dispatch_span) + .await; + if should_exit { + break; } } debug!("Agent loop exited"); } +fn submission_dispatch_span(sub: &Submission) -> tracing::Span { + let dispatch_span = info_span!("submission_dispatch", submission.id = sub.id.as_str()); + if let Some(trace) = sub.trace.as_ref() + && !set_parent_from_w3c_trace_context(&dispatch_span, trace) + { + warn!( + submission.id = sub.id.as_str(), + "ignoring invalid submission trace carrier" + ); + } + dispatch_span +} + /// Operation handlers mod handlers { use crate::codex::Session; @@ -6627,9 +6690,17 @@ mod tests { use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelsResponse; + use codex_protocol::protocol::Submission; + use codex_protocol::protocol::W3cTraceContext; + use opentelemetry::trace::TraceContextExt; + use opentelemetry::trace::TraceId; + use opentelemetry::trace::TracerProvider as _; + use opentelemetry_sdk::trace::SdkTracerProvider; use std::path::Path; use std::time::Duration; use tokio::time::sleep; + use tracing_opentelemetry::OpenTelemetrySpanExt; + use tracing_subscriber::prelude::*; use codex_protocol::mcp::CallToolResult as McpCallToolResult; use pretty_assertions::assert_eq; @@ -8105,6 +8176,12 @@ mod tests { }) } + fn test_tracing_subscriber() -> impl tracing::Subscriber + Send + Sync { + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("codex-core-tests"); + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)) + } + async fn build_test_config(codex_home: &Path) -> Config { ConfigBuilder::default() .codex_home(codex_home.to_path_buf()) @@ -8432,6 +8509,86 @@ mod tests { (session, turn_context) } + #[tokio::test] + async fn submit_with_id_captures_current_span_trace_context() { + let (session, _turn_context) = make_session_and_context().await; + let (tx_sub, rx_sub) = async_channel::bounded(1); + let (_tx_event, rx_event) = async_channel::unbounded(); + let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); + let codex = Codex { + tx_sub, + rx_event, + agent_status, + session: Arc::new(session), + }; + + let subscriber = test_tracing_subscriber(); + let _guard = tracing::subscriber::set_default(subscriber); + + let request_parent = W3cTraceContext { + traceparent: Some("00-00000000000000000000000000000011-0000000000000022-01".into()), + tracestate: Some("vendor=value".into()), + }; + let request_span = info_span!("app_server.request"); + assert!(set_parent_from_w3c_trace_context( + &request_span, + &request_parent + )); + + let expected_trace = async { + let expected_trace = + current_span_w3c_trace_context().expect("current span should have trace context"); + codex + .submit_with_id(Submission { + id: "sub-1".into(), + op: Op::Interrupt, + trace: None, + }) + .await + .expect("submit should succeed"); + expected_trace + } + .instrument(request_span) + .await; + + let submitted = rx_sub.recv().await.expect("submission"); + assert_eq!(submitted.trace, Some(expected_trace)); + } + + #[test] + fn submission_dispatch_span_prefers_submission_trace_context() { + let subscriber = test_tracing_subscriber(); + let _guard = tracing::subscriber::set_default(subscriber); + + let ambient_parent = W3cTraceContext { + traceparent: Some("00-00000000000000000000000000000033-0000000000000044-01".into()), + tracestate: None, + }; + let ambient_span = info_span!("ambient"); + assert!(set_parent_from_w3c_trace_context( + &ambient_span, + &ambient_parent + )); + + let submission_trace = W3cTraceContext { + traceparent: Some("00-00000000000000000000000000000055-0000000000000066-01".into()), + tracestate: Some("vendor=value".into()), + }; + let dispatch_span = ambient_span.in_scope(|| { + submission_dispatch_span(&Submission { + id: "sub-1".into(), + op: Op::Interrupt, + trace: Some(submission_trace), + }) + }); + + let trace_id = dispatch_span.context().span().span_context().trace_id(); + assert_eq!( + trace_id, + TraceId::from_hex("00000000000000000000000000000055").expect("trace id") + ); + } + pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( dynamic_tools: Vec, ) -> ( diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 1e8dd7141..f6ef7f36e 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -154,6 +154,7 @@ pub(crate) async fn run_codex_thread_one_shot( .send(Submission { id: "shutdown".to_string(), op: Op::Shutdown {}, + trace: None, }) .await; child_cancel.cancel(); @@ -298,11 +299,11 @@ async fn forward_ops( cancel_token_ops: CancellationToken, ) { loop { - let op: Op = match rx_ops.recv().or_cancel(&cancel_token_ops).await { - Ok(Ok(Submission { id: _, op })) => op, + let submission = match rx_ops.recv().or_cancel(&cancel_token_ops).await { + Ok(Ok(submission)) => submission, Ok(Err(_)) | Err(_) => break, }; - let _ = codex.submit(op).await; + let _ = codex.submit_with_id(submission).await; } } @@ -550,4 +551,47 @@ mod tests { "expected Shutdown op after cancellation" ); } + + #[tokio::test] + async fn forward_ops_preserves_submission_trace_context() { + let (tx_sub, rx_sub) = bounded(SUBMISSION_CHANNEL_CAPACITY); + let (_tx_events, rx_events) = bounded(SUBMISSION_CHANNEL_CAPACITY); + let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); + let (session, _ctx, _rx_evt) = crate::codex::make_session_and_context_with_rx().await; + let codex = Arc::new(Codex { + tx_sub, + rx_event: rx_events, + agent_status, + session, + }); + let (tx_ops, rx_ops) = bounded(1); + let cancel = CancellationToken::new(); + let forward = tokio::spawn(forward_ops(Arc::clone(&codex), rx_ops, cancel)); + + let submission = Submission { + id: "sub-1".to_string(), + op: Op::Interrupt, + trace: Some(codex_protocol::protocol::W3cTraceContext { + traceparent: Some( + "00-1234567890abcdef1234567890abcdef-1234567890abcdef-01".to_string(), + ), + tracestate: Some("vendor=state".to_string()), + }), + }; + tx_ops.send(submission.clone()).await.unwrap(); + drop(tx_ops); + + let forwarded = timeout(Duration::from_secs(1), rx_sub.recv()) + .await + .expect("forward_ops hung") + .expect("forwarded submission missing"); + assert_eq!(submission.id, forwarded.id); + assert_eq!(submission.op, forwarded.op); + assert_eq!(submission.trace, forwarded.trace); + + timeout(Duration::from_secs(1), forward) + .await + .expect("forward_ops did not exit") + .expect("forward_ops join error"); + } } diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 62493a46c..b4af4e11f 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -78,9 +78,6 @@ impl SessionTask for RegularTask { let sess = session.clone_session(); let run_turn_span = trace_span!("run_turn"); sess.set_server_reasoning_included(false).await; - sess.services - .otel_manager - .apply_traceparent_parent(&run_turn_span); let prewarmed_client_session = self.take_prewarmed_session().await; run_turn( sess, diff --git a/codex-rs/exec/Cargo.toml b/codex-rs/exec/Cargo.toml index f3fa573ab..1eb4f7e27 100644 --- a/codex-rs/exec/Cargo.toml +++ b/codex-rs/exec/Cargo.toml @@ -21,6 +21,7 @@ clap = { workspace = true, features = ["derive"] } codex-arg0 = { workspace = true } codex-cloud-requirements = { workspace = true } codex-core = { workspace = true } +codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-cli = { workspace = true } @@ -55,10 +56,13 @@ codex-apply-patch = { workspace = true } codex-utils-cargo-bin = { workspace = true } core_test_support = { workspace = true } libc = { workspace = true } +opentelemetry = { workspace = true } +opentelemetry_sdk = { workspace = true } predicates = { workspace = true } pretty_assertions = { workspace = true } rmcp = { workspace = true } tempfile = { workspace = true } +tracing-opentelemetry = { workspace = true } uuid = { workspace = true } walkdir = { workspace = true } wiremock = { workspace = true } diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index 5f0ac4865..0c7271c56 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -34,6 +34,8 @@ use codex_core::format_exec_policy_error_with_source; use codex_core::git_info::get_git_repo_root; use codex_core::models_manager::collaboration_mode_presets::CollaborationModesConfig; use codex_core::models_manager::manager::RefreshStrategy; +use codex_otel::set_parent_from_context; +use codex_otel::traceparent_context_from_env; use codex_protocol::approvals::ElicitationAction; use codex_protocol::config_types::SandboxMode; use codex_protocol::protocol::AskForApproval; @@ -58,9 +60,12 @@ use std::path::PathBuf; use std::sync::Arc; use supports_color::Stream; use tokio::sync::Mutex; +use tracing::Instrument; use tracing::debug; use tracing::error; +use tracing::field; use tracing::info; +use tracing::info_span; use tracing::warn; use tracing_subscriber::EnvFilter; use tracing_subscriber::prelude::*; @@ -94,6 +99,32 @@ struct ThreadEventEnvelope { suppress_output: bool, } +struct ExecRunArgs { + command: Option, + config: Config, + cursor_ansi: bool, + dangerously_bypass_approvals_and_sandbox: bool, + exec_span: tracing::Span, + images: Vec, + json_mode: bool, + last_message_file: Option, + model_provider: Option, + oss: bool, + output_schema_path: Option, + prompt: Option, + skip_git_repo_check: bool, + stderr_with_ansi: bool, +} + +fn exec_root_span() -> tracing::Span { + info_span!( + "codex.exec", + otel.kind = "internal", + thread.id = field::Empty, + turn.id = field::Empty, + ) +} + pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { if let Err(err) = set_default_originator("codex_exec".to_string()) { tracing::warn!(?err, "Failed to set codex exec originator override {err:?}"); @@ -347,6 +378,48 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result .with(otel_logger_layer) .try_init(); + let exec_span = exec_root_span(); + if let Some(context) = traceparent_context_from_env() { + set_parent_from_context(&exec_span, context); + } + run_exec_session(ExecRunArgs { + command, + config, + cursor_ansi, + dangerously_bypass_approvals_and_sandbox, + exec_span: exec_span.clone(), + images, + json_mode, + last_message_file, + model_provider, + oss, + output_schema_path, + prompt, + skip_git_repo_check, + stderr_with_ansi, + }) + .instrument(exec_span) + .await +} + +async fn run_exec_session(args: ExecRunArgs) -> anyhow::Result<()> { + let ExecRunArgs { + command, + config, + cursor_ansi, + dangerously_bypass_approvals_and_sandbox, + exec_span, + images, + json_mode, + last_message_file, + model_provider, + oss, + output_schema_path, + prompt, + skip_git_repo_check, + stderr_with_ansi, + } = args; + let mut event_processor: Box = match json_mode { true => Box::new(EventProcessorWithJsonOutput::new(last_message_file.clone())), _ => Box::new(EventProcessorWithHumanOutput::create_with_ansi( @@ -435,6 +508,9 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result } else { thread_manager.start_thread(config.clone()).await? }; + let primary_thread_id_for_span = primary_thread_id.to_string(); + exec_span.record("thread.id", primary_thread_id_for_span.as_str()); + let (initial_operation, prompt_summary) = match (command, prompt, images) { (Some(ExecCommand::Review(review_cli)), _, _) => { let review_request = build_review_request(review_cli)?; @@ -554,7 +630,7 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result }); } - match initial_operation { + let task_id = match initial_operation { InitialOperation::UserTurn { items, output_schema, @@ -583,6 +659,7 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result task_id } }; + exec_span.record("turn.id", task_id.as_str()); // Run the loop until the task is complete. // Track whether a fatal error was reported by the server so we can @@ -934,13 +1011,44 @@ fn build_review_request(args: ReviewArgs) -> anyhow::Result { #[cfg(test)] mod tests { use super::*; + use codex_otel::set_parent_from_w3c_trace_context; + use opentelemetry::trace::TraceContextExt; + use opentelemetry::trace::TraceId; + use opentelemetry::trace::TracerProvider as _; + use opentelemetry_sdk::trace::SdkTracerProvider; use pretty_assertions::assert_eq; + use tracing_opentelemetry::OpenTelemetrySpanExt; + + fn test_tracing_subscriber() -> impl tracing::Subscriber + Send + Sync { + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("codex-exec-tests"); + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)) + } #[test] fn exec_defaults_analytics_to_enabled() { assert_eq!(DEFAULT_ANALYTICS_ENABLED, true); } + #[test] + fn exec_root_span_can_be_parented_from_trace_context() { + let subscriber = test_tracing_subscriber(); + let _guard = tracing::subscriber::set_default(subscriber); + + let parent = codex_protocol::protocol::W3cTraceContext { + traceparent: Some("00-00000000000000000000000000000077-0000000000000088-01".into()), + tracestate: Some("vendor=value".into()), + }; + let exec_span = exec_root_span(); + assert!(set_parent_from_w3c_trace_context(&exec_span, &parent)); + + let trace_id = exec_span.context().span().span_context().trace_id(); + assert_eq!( + trace_id, + TraceId::from_hex("00000000000000000000000000000077").expect("trace id") + ); + } + #[test] fn builds_uncommitted_review_request() { let request = build_review_request(ReviewArgs { diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index e624fe2b1..24896b042 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -115,6 +115,7 @@ pub async fn run_codex_tool_session( }], final_output_json_schema: None, }, + trace: None, }; if let Err(e) = thread.submit_with_id(submission).await { diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index 6b634e194..e51c3d619 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -576,6 +576,7 @@ impl MessageProcessor { .submit_with_id(Submission { id: request_id_string, op: codex_protocol::protocol::Op::Interrupt, + trace: None, }) .await { diff --git a/codex-rs/otel/src/otel_provider.rs b/codex-rs/otel/src/otel_provider.rs index c9233f5d6..dcfb313b6 100644 --- a/codex-rs/otel/src/otel_provider.rs +++ b/codex-rs/otel/src/otel_provider.rs @@ -7,7 +7,6 @@ use crate::trace_context::context_from_trace_headers; use gethostname::gethostname; use opentelemetry::Context; use opentelemetry::KeyValue; -use opentelemetry::context::ContextGuard; use opentelemetry::global; use opentelemetry::trace::TracerProvider as _; use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge; @@ -28,7 +27,6 @@ use opentelemetry_sdk::trace::BatchSpanProcessor; use opentelemetry_sdk::trace::SdkTracerProvider; use opentelemetry_sdk::trace::Tracer; use opentelemetry_semantic_conventions as semconv; -use std::cell::RefCell; use std::env; use std::error::Error; use std::sync::OnceLock; @@ -43,10 +41,6 @@ const HOST_NAME_ATTRIBUTE: &str = "host.name"; const TRACEPARENT_ENV_VAR: &str = "TRACEPARENT"; const TRACESTATE_ENV_VAR: &str = "TRACESTATE"; static TRACEPARENT_CONTEXT: OnceLock> = OnceLock::new(); - -thread_local! { - static TRACEPARENT_GUARD: RefCell> = const { RefCell::new(None) }; -} pub struct OtelProvider { pub logger: Option, pub tracer_provider: Option, @@ -113,10 +107,6 @@ impl OtelProvider { global::set_tracer_provider(provider); global::set_text_map_propagator(TraceContextPropagator::new()); } - if tracer.is_some() { - attach_traceparent_context(); - } - Ok(Some(Self { logger, tracer_provider, @@ -176,18 +166,6 @@ pub fn traceparent_context_from_env() -> Option { .clone() } -fn attach_traceparent_context() { - TRACEPARENT_GUARD.with(|guard| { - let mut guard = guard.borrow_mut(); - if guard.is_some() { - return; - } - if let Some(context) = traceparent_context_from_env() { - *guard = Some(context.attach()); - } - }); -} - fn load_traceparent_context() -> Option { let traceparent = env::var(TRACEPARENT_ENV_VAR).ok()?; let tracestate = env::var(TRACESTATE_ENV_VAR).ok(); diff --git a/codex-rs/otel/src/traces/otel_manager.rs b/codex-rs/otel/src/traces/otel_manager.rs index 32b9fba86..1074b6edd 100644 --- a/codex-rs/otel/src/traces/otel_manager.rs +++ b/codex-rs/otel/src/traces/otel_manager.rs @@ -15,7 +15,6 @@ use crate::metrics::names::WEBSOCKET_EVENT_COUNT_METRIC; use crate::metrics::names::WEBSOCKET_EVENT_DURATION_METRIC; use crate::metrics::names::WEBSOCKET_REQUEST_COUNT_METRIC; use crate::metrics::names::WEBSOCKET_REQUEST_DURATION_METRIC; -use crate::otel_provider::traceparent_context_from_env; use crate::sanitize_metric_tag_value; use chrono::SecondsFormat; use chrono::Utc; @@ -41,7 +40,6 @@ use std::time::Duration; use std::time::Instant; use tokio::time::error::Elapsed; use tracing::Span; -use tracing_opentelemetry::OpenTelemetrySpanExt; pub use crate::OtelEventMetadata; pub use crate::OtelManager; @@ -92,12 +90,6 @@ impl OtelManager { } } - pub fn apply_traceparent_parent(&self, span: &Span) { - if let Some(context) = traceparent_context_from_env() { - let _ = span.set_parent(context); - } - } - pub fn record_responses(&self, handle_responses_span: &Span, event: &ResponseEvent) { handle_responses_span.record("otel.name", OtelManager::responses_type(event)); diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 3f1c39271..2591a8fb1 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -82,6 +82,9 @@ pub struct Submission { pub id: String, /// Payload pub op: Op, + /// Optional W3C trace carrier propagated across async submission handoffs. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub trace: Option, } #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)]