diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index c63103161f..dae06b7236 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -769,6 +769,14 @@ pub(crate) async fn apply_bespoke_event_handling( server_name: request.server_name.clone(), request: request_body, }; + let pause_incremented = + match conversation.increment_out_of_band_elicitation_count().await { + Ok(_) => true, + Err(err) => { + error!("failed to pause timeout accounting for MCP elicitation: {err}"); + false + } + }; let (pending_request_id, rx) = outgoing .send_request(ServerRequestPayload::McpServerElicitationRequest(params)) .await; @@ -781,6 +789,7 @@ pub(crate) async fn apply_bespoke_event_handling( conversation, thread_state, permission_guard, + pause_incremented, ) .await; }); @@ -1772,8 +1781,14 @@ async fn on_mcp_server_elicitation_response( conversation: Arc, thread_state: Arc>, permission_guard: ThreadWatchActiveGuard, + pause_incremented: bool, ) { let response = receiver.await; + if pause_incremented + && let Err(err) = conversation.decrement_out_of_band_elicitation_count().await + { + error!("failed to resume timeout accounting after MCP elicitation: {err}"); + } resolve_server_request_on_thread_listener(&thread_state, pending_request_id).await; drop(permission_guard); let response = mcp_server_elicitation_response_from_client_result(response); diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 25d95e30ea..144cf6f738 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -218,6 +218,10 @@ impl Session { }; } + let _pending_elicitation = turn_context + .turn_metadata_state + .track_pending_mcp_elicitation(); + let (tx_response, rx_response) = oneshot::channel(); let prev_entry = { let mut active = self.active_turn.lock().await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d43724e5db..b20e024ebe 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1145,6 +1145,10 @@ impl Session { self.out_of_band_elicitation_paused.subscribe() } + pub(crate) fn is_out_of_band_elicitation_paused(&self) -> bool { + *self.out_of_band_elicitation_paused.borrow() + } + pub(crate) fn set_out_of_band_elicitation_pause_state(&self, paused: bool) { self.out_of_band_elicitation_paused.send_replace(paused); } diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index ad746a1de2..986cbb3053 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -184,7 +184,7 @@ pub(super) async fn handle_runtime_response( max_output_tokens: Option, started_at: std::time::Instant, ) -> Result { - let script_status = format_script_status(&response); + let script_status = format_script_status(exec, &response); match response { RuntimeResponse::Yielded { content_items, .. } => { @@ -228,8 +228,14 @@ fn sanitize_runtime_image_detail(turn: &TurnContext, items: &mut [FunctionCallOu sanitize_image_detail_items(can_request_original_image_detail(&turn.model_info), items); } -fn format_script_status(response: &RuntimeResponse) -> String { +fn format_script_status(exec: &ExecContext, response: &RuntimeResponse) -> String { match response { + RuntimeResponse::Yielded { cell_id, .. } + if exec.session.is_out_of_band_elicitation_paused() + || exec.turn.turn_metadata_state.has_pending_mcp_elicitation() => + { + format!("Script running with cell ID {cell_id}\nWaiting for user input") + } RuntimeResponse::Yielded { cell_id, .. } => { format!("Script running with cell ID {cell_id}") } diff --git a/codex-rs/core/src/turn_metadata.rs b/codex-rs/core/src/turn_metadata.rs index 0a1f781b38..15be4295e2 100644 --- a/codex-rs/core/src/turn_metadata.rs +++ b/codex-rs/core/src/turn_metadata.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use std::sync::Mutex; use std::sync::RwLock; use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use serde_json::Value; @@ -99,9 +100,21 @@ pub(crate) struct TurnMetadataState { turn_started_at_unix_ms: Arc>>, responsesapi_client_metadata: Arc>>, user_input_requested_during_turn: Arc, + pending_mcp_elicitations: Arc, enrichment_task: Arc>>>, } +pub(crate) struct PendingMcpElicitationGuard { + count: Arc, +} + +impl Drop for PendingMcpElicitationGuard { + fn drop(&mut self) { + let previous = self.count.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "pending MCP elicitation count underflowed"); + } +} + impl TurnMetadataState { #[allow(clippy::too_many_arguments)] pub(crate) fn new( @@ -142,6 +155,7 @@ impl TurnMetadataState { turn_started_at_unix_ms: Arc::new(RwLock::new(None)), responsesapi_client_metadata: Arc::new(RwLock::new(BTreeMap::new())), user_input_requested_during_turn: Arc::new(AtomicBool::new(false)), + pending_mcp_elicitations: Arc::new(AtomicUsize::new(0)), enrichment_task: Arc::new(Mutex::new(None)), } } @@ -203,6 +217,17 @@ impl TurnMetadataState { .store(true, Ordering::Relaxed); } + pub(crate) fn track_pending_mcp_elicitation(&self) -> PendingMcpElicitationGuard { + self.pending_mcp_elicitations.fetch_add(1, Ordering::AcqRel); + PendingMcpElicitationGuard { + count: Arc::clone(&self.pending_mcp_elicitations), + } + } + + pub(crate) fn has_pending_mcp_elicitation(&self) -> bool { + self.pending_mcp_elicitations.load(Ordering::Acquire) > 0 + } + pub(crate) fn set_responsesapi_client_metadata( &self, responsesapi_client_metadata: HashMap, diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 26c041d8f7..721d65e914 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -3,6 +3,7 @@ use anyhow::Result; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_config::types::AppToolApproval; use codex_config::types::McpServerConfig; use codex_config::types::McpServerTransportConfig; use codex_core::config::Config; @@ -12,6 +13,7 @@ use codex_features::CurrentTimeSource; use codex_features::Feature; use codex_login::CodexAuth; use codex_models_manager::bundled_models_response; +use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::WebSearchMode; use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem; use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; @@ -479,6 +481,125 @@ async fn run_code_mode_turn_with_rmcp_config( Ok((test, second_mock)) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_reports_pending_mcp_elicitation_on_yield_and_wait() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let rmcp_test_server_bin = stdio_server_bin()?; + let mut builder = test_codex().with_config(move |config| { + let _ = config.features.enable(Feature::CodeMode); + let _ = config.features.enable(Feature::ToolCallMcpElicitation); + config.approvals_reviewer = ApprovalsReviewer::User; + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + "rmcp".to_string(), + McpServerConfig { + auth: Default::default(), + transport: McpServerTransportConfig::Stdio { + command: rmcp_test_server_bin, + args: Vec::new(), + env: None, + env_vars: Vec::new(), + cwd: None, + }, + environment_id: "local".to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: Some(Duration::from_secs(10)), + tool_timeout_sec: None, + default_tools_approval_mode: Some(AppToolApproval::Prompt), + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), + }, + ); + config + .mcp_servers + .set(servers) + .expect("test mcp servers should accept any configuration"); + }); + let test = builder.build(&server).await?; + wait_for_mcp_server(&test.codex, "rmcp").await?; + + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_custom_tool_call( + "call-1", + "exec", + r#"// @exec: {"yield_time_ms": 1000} +await tools.mcp__rmcp__sync({});"#, + ), + ev_completed("resp-1"), + ]), + ) + .await; + let initial_yield = responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-2"), + responses::ev_function_call( + "call-2", + "wait", + &serde_json::to_string(&serde_json::json!({ + "cell_id": "1", + "yield_time_ms": 10, + }))?, + ), + ev_completed("resp-2"), + ]), + ) + .await; + let later_wait = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-1", "waiting"), + ev_completed("resp-3"), + ]), + ) + .await; + + test.submit_turn_with_approval_and_permission_profile( + "request confirmation from code mode", + AskForApproval::OnRequest, + PermissionProfile::Disabled, + ) + .await?; + + let initial_items = custom_tool_output_items(&initial_yield.single_request(), "call-1"); + assert_eq!(initial_items.len(), 1); + assert_regex_match( + concat!( + r"(?s)\A", + r"Script running with cell ID 1\n", + r"Waiting for user input\n", + r"Wall time \d+\.\d seconds\nOutput:\n\z" + ), + text_item(&initial_items, /*index*/ 0), + ); + + let wait_items = function_tool_output_items(&later_wait.single_request(), "call-2"); + assert_eq!(wait_items.len(), 1); + assert_regex_match( + concat!( + r"(?s)\A", + r"Script running with cell ID 1\n", + r"Waiting for user input\n", + r"Wall time \d+\.\d seconds\nOutput:\n\z" + ), + text_item(&wait_items, /*index*/ 0), + ); + + Ok(()) +} + #[cfg_attr(windows, ignore = "no exec_command on Windows")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_can_return_exec_command_output() -> Result<()> {