Propagate RPC spans into bounded app-server tasks

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

@@ -23,6 +23,15 @@ use tracing::Span;
 use tracing::field;
 use tracing::info_span;
 
+pub(crate) fn spawn_blocking_in_current_span<F, T>(function: F) -> tokio::task::JoinHandle<T>
+where
+    F: FnOnce() -> T + Send + 'static,
+    T: Send + 'static,
+{
+    let span = Span::current();
+    tokio::task::spawn_blocking(move || span.in_scope(function))
+}
+
 pub(crate) fn request_span(
     request: &JSONRPCRequest,
     transport: &AppServerTransport,

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

@@ -118,6 +118,7 @@ use std::time::SystemTime;
 use std::time::UNIX_EPOCH;
 use tokio::sync::Mutex;
 use tokio::sync::oneshot;
+use tracing::Instrument;
 use tracing::error;
 
 enum CommandExecutionApprovalPresentation {
@@ -543,7 +544,7 @@ pub(crate) async fn apply_bespoke_event_handling(
             let (pending_request_id, rx) = outgoing
                 .send_request(ServerRequestPayload::FileChangeRequestApproval(params))
                 .await;
-            tokio::spawn(async move {
+            let response_task = async move {
                 on_file_change_request_approval_response(
                     item_id,
                     pending_request_id,
@@ -553,7 +554,8 @@ pub(crate) async fn apply_bespoke_event_handling(
                     permission_guard,
                 )
                 .await;
-            });
+            };
+            tokio::spawn(response_task.in_current_span());
         }
         EventMsg::ExecApprovalRequest(ev) => {
             let permission_guard = thread_watch_manager
@@ -661,7 +663,7 @@ pub(crate) async fn apply_bespoke_event_handling(
                     params,
                 ))
                 .await;
-            tokio::spawn(async move {
+            let response_task = async move {
                 on_command_execution_request_approval_response(
                     event_turn_id,
                     conversation_id,
@@ -676,7 +678,8 @@ pub(crate) async fn apply_bespoke_event_handling(
                     permission_guard,
                 )
                 .await;
-            });
+            };
+            tokio::spawn(response_task.in_current_span());
         }
         EventMsg::RequestUserInput(request) => {
             let user_input_guard = thread_watch_manager
@@ -712,7 +715,7 @@ pub(crate) async fn apply_bespoke_event_handling(
             let (pending_request_id, rx) = outgoing
                 .send_request(ServerRequestPayload::ToolRequestUserInput(params))
                 .await;
-            tokio::spawn(async move {
+            let response_task = async move {
                 on_request_user_input_response(
                     event_turn_id,
                     pending_request_id,
@@ -722,7 +725,8 @@ pub(crate) async fn apply_bespoke_event_handling(
                     user_input_guard,
                 )
                 .await;
-            });
+            };
+            tokio::spawn(response_task.in_current_span());
         }
         EventMsg::ElicitationRequest(request) => {
             let permission_guard = thread_watch_manager
@@ -769,7 +773,7 @@ pub(crate) async fn apply_bespoke_event_handling(
             let (pending_request_id, rx) = outgoing
                 .send_request(ServerRequestPayload::McpServerElicitationRequest(params))
                 .await;
-            tokio::spawn(async move {
+            let response_task = async move {
                 on_mcp_server_elicitation_response(
                     request.server_name,
                     request.id,
@@ -780,7 +784,8 @@ pub(crate) async fn apply_bespoke_event_handling(
                     permission_guard,
                 )
                 .await;
-            });
+            };
+            tokio::spawn(response_task.in_current_span());
         }
         EventMsg::RequestPermissions(request) => {
             let permission_guard = thread_watch_manager
@@ -815,9 +820,10 @@ pub(crate) async fn apply_bespoke_event_handling(
                 receiver: rx,
                 request_permissions_guard: permission_guard,
             };
-            tokio::spawn(async move {
+            let response_task = async move {
                 on_request_permissions_response(pending_response, conversation, thread_state).await;
-            });
+            };
+            tokio::spawn(response_task.in_current_span());
         }
         EventMsg::DynamicToolCallRequest(_)
         | EventMsg::DynamicToolCallResponse(_)
@@ -972,9 +978,10 @@ pub(crate) async fn apply_bespoke_event_handling(
                 let (_pending_request_id, rx) = outgoing
                     .send_request(ServerRequestPayload::DynamicToolCall(params))
                     .await;
-                tokio::spawn(async move {
+                let response_task = async move {
                     crate::dynamic_tools::on_call_response(call_id, rx, conversation).await;
-                });
+                };
+                tokio::spawn(response_task.in_current_span());
             }
         }
         EventMsg::ItemCompleted(event) => {

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

@@ -33,6 +33,7 @@ use tokio::sync::Mutex;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
 use tokio::sync::watch;
+use tracing::Instrument;
 
 use crate::error_code::internal_error;
 use crate::error_code::invalid_params;
@@ -200,12 +201,21 @@ impl CommandExecManager {
                 );
             }
             let sessions = Arc::clone(&self.sessions);
-            tokio::spawn(async move {
+            let run_span = tracing::info_span!(
+                "app_server.command_exec.run",
+                rpc.method = "command/exec",
+                rpc.request_id = %request_id.request_id,
+                app_server.connection_id = %request_id.connection_id,
+                command_exec.mode = "windows_restricted_token",
+                command_exec.outcome = tracing::field::Empty,
+            );
+            let run_task = async move {
                 let _started_network_proxy = started_network_proxy;
                 match codex_core::sandboxing::execute_env(exec_request, /*stdout_stream*/ None)
                     .await
                 {
                     Ok(output) => {
+                        tracing::Span::current().record("command_exec.outcome", "exited");
                         outgoing
                             .send_response(
                                 request_id,
@@ -218,13 +228,15 @@ impl CommandExecManager {
                             .await;
                     }
                     Err(err) => {
+                        tracing::Span::current().record("command_exec.outcome", "error");
                         outgoing
                             .send_error(request_id, internal_error(format!("exec failed: {err}")))
                             .await;
                     }
                 }
                 sessions.lock().await.remove(&process_key);
-            });
+            };
+            tokio::spawn(run_task.instrument(run_span));
             return Ok(());
         }
 
@@ -290,7 +302,7 @@ impl CommandExecManager {
                 return Err(internal_error(format!("failed to spawn command: {err}")));
             }
         };
-        tokio::spawn(async move {
+        let run_task = async move {
             let _started_network_proxy = started_network_proxy;
             run_command(RunCommandParams {
                 outgoing,
@@ -305,7 +317,8 @@ impl CommandExecManager {
             })
             .await;
             sessions.lock().await.remove(&process_key);
-        });
+        };
+        tokio::spawn(run_task.in_current_span());
         Ok(())
     }
 
@@ -443,6 +456,17 @@ impl CommandExecManager {
     }
 }
 
+#[tracing::instrument(
+    name = "app_server.command_exec.run",
+    level = "info",
+    skip_all,
+    fields(
+        rpc.method = "command/exec",
+        rpc.request_id = %params.request_id.request_id,
+        app_server.connection_id = %params.request_id.connection_id,
+        command_exec.outcome = tracing::field::Empty,
+    )
+)]
 async fn run_command(params: RunCommandParams) {
     let RunCommandParams {
         outgoing,
@@ -536,10 +560,18 @@ async fn run_command(params: RunCommandParams) {
         }
     };
 
-    let timeout_handle = tokio::spawn(async move {
+    let timeout_task = async move {
         tokio::time::sleep(Duration::from_millis(IO_DRAIN_TIMEOUT_MS)).await;
         let _ = stdio_timeout_tx.send(true);
-    });
+    };
+    let timeout_handle = tokio::spawn(timeout_task.in_current_span());
+
+    let outcome = match expiration_outcome {
+        Some(ExecExpirationOutcome::TimedOut) => "timed_out",
+        Some(ExecExpirationOutcome::Cancelled) => "cancelled",
+        None => "exited",
+    };
+    tracing::Span::current().record("command_exec.outcome", outcome);
 
     let stdout = stdout_handle.await.unwrap_or_default();
     let stderr = stderr_handle.await.unwrap_or_default();
@@ -568,7 +600,7 @@ fn spawn_process_output(params: SpawnProcessOutputParams) -> tokio::task::JoinHa
         stream_output,
         output_bytes_cap,
     } = params;
-    tokio::spawn(async move {
+    let output_task = async move {
         let mut buffer: Vec<u8> = Vec::new();
         let mut observed_num_bytes = 0usize;
         loop {
@@ -618,7 +650,8 @@ fn spawn_process_output(params: SpawnProcessOutputParams) -> tokio::task::JoinHa
             }
         }
         bytes_to_string_smart(&buffer)
-    })
+    };
+    tokio::spawn(output_task.in_current_span())
 }
 
 async fn handle_process_write(

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

@@ -1,3 +1,4 @@
+use crate::app_server_tracing::spawn_blocking_in_current_span;
 use crate::config_layer::config_layer_metadata_to_api;
 use crate::config_layer::config_layer_to_api;
 use crate::config_manager::ConfigManager;
@@ -34,7 +35,6 @@ use std::borrow::Cow;
 use std::path::Path;
 use std::path::PathBuf;
 use thiserror::Error;
-use tokio::task;
 use toml::Value as TomlValue;
 use toml_edit::Item as TomlItem;
 
@@ -405,7 +405,7 @@ async fn create_empty_user_layer(
 }
 
 async fn write_empty_user_config(write_path: PathBuf) -> Result<(), ConfigManagerError> {
-    task::spawn_blocking(move || write_atomically(&write_path, ""))
+    spawn_blocking_in_current_span(move || write_atomically(&write_path, ""))
         .await
         .map_err(|err| ConfigManagerError::anyhow("config persistence task panicked", err.into()))?
         .map_err(|err| ConfigManagerError::io("failed to create empty user config.toml", err))

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

@@ -13,6 +13,7 @@ use codex_app_server_protocol::ServerNotification;
 use codex_file_search as file_search;
 use tracing::warn;
 
+use crate::app_server_tracing::spawn_blocking_in_current_span;
 use crate::outgoing_message::OutgoingMessageSender;
 
 const MATCH_LIMIT: usize = 50;
@@ -38,7 +39,7 @@ pub(crate) async fn run_fuzzy_file_search(
     let threads = NonZero::new(threads.max(1)).expect("threads should be non-zero");
     let search_dirs: Vec<PathBuf> = roots.iter().map(PathBuf::from).collect();
 
-    let mut files = match tokio::task::spawn_blocking(move || {
+    let mut files = match spawn_blocking_in_current_span(move || {
         file_search::run(
             query.as_str(),
             search_dirs,

codex-rs/app-server/src/request_processors/apps_processor.rs

@@ -98,7 +98,7 @@ impl AppsRequestProcessor {
         let mcp_manager = self.thread_manager.mcp_manager();
         let plugins_manager = self.thread_manager.plugins_manager();
         let shutdown_token = self.shutdown_token.child_token();
-        tokio::spawn(async move {
+        let apps_list_task = async move {
             tokio::select! {
                 _ = shutdown_token.cancelled() => {}
                 _ = Self::apps_list_task(
@@ -112,7 +112,8 @@ impl AppsRequestProcessor {
                     installed_start,
                 ) => {}
             }
-        });
+        };
+        tokio::spawn(apps_list_task.in_current_span());
         Ok(None)
     }
 
@@ -214,7 +215,7 @@ impl AppsRequestProcessor {
 
         let accessible_config = config.clone();
         let accessible_tx = tx.clone();
-        tokio::spawn(async move {
+        let load_accessible_task = async move {
             let result = connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager(
                 &accessible_config,
                 force_refetch,
@@ -224,11 +225,12 @@ impl AppsRequestProcessor {
             .await
             .map_err(|err| format!("failed to load accessible apps: {err}"));
             let _ = accessible_tx.send(AppListLoadResult::Accessible(result));
-        });
+        };
+        tokio::spawn(load_accessible_task.in_current_span());
 
         let all_config = config.clone();
         let all_plugin_apps = plugin_apps.clone();
-        tokio::spawn(async move {
+        let load_all_task = async move {
             let result = connectors::list_all_connectors_with_options(
                 &all_config,
                 force_refetch,
@@ -237,7 +239,8 @@ impl AppsRequestProcessor {
             .await
             .map_err(|err| format!("failed to list apps: {err}"));
             let _ = tx.send(AppListLoadResult::Directory(result));
-        });
+        };
+        tokio::spawn(load_all_task.in_current_span());
 
         let app_list_deadline = tokio::time::Instant::now() + APP_LIST_LOAD_TIMEOUT;
         let mut accessible_loaded = false;

codex-rs/app-server/src/request_processors/external_agent_session_import.rs

@@ -154,10 +154,12 @@ impl ExternalAgentSessionImporter {
         session: ExternalAgentSessionMigration,
     ) -> Result<Option<PendingSessionImport>, String> {
         let codex_home = self.codex_home.clone();
-        tokio::task::spawn_blocking(move || prepare_validated_session_import(&codex_home, session))
-            .await
-            .map_err(|err| format!("external agent session preparation task failed: {err}"))?
-            .map_err(|err| format!("failed to prepare external agent session: {err}"))
+        crate::app_server_tracing::spawn_blocking_in_current_span(move || {
+            prepare_validated_session_import(&codex_home, session)
+        })
+        .await
+        .map_err(|err| format!("external agent session preparation task failed: {err}"))?
+        .map_err(|err| format!("failed to prepare external agent session: {err}"))
     }
 
     async fn persist_session(

codex-rs/app-server/src/request_processors/feedback_processor.rs

@@ -222,7 +222,7 @@ impl FeedbackRequestProcessor {
 
         let session_source = self.thread_manager.session_source();
 
-        let upload_result = tokio::task::spawn_blocking(move || {
+        let upload_result = crate::app_server_tracing::spawn_blocking_in_current_span(move || {
             let tags = (!upload_tags.is_empty()).then_some(&upload_tags);
             snapshot.upload_feedback(FeedbackUploadOptions {
                 classification: &classification,

codex-rs/app-server/src/request_processors/marketplace_processor.rs

@@ -77,7 +77,7 @@ impl MarketplaceRequestProcessor {
         let MarketplaceUpgradeParams { marketplace_name } = params;
         let plugins_input = config.plugins_config_input();
 
-        let outcome = tokio::task::spawn_blocking(move || {
+        let outcome = crate::app_server_tracing::spawn_blocking_in_current_span(move || {
             plugins_manager.upgrade_configured_marketplaces_for_config(
                 &plugins_input,
                 marketplace_name.as_deref(),

codex-rs/app-server/src/request_processors/mcp_processor.rs

@@ -262,7 +262,7 @@ impl McpRequestProcessor {
             }
         };
 
-        tokio::spawn(async move {
+        let status_task = async move {
             Self::list_mcp_server_status_task(
                 outgoing,
                 request,
@@ -273,7 +273,8 @@ impl McpRequestProcessor {
                 codex_apps_tools_cache,
             )
             .await;
-        });
+        };
+        tokio::spawn(status_task.in_current_span());
         Ok(())
     }
 
@@ -399,10 +400,11 @@ impl McpRequestProcessor {
             let (_, thread) = self.load_thread(&thread_id).await?;
             let request_id = request_id.clone();
 
-            tokio::spawn(async move {
+            let read_task = async move {
                 let result = thread.read_mcp_resource(&server, &uri).await;
                 Self::send_mcp_resource_read_response(outgoing, request_id, result).await;
-            });
+            };
+            tokio::spawn(read_task.in_current_span());
             return Ok(());
         }
 
@@ -419,7 +421,7 @@ impl McpRequestProcessor {
             McpRuntimeContext::new(Arc::clone(&environment_manager), config.cwd.to_path_buf());
         let request_id = request_id.clone();
 
-        tokio::spawn(async move {
+        let read_task = async move {
             let result = read_mcp_resource_without_thread(
                 &mcp_config,
                 auth.as_ref(),
@@ -431,7 +433,8 @@ impl McpRequestProcessor {
             .await
             .and_then(|result| serde_json::to_value(result).map_err(anyhow::Error::from));
             Self::send_mcp_resource_read_response(outgoing, request_id, result).await;
-        });
+        };
+        tokio::spawn(read_task.in_current_span());
         Ok(())
     }
 
@@ -463,14 +466,15 @@ impl McpRequestProcessor {
         let meta = with_mcp_tool_call_thread_id_meta(params.meta, &thread_id);
         let request_id = request_id.clone();
 
-        tokio::spawn(async move {
+        let tool_call_task = async move {
             let result = thread
                 .call_mcp_tool(&params.server, &params.tool, params.arguments, meta)
                 .await
                 .map(McpServerToolCallResponse::from)
                 .map_err(|error| internal_error(format!("{error:#}")));
             outgoing.send_result(request_id, result).await;
-        });
+        };
+        tokio::spawn(tool_call_task.in_current_span());
         Ok(())
     }
 }

codex-rs/app-server/src/request_processors/plugins.rs

@@ -595,7 +595,7 @@ impl PluginRequestProcessor {
             let plugins_manager_for_marketplace_listing = plugins_manager.clone();
             let roots_for_marketplace_listing = roots.clone();
             let shared_plugin_ids_by_local_path = load_shared_plugin_ids_by_local_path(&config)?;
-            match tokio::task::spawn_blocking(move || {
+            match crate::app_server_tracing::spawn_blocking_in_current_span(move || {
                 let outcome = plugins_manager_for_marketplace_listing
                     .list_marketplaces_for_config(
                         &config_for_marketplace_listing,
@@ -876,7 +876,7 @@ impl PluginRequestProcessor {
     > {
         let config_for_marketplace_listing = plugins_input.clone();
         let shared_plugin_ids_by_local_path = load_shared_plugin_ids_by_local_path(config)?;
-        match tokio::task::spawn_blocking(move || {
+        match crate::app_server_tracing::spawn_blocking_in_current_span(move || {
             let outcome = plugins_manager.list_marketplaces_for_config(
                 &config_for_marketplace_listing,
                 &roots,