diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 64abfaac3..addd5ea83 100644 --- a/codex-rs/Cargo.lock +++ b/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", diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 24f5155c1..a69dfab96 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/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 } diff --git a/codex-rs/app-server/src/app_server_tracing.rs b/codex-rs/app-server/src/app_server_tracing.rs index fd32be87c..d1555e28e 100644 --- a/codex-rs/app-server/src/app_server_tracing.rs +++ b/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, diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 0ab906620..fdd05b450 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/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>>>, fuzzy_search_sessions: Arc>>, + background_tasks: TaskTracker, feedback: CodexFeedback, log_db: Option, } @@ -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, + 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 { + self.outgoing.request_trace_context(request_id).await + } + + async fn submit_core_op( + &self, + request_id: &ConnectionRequestId, + thread: &CodexThread, + op: Op, + ) -> CodexResult { + 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, experimental_raw_events: bool, + request_trace: Option, ) { 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) -> 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( diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index c31236f1a..6110fb52a 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/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 = diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 5616dcef3..ffd8eecbf 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/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)"); } }); diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 973a50431..3b85cc77c 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/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::(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::(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( + outgoing: Arc, + request_context: RequestContext, + request_fut: F, + ) where + F: Future, + { + outgoing + .register_request_context(request_context.clone()) + .await; + request_fut.instrument(request_context.span()).await; + } + pub(crate) fn thread_created_receiver(&self) -> broadcast::Receiver { 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; diff --git a/codex-rs/app-server/src/message_processor/tracing_tests.rs b/codex-rs/app-server/src/message_processor/tracing_tests.rs new file mode 100644 index 000000000..622aa4023 --- /dev/null +++ b/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 = 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> = OnceLock::new(); + GUARD.get_or_init(|| tokio::sync::Mutex::new(())) +} + +struct TracingHarness { + _server: MockServer, + _codex_home: TempDir, + processor: MessageProcessor, + outgoing_rx: mpsc::Receiver, + session: ConnectionSessionState, + tracing: &'static TestTracing, +} + +impl TracingHarness { + async fn new() -> Result { + 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(&mut self, request: ClientRequest, trace: Option) -> 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, + ) -> 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 { + 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, +) -> ( + MessageProcessor, + mpsc::Receiver, +) { + 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::>() + .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( + outgoing_rx: &mut mpsc::Receiver, + 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, +) { + 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(tracing: &TestTracing, predicate: F) -> Vec +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(()) +} diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index d615e5151..baf4b65a0 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/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, +} + +impl RequestContext { + pub(crate) fn new( + request_id: ConnectionRequestId, + span: Span, + parent_trace: Option, + ) -> Self { + Self { + request_id, + span, + parent_trace, + } + } + + pub(crate) fn request_trace(&self) -> Option { + 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, request_id_to_callback: Mutex>, + /// 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>, } #[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 { + 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 { + 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, + 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, + 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::(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::(4); @@ -775,6 +907,40 @@ mod tests { } } + #[tokio::test] + async fn connection_closed_clears_registered_request_contexts() { + let (tx, _rx) = mpsc::channel::(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::(4); diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 2747cd070..d7070981c 100644 --- a/codex-rs/core/src/codex.rs +++ b/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, pub(crate) session: Arc, + // 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>; + /// 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, + pub(crate) models_manager: Arc, + pub(crate) skills_manager: Arc, + pub(crate) plugins_manager: Arc, + pub(crate) mcp_manager: Arc, + pub(crate) file_watcher: Arc, + pub(crate) conversation_history: InitialHistory, + pub(crate) session_source: SessionSource, + pub(crate) agent_control: AgentControl, + pub(crate) dynamic_tools: Vec, + pub(crate) persist_extended_history: bool, + pub(crate) metrics_service_name: Option, + pub(crate) inherited_shell_snapshot: Option>, + pub(crate) parent_trace: Option, +} + 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, - models_manager: Arc, - skills_manager: Arc, - plugins_manager: Arc, - mcp_manager: Arc, - file_watcher: Arc, - conversation_history: InitialHistory, - session_source: SessionSource, - agent_control: AgentControl, - dynamic_tools: Vec, - persist_extended_history: bool, - metrics_service_name: Option, - inherited_shell_snapshot: Option>, - ) -> CodexResult { + pub(crate) async fn spawn(args: CodexSpawnArgs) -> CodexResult { + 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 { + 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 { + self.submit_with_trace(op, None).await + } + + pub async fn submit_with_trace( + &self, + op: Op, + trace: Option, + ) -> CodexResult { 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 { 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. diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index f353febd0..3219a1033 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/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(); diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 0265d810b..07881b153 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/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, ) -> ( diff --git a/codex-rs/core/src/codex_tests_guardian.rs b/codex-rs/core/src/codex_tests_guardian.rs index 4f8bfea9e..b0a16fc02 100644 --- a/codex-rs/core/src/codex_tests_guardian.rs +++ b/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"); diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index b33a66c3e..7a243f9f7 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/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, + ) -> CodexResult { + self.codex.submit_with_trace(op, trace).await + } + pub async fn steer_input( &self, input: Vec, diff --git a/codex-rs/core/src/memories/tests.rs b/codex-rs/core/src/memories/tests.rs index 8914beaac..d32564aad 100644 --- a/codex-rs/core/src/memories/tests.rs +++ b/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 { diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 5ee47bf74..405388502 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/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, + pub submit_failed: Vec, + pub timed_out: Vec, +} + +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, persist_extended_history: bool, metrics_service_name: Option, + parent_trace: Option, ) -> CodexResult { 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, + parent_trace: Option, ) -> CodexResult { 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, persist_extended_history: bool, + parent_trace: Option, ) -> CodexResult { 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::>() + }; + + 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::>(); + 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, ) -> CodexResult { 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, persist_extended_history: bool, metrics_service_name: Option, + parent_trace: Option, ) -> CodexResult { 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, inherited_shell_snapshot: Option>, + parent_trace: Option, ) -> CodexResult { 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()); + } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents.rs b/codex-rs/core/src/tools/handlers/multi_agents.rs index 9c9393a18..38594b502 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents.rs +++ b/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"); diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index ffb37a0a5..b2aaec777 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -202,6 +202,7 @@ impl TestCodexBuilder { config.clone(), path, auth_manager, + None, )) .await? } diff --git a/codex-rs/core/tests/suite/compact_resume_fork.rs b/codex-rs/core/tests/suite/compact_resume_fork.rs index 14120bad5..079c49797 100644 --- a/codex-rs/core/tests/suite/compact_resume_fork.rs +++ b/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 { - 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 diff --git a/codex-rs/core/tests/suite/fork_thread.rs b/codex-rs/core/tests/suite/fork_thread.rs index 96457065d..96fef440c 100644 --- a/codex-rs/core/tests/suite/fork_thread.rs +++ b/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"); diff --git a/codex-rs/core/tests/suite/permissions_messages.rs b/codex-rs/core/tests/suite/permissions_messages.rs index 8dfe2b5e6..ceaf65aa6 100644 --- a/codex-rs/core/tests/suite/permissions_messages.rs +++ b/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 diff --git a/codex-rs/core/tests/suite/resume_warning.rs b/codex-rs/core/tests/suite/resume_warning.rs index fcf2bf8e0..b1e6a694c 100644 --- a/codex-rs/core/tests/suite/resume_warning.rs +++ b/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"); diff --git a/codex-rs/core/tests/suite/unstable_features_warning.rs b/codex-rs/core/tests/suite/unstable_features_warning.rs index 9269ed262..94d7b5183 100644 --- a/codex-rs/core/tests/suite/unstable_features_warning.rs +++ b/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"); diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index d5b249fbc..3ad7c579c 100644 --- a/codex-rs/tui/src/app.rs +++ b/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) => {