From bba219255210888d53ae74ee6d34067ff0178e10 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Thu, 9 Jul 2026 04:13:31 +0000 Subject: [PATCH] Propagate RPC spans into bounded app-server tasks --- codex-rs/app-server/src/app_server_tracing.rs | 9 ++++ .../app-server/src/bespoke_event_handling.rs | 31 +++++++----- codex-rs/app-server/src/command_exec.rs | 49 ++++++++++++++++--- .../app-server/src/config_manager_service.rs | 4 +- codex-rs/app-server/src/fuzzy_file_search.rs | 3 +- .../src/request_processors/apps_processor.rs | 15 +++--- .../external_agent_session_import.rs | 10 ++-- .../request_processors/feedback_processor.rs | 2 +- .../marketplace_processor.rs | 2 +- .../src/request_processors/mcp_processor.rs | 20 +++++--- .../src/request_processors/plugins.rs | 4 +- 11 files changed, 104 insertions(+), 45 deletions(-) diff --git a/codex-rs/app-server/src/app_server_tracing.rs b/codex-rs/app-server/src/app_server_tracing.rs index f8c17a3acb..7ff9c24b00 100644 --- a/codex-rs/app-server/src/app_server_tracing.rs +++ b/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(function: F) -> tokio::task::JoinHandle +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, diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 829aa3f126..59616e3479 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/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) => { diff --git a/codex-rs/app-server/src/command_exec.rs b/codex-rs/app-server/src/command_exec.rs index d7f0b90d5c..987890d96e 100644 --- a/codex-rs/app-server/src/command_exec.rs +++ b/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 = 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( diff --git a/codex-rs/app-server/src/config_manager_service.rs b/codex-rs/app-server/src/config_manager_service.rs index 94980bf8e0..920a8bba66 100644 --- a/codex-rs/app-server/src/config_manager_service.rs +++ b/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)) diff --git a/codex-rs/app-server/src/fuzzy_file_search.rs b/codex-rs/app-server/src/fuzzy_file_search.rs index f8cd61e3ad..d7b2b7c23e 100644 --- a/codex-rs/app-server/src/fuzzy_file_search.rs +++ b/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 = 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, diff --git a/codex-rs/app-server/src/request_processors/apps_processor.rs b/codex-rs/app-server/src/request_processors/apps_processor.rs index 5ff7ec40db..6f7ef8ed54 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor.rs +++ b/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; diff --git a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs index 6ec89401ec..a9b9ebcd59 100644 --- a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs +++ b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs @@ -154,10 +154,12 @@ impl ExternalAgentSessionImporter { session: ExternalAgentSessionMigration, ) -> Result, 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( diff --git a/codex-rs/app-server/src/request_processors/feedback_processor.rs b/codex-rs/app-server/src/request_processors/feedback_processor.rs index 5f488f72c3..d52faee73a 100644 --- a/codex-rs/app-server/src/request_processors/feedback_processor.rs +++ b/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, diff --git a/codex-rs/app-server/src/request_processors/marketplace_processor.rs b/codex-rs/app-server/src/request_processors/marketplace_processor.rs index cba73c4ff1..ac8f7166e5 100644 --- a/codex-rs/app-server/src/request_processors/marketplace_processor.rs +++ b/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(), diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index ac2125c2cf..60041e39c0 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/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(¶ms.server, ¶ms.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(()) } } diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index 688cfae917..8d672eaab5 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/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,