From 10f9bfc4cb6cfca6c00df6c92fc50a086e253418 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 25 Jun 2026 11:04:38 +0100 Subject: [PATCH 1/6] Pin MCP runtime to model steps --- .../src/request_processors/mcp_processor.rs | 55 ++++---- codex-rs/core/src/codex_delegate.rs | 61 ++++++--- codex-rs/core/src/codex_delegate_tests.rs | 22 ++-- codex-rs/core/src/codex_thread.rs | 5 + codex-rs/core/src/mcp_skill_dependencies.rs | 11 +- codex-rs/core/src/mcp_tool_call.rs | 120 ++++++++---------- codex-rs/core/src/mcp_tool_call_tests.rs | 54 +++++--- codex-rs/core/src/session/handlers.rs | 4 +- codex-rs/core/src/session/mcp.rs | 105 ++++++++------- codex-rs/core/src/session/mcp_runtime.rs | 94 ++++++++++++++ codex-rs/core/src/session/mod.rs | 39 ++++-- codex-rs/core/src/session/session.rs | 46 +++---- codex-rs/core/src/session/step_context.rs | 5 + codex-rs/core/src/session/tests.rs | 46 ++++--- codex-rs/core/src/session/turn_context.rs | 3 +- codex-rs/core/src/state/service.rs | 38 +++++- codex-rs/core/src/tools/handlers/mcp.rs | 5 +- .../list_mcp_resource_templates.rs | 11 +- .../mcp_resource/list_mcp_resources.rs | 11 +- .../mcp_resource/read_mcp_resource.rs | 6 +- .../tools/handlers/request_plugin_install.rs | 18 ++- codex-rs/core/src/tools/router_tests.rs | 4 +- codex-rs/core/src/tools/spec_plan_tests.rs | 4 +- .../core/tests/suite/mcp_refresh_cleanup.rs | 13 +- 24 files changed, 480 insertions(+), 300 deletions(-) create mode 100644 codex-rs/core/src/session/mcp_runtime.rs diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index 18b49ede3bfc..d38db0d4a302 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -121,17 +121,11 @@ impl McpRequestProcessor { } = params; let auth = self.auth_manager.auth().await; - let (config, mcp_config) = match thread_id.as_deref() { + let (mcp_config, runtime_context) = match thread_id.as_deref() { Some(thread_id) => { let (_, thread) = self.load_thread(thread_id).await?; - let thread_config = thread.config().await; - let config = self - .config_manager - .load_latest_config_for_thread(thread_config.as_ref()) - .await - .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; - let mcp_config = thread.runtime_mcp_config(&config).await; - (config, mcp_config) + let runtime = thread.current_mcp_runtime(); + (runtime.config().clone(), runtime.runtime_context().clone()) } None => { let config = self.load_latest_config(/*fallback_cwd*/ None).await?; @@ -140,7 +134,11 @@ impl McpRequestProcessor { .mcp_manager() .runtime_config(&config) .await; - (config, mcp_config) + let runtime_context = McpRuntimeContext::new( + self.thread_manager.environment_manager(), + config.cwd.to_path_buf(), + ); + (mcp_config, runtime_context) } }; let effective_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref()); @@ -167,10 +165,6 @@ impl McpRequestProcessor { } }; - let runtime_context = McpRuntimeContext::new( - self.thread_manager.environment_manager(), - config.cwd.to_path_buf(), - ); let http_client = runtime_context .resolve_http_client(&name, server) .map_err(|err| { @@ -189,16 +183,16 @@ impl McpRequestProcessor { let handle = perform_oauth_login_return_url_with_http_client( &name, &url, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind(), + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, http_headers, env_http_headers, &resolved_scopes.scopes, server.oauth_client_id(), server.oauth_resource.as_deref(), timeout_secs, - config.mcp_oauth_callback_port, - config.mcp_oauth_callback_url.as_deref(), + mcp_config.mcp_oauth_callback_port, + mcp_config.mcp_oauth_callback_url.as_deref(), http_client, ) .await @@ -249,22 +243,25 @@ impl McpRequestProcessor { } None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None), }; - let mcp_config = match thread { - Some(thread) => thread.runtime_mcp_config(&config).await, + let auth = self.auth_manager.auth().await; + let (mcp_config, runtime_context) = match thread { + Some(thread) => { + let runtime = thread.current_mcp_runtime(); + (runtime.config().clone(), runtime.runtime_context().clone()) + } None => { - self.thread_manager + let mcp_config = self + .thread_manager .mcp_manager() .runtime_config(&config) - .await + .await; + let runtime_context = McpRuntimeContext::new( + self.thread_manager.environment_manager(), + config.cwd.to_path_buf(), + ); + (mcp_config, runtime_context) } }; - let auth = self.auth_manager.auth().await; - let environment_manager = self.thread_manager.environment_manager(); - // This status path has no turn-selected environment. Use config cwd - // as the local stdio fallback; named environment stdio MCPs must - // declare their own absolute cwd. - let runtime_context = - McpRuntimeContext::new(Arc::clone(&environment_manager), config.cwd.to_path_buf()); tokio::spawn(async move { Self::list_mcp_server_status_task( diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 6fbd195fce93..1872ae3ecc9b 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -40,6 +40,7 @@ use crate::guardian::spawn_approval_request_review; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_ACCEPT; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_ACCEPT_FOR_SESSION; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_DECLINE_SYNTHETIC; +use crate::mcp_tool_call::McpToolApprovalMetadata; use crate::mcp_tool_call::build_guardian_mcp_tool_review_request; use crate::mcp_tool_call::is_mcp_tool_approval_question_id; use crate::mcp_tool_call::lookup_mcp_tool_metadata; @@ -60,6 +61,12 @@ use codex_protocol::protocol::MultiAgentVersion; #[cfg(test)] use crate::session::completed_session_loop_termination; +#[derive(Clone)] +struct PendingMcpInvocation { + invocation: McpInvocation, + metadata: Option, +} + /// Start an interactive sub-Codex thread and return IO channels. /// /// The returned `events_rx` yields non-approval events emitted by the sub-agent. @@ -148,10 +155,10 @@ pub(crate) async fn run_codex_thread_interactive( let parent_session_clone = Arc::clone(&parent_session); let parent_ctx_clone = Arc::clone(&parent_ctx); let codex_for_events = Arc::clone(&codex); - // Cache delegated MCP invocations so guardian can recover the full tool call - // context when the later legacy RequestUserInput approval event only carries - // a call_id plus approval question metadata. - let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::::new())); + // Cache the child call's MCP metadata at begin time. The later legacy + // RequestUserInput approval event only carries a call_id and question metadata. + let pending_mcp_invocations = + Arc::new(Mutex::new(HashMap::::new())); tokio::spawn(async move { forward_events( codex_for_events, @@ -269,7 +276,7 @@ async fn forward_events( tx_sub: Sender, parent_session: Arc, parent_ctx: Arc, - pending_mcp_invocations: Arc>>, + pending_mcp_invocations: Arc>>, cancel_token: CancellationToken, ) { let cancelled = cancel_token.cancelled(); @@ -356,10 +363,33 @@ async fn forward_events( id, msg: EventMsg::McpToolCallBegin(event), } => { + // Runtime refreshes are published before a request step is captured, so + // the child runtime at call begin is the one executing this invocation. + // Cache its metadata now; the later approval event has only a call ID. + let metadata = if let Some(turn_context) = + codex.session.turn_context_for_sub_id(&id).await + { + let mcp = codex.session.services.latest_mcp_runtime(); + lookup_mcp_tool_metadata( + turn_context.as_ref(), + mcp.manager(), + &event.invocation.server, + &event.invocation.tool, + ) + .await + } else { + None + }; pending_mcp_invocations .lock() .await - .insert(event.call_id.clone(), event.invocation.clone()); + .insert( + event.call_id.clone(), + PendingMcpInvocation { + invocation: event.invocation.clone(), + metadata, + }, + ); if !forward_event_or_shutdown( &codex, &tx_sub, @@ -647,7 +677,7 @@ async fn handle_request_user_input( id: String, parent_session: &Arc, parent_ctx: &Arc, - pending_mcp_invocations: &Arc>>, + pending_mcp_invocations: &Arc>>, event: RequestUserInputEvent, cancel_token: &CancellationToken, ) { @@ -685,12 +715,12 @@ async fn handle_request_user_input( /// programmatically after running the guardian review. /// /// The RequestUserInput event only carries `call_id` plus approval question -/// metadata, so this helper joins it back to the cached `McpToolCallBegin` -/// invocation in order to rebuild the full guardian review request. +/// metadata, so this helper joins it back to the child runtime metadata cached at +/// `McpToolCallBegin` in order to rebuild the full guardian review request. async fn maybe_auto_review_mcp_request_user_input( parent_session: &Arc, parent_ctx: &Arc, - pending_mcp_invocations: &Arc>>, + pending_mcp_invocations: &Arc>>, event: &RequestUserInputEvent, cancel_token: &CancellationToken, ) -> Option { @@ -701,18 +731,13 @@ async fn maybe_auto_review_mcp_request_user_input( .questions .iter() .find(|question| is_mcp_tool_approval_question_id(&question.id))?; - let invocation = pending_mcp_invocations + let pending = pending_mcp_invocations .lock() .await .get(&event.call_id) .cloned()?; - let metadata = lookup_mcp_tool_metadata( - parent_session.as_ref(), - parent_ctx.as_ref(), - &invocation.server, - &invocation.tool, - ) - .await; + let invocation = pending.invocation; + let metadata = pending.metadata; let approvals_reviewer = mcp_approvals_reviewer(parent_ctx, &invocation.server, metadata.as_ref()); if !routes_approval_to_guardian_with_reviewer(parent_ctx, approvals_reviewer) { diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 6d51a0032b24..2791c887e6cb 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -411,10 +411,13 @@ async fn delegated_mcp_guardian_abort_returns_synthetic_decline_answer() { let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([( "call-1".to_string(), - McpInvocation { - server: "custom_server".to_string(), - tool: "dangerous_tool".to_string(), - arguments: None, + PendingMcpInvocation { + invocation: McpInvocation { + server: "custom_server".to_string(), + tool: "dangerous_tool".to_string(), + arguments: None, + }, + metadata: None, }, )]))); let cancel_token = CancellationToken::new(); @@ -460,10 +463,13 @@ async fn delegated_mcp_user_reviewer_returns_none_without_metadata() { crate::session::tests::make_session_and_context_with_rx().await; let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([( "call-1".to_string(), - McpInvocation { - server: CODEX_APPS_MCP_SERVER_NAME.to_string(), - tool: "dangerous_tool".to_string(), - arguments: None, + PendingMcpInvocation { + invocation: McpInvocation { + server: CODEX_APPS_MCP_SERVER_NAME.to_string(), + tool: "dangerous_tool".to_string(), + arguments: None, + }, + metadata: None, }, )]))); let cancel_token = CancellationToken::new(); diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 25f47ce45d8b..46b38af151ca 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -594,6 +594,11 @@ impl CodexThread { self.codex.session.runtime_mcp_config(config).await } + /// Returns the exact MCP config, environment bindings, and manager most recently published. + pub fn current_mcp_runtime(&self) -> Arc { + self.codex.session.services.latest_mcp_runtime() + } + pub fn multi_agent_version(&self) -> Option { self.codex.session.multi_agent_version() } diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index 0f8eabc08b78..91fed2d41226 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -200,15 +200,8 @@ pub(crate) async fn maybe_install_mcp_dependencies( warn!("failed to refresh MCP dependencies for mentioned skills: {err}"); return false; } - let refresh_servers = sess.runtime_mcp_servers(&refresh_config).await; - sess.refresh_mcp_servers_now( - turn_context, - refresh_servers, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind(), - elicitation_reviewer, - ) - .await; + sess.refresh_mcp_servers_now(turn_context, &refresh_config, elicitation_reviewer) + .await; true } diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index df2a375e351b..34dd9824a271 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -18,6 +18,7 @@ use crate::mcp_openai_file::rewrite_mcp_tool_arguments_for_openai_files; use crate::mcp_tool_approval_templates::RenderedMcpToolApprovalParam; use crate::mcp_tool_approval_templates::render_mcp_tool_approval_template; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::hook_names::HookToolName; use crate::tools::sandboxing::PermissionRequestPayload; @@ -35,6 +36,7 @@ use codex_features::Feature; use codex_hooks::PermissionRequestDecision; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; +use codex_mcp::McpConnectionManager; use codex_mcp::McpPermissionPromptAutoApproveContext; use codex_mcp::SandboxState; use codex_mcp::auth_elicitation_completed_result; @@ -111,13 +113,15 @@ const MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES: usize = DEFAULT_OUTPUT_BYTES_CAP; /// item lifecycle events to the `Session`. pub(crate) async fn handle_mcp_tool_call( sess: Arc, - turn_context: &Arc, + step_context: &Arc, call_id: String, server: String, tool_name: String, hook_tool_name: HookToolName, arguments: String, ) -> HandledMcpToolCall { + let turn_context = &step_context.turn; + let manager = step_context.mcp.manager(); // Parse the `arguments` as JSON. An empty string is OK, but invalid JSON // is not. let arguments_value = if arguments.trim().is_empty() { @@ -142,7 +146,7 @@ pub(crate) async fn handle_mcp_tool_call( }; let metadata = - lookup_mcp_tool_metadata(sess.as_ref(), turn_context.as_ref(), &server, &tool_name).await; + lookup_mcp_tool_metadata(turn_context.as_ref(), manager, &server, &tool_name).await; let item_metadata = McpToolCallItemMetadata::from_tool_metadata(&server, metadata.as_ref()); let app_tool_policy = if server == CODEX_APPS_MCP_SERVER_NAME { let annotations = metadata @@ -169,7 +173,6 @@ pub(crate) async fn handle_mcp_tool_call( } else if let Some(approval_mode) = { // Selected-plugin registrations are absent from config.toml and the legacy plugin manager, // so their resolved catalog entry is the authoritative source for tool approval policy. - let manager = sess.services.mcp_connection_manager.load(); manager .is_selected_plugin_mcp_server(&server) .then(|| manager.tool_approval_mode(&server, &tool_name)) @@ -226,7 +229,7 @@ pub(crate) async fn handle_mcp_tool_call( if let Some(decision) = maybe_request_mcp_tool_approval( &sess, - turn_context, + step_context, &call_id, &invocation, &hook_tool_name, @@ -241,7 +244,7 @@ pub(crate) async fn handle_mcp_tool_call( | McpToolApprovalDecision::AcceptAndRemember => { return handle_approved_mcp_tool_call( sess.as_ref(), - turn_context.as_ref(), + step_context.as_ref(), &call_id, invocation, metadata.as_ref(), @@ -298,7 +301,7 @@ pub(crate) async fn handle_mcp_tool_call( handle_approved_mcp_tool_call( sess.as_ref(), - turn_context.as_ref(), + step_context.as_ref(), &call_id, invocation, metadata.as_ref(), @@ -340,24 +343,21 @@ impl McpToolCallItemMetadata { async fn handle_approved_mcp_tool_call( sess: &Session, - turn_context: &TurnContext, + step_context: &StepContext, call_id: &str, invocation: McpInvocation, metadata: Option<&McpToolApprovalMetadata>, item_metadata: McpToolCallItemMetadata, ) -> HandledMcpToolCall { + let turn_context = step_context.turn.as_ref(); + let manager = step_context.mcp.manager(); let server = invocation.server.clone(); - maybe_mark_thread_memory_mode_polluted(sess, turn_context, &server).await; + maybe_mark_thread_memory_mode_polluted(sess, turn_context, manager, &server).await; let tool_name = invocation.tool.clone(); let arguments_value = invocation.arguments.clone(); let connector_id = metadata.and_then(|metadata| metadata.connector_id.as_deref()); let connector_name = metadata.and_then(|metadata| metadata.connector_name.as_deref()); - let server_origin = { - let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); - mcp_connection_manager - .server_origin(&server) - .map(str::to_string) - }; + let server_origin = manager.server_origin(&server).map(str::to_string); let start = Instant::now(); let rewrite = rewrite_mcp_tool_arguments_for_openai_files( @@ -380,7 +380,7 @@ async fn handle_approved_mcp_tool_call( build_mcp_tool_call_request_meta(turn_context, &server, call_id, metadata); execute_mcp_tool_call( sess, - turn_context, + step_context, call_id, &invocation, rewritten_arguments, @@ -420,7 +420,7 @@ async fn handle_approved_mcp_tool_call( truncate_mcp_tool_result_for_event(&result), ) .await; - maybe_track_codex_app_used(sess, turn_context, &server, &tool_name).await; + maybe_track_codex_app_used(sess, turn_context, manager, &server, &tool_name).await; let outcome = mcp_call_metric_outcome(&result); emit_mcp_call_metrics( @@ -547,17 +547,19 @@ fn truncate_str_to_char_boundary(value: &str, max_chars: usize) -> &str { async fn execute_mcp_tool_call( sess: &Session, - turn_context: &TurnContext, + step_context: &StepContext, call_id: &str, invocation: &McpInvocation, rewritten_arguments: Option, metadata: Option<&McpToolApprovalMetadata>, request_meta: Option, ) -> Result { + let turn_context = step_context.turn.as_ref(); + let manager = step_context.mcp.manager(); let request_meta = with_mcp_tool_call_thread_id_meta(request_meta, &sess.thread_id.to_string()); let request_meta = augment_mcp_tool_request_meta_with_sandbox_state( - sess, - turn_context, + step_context, + manager, &invocation.server, request_meta, ) @@ -568,7 +570,7 @@ async fn execute_mcp_tool_call( .rollout_thread_trace .start_mcp_call_trace(call_id); let request_meta = mcp_call_trace.add_request_meta(request_meta); - let result = sess + let result = manager .call_tool( &invocation.server, &invocation.tool, @@ -587,6 +589,7 @@ async fn execute_mcp_tool_call( Ok(maybe_request_codex_apps_auth_elicitation( sess, turn_context, + manager, call_id, &invocation.server, metadata, @@ -598,17 +601,13 @@ async fn execute_mcp_tool_call( async fn maybe_request_codex_apps_auth_elicitation( sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, call_id: &str, server: &str, metadata: Option<&McpToolApprovalMetadata>, result: CallToolResult, ) -> CallToolResult { - if !sess - .services - .mcp_connection_manager - .load_full() - .is_host_owned_codex_apps_server(server) - { + if !manager.is_host_owned_codex_apps_server(server) { return result; } @@ -666,15 +665,16 @@ async fn maybe_request_codex_apps_auth_elicitation( return result; } - refresh_codex_apps_after_connector_auth(sess, turn_context).await; + refresh_codex_apps_after_connector_auth(sess, turn_context, manager).await; auth_elicitation_completed_result(&plan.auth_failure, result.meta) } -async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: &TurnContext) { - let mcp_tools_result = { - let manager = sess.services.mcp_connection_manager.load_full(); - manager.hard_refresh_codex_apps_tools_cache().await - }; +async fn refresh_codex_apps_after_connector_auth( + sess: &Session, + turn_context: &TurnContext, + manager: &McpConnectionManager, +) { + let mcp_tools_result = manager.hard_refresh_codex_apps_tools_cache().await; match mcp_tools_result { Ok(mcp_tools) => { @@ -692,13 +692,13 @@ async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: & } async fn augment_mcp_tool_request_meta_with_sandbox_state( - sess: &Session, - turn_context: &TurnContext, + step_context: &StepContext, + manager: &McpConnectionManager, server: &str, mut meta: Option, ) -> anyhow::Result> { - let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); - let supports_sandbox_state_meta = mcp_connection_manager + let turn_context = step_context.turn.as_ref(); + let supports_sandbox_state_meta = manager .server_supports_sandbox_state_meta_capability(server) .await .unwrap_or(false); @@ -706,18 +706,18 @@ async fn augment_mcp_tool_request_meta_with_sandbox_state( return Ok(meta); } - let server_environment_id = mcp_connection_manager + let server_environment_id = manager .server_environment_id(server) .unwrap_or(codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID); - let Some(sandbox_cwd) = sandbox_cwd_for_mcp_server(turn_context, server_environment_id) else { + let Some(sandbox_cwd) = sandbox_cwd_for_mcp_server(step_context, server_environment_id) else { return Ok(meta); }; let permission_profile = turn_context.permission_profile(); let sandbox_state = serde_json::to_value(SandboxState { permission_profile, - codex_linux_sandbox_exe: turn_context.config.codex_linux_sandbox_exe.clone(), + codex_linux_sandbox_exe: step_context.mcp.config().codex_linux_sandbox_exe.clone(), sandbox_cwd, - use_legacy_landlock: turn_context.config.features.use_legacy_landlock(), + use_legacy_landlock: step_context.mcp.config().use_legacy_landlock, })?; match meta.as_mut() { @@ -741,8 +741,8 @@ async fn augment_mcp_tool_request_meta_with_sandbox_state( Ok(meta) } -fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str) -> Option { - if let Some(environment) = turn_context +fn sandbox_cwd_for_mcp_server(step_context: &StepContext, environment_id: &str) -> Option { + if let Some(environment) = step_context .environments .turn_environments .iter() @@ -753,7 +753,7 @@ fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str) if environment_id == codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID { #[allow(deprecated)] - return Some(PathUri::from_abs_path(&turn_context.cwd)); + return Some(PathUri::from_abs_path(&step_context.turn.cwd)); } None @@ -762,16 +762,13 @@ fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str) async fn maybe_mark_thread_memory_mode_polluted( sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, server: &str, ) { if !turn_context.config.memories.disable_on_external_context { return; } - let pollutes_memory = sess - .services - .mcp_connection_manager - .load_full() - .server_pollutes_memory(server); + let pollutes_memory = manager.server_pollutes_memory(server); if !pollutes_memory { return; } @@ -937,13 +934,14 @@ struct McpAppUsageMetadata { async fn maybe_track_codex_app_used( sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, server: &str, tool_name: &str, ) { if server != CODEX_APPS_MCP_SERVER_NAME { return; } - let metadata = lookup_mcp_app_usage_metadata(sess, server, tool_name).await; + let metadata = lookup_mcp_app_usage_metadata(manager, server, tool_name).await; let (connector_id, app_name) = metadata .map(|metadata| (metadata.connector_id, metadata.app_name)) .unwrap_or((None, None)); @@ -983,6 +981,7 @@ enum McpToolApprovalDecision { Cancel, } +#[derive(Clone)] pub(crate) struct McpToolApprovalMetadata { annotations: Option, connector_id: Option, @@ -1176,13 +1175,15 @@ fn mcp_tool_approval_prompt_options( async fn maybe_request_mcp_tool_approval( sess: &Arc, - turn_context: &Arc, + step_context: &Arc, call_id: &str, invocation: &McpInvocation, hook_tool_name: &HookToolName, metadata: Option<&McpToolApprovalMetadata>, approval_mode: AppToolApproval, ) -> Option { + let turn_context = &step_context.turn; + let manager = step_context.mcp.manager(); let approvals_reviewer = mcp_approvals_reviewer(turn_context, &invocation.server, metadata); if mcp_permission_prompt_is_auto_approved( turn_context.approval_policy.value(), @@ -1201,12 +1202,7 @@ async fn maybe_request_mcp_tool_approval( } let session_approval_key = session_mcp_tool_approval_key(invocation, metadata, approval_mode); - let persistent_approval_key = if sess - .services - .mcp_connection_manager - .load() - .is_selected_plugin_mcp_server(&invocation.server) - { + let persistent_approval_key = if manager.is_selected_plugin_mcp_server(&invocation.server) { None } else { persistent_mcp_tool_approval_key(invocation, metadata, approval_mode) @@ -1450,12 +1446,11 @@ async fn mcp_tool_approval_decision_from_guardian( } pub(crate) async fn lookup_mcp_tool_metadata( - sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, server: &str, tool_name: &str, ) -> Option { - let manager = sess.services.mcp_connection_manager.load_full(); let plugin_id = manager .plugin_id_for_mcp_server_name(server) .map(str::to_string); @@ -1562,16 +1557,11 @@ fn get_mcp_app_resource_uri( } async fn lookup_mcp_app_usage_metadata( - sess: &Session, + manager: &McpConnectionManager, server: &str, tool_name: &str, ) -> Option { - let tools = sess - .services - .mcp_connection_manager - .load_full() - .list_all_tools() - .await; + let tools = manager.list_all_tools().await; tools.into_iter().find_map(|tool_info| { if tool_info.server_name == server && tool_info.tool.name == tool_name { diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 387987c9f649..2340d3084337 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::config::ConfigBuilder; use crate::config::ManagedFeatures; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::tests::make_session_and_context_with_rx; use crate::session::turn_context::TurnEnvironment; @@ -157,10 +158,12 @@ async fn execute_mcp_tool_call_records_replayable_correlation() -> anyhow::Resul }) }); assert!(dispatch_trace.is_enabled()); + let turn_context = Arc::new(turn_context); + let step_context = StepContext::for_test(Arc::clone(&turn_context)); let result = execute_mcp_tool_call( &session, - &turn_context, + step_context.as_ref(), "mcp-call", &McpInvocation { server: "docs".to_string(), @@ -1147,7 +1150,8 @@ async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Resul /*shell*/ None, )); - let sandbox_cwd = sandbox_cwd_for_mcp_server(&turn_context, "remote"); + let step_context = StepContext::for_test(Arc::new(turn_context)); + let sandbox_cwd = sandbox_cwd_for_mcp_server(&step_context, "remote"); assert_eq!(sandbox_cwd, Some(secondary_cwd)); Ok(()) @@ -1157,7 +1161,8 @@ async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Resul async fn mcp_sandbox_cwd_is_none_for_unselected_server_environment() -> anyhow::Result<()> { let (_, turn_context) = make_session_and_context().await; - let sandbox_cwd = sandbox_cwd_for_mcp_server(&turn_context, "remote"); + let step_context = StepContext::for_test(Arc::new(turn_context)); + let sandbox_cwd = sandbox_cwd_for_mcp_server(&step_context, "remote"); assert_eq!(sandbox_cwd, None); Ok(()) @@ -1370,7 +1375,10 @@ fn codex_apps_auth_failure_metadata() -> McpToolApprovalMetadata { ) } -async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: &TurnContext) { +async fn host_owned_codex_apps_manager( + session: &Session, + turn_context: &TurnContext, +) -> Arc { let auth = session.services.auth_manager.auth().await; let startup_cancellation_token = CancellationToken::new(); startup_cancellation_token.cancel(); @@ -1409,22 +1417,20 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: /*elicitation_reviewer*/ None, ) .await; - session - .services - .mcp_connection_manager - .store(Arc::new(manager)); + Arc::new(manager) } #[tokio::test] async fn codex_apps_auth_elicitation_feature_disabled_returns_original_result() { let (session, turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let result = codex_apps_auth_failure_result(); let metadata = codex_apps_auth_failure_metadata(); let returned = maybe_request_codex_apps_auth_elicitation( &session, &turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1445,10 +1451,12 @@ async fn codex_apps_auth_elicitation_non_host_owned_server_returns_original_resu Arc::make_mut(&mut turn_context.config).features = ManagedFeatures::from(features); let result = codex_apps_auth_failure_result(); let metadata = codex_apps_auth_failure_metadata(); + let manager = session.services.latest_mcp_runtime().manager_arc(); let returned = maybe_request_codex_apps_auth_elicitation( &session, turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1463,7 +1471,7 @@ async fn codex_apps_auth_elicitation_non_host_owned_server_returns_original_resu #[tokio::test] async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_result() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let mut features = Features::with_defaults(); features.enable(Feature::AuthElicitation); let turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); @@ -1478,6 +1486,7 @@ async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_resul let returned = maybe_request_codex_apps_auth_elicitation( &session, turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1492,7 +1501,7 @@ async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_resul #[tokio::test] async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_result() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let mut features = Features::with_defaults(); features.enable(Feature::AuthElicitation); let turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); @@ -1513,6 +1522,7 @@ async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_resu let returned = maybe_request_codex_apps_auth_elicitation( &session, turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1527,7 +1537,7 @@ async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_resu #[tokio::test] async fn codex_apps_auth_elicitation_feature_enabled_requests_elicitation() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; *session.active_turn.lock().await = Some(ActiveTurn::default()); let mut features = Features::with_defaults(); features.enable(Feature::AuthElicitation); @@ -1541,10 +1551,12 @@ async fn codex_apps_auth_elicitation_feature_enabled_requests_elicitation() { let request_task = tokio::spawn({ let session = Arc::clone(&session); let turn_context = Arc::clone(&turn_context); + let manager = Arc::clone(&manager); async move { maybe_request_codex_apps_auth_elicitation( &session, &turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -2465,7 +2477,7 @@ async fn approve_mode_skips_when_annotations_do_not_require_approval() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-1", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2541,7 +2553,7 @@ async fn guardian_mode_skips_auto_when_annotations_do_not_require_approval() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-guardian", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2600,7 +2612,7 @@ async fn permission_request_hook_allows_mcp_tool_call() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-mcp-hook", &invocation, &HookToolName::new("mcp__memory__create_entities"), @@ -2662,7 +2674,7 @@ async fn permission_request_hook_uses_hook_tool_name_without_metadata() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-mcp-hook-no-metadata", &invocation, &HookToolName::new("mcp__memory__create_entities"), @@ -2744,7 +2756,7 @@ async fn permission_request_hook_runs_after_remembered_mcp_approval() { let turn_context = Arc::new(turn_context); let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-mcp-remembered", &invocation, &HookToolName::new("mcp__memory__create_entities"), @@ -2827,7 +2839,7 @@ async fn guardian_mode_mcp_denial_returns_rationale_message() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-guardian-deny", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2887,7 +2899,7 @@ async fn prompt_mode_waits_for_approval_when_annotations_do_not_require_approval tokio::spawn(async move { maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-prompt", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2945,7 +2957,7 @@ async fn full_access_mode_skips_mcp_tool_approval_for_all_approval_modes() { ] { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-2", &invocation, &HookToolName::new("mcp__test__tool"), @@ -3034,7 +3046,7 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() { let turn_context = Arc::new(turn_context); let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-3", &invocation, &HookToolName::new("mcp__test__tool"), diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 5c68893e7ff9..75fbca4c8607 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -598,8 +598,8 @@ async fn shutdown_session_runtime(sess: &Arc) { warn!("failed to shutdown code mode session: {err}"); } sess.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .shutdown() .await; sess.guardian_review_session.shutdown().await; diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index c060ef090b68..9b15bd16879d 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -105,8 +105,8 @@ impl Session { ) -> McpServerElicitationOutcome { if self .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .elicitations_auto_deny() { return McpServerElicitationOutcome { @@ -201,8 +201,8 @@ impl Session { } self.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .resolve_elicitation(server_name, id, response) .await } @@ -213,8 +213,8 @@ impl Session { params: Option, ) -> anyhow::Result { self.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .list_resources(server, params) .await } @@ -225,8 +225,8 @@ impl Session { params: Option, ) -> anyhow::Result { self.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .list_resource_templates(server, params) .await } @@ -237,8 +237,8 @@ impl Session { params: ReadResourceRequestParams, ) -> anyhow::Result { self.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .read_resource(server, params) .await } @@ -251,8 +251,8 @@ impl Session { meta: Option, ) -> anyhow::Result { self.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .call_tool(server, tool, arguments, meta) .await } @@ -260,17 +260,13 @@ impl Session { async fn refresh_mcp_servers_inner( &self, turn_context: &TurnContext, - mcp_servers: HashMap, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + refresh_config: &Config, elicitation_reviewer: Option, ) { let auth = self.services.auth_manager.auth().await; - let config = self.get_config().await; - let mcp_config = self.runtime_mcp_config(config.as_ref()).await; + let mcp_config = Arc::new(self.runtime_mcp_config(refresh_config).await); let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); - let mcp_servers = - effective_mcp_servers_from_configured(mcp_servers, &mcp_config, auth.as_ref()); + let mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref()); let environment_manager = self.services.turn_environments.environment_manager(); // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment cwd // values can be used without falling back to the legacy host cwd. @@ -286,34 +282,35 @@ impl Session { let mcp_runtime_context = McpRuntimeContext::new(environment_manager, cwd); let auth_statuses = compute_auth_statuses( mcp_servers.iter(), - store_mode, - keyring_backend_kind, + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, auth.as_ref(), &mcp_runtime_context, ) .await; let mcp_startup_cancellation_token = { let mut guard = self.services.mcp_startup_cancellation_token.lock().await; - guard.cancel(); + // The previous runtime owns the old token and may still be serving an in-flight step. + // Its manager cancels that token when the last runtime handle is dropped. let cancellation_token = CancellationToken::new(); *guard = cancellation_token.clone(); cancellation_token }; let refreshed_manager = McpConnectionManager::new( &mcp_servers, - store_mode, - keyring_backend_kind, + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, auth_statuses, &turn_context.approval_policy, turn_context.sub_id.clone(), self.get_tx_event(), mcp_startup_cancellation_token, turn_context.permission_profile(), - mcp_runtime_context, - config.codex_home.to_path_buf(), + mcp_runtime_context.clone(), + mcp_config.codex_home.clone(), codex_apps_tools_cache_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, - mcp_config.client_elicitation_capability, + mcp_config.client_elicitation_capability.clone(), self.services .supports_openai_form_elicitation .load(std::sync::atomic::Ordering::Relaxed), @@ -323,14 +320,12 @@ impl Session { ) .await; { - let current_manager = self.services.mcp_connection_manager.load_full(); - refreshed_manager.set_elicitations_auto_deny(current_manager.elicitations_auto_deny()); + let current_manager = self.services.latest_mcp_runtime(); + refreshed_manager + .set_elicitations_auto_deny(current_manager.manager().elicitations_auto_deny()); } - let superseded_manager = self - .services - .mcp_connection_manager - .swap(Arc::new(refreshed_manager)); - superseded_manager.shutdown().await; + self.services + .publish_mcp_runtime(mcp_config, mcp_runtime_context, refreshed_manager); } pub(crate) async fn refresh_mcp_servers_if_requested( @@ -375,14 +370,26 @@ impl Session { } }; - self.refresh_mcp_servers_inner( - turn_context, - mcp_servers, - store_mode, - keyring_backend_kind, - elicitation_reviewer, - ) - .await; + let mut refresh_config = self.get_config().await.as_ref().clone(); + if let Err(err) = refresh_config.mcp_servers.set(mcp_servers) { + warn!("failed to apply MCP server refresh config: {err}"); + return; + } + refresh_config.mcp_oauth_credentials_store_mode = store_mode; + let secret_auth_storage_enabled = match keyring_backend_kind { + AuthKeyringBackendKind::Direct => false, + AuthKeyringBackendKind::Secrets => true, + }; + if let Err(err) = refresh_config + .features + .set_enabled(Feature::SecretAuthStorage, secret_auth_storage_enabled) + { + warn!("failed to apply MCP auth keyring backend refresh config: {err}"); + return; + } + + self.refresh_mcp_servers_inner(turn_context, &refresh_config, elicitation_reviewer) + .await; } pub(crate) async fn set_openai_form_elicitation_support( @@ -416,19 +423,11 @@ impl Session { pub(crate) async fn refresh_mcp_servers_now( &self, turn_context: &TurnContext, - mcp_servers: HashMap, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + refresh_config: &Config, elicitation_reviewer: Option, ) { - self.refresh_mcp_servers_inner( - turn_context, - mcp_servers, - store_mode, - keyring_backend_kind, - elicitation_reviewer, - ) - .await; + self.refresh_mcp_servers_inner(turn_context, refresh_config, elicitation_reviewer) + .await; } #[cfg(test)] diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs new file mode 100644 index 000000000000..dc5ea011ca32 --- /dev/null +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -0,0 +1,94 @@ +use std::fmt; +use std::sync::Arc; + +use codex_mcp::McpConfig; +use codex_mcp::McpConnectionManager; +use codex_mcp::McpRuntimeContext; + +/// MCP config, exact environment bindings, and manager used by one model request. +pub struct McpRuntimeSnapshot { + config: Arc, + manager: Arc, + runtime_context: McpRuntimeContext, +} + +impl McpRuntimeSnapshot { + pub(crate) fn new( + config: Arc, + manager: Arc, + runtime_context: McpRuntimeContext, + ) -> Self { + Self { + config, + manager, + runtime_context, + } + } + + pub fn config(&self) -> &McpConfig { + self.config.as_ref() + } + + pub fn manager(&self) -> &McpConnectionManager { + self.manager.as_ref() + } + + pub(crate) fn manager_arc(&self) -> Arc { + Arc::clone(&self.manager) + } + + pub fn runtime_context(&self) -> &McpRuntimeContext { + &self.runtime_context + } + + #[cfg(test)] + pub(crate) fn new_uninitialized_for_test(config: &crate::config::Config) -> Arc { + use codex_exec_server::EnvironmentManager; + use codex_features::Feature; + use codex_mcp::ResolvedMcpCatalog; + use rmcp::model::ElicitationCapability; + + let mcp_config = McpConfig { + chatgpt_base_url: config.chatgpt_base_url.clone(), + apps_mcp_product_sku: config.apps_mcp_product_sku.clone(), + codex_home: config.codex_home.to_path_buf(), + mcp_oauth_credentials_store_mode: config.mcp_oauth_credentials_store_mode, + auth_keyring_backend_kind: config.auth_keyring_backend_kind(), + mcp_oauth_callback_port: config.mcp_oauth_callback_port, + mcp_oauth_callback_url: config.mcp_oauth_callback_url.clone(), + skill_mcp_dependency_install_enabled: config + .features + .enabled(Feature::SkillMcpDependencyInstall), + approval_policy: config.permissions.approval_policy.clone(), + codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(), + use_legacy_landlock: config.features.use_legacy_landlock(), + apps_enabled: config.features.enabled(Feature::Apps), + prefix_mcp_tool_names: config.prefix_mcp_tool_names(), + client_elicitation_capability: ElicitationCapability::default(), + mcp_server_catalog: ResolvedMcpCatalog::default(), + connector_snapshot: codex_connectors::ConnectorSnapshot::default(), + }; + let manager = McpConnectionManager::new_uninitialized_with_permission_profile( + &config.permissions.approval_policy, + config.permissions.permission_profile(), + config.prefix_mcp_tool_names(), + ); + let runtime_context = McpRuntimeContext::new( + Arc::new(EnvironmentManager::default_for_tests()), + config.cwd.to_path_buf(), + ); + Arc::new(Self::new( + Arc::new(mcp_config), + Arc::new(manager), + runtime_context, + )) + } +} + +impl fmt::Debug for McpRuntimeSnapshot { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("McpRuntimeSnapshot") + .finish_non_exhaustive() + } +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b3a0f1434c51..45dc1e68804e 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -154,15 +154,12 @@ use codex_utils_path_uri::PathUri; use futures::future::BoxFuture; use futures::future::Shared; use futures::prelude::*; -use rmcp::model::ElicitationCapability; -use rmcp::model::FormElicitationCapability; use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; -use rmcp::model::UrlElicitationCapability; use serde_json::Value; use tokio::sync::Mutex; use tokio::sync::RwLock; @@ -210,6 +207,7 @@ mod handlers; mod inject; mod input_queue; mod mcp; +mod mcp_runtime; pub(crate) mod multi_agents; mod review; mod rollout_budget; @@ -231,6 +229,7 @@ use self::handlers::submission_loop; pub(crate) use self::input_queue::InputQueueActivity; pub(crate) use self::input_queue::TurnInput; pub(crate) use self::input_queue::TurnInputQueue; +pub use self::mcp_runtime::McpRuntimeSnapshot; use self::review::spawn_review_thread; use self::session::AppServerClientMetadata; use self::session::Session; @@ -336,7 +335,7 @@ use codex_core_plugins::RecommendedPluginCandidatesInput; use codex_git_utils::get_git_repo_root; use codex_mcp::McpConfig; use codex_mcp::compute_auth_statuses; -use codex_mcp::effective_mcp_servers_from_configured; +use codex_mcp::effective_mcp_servers; use codex_otel::SessionTelemetry; use codex_otel::THREAD_STARTED_METRIC; use codex_otel::TelemetryAuthMode; @@ -819,8 +818,11 @@ impl Codex { ..Default::default() }) .await?; - let mcp_connection_manager = self.session.services.mcp_connection_manager.load_full(); - mcp_connection_manager.set_elicitations_auto_deny(mcp_elicitations_auto_deny); + self.session + .services + .latest_mcp_runtime() + .manager() + .set_elicitations_auto_deny(mcp_elicitations_auto_deny); Ok(()) } @@ -2833,6 +2835,7 @@ impl Session { &environments.captured_environment_availability(), ) .await; + let mcp = self.services.latest_mcp_runtime(); let extra_skill_sources = self .services .thread_extension_data @@ -2853,6 +2856,7 @@ impl Session { environments, selected_capability_roots, skills, + mcp, loaded_agents_md, )) } @@ -3118,6 +3122,17 @@ impl Session { &self, turn_context: &TurnContext, world_state: &WorldState, + ) -> Vec { + let mcp = self.services.latest_mcp_runtime(); + self.build_initial_context_with_world_state_and_mcp(turn_context, world_state, &mcp) + .await + } + + async fn build_initial_context_with_world_state_and_mcp( + &self, + turn_context: &TurnContext, + world_state: &WorldState, + mcp: &McpRuntimeSnapshot, ) -> Vec { let mut developer_sections = Vec::::with_capacity(8); let mut contextual_user_sections = Vec::::with_capacity(2); @@ -3211,10 +3226,9 @@ impl Session { } } if turn_context.config.include_apps_instructions && turn_context.apps_enabled() { - let mcp_connection_manager = self.services.mcp_connection_manager.load_full(); let accessible_and_enabled_connectors = connectors::list_accessible_and_enabled_connectors_from_manager( - &mcp_connection_manager, + mcp.manager(), &turn_context.config, ) .await; @@ -3298,7 +3312,8 @@ impl Session { if turn_context.config.features.enabled(Feature::TokenBudget) && turn_context.model_context_window().is_some() { - let mcp_result = self + let mcp_result = mcp + .manager() .call_tool( "notes", "thread_hint", @@ -3513,7 +3528,11 @@ impl Session { // Full initial context resets the baseline; later turns persist only its changes. let (mut context_items, world_state_item) = if should_inject_full_context { let context_items = self - .build_initial_context_with_world_state(turn_context, world_state.as_ref()) + .build_initial_context_with_world_state_and_mcp( + turn_context, + world_state.as_ref(), + step_context.mcp.as_ref(), + ) .await; let snapshot = world_state.snapshot(); self.state diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 85daa021b58a..33b24a29a018 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -654,10 +654,11 @@ impl Session { let config_for_mcp = Arc::clone(&config); let mcp_manager_for_mcp = Arc::clone(&mcp_manager); let mcp_thread_init_for_startup = &mcp_thread_init; - let mcp_runtime_context_for_auth = McpRuntimeContext::new( + let mcp_runtime_context = McpRuntimeContext::new( Arc::clone(&environment_manager), session_configuration.cwd().to_path_buf(), ); + let mcp_runtime_context_for_auth = mcp_runtime_context.clone(); let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; let mcp_config = mcp_manager_for_mcp @@ -673,7 +674,13 @@ impl Session { &mcp_runtime_context_for_auth, ) .await; - (auth, mcp_servers, auth_statuses, tool_plugin_provenance) + ( + auth, + mcp_config, + mcp_servers, + auth_statuses, + tool_plugin_provenance, + ) } .instrument(info_span!( "session_init.auth_mcp", @@ -684,7 +691,7 @@ impl Session { let ( thread_persistence_result, state_db_ctx, - (auth, mcp_servers, auth_statuses, tool_plugin_provenance), + (auth, mcp_config, mcp_servers, auth_statuses, tool_plugin_provenance), ) = tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut); let mut live_thread_init = @@ -1020,6 +1027,7 @@ impl Session { // changing this to use Option or OnceCell, though the current // setup is straightforward enough and performs well. mcp_connection_manager, + mcp_runtime: arc_swap::ArcSwapOption::empty(), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -1153,14 +1161,6 @@ impl Session { sess.send_event_raw(event).await; } - let client_elicitation_capability = if config.features.enabled(Feature::AuthElicitation) { - ElicitationCapability { - form: Some(FormElicitationCapability::default()), - url: Some(UrlElicitationCapability::default()), - } - } else { - ElicitationCapability::default() - }; let mcp_startup_cancellation_token = { let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; cancel_guard.cancel(); @@ -1168,20 +1168,6 @@ impl Session { *cancel_guard = cancel_token.clone(); cancel_token }; - let mcp_runtime_context = { - let turn_environments = sess.services.turn_environments.snapshot().await; - // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment - // cwd values can be used without falling back to the session host cwd. - let cwd = turn_environments - .primary() - .and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok()) - .map(|cwd| cwd.to_path_buf()) - .unwrap_or_else(|| session_configuration.cwd().to_path_buf()); - McpRuntimeContext::new( - sess.services.turn_environments.environment_manager(), - cwd, - ) - }; let mcp_connection_manager = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, @@ -1192,11 +1178,11 @@ impl Session { tx_event.clone(), mcp_startup_cancellation_token, session_configuration.permission_profile(), - mcp_runtime_context, + mcp_runtime_context.clone(), config.codex_home.to_path_buf(), codex_apps_tools_cache_key(auth), config.prefix_mcp_tool_names(), - client_elicitation_capability, + mcp_config.client_elicitation_capability.clone(), sess.services .supports_openai_form_elicitation .load(std::sync::atomic::Ordering::Relaxed), @@ -1210,7 +1196,11 @@ impl Session { )) .await; sess.services - .install_mcp_connection_manager(mcp_connection_manager) + .install_mcp_connection_manager( + Arc::new(mcp_config), + mcp_runtime_context, + mcp_connection_manager, + ) .await?; sess.schedule_startup_prewarm(session_configuration.base_instructions.clone()) .await; diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index d13110e281c7..ccf94e289398 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use crate::agents_md::LoadedAgentsMd; use crate::environment_selection::TurnEnvironmentSnapshot; +use crate::session::McpRuntimeSnapshot; use crate::session::turn_context::TurnContext; use codex_core_skills::SkillsSnapshot; use codex_exec_server::ResolvedSelectedCapabilityRoot; @@ -15,6 +16,8 @@ pub(crate) struct StepContext { pub(crate) selected_capability_roots: Vec, /// Complete skill catalog and exact read routes captured for this step. pub(crate) skills: Arc, + /// The exact MCP config and manager used to advertise and execute tools for this step. + pub(crate) mcp: Arc, /// The canonical AGENTS.md value observed with this environment snapshot. pub(crate) loaded_agents_md: Option>, } @@ -25,6 +28,7 @@ impl StepContext { environments: TurnEnvironmentSnapshot, selected_capability_roots: Vec, skills: Arc, + mcp: Arc, loaded_agents_md: Option>, ) -> Self { Self { @@ -32,6 +36,7 @@ impl StepContext { environments, selected_capability_roots, skills, + mcp, loaded_agents_md, } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d055277a28e0..6444d2545698 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -193,10 +193,11 @@ impl StepContext { turn.model_context_window(), )); Arc::new(Self::new( - turn, + Arc::clone(&turn), environments, Vec::new(), skills, + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config), /*loaded_agents_md*/ None, )) } @@ -385,8 +386,8 @@ async fn request_mcp_server_elicitation_auto_accepts_when_auto_deny_is_enabled() let (session, turn_context, rx) = make_session_and_context_with_rx().await; session .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .set_elicitations_auto_deny(/*auto_deny*/ true); let response = session @@ -5369,14 +5370,11 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); + let mcp_runtime = + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); let services = SessionServices { - mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( - McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), - )), + mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), + mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -7447,14 +7445,11 @@ where /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); + let mcp_runtime = + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); let services = SessionServices { - mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( - McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), - )), + mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), + mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -7611,8 +7606,14 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( } #[tokio::test] -async fn refresh_mcp_servers_is_deferred_until_next_turn() { +async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { let (session, turn_context) = make_session_and_context().await; + let turn_context = Arc::new(turn_context); + let old_runtime = session.services.latest_mcp_runtime(); + let step_context = session + .capture_step_context(Arc::clone(&turn_context)) + .await; + assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime)); let old_token = session.mcp_startup_cancellation_token().await; assert!(!old_token.is_cancelled()); @@ -7643,7 +7644,7 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() { .refresh_mcp_servers_if_requested(&turn_context, /*elicitation_reviewer*/ None) .await; - assert!(old_token.is_cancelled()); + assert!(!old_token.is_cancelled()); assert!( session .pending_mcp_server_refresh_config @@ -7653,6 +7654,9 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() { ); let new_token = session.mcp_startup_cancellation_token().await; assert!(!new_token.is_cancelled()); + let new_runtime = session.services.latest_mcp_runtime(); + assert!(!Arc::ptr_eq(&old_runtime, &new_runtime)); + assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime)); } #[tokio::test] @@ -10113,8 +10117,8 @@ async fn fatal_tool_error_stops_turn_and_reports_error() { let tools = { session .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .list_all_tools() .await }; diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index a5ae45895750..5efbeb8ef396 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -686,7 +686,8 @@ impl Session { .unwrap_or_else(|| session_configuration.cwd().clone()); let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone()); { - let mcp_connection_manager = self.services.mcp_connection_manager.load_full(); + let mcp_runtime = self.services.latest_mcp_runtime(); + let mcp_connection_manager = mcp_runtime.manager(); mcp_connection_manager.set_approval_policy(&session_configuration.approval_policy); mcp_connection_manager .set_permission_profile(session_configuration.permission_profile()); diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index a685a239c206..b64a228404e6 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -15,6 +15,7 @@ use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejection; use crate::guardian::GuardianRejectionCircuitBreaker; use crate::mcp::McpManager; +use crate::session::McpRuntimeSnapshot; use crate::tools::code_mode::CodeModeService; use crate::tools::handlers::ToolSearchHandlerCache; use crate::tools::network_approval::NetworkApprovalService; @@ -31,7 +32,9 @@ use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; use codex_hooks::Hooks; use codex_login::AuthManager; +use codex_mcp::McpConfig; use codex_mcp::McpConnectionManager; +use codex_mcp::McpRuntimeContext; use codex_models_manager::manager::SharedModelsManager; use codex_otel::SessionTelemetry; use codex_protocol::capabilities::SelectedCapabilityRoot; @@ -45,8 +48,10 @@ use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { - /// The latest manager; callers retain an owned handle while performing MCP I/O. + /// Mirror of the latest manager for extension resource clients that predate runtime snapshots. pub(crate) mcp_connection_manager: Arc>, + /// The latest atomically published MCP config and manager pair. + pub(crate) mcp_runtime: ArcSwapOption, pub(crate) mcp_startup_cancellation_token: Mutex, pub(crate) unified_exec_manager: UnifiedExecProcessManager, #[cfg_attr(not(unix), allow(dead_code))] @@ -103,12 +108,33 @@ impl SessionServices { /// resolve through the session's manager while validation waits. pub(crate) async fn install_mcp_connection_manager( &self, + config: Arc, + runtime_context: McpRuntimeContext, manager: McpConnectionManager, ) -> Result<()> { - self.mcp_connection_manager.store(Arc::new(manager)); - self.mcp_connection_manager - .load_full() - .validate_required_servers() - .await + let runtime = self.publish_mcp_runtime(config, runtime_context, manager); + runtime.manager().validate_required_servers().await + } + + pub(crate) fn publish_mcp_runtime( + &self, + config: Arc, + runtime_context: McpRuntimeContext, + manager: McpConnectionManager, + ) -> Arc { + let manager = Arc::new(manager); + // Publish the manager for legacy resource clients first. Once the paired snapshot is + // visible, every model-scoped consumer observes this exact manager. + self.mcp_connection_manager.store(Arc::clone(&manager)); + let runtime = Arc::new(McpRuntimeSnapshot::new(config, manager, runtime_context)); + self.mcp_runtime.store(Some(Arc::clone(&runtime))); + runtime + } + + pub(crate) fn latest_mcp_runtime(&self) -> Arc { + let Some(runtime) = self.mcp_runtime.load_full() else { + unreachable!("MCP runtime must be installed before handling requests"); + }; + runtime } } diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index b92a2a5655be..c9533fd223ea 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -124,11 +124,12 @@ impl McpHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = Arc::clone(&step_context.turn); let payload = match payload { ToolPayload::Function { arguments } => arguments, @@ -143,7 +144,7 @@ impl McpHandler { // TODO(sayan): Use StepContext for MCP file arguments when MCP follows dynamic environments. let result = handle_mcp_tool_call( Arc::clone(&session), - &turn, + &step_context, call_id.clone(), self.tool_info.server_name.clone(), self.tool_info.tool.name.to_string(), diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs index ce762a0f021b..8c7f21fd17d3 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs @@ -53,11 +53,13 @@ impl ListMcpResourceTemplatesHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = std::sync::Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -89,7 +91,7 @@ impl ListMcpResourceTemplatesHandler { let params = cursor .clone() .map(|value| PaginatedRequestParams::default().with_cursor(Some(value))); - let result = session + let result = manager .list_resource_templates(&server_name, params) .await .map_err(|err| { @@ -108,10 +110,7 @@ impl ListMcpResourceTemplatesHandler { )); } - let templates = session - .services - .mcp_connection_manager - .load_full() + let templates = manager .list_all_resource_templates(|server_name| { model_can_access_mcp_server(turn.as_ref(), server_name) }) diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs index 564fb440fda8..f5e9bae7e3e7 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs @@ -53,11 +53,13 @@ impl ListMcpResourcesHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = std::sync::Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -89,7 +91,7 @@ impl ListMcpResourcesHandler { let params = cursor .clone() .map(|value| PaginatedRequestParams::default().with_cursor(Some(value))); - let result = session + let result = manager .list_resources(&server_name, params) .await .map_err(|err| { @@ -106,10 +108,7 @@ impl ListMcpResourcesHandler { )); } - let resources = session - .services - .mcp_connection_manager - .load_full() + let resources = manager .list_all_resources(|server_name| { model_can_access_mcp_server(turn.as_ref(), server_name) }) diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs index 3909011eb0a2..5e9b85241c8d 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs @@ -52,11 +52,13 @@ impl ReadMcpResourceHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = std::sync::Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -84,7 +86,7 @@ impl ReadMcpResourceHandler { let payload_result: Result = async { ensure_model_can_access_mcp_server(turn.as_ref(), &server)?; - let result = session + let result = manager .read_resource(&server, ReadResourceRequestParams::new(uri.clone())) .await .map_err(|err| { diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index ce7ce10da853..e44ed41fa4d8 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -1,4 +1,5 @@ use std::collections::HashSet; +use std::sync::Arc; use codex_analytics::PluginInstallRequestSource; use codex_analytics::PluginInstallRequested; @@ -94,10 +95,12 @@ impl RequestPluginInstallHandler { let ToolInvocation { payload, session, - turn, + step_context, call_id, .. } = invocation; + let turn = Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -227,7 +230,8 @@ impl RequestPluginInstallHandler { let auth = session.services.auth_manager.auth().await; let completed = if user_confirmed { - verify_request_plugin_install_completed(&session, &turn, &tool, auth.as_ref()).await + verify_request_plugin_install_completed(&session, &turn, manager, &tool, auth.as_ref()) + .await } else { false }; @@ -345,13 +349,14 @@ fn disabled_install_request(tool: &DiscoverableTool) -> ToolSuggestDisabledTool async fn verify_request_plugin_install_completed( session: &crate::session::session::Session, turn: &crate::session::turn_context::TurnContext, + manager: &codex_mcp::McpConnectionManager, tool: &DiscoverableTool, auth: Option<&codex_login::CodexAuth>, ) -> bool { match tool { DiscoverableTool::Connector(connector) => refresh_missing_requested_connectors( - session, turn, + manager, auth, std::slice::from_ref(&connector.id), connector.id.as_str(), @@ -370,8 +375,8 @@ async fn verify_request_plugin_install_completed( plugin.id.as_str(), ), refresh_missing_requested_connectors( - session, turn, + manager, auth, &plugin.app_connector_ids, plugin.id.as_str(), @@ -393,8 +398,8 @@ async fn verify_request_plugin_install_completed( session.services.plugins_manager.as_ref(), ); let _ = refresh_missing_requested_connectors( - session, turn, + manager, auth, &plugin.app_connector_ids, plugin.id.as_str(), @@ -435,8 +440,8 @@ fn is_remote_plugin_install_suggestion(plugin_id: &str) -> bool { } async fn refresh_missing_requested_connectors( - session: &crate::session::session::Session, turn: &crate::session::turn_context::TurnContext, + manager: &codex_mcp::McpConnectionManager, auth: Option<&codex_login::CodexAuth>, expected_connector_ids: &[String], tool_id: &str, @@ -445,7 +450,6 @@ async fn refresh_missing_requested_connectors( return Some(Vec::new()); } - let manager = session.services.mcp_connection_manager.load_full(); let mcp_tools = manager.list_all_tools().await; let accessible_connectors = connectors::with_app_enabled_state( connectors::accessible_connectors_from_mcp_tools(&mcp_tools), diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index db1209cc0357..6f399fd5fc85 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -110,8 +110,8 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow let step_context = StepContext::for_test(Arc::clone(&turn)); let mcp_tools = session .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .list_all_tools() .await; let router = ToolRouter::from_context( diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 54eabec83c6d..7342e5cfbed6 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -690,11 +690,13 @@ async fn environment_tools_follow_the_step_context() { turn.model_context_window(), )); turn.environments.turn_environments.clear(); + let turn = Arc::new(turn); let step_context = Arc::new(StepContext::new( - Arc::new(turn), + Arc::clone(&turn), environments, Vec::new(), skills, + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config), /*loaded_agents_md*/ None, )); diff --git a/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs index 7df18727b9f5..620b6c734c2c 100644 --- a/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs +++ b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs @@ -16,7 +16,7 @@ use core_test_support::test_codex::test_codex; use core_test_support::wait_for_mcp_server; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> { +async fn refresh_keeps_superseded_mcp_server_alive_for_in_flight_calls() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -84,7 +84,7 @@ async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> "sync", Some(serde_json::json!({ "barrier": barrier, - "sleep_after_ms": 30_000 + "sleep_after_ms": 300_000 })), /*meta*/ None, ) @@ -119,9 +119,16 @@ async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> let replacement_pid = wait_for_pid_file(&pid_file).await?; assert_ne!(replacement_pid, superseded_pid); + assert!(process_is_alive(&superseded_pid)?); + long_call.abort(); + assert!( + long_call + .await + .expect_err("call should be aborted") + .is_cancelled() + ); wait_for_process_exit(&superseded_pid).await?; assert!(process_is_alive(&replacement_pid)?); - assert!(long_call.await?.is_err()); fixture.codex.shutdown_and_wait().await?; wait_for_process_exit(&replacement_pid).await From e02a5d32a2f848236f6d3af6b4ebf52e6428fce2 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 25 Jun 2026 11:36:07 +0100 Subject: [PATCH 2/6] Unify executor plugin runtime projection --- codex-rs/Cargo.lock | 16 - codex-rs/Cargo.toml | 1 - .../codex-mcp/src/connection_manager_tests.rs | 95 ++++ codex-rs/codex-mcp/src/elicitation.rs | 12 +- codex-rs/core-plugins/src/executor_runtime.rs | 224 ++++++++++ codex-rs/core-plugins/src/lib.rs | 3 + codex-rs/core-plugins/src/provider.rs | 15 + codex-rs/ext/connectors/BUILD.bazel | 6 - codex-rs/ext/connectors/Cargo.toml | 22 - .../ext/connectors/src/executor_plugin.rs | 64 --- codex-rs/ext/connectors/src/lib.rs | 6 - codex-rs/ext/mcp/Cargo.toml | 8 +- codex-rs/ext/mcp/src/executor_plugin.rs | 75 ++-- .../ext/mcp/src/executor_plugin/provider.rs | 138 ------ .../mcp/src/executor_plugin/provider_tests.rs | 421 ------------------ 15 files changed, 380 insertions(+), 726 deletions(-) create mode 100644 codex-rs/core-plugins/src/executor_runtime.rs delete mode 100644 codex-rs/ext/connectors/BUILD.bazel delete mode 100644 codex-rs/ext/connectors/Cargo.toml delete mode 100644 codex-rs/ext/connectors/src/executor_plugin.rs delete mode 100644 codex-rs/ext/connectors/src/lib.rs delete mode 100644 codex-rs/ext/mcp/src/executor_plugin/provider.rs delete mode 100644 codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index eb151ebe703c..00c6ef08d56c 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2596,18 +2596,6 @@ dependencies = [ "urlencoding", ] -[[package]] -name = "codex-connectors-extension" -version = "0.0.0" -dependencies = [ - "codex-connectors", - "codex-core-plugins", - "codex-plugin", - "codex-utils-path-uri", - "serde_json", - "thiserror 2.0.18", -] - [[package]] name = "codex-context-fragments" version = "0.0.0" @@ -3361,14 +3349,10 @@ dependencies = [ "codex-features", "codex-login", "codex-mcp", - "codex-plugin", "codex-protocol", - "codex-utils-absolute-path", "codex-utils-path-uri", "pretty_assertions", - "serde_json", "tempfile", - "thiserror 2.0.18", "tokio", "tracing", ] diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index ef17cb9ce5f1..abf5e0411a11 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -48,7 +48,6 @@ members = [ "exec-server", "execpolicy", "execpolicy-legacy", - "ext/connectors", "ext/extension-api", "ext/goal", "ext/guardian", diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index d66474c7bb30..750938708ff6 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -310,6 +310,101 @@ async fn disabled_permissions_do_not_auto_accept_elicitation_with_requested_fiel ); } +#[tokio::test] +async fn elicitation_route_ids_are_unique_across_managers() { + let first_manager = ElicitationRequestManager::new( + AskForApproval::OnRequest, + PermissionProfile::default(), + /*reviewer*/ None, + ); + let second_manager = ElicitationRequestManager::new( + AskForApproval::OnRequest, + PermissionProfile::default(), + /*reviewer*/ None, + ); + let (first_tx, first_rx) = async_channel::bounded(1); + let (second_tx, second_rx) = async_channel::bounded(1); + let first_sender = first_manager.make_sender("server".to_string(), first_tx); + let second_sender = second_manager.make_sender("server".to_string(), second_tx); + + let first_task = tokio::spawn(first_sender( + NumberOrString::Number(1), + codex_rmcp_client::Elicitation::OpenAiForm { + meta: None, + message: "First?".to_string(), + requested_schema: serde_json::json!({}), + }, + )); + let second_task = tokio::spawn(second_sender( + NumberOrString::Number(1), + codex_rmcp_client::Elicitation::OpenAiForm { + meta: None, + message: "Second?".to_string(), + requested_schema: serde_json::json!({}), + }, + )); + + let EventMsg::ElicitationRequest(first_request) = first_rx + .recv() + .await + .expect("first elicitation should be emitted") + .msg + else { + panic!("expected first elicitation request"); + }; + let EventMsg::ElicitationRequest(second_request) = second_rx + .recv() + .await + .expect("second elicitation should be emitted") + .msg + else { + panic!("expected second elicitation request"); + }; + assert_ne!(first_request.id, second_request.id); + + let first_response = ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(serde_json::json!({"manager": "first"})), + meta: None, + }; + let second_response = ElicitationResponse { + action: ElicitationAction::Decline, + content: Some(serde_json::json!({"manager": "second"})), + meta: None, + }; + let first_id = match first_request.id { + codex_protocol::mcp::RequestId::String(value) => NumberOrString::String(Arc::from(value)), + codex_protocol::mcp::RequestId::Integer(value) => NumberOrString::Number(value), + }; + let second_id = match second_request.id { + codex_protocol::mcp::RequestId::String(value) => NumberOrString::String(Arc::from(value)), + codex_protocol::mcp::RequestId::Integer(value) => NumberOrString::Number(value), + }; + first_manager + .resolve("server".to_string(), first_id, first_response.clone()) + .await + .expect("first manager should resolve its routed request"); + second_manager + .resolve("server".to_string(), second_id, second_response.clone()) + .await + .expect("second manager should resolve its routed request"); + + assert_eq!( + first_task + .await + .expect("first elicitation task should join") + .expect("first elicitation should resolve"), + first_response + ); + assert_eq!( + second_task + .await + .expect("second elicitation task should join") + .expect("second elicitation should resolve"), + second_response + ); +} + #[test] fn test_normalize_tools_short_non_duplicated_names() { let tools = vec![ diff --git a/codex-rs/codex-mcp/src/elicitation.rs b/codex-rs/codex-mcp/src/elicitation.rs index d6a685b50203..799852437eec 100644 --- a/codex-rs/codex-mcp/src/elicitation.rs +++ b/codex-rs/codex-mcp/src/elicitation.rs @@ -8,6 +8,8 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use crate::mcp::McpPermissionPromptAutoApproveContext; use crate::mcp::mcp_permission_prompt_is_auto_approved; @@ -48,6 +50,8 @@ pub trait ElicitationReviewer: Send + Sync { pub type ElicitationReviewerHandle = Arc; +static NEXT_ELICITATION_ROUTE_ID: AtomicU64 = AtomicU64::new(1); + #[derive(Clone)] pub(crate) struct ElicitationRequestManager { requests: Arc>, @@ -55,6 +59,7 @@ pub(crate) struct ElicitationRequestManager { pub(crate) permission_profile: Arc>, auto_deny: Arc>, reviewer: Option, + route_id: u64, } impl ElicitationRequestManager { @@ -69,6 +74,7 @@ impl ElicitationRequestManager { permission_profile: Arc::new(StdMutex::new(permission_profile)), auto_deny: Arc::new(StdMutex::new(false)), reviewer, + route_id: NEXT_ELICITATION_ROUTE_ID.fetch_add(1, Ordering::Relaxed), } } @@ -110,6 +116,7 @@ impl ElicitationRequestManager { let permission_profile = self.permission_profile.clone(); let auto_deny = self.auto_deny.clone(); let reviewer = self.reviewer.clone(); + let route_id = self.route_id; Box::new(move |id, elicitation| { let elicitation_requests = elicitation_requests.clone(); let tx_event = tx_event.clone(); @@ -214,9 +221,10 @@ impl ElicitationRequestManager { }, }; let (tx, rx) = oneshot::channel(); + let response_id = RequestId::String(Arc::from(format!("{route_id}:{id}"))); { let mut lock = elicitation_requests.lock().await; - lock.insert((server_name.clone(), id.clone()), tx); + lock.insert((server_name.clone(), response_id.clone()), tx); } let _ = tx_event .send(Event { @@ -224,7 +232,7 @@ impl ElicitationRequestManager { msg: EventMsg::ElicitationRequest(ElicitationRequestEvent { turn_id: None, server_name, - id: match id.clone() { + id: match response_id { rmcp::model::NumberOrString::String(value) => { ProtocolRequestId::String(value.to_string()) } diff --git a/codex-rs/core-plugins/src/executor_runtime.rs b/codex-rs/core-plugins/src/executor_runtime.rs new file mode 100644 index 000000000000..c70750805643 --- /dev/null +++ b/codex-rs/core-plugins/src/executor_runtime.rs @@ -0,0 +1,224 @@ +use codex_config::McpServerConfig; +use codex_connectors::parse_plugin_app_config; +use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::ResolvedSelectedCapabilityRoot; +use codex_mcp::parse_executor_plugin_mcp_config; +use codex_plugin::AppDeclaration; +use codex_plugin::PluginResourceLocator; +use codex_plugin::ResolvedPlugin; +use codex_plugin::ResolvedPluginLocation; +use codex_plugin::manifest::PluginManifestMcpServers; +use codex_utils_path_uri::PathUri; +use codex_utils_path_uri::PathUriParseError; +use std::io; +use thiserror::Error; + +use crate::ExecutorPluginProvider; +use crate::ExecutorPluginProviderError; +use crate::ResolvedExecutorPlugin; + +const DEFAULT_MCP_CONFIG_FILE: &str = ".mcp.json"; + +/// MCP and connector declarations read from one exact executor binding. +#[derive(Clone, Debug)] +pub struct ExecutorPluginRuntime { + plugin: ResolvedPlugin, + mcp_servers: Vec<(String, McpServerConfig)>, + apps: Vec, +} + +/// Failure to project runtime capabilities from an executor plugin. +#[derive(Debug, Error)] +pub enum ExecutorPluginRuntimeError { + #[error(transparent)] + Resolve(#[from] ExecutorPluginProviderError), + #[error("failed to read MCP config for selected plugin `{plugin_id}` at `{path}`: {source}")] + ReadConfig { + plugin_id: String, + path: PathUri, + #[source] + source: io::Error, + }, + #[error( + "failed to resolve MCP config path `{relative_path}` below selected plugin `{plugin_id}` at `{root}`: {source}" + )] + InvalidConfigPath { + plugin_id: String, + root: PathUri, + relative_path: &'static str, + #[source] + source: PathUriParseError, + }, + #[error("failed to parse MCP config for selected plugin `{plugin_id}` at `{path}`: {source}")] + ParseConfig { + plugin_id: String, + path: PathUri, + #[source] + source: serde_json::Error, + }, + #[error("failed to read app config for selected plugin `{plugin_id}` at `{path}`: {source}")] + ReadAppConfig { + plugin_id: String, + path: PathUri, + #[source] + source: io::Error, + }, + #[error("failed to parse app config for selected plugin `{plugin_id}` at `{path}`: {source}")] + ParseAppConfig { + plugin_id: String, + path: PathUri, + #[source] + source: serde_json::Error, + }, +} + +impl ExecutorPluginRuntime { + /// Reads both runtime declaration files through the root's pinned filesystem. + /// + /// `Ok(None)` is intentionally not cacheable: the plugin manifest may appear + /// once a deferred executor finishes starting. + pub async fn project( + root: &ResolvedSelectedCapabilityRoot, + ) -> Result, ExecutorPluginRuntimeError> { + let Some(plugin) = ExecutorPluginProvider::resolve_pinned(root).await? else { + return Ok(None); + }; + let ResolvedPluginLocation::Environment { + root: plugin_root, .. + } = plugin.plugin().location(); + let mcp_servers = + load_from_file_system(plugin.plugin(), plugin_root, plugin.file_system()).await?; + let apps = match load_apps(&plugin).await { + Ok(apps) => apps, + Err(err) => { + tracing::warn!( + plugin = plugin.plugin().selected_root_id(), + error = %err, + "ignoring invalid executor plugin app declarations" + ); + Vec::new() + } + }; + Ok(Some(Self { + plugin: plugin.plugin().clone(), + mcp_servers, + apps, + })) + } + + pub fn plugin(&self) -> &ResolvedPlugin { + &self.plugin + } + + pub fn mcp_servers(&self) -> &[(String, McpServerConfig)] { + &self.mcp_servers + } + + pub fn apps(&self) -> &[AppDeclaration] { + &self.apps + } +} + +async fn load_from_file_system( + plugin: &ResolvedPlugin, + plugin_root: &PathUri, + file_system: &dyn ExecutorFileSystem, +) -> Result, ExecutorPluginRuntimeError> { + let ResolvedPluginLocation::Environment { environment_id, .. } = plugin.location(); + let plugin_id = plugin.selected_root_id(); + let (contents, config_path) = match plugin.manifest().paths.mcp_servers.as_ref() { + Some(PluginManifestMcpServers::Path(PluginResourceLocator::Environment { + path, .. + })) => { + ( + file_system + .read_file_text(path, /*sandbox*/ None) + .await + .map_err(|source| ExecutorPluginRuntimeError::ReadConfig { + plugin_id: plugin_id.to_string(), + path: path.clone(), + source, + })?, + path.clone(), + ) + } + Some(PluginManifestMcpServers::Object(object_config)) => { + let PluginResourceLocator::Environment { path, .. } = plugin.manifest_path(); + (object_config.clone(), path.clone()) + } + None => { + let config_path = plugin_root + .join(DEFAULT_MCP_CONFIG_FILE) + .map_err(|source| ExecutorPluginRuntimeError::InvalidConfigPath { + plugin_id: plugin_id.to_string(), + root: plugin_root.clone(), + relative_path: DEFAULT_MCP_CONFIG_FILE, + source, + })?; + let contents = match file_system + .read_file_text(&config_path, /*sandbox*/ None) + .await + { + Ok(contents) => contents, + Err(source) if source.kind() == io::ErrorKind::NotFound => { + return Ok(Vec::new()); + } + Err(source) => { + return Err(ExecutorPluginRuntimeError::ReadConfig { + plugin_id: plugin_id.to_string(), + path: config_path.clone(), + source, + }); + } + }; + (contents, config_path) + } + }; + let parsed = parse_executor_plugin_mcp_config(plugin_root, &contents, environment_id).map_err( + |source| ExecutorPluginRuntimeError::ParseConfig { + plugin_id: plugin_id.to_string(), + path: config_path, + source, + }, + )?; + + for error in parsed.errors { + tracing::warn!( + plugin = plugin_id, + server = error.name, + error = error.message, + "ignoring invalid executor plugin MCP server" + ); + } + + Ok(parsed.servers.into_iter().collect()) +} + +async fn load_apps( + plugin: &ResolvedExecutorPlugin, +) -> Result, ExecutorPluginRuntimeError> { + let resolved_plugin = plugin.plugin(); + let plugin_id = resolved_plugin.selected_root_id(); + let Some(PluginResourceLocator::Environment { + path: config_path, .. + }) = resolved_plugin.manifest().paths.apps.as_ref() + else { + return Ok(Vec::new()); + }; + let contents = plugin + .file_system() + .read_file_text(config_path, /*sandbox*/ None) + .await + .map_err(|source| ExecutorPluginRuntimeError::ReadAppConfig { + plugin_id: plugin_id.to_string(), + path: config_path.clone(), + source, + })?; + parse_plugin_app_config(&contents).map_err(|source| { + ExecutorPluginRuntimeError::ParseAppConfig { + plugin_id: plugin_id.to_string(), + path: config_path.clone(), + source, + } + }) +} diff --git a/codex-rs/core-plugins/src/lib.rs b/codex-rs/core-plugins/src/lib.rs index 7ee5ddd117ca..ad6d7ba352c4 100644 --- a/codex-rs/core-plugins/src/lib.rs +++ b/codex-rs/core-plugins/src/lib.rs @@ -1,5 +1,6 @@ mod app_mcp_routing; mod discoverable; +mod executor_runtime; pub mod installed_marketplaces; pub mod loader; mod manager; @@ -38,6 +39,8 @@ pub type PluginLoadOutcome = codex_plugin::PluginLoadOutcome Result, ExecutorPluginProviderError> { + let selected_root = resolved_root.selected_root(); + let plugin_root = selected_plugin_root(selected_root); + let file_system = resolved_root.file_system(); + let plugin = resolve_plugin_root(selected_root, plugin_root, file_system.as_ref()).await?; + Ok(plugin.map(|plugin| ResolvedExecutorPlugin { + plugin, + file_system, + })) + } } impl PluginProvider for ExecutorPluginProvider { diff --git a/codex-rs/ext/connectors/BUILD.bazel b/codex-rs/ext/connectors/BUILD.bazel deleted file mode 100644 index 304349b8a17c..000000000000 --- a/codex-rs/ext/connectors/BUILD.bazel +++ /dev/null @@ -1,6 +0,0 @@ -load("//:defs.bzl", "codex_rust_crate") - -codex_rust_crate( - name = "connectors", - crate_name = "codex_connectors_extension", -) diff --git a/codex-rs/ext/connectors/Cargo.toml b/codex-rs/ext/connectors/Cargo.toml deleted file mode 100644 index 1103ae4e2b34..000000000000 --- a/codex-rs/ext/connectors/Cargo.toml +++ /dev/null @@ -1,22 +0,0 @@ -[package] -edition.workspace = true -license.workspace = true -name = "codex-connectors-extension" -version.workspace = true - -[lib] -name = "codex_connectors_extension" -path = "src/lib.rs" -doctest = false -test = false - -[lints] -workspace = true - -[dependencies] -codex-connectors = { workspace = true } -codex-core-plugins = { workspace = true } -codex-plugin = { workspace = true } -codex-utils-path-uri = { workspace = true } -serde_json = { workspace = true } -thiserror = { workspace = true } diff --git a/codex-rs/ext/connectors/src/executor_plugin.rs b/codex-rs/ext/connectors/src/executor_plugin.rs deleted file mode 100644 index c27ae3ac6221..000000000000 --- a/codex-rs/ext/connectors/src/executor_plugin.rs +++ /dev/null @@ -1,64 +0,0 @@ -use codex_connectors::parse_plugin_app_config; -use codex_core_plugins::ResolvedExecutorPlugin; -use codex_plugin::AppDeclaration; -use codex_plugin::PluginResourceLocator; -use codex_utils_path_uri::PathUri; -use std::io; -use thiserror::Error; - -/// Loads connector declarations from a resolved plugin through its owning executor. -#[derive(Clone, Copy, Debug, Default)] -pub struct ExecutorPluginConnectorProvider; - -/// Failure to load connector declarations from an executor plugin. -#[derive(Debug, Error)] -pub enum ExecutorPluginConnectorProviderError { - #[error("failed to read app config for selected plugin `{plugin_id}` at `{path}`: {source}")] - ReadConfig { - plugin_id: String, - path: PathUri, - #[source] - source: io::Error, - }, - #[error("failed to parse app config for selected plugin `{plugin_id}` at `{path}`: {source}")] - ParseConfig { - plugin_id: String, - path: PathUri, - #[source] - source: serde_json::Error, - }, -} - -impl ExecutorPluginConnectorProvider { - /// Returns the connector declarations contributed by `plugin`. - pub async fn load( - &self, - plugin: &ResolvedExecutorPlugin, - ) -> Result, ExecutorPluginConnectorProviderError> { - let resolved_plugin = plugin.plugin(); - let plugin_id = resolved_plugin.selected_root_id(); - let Some(PluginResourceLocator::Environment { - path: config_path, .. - }) = resolved_plugin.manifest().paths.apps.as_ref() - else { - return Ok(Vec::new()); - }; - let contents = plugin - .file_system() - .read_file_text(config_path, /*sandbox*/ None) - .await - .map_err(|source| ExecutorPluginConnectorProviderError::ReadConfig { - plugin_id: plugin_id.to_string(), - path: config_path.clone(), - source, - })?; - - parse_plugin_app_config(&contents).map_err(|source| { - ExecutorPluginConnectorProviderError::ParseConfig { - plugin_id: plugin_id.to_string(), - path: config_path.clone(), - source, - } - }) - } -} diff --git a/codex-rs/ext/connectors/src/lib.rs b/codex-rs/ext/connectors/src/lib.rs deleted file mode 100644 index f60e5f916a39..000000000000 --- a/codex-rs/ext/connectors/src/lib.rs +++ /dev/null @@ -1,6 +0,0 @@ -//! Executor-backed connector declaration loading. - -mod executor_plugin; - -pub use executor_plugin::ExecutorPluginConnectorProvider; -pub use executor_plugin::ExecutorPluginConnectorProviderError; diff --git a/codex-rs/ext/mcp/Cargo.toml b/codex-rs/ext/mcp/Cargo.toml index b98c0905dbc2..13c219edb9f7 100644 --- a/codex-rs/ext/mcp/Cargo.toml +++ b/codex-rs/ext/mcp/Cargo.toml @@ -15,22 +15,18 @@ workspace = true [dependencies] codex-core = { workspace = true } codex-core-plugins = { workspace = true } -codex-config = { workspace = true } codex-exec-server = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } codex-mcp = { workspace = true } -codex-plugin = { workspace = true } codex-protocol = { workspace = true } -codex-utils-absolute-path = { workspace = true } -codex-utils-path-uri = { workspace = true } -serde_json = { workspace = true } -thiserror = { workspace = true } tracing = { workspace = true } tokio = { workspace = true, features = ["sync"] } [dev-dependencies] +codex-config = { workspace = true } codex-login = { workspace = true } +codex-utils-path-uri = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/ext/mcp/src/executor_plugin.rs b/codex-rs/ext/mcp/src/executor_plugin.rs index 69d1a47b31b3..e8eb7aacba1d 100644 --- a/codex-rs/ext/mcp/src/executor_plugin.rs +++ b/codex-rs/ext/mcp/src/executor_plugin.rs @@ -1,5 +1,5 @@ use codex_core::config::Config; -use codex_core_plugins::ExecutorPluginProvider; +use codex_core_plugins::ExecutorPluginRuntime; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionFuture; @@ -11,25 +11,16 @@ use std::collections::HashMap; use std::sync::Arc; use tokio::sync::OnceCell; -use self::provider::ExecutorPluginMcpProvider; - -mod provider; - -/// Frozen MCP declarations for one selected package. -/// -/// Each server config retains the stable logical environment ID. Reconnection may replace the -/// concrete environment instance without changing that authority. +/// The runtime declarations frozen for one selected package at thread start. #[derive(Clone)] -struct SelectedPluginMcpServers { - plugin_id: String, - plugin_display_name: String, +struct SelectedPluginRuntime { selection_order: usize, - servers: Vec<(String, codex_config::McpServerConfig)>, + runtime: ExecutorPluginRuntime, } #[derive(Default)] pub(crate) struct SelectedExecutorPluginMcpState { - snapshot: OnceCell>, + snapshot: OnceCell>, } pub(crate) fn seed_thread_state(thread_init: &mut ExtensionDataInit) { @@ -37,49 +28,38 @@ pub(crate) fn seed_thread_state(thread_init: &mut ExtensionDataInit) { } pub(crate) struct SelectedExecutorPluginMcpContributor { - plugin_provider: ExecutorPluginProvider, - mcp_provider: ExecutorPluginMcpProvider, + environment_manager: Arc, } impl SelectedExecutorPluginMcpContributor { pub(crate) fn new(environment_manager: Arc) -> Self { Self { - plugin_provider: ExecutorPluginProvider::new(Arc::clone(&environment_manager)), - mcp_provider: ExecutorPluginMcpProvider, + environment_manager, } } async fn resolve_snapshot( &self, selected_roots: &[SelectedCapabilityRoot], - ) -> Vec { + ) -> Vec { let mut snapshot = Vec::new(); + let resolved_roots = self + .environment_manager + .bind_selected_capability_roots(selected_roots); - for (selection_order, selected_root) in selected_roots.iter().enumerate() { - let plugin = match self.plugin_provider.resolve_bound(selected_root).await { - Ok(Some(plugin)) => plugin, - Ok(None) => continue, - Err(err) => { - tracing::warn!( - selected_root = selected_root.id, - error = %err, - "failed to resolve selected executor plugin for MCP discovery" - ); - continue; - } - }; - match self.mcp_provider.load(&plugin).await { - Ok(servers) => snapshot.push(SelectedPluginMcpServers { - plugin_id: plugin.plugin().selected_root_id().to_string(), - plugin_display_name: plugin.plugin().manifest().display_name().to_string(), + for (selection_order, resolved_root) in resolved_roots.iter().enumerate() { + let selected_root = resolved_root.selected_root(); + match ExecutorPluginRuntime::project(resolved_root).await { + Ok(Some(runtime)) => snapshot.push(SelectedPluginRuntime { selection_order, - servers, + runtime, }), + Ok(None) => {} Err(err) => { tracing::warn!( selected_root = selected_root.id, error = %err, - "failed to load selected executor plugin MCP servers" + "failed to project selected executor plugin runtime" ); } } @@ -115,19 +95,26 @@ impl McpServerContributor for SelectedExecutorPluginMcpContributor { .await; let mut contributions = Vec::new(); - for plugin in snapshot { - let mut servers = plugin.servers.iter().cloned().collect::>(); + for selected in snapshot { + let plugin = selected.runtime.plugin(); + let plugin_id = plugin.selected_root_id(); + let mut servers = selected + .runtime + .mcp_servers() + .iter() + .cloned() + .collect::>(); context .config() - .apply_plugin_mcp_server_requirements(&plugin.plugin_id, &mut servers); + .apply_plugin_mcp_server_requirements(plugin_id, &mut servers); let mut servers = servers.into_iter().collect::>(); servers.sort_unstable_by(|left, right| left.0.cmp(&right.0)); contributions.extend(servers.into_iter().map(|(name, config)| { McpServerContribution::SelectedPlugin { name, - plugin_id: plugin.plugin_id.clone(), - plugin_display_name: plugin.plugin_display_name.clone(), - selection_order: plugin.selection_order, + plugin_id: plugin_id.to_string(), + plugin_display_name: plugin.manifest().display_name().to_string(), + selection_order: selected.selection_order, config: Box::new(config), } })); diff --git a/codex-rs/ext/mcp/src/executor_plugin/provider.rs b/codex-rs/ext/mcp/src/executor_plugin/provider.rs deleted file mode 100644 index a6972df0969e..000000000000 --- a/codex-rs/ext/mcp/src/executor_plugin/provider.rs +++ /dev/null @@ -1,138 +0,0 @@ -use codex_config::McpServerConfig; -use codex_core_plugins::ResolvedExecutorPlugin; -use codex_exec_server::ExecutorFileSystem; -use codex_mcp::parse_executor_plugin_mcp_config; -use codex_plugin::PluginResourceLocator; -use codex_plugin::ResolvedPlugin; -use codex_plugin::ResolvedPluginLocation; -use codex_plugin::manifest::PluginManifestMcpServers; -use codex_utils_path_uri::PathUri; -use codex_utils_path_uri::PathUriParseError; -use std::io; -use thiserror::Error; - -const DEFAULT_MCP_CONFIG_FILE: &str = ".mcp.json"; - -/// Loads MCP declarations from resolved plugins through their owning executor. -#[derive(Clone, Copy, Debug, Default)] -pub(super) struct ExecutorPluginMcpProvider; - -/// Failure to load an executor plugin's MCP declarations. -#[derive(Debug, Error)] -pub(super) enum ExecutorPluginMcpProviderError { - #[error("failed to read MCP config for selected plugin `{plugin_id}` at `{path}`: {source}")] - ReadConfig { - plugin_id: String, - path: PathUri, - #[source] - source: io::Error, - }, - #[error( - "failed to resolve MCP config path `{relative_path}` below selected plugin `{plugin_id}` at `{root}`: {source}" - )] - InvalidConfigPath { - plugin_id: String, - root: PathUri, - relative_path: &'static str, - #[source] - source: PathUriParseError, - }, - #[error("failed to parse MCP config for selected plugin `{plugin_id}` at `{path}`: {source}")] - ParseConfig { - plugin_id: String, - path: PathUri, - #[source] - source: serde_json::Error, - }, -} - -impl ExecutorPluginMcpProvider { - /// Returns MCP servers declared by `plugin`, bound to its environment. - pub(super) async fn load( - &self, - plugin: &ResolvedExecutorPlugin, - ) -> Result, ExecutorPluginMcpProviderError> { - let ResolvedPluginLocation::Environment { root, .. } = plugin.plugin().location(); - - load_from_file_system(plugin.plugin(), root, plugin.file_system()).await - } -} - -async fn load_from_file_system( - plugin: &ResolvedPlugin, - plugin_root: &PathUri, - file_system: &dyn ExecutorFileSystem, -) -> Result, ExecutorPluginMcpProviderError> { - let ResolvedPluginLocation::Environment { environment_id, .. } = plugin.location(); - let plugin_id = plugin.selected_root_id(); - let (contents, config_path) = match plugin.manifest().paths.mcp_servers.as_ref() { - Some(PluginManifestMcpServers::Path(PluginResourceLocator::Environment { - path, .. - })) => { - ( - file_system - .read_file_text(path, /*sandbox*/ None) - .await - .map_err(|source| ExecutorPluginMcpProviderError::ReadConfig { - plugin_id: plugin_id.to_string(), - path: path.clone(), - source, - })?, - path.clone(), - ) - } - Some(PluginManifestMcpServers::Object(object_config)) => { - let PluginResourceLocator::Environment { path, .. } = plugin.manifest_path(); - (object_config.clone(), path.clone()) - } - None => { - let config_path = plugin_root - .join(DEFAULT_MCP_CONFIG_FILE) - .map_err(|source| ExecutorPluginMcpProviderError::InvalidConfigPath { - plugin_id: plugin_id.to_string(), - root: plugin_root.clone(), - relative_path: DEFAULT_MCP_CONFIG_FILE, - source, - })?; - let contents = match file_system - .read_file_text(&config_path, /*sandbox*/ None) - .await - { - Ok(contents) => contents, - Err(source) if source.kind() == io::ErrorKind::NotFound => { - return Ok(Vec::new()); - } - Err(source) => { - return Err(ExecutorPluginMcpProviderError::ReadConfig { - plugin_id: plugin_id.to_string(), - path: config_path.clone(), - source, - }); - } - }; - (contents, config_path) - } - }; - let parsed = parse_executor_plugin_mcp_config(plugin_root, &contents, environment_id).map_err( - |source| ExecutorPluginMcpProviderError::ParseConfig { - plugin_id: plugin_id.to_string(), - path: config_path, - source, - }, - )?; - - for error in parsed.errors { - tracing::warn!( - plugin = plugin_id, - server = error.name, - error = error.message, - "ignoring invalid executor plugin MCP server" - ); - } - - Ok(parsed.servers.into_iter().collect()) -} - -#[cfg(test)] -#[path = "provider_tests.rs"] -mod tests; diff --git a/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs b/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs deleted file mode 100644 index bc59e104a01f..000000000000 --- a/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs +++ /dev/null @@ -1,421 +0,0 @@ -use super::DEFAULT_MCP_CONFIG_FILE; -use super::ExecutorPluginMcpProviderError; -use super::load_from_file_system; -use codex_config::McpServerConfig; -use codex_config::McpServerTransportConfig; -use codex_exec_server::CopyOptions; -use codex_exec_server::CreateDirectoryOptions; -use codex_exec_server::ExecutorFileSystem; -use codex_exec_server::ExecutorFileSystemFuture; -use codex_exec_server::FileMetadata; -use codex_exec_server::FileSystemReadStream; -use codex_exec_server::FileSystemResult; -use codex_exec_server::FileSystemSandboxContext; -use codex_exec_server::ReadDirectoryEntry; -use codex_exec_server::RemoveOptions; -use codex_plugin::ResolvedPlugin; -use codex_plugin::manifest::PluginManifest; -use codex_plugin::manifest::PluginManifestMcpServers; -use codex_plugin::manifest::PluginManifestPaths; -use codex_utils_absolute_path::AbsolutePathBuf; -use codex_utils_path_uri::LegacyAppPathString; -use codex_utils_path_uri::PathUri; -use pretty_assertions::assert_eq; -use std::collections::HashMap; -use std::io; -use std::sync::Mutex; - -const MCP_CONFIG_CONTENTS: &str = r#"{ - "mcpServers": { - "demo": {"command": "demo-mcp", "environment_id": "local"}, - "hosted": {"url": "https://example.com/mcp"} - } -}"#; - -struct SyntheticExecutorFileSystem { - config_path: AbsolutePathBuf, - config_contents: Option<&'static str>, - reads: Mutex>, -} - -impl SyntheticExecutorFileSystem { - fn unsupported() -> FileSystemResult { - Err(io::Error::new( - io::ErrorKind::Unsupported, - "operation is not used by executor MCP provider tests", - )) - } -} - -impl ExecutorFileSystem for SyntheticExecutorFileSystem { - fn canonicalize<'a>( - &'a self, - _path: &'a PathUri, - _sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, PathUri> { - Box::pin(async { Self::unsupported() }) - } - - fn read_file<'a>( - &'a self, - path: &'a PathUri, - _sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, Vec> { - Box::pin(async move { - let path = path.to_abs_path()?; - self.reads - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(path.clone()); - if path != self.config_path { - return Err(io::Error::new(io::ErrorKind::NotFound, "not found")); - } - self.config_contents - .map(|contents| contents.as_bytes().to_vec()) - .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "not found")) - }) - } - - fn read_file_stream<'a>( - &'a self, - _path: &'a PathUri, - _sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> { - Box::pin(async { Self::unsupported() }) - } - - fn write_file<'a>( - &'a self, - _path: &'a PathUri, - _contents: Vec, - _sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, ()> { - Box::pin(async { Self::unsupported() }) - } - - fn create_directory<'a>( - &'a self, - _path: &'a PathUri, - _options: CreateDirectoryOptions, - _sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, ()> { - Box::pin(async { Self::unsupported() }) - } - - fn get_metadata<'a>( - &'a self, - _path: &'a PathUri, - _sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, FileMetadata> { - Box::pin(async { Self::unsupported() }) - } - - fn read_directory<'a>( - &'a self, - _path: &'a PathUri, - _sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, Vec> { - Box::pin(async { Self::unsupported() }) - } - - fn remove<'a>( - &'a self, - _path: &'a PathUri, - _options: RemoveOptions, - _sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, ()> { - Box::pin(async { Self::unsupported() }) - } - - fn copy<'a>( - &'a self, - _source_path: &'a PathUri, - _destination_path: &'a PathUri, - _options: CopyOptions, - _sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, ()> { - Box::pin(async { Self::unsupported() }) - } -} - -#[tokio::test] -async fn reads_declared_config_only_through_executor_file_system() { - let temp_dir = tempfile::tempdir().expect("tempdir"); - let plugin_root = - AbsolutePathBuf::from_absolute_path_checked(temp_dir.path().join("executor-only-plugin")) - .expect("absolute plugin root"); - assert!(!plugin_root.as_path().exists()); - let config_path = plugin_root.join("config/mcp.json"); - let plugin = resolved_plugin( - &plugin_root, - Some(PluginManifestMcpServers::Path(config_path.clone())), - ); - let file_system = SyntheticExecutorFileSystem { - config_path: config_path.clone(), - config_contents: Some(MCP_CONFIG_CONTENTS), - reads: Mutex::new(Vec::new()), - }; - - let plugin_root_uri = PathUri::from_abs_path(&plugin_root); - let servers = load_from_file_system(&plugin, &plugin_root_uri, &file_system) - .await - .expect("load executor MCP config"); - - assert_eq!( - servers, - vec![ - ( - "demo".to_string(), - McpServerConfig { - auth: Default::default(), - transport: McpServerTransportConfig::Stdio { - command: "demo-mcp".to_string(), - args: Vec::new(), - env: None, - env_vars: Vec::new(), - cwd: Some(LegacyAppPathString::from_path(plugin_root.as_path())), - }, - environment_id: "executor-test".to_string(), - enabled: true, - required: false, - supports_parallel_tool_calls: false, - disabled_reason: None, - startup_timeout_sec: None, - tool_timeout_sec: None, - default_tools_approval_mode: None, - enabled_tools: None, - disabled_tools: None, - scopes: None, - oauth: None, - oauth_resource: None, - tools: HashMap::new(), - }, - ), - ( - "hosted".to_string(), - McpServerConfig { - auth: Default::default(), - transport: McpServerTransportConfig::StreamableHttp { - url: "https://example.com/mcp".to_string(), - bearer_token_env_var: None, - http_headers: None, - env_http_headers: None, - }, - environment_id: "executor-test".to_string(), - enabled: true, - required: false, - supports_parallel_tool_calls: false, - disabled_reason: None, - startup_timeout_sec: None, - tool_timeout_sec: None, - default_tools_approval_mode: None, - enabled_tools: None, - disabled_tools: None, - scopes: None, - oauth: None, - oauth_resource: None, - tools: HashMap::new(), - }, - ), - ] - ); - assert_eq!(reads(&file_system), vec![config_path]); -} - -#[tokio::test] -async fn reads_manifest_object_config_without_executor_file_system_access() { - let temp_dir = tempfile::tempdir().expect("tempdir"); - let plugin_root = AbsolutePathBuf::from_absolute_path_checked(temp_dir.path().join("plugin")) - .expect("absolute plugin root"); - let config_path = plugin_root.join(DEFAULT_MCP_CONFIG_FILE); - let plugin = resolved_plugin( - &plugin_root, - Some(PluginManifestMcpServers::Object( - r#"{"counter":{"command":"counter-mcp","environment_id":"local"}}"#.to_string(), - )), - ); - let file_system = SyntheticExecutorFileSystem { - config_path, - config_contents: None, - reads: Mutex::new(Vec::new()), - }; - - let plugin_root_uri = PathUri::from_abs_path(&plugin_root); - let servers = load_from_file_system(&plugin, &plugin_root_uri, &file_system) - .await - .expect("load manifest object executor MCP config"); - - assert_eq!( - servers, - vec![( - "counter".to_string(), - McpServerConfig { - auth: Default::default(), - transport: McpServerTransportConfig::Stdio { - command: "counter-mcp".to_string(), - args: Vec::new(), - env: None, - env_vars: Vec::new(), - cwd: Some(LegacyAppPathString::from_path(plugin_root.as_path())), - }, - environment_id: "executor-test".to_string(), - enabled: true, - required: false, - supports_parallel_tool_calls: false, - disabled_reason: None, - startup_timeout_sec: None, - tool_timeout_sec: None, - default_tools_approval_mode: None, - enabled_tools: None, - disabled_tools: None, - scopes: None, - oauth: None, - oauth_resource: None, - tools: HashMap::new(), - }, - )] - ); - assert_eq!(reads(&file_system), Vec::new()); -} - -#[tokio::test] -async fn missing_default_config_is_empty() { - let temp_dir = tempfile::tempdir().expect("tempdir"); - let plugin_root = AbsolutePathBuf::from_absolute_path_checked(temp_dir.path().join("plugin")) - .expect("absolute plugin root"); - let config_path = plugin_root.join(DEFAULT_MCP_CONFIG_FILE); - let plugin = resolved_plugin(&plugin_root, /*mcp_servers*/ None); - let file_system = SyntheticExecutorFileSystem { - config_path: config_path.clone(), - config_contents: None, - reads: Mutex::new(Vec::new()), - }; - - let plugin_root_uri = PathUri::from_abs_path(&plugin_root); - let servers = load_from_file_system(&plugin, &plugin_root_uri, &file_system) - .await - .expect("missing default config should be ignored"); - - assert_eq!(servers, Vec::new()); - assert_eq!(reads(&file_system), vec![config_path]); -} - -#[tokio::test] -async fn malformed_declared_config_is_an_error() { - let temp_dir = tempfile::tempdir().expect("tempdir"); - let plugin_root = AbsolutePathBuf::from_absolute_path_checked(temp_dir.path().join("plugin")) - .expect("absolute plugin root"); - let config_path = plugin_root.join("mcp.json"); - let plugin = resolved_plugin( - &plugin_root, - Some(PluginManifestMcpServers::Path(config_path.clone())), - ); - let file_system = SyntheticExecutorFileSystem { - config_path: config_path.clone(), - config_contents: Some("{not-json"), - reads: Mutex::new(Vec::new()), - }; - - let plugin_root_uri = PathUri::from_abs_path(&plugin_root); - let err = load_from_file_system(&plugin, &plugin_root_uri, &file_system) - .await - .expect_err("malformed declared config should fail"); - - let ExecutorPluginMcpProviderError::ParseConfig { - plugin_id, - path, - source: _, - } = err - else { - panic!("expected parse error"); - }; - assert_eq!( - (plugin_id, path), - ( - "selected-root".to_string(), - PathUri::from_abs_path(&config_path) - ) - ); - assert_eq!(reads(&file_system), vec![config_path]); -} - -#[tokio::test] -async fn malformed_manifest_object_config_reports_actual_manifest_path() { - let temp_dir = tempfile::tempdir().expect("tempdir"); - let plugin_root = AbsolutePathBuf::from_absolute_path_checked(temp_dir.path().join("plugin")) - .expect("absolute plugin root"); - let plugin = resolved_plugin( - &plugin_root, - Some(PluginManifestMcpServers::Object("{not-json".to_string())), - ); - let file_system = SyntheticExecutorFileSystem { - config_path: plugin_root.join(DEFAULT_MCP_CONFIG_FILE), - config_contents: None, - reads: Mutex::new(Vec::new()), - }; - - let plugin_root_uri = PathUri::from_abs_path(&plugin_root); - let err = load_from_file_system(&plugin, &plugin_root_uri, &file_system) - .await - .expect_err("malformed manifest object config should fail"); - - let ExecutorPluginMcpProviderError::ParseConfig { - plugin_id, - path, - source: _, - } = err - else { - panic!("expected parse error"); - }; - assert_eq!( - (plugin_id, path), - ( - "selected-root".to_string(), - PathUri::from_abs_path(&plugin_root.join(".claude-plugin/plugin.json")) - ) - ); - assert_eq!(reads(&file_system), Vec::new()); -} - -fn resolved_plugin( - plugin_root: &AbsolutePathBuf, - mcp_servers: Option>, -) -> ResolvedPlugin { - let plugin_root_uri = PathUri::from_abs_path(plugin_root); - let mcp_servers = mcp_servers.map(|mcp_servers| match mcp_servers { - PluginManifestMcpServers::Path(path) => { - PluginManifestMcpServers::Path(PathUri::from_abs_path(&path)) - } - PluginManifestMcpServers::Object(config) => PluginManifestMcpServers::Object(config), - }); - ResolvedPlugin::from_environment( - "selected-root".to_string(), - "executor-test".to_string(), - plugin_root_uri.clone(), - plugin_root_uri - .join(".claude-plugin/plugin.json") - .expect("manifest URI"), - PluginManifest { - name: "demo-plugin".to_string(), - version: None, - description: None, - keywords: Vec::new(), - paths: PluginManifestPaths { - skills: Vec::new(), - mcp_servers, - apps: None, - hooks: None, - }, - interface: None, - }, - ) - .expect("valid plugin descriptor") -} - -fn reads(file_system: &SyntheticExecutorFileSystem) -> Vec { - file_system - .reads - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clone() -} From 88296d82047317584d07e5b62d5599b3a77682e9 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 25 Jun 2026 11:45:34 +0100 Subject: [PATCH 3/6] Project selected MCP runtime per model step --- codex-rs/Cargo.lock | 3 - codex-rs/app-server/src/mcp_refresh.rs | 1 - codex-rs/app-server/src/request_processors.rs | 2 - .../src/request_processors/apps_processor.rs | 89 ++--- .../src/request_processors/mcp_processor.rs | 38 ++- .../request_processors/thread_processor.rs | 8 +- .../app-server/tests/suite/v2/app_list.rs | 4 +- .../app-server/tests/suite/v2/executor_mcp.rs | 295 +++++++++++++++- codex-rs/codex-mcp/src/connection_manager.rs | 12 + .../codex-mcp/src/connection_manager_tests.rs | 95 ------ codex-rs/codex-mcp/src/elicitation.rs | 25 +- codex-rs/codex-mcp/src/lib.rs | 1 + codex-rs/codex-mcp/src/resource_client.rs | 14 +- codex-rs/codex-mcp/src/runtime.rs | 118 ++++++- codex-rs/core-plugins/src/executor_runtime.rs | 125 ++----- codex-rs/core-plugins/src/lib.rs | 1 - codex-rs/core/src/agent/control_tests.rs | 1 + codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/codex_thread.rs | 17 +- codex-rs/core/src/connectors.rs | 15 +- codex-rs/core/src/mcp.rs | 118 ++++--- codex-rs/core/src/mcp_skill_dependencies.rs | 9 +- codex-rs/core/src/prompt_debug.rs | 4 +- codex-rs/core/src/session/mcp.rs | 138 ++------ codex-rs/core/src/session/mcp_projection.rs | 319 ++++++++++++++++++ codex-rs/core/src/session/mcp_runtime.rs | 143 ++++---- codex-rs/core/src/session/mod.rs | 28 +- codex-rs/core/src/session/session.rs | 43 +-- codex-rs/core/src/session/tests.rs | 74 ++-- .../core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/src/session/turn.rs | 146 +++++--- codex-rs/core/src/state/service.rs | 61 +++- codex-rs/core/src/thread_manager.rs | 15 + codex-rs/core/src/thread_manager_tests.rs | 296 +--------------- codex-rs/core/src/tools/spec_plan_tests.rs | 2 +- codex-rs/core/tests/common/test_codex.rs | 17 +- codex-rs/core/tests/suite/agents_md.rs | 2 + .../tests/suite/subagent_notifications.rs | 1 + .../ext/extension-api/src/contributors/mcp.rs | 34 +- codex-rs/ext/mcp/Cargo.toml | 1 - codex-rs/ext/mcp/src/executor_plugin.rs | 126 ------- codex-rs/ext/mcp/src/lib.rs | 19 -- codex-rs/ext/mcp/tests/executor_plugin_mcp.rs | 140 -------- codex-rs/ext/skills/Cargo.toml | 6 - codex-rs/memories/write/src/runtime.rs | 1 + 45 files changed, 1349 insertions(+), 1260 deletions(-) create mode 100644 codex-rs/core/src/session/mcp_projection.rs delete mode 100644 codex-rs/ext/mcp/src/executor_plugin.rs delete mode 100644 codex-rs/ext/mcp/tests/executor_plugin_mcp.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 00c6ef08d56c..4f6795b1696b 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3915,10 +3915,7 @@ dependencies = [ "codex-mcp", "codex-protocol", "codex-tools", - "codex-utils-absolute-path", - "codex-utils-path-uri", "codex-utils-string", - "pretty_assertions", "schemars 0.8.22", "serde", "serde_json", diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 426d8fcc67b6..680cf303c371 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -231,7 +231,6 @@ mod tests { analytics_events_client: codex_analytics::AnalyticsEventsClient::disabled(), thread_manager: thread_manager.clone(), goal_service: Arc::new(codex_goal_extension::GoalService::new()), - environment_manager: Arc::clone(&environment_manager), thread_store: Arc::clone(&thread_store), }, ), diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 9b56957a53f5..e6a97d43b9ae 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -305,7 +305,6 @@ use codex_connectors::AppInfo; use codex_core::CodexThread; use codex_core::CodexThreadSettingsOverrides; use codex_core::ForkSnapshot; -use codex_core::McpManager; use codex_core::NewThread; #[cfg(test)] use codex_core::SessionMeta; @@ -335,7 +334,6 @@ use codex_core_plugins::PluginInstallError as CorePluginInstallError; use codex_core_plugins::PluginInstallRequest; use codex_core_plugins::PluginReadRequest; use codex_core_plugins::PluginUninstallError as CorePluginUninstallError; -use codex_core_plugins::PluginsManager; use codex_core_plugins::loader::load_plugin_apps; use codex_core_plugins::loader::load_plugin_mcp_servers; use codex_core_plugins::manifest::PluginManifestInterface; diff --git a/codex-rs/app-server/src/request_processors/apps_processor.rs b/codex-rs/app-server/src/request_processors/apps_processor.rs index 67d768fa57b0..ea4629101de4 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor.rs +++ b/codex-rs/app-server/src/request_processors/apps_processor.rs @@ -59,12 +59,11 @@ impl AppsRequestProcessor { }; let mut config = self.load_latest_config(fallback_cwd).await?; - if let Some(thread) = thread { + if let Some(thread) = thread.as_ref() { let _ = config .features .set_enabled(Feature::Apps, thread.enabled(Feature::Apps)); } - let auth = self.auth_manager.auth().await; if !config .features @@ -86,11 +85,24 @@ impl AppsRequestProcessor { })); } + let mcp_projection = match thread.as_ref() { + Some(thread) => thread.project_mcp_config(&config).await, + None => { + let mcp_config = self + .thread_manager + .mcp_manager() + .runtime_config(&config) + .await; + let runtime_context = McpRuntimeContext::new( + self.thread_manager.environment_manager(), + config.cwd.to_path_buf(), + ); + (mcp_config, runtime_context) + } + }; + let request = request_id.clone(); let outgoing = Arc::clone(&self.outgoing); - let environment_manager = self.thread_manager.environment_manager(); - let mcp_manager = self.thread_manager.mcp_manager(); - let plugins_manager = self.thread_manager.plugins_manager(); let shutdown_token = self.shutdown_token.child_token(); tokio::spawn(async move { tokio::select! { @@ -100,9 +112,7 @@ impl AppsRequestProcessor { request, params, config, - environment_manager, - mcp_manager, - plugins_manager, + mcp_projection, ) => {} } }); @@ -118,24 +128,12 @@ impl AppsRequestProcessor { request_id: ConnectionRequestId, params: AppsListParams, config: Config, - environment_manager: Arc, - mcp_manager: Arc, - plugins_manager: Arc, + mcp_projection: (codex_mcp::McpConfig, codex_mcp::McpRuntimeContext), ) { let retry_params = params.clone(); let retry_config = config.clone(); - let retry_environment_manager = Arc::clone(&environment_manager); - let retry_mcp_manager = Arc::clone(&mcp_manager); - let retry_plugins_manager = Arc::clone(&plugins_manager); - let result = Self::apps_list_response( - &outgoing, - params, - config, - environment_manager, - mcp_manager, - plugins_manager, - ) - .await; + let retry_mcp_projection = mcp_projection.clone(); + let result = Self::apps_list_response(&outgoing, params, config, mcp_projection).await; let should_retry = result .as_ref() .is_ok_and(|(_, codex_apps_ready)| !codex_apps_ready); @@ -150,9 +148,7 @@ impl AppsRequestProcessor { &outgoing, retry_params, retry_config, - retry_environment_manager, - retry_mcp_manager, - retry_plugins_manager, + retry_mcp_projection, ) .await { @@ -165,9 +161,7 @@ impl AppsRequestProcessor { outgoing: &Arc, params: AppsListParams, config: Config, - environment_manager: Arc, - mcp_manager: Arc, - plugins_manager: Arc, + mcp_projection: (codex_mcp::McpConfig, codex_mcp::McpRuntimeContext), ) -> Result<(AppsListResponse, bool), JSONRPCErrorError> { let AppsListParams { cursor, @@ -183,13 +177,7 @@ impl AppsRequestProcessor { None => 0, }; - let loaded_plugins = plugins_manager - .plugins_for_config(&config.plugins_config_input()) - .await; - let connector_snapshot = - codex_connectors::ConnectorSnapshot::from_plugin_capability_summaries( - loaded_plugins.capability_summaries(), - ); + let connector_snapshot = mcp_projection.0.connector_snapshot.clone(); let plugin_apps = connector_snapshot.connector_ids().to_vec(); let (mut accessible_connectors, mut all_connectors) = tokio::join!( connectors::list_cached_accessible_connectors_from_mcp_tools(&config), @@ -202,11 +190,12 @@ impl AppsRequestProcessor { let accessible_config = config.clone(); let accessible_tx = tx.clone(); tokio::spawn(async move { - let result = connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager( + let (mcp_config, runtime_context) = mcp_projection; + let result = codex_core::connectors::list_accessible_connectors_from_mcp_config( &accessible_config, force_refetch, - Arc::clone(&environment_manager), - mcp_manager, + mcp_config, + runtime_context, ) .await .map_err(|err| format!("failed to load accessible apps: {err}")); @@ -234,7 +223,11 @@ impl AppsRequestProcessor { if accessible_connectors.is_some() || all_connectors.is_some() { let merged = connectors::with_app_enabled_state( - merge_loaded_apps(all_connectors.as_deref(), accessible_connectors.as_deref()), + merge_loaded_apps( + all_connectors.as_deref(), + accessible_connectors.as_deref(), + &connector_snapshot, + ), &config, ); if should_send_app_list_updated_notification( @@ -293,7 +286,11 @@ impl AppsRequestProcessor { accessible_connectors.as_deref() }; let merged = connectors::with_app_enabled_state( - merge_loaded_apps(all_connectors_for_update, accessible_connectors_for_update), + merge_loaded_apps( + all_connectors_for_update, + accessible_connectors_for_update, + &connector_snapshot, + ), &config, ); if should_send_app_list_updated_notification( @@ -372,11 +369,19 @@ enum AppListLoadResult { fn merge_loaded_apps( all_connectors: Option<&[AppInfo]>, accessible_connectors: Option<&[AppInfo]>, + connector_snapshot: &codex_connectors::ConnectorSnapshot, ) -> Vec { let all_connectors_loaded = all_connectors.is_some(); let all = all_connectors.map_or_else(Vec::new, <[AppInfo]>::to_vec); let accessible = accessible_connectors.map_or_else(Vec::new, <[AppInfo]>::to_vec); - connectors::merge_connectors_with_accessible(all, accessible, all_connectors_loaded) + let mut merged = + connectors::merge_connectors_with_accessible(all, accessible, all_connectors_loaded); + for connector in &mut merged { + connector.plugin_display_names = connector_snapshot + .plugin_display_names_for_connector_id(&connector.id) + .to_vec(); + } + merged } fn should_send_app_list_updated_notification( diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index d38db0d4a302..f7d5f1b2d9e0 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -124,8 +124,13 @@ impl McpRequestProcessor { let (mcp_config, runtime_context) = match thread_id.as_deref() { Some(thread_id) => { let (_, thread) = self.load_thread(thread_id).await?; - let runtime = thread.current_mcp_runtime(); - (runtime.config().clone(), runtime.runtime_context().clone()) + let thread_config = thread.config().await; + let config = self + .config_manager + .load_latest_config_for_thread(thread_config.as_ref()) + .await + .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; + thread.project_mcp_config(&config).await } None => { let config = self.load_latest_config(/*fallback_cwd*/ None).await?; @@ -245,10 +250,7 @@ impl McpRequestProcessor { }; let auth = self.auth_manager.auth().await; let (mcp_config, runtime_context) = match thread { - Some(thread) => { - let runtime = thread.current_mcp_runtime(); - (runtime.config().clone(), runtime.runtime_context().clone()) - } + Some(thread) => thread.project_mcp_config(&config).await, None => { let mcp_config = self .thread_manager @@ -393,10 +395,20 @@ impl McpRequestProcessor { if let Some(thread_id) = thread_id { let (_, thread) = self.load_thread(&thread_id).await?; + let thread_config = thread.config().await; + let config = self + .config_manager + .load_latest_config_for_thread(thread_config.as_ref()) + .await + .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; + let runtime = thread.project_mcp_runtime(&config).await; let request_id = request_id.clone(); tokio::spawn(async move { - let result = thread.read_mcp_resource(&server, &uri).await; + let result = runtime + .read_resource(&server, uri) + .await + .and_then(|result| serde_json::to_value(result).map_err(anyhow::Error::from)); Self::send_mcp_resource_read_response(outgoing, request_id, result).await; }); return Ok(()); @@ -457,12 +469,20 @@ impl McpRequestProcessor { let outgoing = Arc::clone(&self.outgoing); let thread_id = params.thread_id.clone(); let (_, thread) = self.load_thread(&thread_id).await?; + let thread_config = thread.config().await; + let config = self + .config_manager + .load_latest_config_for_thread(thread_config.as_ref()) + .await + .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; + let runtime = thread.project_mcp_runtime(&config).await; let meta = with_mcp_tool_call_thread_id_meta(params.meta, &thread_id); let request_id = request_id.clone(); tokio::spawn(async move { - let result = thread - .call_mcp_tool(¶ms.server, ¶ms.tool, params.arguments, meta) + let result = runtime + .manager_arc() + .call_tool(¶ms.server, ¶ms.tool, params.arguments, meta) .await .map(McpServerToolCallResponse::from) .map_err(|error| internal_error(format!("{error:#}"))); diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 7242e816beda..588b8dfe7aa6 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1129,11 +1129,6 @@ impl ThreadRequestProcessor { DynamicToolSpec::Namespace(namespace) => namespace.tools.len(), }) .sum(); - let mut thread_extension_init = ExtensionDataInit::new(); - if !selected_capability_roots.is_empty() { - thread_extension_init.insert(selected_capability_roots); - codex_mcp_extension::initialize_executor_plugin_thread_data(&mut thread_extension_init); - } let create_thread_started_at = std::time::Instant::now(); let NewThread { thread_id, @@ -1156,7 +1151,8 @@ impl ThreadRequestProcessor { metrics_service_name: service_name, parent_trace: request_trace, environments, - thread_extension_init, + selected_capability_roots, + thread_extension_init: ExtensionDataInit::new(), supports_openai_form_elicitation, }) .instrument(tracing::info_span!( diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index 1ddf355269f8..5fad04929f7f 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -1533,7 +1533,7 @@ impl ServerHandler for AppListMcpServer { } } -async fn start_apps_server_with_delays( +pub(super) async fn start_apps_server_with_delays( connectors: Vec, tools: Vec, directory_delay: Duration, @@ -1693,7 +1693,7 @@ async fn list_directory_connectors( } } -fn connector_tool(connector_id: &str, connector_name: &str) -> Result { +pub(super) fn connector_tool(connector_id: &str, connector_name: &str) -> Result { let schema: JsonObject = serde_json::from_value(json!({ "type": "object", "additionalProperties": false diff --git a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs index 67d97bc248a1..f23bec8e060f 100644 --- a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs +++ b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs @@ -1,12 +1,18 @@ use anyhow::Result; +use app_test_support::ChatGptAuthFixture; use app_test_support::TestAppServer; use app_test_support::to_response; +use app_test_support::write_chatgpt_auth; use app_test_support::write_mock_responses_config_toml; +use app_test_support::write_mock_responses_config_toml_with_chatgpt_base_url; use axum::Json; use axum::Router; use axum::body::Bytes; use axum::routing::get; use axum::routing::post; +use codex_app_server_protocol::AppInfo; +use codex_app_server_protocol::AppsListParams; +use codex_app_server_protocol::AppsListResponse; use codex_app_server_protocol::CapabilityRootLocation; use codex_app_server_protocol::ListMcpServerStatusParams; use codex_app_server_protocol::ListMcpServerStatusResponse; @@ -16,10 +22,13 @@ use codex_app_server_protocol::McpServerToolCallParams; use codex_app_server_protocol::McpServerToolCallResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::SelectedCapabilityRoot; +use codex_app_server_protocol::ThreadResumeParams; +use codex_app_server_protocol::ThreadResumeResponse; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::UserInput; +use codex_config::types::AuthCredentialsStoreMode; use codex_utils_path_uri::PathUri; use core_test_support::responses; use core_test_support::stdio_server_bin; @@ -59,6 +68,14 @@ const EXECUTOR_ENV_VALUE: &str = "executor-only"; const EXECUTOR_ID: &str = "executor-1"; const REFRESH_PROBE_SERVER_NAME: &str = "refresh_probe"; const TOOL_CALL_ID: &str = "executor-mcp-call"; +const FALLBACK_MCP_SERVER_NAME: &str = "calendar"; +const DYNAMIC_MCP_SERVER_NAME: &str = "executor_probe"; +const CONNECTOR_ID: &str = "calendar"; +const PLUGIN_DISPLAY_NAME: &str = "Executor Demo"; +const DYNAMIC_TOOL_CALL_ID: &str = "dynamic-executor-mcp-call"; + +use super::app_list::connector_tool; +use super::app_list::start_apps_server_with_delays; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn selected_executor_plugin_exposes_its_mcps_only_to_that_thread() -> Result<()> { @@ -191,7 +208,6 @@ HTTP_PROXY = {http_proxy} ) .await?; - std::fs::write(plugin.path().join(".mcp.json"), r#"{"mcpServers":{}}"#)?; let config_path = codex_home.path().join("config.toml"); let mut config = std::fs::read_to_string(&config_path)?; config.push_str(&format!( @@ -412,6 +428,283 @@ startup_timeout_sec = 10 Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn selected_executor_plugin_activates_after_it_becomes_ready() -> Result<()> { + let responses_server = responses::start_mock_server().await; + let (apps_url, apps_server_handle) = start_apps_server_with_delays( + vec![AppInfo { + id: CONNECTOR_ID.to_string(), + name: "Calendar".to_string(), + description: None, + logo_url: None, + logo_url_dark: None, + icon_assets: None, + icon_dark_assets: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: None, + is_accessible: false, + is_enabled: true, + plugin_display_names: Vec::new(), + }], + vec![connector_tool(CONNECTOR_ID, "Calendar")?], + Duration::ZERO, + Duration::ZERO, + ) + .await?; + let codex_home = TempDir::new()?; + write_mock_responses_config_toml_with_chatgpt_base_url( + codex_home.path(), + &responses_server.uri(), + &apps_url, + )?; + let config_path = codex_home.path().join("config.toml"); + let config = std::fs::read_to_string(&config_path)?.replacen( + "model_provider = \"mock_provider\"", + "mcp_oauth_credentials_store = \"file\"\nmodel_provider = \"mock_provider\"", + 1, + ); + std::fs::write(config_path, format!("{config}\n[features]\napps = true\n"))?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("chatgpt-token") + .account_id("account-123") + .email("executor-runtime@example.com") + .plan_type("pro") + .chatgpt_account_id("account-123"), + AuthCredentialsStoreMode::File, + )?; + std::fs::write( + codex_home.path().join("environments.toml"), + format!( + r#" +include_local = true + +[[environments]] +id = "{EXECUTOR_ID}" +program = {} +args = ["exec-server", "--listen", "stdio"] +[environments.env] +{EXECUTOR_ENV_NAME} = "{EXECUTOR_ENV_VALUE}" +"#, + toml::Value::String( + codex_utils_cargo_bin::cargo_bin("codex")? + .to_string_lossy() + .into_owned() + ) + ), + )?; + + // The selected root exists, but its manifest arrives after the first model step. + let plugin = TempDir::new()?; + let mut app_server = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; + let selected_thread = start_thread( + &mut app_server, + Some(vec![SelectedCapabilityRoot { + id: "executor-demo@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: EXECUTOR_ID.to_string(), + path: PathUri::from_host_native_path(plugin.path())?, + }, + }]), + ) + .await?; + + let before_ready = responses::mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("before-ready"), + responses::ev_assistant_message("before-ready-message", "Waiting"), + responses::ev_completed("before-ready"), + ]), + ) + .await; + run_turn(&mut app_server, &selected_thread, "Check available tools").await?; + let namespace = format!("mcp__{DYNAMIC_MCP_SERVER_NAME}"); + assert!( + before_ready + .single_request() + .tool_by_name(&namespace, "echo") + .is_none() + ); + + std::fs::create_dir_all(plugin.path().join(".codex-plugin"))?; + std::fs::write( + plugin.path().join(".codex-plugin/plugin.json"), + r#"{"name":"executor-demo","apps":"./.app.json","interface":{"displayName":"Executor Demo"}}"#, + )?; + std::fs::write( + plugin.path().join(".app.json"), + format!(r#"{{"apps":{{"{FALLBACK_MCP_SERVER_NAME}":{{"id":"{CONNECTOR_ID}"}}}}}}"#), + )?; + std::fs::write( + plugin.path().join(".mcp.json"), + serde_json::to_vec_pretty(&json!({ + "mcpServers": { + (FALLBACK_MCP_SERVER_NAME): { + "command": stdio_server_bin()?, + }, + (DYNAMIC_MCP_SERVER_NAME): { + "command": stdio_server_bin()?, + "env_vars": [EXECUTOR_ENV_NAME], + "startup_timeout_sec": 10, + } + } + }))?, + )?; + + let response_mock = responses::mount_sse_sequence( + &responses_server, + vec![ + responses::sse(vec![ + responses::ev_response_created("executor-call"), + responses::ev_function_call_with_namespace( + DYNAMIC_TOOL_CALL_ID, + &namespace, + "echo", + &json!({ + "message": "hello from executor", + "env_var": EXECUTOR_ENV_NAME, + }) + .to_string(), + ), + responses::ev_completed("executor-call"), + ]), + responses::sse(vec![ + responses::ev_response_created("executor-done"), + responses::ev_assistant_message("executor-done-message", "Done"), + responses::ev_completed("executor-done"), + ]), + ], + ) + .await; + run_turn( + &mut app_server, + &selected_thread, + "Call the executor MCP echo tool", + ) + .await?; + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + assert!(requests[0].tool_by_name(&namespace, "echo").is_some()); + assert!( + requests[0] + .tool_by_name(&format!("mcp__{FALLBACK_MCP_SERVER_NAME}"), "echo") + .is_none() + ); + let connector = requests[0] + .tool_by_name("mcp__codex_apps__calendar", "connector_calendar") + .expect("selected connector should be model-visible"); + assert!( + connector["description"] + .as_str() + .is_some_and(|description| description.contains(PLUGIN_DISPLAY_NAME)) + ); + let tool_output = requests[1].function_call_output(DYNAMIC_TOOL_CALL_ID); + let output = tool_output["output"] + .as_str() + .expect("MCP function output should be text"); + assert!(output.contains("ECHOING: hello from executor")); + assert!(output.contains(EXECUTOR_ENV_VALUE)); + + assert_eq!( + connector_plugin_names(&mut app_server, &selected_thread).await?, + vec![PLUGIN_DISPLAY_NAME.to_string()] + ); + + drop(app_server); + let mut app_server = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; + let request_id = app_server + .send_thread_resume_request(ThreadResumeParams { + thread_id: selected_thread.clone(), + ..Default::default() + }) + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let ThreadResumeResponse { thread, .. } = to_response(response)?; + assert_eq!(thread.id, selected_thread); + let resumed = responses::mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("resumed"), + responses::ev_assistant_message("resumed-message", "Ready"), + responses::ev_completed("resumed"), + ]), + ) + .await; + run_turn(&mut app_server, &selected_thread, "Check resumed tools").await?; + let resumed_request = resumed.single_request(); + assert!(resumed_request.tool_by_name(&namespace, "echo").is_some()); + assert!( + resumed_request + .tool_by_name("mcp__codex_apps__calendar", "connector_calendar") + .is_some() + ); + + apps_server_handle.abort(); + let _ = apps_server_handle.await; + Ok(()) +} + +async fn run_turn(app_server: &mut TestAppServer, thread_id: &str, text: &str) -> Result<()> { + let request_id = app_server + .send_turn_start_request(TurnStartParams { + thread_id: thread_id.to_string(), + input: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_notification_message("turn/completed"), + ) + .await??; + Ok(()) +} + +async fn connector_plugin_names( + app_server: &mut TestAppServer, + thread_id: &str, +) -> Result> { + let request_id = app_server + .send_apps_list_request(AppsListParams { + cursor: None, + limit: None, + thread_id: Some(thread_id.to_string()), + force_refetch: true, + }) + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let response: AppsListResponse = to_response(response)?; + Ok(response + .data + .into_iter() + .find(|app| app.id == CONNECTOR_ID) + .map(|app| app.plugin_display_names) + .unwrap_or_default()) +} + #[derive(Clone, Copy)] struct ExecutorHttpMcpServer; diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 55c72b046723..36d6cfe9f43a 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -429,6 +429,18 @@ impl McpConnectionManager { .await } + /// Resolves an elicitation only if its routed ID belongs to this manager. + pub async fn try_resolve_elicitation( + &self, + server_name: String, + id: RequestId, + response: ElicitationResponse, + ) -> Result { + self.elicitation_requests + .try_resolve(server_name, id, response) + .await + } + pub async fn wait_for_server_ready(&self, server_name: &str, timeout: Duration) -> bool { let Some(async_managed_client) = self.clients.get(server_name) else { return false; diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 750938708ff6..d66474c7bb30 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -310,101 +310,6 @@ async fn disabled_permissions_do_not_auto_accept_elicitation_with_requested_fiel ); } -#[tokio::test] -async fn elicitation_route_ids_are_unique_across_managers() { - let first_manager = ElicitationRequestManager::new( - AskForApproval::OnRequest, - PermissionProfile::default(), - /*reviewer*/ None, - ); - let second_manager = ElicitationRequestManager::new( - AskForApproval::OnRequest, - PermissionProfile::default(), - /*reviewer*/ None, - ); - let (first_tx, first_rx) = async_channel::bounded(1); - let (second_tx, second_rx) = async_channel::bounded(1); - let first_sender = first_manager.make_sender("server".to_string(), first_tx); - let second_sender = second_manager.make_sender("server".to_string(), second_tx); - - let first_task = tokio::spawn(first_sender( - NumberOrString::Number(1), - codex_rmcp_client::Elicitation::OpenAiForm { - meta: None, - message: "First?".to_string(), - requested_schema: serde_json::json!({}), - }, - )); - let second_task = tokio::spawn(second_sender( - NumberOrString::Number(1), - codex_rmcp_client::Elicitation::OpenAiForm { - meta: None, - message: "Second?".to_string(), - requested_schema: serde_json::json!({}), - }, - )); - - let EventMsg::ElicitationRequest(first_request) = first_rx - .recv() - .await - .expect("first elicitation should be emitted") - .msg - else { - panic!("expected first elicitation request"); - }; - let EventMsg::ElicitationRequest(second_request) = second_rx - .recv() - .await - .expect("second elicitation should be emitted") - .msg - else { - panic!("expected second elicitation request"); - }; - assert_ne!(first_request.id, second_request.id); - - let first_response = ElicitationResponse { - action: ElicitationAction::Accept, - content: Some(serde_json::json!({"manager": "first"})), - meta: None, - }; - let second_response = ElicitationResponse { - action: ElicitationAction::Decline, - content: Some(serde_json::json!({"manager": "second"})), - meta: None, - }; - let first_id = match first_request.id { - codex_protocol::mcp::RequestId::String(value) => NumberOrString::String(Arc::from(value)), - codex_protocol::mcp::RequestId::Integer(value) => NumberOrString::Number(value), - }; - let second_id = match second_request.id { - codex_protocol::mcp::RequestId::String(value) => NumberOrString::String(Arc::from(value)), - codex_protocol::mcp::RequestId::Integer(value) => NumberOrString::Number(value), - }; - first_manager - .resolve("server".to_string(), first_id, first_response.clone()) - .await - .expect("first manager should resolve its routed request"); - second_manager - .resolve("server".to_string(), second_id, second_response.clone()) - .await - .expect("second manager should resolve its routed request"); - - assert_eq!( - first_task - .await - .expect("first elicitation task should join") - .expect("first elicitation should resolve"), - first_response - ); - assert_eq!( - second_task - .await - .expect("second elicitation task should join") - .expect("second elicitation should resolve"), - second_response - ); -} - #[test] fn test_normalize_tools_short_non_duplicated_names() { let tools = vec![ diff --git a/codex-rs/codex-mcp/src/elicitation.rs b/codex-rs/codex-mcp/src/elicitation.rs index 799852437eec..ca0fc9b2161c 100644 --- a/codex-rs/codex-mcp/src/elicitation.rs +++ b/codex-rs/codex-mcp/src/elicitation.rs @@ -97,13 +97,26 @@ impl ElicitationRequestManager { id: RequestId, response: ElicitationResponse, ) -> Result<()> { - self.requests - .lock() - .await - .remove(&(server_name, id)) - .ok_or_else(|| anyhow!("elicitation request not found"))? + if self.try_resolve(server_name, id, response).await? { + Ok(()) + } else { + Err(anyhow!("elicitation request not found")) + } + } + + pub(crate) async fn try_resolve( + &self, + server_name: String, + id: RequestId, + response: ElicitationResponse, + ) -> Result { + let Some(responder) = self.requests.lock().await.remove(&(server_name, id)) else { + return Ok(false); + }; + responder .send(response) - .map_err(|e| anyhow!("failed to send elicitation response: {e:?}")) + .map_err(|e| anyhow!("failed to send elicitation response: {e:?}"))?; + Ok(true) } pub(crate) fn make_sender( diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 98efdc1b5654..fd527881af1d 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -9,6 +9,7 @@ pub use resource_client::McpResourcePage; pub use resource_client::McpResourceReadResult; pub use rmcp_client::MCP_SANDBOX_STATE_META_CAPABILITY; pub use runtime::McpRuntimeContext; +pub use runtime::McpRuntimeSnapshot; pub use runtime::SandboxState; pub use tools::ToolInfo; diff --git a/codex-rs/codex-mcp/src/resource_client.rs b/codex-rs/codex-mcp/src/resource_client.rs index 6f479c4027d8..996ed9de696e 100644 --- a/codex-rs/codex-mcp/src/resource_client.rs +++ b/codex-rs/codex-mcp/src/resource_client.rs @@ -3,13 +3,14 @@ use std::sync::Weak; use anyhow::Context; use anyhow::Result; -use arc_swap::ArcSwap; +use arc_swap::ArcSwapOption; use codex_protocol::mcp::Resource; use codex_protocol::mcp::ResourceContent; use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; use crate::McpConnectionManager; +use crate::McpRuntimeSnapshot; /// One page of resources returned by an MCP server. #[derive(Clone, Debug, PartialEq)] @@ -38,7 +39,7 @@ pub struct McpResourceClient { #[derive(Clone)] enum ResourceManager { - Live(Arc>), + Live(Arc>), Snapshot(Arc), } @@ -64,9 +65,9 @@ impl std::fmt::Debug for McpResourceClient { impl McpResourceClient { /// Creates a resource client backed by the session's replaceable MCP manager. - pub fn new(manager: Arc>) -> Self { + pub fn new(runtime: Arc>) -> Self { Self { - manager: ResourceManager::Live(manager), + manager: ResourceManager::Live(runtime), } } @@ -80,7 +81,10 @@ impl McpResourceClient { /// Returns the exact manager generation used by this client right now. pub fn manager_snapshot(&self) -> Arc { match &self.manager { - ResourceManager::Live(manager) => manager.load_full(), + ResourceManager::Live(runtime) => runtime + .load_full() + .expect("MCP runtime must be installed before reading resources") + .manager_arc(), ResourceManager::Snapshot(manager) => Arc::clone(manager), } } diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index b9b769ffb4fe..2d1b5403b450 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -5,6 +5,8 @@ //! tiny shared metrics helper. Transport startup and orchestration live in //! [`crate::rmcp_client`] and [`crate::connection_manager`]. +use std::collections::HashMap; +use std::fmt; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -15,9 +17,87 @@ use codex_exec_server::HttpClient; use codex_exec_server::ReqwestHttpClient; use codex_protocol::models::PermissionProfile; use codex_utils_path_uri::PathUri; +use rmcp::model::ReadResourceRequestParams; +use rmcp::model::ReadResourceResult; use serde::Deserialize; use serde::Serialize; +use crate::connection_manager::McpConnectionManager; +use crate::mcp::McpConfig; + +/// MCP config, exact environment bindings, and manager used by one model request. +pub struct McpRuntimeSnapshot { + config: Arc, + manager: Arc, + runtime_context: McpRuntimeContext, +} + +impl McpRuntimeSnapshot { + pub fn new( + config: Arc, + manager: Arc, + runtime_context: McpRuntimeContext, + ) -> Self { + Self { + config, + manager, + runtime_context, + } + } + + pub fn config(&self) -> &McpConfig { + self.config.as_ref() + } + + pub fn manager(&self) -> &McpConnectionManager { + self.manager.as_ref() + } + + pub fn manager_arc(&self) -> Arc { + Arc::clone(&self.manager) + } + + pub fn runtime_context(&self) -> &McpRuntimeContext { + &self.runtime_context + } + + pub async fn read_resource( + &self, + server: &str, + uri: String, + ) -> anyhow::Result { + self.manager + .read_resource(server, ReadResourceRequestParams::new(uri)) + .await + } + + pub fn matches_projection( + &self, + config: &McpConfig, + runtime_context: &McpRuntimeContext, + ) -> bool { + crate::configured_mcp_servers(self.config()) == crate::configured_mcp_servers(config) + && crate::tool_plugin_provenance(self.config()) == crate::tool_plugin_provenance(config) + && self.config.mcp_oauth_credentials_store_mode + == config.mcp_oauth_credentials_store_mode + && self.config.auth_keyring_backend_kind == config.auth_keyring_backend_kind + && self.config.codex_home == config.codex_home + && self.config.chatgpt_base_url == config.chatgpt_base_url + && self.config.apps_enabled == config.apps_enabled + && self.config.prefix_mcp_tool_names == config.prefix_mcp_tool_names + && self.config.client_elicitation_capability == config.client_elicitation_capability + && self.runtime_context.same_bindings(runtime_context) + } +} + +impl fmt::Debug for McpRuntimeSnapshot { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("McpRuntimeSnapshot") + .finish_non_exhaustive() + } +} + #[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct SandboxState { @@ -30,12 +110,13 @@ pub struct SandboxState { /// Runtime context used when resolving per-server MCP environments. /// -/// `McpConfig` describes what servers exist. This value carries the canonical -/// environment registry plus the local stdio fallback cwd used when a local -/// stdio server omits its own working directory. +/// `McpConfig` describes what servers exist. This value carries the exact +/// step bindings when available, the ambient registry for other servers, and +/// the local stdio fallback cwd. #[derive(Clone)] pub struct McpRuntimeContext { environment_manager: Arc, + pinned_environments: Arc>>, local_stdio_fallback_cwd: PathBuf, } @@ -46,14 +127,36 @@ impl McpRuntimeContext { ) -> Self { Self { environment_manager, + pinned_environments: Arc::new(HashMap::new()), local_stdio_fallback_cwd, } } + /// Pins named environments for servers projected from one model step. + pub fn with_pinned_environments( + mut self, + environments: impl IntoIterator)>, + ) -> Self { + self.pinned_environments = Arc::new(environments.into_iter().collect()); + self + } + pub(crate) fn local_stdio_fallback_cwd(&self) -> PathBuf { self.local_stdio_fallback_cwd.clone() } + fn same_bindings(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.environment_manager, &other.environment_manager) + && self.local_stdio_fallback_cwd == other.local_stdio_fallback_cwd + && self.pinned_environments.len() == other.pinned_environments.len() + && self.pinned_environments.iter().all(|(id, environment)| { + other + .pinned_environments + .get(id) + .is_some_and(|other| Arc::ptr_eq(environment, other)) + }) + } + pub(crate) fn resolve_server_environment( &self, server_name: &str, @@ -63,8 +166,13 @@ impl McpRuntimeContext { // HTTP is the one current exception: it can use the ambient HTTP client // even when no local Environment is configured. if let Some(environment) = self - .environment_manager - .get_environment(&config.environment_id) + .pinned_environments + .get(&config.environment_id) + .cloned() + .or_else(|| { + self.environment_manager + .get_environment(&config.environment_id) + }) { return Ok(Some(environment)); } diff --git a/codex-rs/core-plugins/src/executor_runtime.rs b/codex-rs/core-plugins/src/executor_runtime.rs index c70750805643..5c6becf9973c 100644 --- a/codex-rs/core-plugins/src/executor_runtime.rs +++ b/codex-rs/core-plugins/src/executor_runtime.rs @@ -1,3 +1,4 @@ +use anyhow::Context; use codex_config::McpServerConfig; use codex_connectors::parse_plugin_app_config; use codex_exec_server::ExecutorFileSystem; @@ -9,12 +10,9 @@ use codex_plugin::ResolvedPlugin; use codex_plugin::ResolvedPluginLocation; use codex_plugin::manifest::PluginManifestMcpServers; use codex_utils_path_uri::PathUri; -use codex_utils_path_uri::PathUriParseError; use std::io; -use thiserror::Error; use crate::ExecutorPluginProvider; -use crate::ExecutorPluginProviderError; use crate::ResolvedExecutorPlugin; const DEFAULT_MCP_CONFIG_FILE: &str = ".mcp.json"; @@ -27,59 +25,12 @@ pub struct ExecutorPluginRuntime { apps: Vec, } -/// Failure to project runtime capabilities from an executor plugin. -#[derive(Debug, Error)] -pub enum ExecutorPluginRuntimeError { - #[error(transparent)] - Resolve(#[from] ExecutorPluginProviderError), - #[error("failed to read MCP config for selected plugin `{plugin_id}` at `{path}`: {source}")] - ReadConfig { - plugin_id: String, - path: PathUri, - #[source] - source: io::Error, - }, - #[error( - "failed to resolve MCP config path `{relative_path}` below selected plugin `{plugin_id}` at `{root}`: {source}" - )] - InvalidConfigPath { - plugin_id: String, - root: PathUri, - relative_path: &'static str, - #[source] - source: PathUriParseError, - }, - #[error("failed to parse MCP config for selected plugin `{plugin_id}` at `{path}`: {source}")] - ParseConfig { - plugin_id: String, - path: PathUri, - #[source] - source: serde_json::Error, - }, - #[error("failed to read app config for selected plugin `{plugin_id}` at `{path}`: {source}")] - ReadAppConfig { - plugin_id: String, - path: PathUri, - #[source] - source: io::Error, - }, - #[error("failed to parse app config for selected plugin `{plugin_id}` at `{path}`: {source}")] - ParseAppConfig { - plugin_id: String, - path: PathUri, - #[source] - source: serde_json::Error, - }, -} - impl ExecutorPluginRuntime { /// Reads both runtime declaration files through the root's pinned filesystem. /// /// `Ok(None)` is intentionally not cacheable: the plugin manifest may appear /// once a deferred executor finishes starting. - pub async fn project( - root: &ResolvedSelectedCapabilityRoot, - ) -> Result, ExecutorPluginRuntimeError> { + pub async fn project(root: &ResolvedSelectedCapabilityRoot) -> anyhow::Result> { let Some(plugin) = ExecutorPluginProvider::resolve_pinned(root).await? else { return Ok(None); }; @@ -88,17 +39,7 @@ impl ExecutorPluginRuntime { } = plugin.plugin().location(); let mcp_servers = load_from_file_system(plugin.plugin(), plugin_root, plugin.file_system()).await?; - let apps = match load_apps(&plugin).await { - Ok(apps) => apps, - Err(err) => { - tracing::warn!( - plugin = plugin.plugin().selected_root_id(), - error = %err, - "ignoring invalid executor plugin app declarations" - ); - Vec::new() - } - }; + let apps = load_apps(&plugin).await?; Ok(Some(Self { plugin: plugin.plugin().clone(), mcp_servers, @@ -123,7 +64,7 @@ async fn load_from_file_system( plugin: &ResolvedPlugin, plugin_root: &PathUri, file_system: &dyn ExecutorFileSystem, -) -> Result, ExecutorPluginRuntimeError> { +) -> anyhow::Result> { let ResolvedPluginLocation::Environment { environment_id, .. } = plugin.location(); let plugin_id = plugin.selected_root_id(); let (contents, config_path) = match plugin.manifest().paths.mcp_servers.as_ref() { @@ -134,10 +75,10 @@ async fn load_from_file_system( file_system .read_file_text(path, /*sandbox*/ None) .await - .map_err(|source| ExecutorPluginRuntimeError::ReadConfig { - plugin_id: plugin_id.to_string(), - path: path.clone(), - source, + .with_context(|| { + format!( + "failed to read MCP config for selected plugin `{plugin_id}` at `{path}`" + ) })?, path.clone(), ) @@ -149,11 +90,10 @@ async fn load_from_file_system( None => { let config_path = plugin_root .join(DEFAULT_MCP_CONFIG_FILE) - .map_err(|source| ExecutorPluginRuntimeError::InvalidConfigPath { - plugin_id: plugin_id.to_string(), - root: plugin_root.clone(), - relative_path: DEFAULT_MCP_CONFIG_FILE, - source, + .with_context(|| { + format!( + "failed to resolve `{DEFAULT_MCP_CONFIG_FILE}` below selected plugin `{plugin_id}` at `{plugin_root}`" + ) })?; let contents = match file_system .read_file_text(&config_path, /*sandbox*/ None) @@ -164,23 +104,22 @@ async fn load_from_file_system( return Ok(Vec::new()); } Err(source) => { - return Err(ExecutorPluginRuntimeError::ReadConfig { - plugin_id: plugin_id.to_string(), - path: config_path.clone(), - source, + return Err(source).with_context(|| { + format!( + "failed to read MCP config for selected plugin `{plugin_id}` at `{config_path}`" + ) }); } }; (contents, config_path) } }; - let parsed = parse_executor_plugin_mcp_config(plugin_root, &contents, environment_id).map_err( - |source| ExecutorPluginRuntimeError::ParseConfig { - plugin_id: plugin_id.to_string(), - path: config_path, - source, - }, - )?; + let parsed = parse_executor_plugin_mcp_config(plugin_root, &contents, environment_id) + .with_context(|| { + format!( + "failed to parse MCP config for selected plugin `{plugin_id}` at `{config_path}`" + ) + })?; for error in parsed.errors { tracing::warn!( @@ -194,9 +133,7 @@ async fn load_from_file_system( Ok(parsed.servers.into_iter().collect()) } -async fn load_apps( - plugin: &ResolvedExecutorPlugin, -) -> Result, ExecutorPluginRuntimeError> { +async fn load_apps(plugin: &ResolvedExecutorPlugin) -> anyhow::Result> { let resolved_plugin = plugin.plugin(); let plugin_id = resolved_plugin.selected_root_id(); let Some(PluginResourceLocator::Environment { @@ -209,16 +146,12 @@ async fn load_apps( .file_system() .read_file_text(config_path, /*sandbox*/ None) .await - .map_err(|source| ExecutorPluginRuntimeError::ReadAppConfig { - plugin_id: plugin_id.to_string(), - path: config_path.clone(), - source, + .with_context(|| { + format!( + "failed to read app config for selected plugin `{plugin_id}` at `{config_path}`" + ) })?; - parse_plugin_app_config(&contents).map_err(|source| { - ExecutorPluginRuntimeError::ParseAppConfig { - plugin_id: plugin_id.to_string(), - path: config_path.clone(), - source, - } + parse_plugin_app_config(&contents).with_context(|| { + format!("failed to parse app config for selected plugin `{plugin_id}` at `{config_path}`") }) } diff --git a/codex-rs/core-plugins/src/lib.rs b/codex-rs/core-plugins/src/lib.rs index ad6d7ba352c4..585318d5e82c 100644 --- a/codex-rs/core-plugins/src/lib.rs +++ b/codex-rs/core-plugins/src/lib.rs @@ -40,7 +40,6 @@ pub use app_mcp_routing::apps_route_available; pub use discoverable::ToolSuggestDiscoverablePlugin; pub use discoverable::ToolSuggestPluginDiscoveryInput; pub use executor_runtime::ExecutorPluginRuntime; -pub use executor_runtime::ExecutorPluginRuntimeError; pub use loader::PluginHookLoadOutcome; pub use manager::ConfiguredMarketplace; pub use manager::ConfiguredMarketplaceListOutcome; diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index de8754ad920f..e39c9579dc8d 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -2257,6 +2257,7 @@ async fn spawn_thread_subagents_persist_parent_originator_across_new_and_truncat metrics_service_name: Some("codex_work_desktop".to_string()), parent_trace: None, environments: Vec::new(), + selected_capability_roots: Vec::new(), thread_extension_init: ExtensionDataInit::default(), supports_openai_form_elicitation: false, }) diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 1872ae3ecc9b..bbff67622eb3 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -120,6 +120,7 @@ pub(crate) async fn run_codex_thread_interactive( parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), parent_trace: None, environment_selections: parent_ctx.environments.to_selections(), + selected_capability_roots: Vec::new(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), supports_openai_form_elicitation: parent_session .services diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 46b38af151ca..3bdbd9503952 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -594,9 +594,20 @@ impl CodexThread { self.codex.session.runtime_mcp_config(config).await } - /// Returns the exact MCP config, environment bindings, and manager most recently published. - pub fn current_mcp_runtime(&self) -> Arc { - self.codex.session.services.latest_mcp_runtime() + /// Projects the latest host config through this thread's current selected executor bindings. + pub async fn project_mcp_config( + &self, + config: &crate::config::Config, + ) -> (codex_mcp::McpConfig, codex_mcp::McpRuntimeContext) { + self.codex.session.project_mcp_config(config).await + } + + /// Builds an MCP runtime from the latest host config and current selected executors. + pub async fn project_mcp_runtime( + &self, + config: &crate::config::Config, + ) -> Arc { + self.codex.session.project_mcp_runtime(config).await } pub fn multi_agent_version(&self) -> Option { diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 76b4a84830c6..01d0368e99e6 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -212,6 +212,18 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( force_refetch: bool, environment_manager: Arc, mcp_manager: Arc, +) -> anyhow::Result { + let mcp_config = mcp_manager.runtime_config(config).await; + let runtime_context = McpRuntimeContext::new(environment_manager, config.cwd.to_path_buf()); + list_accessible_connectors_from_mcp_config(config, force_refetch, mcp_config, runtime_context) + .await +} + +pub async fn list_accessible_connectors_from_mcp_config( + config: &Config, + force_refetch: bool, + mcp_config: codex_mcp::McpConfig, + runtime_context: McpRuntimeContext, ) -> anyhow::Result { let auth_manager = AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; @@ -226,7 +238,6 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( }); } let cache_key = accessible_connectors_cache_key(config, auth.as_ref()); - let mcp_config = mcp_manager.runtime_config(config).await; let tool_plugin_provenance = tool_plugin_provenance(&mcp_config); if !force_refetch && let Some(cached_connectors) = read_cached_accessible_connectors(&cache_key) { @@ -246,8 +257,6 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( }); } - let runtime_context = - McpRuntimeContext::new(Arc::clone(&environment_manager), config.cwd.to_path_buf()); let auth_status_entries = compute_auth_statuses( mcp_servers.iter(), config.mcp_oauth_credentials_store_mode, diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 90e6b01e758a..fa97a9d8ab31 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -3,8 +3,11 @@ use std::sync::Arc; use crate::config::Config; use codex_config::McpServerConfig; +use codex_connectors::ConnectorSnapshot; +use codex_connectors::PluginConnectorSource; +use codex_core_plugins::ExecutorPluginRuntime; use codex_core_plugins::PluginsManager; -use codex_extension_api::ExtensionDataInit; +use codex_core_plugins::apps_route_available; use codex_extension_api::ExtensionRegistry; use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; @@ -62,29 +65,7 @@ impl McpManager { /// Returns the MCP config after applying compatibility built-ins and /// runtime-only extension overlays. pub async fn runtime_config(&self, config: &Config) -> McpConfig { - self.runtime_config_with_context(config, /*thread_init*/ None) - .await - } - - pub(crate) async fn runtime_config_for_thread( - &self, - config: &Config, - thread_init: &ExtensionDataInit, - ) -> McpConfig { - self.runtime_config_with_context(config, Some(thread_init)) - .await - } - - async fn runtime_config_with_context( - &self, - config: &Config, - thread_init: Option<&ExtensionDataInit>, - ) -> McpConfig { - let context = match thread_init { - Some(thread_init) => McpServerContributionContext::for_thread(config, thread_init), - None => McpServerContributionContext::global(config), - }; - let mut selected_plugin_registrations = Vec::new(); + let context = McpServerContributionContext::global(config); let mut overlays = Vec::new(); // A contributor can emit multiple ordered actions, so order each action globally rather // than enumerating contributors. @@ -100,20 +81,6 @@ impl McpManager { config, }); } - McpServerContribution::SelectedPlugin { - name, - plugin_id, - plugin_display_name, - selection_order, - config, - } => selected_plugin_registrations.push( - McpServerRegistration::from_selected_plugin( - name, - McpPluginAttribution::new(plugin_id, plugin_display_name), - selection_order, - *config, - ), - ), McpServerContribution::Remove { name } => { overlays.push(OrderedMcpOverlay::Remove { contributor_id: contributor.id(), @@ -126,12 +93,7 @@ impl McpManager { } } - let mut mcp_config = config - .to_mcp_config_with_plugin_registrations( - self.plugins_manager.as_ref(), - selected_plugin_registrations, - ) - .await; + let mut mcp_config = config.to_mcp_config(self.plugins_manager.as_ref()).await; let mut catalog = mcp_config.mcp_server_catalog.to_builder(); if mcp_config.apps_enabled { catalog.register(McpServerRegistration::from_compatibility( @@ -182,6 +144,74 @@ impl McpManager { mcp_config } + /// Adds capabilities read from the executor bindings captured for one model step. + pub(crate) fn runtime_config_for_executor_plugins( + &self, + base_config: &McpConfig, + config: &Config, + plugins: &[(usize, ExecutorPluginRuntime)], + ) -> McpConfig { + // Step runtimes start from the clean global config, never the thread bootstrap view. + let mut mcp_config = base_config.clone(); + let mut catalog = mcp_config.mcp_server_catalog.to_builder(); + let selected_connector_snapshot = + ConnectorSnapshot::from_plugin_sources(plugins.iter().map(|(_, runtime)| { + let plugin = runtime.plugin(); + PluginConnectorSource::new( + plugin.manifest().display_name(), + runtime.apps().iter().cloned(), + ) + })); + let route_apps_over_mcp = config.orchestrator_mcp_enabled + && config + .features + .apps_enabled_for_auth(apps_route_available(self.plugins_manager.auth_mode())); + + for (selection_order, runtime) in plugins { + let plugin = runtime.plugin(); + let plugin_id = plugin.selected_root_id().to_string(); + let display_name = plugin.manifest().display_name().to_string(); + let mut servers = runtime + .mcp_servers() + .iter() + .cloned() + .collect::>(); + config.apply_plugin_mcp_server_requirements(&plugin_id, &mut servers); + let mut servers = servers.into_iter().collect::>(); + servers.sort_unstable_by(|left, right| left.0.cmp(&right.0)); + let attribution = McpPluginAttribution::new(plugin_id.clone(), display_name.clone()); + for (name, server) in servers { + let has_app = runtime + .apps() + .iter() + .any(|app| app.name == name && !app.connector_id.0.trim().is_empty()); + if !route_apps_over_mcp || !has_app { + catalog.register(McpServerRegistration::from_selected_plugin( + name, + attribution.clone(), + *selection_order, + server, + )); + } + } + } + + let catalog = catalog.build(); + for conflict in catalog.conflicts() { + tracing::warn!( + server = conflict.name, + outcome = ?conflict.outcome, + contenders = ?conflict.contenders, + "conflicting selected MCP server actions; using resolved catalog outcome" + ); + } + mcp_config.mcp_server_catalog = catalog; + mcp_config.connector_snapshot = mcp_config + .connector_snapshot + .merged_with(&selected_connector_snapshot); + mcp_config + } + /// Returns config- and plugin-backed servers without runtime contributions. pub async fn configured_servers(&self, config: &Config) -> HashMap { let mcp_config = config.to_mcp_config(self.plugins_manager.as_ref()).await; diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index 91fed2d41226..0be53affdea2 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -36,6 +36,7 @@ pub(crate) async fn maybe_prompt_and_install_mcp_dependencies( turn_context: &TurnContext, cancellation_token: &CancellationToken, mentioned_skills: &[SkillCatalogEntry], + available_mcp_servers: &HashMap, elicitation_reviewer: Option, ) -> bool { let originator_value = originator().value; @@ -53,8 +54,7 @@ pub(crate) async fn maybe_prompt_and_install_mcp_dependencies( return false; } - let installed = sess.runtime_mcp_servers(config.as_ref()).await; - let missing = collect_missing_mcp_dependencies(mentioned_skills, &installed); + let missing = collect_missing_mcp_dependencies(mentioned_skills, available_mcp_servers); if missing.is_empty() { return false; } @@ -72,6 +72,7 @@ pub(crate) async fn maybe_prompt_and_install_mcp_dependencies( turn_context, config.as_ref(), mentioned_skills, + available_mcp_servers, elicitation_reviewer, ) .await; @@ -84,6 +85,7 @@ pub(crate) async fn maybe_install_mcp_dependencies( turn_context: &TurnContext, config: &crate::config::Config, mentioned_skills: &[SkillCatalogEntry], + available_mcp_servers: &HashMap, elicitation_reviewer: Option, ) -> bool { if mentioned_skills.is_empty() @@ -95,8 +97,7 @@ pub(crate) async fn maybe_install_mcp_dependencies( } let codex_home = config.codex_home.clone(); - let installed = sess.runtime_mcp_servers(config).await; - let missing = collect_missing_mcp_dependencies(mentioned_skills, &installed); + let missing = collect_missing_mcp_dependencies(mentioned_skills, available_mcp_servers); if missing.is_empty() { return false; } diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index a866116a2484..cac9bb41ae0d 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -64,7 +64,7 @@ pub async fn build_prompt_input( ); let thread = thread_manager.start_thread(config).await?; - let output = build_prompt_input_from_session(thread.thread.codex.session.as_ref(), input).await; + let output = build_prompt_input_from_session(&thread.thread.codex.session, input).await; let shutdown = thread.thread.shutdown_and_wait().await; let _removed = thread_manager.remove_thread(&thread.thread_id).await; @@ -73,7 +73,7 @@ pub async fn build_prompt_input( } pub(crate) async fn build_prompt_input_from_session( - sess: &Session, + sess: &Arc, input: Vec, ) -> CodexResult> { let turn_context = sess.new_default_turn().await; diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 9b15bd16879d..c4ed7fe6bcbc 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -1,3 +1,4 @@ +use super::mcp_projection::McpRuntimeScope; use super::*; use codex_mcp::ElicitationReviewRequest; use codex_mcp::ElicitationReviewer; @@ -75,17 +76,7 @@ impl ElicitationReviewer for GuardianMcpElicitationReviewer { impl Session { pub(crate) async fn runtime_mcp_config(&self, config: &Config) -> McpConfig { - self.services - .mcp_manager - .runtime_config_for_thread(config, &self.services.mcp_thread_init) - .await - } - - pub(crate) async fn runtime_mcp_servers( - &self, - config: &Config, - ) -> HashMap { - codex_mcp::configured_mcp_servers(&self.runtime_mcp_config(config).await) + self.services.mcp_manager.runtime_config(config).await } pub(crate) fn mcp_elicitation_reviewer(self: &Arc) -> ElicitationReviewerHandle { @@ -200,35 +191,27 @@ impl Session { return Ok(()); } - self.services - .latest_mcp_runtime() - .manager_arc() - .resolve_elicitation(server_name, id, response) - .await - } - - pub async fn list_resources( - &self, - server: &str, - params: Option, - ) -> anyhow::Result { - self.services - .latest_mcp_runtime() - .manager_arc() - .list_resources(server, params) - .await - } - - pub async fn list_resource_templates( - &self, - server: &str, - params: Option, - ) -> anyhow::Result { - self.services - .latest_mcp_runtime() - .manager_arc() - .list_resource_templates(server, params) - .await + let managers = { + let mut managers = self + .services + .mcp_elicitation_managers + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + managers.retain(|manager| manager.strong_count() > 0); + managers + .iter() + .filter_map(std::sync::Weak::upgrade) + .collect::>() + }; + for manager in managers { + if manager + .try_resolve_elicitation(server_name.clone(), id.clone(), response.clone()) + .await? + { + return Ok(()); + } + } + Err(anyhow::anyhow!("elicitation request not found")) } pub async fn read_resource( @@ -263,69 +246,20 @@ impl Session { refresh_config: &Config, elicitation_reviewer: Option, ) { - let auth = self.services.auth_manager.auth().await; let mcp_config = Arc::new(self.runtime_mcp_config(refresh_config).await); - let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); - let mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref()); - let environment_manager = self.services.turn_environments.environment_manager(); - // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment cwd - // values can be used without falling back to the legacy host cwd. - let cwd = turn_context - .environments - .primary() - .and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok()) - .map(|cwd| cwd.to_path_buf()) - .unwrap_or_else(|| { - #[allow(deprecated)] - turn_context.cwd.to_path_buf() - }); - let mcp_runtime_context = McpRuntimeContext::new(environment_manager, cwd); - let auth_statuses = compute_auth_statuses( - mcp_servers.iter(), - mcp_config.mcp_oauth_credentials_store_mode, - mcp_config.auth_keyring_backend_kind, - auth.as_ref(), - &mcp_runtime_context, - ) - .await; - let mcp_startup_cancellation_token = { - let mut guard = self.services.mcp_startup_cancellation_token.lock().await; - // The previous runtime owns the old token and may still be serving an in-flight step. - // Its manager cancels that token when the last runtime handle is dropped. - let cancellation_token = CancellationToken::new(); - *guard = cancellation_token.clone(); - cancellation_token - }; - let refreshed_manager = McpConnectionManager::new( - &mcp_servers, - mcp_config.mcp_oauth_credentials_store_mode, - mcp_config.auth_keyring_backend_kind, - auth_statuses, - &turn_context.approval_policy, - turn_context.sub_id.clone(), - self.get_tx_event(), - mcp_startup_cancellation_token, - turn_context.permission_profile(), - mcp_runtime_context.clone(), - mcp_config.codex_home.clone(), - codex_apps_tools_cache_key(auth.as_ref()), - mcp_config.prefix_mcp_tool_names, - mcp_config.client_elicitation_capability.clone(), - self.services - .supports_openai_form_elicitation - .load(std::sync::atomic::Ordering::Relaxed), - tool_plugin_provenance, - auth.as_ref(), - elicitation_reviewer, - ) - .await; - { - let current_manager = self.services.latest_mcp_runtime(); - refreshed_manager - .set_elicitations_auto_deny(current_manager.manager().elicitations_auto_deny()); - } - self.services - .publish_mcp_runtime(mcp_config, mcp_runtime_context, refreshed_manager); + let configured_servers = codex_mcp::configured_mcp_servers(mcp_config.as_ref()); + let runtime_context = + self.mcp_runtime_context(&turn_context.config, &turn_context.environments, &[]); + let runtime = self + .build_mcp_runtime( + McpRuntimeScope::Turn(turn_context), + mcp_config, + configured_servers, + runtime_context, + elicitation_reviewer, + ) + .await; + self.services.replace_base_with_runtime(runtime).await; } pub(crate) async fn refresh_mcp_servers_if_requested( diff --git a/codex-rs/core/src/session/mcp_projection.rs b/codex-rs/core/src/session/mcp_projection.rs new file mode 100644 index 000000000000..30713e745a15 --- /dev/null +++ b/codex-rs/core/src/session/mcp_projection.rs @@ -0,0 +1,319 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use super::INITIAL_SUBMIT_ID; +use super::McpRuntimeSnapshot; +use super::Session; +use super::TurnContext; +use crate::config::Config; +use crate::environment_selection::TurnEnvironmentSnapshot; +use codex_config::McpServerConfig; +use codex_core_plugins::ExecutorPluginRuntime; +use codex_exec_server::ResolvedSelectedCapabilityRoot; +use codex_mcp::ElicitationReviewerHandle; +use codex_mcp::McpConfig; +use codex_mcp::McpConnectionManager; +use codex_mcp::McpRuntimeContext; +use codex_mcp::codex_apps_tools_cache_key; +use codex_mcp::compute_auth_statuses; +use codex_mcp::effective_mcp_servers_from_configured; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use tokio_util::sync::CancellationToken; + +pub(super) enum McpRuntimeScope<'a> { + Turn(&'a TurnContext), + Thread(&'a Config), +} + +struct ProjectedMcpConfig { + bindings: Vec<(usize, ResolvedSelectedCapabilityRoot)>, + plugins: Vec<(usize, ExecutorPluginRuntime)>, + config: McpConfig, + runtime_context: McpRuntimeContext, +} + +impl Session { + pub(crate) async fn mcp_runtime_for_step( + self: &Arc, + turn_context: &TurnContext, + environments: &TurnEnvironmentSnapshot, + selected_roots: &[SelectedCapabilityRoot], + resolved_roots: &[ResolvedSelectedCapabilityRoot], + ) -> Arc { + let mut cache = self.services.selected_mcp_runtime.lock().await; + if selected_roots.is_empty() { + let base = cache.base_runtime(); + self.services + .publish_existing_mcp_runtime(Arc::clone(&base)); + return base; + } + let bindings = selected_bindings(selected_roots, resolved_roots); + let mut plugins = cache.plugins(&bindings).unwrap_or_default(); + if let Some(runtime) = cache.runtime(&bindings) { + let unresolved = bindings + .iter() + .filter(|(order, _)| { + !plugins + .iter() + .any(|(plugin_order, _)| plugin_order == order) + }) + .cloned() + .collect::>(); + let discovered = project_executor_plugins(&unresolved).await; + if discovered.is_empty() { + self.services + .publish_existing_mcp_runtime(Arc::clone(&runtime)); + return runtime; + } + plugins.extend(discovered); + plugins.sort_unstable_by_key(|(order, _)| *order); + } else { + plugins = project_executor_plugins(&bindings).await; + } + + let base = cache.base_runtime(); + let mcp_config = Arc::new( + self.services + .mcp_manager + .runtime_config_for_executor_plugins(base.config(), &turn_context.config, &plugins), + ); + let configured_servers = codex_mcp::configured_mcp_servers(mcp_config.as_ref()); + let pinned_roots = bindings + .iter() + .map(|(_, root)| root.clone()) + .collect::>(); + let runtime_context = + self.mcp_runtime_context(&turn_context.config, environments, &pinned_roots); + let runtime = if plugins.is_empty() { + base + } else { + self.build_mcp_runtime( + McpRuntimeScope::Turn(turn_context), + mcp_config, + configured_servers, + runtime_context, + Some(self.mcp_elicitation_reviewer()), + ) + .await + }; + cache.insert_runtime(bindings, plugins, Arc::clone(&runtime)); + self.services + .publish_existing_mcp_runtime(Arc::clone(&runtime)); + runtime + } + + pub(crate) async fn project_mcp_config( + &self, + config: &Config, + ) -> (McpConfig, McpRuntimeContext) { + let projection = self.project_mcp_config_inner(config).await; + (projection.config, projection.runtime_context) + } + + async fn project_mcp_config_inner(&self, config: &Config) -> ProjectedMcpConfig { + let environments = self.services.turn_environments.snapshot().await; + let selected_roots = &self.services.selected_capability_roots; + let resolved_roots = self + .services + .turn_environments + .environment_manager() + .resolve_selected_capability_roots( + selected_roots, + &environments.captured_environments(), + ) + .await; + let bindings = selected_bindings(selected_roots, &resolved_roots); + let plugins = project_executor_plugins(&bindings).await; + let base_config = self.services.mcp_manager.runtime_config(config).await; + let mcp_config = self + .services + .mcp_manager + .runtime_config_for_executor_plugins(&base_config, config, &plugins); + let runtime_context = self.mcp_runtime_context(config, &environments, &resolved_roots); + ProjectedMcpConfig { + bindings, + plugins, + config: mcp_config, + runtime_context, + } + } + + pub(crate) async fn project_mcp_runtime( + self: &Arc, + config: &Config, + ) -> Arc { + let mut cache = self.services.selected_mcp_runtime.lock().await; + let projection = self.project_mcp_config_inner(config).await; + let current = cache + .runtime(&projection.bindings) + .unwrap_or_else(|| self.services.latest_mcp_runtime()); + if current.matches_projection(&projection.config, &projection.runtime_context) { + current + .manager() + .set_approval_policy(&config.permissions.approval_policy); + current + .manager() + .set_permission_profile(config.permissions.permission_profile().clone()); + return current; + } + let mcp_config = Arc::new(projection.config); + let configured_servers = codex_mcp::configured_mcp_servers(mcp_config.as_ref()); + let runtime = self + .build_mcp_runtime( + McpRuntimeScope::Thread(config), + mcp_config, + configured_servers, + projection.runtime_context, + Some(self.mcp_elicitation_reviewer()), + ) + .await; + if projection.bindings.is_empty() { + cache.replace_base(Arc::clone(&runtime)); + } else { + cache.insert_runtime( + projection.bindings, + projection.plugins, + Arc::clone(&runtime), + ); + } + self.services + .publish_existing_mcp_runtime(Arc::clone(&runtime)); + runtime + } + + pub(super) async fn build_mcp_runtime( + &self, + scope: McpRuntimeScope<'_>, + mcp_config: Arc, + configured_servers: HashMap, + mcp_runtime_context: McpRuntimeContext, + elicitation_reviewer: Option, + ) -> Arc { + let auth = self.services.auth_manager.auth().await; + let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); + let mcp_servers = effective_mcp_servers_from_configured( + configured_servers, + mcp_config.as_ref(), + auth.as_ref(), + ); + let auth_statuses = compute_auth_statuses( + mcp_servers.iter(), + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, + auth.as_ref(), + &mcp_runtime_context, + ) + .await; + let (approval_policy, submit_id, permission_profile, publish_startup_token) = match scope { + McpRuntimeScope::Turn(turn_context) => ( + &turn_context.approval_policy, + turn_context.sub_id.clone(), + turn_context.permission_profile(), + true, + ), + McpRuntimeScope::Thread(config) => ( + &config.permissions.approval_policy, + INITIAL_SUBMIT_ID.to_string(), + config.permissions.permission_profile().clone(), + false, + ), + }; + let mcp_startup_cancellation_token = CancellationToken::new(); + if publish_startup_token { + let mut guard = self.services.mcp_startup_cancellation_token.lock().await; + // The previous runtime owns the old token and may still be serving an in-flight step. + // Its manager cancels that token when the last runtime handle is dropped. + *guard = mcp_startup_cancellation_token.clone(); + } + let manager = McpConnectionManager::new( + &mcp_servers, + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, + auth_statuses, + approval_policy, + submit_id, + self.get_tx_event(), + mcp_startup_cancellation_token, + permission_profile, + mcp_runtime_context.clone(), + mcp_config.codex_home.clone(), + codex_apps_tools_cache_key(auth.as_ref()), + mcp_config.prefix_mcp_tool_names, + mcp_config.client_elicitation_capability.clone(), + self.services + .supports_openai_form_elicitation + .load(std::sync::atomic::Ordering::Relaxed), + tool_plugin_provenance, + auth.as_ref(), + elicitation_reviewer, + ) + .await; + let current_manager = self.services.latest_mcp_runtime(); + manager.set_elicitations_auto_deny(current_manager.manager().elicitations_auto_deny()); + let runtime = Arc::new(McpRuntimeSnapshot::new( + mcp_config, + Arc::new(manager), + mcp_runtime_context, + )); + self.services + .track_mcp_elicitation_manager(&runtime.manager_arc()); + runtime + } + + pub(super) fn mcp_runtime_context( + &self, + config: &Config, + environments: &TurnEnvironmentSnapshot, + pinned_roots: &[ResolvedSelectedCapabilityRoot], + ) -> McpRuntimeContext { + // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment cwd + // values can be used without falling back to the legacy host cwd. + let cwd = environments + .primary() + .and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok()) + .map(|cwd| cwd.to_path_buf()) + .unwrap_or_else(|| config.cwd.to_path_buf()); + McpRuntimeContext::new(self.services.turn_environments.environment_manager(), cwd) + .with_pinned_environments(pinned_roots.iter().map(|root| { + let CapabilityRootLocation::Environment { environment_id, .. } = + &root.selected_root().location; + (environment_id.clone(), Arc::clone(root.environment())) + })) + } +} + +fn selected_bindings( + selected_roots: &[SelectedCapabilityRoot], + resolved_roots: &[ResolvedSelectedCapabilityRoot], +) -> Vec<(usize, ResolvedSelectedCapabilityRoot)> { + resolved_roots + .iter() + .filter_map(|root| { + selected_roots + .iter() + .position(|selected| selected == root.selected_root()) + .map(|selection_order| (selection_order, root.clone())) + }) + .collect() +} + +async fn project_executor_plugins( + bindings: &[(usize, ResolvedSelectedCapabilityRoot)], +) -> Vec<(usize, ExecutorPluginRuntime)> { + let mut plugins = Vec::new(); + for (selection_order, root) in bindings { + match ExecutorPluginRuntime::project(root).await { + Ok(Some(runtime)) => plugins.push((*selection_order, runtime)), + Ok(None) => {} + Err(err) => { + tracing::warn!( + selected_root = root.selected_root().id, + error = %err, + "failed to project selected executor plugin runtime" + ); + } + } + } + plugins +} diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index dc5ea011ca32..959607bb5537 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -1,94 +1,85 @@ -use std::fmt; use std::sync::Arc; -use codex_mcp::McpConfig; -use codex_mcp::McpConnectionManager; -use codex_mcp::McpRuntimeContext; +use codex_core_plugins::ExecutorPluginRuntime; +use codex_exec_server::ResolvedSelectedCapabilityRoot; +use codex_mcp::McpRuntimeSnapshot; -/// MCP config, exact environment bindings, and manager used by one model request. -pub struct McpRuntimeSnapshot { - config: Arc, - manager: Arc, - runtime_context: McpRuntimeContext, +#[derive(Default)] +pub(crate) struct SelectedMcpRuntimeCache { + base_runtime: Option>, + runtime: Option, } -impl McpRuntimeSnapshot { - pub(crate) fn new( - config: Arc, - manager: Arc, - runtime_context: McpRuntimeContext, - ) -> Self { - Self { - config, - manager, - runtime_context, - } - } +struct CachedSelectedRuntime { + bindings: Vec<(usize, ResolvedSelectedCapabilityRoot)>, + plugins: Vec<(usize, ExecutorPluginRuntime)>, + runtime: Arc, +} - pub fn config(&self) -> &McpConfig { - self.config.as_ref() +impl SelectedMcpRuntimeCache { + pub(crate) fn replace_base(&mut self, runtime: Arc) { + self.base_runtime = Some(runtime); + self.runtime = None; } - pub fn manager(&self) -> &McpConnectionManager { - self.manager.as_ref() + pub(crate) fn base_runtime(&self) -> Arc { + self.base_runtime + .as_ref() + .map(Arc::clone) + .expect("base MCP runtime must be installed before capturing a step") } - pub(crate) fn manager_arc(&self) -> Arc { - Arc::clone(&self.manager) + pub(crate) fn runtime( + &self, + bindings: &[(usize, ResolvedSelectedCapabilityRoot)], + ) -> Option> { + self.runtime + .as_ref() + .filter(|cached| same_bindings(&cached.bindings, bindings)) + .map(|cached| Arc::clone(&cached.runtime)) } - pub fn runtime_context(&self) -> &McpRuntimeContext { - &self.runtime_context + pub(crate) fn plugins( + &self, + bindings: &[(usize, ResolvedSelectedCapabilityRoot)], + ) -> Option> { + self.runtime + .as_ref() + .filter(|cached| same_bindings(&cached.bindings, bindings)) + .map(|cached| cached.plugins.clone()) } - #[cfg(test)] - pub(crate) fn new_uninitialized_for_test(config: &crate::config::Config) -> Arc { - use codex_exec_server::EnvironmentManager; - use codex_features::Feature; - use codex_mcp::ResolvedMcpCatalog; - use rmcp::model::ElicitationCapability; - - let mcp_config = McpConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - apps_mcp_product_sku: config.apps_mcp_product_sku.clone(), - codex_home: config.codex_home.to_path_buf(), - mcp_oauth_credentials_store_mode: config.mcp_oauth_credentials_store_mode, - auth_keyring_backend_kind: config.auth_keyring_backend_kind(), - mcp_oauth_callback_port: config.mcp_oauth_callback_port, - mcp_oauth_callback_url: config.mcp_oauth_callback_url.clone(), - skill_mcp_dependency_install_enabled: config - .features - .enabled(Feature::SkillMcpDependencyInstall), - approval_policy: config.permissions.approval_policy.clone(), - codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(), - use_legacy_landlock: config.features.use_legacy_landlock(), - apps_enabled: config.features.enabled(Feature::Apps), - prefix_mcp_tool_names: config.prefix_mcp_tool_names(), - client_elicitation_capability: ElicitationCapability::default(), - mcp_server_catalog: ResolvedMcpCatalog::default(), - connector_snapshot: codex_connectors::ConnectorSnapshot::default(), - }; - let manager = McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ); - let runtime_context = McpRuntimeContext::new( - Arc::new(EnvironmentManager::default_for_tests()), - config.cwd.to_path_buf(), - ); - Arc::new(Self::new( - Arc::new(mcp_config), - Arc::new(manager), - runtime_context, - )) + pub(crate) fn insert_runtime( + &mut self, + bindings: Vec<(usize, ResolvedSelectedCapabilityRoot)>, + plugins: Vec<(usize, ExecutorPluginRuntime)>, + runtime: Arc, + ) { + self.runtime = Some(CachedSelectedRuntime { + bindings, + plugins, + runtime, + }); } } -impl fmt::Debug for McpRuntimeSnapshot { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("McpRuntimeSnapshot") - .finish_non_exhaustive() - } +fn same_bindings( + left: &[(usize, ResolvedSelectedCapabilityRoot)], + right: &[(usize, ResolvedSelectedCapabilityRoot)], +) -> bool { + left.len() == right.len() + && left + .iter() + .zip(right) + .all(|((left_order, left), (right_order, right))| { + left_order == right_order && same_binding(left, right) + }) +} + +fn same_binding( + left: &ResolvedSelectedCapabilityRoot, + right: &ResolvedSelectedCapabilityRoot, +) -> bool { + left.selected_root() == right.selected_root() + && Arc::ptr_eq(left.environment(), right.environment()) } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 45dc1e68804e..ccf27705919f 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -85,6 +85,7 @@ use codex_protocol::approvals::ElicitationRequestEvent; use codex_protocol::approvals::ExecPolicyAmendment; use codex_protocol::approvals::NetworkPolicyAmendment; use codex_protocol::approvals::NetworkPolicyRuleAction; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::AutoCompactTokenLimitScope; use codex_protocol::config_types::ModeKind; @@ -154,9 +155,6 @@ use codex_utils_path_uri::PathUri; use futures::future::BoxFuture; use futures::future::Shared; use futures::prelude::*; -use rmcp::model::ListResourceTemplatesResult; -use rmcp::model::ListResourcesResult; -use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; @@ -207,6 +205,7 @@ mod handlers; mod inject; mod input_queue; mod mcp; +mod mcp_projection; mod mcp_runtime; pub(crate) mod multi_agents; mod review; @@ -229,7 +228,7 @@ use self::handlers::submission_loop; pub(crate) use self::input_queue::InputQueueActivity; pub(crate) use self::input_queue::TurnInput; pub(crate) use self::input_queue::TurnInputQueue; -pub use self::mcp_runtime::McpRuntimeSnapshot; +pub(crate) use self::mcp_runtime::SelectedMcpRuntimeCache; use self::review::spawn_review_thread; use self::session::AppServerClientMetadata; use self::session::Session; @@ -242,6 +241,9 @@ use self::turn::collect_explicit_app_ids_from_skill_items; use self::turn::realtime_text_for_event; use self::turn_context::TurnContext; use self::turn_context::TurnSkillsContext; +pub(crate) use codex_mcp::McpRuntimeSnapshot; +#[cfg(test)] +pub(crate) use tests::uninitialized_mcp_runtime; #[cfg(test)] mod rollout_reconstruction_tests; @@ -335,7 +337,6 @@ use codex_core_plugins::RecommendedPluginCandidatesInput; use codex_git_utils::get_git_repo_root; use codex_mcp::McpConfig; use codex_mcp::compute_auth_statuses; -use codex_mcp::effective_mcp_servers; use codex_otel::SessionTelemetry; use codex_otel::THREAD_STARTED_METRIC; use codex_otel::TelemetryAuthMode; @@ -437,6 +438,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) user_shell_override: Option, pub(crate) parent_trace: Option, pub(crate) environment_selections: Vec, + pub(crate) selected_capability_roots: Vec, pub(crate) thread_extension_init: ExtensionDataInit, pub(crate) supports_openai_form_elicitation: bool, pub(crate) analytics_events_client: Option, @@ -522,6 +524,7 @@ impl Codex { parent_rollout_thread_trace, parent_trace: _, environment_selections, + selected_capability_roots, thread_extension_init, supports_openai_form_elicitation, analytics_events_client, @@ -666,6 +669,7 @@ impl Codex { plugins_manager, mcp_manager.clone(), extensions, + selected_capability_roots, thread_extension_init, supports_openai_form_elicitation, agent_control, @@ -2806,7 +2810,7 @@ impl Session { /// `run_turn` and pass the result down; standalone request or history boundaries may capture /// their own step. pub(crate) async fn capture_step_context( - &self, + self: &Arc, turn_context: Arc, ) -> Arc { let deferred_executor_enabled = turn_context @@ -2826,16 +2830,24 @@ impl Session { .await; } let loaded_agents_md = self.services.agents_md_manager.get_loaded().await; + let selected_roots = &self.services.selected_capability_roots; let selected_capability_roots = self .services .turn_environments .environment_manager() .resolve_selected_capability_roots( - &self.services.selected_capability_roots, + selected_roots, &environments.captured_environment_availability(), ) .await; - let mcp = self.services.latest_mcp_runtime(); + let mcp = self + .mcp_runtime_for_step( + turn_context.as_ref(), + &environments, + selected_roots, + &selected_capability_roots, + ) + .await; let extra_skill_sources = self .services .thread_extension_data diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 33b24a29a018..8ab001294a1b 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -490,6 +490,7 @@ impl Session { plugins_manager: Arc, mcp_manager: Arc, extensions: Arc>, + selected_capability_roots: Vec, thread_extension_init: ExtensionDataInit, supports_openai_form_elicitation: bool, agent_control: AgentControl, @@ -556,11 +557,11 @@ impl Session { config.current_time_reminder.as_ref(), external_time_provider, )?; - let selected_capability_roots = thread_extension_init - .get::>() - .map(|roots| roots.as_ref().clone()) - .unwrap_or_else(|| initial_history.get_selected_capability_roots()); - let mcp_thread_init = thread_extension_init.clone(); + let selected_capability_roots = if selected_capability_roots.is_empty() { + initial_history.get_selected_capability_roots() + } else { + selected_capability_roots + }; let thread_extension_data = codex_extension_api::ExtensionData::new_with_init( thread_id.to_string(), thread_extension_init, @@ -653,7 +654,6 @@ impl Session { let auth_manager_clone = Arc::clone(&auth_manager); let config_for_mcp = Arc::clone(&config); let mcp_manager_for_mcp = Arc::clone(&mcp_manager); - let mcp_thread_init_for_startup = &mcp_thread_init; let mcp_runtime_context = McpRuntimeContext::new( Arc::clone(&environment_manager), session_configuration.cwd().to_path_buf(), @@ -661,9 +661,7 @@ impl Session { let mcp_runtime_context_for_auth = mcp_runtime_context.clone(); let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; - let mcp_config = mcp_manager_for_mcp - .runtime_config_for_thread(&config_for_mcp, mcp_thread_init_for_startup) - .await; + let mcp_config = mcp_manager_for_mcp.runtime_config(&config_for_mcp).await; let mcp_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref()); let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); let auth_statuses = compute_auth_statuses( @@ -993,20 +991,10 @@ impl Session { config.analytics_enabled, ) }); - // Keep one stable manager handle for the session so extension resource clients - // automatically observe the manager installed at startup and on later refreshes. - let mcp_connection_manager = Arc::new(arc_swap::ArcSwap::from_pointee( - McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), - )); + let mcp_runtime = Arc::new(arc_swap::ArcSwapOption::empty()); let session_extension_data = codex_extension_api::ExtensionData::new(session_id.to_string()); - session_extension_data.insert(McpResourceClient::new(Arc::clone( - &mcp_connection_manager, - ))); + session_extension_data.insert(McpResourceClient::new(Arc::clone(&mcp_runtime))); for contributor in extensions.thread_lifecycle_contributors() { contributor.on_thread_start(codex_extension_api::ThreadStartInput { config: config.as_ref(), @@ -1019,15 +1007,9 @@ impl Session { } let services = SessionServices { - // Initialize the MCP connection manager with an uninitialized - // instance. It will be replaced with one created via - // McpConnectionManager::new() once all its constructor args are - // available. This also ensures `SessionConfigured` is emitted - // before any MCP-related events. It is reasonable to consider - // changing this to use Option or OnceCell, though the current - // setup is straightforward enough and performs well. - mcp_connection_manager, - mcp_runtime: arc_swap::ArcSwapOption::empty(), + mcp_runtime, + mcp_elicitation_managers: std::sync::Mutex::new(Vec::new()), + selected_mcp_runtime: Mutex::new(Default::default()), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -1057,7 +1039,6 @@ impl Session { session_extension_data, thread_extension_data, selected_capability_roots, - mcp_thread_init, supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new( supports_openai_form_elicitation, ), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 6444d2545698..d931216a9299 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -197,12 +197,51 @@ impl StepContext { environments, Vec::new(), skills, - crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config), + uninitialized_mcp_runtime(&turn.config), /*loaded_agents_md*/ None, )) } } +pub(crate) fn uninitialized_mcp_runtime( + config: &Config, +) -> Arc { + let mcp_config = codex_mcp::McpConfig { + chatgpt_base_url: config.chatgpt_base_url.clone(), + apps_mcp_product_sku: config.apps_mcp_product_sku.clone(), + codex_home: config.codex_home.to_path_buf(), + mcp_oauth_credentials_store_mode: config.mcp_oauth_credentials_store_mode, + auth_keyring_backend_kind: config.auth_keyring_backend_kind(), + mcp_oauth_callback_port: config.mcp_oauth_callback_port, + mcp_oauth_callback_url: config.mcp_oauth_callback_url.clone(), + skill_mcp_dependency_install_enabled: config + .features + .enabled(Feature::SkillMcpDependencyInstall), + approval_policy: config.permissions.approval_policy.clone(), + codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(), + use_legacy_landlock: config.features.use_legacy_landlock(), + apps_enabled: config.features.enabled(Feature::Apps), + prefix_mcp_tool_names: config.prefix_mcp_tool_names(), + client_elicitation_capability: rmcp::model::ElicitationCapability::default(), + mcp_server_catalog: codex_mcp::ResolvedMcpCatalog::default(), + connector_snapshot: codex_connectors::ConnectorSnapshot::default(), + }; + let manager = codex_mcp::McpConnectionManager::new_uninitialized_with_permission_profile( + &config.permissions.approval_policy, + config.permissions.permission_profile(), + config.prefix_mcp_tool_names(), + ); + let runtime_context = codex_mcp::McpRuntimeContext::new( + Arc::new(EnvironmentManager::default_for_tests()), + config.cwd.to_path_buf(), + ); + Arc::new(codex_mcp::McpRuntimeSnapshot::new( + Arc::new(mcp_config), + Arc::new(manager), + runtime_context, + )) +} + mod guardian_tests; struct InstructionsTestCase { @@ -5370,11 +5409,13 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); - let mcp_runtime = - crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); + let mcp_runtime = uninitialized_mcp_runtime(config.as_ref()); + let mut selected_mcp_runtime = crate::session::SelectedMcpRuntimeCache::default(); + selected_mcp_runtime.replace_base(Arc::clone(&mcp_runtime)); let services = SessionServices { - mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), - mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), + mcp_runtime: Arc::new(arc_swap::ArcSwapOption::from(Some(mcp_runtime))), + mcp_elicitation_managers: std::sync::Mutex::new(Vec::new()), + selected_mcp_runtime: Mutex::new(selected_mcp_runtime), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -5412,7 +5453,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { ), thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()), selected_capability_roots: Vec::new(), - mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new(false), agent_control, network_proxy: arc_swap::ArcSwapOption::from(None), @@ -7445,11 +7485,13 @@ where /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); - let mcp_runtime = - crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); + let mcp_runtime = uninitialized_mcp_runtime(config.as_ref()); + let mut selected_mcp_runtime = crate::session::SelectedMcpRuntimeCache::default(); + selected_mcp_runtime.replace_base(Arc::clone(&mcp_runtime)); let services = SessionServices { - mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), - mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), + mcp_runtime: Arc::new(arc_swap::ArcSwapOption::from(Some(mcp_runtime))), + mcp_elicitation_managers: std::sync::Mutex::new(Vec::new()), + selected_mcp_runtime: Mutex::new(selected_mcp_runtime), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -7487,7 +7529,6 @@ where ), thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()), selected_capability_roots: Vec::new(), - mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new(false), agent_control, network_proxy: arc_swap::ArcSwapOption::from(None), @@ -7606,14 +7647,8 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( } #[tokio::test] -async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { +async fn refresh_mcp_servers_is_deferred_until_next_turn() { let (session, turn_context) = make_session_and_context().await; - let turn_context = Arc::new(turn_context); - let old_runtime = session.services.latest_mcp_runtime(); - let step_context = session - .capture_step_context(Arc::clone(&turn_context)) - .await; - assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime)); let old_token = session.mcp_startup_cancellation_token().await; assert!(!old_token.is_cancelled()); @@ -7654,9 +7689,6 @@ async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { ); let new_token = session.mcp_startup_cancellation_token().await; assert!(!new_token.is_cancelled()); - let new_runtime = session.services.latest_mcp_runtime(); - assert!(!Arc::ptr_eq(&old_runtime, &new_runtime)); - assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime)); } #[tokio::test] diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 64c4ba29f2bd..aedbd0ebc133 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -744,6 +744,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { user_shell_override: None, parent_trace: None, environment_selections: Vec::new(), + selected_capability_roots: Vec::new(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), supports_openai_form_elicitation: false, analytics_events_client: None, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 448fb7186966..c7d37da387e0 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -180,14 +180,13 @@ pub(crate) async fn run_turn( } let mut accepted_user_input = user_input_from_turn_input(&accepted_input); - let Some((injection_items, explicitly_enabled_connectors, available_connectors)) = - build_skills_and_plugins( - &sess, - turn_context.as_ref(), - &accepted_input, - &cancellation_token, - ) - .await + let Some((injection_items, explicitly_enabled_connectors)) = build_skills_and_plugins( + &sess, + first_step_context.as_ref(), + &accepted_input, + &cancellation_token, + ) + .await else { return Ok(None); }; @@ -251,7 +250,7 @@ pub(crate) async fn run_turn( .await; // Capture once so context, advertised tools, and tool calls share one request view. - let step_context = match next_step_context.take() { + let mut step_context = match next_step_context.take() { Some(step_context) => step_context, None => sess.capture_step_context(Arc::clone(&turn_context)).await, }; @@ -266,7 +265,19 @@ pub(crate) async fn run_turn( world_state = sess .record_step_world_state_if_changed(&world_state, step_context.as_ref()) .await; - let _mcp_refreshed = record_skill_injections( + let mcp_tools = if step_context.turn.apps_enabled() { + step_context + .mcp + .manager() + .list_all_tools() + .or_cancel(&cancellation_token) + .await? + } else { + Vec::new() + }; + let mut available_connectors = + available_connectors_for_step(step_context.as_ref(), &mcp_tools); + let (mcp_refreshed, skill_items) = record_skill_injections( &sess, turn_context.as_ref(), step_context.as_ref(), @@ -276,6 +287,42 @@ pub(crate) async fn run_turn( &mut skill_injection_state, ) .await; + if mcp_refreshed { + let mcp = sess + .mcp_runtime_for_step( + step_context.turn.as_ref(), + &step_context.environments, + &sess.services.selected_capability_roots, + &step_context.selected_capability_roots, + ) + .await; + step_context = Arc::new(StepContext::new( + Arc::clone(&step_context.turn), + step_context.environments.clone(), + step_context.selected_capability_roots.clone(), + Arc::clone(&step_context.skills), + mcp, + step_context.loaded_agents_md.clone(), + )); + let mcp_tools = if step_context.turn.apps_enabled() { + step_context + .mcp + .manager() + .list_all_tools() + .or_cancel(&cancellation_token) + .await? + } else { + Vec::new() + }; + available_connectors = + available_connectors_for_step(step_context.as_ref(), &mcp_tools); + } + let connector_ids = collect_explicit_app_ids_from_skill_items( + &skill_items, + &available_connectors, + &step_context.skills.skill_name_counts_lower(), + ); + sess.merge_connector_selection(connector_ids).await; // Construct the input that we will send to the model. let sampling_request_input: Vec = async { @@ -550,9 +597,9 @@ async fn record_skill_injections( available_connectors: &[connectors::AppInfo], cancellation_token: &CancellationToken, state: &mut SkillInjectionState, -) -> bool { +) -> (bool, Vec) { if crate::guardian::is_guardian_reviewer_source(&turn_context.session_source) { - return false; + return (false, Vec::new()); } for warning in step_context.skills.warnings() { emit_skill_warning_once( @@ -641,18 +688,20 @@ async fn record_skill_injections( if !items.is_empty() { sess.record_conversation_items(turn_context, &items).await; } - return false; + return (false, items); } let selected_entries = accepted .iter() .map(|(injection, _, _)| injection.entry.clone()) .collect::>(); + let available_mcp_servers = codex_mcp::configured_mcp_servers(step_context.mcp.config()); let mcp_refreshed = maybe_prompt_and_install_mcp_dependencies( sess, turn_context, cancellation_token, &selected_entries, + &available_mcp_servers, Some(sess.mcp_elicitation_reviewer()), ) .await; @@ -662,13 +711,6 @@ async fn record_skill_injections( .iter() .map(|(_, instructions, _)| ContextualUserFragment::into(instructions.clone())), ); - let connector_ids = collect_explicit_app_ids_from_skill_items( - &items, - available_connectors, - &step_context.skills.skill_name_counts_lower(), - ); - sess.merge_connector_selection(connector_ids).await; - let skill_invocations = accepted .iter() .filter_map(|(injection, _, _)| { @@ -710,7 +752,7 @@ async fn record_skill_injections( if !items.is_empty() { sess.record_conversation_items(turn_context, &items).await; } - mcp_refreshed + (mcp_refreshed, items) } async fn emit_skill_warning_once( @@ -725,17 +767,38 @@ async fn emit_skill_warning_once( } } +fn available_connectors_for_step( + step_context: &StepContext, + mcp_tools: &[codex_mcp::ToolInfo], +) -> Vec { + if !step_context.turn.apps_enabled() { + return Vec::new(); + } + let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible( + step_context + .mcp + .config() + .connector_snapshot + .connector_ids() + .iter() + .map(|connector_id| connector_id.0.clone()), + connectors::accessible_connectors_from_mcp_tools(mcp_tools), + ); + connectors::with_app_enabled_state(connectors, &step_context.turn.config) +} + #[instrument(level = "trace", skip_all)] async fn build_skills_and_plugins( sess: &Arc, - turn_context: &TurnContext, + step_context: &StepContext, input: &[TurnInput], cancellation_token: &CancellationToken, -) -> Option<(Vec, HashSet, Vec)> { +) -> Option<(Vec, HashSet)> { + let turn_context = step_context.turn.as_ref(); // Guardian input embeds the parent transcript as untrusted evidence. Do not interpret skill or // plugin mentions from that generated prompt as requests to inject additional instructions. if crate::guardian::is_guardian_reviewer_source(&turn_context.session_source) { - return Some((Vec::new(), HashSet::new(), Vec::new())); + return Some((Vec::new(), HashSet::new())); } let user_input = user_input_from_turn_input(input); @@ -754,17 +817,13 @@ async fn build_skills_and_plugins( // enabled plugins, then converted into turn-scoped guidance below. let mentioned_plugins = collect_explicit_plugin_mentions(&user_input, loaded_plugins.capability_summaries()); - let connector_snapshot = codex_connectors::ConnectorSnapshot::from_plugin_capability_summaries( - loaded_plugins.capability_summaries(), - ); let mcp_tools = if turn_context.apps_enabled() || !mentioned_plugins.is_empty() { // Plugin mentions need raw MCP/app inventory even when app tools // are normally hidden so we can describe the plugin's currently // usable capabilities for this turn. - match sess - .services - .mcp_connection_manager - .load_full() + match step_context + .mcp + .manager() .list_all_tools() .or_cancel(cancellation_token) .await @@ -776,18 +835,7 @@ async fn build_skills_and_plugins( } else { Vec::new() }; - let available_connectors = if turn_context.apps_enabled() { - let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible( - connector_snapshot - .connector_ids() - .iter() - .map(|connector_id| connector_id.0.clone()), - connectors::accessible_connectors_from_mcp_tools(&mcp_tools), - ); - connectors::with_app_enabled_state(connectors, &turn_context.config) - } else { - Vec::new() - }; + let available_connectors = available_connectors_for_step(step_context, &mcp_tools); let extension_injection_items = build_extension_turn_input_items(sess, turn_context, &user_input, cancellation_token) .await?; @@ -825,11 +873,7 @@ async fn build_skills_and_plugins( let mut injection_items = plugin_items; injection_items.extend(extension_injection_items); - Some(( - injection_items, - explicitly_enabled_connectors, - available_connectors, - )) + Some((injection_items, explicitly_enabled_connectors)) } #[tracing::instrument( @@ -1338,7 +1382,7 @@ pub(crate) async fn built_tools( cancellation_token: &CancellationToken, ) -> CodexResult> { let turn_context = step_context.turn.as_ref(); - let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); + let mcp_connection_manager = step_context.mcp.manager(); let has_mcp_servers = mcp_connection_manager.has_servers(); let all_mcp_tools = mcp_connection_manager .list_all_tools() @@ -1350,9 +1394,7 @@ pub(crate) async fn built_tools( .plugins_for_config(&turn_context.config.plugins_config_input()) .instrument(trace_span!("built_tools.load_plugins")) .await; - let connector_snapshot = codex_connectors::ConnectorSnapshot::from_plugin_capability_summaries( - loaded_plugins.capability_summaries(), - ); + let connector_snapshot = &step_context.mcp.config().connector_snapshot; let apps_enabled = turn_context.apps_enabled(); let accessible_connectors = diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index b64a228404e6..39ddc5442796 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -1,5 +1,7 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::Weak; use std::sync::atomic::AtomicBool; use crate::SkillsService; @@ -16,6 +18,7 @@ use crate::guardian::GuardianRejection; use crate::guardian::GuardianRejectionCircuitBreaker; use crate::mcp::McpManager; use crate::session::McpRuntimeSnapshot; +use crate::session::SelectedMcpRuntimeCache; use crate::tools::code_mode::CodeModeService; use crate::tools::handlers::ToolSearchHandlerCache; use crate::tools::network_approval::NetworkApprovalService; @@ -28,7 +31,6 @@ use codex_analytics::AnalyticsEventsClient; use codex_core_plugins::PluginsManager; use codex_core_skills::ExecutorSkillCatalogCache; use codex_extension_api::ExtensionData; -use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; use codex_hooks::Hooks; use codex_login::AuthManager; @@ -48,10 +50,12 @@ use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { - /// Mirror of the latest manager for extension resource clients that predate runtime snapshots. - pub(crate) mcp_connection_manager: Arc>, /// The latest atomically published MCP config and manager pair. - pub(crate) mcp_runtime: ArcSwapOption, + pub(crate) mcp_runtime: Arc>, + /// Managers that may still own an outstanding elicitation request. + pub(crate) mcp_elicitation_managers: StdMutex>>, + /// Successful executor projections and the augmented runtime built from them. + pub(crate) selected_mcp_runtime: Mutex, pub(crate) mcp_startup_cancellation_token: Mutex, pub(crate) unified_exec_manager: UnifiedExecProcessManager, #[cfg_attr(not(unix), allow(dead_code))] @@ -85,7 +89,6 @@ pub(crate) struct SessionServices { /// Raw capability selections for this thread. Each model step resolves them against its /// current executor environments before using them. pub(crate) selected_capability_roots: Vec, - pub(crate) mcp_thread_init: ExtensionDataInit, pub(crate) agent_control: AgentControl, pub(crate) network_proxy: ArcSwapOption, pub(crate) network_proxy_audit_metadata: NetworkProxyAuditMetadata, @@ -108,29 +111,57 @@ impl SessionServices { /// resolve through the session's manager while validation waits. pub(crate) async fn install_mcp_connection_manager( &self, - config: Arc, + runtime_config: Arc, runtime_context: McpRuntimeContext, manager: McpConnectionManager, ) -> Result<()> { - let runtime = self.publish_mcp_runtime(config, runtime_context, manager); + let runtime = self + .replace_base_mcp_runtime(runtime_config, runtime_context, manager) + .await; runtime.manager().validate_required_servers().await } - pub(crate) fn publish_mcp_runtime( + pub(crate) async fn replace_base_mcp_runtime( &self, - config: Arc, + runtime_config: Arc, runtime_context: McpRuntimeContext, manager: McpConnectionManager, ) -> Arc { - let manager = Arc::new(manager); - // Publish the manager for legacy resource clients first. Once the paired snapshot is - // visible, every model-scoped consumer observes this exact manager. - self.mcp_connection_manager.store(Arc::clone(&manager)); - let runtime = Arc::new(McpRuntimeSnapshot::new(config, manager, runtime_context)); - self.mcp_runtime.store(Some(Arc::clone(&runtime))); + let runtime = Arc::new(McpRuntimeSnapshot::new( + runtime_config, + Arc::new(manager), + runtime_context, + )); + self.replace_base_with_runtime(Arc::clone(&runtime)).await; runtime } + pub(crate) async fn replace_base_with_runtime(&self, runtime: Arc) { + self.selected_mcp_runtime + .lock() + .await + .replace_base(Arc::clone(&runtime)); + self.publish_existing_mcp_runtime(runtime); + } + + pub(crate) fn publish_existing_mcp_runtime(&self, runtime: Arc) { + let manager = runtime.manager_arc(); + self.track_mcp_elicitation_manager(&manager); + self.mcp_runtime.store(Some(runtime)); + } + + pub(crate) fn track_mcp_elicitation_manager(&self, manager: &Arc) { + let weak_manager = Arc::downgrade(&manager); + let mut managers = self + .mcp_elicitation_managers + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + managers.retain(|manager| manager.strong_count() > 0); + if !managers.iter().any(|manager| manager.ptr_eq(&weak_manager)) { + managers.push(weak_manager); + } + } + pub(crate) fn latest_mcp_runtime(&self) -> Arc { let Some(runtime) = self.mcp_runtime.load_full() else { unreachable!("MCP runtime must be installed before handling requests"); diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index b2add06024fc..be0110789756 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -39,6 +39,7 @@ use codex_model_provider_info::OPENAI_PROVIDER_ID; use codex_models_manager::manager::RefreshStrategy; use codex_models_manager::manager::SharedModelsManager; use codex_protocol::ThreadId; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; @@ -189,6 +190,7 @@ pub struct StartThreadOptions { pub metrics_service_name: Option, pub parent_trace: Option, pub environments: Vec, + pub selected_capability_roots: Vec, pub thread_extension_init: ExtensionDataInit, pub supports_openai_form_elicitation: bool, } @@ -639,6 +641,7 @@ impl ThreadManager { metrics_service_name: None, parent_trace: None, environments, + selected_capability_roots: Vec::new(), thread_extension_init: ExtensionDataInit::default(), supports_openai_form_elicitation: false, })) @@ -680,6 +683,7 @@ impl ThreadManager { /*inherited_exec_policy*/ None, options.parent_trace, options.environments, + options.selected_capability_roots, options.thread_extension_init, options.supports_openai_form_elicitation, /*user_shell_override*/ None, @@ -775,6 +779,7 @@ impl ThreadManager { /*inherited_exec_policy*/ None, parent_trace, environments, + /*selected_capability_roots*/ Vec::new(), /*thread_extension_init*/ ExtensionDataInit::default(), supports_openai_form_elicitation, /*user_shell_override*/ None, @@ -805,6 +810,7 @@ impl ThreadManager { /*metrics_service_name*/ None, /*parent_trace*/ None, environments, + /*selected_capability_roots*/ Vec::new(), /*thread_extension_init*/ ExtensionDataInit::default(), supports_openai_form_elicitation, /*user_shell_override*/ Some(user_shell_override), @@ -844,6 +850,7 @@ impl ThreadManager { /*inherited_exec_policy*/ None, /*parent_trace*/ None, environments, + /*selected_capability_roots*/ Vec::new(), /*thread_extension_init*/ ExtensionDataInit::default(), supports_openai_form_elicitation, /*user_shell_override*/ Some(user_shell_override), @@ -1025,6 +1032,7 @@ impl ThreadManager { /*metrics_service_name*/ None, parent_trace, environments, + /*selected_capability_roots*/ Vec::new(), /*thread_extension_init*/ ExtensionDataInit::default(), supports_openai_form_elicitation, /*user_shell_override*/ None, @@ -1344,6 +1352,7 @@ impl ThreadManagerState { inherited_exec_policy, /*parent_trace*/ None, environments, + /*selected_capability_roots*/ Vec::new(), /*thread_extension_init*/ ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, @@ -1382,6 +1391,7 @@ impl ThreadManagerState { inherited_exec_policy, /*parent_trace*/ None, environments, + /*selected_capability_roots*/ Vec::new(), /*thread_extension_init*/ ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, @@ -1422,6 +1432,7 @@ impl ThreadManagerState { inherited_exec_policy, /*parent_trace*/ None, environments, + /*selected_capability_roots*/ Vec::new(), thread_extension_init, /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, @@ -1444,6 +1455,7 @@ impl ThreadManagerState { metrics_service_name: Option, parent_trace: Option, environments: Vec, + selected_capability_roots: Vec, thread_extension_init: ExtensionDataInit, supports_openai_form_elicitation: bool, user_shell_override: Option, @@ -1463,6 +1475,7 @@ impl ThreadManagerState { /*inherited_exec_policy*/ None, parent_trace, environments, + selected_capability_roots, thread_extension_init, supports_openai_form_elicitation, user_shell_override, @@ -1487,6 +1500,7 @@ impl ThreadManagerState { inherited_exec_policy: Option>, parent_trace: Option, environments: Vec, + selected_capability_roots: Vec, thread_extension_init: ExtensionDataInit, supports_openai_form_elicitation: bool, user_shell_override: Option, @@ -1565,6 +1579,7 @@ impl ThreadManagerState { user_shell_override, parent_trace, environment_selections: environments, + selected_capability_roots, thread_extension_init, supports_openai_form_elicitation, analytics_events_client: self.analytics_events_client.clone(), diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index b230eb12bb3d..d01db59aa6c1 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -10,8 +10,6 @@ use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; use codex_extension_api::empty_extension_registry; use codex_models_manager::manager::RefreshStrategy; -use codex_protocol::capabilities::CapabilityRootLocation; -use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::models::ContentItem; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::ResponseItem; @@ -20,8 +18,6 @@ use codex_protocol::protocol::AgentMessageEvent; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::InternalSessionSource; use codex_protocol::protocol::ResumedHistory; -use codex_protocol::protocol::SessionMeta; -use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnStartedEvent; @@ -409,6 +405,7 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { metrics_service_name: None, parent_trace: None, environments: Vec::new(), + selected_capability_roots: Vec::new(), thread_extension_init: Default::default(), supports_openai_form_elicitation: false, }) @@ -427,295 +424,6 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { assert!(manager.list_thread_ids().await.is_empty()); } -#[tokio::test] -async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() { - struct InitialDataRecorder { - lifecycle_observed: Arc>>, - mcp_observed: Arc>>, - } - - impl codex_extension_api::ThreadLifecycleContributor for InitialDataRecorder { - fn on_thread_start<'a>( - &'a self, - input: codex_extension_api::ThreadStartInput<'a, Config>, - ) -> codex_extension_api::ExtensionFuture<'a, ()> { - Box::pin(async move { - let selected_root = input - .thread_store - .get::>() - .and_then(|roots| roots.first().cloned()) - .expect("selected root should be available"); - self.lifecycle_observed - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push((input.thread_store.level_id().to_string(), selected_root.id)); - input - .thread_store - .insert(Vec::::new()); - }) - } - } - - impl codex_extension_api::McpServerContributor for InitialDataRecorder { - fn id(&self) -> &'static str { - "selected_root_test" - } - - fn contribute<'a>( - &'a self, - context: codex_extension_api::McpServerContributionContext<'a, Config>, - ) -> codex_extension_api::ExtensionFuture<'a, Vec> - { - Box::pin(async move { - let thread_init = context - .thread_init() - .expect("initial MCP resolution should be thread-scoped"); - let selected_root = thread_init - .get::>() - .and_then(|roots| roots.first().cloned()) - .expect("selected root should be available"); - self.mcp_observed - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(selected_root.id.clone()); - let mut server = codex_mcp::codex_apps_mcp_server_config( - "https://selected.invalid", - /*apps_mcp_product_sku*/ None, - ); - let CapabilityRootLocation::Environment { environment_id, .. } = - &selected_root.location; - server.environment_id = environment_id.clone(); - server.enabled = false; - let plugin_id = selected_root.id; - vec![codex_extension_api::McpServerContribution::SelectedPlugin { - name: plugin_id.clone(), - plugin_display_name: plugin_id.clone(), - plugin_id, - selection_order: 0, - config: Box::new(server), - }] - }) - } - } - - let temp_dir = tempdir().expect("tempdir"); - let mut config = test_config().await; - config.codex_home = temp_dir.path().join("codex-home").abs(); - config.cwd = config.codex_home.abs(); - std::fs::create_dir_all(&config.codex_home).expect("create codex home"); - - let lifecycle_observed = Arc::new(std::sync::Mutex::new(Vec::new())); - let mcp_observed = Arc::new(std::sync::Mutex::new(Vec::new())); - let recorder = Arc::new(InitialDataRecorder { - lifecycle_observed: Arc::clone(&lifecycle_observed), - mcp_observed: Arc::clone(&mcp_observed), - }); - let mut extensions = codex_extension_api::ExtensionRegistryBuilder::new(); - extensions.thread_lifecycle_contributor(recorder.clone()); - extensions.mcp_server_contributor(recorder); - let manager = ThreadManager::new( - &config, - AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()), - SessionSource::Exec, - Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - Arc::new(extensions.build()), - Arc::new(crate::test_support::EmptyUserInstructionsProvider), - /*analytics_events_client*/ None, - thread_store_from_config(&config, /*state_db*/ None), - /*agent_graph_store*/ None, - TEST_INSTALLATION_ID.to_string(), - /*attestation_provider*/ None, - /*external_time_provider*/ None, - ); - let selected_root_init = |id: &str, environment_id: &str| { - let mut init = codex_extension_api::ExtensionDataInit::new(); - init.insert(vec![SelectedCapabilityRoot { - id: id.to_string(), - location: CapabilityRootLocation::Environment { - environment_id: environment_id.to_string(), - path: PathUri::parse(&format!("file:///plugins/{id}")).expect("plugin root URI"), - }, - }]); - init - }; - - let first_thread = manager - .start_thread_with_options(StartThreadOptions { - config: config.clone(), - initial_history: InitialHistory::New, - session_source: None, - thread_source: None, - dynamic_tools: Vec::new(), - metrics_service_name: None, - parent_trace: None, - environments: Vec::new(), - thread_extension_init: selected_root_init("selected-a", "env-a"), - supports_openai_form_elicitation: false, - }) - .await - .expect("start first thread"); - let second_thread = manager - .start_thread_with_options(StartThreadOptions { - config: config.clone(), - initial_history: InitialHistory::New, - session_source: None, - thread_source: None, - dynamic_tools: Vec::new(), - metrics_service_name: None, - parent_trace: None, - environments: Vec::new(), - thread_extension_init: selected_root_init("selected-b", "env-b"), - supports_openai_form_elicitation: false, - }) - .await - .expect("start second thread"); - let first_resolved = first_thread.thread.runtime_mcp_config(&config).await; - let second_resolved = second_thread.thread.runtime_mcp_config(&config).await; - - assert_eq!( - *lifecycle_observed - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner), - vec![ - (first_thread.thread_id.to_string(), "selected-a".to_string()), - ( - second_thread.thread_id.to_string(), - "selected-b".to_string() - ), - ] - ); - assert_eq!( - *mcp_observed - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner), - vec![ - "selected-a".to_string(), - "selected-b".to_string(), - "selected-a".to_string(), - "selected-b".to_string(), - ] - ); - let selected_servers = |config: &codex_mcp::McpConfig| { - codex_mcp::configured_mcp_servers(config) - .into_iter() - .filter(|(name, _)| name.starts_with("selected-")) - .map(|(name, server)| (name, server.environment_id)) - .collect::>() - }; - assert_eq!( - selected_servers(&first_resolved), - std::collections::BTreeMap::from([("selected-a".to_string(), "env-a".to_string())]) - ); - assert_eq!( - selected_servers(&second_resolved), - std::collections::BTreeMap::from([("selected-b".to_string(), "env-b".to_string())]) - ); -} - -#[tokio::test] -async fn selected_capability_roots_round_trip_through_fork_and_explicit_empty_wins() { - let temp_dir = tempdir().expect("tempdir"); - let mut config = test_config().await; - config.codex_home = temp_dir.path().join("codex-home").abs(); - config.cwd = config.codex_home.abs(); - 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.to_path_buf(), - Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), - ); - let selected_roots = vec![SelectedCapabilityRoot { - id: "demo@1".to_string(), - location: CapabilityRootLocation::Environment { - environment_id: "build".to_string(), - path: PathUri::parse("file:///plugins/demo").expect("plugin root URI"), - }, - }]; - let fork_history = || { - InitialHistory::Forked(vec![RolloutItem::SessionMeta(SessionMetaLine { - meta: SessionMeta { - selected_capability_roots: selected_roots.clone(), - ..SessionMeta::default() - }, - git: None, - })]) - }; - - let inherited = manager - .start_thread_with_options(StartThreadOptions { - config: config.clone(), - initial_history: fork_history(), - session_source: None, - thread_source: None, - dynamic_tools: Vec::new(), - metrics_service_name: None, - parent_trace: None, - environments: Vec::new(), - thread_extension_init: Default::default(), - supports_openai_form_elicitation: false, - }) - .await - .expect("start inherited fork"); - inherited.thread.ensure_rollout_materialized().await; - inherited - .thread - .flush_rollout() - .await - .expect("flush inherited fork"); - let inherited_history = RolloutRecorder::get_rollout_history( - &inherited - .thread - .rollout_path() - .expect("inherited fork rollout path"), - ) - .await - .expect("read inherited fork rollout"); - - assert_eq!( - inherited_history.get_selected_capability_roots(), - selected_roots - ); - - let mut explicit_empty = codex_extension_api::ExtensionDataInit::new(); - explicit_empty.insert(Vec::::new()); - let overridden = manager - .start_thread_with_options(StartThreadOptions { - config, - initial_history: fork_history(), - session_source: None, - thread_source: None, - dynamic_tools: Vec::new(), - metrics_service_name: None, - parent_trace: None, - environments: Vec::new(), - thread_extension_init: explicit_empty, - supports_openai_form_elicitation: false, - }) - .await - .expect("start explicitly empty fork"); - overridden.thread.ensure_rollout_materialized().await; - overridden - .thread - .flush_rollout() - .await - .expect("flush explicitly empty fork"); - let overridden_history = RolloutRecorder::get_rollout_history( - &overridden - .thread - .rollout_path() - .expect("explicitly empty fork rollout path"), - ) - .await - .expect("read explicitly empty fork rollout"); - - assert_eq!( - overridden_history.get_selected_capability_roots(), - Vec::new() - ); -} - #[tokio::test] async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { let temp_dir = tempdir().expect("tempdir"); @@ -760,6 +468,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { metrics_service_name: None, parent_trace: None, environments: environments.clone(), + selected_capability_roots: Vec::new(), thread_extension_init: Default::default(), supports_openai_form_elicitation: false, }) @@ -1043,6 +752,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { metrics_service_name: None, parent_trace: None, environments: Vec::new(), + selected_capability_roots: Vec::new(), thread_extension_init: Default::default(), supports_openai_form_elicitation: false, }) diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 7342e5cfbed6..30fd9f1f21e9 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -696,7 +696,7 @@ async fn environment_tools_follow_the_step_context() { environments, Vec::new(), skills, - crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config), + crate::session::uninitialized_mcp_runtime(&turn.config), /*loaded_agents_md*/ None, )); diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 8a4f5cca7dc1..c62fe44c63ac 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -26,7 +26,6 @@ use codex_core::thread_store_from_config; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::ExecutorFileSystem; use codex_exec_server::RemoveOptions; -use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; use codex_extension_api::LoadUserInstructionsFuture; use codex_extension_api::UserInstructionsProvider; @@ -289,7 +288,7 @@ pub struct TestCodexBuilder { user_shell_override: Option, exec_server_url: Option, extensions: Arc>, - thread_extension_init: ExtensionDataInit, + selected_capability_roots: Vec, user_instructions_provider: Option>, supports_openai_form_elicitation: bool, external_time_provider: Option>, @@ -381,8 +380,11 @@ impl TestCodexBuilder { self } - pub fn with_thread_extension_init(mut self, thread_extension_init: ExtensionDataInit) -> Self { - self.thread_extension_init = thread_extension_init; + pub fn with_selected_capability_roots( + mut self, + selected_capability_roots: Vec, + ) -> Self { + self.selected_capability_roots = selected_capability_roots; self } @@ -681,7 +683,10 @@ impl TestCodexBuilder { metrics_service_name: None, parent_trace: None, environments, - thread_extension_init: std::mem::take(&mut self.thread_extension_init), + selected_capability_roots: std::mem::take( + &mut self.selected_capability_roots, + ), + thread_extension_init: Default::default(), supports_openai_form_elicitation: self.supports_openai_form_elicitation, }), ) @@ -1230,7 +1235,7 @@ pub fn test_codex() -> TestCodexBuilder { user_shell_override: None, exec_server_url: None, extensions: empty_extension_registry(), - thread_extension_init: ExtensionDataInit::new(), + selected_capability_roots: Vec::new(), user_instructions_provider: None, supports_openai_form_elicitation: false, external_time_provider: None, diff --git a/codex-rs/core/tests/suite/agents_md.rs b/codex-rs/core/tests/suite/agents_md.rs index 605e0ada7dfb..62a26910449a 100644 --- a/codex-rs/core/tests/suite/agents_md.rs +++ b/codex-rs/core/tests/suite/agents_md.rs @@ -493,6 +493,7 @@ async fn loads_user_instructions_without_a_primary_environment() -> Result<()> { metrics_service_name: None, parent_trace: None, environments: Vec::new(), + selected_capability_roots: Vec::new(), thread_extension_init: Default::default(), supports_openai_form_elicitation: false, }) @@ -707,6 +708,7 @@ async fn multi_environment_thread_loads_every_project_and_keeps_creation_snapsho cwd: PathUri::from_host_native_path(local_root.path())?, }, ], + selected_capability_roots: Vec::new(), thread_extension_init: Default::default(), supports_openai_form_elicitation: false, }) diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index aad8fc494bb9..034c97963285 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -768,6 +768,7 @@ async fn subagent_stop_replaces_stop_and_skips_internal_subagents() -> Result<() metrics_service_name: None, parent_trace: None, environments: Vec::new(), + selected_capability_roots: Vec::new(), thread_extension_init: Default::default(), supports_openai_form_elicitation: false, }) diff --git a/codex-rs/ext/extension-api/src/contributors/mcp.rs b/codex-rs/ext/extension-api/src/contributors/mcp.rs index 91e9cb459cca..7d759c8a2cfa 100644 --- a/codex-rs/ext/extension-api/src/contributors/mcp.rs +++ b/codex-rs/ext/extension-api/src/contributors/mcp.rs @@ -1,17 +1,9 @@ use codex_config::McpServerConfig; -use crate::ExtensionDataInit; - /// Input supplied while resolving MCP server contributions. -/// -/// Thread-scoped implementations can read the immutable host-seeded inputs -/// through [`Self::thread_init`]. Implementations should not retain borrowed -/// context after contribution completes. pub struct McpServerContributionContext<'a, C> { /// Host configuration visible during MCP resolution. config: &'a C, - /// Initial inputs for the active thread, when resolution is thread-scoped. - thread_init: Option<&'a ExtensionDataInit>, } impl Clone for McpServerContributionContext<'_, C> { @@ -25,29 +17,13 @@ impl Copy for McpServerContributionContext<'_, C> {} impl<'a, C> McpServerContributionContext<'a, C> { /// Creates context for resolution that is not associated with a running thread. pub fn global(config: &'a C) -> Self { - Self { - config, - thread_init: None, - } - } - - /// Creates context for one active thread runtime. - pub fn for_thread(config: &'a C, thread_init: &'a ExtensionDataInit) -> Self { - Self { - config, - thread_init: Some(thread_init), - } + Self { config } } /// Returns the host configuration visible during resolution. pub fn config(&self) -> &'a C { self.config } - - /// Returns the frozen initial inputs when resolving for a running thread. - pub fn thread_init(&self) -> Option<&'a ExtensionDataInit> { - self.thread_init - } } /// One extension-owned overlay for the runtime MCP server configuration. @@ -58,14 +34,6 @@ pub enum McpServerContribution { name: String, config: Box, }, - /// Registers a server declared by a plugin selected for this thread. - SelectedPlugin { - name: String, - plugin_id: String, - plugin_display_name: String, - selection_order: usize, - config: Box, - }, /// Removes a named MCP server. Remove { name: String }, } diff --git a/codex-rs/ext/mcp/Cargo.toml b/codex-rs/ext/mcp/Cargo.toml index 13c219edb9f7..c3575042d6d3 100644 --- a/codex-rs/ext/mcp/Cargo.toml +++ b/codex-rs/ext/mcp/Cargo.toml @@ -21,7 +21,6 @@ codex-features = { workspace = true } codex-mcp = { workspace = true } codex-protocol = { workspace = true } tracing = { workspace = true } -tokio = { workspace = true, features = ["sync"] } [dev-dependencies] codex-config = { workspace = true } diff --git a/codex-rs/ext/mcp/src/executor_plugin.rs b/codex-rs/ext/mcp/src/executor_plugin.rs deleted file mode 100644 index e8eb7aacba1d..000000000000 --- a/codex-rs/ext/mcp/src/executor_plugin.rs +++ /dev/null @@ -1,126 +0,0 @@ -use codex_core::config::Config; -use codex_core_plugins::ExecutorPluginRuntime; -use codex_exec_server::EnvironmentManager; -use codex_extension_api::ExtensionDataInit; -use codex_extension_api::ExtensionFuture; -use codex_extension_api::McpServerContribution; -use codex_extension_api::McpServerContributionContext; -use codex_extension_api::McpServerContributor; -use codex_protocol::capabilities::SelectedCapabilityRoot; -use std::collections::HashMap; -use std::sync::Arc; -use tokio::sync::OnceCell; - -/// The runtime declarations frozen for one selected package at thread start. -#[derive(Clone)] -struct SelectedPluginRuntime { - selection_order: usize, - runtime: ExecutorPluginRuntime, -} - -#[derive(Default)] -pub(crate) struct SelectedExecutorPluginMcpState { - snapshot: OnceCell>, -} - -pub(crate) fn seed_thread_state(thread_init: &mut ExtensionDataInit) { - thread_init.insert(SelectedExecutorPluginMcpState::default()); -} - -pub(crate) struct SelectedExecutorPluginMcpContributor { - environment_manager: Arc, -} - -impl SelectedExecutorPluginMcpContributor { - pub(crate) fn new(environment_manager: Arc) -> Self { - Self { - environment_manager, - } - } - - async fn resolve_snapshot( - &self, - selected_roots: &[SelectedCapabilityRoot], - ) -> Vec { - let mut snapshot = Vec::new(); - let resolved_roots = self - .environment_manager - .bind_selected_capability_roots(selected_roots); - - for (selection_order, resolved_root) in resolved_roots.iter().enumerate() { - let selected_root = resolved_root.selected_root(); - match ExecutorPluginRuntime::project(resolved_root).await { - Ok(Some(runtime)) => snapshot.push(SelectedPluginRuntime { - selection_order, - runtime, - }), - Ok(None) => {} - Err(err) => { - tracing::warn!( - selected_root = selected_root.id, - error = %err, - "failed to project selected executor plugin runtime" - ); - } - } - } - - snapshot - } -} - -impl McpServerContributor for SelectedExecutorPluginMcpContributor { - fn id(&self) -> &'static str { - "selected_executor_plugin_mcp" - } - - fn contribute<'a>( - &'a self, - context: McpServerContributionContext<'a, Config>, - ) -> ExtensionFuture<'a, Vec> { - Box::pin(async move { - let Some(thread_init) = context.thread_init() else { - return Vec::new(); - }; - let Some(selected_roots) = thread_init.get::>() else { - return Vec::new(); - }; - let Some(state) = thread_init.get::() else { - tracing::warn!("selected executor plugin MCP state was not initialized"); - return Vec::new(); - }; - let snapshot = state - .snapshot - .get_or_init(|| self.resolve_snapshot(selected_roots.as_ref())) - .await; - let mut contributions = Vec::new(); - - for selected in snapshot { - let plugin = selected.runtime.plugin(); - let plugin_id = plugin.selected_root_id(); - let mut servers = selected - .runtime - .mcp_servers() - .iter() - .cloned() - .collect::>(); - context - .config() - .apply_plugin_mcp_server_requirements(plugin_id, &mut servers); - let mut servers = servers.into_iter().collect::>(); - servers.sort_unstable_by(|left, right| left.0.cmp(&right.0)); - contributions.extend(servers.into_iter().map(|(name, config)| { - McpServerContribution::SelectedPlugin { - name, - plugin_id: plugin_id.to_string(), - plugin_display_name: plugin.manifest().display_name().to_string(), - selection_order: selected.selection_order, - config: Box::new(config), - } - })); - } - - contributions - }) - } -} diff --git a/codex-rs/ext/mcp/src/lib.rs b/codex-rs/ext/mcp/src/lib.rs index e32316ac36fb..e1a1007d4002 100644 --- a/codex-rs/ext/mcp/src/lib.rs +++ b/codex-rs/ext/mcp/src/lib.rs @@ -7,8 +7,6 @@ use codex_extension_api::McpServerContributor; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::hosted_plugin_runtime_mcp_server_config; -mod executor_plugin; - struct HostedPluginRuntimeExtension; impl McpServerContributor for HostedPluginRuntimeExtension { @@ -41,20 +39,3 @@ impl McpServerContributor for HostedPluginRuntimeExtension { pub fn install(builder: &mut ExtensionRegistryBuilder) { builder.mcp_server_contributor(std::sync::Arc::new(HostedPluginRuntimeExtension)); } - -/// Installs discovery for MCP servers declared by thread-selected executor plugins. -pub fn install_executor_plugins( - builder: &mut ExtensionRegistryBuilder, - environment_manager: std::sync::Arc, -) { - builder.mcp_server_contributor(std::sync::Arc::new( - executor_plugin::SelectedExecutorPluginMcpContributor::new(environment_manager), - )); -} - -/// Seeds the per-thread snapshot used by selected executor plugin MCP discovery. -pub fn initialize_executor_plugin_thread_data( - thread_init: &mut codex_extension_api::ExtensionDataInit, -) { - executor_plugin::seed_thread_state(thread_init); -} diff --git a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs b/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs deleted file mode 100644 index 7bdbc83d62e4..000000000000 --- a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs +++ /dev/null @@ -1,140 +0,0 @@ -use codex_config::test_support::CloudConfigBundleFixture; -use codex_core::config::Config; -use codex_core::config::ConfigBuilder; -use codex_exec_server::EnvironmentManager; -use codex_exec_server::LOCAL_ENVIRONMENT_ID; -use codex_extension_api::ExtensionDataInit; -use codex_extension_api::ExtensionRegistryBuilder; -use codex_extension_api::McpServerContribution; -use codex_extension_api::McpServerContributionContext; -use codex_protocol::capabilities::CapabilityRootLocation; -use codex_protocol::capabilities::SelectedCapabilityRoot; -use codex_utils_path_uri::PathUri; -use pretty_assertions::assert_eq; -use std::sync::Arc; - -type TestResult = Result<(), Box>; - -#[derive(Debug, PartialEq, Eq)] -struct ContributionSummary { - name: String, - plugin_id: String, - plugin_display_name: String, - selection_order: usize, - enabled: bool, -} - -#[tokio::test] -async fn selected_plugin_servers_use_managed_requirements_for_the_selected_root_id() -> TestResult { - let codex_home = tempfile::tempdir()?; - let plugin_root = tempfile::tempdir()?; - std::fs::create_dir_all(plugin_root.path().join(".codex-plugin"))?; - std::fs::write( - plugin_root.path().join(".codex-plugin/plugin.json"), - r#"{"name":"different-manifest-name","interface":{"displayName":"Selected Demo"}}"#, - )?; - std::fs::write( - plugin_root.path().join(".mcp.json"), - r#"{ - "mcpServers": { - "allowed": {"command":"allowed-command"}, - "mismatched": {"command":"wrong-command"}, - "unlisted": {"command":"unlisted-command"} - } -}"#, - )?; - let config = ConfigBuilder::default() - .codex_home(codex_home.path().to_path_buf()) - .fallback_cwd(Some(codex_home.path().to_path_buf())) - .cloud_config_bundle( - CloudConfigBundleFixture::loader_with_enterprise_requirement( - r#" -[plugins."selected-root".mcp_servers.allowed.identity] -command = "allowed-command" - -[plugins."selected-root".mcp_servers.mismatched.identity] -command = "expected-command" -"#, - ), - ) - .build() - .await?; - - let contributions = selected_plugin_contributions(&config, plugin_root.path()).await?; - - assert_eq!( - contributions, - vec![ - ContributionSummary { - name: "allowed".to_string(), - plugin_id: "selected-root".to_string(), - plugin_display_name: "Selected Demo".to_string(), - selection_order: 0, - enabled: true, - }, - ContributionSummary { - name: "mismatched".to_string(), - plugin_id: "selected-root".to_string(), - plugin_display_name: "Selected Demo".to_string(), - selection_order: 0, - enabled: false, - }, - ContributionSummary { - name: "unlisted".to_string(), - plugin_id: "selected-root".to_string(), - plugin_display_name: "Selected Demo".to_string(), - selection_order: 0, - enabled: false, - }, - ] - ); - Ok(()) -} - -async fn selected_plugin_contributions( - config: &Config, - plugin_root: &std::path::Path, -) -> Result, Box> { - let mut builder = ExtensionRegistryBuilder::new(); - codex_mcp_extension::install_executor_plugins( - &mut builder, - Arc::new(EnvironmentManager::default_for_tests()), - ); - let registry = builder.build(); - let mut thread_init = ExtensionDataInit::new(); - thread_init.insert(vec![SelectedCapabilityRoot { - id: "selected-root".to_string(), - location: CapabilityRootLocation::Environment { - environment_id: LOCAL_ENVIRONMENT_ID.to_string(), - path: PathUri::from_host_native_path(plugin_root)?, - }, - }]); - codex_mcp_extension::initialize_executor_plugin_thread_data(&mut thread_init); - - Ok(registry.mcp_server_contributors()[0] - .contribute(McpServerContributionContext::for_thread( - config, - &thread_init, - )) - .await - .into_iter() - .map(|contribution| match contribution { - McpServerContribution::SelectedPlugin { - name, - plugin_id, - plugin_display_name, - selection_order, - config, - } => ContributionSummary { - name, - plugin_id, - plugin_display_name, - selection_order, - enabled: config.enabled, - }, - McpServerContribution::Set { .. } | McpServerContribution::Remove { .. } => { - panic!("expected selected plugin contribution") - } - }) - .collect()) -} diff --git a/codex-rs/ext/skills/Cargo.toml b/codex-rs/ext/skills/Cargo.toml index 35827ac5788d..88a62d9a53fc 100644 --- a/codex-rs/ext/skills/Cargo.toml +++ b/codex-rs/ext/skills/Cargo.toml @@ -20,7 +20,6 @@ codex-extension-api = { workspace = true } codex-mcp = { workspace = true } codex-protocol = { workspace = true } codex-tools = { workspace = true } -codex-utils-path-uri = { workspace = true } codex-utils-string = { workspace = true } schemars = { workspace = true } serde = { workspace = true, features = ["derive"] } @@ -28,8 +27,3 @@ serde_json = { workspace = true } tokio = { workspace = true, features = ["sync", "time"] } tracing = { workspace = true } url = { workspace = true } - -[dev-dependencies] -codex-utils-absolute-path = { workspace = true } -pretty_assertions = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 0ee0313b5845..390f3b56725a 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -338,6 +338,7 @@ impl MemoryStartupContext { metrics_service_name: None, parent_trace: None, environments, + selected_capability_roots: Vec::new(), thread_extension_init: Default::default(), supports_openai_form_elicitation: false, }) From 102ffbc704f7e67511fc291dfcb755be0dd6c0ce Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 25 Jun 2026 17:10:39 +0100 Subject: [PATCH 4/6] Make selected MCP cache invalidation explicit --- codex-rs/core/src/session/mcp_projection.rs | 34 ++++++++++++++------- codex-rs/core/src/session/mcp_runtime.rs | 31 ++++++++++++++++--- codex-rs/core/src/session/tests.rs | 7 +++-- codex-rs/core/src/state/service.rs | 2 +- codex-rs/core/tests/common/test_codex.rs | 8 ----- 5 files changed, 56 insertions(+), 26 deletions(-) diff --git a/codex-rs/core/src/session/mcp_projection.rs b/codex-rs/core/src/session/mcp_projection.rs index 30713e745a15..7094e76b91b2 100644 --- a/codex-rs/core/src/session/mcp_projection.rs +++ b/codex-rs/core/src/session/mcp_projection.rs @@ -49,8 +49,10 @@ impl Session { return base; } let bindings = selected_bindings(selected_roots, resolved_roots); - let mut plugins = cache.plugins(&bindings).unwrap_or_default(); - if let Some(runtime) = cache.runtime(&bindings) { + let cached_runtime = cache.runtime_for_bindings(&bindings); + let mut plugins = cache.plugins_for_bindings(&bindings).unwrap_or_default(); + let mut discovered_plugin = false; + if cached_runtime.is_some() { let unresolved = bindings .iter() .filter(|(order, _)| { @@ -61,11 +63,7 @@ impl Session { .cloned() .collect::>(); let discovered = project_executor_plugins(&unresolved).await; - if discovered.is_empty() { - self.services - .publish_existing_mcp_runtime(Arc::clone(&runtime)); - return runtime; - } + discovered_plugin = !discovered.is_empty(); plugins.extend(discovered); plugins.sort_unstable_by_key(|(order, _)| *order); } else { @@ -85,6 +83,20 @@ impl Session { .collect::>(); let runtime_context = self.mcp_runtime_context(&turn_context.config, environments, &pinned_roots); + if !discovered_plugin + && let Some(runtime) = cached_runtime + && runtime.matches_projection(mcp_config.as_ref(), &runtime_context) + { + runtime + .manager() + .set_approval_policy(&turn_context.approval_policy); + runtime + .manager() + .set_permission_profile(turn_context.permission_profile()); + self.services + .publish_existing_mcp_runtime(Arc::clone(&runtime)); + return runtime; + } let runtime = if plugins.is_empty() { base } else { @@ -97,7 +109,7 @@ impl Session { ) .await }; - cache.insert_runtime(bindings, plugins, Arc::clone(&runtime)); + cache.replace_selected_runtime(bindings, plugins, Arc::clone(&runtime)); self.services .publish_existing_mcp_runtime(Arc::clone(&runtime)); runtime @@ -146,7 +158,7 @@ impl Session { let mut cache = self.services.selected_mcp_runtime.lock().await; let projection = self.project_mcp_config_inner(config).await; let current = cache - .runtime(&projection.bindings) + .runtime_for_bindings(&projection.bindings) .unwrap_or_else(|| self.services.latest_mcp_runtime()); if current.matches_projection(&projection.config, &projection.runtime_context) { current @@ -169,9 +181,9 @@ impl Session { ) .await; if projection.bindings.is_empty() { - cache.replace_base(Arc::clone(&runtime)); + cache.replace_base_and_invalidate_selected(Arc::clone(&runtime)); } else { - cache.insert_runtime( + cache.replace_selected_runtime( projection.bindings, projection.plugins, Arc::clone(&runtime), diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index 959607bb5537..a58d526c7ba2 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -4,6 +4,23 @@ use codex_core_plugins::ExecutorPluginRuntime; use codex_exec_server::ResolvedSelectedCapabilityRoot; use codex_mcp::McpRuntimeSnapshot; +/// One live selected-plugin MCP runtime retained between model steps. +/// +/// A cached runtime is a reuse candidate only for the same ordered selected roots and the same +/// process-local environment instances. The caller additionally compares the effective MCP config +/// and runtime context before reuse. Selected environment contents are treated as stable, so +/// manifest and MCP config file changes do not invalidate this cache. +/// +/// Within a live session, the selected runtime is invalidated in exactly two ways: +/// +/// 1. [`Self::replace_base_and_invalidate_selected`] installs a new base MCP runtime. +/// 2. [`Self::replace_selected_runtime`] stores a newly projected runtime. This happens when the +/// bindings change, when a previously unavailable plugin appears, or when the effective config +/// or runtime context changes. An unavailable environment disappears from the binding list and +/// therefore follows this path; returning with a new environment instance rebuilds the live +/// runtime even when the stable environment ID is unchanged. +/// +/// In-flight [`McpRuntimeSnapshot`] values retain their manager until their model step finishes. #[derive(Default)] pub(crate) struct SelectedMcpRuntimeCache { base_runtime: Option>, @@ -17,7 +34,10 @@ struct CachedSelectedRuntime { } impl SelectedMcpRuntimeCache { - pub(crate) fn replace_base(&mut self, runtime: Arc) { + pub(crate) fn replace_base_and_invalidate_selected( + &mut self, + runtime: Arc, + ) { self.base_runtime = Some(runtime); self.runtime = None; } @@ -29,7 +49,7 @@ impl SelectedMcpRuntimeCache { .expect("base MCP runtime must be installed before capturing a step") } - pub(crate) fn runtime( + pub(crate) fn runtime_for_bindings( &self, bindings: &[(usize, ResolvedSelectedCapabilityRoot)], ) -> Option> { @@ -39,7 +59,7 @@ impl SelectedMcpRuntimeCache { .map(|cached| Arc::clone(&cached.runtime)) } - pub(crate) fn plugins( + pub(crate) fn plugins_for_bindings( &self, bindings: &[(usize, ResolvedSelectedCapabilityRoot)], ) -> Option> { @@ -49,7 +69,7 @@ impl SelectedMcpRuntimeCache { .map(|cached| cached.plugins.clone()) } - pub(crate) fn insert_runtime( + pub(crate) fn replace_selected_runtime( &mut self, bindings: Vec<(usize, ResolvedSelectedCapabilityRoot)>, plugins: Vec<(usize, ExecutorPluginRuntime)>, @@ -67,6 +87,9 @@ fn same_bindings( left: &[(usize, ResolvedSelectedCapabilityRoot)], right: &[(usize, ResolvedSelectedCapabilityRoot)], ) -> bool { + // Order is part of the key because later selected roots can be renamed when MCP server names + // collide. Arc identity is part of the key because live processes and connections belong to + // one exact environment instance, even when a replacement reuses the same stable ID. left.len() == right.len() && left .iter() diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d931216a9299..9179b967052f 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5268,6 +5268,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { plugins_manager, mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), + /*selected_capability_roots*/ Vec::new(), codex_extension_api::ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, AgentControl::default(), @@ -5411,7 +5412,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { let network_approval = Arc::new(NetworkApprovalService::default()); let mcp_runtime = uninitialized_mcp_runtime(config.as_ref()); let mut selected_mcp_runtime = crate::session::SelectedMcpRuntimeCache::default(); - selected_mcp_runtime.replace_base(Arc::clone(&mcp_runtime)); + selected_mcp_runtime.replace_base_and_invalidate_selected(Arc::clone(&mcp_runtime)); let services = SessionServices { mcp_runtime: Arc::new(arc_swap::ArcSwapOption::from(Some(mcp_runtime))), mcp_elicitation_managers: std::sync::Mutex::new(Vec::new()), @@ -5646,6 +5647,7 @@ async fn make_session_with_config_and_rx( plugins_manager, mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), + /*selected_capability_roots*/ Vec::new(), codex_extension_api::ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, AgentControl::default(), @@ -5752,6 +5754,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( plugins_manager, mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), + /*selected_capability_roots*/ Vec::new(), codex_extension_api::ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, agent_control, @@ -7487,7 +7490,7 @@ where let network_approval = Arc::new(NetworkApprovalService::default()); let mcp_runtime = uninitialized_mcp_runtime(config.as_ref()); let mut selected_mcp_runtime = crate::session::SelectedMcpRuntimeCache::default(); - selected_mcp_runtime.replace_base(Arc::clone(&mcp_runtime)); + selected_mcp_runtime.replace_base_and_invalidate_selected(Arc::clone(&mcp_runtime)); let services = SessionServices { mcp_runtime: Arc::new(arc_swap::ArcSwapOption::from(Some(mcp_runtime))), mcp_elicitation_managers: std::sync::Mutex::new(Vec::new()), diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 39ddc5442796..313c447a12f7 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -140,7 +140,7 @@ impl SessionServices { self.selected_mcp_runtime .lock() .await - .replace_base(Arc::clone(&runtime)); + .replace_base_and_invalidate_selected(Arc::clone(&runtime)); self.publish_existing_mcp_runtime(runtime); } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index c62fe44c63ac..dd7adf8adda7 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -388,14 +388,6 @@ impl TestCodexBuilder { self } - pub fn with_selected_capability_roots( - mut self, - selected_capability_roots: Vec, - ) -> Self { - self.thread_extension_init.insert(selected_capability_roots); - self - } - pub fn with_user_instructions_provider( mut self, provider: Arc, From 338cbd0bf93260f4bc98de9549b7d80b3de6a54c Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 25 Jun 2026 17:42:10 +0100 Subject: [PATCH 5/6] Cache selected plugin metadata by stable root --- .../app-server/tests/suite/v2/executor_mcp.rs | 51 ++++-------- codex-rs/core/src/session/mcp_projection.rs | 67 +++------------- codex-rs/core/src/session/mcp_runtime.rs | 79 +++++++++++++------ 3 files changed, 85 insertions(+), 112 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs index f23bec8e060f..eff275dfc139 100644 --- a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs +++ b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs @@ -429,7 +429,7 @@ startup_timeout_sec = 10 } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn selected_executor_plugin_activates_after_it_becomes_ready() -> Result<()> { +async fn selected_executor_plugin_runtime_survives_thread_resume() -> Result<()> { let responses_server = responses::start_mock_server().await; let (apps_url, apps_server_handle) = start_apps_server_with_delays( vec![AppInfo { @@ -497,40 +497,8 @@ args = ["exec-server", "--listen", "stdio"] ), )?; - // The selected root exists, but its manifest arrives after the first model step. + // Environment contents are complete before selection and remain stable for the thread. let plugin = TempDir::new()?; - let mut app_server = TestAppServer::new(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; - let selected_thread = start_thread( - &mut app_server, - Some(vec![SelectedCapabilityRoot { - id: "executor-demo@1".to_string(), - location: CapabilityRootLocation::Environment { - environment_id: EXECUTOR_ID.to_string(), - path: PathUri::from_host_native_path(plugin.path())?, - }, - }]), - ) - .await?; - - let before_ready = responses::mount_sse_once( - &responses_server, - responses::sse(vec![ - responses::ev_response_created("before-ready"), - responses::ev_assistant_message("before-ready-message", "Waiting"), - responses::ev_completed("before-ready"), - ]), - ) - .await; - run_turn(&mut app_server, &selected_thread, "Check available tools").await?; - let namespace = format!("mcp__{DYNAMIC_MCP_SERVER_NAME}"); - assert!( - before_ready - .single_request() - .tool_by_name(&namespace, "echo") - .is_none() - ); - std::fs::create_dir_all(plugin.path().join(".codex-plugin"))?; std::fs::write( plugin.path().join(".codex-plugin/plugin.json"), @@ -556,6 +524,21 @@ args = ["exec-server", "--listen", "stdio"] }))?, )?; + let mut app_server = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; + let selected_thread = start_thread( + &mut app_server, + Some(vec![SelectedCapabilityRoot { + id: "executor-demo@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: EXECUTOR_ID.to_string(), + path: PathUri::from_host_native_path(plugin.path())?, + }, + }]), + ) + .await?; + let namespace = format!("mcp__{DYNAMIC_MCP_SERVER_NAME}"); + let response_mock = responses::mount_sse_sequence( &responses_server, vec![ diff --git a/codex-rs/core/src/session/mcp_projection.rs b/codex-rs/core/src/session/mcp_projection.rs index 7094e76b91b2..143a779816f5 100644 --- a/codex-rs/core/src/session/mcp_projection.rs +++ b/codex-rs/core/src/session/mcp_projection.rs @@ -8,7 +8,6 @@ use super::TurnContext; use crate::config::Config; use crate::environment_selection::TurnEnvironmentSnapshot; use codex_config::McpServerConfig; -use codex_core_plugins::ExecutorPluginRuntime; use codex_exec_server::ResolvedSelectedCapabilityRoot; use codex_mcp::ElicitationReviewerHandle; use codex_mcp::McpConfig; @@ -28,7 +27,6 @@ pub(super) enum McpRuntimeScope<'a> { struct ProjectedMcpConfig { bindings: Vec<(usize, ResolvedSelectedCapabilityRoot)>, - plugins: Vec<(usize, ExecutorPluginRuntime)>, config: McpConfig, runtime_context: McpRuntimeContext, } @@ -50,25 +48,7 @@ impl Session { } let bindings = selected_bindings(selected_roots, resolved_roots); let cached_runtime = cache.runtime_for_bindings(&bindings); - let mut plugins = cache.plugins_for_bindings(&bindings).unwrap_or_default(); - let mut discovered_plugin = false; - if cached_runtime.is_some() { - let unresolved = bindings - .iter() - .filter(|(order, _)| { - !plugins - .iter() - .any(|(plugin_order, _)| plugin_order == order) - }) - .cloned() - .collect::>(); - let discovered = project_executor_plugins(&unresolved).await; - discovered_plugin = !discovered.is_empty(); - plugins.extend(discovered); - plugins.sort_unstable_by_key(|(order, _)| *order); - } else { - plugins = project_executor_plugins(&bindings).await; - } + let plugins = cache.project_plugins(&bindings).await; let base = cache.base_runtime(); let mcp_config = Arc::new( @@ -83,8 +63,7 @@ impl Session { .collect::>(); let runtime_context = self.mcp_runtime_context(&turn_context.config, environments, &pinned_roots); - if !discovered_plugin - && let Some(runtime) = cached_runtime + if let Some(runtime) = cached_runtime && runtime.matches_projection(mcp_config.as_ref(), &runtime_context) { runtime @@ -109,7 +88,7 @@ impl Session { ) .await }; - cache.replace_selected_runtime(bindings, plugins, Arc::clone(&runtime)); + cache.replace_selected_runtime(bindings, Arc::clone(&runtime)); self.services .publish_existing_mcp_runtime(Arc::clone(&runtime)); runtime @@ -119,11 +98,16 @@ impl Session { &self, config: &Config, ) -> (McpConfig, McpRuntimeContext) { - let projection = self.project_mcp_config_inner(config).await; + let mut cache = self.services.selected_mcp_runtime.lock().await; + let projection = self.project_mcp_config_inner(config, &mut cache).await; (projection.config, projection.runtime_context) } - async fn project_mcp_config_inner(&self, config: &Config) -> ProjectedMcpConfig { + async fn project_mcp_config_inner( + &self, + config: &Config, + cache: &mut super::SelectedMcpRuntimeCache, + ) -> ProjectedMcpConfig { let environments = self.services.turn_environments.snapshot().await; let selected_roots = &self.services.selected_capability_roots; let resolved_roots = self @@ -136,7 +120,7 @@ impl Session { ) .await; let bindings = selected_bindings(selected_roots, &resolved_roots); - let plugins = project_executor_plugins(&bindings).await; + let plugins = cache.project_plugins(&bindings).await; let base_config = self.services.mcp_manager.runtime_config(config).await; let mcp_config = self .services @@ -145,7 +129,6 @@ impl Session { let runtime_context = self.mcp_runtime_context(config, &environments, &resolved_roots); ProjectedMcpConfig { bindings, - plugins, config: mcp_config, runtime_context, } @@ -156,7 +139,7 @@ impl Session { config: &Config, ) -> Arc { let mut cache = self.services.selected_mcp_runtime.lock().await; - let projection = self.project_mcp_config_inner(config).await; + let projection = self.project_mcp_config_inner(config, &mut cache).await; let current = cache .runtime_for_bindings(&projection.bindings) .unwrap_or_else(|| self.services.latest_mcp_runtime()); @@ -183,11 +166,7 @@ impl Session { if projection.bindings.is_empty() { cache.replace_base_and_invalidate_selected(Arc::clone(&runtime)); } else { - cache.replace_selected_runtime( - projection.bindings, - projection.plugins, - Arc::clone(&runtime), - ); + cache.replace_selected_runtime(projection.bindings, Arc::clone(&runtime)); } self.services .publish_existing_mcp_runtime(Arc::clone(&runtime)); @@ -309,23 +288,3 @@ fn selected_bindings( }) .collect() } - -async fn project_executor_plugins( - bindings: &[(usize, ResolvedSelectedCapabilityRoot)], -) -> Vec<(usize, ExecutorPluginRuntime)> { - let mut plugins = Vec::new(); - for (selection_order, root) in bindings { - match ExecutorPluginRuntime::project(root).await { - Ok(Some(runtime)) => plugins.push((*selection_order, runtime)), - Ok(None) => {} - Err(err) => { - tracing::warn!( - selected_root = root.selected_root().id, - error = %err, - "failed to project selected executor plugin runtime" - ); - } - } - } - plugins -} diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index a58d526c7ba2..263bf96dd9ff 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -3,36 +3,43 @@ use std::sync::Arc; use codex_core_plugins::ExecutorPluginRuntime; use codex_exec_server::ResolvedSelectedCapabilityRoot; use codex_mcp::McpRuntimeSnapshot; +use codex_protocol::capabilities::SelectedCapabilityRoot; /// One live selected-plugin MCP runtime retained between model steps. /// -/// A cached runtime is a reuse candidate only for the same ordered selected roots and the same -/// process-local environment instances. The caller additionally compares the effective MCP config -/// and runtime context before reuse. Selected environment contents are treated as stable, so -/// manifest and MCP config file changes do not invalidate this cache. +/// Selected environment identity and contents are stable. Plugin manifests, MCP declarations, and +/// app declarations are therefore cached by the complete selected root for the session lifetime. +/// Missing or failed projections are not cached, so a deferred environment can recover later. +/// +/// A live runtime is a separate cache: it is reusable only for the same ordered selected roots and +/// the same process-local environment handles. The caller additionally compares the effective MCP +/// config and runtime context. Replacing a connection handle may rebuild live processes, but it +/// reuses the stable plugin projection and does not reread capability files. /// /// Within a live session, the selected runtime is invalidated in exactly two ways: /// /// 1. [`Self::replace_base_and_invalidate_selected`] installs a new base MCP runtime. /// 2. [`Self::replace_selected_runtime`] stores a newly projected runtime. This happens when the -/// bindings change, when a previously unavailable plugin appears, or when the effective config -/// or runtime context changes. An unavailable environment disappears from the binding list and -/// therefore follows this path; returning with a new environment instance rebuilds the live -/// runtime even when the stable environment ID is unchanged. +/// active bindings or effective runtime configuration change. /// /// In-flight [`McpRuntimeSnapshot`] values retain their manager until their model step finishes. #[derive(Default)] pub(crate) struct SelectedMcpRuntimeCache { base_runtime: Option>, + plugin_projections: Vec, runtime: Option, } struct CachedSelectedRuntime { bindings: Vec<(usize, ResolvedSelectedCapabilityRoot)>, - plugins: Vec<(usize, ExecutorPluginRuntime)>, runtime: Arc, } +struct CachedExecutorPluginProjection { + selected_root: SelectedCapabilityRoot, + plugin: ExecutorPluginRuntime, +} + impl SelectedMcpRuntimeCache { pub(crate) fn replace_base_and_invalidate_selected( &mut self, @@ -59,27 +66,51 @@ impl SelectedMcpRuntimeCache { .map(|cached| Arc::clone(&cached.runtime)) } - pub(crate) fn plugins_for_bindings( - &self, + pub(crate) async fn project_plugins( + &mut self, bindings: &[(usize, ResolvedSelectedCapabilityRoot)], - ) -> Option> { - self.runtime - .as_ref() - .filter(|cached| same_bindings(&cached.bindings, bindings)) - .map(|cached| cached.plugins.clone()) + ) -> Vec<(usize, ExecutorPluginRuntime)> { + let mut plugins = Vec::new(); + for (selection_order, root) in bindings { + let selected_root = root.selected_root(); + if let Some(plugin) = self + .plugin_projections + .iter() + .find(|cached| &cached.selected_root == selected_root) + .map(|cached| cached.plugin.clone()) + { + plugins.push((*selection_order, plugin)); + continue; + } + + match ExecutorPluginRuntime::project(root).await { + Ok(Some(plugin)) => { + self.plugin_projections + .push(CachedExecutorPluginProjection { + selected_root: selected_root.clone(), + plugin: plugin.clone(), + }); + plugins.push((*selection_order, plugin)); + } + Ok(None) => {} + Err(err) => { + tracing::warn!( + selected_root = selected_root.id, + error = %err, + "failed to project selected executor plugin runtime" + ); + } + } + } + plugins } pub(crate) fn replace_selected_runtime( &mut self, bindings: Vec<(usize, ResolvedSelectedCapabilityRoot)>, - plugins: Vec<(usize, ExecutorPluginRuntime)>, runtime: Arc, ) { - self.runtime = Some(CachedSelectedRuntime { - bindings, - plugins, - runtime, - }); + self.runtime = Some(CachedSelectedRuntime { bindings, runtime }); } } @@ -88,8 +119,8 @@ fn same_bindings( right: &[(usize, ResolvedSelectedCapabilityRoot)], ) -> bool { // Order is part of the key because later selected roots can be renamed when MCP server names - // collide. Arc identity is part of the key because live processes and connections belong to - // one exact environment instance, even when a replacement reuses the same stable ID. + // collide. Arc identity is only a live-connection key: stable plugin metadata is cached above + // by selected root and survives connection-handle replacement. left.len() == right.len() && left .iter() From 1cef19d72cf4743539beecde54787324626dffe1 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 25 Jun 2026 17:44:16 +0100 Subject: [PATCH 6/6] Apply stable environment identity across MCP projection --- codex-rs/core-plugins/src/executor_runtime.rs | 8 +++----- codex-rs/core/src/session/mcp_projection.rs | 2 +- codex-rs/core/src/session/mcp_runtime.rs | 19 +++++++++++-------- codex-rs/core/src/session/session.rs | 10 +++++++--- 4 files changed, 22 insertions(+), 17 deletions(-) diff --git a/codex-rs/core-plugins/src/executor_runtime.rs b/codex-rs/core-plugins/src/executor_runtime.rs index 5c6becf9973c..0f7e7a5e8e24 100644 --- a/codex-rs/core-plugins/src/executor_runtime.rs +++ b/codex-rs/core-plugins/src/executor_runtime.rs @@ -17,7 +17,7 @@ use crate::ResolvedExecutorPlugin; const DEFAULT_MCP_CONFIG_FILE: &str = ".mcp.json"; -/// MCP and connector declarations read from one exact executor binding. +/// MCP and connector declarations read from one stable executor capability root. #[derive(Clone, Debug)] pub struct ExecutorPluginRuntime { plugin: ResolvedPlugin, @@ -26,10 +26,8 @@ pub struct ExecutorPluginRuntime { } impl ExecutorPluginRuntime { - /// Reads both runtime declaration files through the root's pinned filesystem. - /// - /// `Ok(None)` is intentionally not cacheable: the plugin manifest may appear - /// once a deferred executor finishes starting. + /// Reads both runtime declaration files through the root's current ready filesystem. + /// `Ok(None)` means the stable root is not a plugin and may be cached by the caller. pub async fn project(root: &ResolvedSelectedCapabilityRoot) -> anyhow::Result> { let Some(plugin) = ExecutorPluginProvider::resolve_pinned(root).await? else { return Ok(None); diff --git a/codex-rs/core/src/session/mcp_projection.rs b/codex-rs/core/src/session/mcp_projection.rs index 143a779816f5..d1dcf6956d55 100644 --- a/codex-rs/core/src/session/mcp_projection.rs +++ b/codex-rs/core/src/session/mcp_projection.rs @@ -116,7 +116,7 @@ impl Session { .environment_manager() .resolve_selected_capability_roots( selected_roots, - &environments.captured_environments(), + &environments.captured_environment_availability(), ) .await; let bindings = selected_bindings(selected_roots, &resolved_roots); diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index 263bf96dd9ff..2bec6cb1ae7b 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -9,7 +9,8 @@ use codex_protocol::capabilities::SelectedCapabilityRoot; /// /// Selected environment identity and contents are stable. Plugin manifests, MCP declarations, and /// app declarations are therefore cached by the complete selected root for the session lifetime. -/// Missing or failed projections are not cached, so a deferred environment can recover later. +/// Successful projections and stable non-plugin roots are cached. Failed reads are not cached, so +/// a transient executor error can recover later. /// /// A live runtime is a separate cache: it is reusable only for the same ordered selected roots and /// the same process-local environment handles. The caller additionally compares the effective MCP @@ -37,7 +38,7 @@ struct CachedSelectedRuntime { struct CachedExecutorPluginProjection { selected_root: SelectedCapabilityRoot, - plugin: ExecutorPluginRuntime, + plugin: Option, } impl SelectedMcpRuntimeCache { @@ -73,26 +74,28 @@ impl SelectedMcpRuntimeCache { let mut plugins = Vec::new(); for (selection_order, root) in bindings { let selected_root = root.selected_root(); - if let Some(plugin) = self + if let Some(cached) = self .plugin_projections .iter() .find(|cached| &cached.selected_root == selected_root) - .map(|cached| cached.plugin.clone()) { - plugins.push((*selection_order, plugin)); + if let Some(plugin) = &cached.plugin { + plugins.push((*selection_order, plugin.clone())); + } continue; } match ExecutorPluginRuntime::project(root).await { - Ok(Some(plugin)) => { + Ok(plugin) => { self.plugin_projections .push(CachedExecutorPluginProjection { selected_root: selected_root.clone(), plugin: plugin.clone(), }); - plugins.push((*selection_order, plugin)); + if let Some(plugin) = plugin { + plugins.push((*selection_order, plugin)); + } } - Ok(None) => {} Err(err) => { tracing::warn!( selected_root = selected_root.id, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 8ab001294a1b..761f74a07c4e 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -557,10 +557,14 @@ impl Session { config.current_time_reminder.as_ref(), external_time_provider, )?; - let selected_capability_roots = if selected_capability_roots.is_empty() { - initial_history.get_selected_capability_roots() - } else { + let selected_capability_roots = if !selected_capability_roots.is_empty() { selected_capability_roots + } else if let Some(selected_capability_roots) = + thread_extension_init.get::>() + { + selected_capability_roots.as_ref().clone() + } else { + initial_history.get_selected_capability_roots() }; let thread_extension_data = codex_extension_api::ExtensionData::new_with_init( thread_id.to_string(),