feat(app-server): propagate traces across tasks and core ops (#14387)

## Summary

This PR keeps app-server RPC request trace context alive for the full
lifetime of the work that request kicks off (e.g. for `thread/start`,
this is `app-server rpc handler -> tokio background task -> core op
submissions`). Previously we lose trace lineage once the request handler
returns or hands work off to background tasks.

This approach is especially relevant for `thread/start` and other RPC
handlers that run in a non-blocking way. In the near future we'll most
likely want to make all app-server handlers run in a non-blocking way by
default, and only queue operations that must operate in order (e.g.
thread RPCs per thread?), so we want to make sure tracing in app-server
just generally works.

Depends on https://github.com/openai/codex/pull/14300

**Before**
<img width="155" height="207" alt="image"
src="https://github.com/user-attachments/assets/c9487459-36f1-436c-beb7-fafeb40737af"
/>


**After**
<img width="299" height="337" alt="image"
src="https://github.com/user-attachments/assets/727392b2-d072-4427-9dc4-0502d8652dea"
/>

## What changed

- Keep request-scoped trace context around until we send the final
response or error, or the connection closes.
- Thread that trace context through detached `thread/start` work so
background startup stays attached to the originating request.
- Pass request trace context through to downstream core operations,
including:
  - thread creation
  - resume/fork flows
  - turn submission
  - review
  - interrupt
  - realtime conversation operations
- Add tracing tests that verify:
  - remote W3C trace context is preserved for `thread/start`
  - remote W3C trace context is preserved for `turn/start`
  - downstream core spans stay under the originating request span
  - request-scoped tracing state is cleaned up correctly
- Clean up shutdown behavior so detached background tasks and spawned
threads are drained before process exit.

codex-rs/Cargo.lock

@@ -1442,6 +1442,8 @@ dependencies = [
  "codex-utils-pty",
  "core_test_support",
  "futures",
+ "opentelemetry",
+ "opentelemetry_sdk",
  "owo-colors",
  "pretty_assertions",
  "reqwest",
@@ -1457,6 +1459,7 @@ dependencies = [
  "tokio-util",
  "toml 0.9.11+spec-1.1.0",
  "tracing",
+ "tracing-opentelemetry",
  "tracing-subscriber",
  "uuid",
  "wiremock",

codex-rs/app-server/Cargo.toml

@@ -79,6 +79,8 @@ axum = { workspace = true, default-features = false, features = [
 ] }
 core_test_support = { workspace = true }
 codex-utils-cargo-bin = { workspace = true }
+opentelemetry = { workspace = true }
+opentelemetry_sdk = { workspace = true }
 pretty_assertions = { workspace = true }
 reqwest = { workspace = true, features = ["rustls-tls"] }
 rmcp = { workspace = true, default-features = false, features = [
@@ -88,5 +90,6 @@ rmcp = { workspace = true, default-features = false, features = [
 ] }
 serial_test = { workspace = true }
 tokio-tungstenite = { workspace = true }
+tracing-opentelemetry = { workspace = true }
 wiremock = { workspace = true }
 shlex = { workspace = true }

codex-rs/app-server/src/app_server_tracing.rs

@@ -27,50 +27,29 @@ pub(crate) fn request_span(
     connection_id: ConnectionId,
     session: &ConnectionSessionState,
 ) -> Span {
-    let span = info_span!(
-        "app_server.request",
-        otel.kind = "server",
-        otel.name = request.method.as_str(),
-        rpc.system = "jsonrpc",
-        rpc.method = request.method.as_str(),
-        rpc.transport = transport_name(transport),
-        rpc.request_id = ?request.id,
-        app_server.connection_id = ?connection_id,
-        app_server.api_version = "v2",
-        app_server.client_name = field::Empty,
-        app_server.client_version = field::Empty,
+    let initialize_client_info = initialize_client_info(request);
+    let method = request.method.as_str();
+    let span = app_server_request_span_template(
+        method,
+        transport_name(transport),
+        &request.id,
+        connection_id,
     );
 
-    let initialize_client_info = initialize_client_info(request);
-    if let Some(client_name) = client_name(initialize_client_info.as_ref(), session) {
-        span.record("app_server.client_name", client_name);
-    }
-    if let Some(client_version) = client_version(initialize_client_info.as_ref(), session) {
-        span.record("app_server.client_version", client_version);
-    }
+    record_client_info(
+        &span,
+        client_name(initialize_client_info.as_ref(), session),
+        client_version(initialize_client_info.as_ref(), session),
+    );
 
-    if let Some(traceparent) = request
-        .trace
-        .as_ref()
-        .and_then(|trace| trace.traceparent.as_deref())
-    {
-        let trace = W3cTraceContext {
-            traceparent: Some(traceparent.to_string()),
-            tracestate: request
-                .trace
-                .as_ref()
-                .and_then(|value| value.tracestate.clone()),
-        };
-        if !set_parent_from_w3c_trace_context(&span, &trace) {
-            tracing::warn!(
-                rpc_method = request.method.as_str(),
-                rpc_request_id = ?request.id,
-                "ignoring invalid inbound request trace carrier"
-            );
-        }
-    } else if let Some(context) = traceparent_context_from_env() {
-        set_parent_from_context(&span, context);
-    }
+    let parent_trace = request.trace.as_ref().and_then(|trace| {
+        trace.traceparent.as_ref()?;
+        Some(W3cTraceContext {
+            traceparent: trace.traceparent.clone(),
+            tracestate: trace.tracestate.clone(),
+        })
+    });
+    attach_parent_context(&span, method, &request.id, parent_trace.as_ref());
 
     span
 }
@@ -86,37 +65,20 @@ pub(crate) fn typed_request_span(
     session: &ConnectionSessionState,
 ) -> Span {
     let method = request.method();
-    let span = info_span!(
-        "app_server.request",
-        otel.kind = "server",
-        otel.name = method,
-        rpc.system = "jsonrpc",
-        rpc.method = method,
-        rpc.transport = "in-process",
-        rpc.request_id = ?request.id(),
-        app_server.connection_id = ?connection_id,
-        app_server.api_version = "v2",
-        app_server.client_name = field::Empty,
-        app_server.client_version = field::Empty,
+    let span = app_server_request_span_template(&method, "in-process", request.id(), connection_id);
+
+    let client_info = initialize_client_info_from_typed_request(request);
+    record_client_info(
+        &span,
+        client_info
+            .map(|(client_name, _)| client_name)
+            .or(session.app_server_client_name.as_deref()),
+        client_info
+            .map(|(_, client_version)| client_version)
+            .or(session.client_version.as_deref()),
     );
 
-    if let Some((client_name, client_version)) = initialize_client_info_from_typed_request(request)
-    {
-        span.record("app_server.client_name", client_name);
-        span.record("app_server.client_version", client_version);
-    } else {
-        if let Some(client_name) = session.app_server_client_name.as_deref() {
-            span.record("app_server.client_name", client_name);
-        }
-        if let Some(client_version) = session.client_version.as_deref() {
-            span.record("app_server.client_version", client_version);
-        }
-    }
-
-    if let Some(context) = traceparent_context_from_env() {
-        set_parent_from_context(&span, context);
-    }
-
+    attach_parent_context(&span, &method, request.id(), None);
     span
 }
 
@@ -127,6 +89,55 @@ fn transport_name(transport: AppServerTransport) -> &'static str {
     }
 }
 
+fn app_server_request_span_template(
+    method: &str,
+    transport: &'static str,
+    request_id: &impl std::fmt::Debug,
+    connection_id: ConnectionId,
+) -> Span {
+    info_span!(
+        "app_server.request",
+        otel.kind = "server",
+        otel.name = method,
+        rpc.system = "jsonrpc",
+        rpc.method = method,
+        rpc.transport = transport,
+        rpc.request_id = ?request_id,
+        app_server.connection_id = ?connection_id,
+        app_server.api_version = "v2",
+        app_server.client_name = field::Empty,
+        app_server.client_version = field::Empty,
+    )
+}
+
+fn record_client_info(span: &Span, client_name: Option<&str>, client_version: Option<&str>) {
+    if let Some(client_name) = client_name {
+        span.record("app_server.client_name", client_name);
+    }
+    if let Some(client_version) = client_version {
+        span.record("app_server.client_version", client_version);
+    }
+}
+
+fn attach_parent_context(
+    span: &Span,
+    method: &str,
+    request_id: &impl std::fmt::Debug,
+    parent_trace: Option<&W3cTraceContext>,
+) {
+    if let Some(trace) = parent_trace {
+        if !set_parent_from_w3c_trace_context(span, trace) {
+            tracing::warn!(
+                rpc_method = method,
+                rpc_request_id = ?request_id,
+                "ignoring invalid inbound request trace carrier"
+            );
+        }
+    } else if let Some(context) = traceparent_context_from_env() {
+        set_parent_from_context(span, context);
+    }
+}
+
 fn client_name<'a>(
     initialize_client_info: Option<&'a InitializeParams>,
     session: &'a ConnectionSessionState,

codex-rs/app-server/src/codex_message_processor.rs

@@ -13,6 +13,7 @@ use crate::outgoing_message::ConnectionId;
 use crate::outgoing_message::ConnectionRequestId;
 use crate::outgoing_message::OutgoingMessageSender;
 use crate::outgoing_message::OutgoingNotification;
+use crate::outgoing_message::RequestContext;
 use crate::outgoing_message::ThreadScopedOutgoingMessageSender;
 use crate::thread_status::ThreadWatchManager;
 use crate::thread_status::resolve_thread_status;
@@ -203,6 +204,7 @@ use codex_core::connectors::filter_disallowed_connectors;
 use codex_core::connectors::merge_plugin_apps;
 use codex_core::default_client::set_default_client_residency_requirement;
 use codex_core::error::CodexErr;
+use codex_core::error::Result as CodexResult;
 use codex_core::exec::ExecExpiration;
 use codex_core::exec::ExecParams;
 use codex_core::exec_env::create_env;
@@ -269,6 +271,7 @@ use codex_protocol::protocol::RolloutItem;
 use codex_protocol::protocol::SessionConfiguredEvent;
 use codex_protocol::protocol::SessionMetaLine;
 use codex_protocol::protocol::USER_MESSAGE_BEGIN;
+use codex_protocol::protocol::W3cTraceContext;
 use codex_protocol::user_input::MAX_USER_INPUT_TEXT_CHARS;
 use codex_protocol::user_input::UserInput as CoreInputItem;
 use codex_rmcp_client::perform_oauth_login_return_url;
@@ -296,7 +299,9 @@ use tokio::sync::broadcast;
 use tokio::sync::oneshot;
 use tokio::sync::watch;
 use tokio_util::sync::CancellationToken;
+use tokio_util::task::TaskTracker;
 use toml::Value as TomlValue;
+use tracing::Instrument;
 use tracing::error;
 use tracing::info;
 use tracing::warn;
@@ -386,6 +391,7 @@ pub(crate) struct CodexMessageProcessor {
     command_exec_manager: CommandExecManager,
     pending_fuzzy_searches: Arc<Mutex<HashMap<String, Arc<AtomicBool>>>>,
     fuzzy_search_sessions: Arc<Mutex<HashMap<String, FuzzyFileSearchSession>>>,
+    background_tasks: TaskTracker,
     feedback: CodexFeedback,
     log_db: Option<LogDbLayer>,
 }
@@ -500,6 +506,7 @@ impl CodexMessageProcessor {
             command_exec_manager: CommandExecManager::default(),
             pending_fuzzy_searches: Arc::new(Mutex::new(HashMap::new())),
             fuzzy_search_sessions: Arc::new(Mutex::new(HashMap::new())),
+            background_tasks: TaskTracker::new(),
             feedback,
             log_db,
         }
@@ -620,6 +627,7 @@ impl CodexMessageProcessor {
         connection_id: ConnectionId,
         request: ClientRequest,
         app_server_client_name: Option<String>,
+        request_context: RequestContext,
     ) {
         let to_connection_request_id = |request_id| ConnectionRequestId {
             connection_id,
@@ -632,8 +640,12 @@ impl CodexMessageProcessor {
             }
             // === v2 Thread/Turn APIs ===
             ClientRequest::ThreadStart { request_id, params } => {
-                self.thread_start(to_connection_request_id(request_id), params)
-                    .await;
+                self.thread_start(
+                    to_connection_request_id(request_id),
+                    params,
+                    request_context,
+                )
+                .await;
             }
             ClientRequest::ThreadUnsubscribe { request_id, params } => {
                 self.thread_unsubscribe(to_connection_request_id(request_id), params)
@@ -1806,7 +1818,12 @@ impl CodexMessageProcessor {
         }
     }
 
-    async fn thread_start(&self, request_id: ConnectionRequestId, params: ThreadStartParams) {
+    async fn thread_start(
+        &self,
+        request_id: ConnectionRequestId,
+        params: ThreadStartParams,
+        request_context: RequestContext,
+    ) {
         let ThreadStartParams {
             model,
             model_provider,
@@ -1847,8 +1864,8 @@ impl CodexMessageProcessor {
             fallback_model_provider: self.config.model_provider_id.clone(),
             codex_home: self.config.codex_home.clone(),
         };
-
-        tokio::spawn(async move {
+        let request_trace = request_context.request_trace();
+        let thread_start_task = async move {
             Self::thread_start_task(
                 listener_task_context,
                 cli_overrides,
@@ -1860,9 +1877,53 @@ impl CodexMessageProcessor {
                 persist_extended_history,
                 service_name,
                 experimental_raw_events,
+                request_trace,
             )
             .await;
-        });
+        };
+        self.background_tasks
+            .spawn(thread_start_task.instrument(request_context.span()));
+    }
+
+    pub(crate) async fn drain_background_tasks(&self) {
+        self.background_tasks.close();
+        if tokio::time::timeout(Duration::from_secs(10), self.background_tasks.wait())
+            .await
+            .is_err()
+        {
+            warn!("timed out waiting for background tasks to shut down; proceeding");
+        }
+    }
+
+    pub(crate) async fn shutdown_threads(&self) {
+        let report = self
+            .thread_manager
+            .shutdown_all_threads_bounded(Duration::from_secs(10))
+            .await;
+        for thread_id in report.submit_failed {
+            warn!("failed to submit Shutdown to thread {thread_id}");
+        }
+        for thread_id in report.timed_out {
+            warn!("timed out waiting for thread {thread_id} to shut down");
+        }
+    }
+
+    async fn request_trace_context(
+        &self,
+        request_id: &ConnectionRequestId,
+    ) -> Option<codex_protocol::protocol::W3cTraceContext> {
+        self.outgoing.request_trace_context(request_id).await
+    }
+
+    async fn submit_core_op(
+        &self,
+        request_id: &ConnectionRequestId,
+        thread: &CodexThread,
+        op: Op,
+    ) -> CodexResult<String> {
+        thread
+            .submit_with_trace(op, self.request_trace_context(request_id).await)
+            .await
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -1877,6 +1938,7 @@ impl CodexMessageProcessor {
         persist_extended_history: bool,
         service_name: Option<String>,
         experimental_raw_events: bool,
+        request_trace: Option<W3cTraceContext>,
     ) {
         let config = match derive_config_from_params(
             &cli_overrides,
@@ -1934,6 +1996,7 @@ impl CodexMessageProcessor {
                 core_dynamic_tools,
                 persist_extended_history,
                 service_name,
+                request_trace,
             )
             .await
         {
@@ -2199,7 +2262,10 @@ impl CodexMessageProcessor {
         };
 
         if let Ok(thread) = self.thread_manager.get_thread(thread_id).await {
-            if let Err(err) = thread.submit(Op::SetThreadName { name }).await {
+            if let Err(err) = self
+                .submit_core_op(&request_id, thread.as_ref(), Op::SetThreadName { name })
+                .await
+            {
                 self.send_internal_error(request_id, format!("failed to set thread name: {err}"))
                     .await;
                 return;
@@ -2784,7 +2850,14 @@ impl CodexMessageProcessor {
             return;
         }
 
-        if let Err(err) = thread.submit(Op::ThreadRollback { num_turns }).await {
+        if let Err(err) = self
+            .submit_core_op(
+                &request_id,
+                thread.as_ref(),
+                Op::ThreadRollback { num_turns },
+            )
+            .await
+        {
             // No ThreadRollback event will arrive if an error occurs.
             // Clean up and reply immediately.
             let thread_state = self.thread_state_manager.thread_state(thread_id).await;
@@ -2812,7 +2885,10 @@ impl CodexMessageProcessor {
             }
         };
 
-        match thread.submit(Op::Compact).await {
+        match self
+            .submit_core_op(&request_id, thread.as_ref(), Op::Compact)
+            .await
+        {
             Ok(_) => {
                 self.outgoing
                     .send_response(request_id, ThreadCompactStartResponse {})
@@ -2840,7 +2916,10 @@ impl CodexMessageProcessor {
             }
         };
 
-        match thread.submit(Op::CleanBackgroundTerminals).await {
+        match self
+            .submit_core_op(&request_id, thread.as_ref(), Op::CleanBackgroundTerminals)
+            .await
+        {
             Ok(_) => {
                 self.outgoing
                     .send_response(request_id, ThreadBackgroundTerminalsCleanResponse {})
@@ -3298,6 +3377,7 @@ impl CodexMessageProcessor {
                 thread_history,
                 self.auth_manager.clone(),
                 persist_extended_history,
+                self.request_trace_context(&request_id).await,
             )
             .await
         {
@@ -3823,6 +3903,7 @@ impl CodexMessageProcessor {
                 config,
                 rollout_path.clone(),
                 persist_extended_history,
+                self.request_trace_context(&request_id).await,
             )
             .await
         {
@@ -4694,26 +4775,10 @@ impl CodexMessageProcessor {
     }
 
     async fn wait_for_thread_shutdown(thread: &Arc<CodexThread>) -> ThreadShutdownResult {
-        match thread.submit(Op::Shutdown).await {
-            Ok(_) => {
-                let wait_for_shutdown = async {
-                    loop {
-                        if matches!(thread.agent_status().await, AgentStatus::Shutdown) {
-                            break;
-                        }
-                        tokio::time::sleep(Duration::from_millis(50)).await;
-                    }
-                };
-                if tokio::time::timeout(Duration::from_secs(10), wait_for_shutdown)
-                    .await
-                    .is_err()
-                {
-                    ThreadShutdownResult::TimedOut
-                } else {
-                    ThreadShutdownResult::Complete
-                }
-            }
-            Err(_) => ThreadShutdownResult::SubmitFailed,
+        match tokio::time::timeout(Duration::from_secs(10), thread.shutdown_and_wait()).await {
+            Ok(Ok(())) => ThreadShutdownResult::Complete,
+            Ok(Err(_)) => ThreadShutdownResult::SubmitFailed,
+            Err(_) => ThreadShutdownResult::TimedOut,
         }
     }
 
@@ -5799,28 +5864,36 @@ impl CodexMessageProcessor {
 
         // If any overrides are provided, update the session turn context first.
         if has_any_overrides {
-            let _ = thread
-                .submit(Op::OverrideTurnContext {
-                    cwd: params.cwd,
-                    approval_policy: params.approval_policy.map(AskForApproval::to_core),
-                    sandbox_policy: params.sandbox_policy.map(|p| p.to_core()),
-                    windows_sandbox_level: None,
-                    model: params.model,
-                    effort: params.effort.map(Some),
-                    summary: params.summary,
-                    service_tier: params.service_tier,
-                    collaboration_mode,
-                    personality: params.personality,
-                })
+            let _ = self
+                .submit_core_op(
+                    &request_id,
+                    thread.as_ref(),
+                    Op::OverrideTurnContext {
+                        cwd: params.cwd,
+                        approval_policy: params.approval_policy.map(AskForApproval::to_core),
+                        sandbox_policy: params.sandbox_policy.map(|p| p.to_core()),
+                        windows_sandbox_level: None,
+                        model: params.model,
+                        effort: params.effort.map(Some),
+                        summary: params.summary,
+                        service_tier: params.service_tier,
+                        collaboration_mode,
+                        personality: params.personality,
+                    },
+                )
                 .await;
         }
 
         // Start the turn by submitting the user input. Return its submission id as turn_id.
-        let turn_id = thread
-            .submit(Op::UserInput {
-                items: mapped_items,
-                final_output_json_schema: params.output_schema,
-            })
+        let turn_id = self
+            .submit_core_op(
+                &request_id,
+                thread.as_ref(),
+                Op::UserInput {
+                    items: mapped_items,
+                    final_output_json_schema: params.output_schema,
+                },
+            )
             .await;
 
         match turn_id {
@@ -5977,11 +6050,15 @@ impl CodexMessageProcessor {
             return;
         };
 
-        let submit = thread
-            .submit(Op::RealtimeConversationStart(ConversationStartParams {
-                prompt: params.prompt,
-                session_id: params.session_id,
-            }))
+        let submit = self
+            .submit_core_op(
+                &request_id,
+                thread.as_ref(),
+                Op::RealtimeConversationStart(ConversationStartParams {
+                    prompt: params.prompt,
+                    session_id: params.session_id,
+                }),
+            )
             .await;
 
         match submit {
@@ -6012,10 +6089,14 @@ impl CodexMessageProcessor {
             return;
         };
 
-        let submit = thread
-            .submit(Op::RealtimeConversationAudio(ConversationAudioParams {
-                frame: params.audio.into(),
-            }))
+        let submit = self
+            .submit_core_op(
+                &request_id,
+                thread.as_ref(),
+                Op::RealtimeConversationAudio(ConversationAudioParams {
+                    frame: params.audio.into(),
+                }),
+            )
             .await;
 
         match submit {
@@ -6046,10 +6127,12 @@ impl CodexMessageProcessor {
             return;
         };
 
-        let submit = thread
-            .submit(Op::RealtimeConversationText(ConversationTextParams {
-                text: params.text,
-            }))
+        let submit = self
+            .submit_core_op(
+                &request_id,
+                thread.as_ref(),
+                Op::RealtimeConversationText(ConversationTextParams { text: params.text }),
+            )
             .await;
 
         match submit {
@@ -6080,7 +6163,9 @@ impl CodexMessageProcessor {
             return;
         };
 
-        let submit = thread.submit(Op::RealtimeConversationClose).await;
+        let submit = self
+            .submit_core_op(&request_id, thread.as_ref(), Op::RealtimeConversationClose)
+            .await;
 
         match submit {
             Ok(_) => {
@@ -6143,7 +6228,13 @@ impl CodexMessageProcessor {
         display_text: &str,
         parent_thread_id: String,
     ) -> std::result::Result<(), JSONRPCErrorError> {
-        let turn_id = parent_thread.submit(Op::Review { review_request }).await;
+        let turn_id = self
+            .submit_core_op(
+                request_id,
+                parent_thread.as_ref(),
+                Op::Review { review_request },
+            )
+            .await;
 
         match turn_id {
             Ok(turn_id) => {
@@ -6197,7 +6288,13 @@ impl CodexMessageProcessor {
             ..
         } = self
             .thread_manager
-            .fork_thread(usize::MAX, config, rollout_path, false)
+            .fork_thread(
+                usize::MAX,
+                config,
+                rollout_path,
+                false,
+                self.request_trace_context(request_id).await,
+            )
             .await
             .map_err(|err| JSONRPCErrorError {
                 code: INTERNAL_ERROR_CODE,
@@ -6252,8 +6349,12 @@ impl CodexMessageProcessor {
             );
         }
 
-        let turn_id = review_thread
-            .submit(Op::Review { review_request })
+        let turn_id = self
+            .submit_core_op(
+                request_id,
+                review_thread.as_ref(),
+                Op::Review { review_request },
+            )
             .await
             .map_err(|err| JSONRPCErrorError {
                 code: INTERNAL_ERROR_CODE,
@@ -6351,7 +6452,9 @@ impl CodexMessageProcessor {
         }
 
         // Submit the interrupt; we'll respond upon TurnAborted.
-        let _ = thread.submit(Op::Interrupt).await;
+        let _ = self
+            .submit_core_op(&request_id, thread.as_ref(), Op::Interrupt)
+            .await;
     }
 
     async fn ensure_conversation_listener(

codex-rs/app-server/src/in_process.rs

@@ -475,6 +475,8 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
                 }
             }
 
+            processor.drain_background_tasks().await;
+            processor.shutdown_threads().await;
             processor.connection_closed(IN_PROCESS_CONNECTION_ID).await;
         });
         let mut pending_request_responses =

codex-rs/app-server/src/lib.rs

@@ -809,6 +809,10 @@ pub async fn run_main_with_transport(
                 }
             }
 
+            if !shutdown_state.forced() {
+                processor.drain_background_tasks().await;
+                processor.shutdown_threads().await;
+            }
             info!("processor task exited (channel closed)");
         }
     });

codex-rs/app-server/src/message_processor.rs

@@ -1,4 +1,5 @@
 use std::collections::HashSet;
+use std::future::Future;
 use std::sync::Arc;
 use std::sync::RwLock;
 use std::sync::atomic::AtomicBool;
@@ -12,6 +13,7 @@ use crate::external_agent_config_api::ExternalAgentConfigApi;
 use crate::outgoing_message::ConnectionId;
 use crate::outgoing_message::ConnectionRequestId;
 use crate::outgoing_message::OutgoingMessageSender;
+use crate::outgoing_message::RequestContext;
 use crate::transport::AppServerTransport;
 use async_trait::async_trait;
 use codex_app_server_protocol::ChatgptAuthTokensRefreshParams;
@@ -55,6 +57,7 @@ use codex_core::models_manager::collaboration_mode_presets::CollaborationModesCo
 use codex_feedback::CodexFeedback;
 use codex_protocol::ThreadId;
 use codex_protocol::protocol::SessionSource;
+use codex_protocol::protocol::W3cTraceContext;
 use codex_state::log_db::LogDbLayer;
 use futures::FutureExt;
 use tokio::sync::broadcast;
@@ -240,53 +243,66 @@ impl MessageProcessor {
         transport: AppServerTransport,
         session: &mut ConnectionSessionState,
     ) {
+        let request_method = request.method.as_str();
+        tracing::trace!(
+            ?connection_id,
+            request_id = ?request.id,
+            "app-server request: {request_method}"
+        );
+        let request_id = ConnectionRequestId {
+            connection_id,
+            request_id: request.id.clone(),
+        };
         let request_span =
             crate::app_server_tracing::request_span(&request, transport, connection_id, session);
-        async {
-            let request_method = request.method.as_str();
-            tracing::trace!(
-                ?connection_id,
-                request_id = ?request.id,
-                "app-server request: {request_method}"
-            );
-            let request_id = ConnectionRequestId {
-                connection_id,
-                request_id: request.id.clone(),
-            };
-            let request_json = match serde_json::to_value(&request) {
-                Ok(request_json) => request_json,
-                Err(err) => {
-                    let error = JSONRPCErrorError {
-                        code: INVALID_REQUEST_ERROR_CODE,
-                        message: format!("Invalid request: {err}"),
-                        data: None,
-                    };
-                    self.outgoing.send_error(request_id, error).await;
-                    return;
-                }
-            };
+        let request_trace = request.trace.as_ref().map(|trace| W3cTraceContext {
+            traceparent: trace.traceparent.clone(),
+            tracestate: trace.tracestate.clone(),
+        });
+        let request_context = RequestContext::new(request_id.clone(), request_span, request_trace);
+        Self::run_request_with_context(
+            Arc::clone(&self.outgoing),
+            request_context.clone(),
+            async {
+                let request_json = match serde_json::to_value(&request) {
+                    Ok(request_json) => request_json,
+                    Err(err) => {
+                        let error = JSONRPCErrorError {
+                            code: INVALID_REQUEST_ERROR_CODE,
+                            message: format!("Invalid request: {err}"),
+                            data: None,
+                        };
+                        self.outgoing.send_error(request_id.clone(), error).await;
+                        return;
+                    }
+                };
 
-            let codex_request = match serde_json::from_value::<ClientRequest>(request_json) {
-                Ok(codex_request) => codex_request,
-                Err(err) => {
-                    let error = JSONRPCErrorError {
-                        code: INVALID_REQUEST_ERROR_CODE,
-                        message: format!("Invalid request: {err}"),
-                        data: None,
-                    };
-                    self.outgoing.send_error(request_id, error).await;
-                    return;
-                }
-            };
-
-            // Websocket callers finalize outbound readiness in lib.rs after mirroring
-            // session state into outbound state and sending initialize notifications to
-            // this specific connection. Passing `None` avoids marking the connection
-            // ready too early from inside the shared request handler.
-            self.handle_client_request(connection_id, request_id, codex_request, session, None)
+                let codex_request = match serde_json::from_value::<ClientRequest>(request_json) {
+                    Ok(codex_request) => codex_request,
+                    Err(err) => {
+                        let error = JSONRPCErrorError {
+                            code: INVALID_REQUEST_ERROR_CODE,
+                            message: format!("Invalid request: {err}"),
+                            data: None,
+                        };
+                        self.outgoing.send_error(request_id.clone(), error).await;
+                        return;
+                    }
+                };
+                // Websocket callers finalize outbound readiness in lib.rs after mirroring
+                // session state into outbound state and sending initialize notifications to
+                // this specific connection. Passing `None` avoids marking the connection
+                // ready too early from inside the shared request handler.
+                self.handle_client_request(
+                    request_id.clone(),
+                    codex_request,
+                    session,
+                    None,
+                    request_context.clone(),
+                )
                 .await;
-        }
-        .instrument(request_span)
+            },
+        )
         .await;
     }
 
@@ -301,31 +317,35 @@ impl MessageProcessor {
         session: &mut ConnectionSessionState,
         outbound_initialized: &AtomicBool,
     ) {
+        let request_id = ConnectionRequestId {
+            connection_id,
+            request_id: request.id().clone(),
+        };
         let request_span =
             crate::app_server_tracing::typed_request_span(&request, connection_id, session);
-        async {
-            let request_id = ConnectionRequestId {
-                connection_id,
-                request_id: request.id().clone(),
-            };
-            tracing::trace!(
-                ?connection_id,
-                request_id = ?request_id.request_id,
-                "app-server typed request"
-            );
-            // In-process clients do not have the websocket transport loop that performs
-            // post-initialize bookkeeping, so they still finalize outbound readiness in
-            // the shared request handler.
-            self.handle_client_request(
-                connection_id,
-                request_id,
-                request,
-                session,
-                Some(outbound_initialized),
-            )
-            .await;
-        }
-        .instrument(request_span)
+        let request_context = RequestContext::new(request_id.clone(), request_span, None);
+        tracing::trace!(
+            ?connection_id,
+            request_id = ?request_id.request_id,
+            "app-server typed request"
+        );
+        Self::run_request_with_context(
+            Arc::clone(&self.outgoing),
+            request_context.clone(),
+            async {
+                // In-process clients do not have the websocket transport loop that performs
+                // post-initialize bookkeeping, so they still finalize outbound readiness in
+                // the shared request handler.
+                self.handle_client_request(
+                    request_id.clone(),
+                    request,
+                    session,
+                    Some(outbound_initialized),
+                    request_context.clone(),
+                )
+                .await;
+            },
+        )
         .await;
     }
 
@@ -342,6 +362,19 @@ impl MessageProcessor {
         tracing::info!("<- typed notification: {:?}", notification);
     }
 
+    async fn run_request_with_context<F>(
+        outgoing: Arc<OutgoingMessageSender>,
+        request_context: RequestContext,
+        request_fut: F,
+    ) where
+        F: Future<Output = ()>,
+    {
+        outgoing
+            .register_request_context(request_context.clone())
+            .await;
+        request_fut.instrument(request_context.span()).await;
+    }
+
     pub(crate) fn thread_created_receiver(&self) -> broadcast::Receiver<ThreadId> {
         self.codex_message_processor.thread_created_receiver()
     }
@@ -384,7 +417,16 @@ impl MessageProcessor {
             .await;
     }
 
+    pub(crate) async fn drain_background_tasks(&self) {
+        self.codex_message_processor.drain_background_tasks().await;
+    }
+
+    pub(crate) async fn shutdown_threads(&self) {
+        self.codex_message_processor.shutdown_threads().await;
+    }
+
     pub(crate) async fn connection_closed(&mut self, connection_id: ConnectionId) {
+        self.outgoing.connection_closed(connection_id).await;
         self.codex_message_processor
             .connection_closed(connection_id)
             .await;
@@ -410,20 +452,21 @@ impl MessageProcessor {
 
     async fn handle_client_request(
         &mut self,
-        connection_id: ConnectionId,
-        request_id: ConnectionRequestId,
+        connection_request_id: ConnectionRequestId,
         codex_request: ClientRequest,
         session: &mut ConnectionSessionState,
         // `Some(...)` means the caller wants initialize to immediately mark the
         // connection outbound-ready. Websocket JSON-RPC calls pass `None` so
         // lib.rs can deliver connection-scoped initialize notifications first.
         outbound_initialized: Option<&AtomicBool>,
+        request_context: RequestContext,
     ) {
+        let connection_id = connection_request_id.connection_id;
         match codex_request {
             // Handle Initialize internally so CodexMessageProcessor does not have to concern
             // itself with the `initialized` bool.
             ClientRequest::Initialize { request_id, params } => {
-                let request_id = ConnectionRequestId {
+                let connection_request_id = ConnectionRequestId {
                     connection_id,
                     request_id,
                 };
@@ -433,7 +476,7 @@ impl MessageProcessor {
                         message: "Already initialized".to_string(),
                         data: None,
                     };
-                    self.outgoing.send_error(request_id, error).await;
+                    self.outgoing.send_error(connection_request_id, error).await;
                     return;
                 }
 
@@ -473,7 +516,9 @@ impl MessageProcessor {
                                 ),
                                 data: None,
                             };
-                            self.outgoing.send_error(request_id.clone(), error).await;
+                            self.outgoing
+                                .send_error(connection_request_id.clone(), error)
+                                .await;
                             return;
                         }
                         SetOriginatorError::AlreadyInitialized => {
@@ -492,7 +537,9 @@ impl MessageProcessor {
 
                 let user_agent = get_codex_user_agent();
                 let response = InitializeResponse { user_agent };
-                self.outgoing.send_response(request_id, response).await;
+                self.outgoing
+                    .send_response(connection_request_id, response)
+                    .await;
 
                 session.initialized = true;
                 if let Some(outbound_initialized) = outbound_initialized {
@@ -513,7 +560,7 @@ impl MessageProcessor {
                         message: "Not initialized".to_string(),
                         data: None,
                     };
-                    self.outgoing.send_error(request_id, error).await;
+                    self.outgoing.send_error(connection_request_id, error).await;
                     return;
                 }
             }
@@ -526,7 +573,7 @@ impl MessageProcessor {
                 message: experimental_required_message(reason),
                 data: None,
             };
-            self.outgoing.send_error(request_id, error).await;
+            self.outgoing.send_error(connection_request_id, error).await;
             return;
         }
 
@@ -596,7 +643,12 @@ impl MessageProcessor {
                 // inline the full `CodexMessageProcessor::process_request` future, which
                 // can otherwise push worker-thread stack usage over the edge.
                 self.codex_message_processor
-                    .process_request(connection_id, other, session.app_server_client_name.clone())
+                    .process_request(
+                        connection_id,
+                        other,
+                        session.app_server_client_name.clone(),
+                        request_context,
+                    )
                     .boxed()
                     .await;
             }
@@ -673,3 +725,6 @@ impl MessageProcessor {
         }
     }
 }
+
+#[cfg(test)]
+mod tracing_tests;

codex-rs/app-server/src/message_processor/tracing_tests.rs

@@ -0,0 +1,544 @@
+use super::ConnectionSessionState;
+use super::MessageProcessor;
+use super::MessageProcessorArgs;
+use crate::outgoing_message::ConnectionId;
+use crate::outgoing_message::OutgoingMessageSender;
+use crate::transport::AppServerTransport;
+use anyhow::Result;
+use app_test_support::create_mock_responses_server_repeating_assistant;
+use app_test_support::write_mock_responses_config_toml;
+use codex_app_server_protocol::ClientInfo;
+use codex_app_server_protocol::ClientRequest;
+use codex_app_server_protocol::InitializeCapabilities;
+use codex_app_server_protocol::InitializeParams;
+use codex_app_server_protocol::InitializeResponse;
+use codex_app_server_protocol::JSONRPCRequest;
+use codex_app_server_protocol::RequestId;
+use codex_app_server_protocol::ThreadStartParams;
+use codex_app_server_protocol::ThreadStartResponse;
+use codex_app_server_protocol::TurnStartParams;
+use codex_app_server_protocol::TurnStartResponse;
+use codex_app_server_protocol::UserInput;
+use codex_arg0::Arg0DispatchPaths;
+use codex_core::config::Config;
+use codex_core::config::ConfigBuilder;
+use codex_core::config_loader::CloudRequirementsLoader;
+use codex_core::config_loader::LoaderOverrides;
+use codex_feedback::CodexFeedback;
+use codex_protocol::protocol::SessionSource;
+use codex_protocol::protocol::W3cTraceContext;
+use opentelemetry::global;
+use opentelemetry::trace::SpanId;
+use opentelemetry::trace::SpanKind;
+use opentelemetry::trace::TraceId;
+use opentelemetry::trace::TracerProvider as _;
+use opentelemetry_sdk::propagation::TraceContextPropagator;
+use opentelemetry_sdk::trace::InMemorySpanExporter;
+use opentelemetry_sdk::trace::SdkTracerProvider;
+use opentelemetry_sdk::trace::SpanData;
+use pretty_assertions::assert_eq;
+use std::collections::BTreeMap;
+use std::path::Path;
+use std::sync::Arc;
+use std::sync::OnceLock;
+use tempfile::TempDir;
+use tokio::sync::mpsc;
+use tracing_subscriber::layer::SubscriberExt;
+use wiremock::MockServer;
+
+const TEST_CONNECTION_ID: ConnectionId = ConnectionId(7);
+
+struct TestTracing {
+    exporter: InMemorySpanExporter,
+    provider: SdkTracerProvider,
+}
+
+struct RemoteTrace {
+    trace_id: TraceId,
+    parent_span_id: SpanId,
+    context: W3cTraceContext,
+}
+
+impl RemoteTrace {
+    fn new(trace_id: &str, parent_span_id: &str) -> Self {
+        let trace_id = TraceId::from_hex(trace_id).expect("trace id");
+        let parent_span_id = SpanId::from_hex(parent_span_id).expect("parent span id");
+        let context = W3cTraceContext {
+            traceparent: Some(format!("00-{trace_id}-{parent_span_id}-01")),
+            tracestate: Some("vendor=value".to_string()),
+        };
+
+        Self {
+            trace_id,
+            parent_span_id,
+            context,
+        }
+    }
+}
+
+fn init_test_tracing() -> &'static TestTracing {
+    static TEST_TRACING: OnceLock<TestTracing> = OnceLock::new();
+    TEST_TRACING.get_or_init(|| {
+        let exporter = InMemorySpanExporter::default();
+        let provider = SdkTracerProvider::builder()
+            .with_simple_exporter(exporter.clone())
+            .build();
+        let tracer = provider.tracer("codex-app-server-message-processor-tests");
+        global::set_text_map_propagator(TraceContextPropagator::new());
+        let subscriber =
+            tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer));
+        tracing::subscriber::set_global_default(subscriber)
+            .expect("global tracing subscriber should only be installed once");
+        TestTracing { exporter, provider }
+    })
+}
+
+fn request_from_client_request(request: ClientRequest) -> JSONRPCRequest {
+    serde_json::from_value(serde_json::to_value(request).expect("serialize client request"))
+        .expect("client request should convert to JSON-RPC")
+}
+
+fn tracing_test_guard() -> &'static tokio::sync::Mutex<()> {
+    static GUARD: OnceLock<tokio::sync::Mutex<()>> = OnceLock::new();
+    GUARD.get_or_init(|| tokio::sync::Mutex::new(()))
+}
+
+struct TracingHarness {
+    _server: MockServer,
+    _codex_home: TempDir,
+    processor: MessageProcessor,
+    outgoing_rx: mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
+    session: ConnectionSessionState,
+    tracing: &'static TestTracing,
+}
+
+impl TracingHarness {
+    async fn new() -> Result<Self> {
+        let server = create_mock_responses_server_repeating_assistant("Done").await;
+        let codex_home = TempDir::new()?;
+        let config = Arc::new(build_test_config(codex_home.path(), &server.uri()).await?);
+        let (processor, outgoing_rx) = build_test_processor(config);
+        let tracing = init_test_tracing();
+        tracing.exporter.reset();
+        tracing::callsite::rebuild_interest_cache();
+        let mut harness = Self {
+            _server: server,
+            _codex_home: codex_home,
+            processor,
+            outgoing_rx,
+            session: ConnectionSessionState::default(),
+            tracing,
+        };
+
+        let _: InitializeResponse = harness
+            .request(
+                ClientRequest::Initialize {
+                    request_id: RequestId::Integer(1),
+                    params: InitializeParams {
+                        client_info: ClientInfo {
+                            name: "codex-app-server-tests".to_string(),
+                            title: None,
+                            version: "0.1.0".to_string(),
+                        },
+                        capabilities: Some(InitializeCapabilities {
+                            experimental_api: true,
+                            ..Default::default()
+                        }),
+                    },
+                },
+                None,
+            )
+            .await;
+        assert!(harness.session.initialized);
+
+        Ok(harness)
+    }
+
+    fn reset_tracing(&self) {
+        self.tracing.exporter.reset();
+    }
+
+    async fn request<T>(&mut self, request: ClientRequest, trace: Option<W3cTraceContext>) -> T
+    where
+        T: serde::de::DeserializeOwned,
+    {
+        let request_id = match request.id() {
+            RequestId::Integer(request_id) => *request_id,
+            request_id => panic!("expected integer request id in test harness, got {request_id:?}"),
+        };
+        let mut request = request_from_client_request(request);
+        request.trace = trace;
+
+        self.processor
+            .process_request(
+                TEST_CONNECTION_ID,
+                request,
+                AppServerTransport::Stdio,
+                &mut self.session,
+            )
+            .await;
+        read_response(&mut self.outgoing_rx, request_id).await
+    }
+
+    async fn start_thread(
+        &mut self,
+        request_id: i64,
+        trace: Option<W3cTraceContext>,
+    ) -> ThreadStartResponse {
+        let response = self
+            .request(
+                ClientRequest::ThreadStart {
+                    request_id: RequestId::Integer(request_id),
+                    params: ThreadStartParams {
+                        ephemeral: Some(true),
+                        ..ThreadStartParams::default()
+                    },
+                },
+                trace,
+            )
+            .await;
+        read_thread_started_notification(&mut self.outgoing_rx).await;
+        response
+    }
+}
+
+async fn build_test_config(codex_home: &Path, server_uri: &str) -> Result<Config> {
+    write_mock_responses_config_toml(
+        codex_home,
+        server_uri,
+        &BTreeMap::new(),
+        8_192,
+        Some(false),
+        "mock_provider",
+        "compact",
+    )?;
+
+    Ok(ConfigBuilder::default()
+        .codex_home(codex_home.to_path_buf())
+        .build()
+        .await?)
+}
+
+fn build_test_processor(
+    config: Arc<Config>,
+) -> (
+    MessageProcessor,
+    mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
+) {
+    let (outgoing_tx, outgoing_rx) = mpsc::channel(16);
+    let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx));
+    let processor = MessageProcessor::new(MessageProcessorArgs {
+        outgoing,
+        arg0_paths: Arg0DispatchPaths::default(),
+        config,
+        cli_overrides: Vec::new(),
+        loader_overrides: LoaderOverrides::default(),
+        cloud_requirements: CloudRequirementsLoader::default(),
+        feedback: CodexFeedback::new(),
+        log_db: None,
+        config_warnings: Vec::new(),
+        session_source: SessionSource::VSCode,
+        enable_codex_api_key_env: false,
+    });
+    (processor, outgoing_rx)
+}
+
+fn span_attr<'a>(span: &'a SpanData, key: &str) -> Option<&'a str> {
+    span.attributes
+        .iter()
+        .find(|kv| kv.key.as_str() == key)
+        .and_then(|kv| match &kv.value {
+            opentelemetry::Value::String(value) => Some(value.as_str()),
+            _ => None,
+        })
+}
+
+fn find_rpc_span_with_trace<'a>(
+    spans: &'a [SpanData],
+    kind: SpanKind,
+    method: &str,
+    trace_id: TraceId,
+) -> &'a SpanData {
+    spans
+        .iter()
+        .find(|span| {
+            span.span_kind == kind
+                && span_attr(span, "rpc.system") == Some("jsonrpc")
+                && span_attr(span, "rpc.method") == Some(method)
+                && span.span_context.trace_id() == trace_id
+        })
+        .unwrap_or_else(|| {
+            panic!(
+                "missing {kind:?} span for rpc.method={method} trace={trace_id}; exported spans:\n{}",
+                format_spans(spans)
+            )
+        })
+}
+
+fn find_span_by_name_with_trace<'a>(
+    spans: &'a [SpanData],
+    name: &str,
+    trace_id: TraceId,
+) -> &'a SpanData {
+    spans
+        .iter()
+        .find(|span| span.name.as_ref() == name && span.span_context.trace_id() == trace_id)
+        .unwrap_or_else(|| {
+            panic!(
+                "missing span named {name} for trace={trace_id}; exported spans:\n{}",
+                format_spans(spans)
+            )
+        })
+}
+
+fn format_spans(spans: &[SpanData]) -> String {
+    spans
+        .iter()
+        .map(|span| {
+            let rpc_method = span_attr(span, "rpc.method").unwrap_or("-");
+            format!(
+                "name={} span_id={} kind={:?} parent={} trace={} rpc.method={}",
+                span.name,
+                span.span_context.span_id(),
+                span.span_kind,
+                span.parent_span_id,
+                span.span_context.trace_id(),
+                rpc_method
+            )
+        })
+        .collect::<Vec<_>>()
+        .join("\n")
+}
+
+fn assert_span_descends_from(spans: &[SpanData], child: &SpanData, ancestor: &SpanData) {
+    let ancestor_span_id = ancestor.span_context.span_id();
+    let mut parent_span_id = child.parent_span_id;
+    while parent_span_id != SpanId::INVALID {
+        if parent_span_id == ancestor_span_id {
+            return;
+        }
+        let Some(parent_span) = spans
+            .iter()
+            .find(|span| span.span_context.span_id() == parent_span_id)
+        else {
+            break;
+        };
+        parent_span_id = parent_span.parent_span_id;
+    }
+
+    panic!(
+        "span {} does not descend from {}; exported spans:\n{}",
+        child.name,
+        ancestor.name,
+        format_spans(spans)
+    );
+}
+
+async fn read_response<T: serde::de::DeserializeOwned>(
+    outgoing_rx: &mut mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
+    request_id: i64,
+) -> T {
+    loop {
+        let envelope = tokio::time::timeout(std::time::Duration::from_secs(5), outgoing_rx.recv())
+            .await
+            .expect("timed out waiting for response")
+            .expect("outgoing channel closed");
+        let crate::outgoing_message::OutgoingEnvelope::ToConnection {
+            connection_id,
+            message,
+        } = envelope
+        else {
+            continue;
+        };
+        if connection_id != TEST_CONNECTION_ID {
+            continue;
+        }
+        let crate::outgoing_message::OutgoingMessage::Response(response) = message else {
+            continue;
+        };
+        if response.id != RequestId::Integer(request_id) {
+            continue;
+        }
+        return serde_json::from_value(response.result)
+            .expect("response payload should deserialize");
+    }
+}
+
+async fn read_thread_started_notification(
+    outgoing_rx: &mut mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
+) {
+    loop {
+        let envelope = tokio::time::timeout(std::time::Duration::from_secs(5), outgoing_rx.recv())
+            .await
+            .expect("timed out waiting for thread/started notification")
+            .expect("outgoing channel closed");
+        match envelope {
+            crate::outgoing_message::OutgoingEnvelope::ToConnection {
+                connection_id,
+                message,
+            } => {
+                if connection_id != TEST_CONNECTION_ID {
+                    continue;
+                }
+                let crate::outgoing_message::OutgoingMessage::AppServerNotification(notification) =
+                    message
+                else {
+                    continue;
+                };
+                if matches!(
+                    notification,
+                    codex_app_server_protocol::ServerNotification::ThreadStarted(_)
+                ) {
+                    return;
+                }
+            }
+            crate::outgoing_message::OutgoingEnvelope::Broadcast { message } => {
+                let crate::outgoing_message::OutgoingMessage::AppServerNotification(notification) =
+                    message
+                else {
+                    continue;
+                };
+                if matches!(
+                    notification,
+                    codex_app_server_protocol::ServerNotification::ThreadStarted(_)
+                ) {
+                    return;
+                }
+            }
+        }
+    }
+}
+
+async fn wait_for_exported_spans<F>(tracing: &TestTracing, predicate: F) -> Vec<SpanData>
+where
+    F: Fn(&[SpanData]) -> bool,
+{
+    let mut last_spans = Vec::new();
+    for _ in 0..200 {
+        tokio::task::yield_now().await;
+        tracing
+            .provider
+            .force_flush()
+            .expect("force flush should succeed");
+        let spans = tracing.exporter.get_finished_spans().expect("span export");
+        last_spans = spans.clone();
+        if predicate(&spans) {
+            return spans;
+        }
+        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+    }
+
+    panic!(
+        "timed out waiting for expected exported spans:\n{}",
+        format_spans(&last_spans)
+    );
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn thread_start_jsonrpc_span_exports_server_span_and_parents_children() -> Result<()> {
+    let _guard = tracing_test_guard().lock().await;
+    let mut harness = TracingHarness::new().await?;
+
+    let RemoteTrace {
+        trace_id: remote_trace_id,
+        context: remote_trace,
+        ..
+    } = RemoteTrace::new("00000000000000000000000000000011", "0000000000000022");
+
+    let _: ThreadStartResponse = harness.start_thread(2, Some(remote_trace)).await;
+    drop(harness.processor);
+    let spans = wait_for_exported_spans(harness.tracing, |spans| {
+        spans.iter().any(|span| {
+            span.span_kind == SpanKind::Server
+                && span_attr(span, "rpc.method") == Some("thread/start")
+                && span.span_context.trace_id() == remote_trace_id
+        }) && spans
+            .iter()
+            .any(|span| span.name.as_ref() == "thread_spawn")
+    })
+    .await;
+
+    let server_request_span =
+        find_rpc_span_with_trace(&spans, SpanKind::Server, "thread/start", remote_trace_id);
+    let thread_spawn_span = find_span_by_name_with_trace(&spans, "thread_spawn", remote_trace_id);
+    let session_init_span = find_span_by_name_with_trace(&spans, "session_init", remote_trace_id);
+    assert_eq!(server_request_span.name.as_ref(), "thread/start");
+    assert_eq!(server_request_span.span_context.trace_id(), remote_trace_id);
+    assert_ne!(server_request_span.span_context.span_id(), SpanId::INVALID);
+    assert_span_descends_from(&spans, thread_spawn_span, server_request_span);
+    assert_span_descends_from(&spans, session_init_span, server_request_span);
+
+    Ok(())
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn turn_start_jsonrpc_span_parents_core_turn_spans() -> Result<()> {
+    let _guard = tracing_test_guard().lock().await;
+    let mut harness = TracingHarness::new().await?;
+    let thread_start_response = harness.start_thread(2, None).await;
+    let thread_id = thread_start_response.thread.id.clone();
+
+    harness.reset_tracing();
+
+    let RemoteTrace {
+        trace_id: remote_trace_id,
+        parent_span_id: remote_parent_span_id,
+        context: remote_trace,
+    } = RemoteTrace::new("00000000000000000000000000000077", "0000000000000088");
+    let _: TurnStartResponse = harness
+        .request(
+            ClientRequest::TurnStart {
+                request_id: RequestId::Integer(3),
+                params: TurnStartParams {
+                    thread_id,
+                    input: vec![UserInput::Text {
+                        text: "hello".to_string(),
+                        text_elements: Vec::new(),
+                    }],
+                    cwd: None,
+                    approval_policy: None,
+                    sandbox_policy: None,
+                    model: None,
+                    service_tier: None,
+                    effort: None,
+                    summary: None,
+                    personality: None,
+                    output_schema: None,
+                    collaboration_mode: None,
+                },
+            },
+            Some(remote_trace),
+        )
+        .await;
+    let spans = wait_for_exported_spans(harness.tracing, |spans| {
+        spans
+            .iter()
+            .any(|span| span.name.as_ref() == "submission_dispatch")
+            && spans
+                .iter()
+                .any(|span| span.name.as_ref() == "session_task.turn")
+            && spans.iter().any(|span| span.name.as_ref() == "run_turn")
+    })
+    .await;
+    drop(harness.processor);
+    tokio::task::yield_now().await;
+
+    let server_request_span =
+        find_rpc_span_with_trace(&spans, SpanKind::Server, "turn/start", remote_trace_id);
+    let submission_dispatch_span =
+        find_span_by_name_with_trace(&spans, "submission_dispatch", remote_trace_id);
+    let session_task_turn_span =
+        find_span_by_name_with_trace(&spans, "session_task.turn", remote_trace_id);
+    let run_turn_span = find_span_by_name_with_trace(&spans, "run_turn", remote_trace_id);
+
+    assert_eq!(server_request_span.parent_span_id, remote_parent_span_id);
+    assert!(server_request_span.parent_span_is_remote);
+    assert_eq!(server_request_span.span_context.trace_id(), remote_trace_id);
+    assert_span_descends_from(&spans, submission_dispatch_span, server_request_span);
+    assert_span_descends_from(&spans, session_task_turn_span, server_request_span);
+    assert_span_descends_from(&spans, run_turn_span, server_request_span);
+    assert_span_descends_from(&spans, session_task_turn_span, submission_dispatch_span);
+    assert_span_descends_from(&spans, run_turn_span, session_task_turn_span);
+
+    Ok(())
+}

codex-rs/app-server/src/outgoing_message.rs

@@ -9,11 +9,15 @@ use codex_app_server_protocol::Result;
 use codex_app_server_protocol::ServerNotification;
 use codex_app_server_protocol::ServerRequest;
 use codex_app_server_protocol::ServerRequestPayload;
+use codex_otel::span_w3c_trace_context;
 use codex_protocol::ThreadId;
+use codex_protocol::protocol::W3cTraceContext;
 use serde::Serialize;
 use tokio::sync::Mutex;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
+use tracing::Instrument;
+use tracing::Span;
 use tracing::warn;
 
 use crate::error_code::INTERNAL_ERROR_CODE;
@@ -35,6 +39,37 @@ pub(crate) struct ConnectionRequestId {
     pub(crate) request_id: RequestId,
 }
 
+/// Trace data we keep for an incoming request until we send its final
+/// response or error.
+#[derive(Clone)]
+pub(crate) struct RequestContext {
+    request_id: ConnectionRequestId,
+    span: Span,
+    parent_trace: Option<W3cTraceContext>,
+}
+
+impl RequestContext {
+    pub(crate) fn new(
+        request_id: ConnectionRequestId,
+        span: Span,
+        parent_trace: Option<W3cTraceContext>,
+    ) -> Self {
+        Self {
+            request_id,
+            span,
+            parent_trace,
+        }
+    }
+
+    pub(crate) fn request_trace(&self) -> Option<W3cTraceContext> {
+        span_w3c_trace_context(&self.span).or_else(|| self.parent_trace.clone())
+    }
+
+    pub(crate) fn span(&self) -> Span {
+        self.span.clone()
+    }
+}
+
 #[derive(Debug, Clone)]
 pub(crate) enum OutgoingEnvelope {
     ToConnection {
@@ -51,6 +86,10 @@ pub(crate) struct OutgoingMessageSender {
     next_server_request_id: AtomicI64,
     sender: mpsc::Sender<OutgoingEnvelope>,
     request_id_to_callback: Mutex<HashMap<RequestId, PendingCallbackEntry>>,
+    /// Incoming requests that are still waiting on a final response or error.
+    /// We keep them here because this is where responses, errors, and
+    /// disconnect cleanup all get handled.
+    request_contexts: Mutex<HashMap<ConnectionRequestId, RequestContext>>,
 }
 
 #[derive(Clone)]
@@ -142,9 +181,48 @@ impl OutgoingMessageSender {
             next_server_request_id: AtomicI64::new(0),
             sender,
             request_id_to_callback: Mutex::new(HashMap::new()),
+            request_contexts: Mutex::new(HashMap::new()),
         }
     }
 
+    pub(crate) async fn register_request_context(&self, request_context: RequestContext) {
+        let mut request_contexts = self.request_contexts.lock().await;
+        if request_contexts
+            .insert(request_context.request_id.clone(), request_context)
+            .is_some()
+        {
+            warn!("replaced unresolved request context");
+        }
+    }
+
+    pub(crate) async fn connection_closed(&self, connection_id: ConnectionId) {
+        let mut request_contexts = self.request_contexts.lock().await;
+        request_contexts.retain(|request_id, _| request_id.connection_id != connection_id);
+    }
+
+    pub(crate) async fn request_trace_context(
+        &self,
+        request_id: &ConnectionRequestId,
+    ) -> Option<W3cTraceContext> {
+        let request_contexts = self.request_contexts.lock().await;
+        request_contexts
+            .get(request_id)
+            .and_then(RequestContext::request_trace)
+    }
+
+    async fn take_request_context(
+        &self,
+        request_id: &ConnectionRequestId,
+    ) -> Option<RequestContext> {
+        let mut request_contexts = self.request_contexts.lock().await;
+        request_contexts.remove(request_id)
+    }
+
+    #[cfg(test)]
+    async fn request_context_count(&self) -> usize {
+        self.request_contexts.lock().await.len()
+    }
+
     pub(crate) async fn send_request(
         &self,
         request: ServerRequestPayload,
@@ -353,25 +431,24 @@ impl OutgoingMessageSender {
         request_id: ConnectionRequestId,
         response: T,
     ) {
+        let request_context = self.take_request_context(&request_id).await;
         match serde_json::to_value(response) {
             Ok(result) => {
                 let outgoing_message = OutgoingMessage::Response(OutgoingResponse {
-                    id: request_id.request_id,
+                    id: request_id.request_id.clone(),
                     result,
                 });
-                if let Err(err) = self
-                    .sender
-                    .send(OutgoingEnvelope::ToConnection {
-                        connection_id: request_id.connection_id,
-                        message: outgoing_message,
-                    })
-                    .await
-                {
-                    warn!("failed to send response to client: {err:?}");
-                }
+                self.send_outgoing_message_to_connection(
+                    request_context,
+                    request_id.connection_id,
+                    outgoing_message,
+                    "response",
+                )
+                .await;
             }
             Err(err) => {
-                self.send_error(
+                self.send_error_inner(
+                    request_context,
                     request_id,
                     JSONRPCErrorError {
                         code: INTERNAL_ERROR_CODE,
@@ -461,20 +538,50 @@ impl OutgoingMessageSender {
         &self,
         request_id: ConnectionRequestId,
         error: JSONRPCErrorError,
+    ) {
+        let request_context = self.take_request_context(&request_id).await;
+        self.send_error_inner(request_context, request_id, error)
+            .await;
+    }
+
+    async fn send_error_inner(
+        &self,
+        request_context: Option<RequestContext>,
+        request_id: ConnectionRequestId,
+        error: JSONRPCErrorError,
     ) {
         let outgoing_message = OutgoingMessage::Error(OutgoingError {
             id: request_id.request_id,
             error,
         });
-        if let Err(err) = self
-            .sender
-            .send(OutgoingEnvelope::ToConnection {
-                connection_id: request_id.connection_id,
-                message: outgoing_message,
-            })
-            .await
-        {
-            warn!("failed to send error to client: {err:?}");
+        self.send_outgoing_message_to_connection(
+            request_context,
+            request_id.connection_id,
+            outgoing_message,
+            "error",
+        )
+        .await;
+    }
+
+    async fn send_outgoing_message_to_connection(
+        &self,
+        request_context: Option<RequestContext>,
+        connection_id: ConnectionId,
+        message: OutgoingMessage,
+        message_kind: &'static str,
+    ) {
+        let send_fut = self.sender.send(OutgoingEnvelope::ToConnection {
+            connection_id,
+            message,
+        });
+        let send_result = if let Some(request_context) = request_context {
+            send_fut.instrument(request_context.span()).await
+        } else {
+            send_fut.await
+        };
+
+        if let Err(err) = send_result {
+            warn!("failed to send {message_kind} to client: {err:?}");
         }
     }
 }
@@ -738,6 +845,31 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn send_response_clears_registered_request_context() {
+        let (tx, _rx) = mpsc::channel::<OutgoingEnvelope>(4);
+        let outgoing = OutgoingMessageSender::new(tx);
+        let request_id = ConnectionRequestId {
+            connection_id: ConnectionId(42),
+            request_id: RequestId::Integer(7),
+        };
+
+        outgoing
+            .register_request_context(RequestContext::new(
+                request_id.clone(),
+                tracing::info_span!("app_server.request", rpc.method = "thread/start"),
+                None,
+            ))
+            .await;
+        assert_eq!(outgoing.request_context_count().await, 1);
+
+        outgoing
+            .send_response(request_id, json!({ "ok": true }))
+            .await;
+
+        assert_eq!(outgoing.request_context_count().await, 0);
+    }
+
     #[tokio::test]
     async fn send_error_routes_to_target_connection() {
         let (tx, mut rx) = mpsc::channel::<OutgoingEnvelope>(4);
@@ -775,6 +907,40 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn connection_closed_clears_registered_request_contexts() {
+        let (tx, _rx) = mpsc::channel::<OutgoingEnvelope>(4);
+        let outgoing = OutgoingMessageSender::new(tx);
+        let closed_connection_request = ConnectionRequestId {
+            connection_id: ConnectionId(9),
+            request_id: RequestId::Integer(3),
+        };
+        let open_connection_request = ConnectionRequestId {
+            connection_id: ConnectionId(10),
+            request_id: RequestId::Integer(4),
+        };
+
+        outgoing
+            .register_request_context(RequestContext::new(
+                closed_connection_request,
+                tracing::info_span!("app_server.request", rpc.method = "turn/interrupt"),
+                None,
+            ))
+            .await;
+        outgoing
+            .register_request_context(RequestContext::new(
+                open_connection_request,
+                tracing::info_span!("app_server.request", rpc.method = "turn/start"),
+                None,
+            ))
+            .await;
+        assert_eq!(outgoing.request_context_count().await, 2);
+
+        outgoing.connection_closed(ConnectionId(9)).await;
+
+        assert_eq!(outgoing.request_context_count().await, 1);
+    }
+
     #[tokio::test]
     async fn notify_client_error_forwards_error_to_waiter() {
         let (tx, _rx) = mpsc::channel::<OutgoingEnvelope>(4);

codex-rs/core/src/codex.rs

@@ -104,6 +104,7 @@ use codex_protocol::protocol::TurnAbortReason;
 use codex_protocol::protocol::TurnContextItem;
 use codex_protocol::protocol::TurnContextNetworkItem;
 use codex_protocol::protocol::TurnStartedEvent;
+use codex_protocol::protocol::W3cTraceContext;
 use codex_protocol::request_permissions::PermissionGrantScope;
 use codex_protocol::request_permissions::RequestPermissionsArgs;
 use codex_protocol::request_permissions::RequestPermissionsEvent;
@@ -118,6 +119,7 @@ use codex_utils_stream_parser::ProposedPlanSegment;
 use codex_utils_stream_parser::extract_proposed_plan_text;
 use codex_utils_stream_parser::strip_citations;
 use futures::future::BoxFuture;
+use futures::future::Shared;
 use futures::prelude::*;
 use futures::stream::FuturesOrdered;
 use rmcp::model::ListResourceTemplatesResult;
@@ -330,8 +332,13 @@ pub struct Codex {
     // Last known status of the agent.
     pub(crate) agent_status: watch::Receiver<AgentStatus>,
     pub(crate) session: Arc<Session>,
+    // Shared future for the background submission loop completion so multiple
+    // callers can wait for shutdown.
+    pub(crate) session_loop_termination: SessionLoopTermination,
 }
 
+pub(crate) type SessionLoopTermination = Shared<BoxFuture<'static, ()>>;
+
 /// Wrapper returned by [`Codex::spawn`] containing the spawned [`Codex`],
 /// the submission id for the initial `ConfigureSession` request and the
 /// unique session id.
@@ -342,6 +349,24 @@ pub struct CodexSpawnOk {
     pub conversation_id: ThreadId,
 }
 
+pub(crate) struct CodexSpawnArgs {
+    pub(crate) config: Config,
+    pub(crate) auth_manager: Arc<AuthManager>,
+    pub(crate) models_manager: Arc<ModelsManager>,
+    pub(crate) skills_manager: Arc<SkillsManager>,
+    pub(crate) plugins_manager: Arc<PluginsManager>,
+    pub(crate) mcp_manager: Arc<McpManager>,
+    pub(crate) file_watcher: Arc<FileWatcher>,
+    pub(crate) conversation_history: InitialHistory,
+    pub(crate) session_source: SessionSource,
+    pub(crate) agent_control: AgentControl,
+    pub(crate) dynamic_tools: Vec<DynamicToolSpec>,
+    pub(crate) persist_extended_history: bool,
+    pub(crate) metrics_service_name: Option<String>,
+    pub(crate) inherited_shell_snapshot: Option<Arc<ShellSnapshot>>,
+    pub(crate) parent_trace: Option<W3cTraceContext>,
+}
+
 pub(crate) const INITIAL_SUBMIT_ID: &str = "";
 pub(crate) const SUBMISSION_CHANNEL_CAPACITY: usize = 512;
 const CYBER_VERIFY_URL: &str = "https://chatgpt.com/cyber";
@@ -349,23 +374,48 @@ const CYBER_SAFETY_URL: &str = "https://developers.openai.com/codex/concepts/cyb
 
 impl Codex {
     /// Spawn a new [`Codex`] and initialize the session.
-    #[allow(clippy::too_many_arguments)]
-    pub(crate) async fn spawn(
-        mut config: Config,
-        auth_manager: Arc<AuthManager>,
-        models_manager: Arc<ModelsManager>,
-        skills_manager: Arc<SkillsManager>,
-        plugins_manager: Arc<PluginsManager>,
-        mcp_manager: Arc<McpManager>,
-        file_watcher: Arc<FileWatcher>,
-        conversation_history: InitialHistory,
-        session_source: SessionSource,
-        agent_control: AgentControl,
-        dynamic_tools: Vec<DynamicToolSpec>,
-        persist_extended_history: bool,
-        metrics_service_name: Option<String>,
-        inherited_shell_snapshot: Option<Arc<ShellSnapshot>>,
-    ) -> CodexResult<CodexSpawnOk> {
+    pub(crate) async fn spawn(args: CodexSpawnArgs) -> CodexResult<CodexSpawnOk> {
+        let parent_trace = match args.parent_trace {
+            Some(trace) => {
+                if codex_otel::context_from_w3c_trace_context(&trace).is_some() {
+                    Some(trace)
+                } else {
+                    warn!("ignoring invalid thread spawn trace carrier");
+                    None
+                }
+            }
+            None => None,
+        };
+        let thread_spawn_span = info_span!("thread_spawn", otel.name = "thread_spawn");
+        if let Some(trace) = parent_trace.as_ref() {
+            let _ = set_parent_from_w3c_trace_context(&thread_spawn_span, trace);
+        }
+        Self::spawn_internal(CodexSpawnArgs {
+            parent_trace,
+            ..args
+        })
+        .instrument(thread_spawn_span)
+        .await
+    }
+
+    async fn spawn_internal(args: CodexSpawnArgs) -> CodexResult<CodexSpawnOk> {
+        let CodexSpawnArgs {
+            mut config,
+            auth_manager,
+            models_manager,
+            skills_manager,
+            plugins_manager,
+            mcp_manager,
+            file_watcher,
+            conversation_history,
+            session_source,
+            agent_control,
+            dynamic_tools,
+            persist_extended_history,
+            metrics_service_name,
+            inherited_shell_snapshot,
+            parent_trace: _,
+        } = args;
         let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
         let (tx_event, rx_event) = async_channel::unbounded();
 
@@ -557,15 +607,18 @@ impl Codex {
         let thread_id = session.conversation_id;
 
         // This task will run until Op::Shutdown is received.
-        let session_loop_span = info_span!("session_loop", thread_id = %thread_id);
-        tokio::spawn(
-            submission_loop(Arc::clone(&session), config, rx_sub).instrument(session_loop_span),
-        );
+        let session_for_loop = Arc::clone(&session);
+        let session_loop_handle = tokio::spawn(async move {
+            submission_loop(session_for_loop, config, rx_sub)
+                .instrument(info_span!("session_loop", thread_id = %thread_id))
+                .await;
+        });
         let codex = Codex {
             tx_sub,
             rx_event,
             agent_status: agent_status_rx,
             session,
+            session_loop_termination: session_loop_termination_from_handle(session_loop_handle),
         };
 
         #[allow(deprecated)]
@@ -578,11 +631,19 @@ impl Codex {
 
     /// Submit the `op` wrapped in a `Submission` with a unique ID.
     pub async fn submit(&self, op: Op) -> CodexResult<String> {
+        self.submit_with_trace(op, None).await
+    }
+
+    pub async fn submit_with_trace(
+        &self,
+        op: Op,
+        trace: Option<W3cTraceContext>,
+    ) -> CodexResult<String> {
         let id = Uuid::now_v7().to_string();
         let sub = Submission {
             id: id.clone(),
             op,
-            trace: None,
+            trace,
         };
         self.submit_with_id(sub).await?;
         Ok(id)
@@ -601,6 +662,17 @@ impl Codex {
         Ok(())
     }
 
+    pub async fn shutdown_and_wait(&self) -> CodexResult<()> {
+        let session_loop_termination = self.session_loop_termination.clone();
+        match self.submit(Op::Shutdown).await {
+            Ok(_) => {}
+            Err(CodexErr::InternalAgentDied) => {}
+            Err(err) => return Err(err),
+        }
+        session_loop_termination.await;
+        Ok(())
+    }
+
     pub async fn next_event(&self) -> CodexResult<Event> {
         let event = self
             .rx_event
@@ -648,6 +720,21 @@ impl Codex {
     }
 }
 
+#[cfg(test)]
+pub(crate) fn completed_session_loop_termination() -> SessionLoopTermination {
+    futures::future::ready(()).boxed().shared()
+}
+
+pub(crate) fn session_loop_termination_from_handle(
+    handle: JoinHandle<()>,
+) -> SessionLoopTermination {
+    async move {
+        let _ = handle.await;
+    }
+    .boxed()
+    .shared()
+}
+
 /// Context for an initialized model agent
 ///
 /// A session has at most 1 running task at a time, and can be interrupted by user input.

codex-rs/core/src/codex_delegate.rs

@@ -27,6 +27,7 @@ use tokio_util::sync::CancellationToken;
 
 use crate::AuthManager;
 use crate::codex::Codex;
+use crate::codex::CodexSpawnArgs;
 use crate::codex::CodexSpawnOk;
 use crate::codex::SUBMISSION_CHANNEL_CAPACITY;
 use crate::codex::Session;
@@ -36,6 +37,9 @@ use crate::error::CodexErr;
 use crate::models_manager::manager::ModelsManager;
 use codex_protocol::protocol::InitialHistory;
 
+#[cfg(test)]
+use crate::codex::completed_session_loop_termination;
+
 /// Start an interactive sub-Codex thread and return IO channels.
 ///
 /// The returned `events_rx` yields non-approval events emitted by the sub-agent.
@@ -55,22 +59,23 @@ pub(crate) async fn run_codex_thread_interactive(
     let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
     let (tx_ops, rx_ops) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
 
-    let CodexSpawnOk { codex, .. } = Codex::spawn(
+    let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs {
         config,
         auth_manager,
         models_manager,
-        Arc::clone(&parent_session.services.skills_manager),
-        Arc::clone(&parent_session.services.plugins_manager),
-        Arc::clone(&parent_session.services.mcp_manager),
-        Arc::clone(&parent_session.services.file_watcher),
-        initial_history.unwrap_or(InitialHistory::New),
-        SessionSource::SubAgent(subagent_source),
-        parent_session.services.agent_control.clone(),
-        Vec::new(),
-        false,
-        None,
-        None,
-    )
+        skills_manager: Arc::clone(&parent_session.services.skills_manager),
+        plugins_manager: Arc::clone(&parent_session.services.plugins_manager),
+        mcp_manager: Arc::clone(&parent_session.services.mcp_manager),
+        file_watcher: Arc::clone(&parent_session.services.file_watcher),
+        conversation_history: initial_history.unwrap_or(InitialHistory::New),
+        session_source: SessionSource::SubAgent(subagent_source),
+        agent_control: parent_session.services.agent_control.clone(),
+        dynamic_tools: Vec::new(),
+        persist_extended_history: false,
+        metrics_service_name: None,
+        inherited_shell_snapshot: None,
+        parent_trace: None,
+    })
     .await?;
     let codex = Arc::new(codex);
 
@@ -105,6 +110,7 @@ pub(crate) async fn run_codex_thread_interactive(
         rx_event: rx_sub,
         agent_status: codex.agent_status.clone(),
         session: Arc::clone(&codex.session),
+        session_loop_termination: codex.session_loop_termination.clone(),
     })
 }
 
@@ -151,6 +157,7 @@ pub(crate) async fn run_codex_thread_one_shot(
     let ops_tx = io.tx_sub.clone();
     let agent_status = io.agent_status.clone();
     let session = Arc::clone(&io.session);
+    let session_loop_termination = io.session_loop_termination.clone();
     let io_for_bridge = io;
     tokio::spawn(async move {
         while let Ok(event) = io_for_bridge.next_event().await {
@@ -184,6 +191,7 @@ pub(crate) async fn run_codex_thread_one_shot(
         tx_sub: tx_closed,
         agent_status,
         session,
+        session_loop_termination,
     })
 }
 
@@ -572,6 +580,7 @@ mod tests {
             rx_event: rx_events,
             agent_status,
             session: Arc::clone(&session),
+            session_loop_termination: completed_session_loop_termination(),
         });
 
         let (tx_out, rx_out) = bounded(1);
@@ -645,6 +654,7 @@ mod tests {
             rx_event: rx_events,
             agent_status,
             session,
+            session_loop_termination: completed_session_loop_termination(),
         });
         let (tx_ops, rx_ops) = bounded(1);
         let cancel = CancellationToken::new();
@@ -691,6 +701,7 @@ mod tests {
             rx_event: rx_events_child,
             agent_status,
             session: Arc::clone(&parent_session),
+            session_loop_termination: completed_session_loop_termination(),
         });
 
         let call_id = "tool-call-1".to_string();

codex-rs/core/src/codex_tests.rs

@@ -2362,6 +2362,7 @@ async fn submit_with_id_captures_current_span_trace_context() {
         rx_event,
         agent_status,
         session: Arc::new(session),
+        session_loop_termination: completed_session_loop_termination(),
     };
 
     init_test_tracing();
@@ -2589,6 +2590,81 @@ async fn spawn_task_turn_span_inherits_dispatch_trace_context() {
     );
 }
 
+#[tokio::test]
+async fn shutdown_and_wait_allows_multiple_waiters() {
+    let (session, _turn_context) = make_session_and_context().await;
+    let (tx_sub, rx_sub) = async_channel::bounded(4);
+    let (_tx_event, rx_event) = async_channel::unbounded();
+    let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit);
+    let session_loop_handle = tokio::spawn(async move {
+        let shutdown: Submission = rx_sub.recv().await.expect("shutdown submission");
+        assert_eq!(shutdown.op, Op::Shutdown);
+        tokio::time::sleep(StdDuration::from_millis(50)).await;
+    });
+    let codex = Arc::new(Codex {
+        tx_sub,
+        rx_event,
+        agent_status,
+        session: Arc::new(session),
+        session_loop_termination: session_loop_termination_from_handle(session_loop_handle),
+    });
+
+    let waiter_1 = {
+        let codex = Arc::clone(&codex);
+        tokio::spawn(async move { codex.shutdown_and_wait().await })
+    };
+    let waiter_2 = {
+        let codex = Arc::clone(&codex);
+        tokio::spawn(async move { codex.shutdown_and_wait().await })
+    };
+
+    waiter_1
+        .await
+        .expect("first shutdown waiter join")
+        .expect("first shutdown waiter");
+    waiter_2
+        .await
+        .expect("second shutdown waiter join")
+        .expect("second shutdown waiter");
+}
+
+#[tokio::test]
+async fn shutdown_and_wait_waits_when_shutdown_is_already_in_progress() {
+    let (session, _turn_context) = make_session_and_context().await;
+    let (tx_sub, rx_sub) = async_channel::bounded(4);
+    drop(rx_sub);
+    let (_tx_event, rx_event) = async_channel::unbounded();
+    let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit);
+    let (shutdown_complete_tx, shutdown_complete_rx) = tokio::sync::oneshot::channel();
+    let session_loop_handle = tokio::spawn(async move {
+        let _ = shutdown_complete_rx.await;
+    });
+    let codex = Arc::new(Codex {
+        tx_sub,
+        rx_event,
+        agent_status,
+        session: Arc::new(session),
+        session_loop_termination: session_loop_termination_from_handle(session_loop_handle),
+    });
+
+    let waiter = {
+        let codex = Arc::clone(&codex);
+        tokio::spawn(async move { codex.shutdown_and_wait().await })
+    };
+
+    tokio::time::sleep(StdDuration::from_millis(10)).await;
+    assert!(!waiter.is_finished());
+
+    shutdown_complete_tx
+        .send(())
+        .expect("session loop should still be waiting to terminate");
+
+    waiter
+        .await
+        .expect("shutdown waiter join")
+        .expect("shutdown waiter");
+}
+
 pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
     dynamic_tools: Vec<DynamicToolSpec>,
 ) -> (

codex-rs/core/src/codex_tests_guardian.rs

@@ -289,7 +289,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
     let mcp_manager = Arc::new(McpManager::new(Arc::clone(&plugins_manager)));
     let file_watcher = Arc::new(FileWatcher::noop());
 
-    let CodexSpawnOk { codex, .. } = Codex::spawn(
+    let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs {
         config,
         auth_manager,
         models_manager,
@@ -297,14 +297,17 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
         plugins_manager,
         mcp_manager,
         file_watcher,
-        InitialHistory::New,
-        SessionSource::SubAgent(SubAgentSource::Other(GUARDIAN_SUBAGENT_NAME.to_string())),
-        AgentControl::default(),
-        Vec::new(),
-        false,
-        None,
-        None,
-    )
+        conversation_history: InitialHistory::New,
+        session_source: SessionSource::SubAgent(SubAgentSource::Other(
+            GUARDIAN_SUBAGENT_NAME.to_string(),
+        )),
+        agent_control: AgentControl::default(),
+        dynamic_tools: Vec::new(),
+        persist_extended_history: false,
+        metrics_service_name: None,
+        inherited_shell_snapshot: None,
+        parent_trace: None,
+    })
     .await
     .expect("spawn guardian subagent");
 

codex-rs/core/src/codex_thread.rs

@@ -19,6 +19,7 @@ use codex_protocol::protocol::AskForApproval;
 use codex_protocol::protocol::SandboxPolicy;
 use codex_protocol::protocol::SessionSource;
 use codex_protocol::protocol::TokenUsage;
+use codex_protocol::protocol::W3cTraceContext;
 use codex_protocol::user_input::UserInput;
 use std::path::PathBuf;
 use tokio::sync::Mutex;
@@ -67,6 +68,18 @@ impl CodexThread {
         self.codex.submit(op).await
     }
 
+    pub async fn shutdown_and_wait(&self) -> CodexResult<()> {
+        self.codex.shutdown_and_wait().await
+    }
+
+    pub async fn submit_with_trace(
+        &self,
+        op: Op,
+        trace: Option<W3cTraceContext>,
+    ) -> CodexResult<String> {
+        self.codex.submit_with_trace(op, trace).await
+    }
+
     pub async fn steer_input(
         &self,
         input: Vec<UserInput>,

codex-rs/core/src/memories/tests.rs

@@ -547,10 +547,12 @@ mod phase2 {
         }
 
         async fn shutdown_threads(&self) {
-            self.manager
-                .remove_and_close_all_threads()
-                .await
-                .expect("shutdown spawned threads");
+            let report = self
+                .manager
+                .shutdown_all_threads_bounded(std::time::Duration::from_secs(10))
+                .await;
+            assert!(report.submit_failed.is_empty());
+            assert!(report.timed_out.is_empty());
         }
 
         fn user_input_ops_count(&self) -> usize {

codex-rs/core/src/thread_manager.rs

@@ -3,6 +3,7 @@ use crate::CodexAuth;
 use crate::ModelProviderInfo;
 use crate::agent::AgentControl;
 use crate::codex::Codex;
+use crate::codex::CodexSpawnArgs;
 use crate::codex::CodexSpawnOk;
 use crate::codex::INITIAL_SUBMIT_ID;
 use crate::codex_thread::CodexThread;
@@ -30,11 +31,15 @@ use codex_protocol::protocol::McpServerRefreshConfig;
 use codex_protocol::protocol::Op;
 use codex_protocol::protocol::RolloutItem;
 use codex_protocol::protocol::SessionSource;
+use codex_protocol::protocol::W3cTraceContext;
+use futures::StreamExt;
+use futures::stream::FuturesUnordered;
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
+use std::time::Duration;
 use tokio::runtime::Handle;
 use tokio::runtime::RuntimeFlavor;
 use tokio::sync::RwLock;
@@ -118,6 +123,19 @@ pub struct NewThread {
     pub session_configured: SessionConfiguredEvent,
 }
 
+#[derive(Debug, Default, PartialEq, Eq)]
+pub struct ThreadShutdownReport {
+    pub completed: Vec<ThreadId>,
+    pub submit_failed: Vec<ThreadId>,
+    pub timed_out: Vec<ThreadId>,
+}
+
+enum ShutdownOutcome {
+    Complete,
+    SubmitFailed,
+    TimedOut,
+}
+
 /// [`ThreadManager`] is responsible for creating threads and maintaining
 /// them in memory.
 pub struct ThreadManager {
@@ -329,6 +347,7 @@ impl ThreadManager {
             dynamic_tools,
             persist_extended_history,
             None,
+            None,
         ))
         .await
     }
@@ -339,6 +358,7 @@ impl ThreadManager {
         dynamic_tools: Vec<codex_protocol::dynamic_tools::DynamicToolSpec>,
         persist_extended_history: bool,
         metrics_service_name: Option<String>,
+        parent_trace: Option<W3cTraceContext>,
     ) -> CodexResult<NewThread> {
         Box::pin(self.state.spawn_thread(
             config,
@@ -348,6 +368,7 @@ impl ThreadManager {
             dynamic_tools,
             persist_extended_history,
             metrics_service_name,
+            parent_trace,
         ))
         .await
     }
@@ -357,10 +378,17 @@ impl ThreadManager {
         config: Config,
         rollout_path: PathBuf,
         auth_manager: Arc<AuthManager>,
+        parent_trace: Option<W3cTraceContext>,
     ) -> CodexResult<NewThread> {
         let initial_history = RolloutRecorder::get_rollout_history(&rollout_path).await?;
-        Box::pin(self.resume_thread_with_history(config, initial_history, auth_manager, false))
-            .await
+        Box::pin(self.resume_thread_with_history(
+            config,
+            initial_history,
+            auth_manager,
+            false,
+            parent_trace,
+        ))
+        .await
     }
 
     pub async fn resume_thread_with_history(
@@ -369,6 +397,7 @@ impl ThreadManager {
         initial_history: InitialHistory,
         auth_manager: Arc<AuthManager>,
         persist_extended_history: bool,
+        parent_trace: Option<W3cTraceContext>,
     ) -> CodexResult<NewThread> {
         Box::pin(self.state.spawn_thread(
             config,
@@ -378,6 +407,7 @@ impl ThreadManager {
             Vec::new(),
             persist_extended_history,
             None,
+            parent_trace,
         ))
         .await
     }
@@ -389,13 +419,55 @@ impl ThreadManager {
         self.state.threads.write().await.remove(thread_id)
     }
 
-    /// Closes all threads open in this ThreadManager
-    pub async fn remove_and_close_all_threads(&self) -> CodexResult<()> {
-        for thread in self.state.threads.read().await.values() {
-            thread.submit(Op::Shutdown).await?;
+    /// Tries to shut down all tracked threads concurrently within the provided timeout.
+    /// Threads that complete shutdown are removed from the manager; incomplete shutdowns
+    /// remain tracked so callers can retry or inspect them later.
+    pub async fn shutdown_all_threads_bounded(&self, timeout: Duration) -> ThreadShutdownReport {
+        let threads = {
+            let threads = self.state.threads.read().await;
+            threads
+                .iter()
+                .map(|(thread_id, thread)| (*thread_id, Arc::clone(thread)))
+                .collect::<Vec<_>>()
+        };
+
+        let mut shutdowns = threads
+            .into_iter()
+            .map(|(thread_id, thread)| async move {
+                let outcome = match tokio::time::timeout(timeout, thread.shutdown_and_wait()).await
+                {
+                    Ok(Ok(())) => ShutdownOutcome::Complete,
+                    Ok(Err(_)) => ShutdownOutcome::SubmitFailed,
+                    Err(_) => ShutdownOutcome::TimedOut,
+                };
+                (thread_id, outcome)
+            })
+            .collect::<FuturesUnordered<_>>();
+        let mut report = ThreadShutdownReport::default();
+
+        while let Some((thread_id, outcome)) = shutdowns.next().await {
+            match outcome {
+                ShutdownOutcome::Complete => report.completed.push(thread_id),
+                ShutdownOutcome::SubmitFailed => report.submit_failed.push(thread_id),
+                ShutdownOutcome::TimedOut => report.timed_out.push(thread_id),
+            }
         }
-        self.state.threads.write().await.clear();
-        Ok(())
+
+        let mut tracked_threads = self.state.threads.write().await;
+        for thread_id in &report.completed {
+            tracked_threads.remove(thread_id);
+        }
+
+        report
+            .completed
+            .sort_by_key(std::string::ToString::to_string);
+        report
+            .submit_failed
+            .sort_by_key(std::string::ToString::to_string);
+        report
+            .timed_out
+            .sort_by_key(std::string::ToString::to_string);
+        report
     }
 
     /// Fork an existing thread by taking messages up to the given position (not including
@@ -408,6 +480,7 @@ impl ThreadManager {
         config: Config,
         path: PathBuf,
         persist_extended_history: bool,
+        parent_trace: Option<W3cTraceContext>,
     ) -> CodexResult<NewThread> {
         let history = RolloutRecorder::get_rollout_history(&path).await?;
         let history = truncate_before_nth_user_message(history, nth_user_message);
@@ -419,6 +492,7 @@ impl ThreadManager {
             Vec::new(),
             persist_extended_history,
             None,
+            parent_trace,
         ))
         .await
     }
@@ -503,6 +577,7 @@ impl ThreadManagerState {
             persist_extended_history,
             metrics_service_name,
             inherited_shell_snapshot,
+            None,
         ))
         .await
     }
@@ -526,6 +601,7 @@ impl ThreadManagerState {
             false,
             None,
             inherited_shell_snapshot,
+            None,
         ))
         .await
     }
@@ -549,6 +625,7 @@ impl ThreadManagerState {
             persist_extended_history,
             None,
             inherited_shell_snapshot,
+            None,
         ))
         .await
     }
@@ -564,6 +641,7 @@ impl ThreadManagerState {
         dynamic_tools: Vec<codex_protocol::dynamic_tools::DynamicToolSpec>,
         persist_extended_history: bool,
         metrics_service_name: Option<String>,
+        parent_trace: Option<W3cTraceContext>,
     ) -> CodexResult<NewThread> {
         Box::pin(self.spawn_thread_with_source(
             config,
@@ -575,6 +653,7 @@ impl ThreadManagerState {
             persist_extended_history,
             metrics_service_name,
             None,
+            parent_trace,
         ))
         .await
     }
@@ -591,28 +670,30 @@ impl ThreadManagerState {
         persist_extended_history: bool,
         metrics_service_name: Option<String>,
         inherited_shell_snapshot: Option<Arc<ShellSnapshot>>,
+        parent_trace: Option<W3cTraceContext>,
     ) -> CodexResult<NewThread> {
         let watch_registration = self
             .file_watcher
             .register_config(&config, self.skills_manager.as_ref());
         let CodexSpawnOk {
             codex, thread_id, ..
-        } = Codex::spawn(
+        } = Codex::spawn(CodexSpawnArgs {
             config,
             auth_manager,
-            Arc::clone(&self.models_manager),
-            Arc::clone(&self.skills_manager),
-            Arc::clone(&self.plugins_manager),
-            Arc::clone(&self.mcp_manager),
-            Arc::clone(&self.file_watcher),
-            initial_history,
+            models_manager: Arc::clone(&self.models_manager),
+            skills_manager: Arc::clone(&self.skills_manager),
+            plugins_manager: Arc::clone(&self.plugins_manager),
+            mcp_manager: Arc::clone(&self.mcp_manager),
+            file_watcher: Arc::clone(&self.file_watcher),
+            conversation_history: initial_history,
             session_source,
             agent_control,
             dynamic_tools,
             persist_extended_history,
             metrics_service_name,
             inherited_shell_snapshot,
-        )
+            parent_trace,
+        })
         .await?;
         self.finalize_thread_spawn(codex, thread_id, watch_registration)
             .await
@@ -672,11 +753,14 @@ fn truncate_before_nth_user_message(history: InitialHistory, n: usize) -> Initia
 mod tests {
     use super::*;
     use crate::codex::make_session_and_context;
+    use crate::config::test_config;
     use assert_matches::assert_matches;
     use codex_protocol::models::ContentItem;
     use codex_protocol::models::ReasoningItemReasoningSummary;
     use codex_protocol::models::ResponseItem;
     use pretty_assertions::assert_eq;
+    use std::time::Duration;
+    use tempfile::tempdir;
 
     fn user_msg(text: &str) -> ResponseItem {
         ResponseItem::Message {
@@ -783,4 +867,40 @@ mod tests {
             serde_json::to_value(&expected).unwrap()
         );
     }
+
+    #[tokio::test]
+    async fn shutdown_all_threads_bounded_submits_shutdown_to_every_thread() {
+        let temp_dir = tempdir().expect("tempdir");
+        let mut config = test_config();
+        config.codex_home = temp_dir.path().join("codex-home");
+        config.cwd = config.codex_home.clone();
+        std::fs::create_dir_all(&config.codex_home).expect("create codex home");
+
+        let manager = ThreadManager::with_models_provider_and_home_for_tests(
+            CodexAuth::from_api_key("dummy"),
+            config.model_provider.clone(),
+            config.codex_home.clone(),
+        );
+        let thread_1 = manager
+            .start_thread(config.clone())
+            .await
+            .expect("start first thread")
+            .thread_id;
+        let thread_2 = manager
+            .start_thread(config)
+            .await
+            .expect("start second thread")
+            .thread_id;
+
+        let report = manager
+            .shutdown_all_threads_bounded(Duration::from_secs(10))
+            .await;
+
+        let mut expected_completed = vec![thread_1, thread_2];
+        expected_completed.sort_by_key(std::string::ToString::to_string);
+        assert_eq!(report.completed, expected_completed);
+        assert!(report.submit_failed.is_empty());
+        assert!(report.timed_out.is_empty());
+        assert!(manager.list_thread_ids().await.is_empty());
+    }
 }

codex-rs/core/src/tools/handlers/multi_agents.rs

@@ -1725,6 +1725,7 @@ mod tests {
                 })]),
                 AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")),
                 false,
+                None,
             )
             .await
             .expect("start thread");

codex-rs/core/tests/common/test_codex.rs

@@ -202,6 +202,7 @@ impl TestCodexBuilder {
                     config.clone(),
                     path,
                     auth_manager,
+                    None,
                 ))
                 .await?
             }

codex-rs/core/tests/suite/compact_resume_fork.rs

@@ -687,7 +687,7 @@ async fn resume_conversation(
     let auth_manager = codex_core::test_support::auth_manager_from_auth(
         codex_core::CodexAuth::from_api_key("dummy"),
     );
-    Box::pin(manager.resume_thread_from_rollout(config.clone(), path, auth_manager))
+    Box::pin(manager.resume_thread_from_rollout(config.clone(), path, auth_manager, None))
         .await
         .expect("resume conversation")
         .thread
@@ -700,7 +700,7 @@ async fn fork_thread(
     path: std::path::PathBuf,
     nth_user_message: usize,
 ) -> Arc<CodexThread> {
-    Box::pin(manager.fork_thread(nth_user_message, config.clone(), path, false))
+    Box::pin(manager.fork_thread(nth_user_message, config.clone(), path, false, None))
         .await
         .expect("fork conversation")
         .thread

codex-rs/core/tests/suite/fork_thread.rs

@@ -110,7 +110,7 @@ async fn fork_thread_twice_drops_to_first_message() {
         thread: codex_fork1,
         ..
     } = thread_manager
-        .fork_thread(1, config_for_fork.clone(), base_path.clone(), false)
+        .fork_thread(1, config_for_fork.clone(), base_path.clone(), false, None)
         .await
         .expect("fork 1");
 
@@ -129,7 +129,7 @@ async fn fork_thread_twice_drops_to_first_message() {
         thread: codex_fork2,
         ..
     } = thread_manager
-        .fork_thread(0, config_for_fork.clone(), fork1_path.clone(), false)
+        .fork_thread(0, config_for_fork.clone(), fork1_path.clone(), false, None)
         .await
         .expect("fork 2");
 

codex-rs/core/tests/suite/permissions_messages.rs

@@ -416,7 +416,7 @@ async fn resume_and_fork_append_permissions_messages() -> Result<()> {
     fork_config.permissions.approval_policy = Constrained::allow_any(AskForApproval::UnlessTrusted);
     let forked = initial
         .thread_manager
-        .fork_thread(usize::MAX, fork_config, rollout_path, false)
+        .fork_thread(usize::MAX, fork_config, rollout_path, false, None)
         .await?;
     forked
         .thread

codex-rs/core/tests/suite/resume_warning.rs

@@ -98,7 +98,7 @@ async fn emits_warning_when_resumed_model_differs() {
         thread: conversation,
         ..
     } = thread_manager
-        .resume_thread_with_history(config, initial_history, auth_manager, false)
+        .resume_thread_with_history(config, initial_history, auth_manager, false, None)
         .await
         .expect("resume conversation");
 

codex-rs/core/tests/suite/unstable_features_warning.rs

@@ -42,7 +42,7 @@ async fn emits_warning_when_unstable_features_enabled_via_config() {
         thread: conversation,
         ..
     } = thread_manager
-        .resume_thread_with_history(config, InitialHistory::New, auth_manager, false)
+        .resume_thread_with_history(config, InitialHistory::New, auth_manager, false, None)
         .await
         .expect("spawn conversation");
 
@@ -83,7 +83,7 @@ async fn suppresses_warning_when_configured() {
         thread: conversation,
         ..
     } = thread_manager
-        .resume_thread_with_history(config, InitialHistory::New, auth_manager, false)
+        .resume_thread_with_history(config, InitialHistory::New, auth_manager, false, None)
         .await
         .expect("spawn conversation");
 

codex-rs/tui/src/app.rs

@@ -1567,8 +1567,16 @@ impl App {
             self.chat_widget.thread_name(),
         );
         self.shutdown_current_thread().await;
-        if let Err(err) = self.server.remove_and_close_all_threads().await {
-            tracing::warn!(error = %err, "failed to close all threads");
+        let report = self
+            .server
+            .shutdown_all_threads_bounded(Duration::from_secs(10))
+            .await;
+        if !report.submit_failed.is_empty() || !report.timed_out.is_empty() {
+            tracing::warn!(
+                submit_failed = report.submit_failed.len(),
+                timed_out = report.timed_out.len(),
+                "failed to close all threads"
+            );
         }
         let init = crate::chatwidget::ChatWidgetInit {
             config,
@@ -1834,6 +1842,7 @@ impl App {
                         config.clone(),
                         target_session.path.clone(),
                         auth_manager.clone(),
+                        None,
                     )
                     .await
                     .wrap_err_with(|| {
@@ -1871,6 +1880,7 @@ impl App {
                         config.clone(),
                         target_session.path.clone(),
                         false,
+                        None,
                     )
                     .await
                     .wrap_err_with(|| {
@@ -2182,6 +2192,7 @@ impl App {
                                 resume_config.clone(),
                                 target_session.path.clone(),
                                 self.auth_manager.clone(),
+                                None,
                             )
                             .await
                         {
@@ -2250,7 +2261,7 @@ impl App {
                     if path.exists() {
                         match self
                             .server
-                            .fork_thread(usize::MAX, self.config.clone(), path.clone(), false)
+                            .fork_thread(usize::MAX, self.config.clone(), path.clone(), false, None)
                             .await
                         {
                             Ok(forked) => {