diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 8cdaa316a8..98e5031426 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3386,6 +3386,7 @@ dependencies = [ "codex-utils-path-uri", "codex-utils-plugins", "futures", + "lru 0.16.3", "pretty_assertions", "regex-lite", "rmcp", 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 34ec3887ce..1cfbc65672 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -1,6 +1,5 @@ use super::*; -use codex_connectors::ConnectorRuntimeManager; -use codex_mcp::ToolInfo; +use codex_core::McpManager; const MCP_TOOL_THREAD_ID_META_KEY: &str = "threadId"; @@ -246,7 +245,6 @@ impl McpRequestProcessor { None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None), }; let mcp_manager = self.thread_manager.mcp_manager(); - let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); let auth = self.auth_manager.auth().await; let (mcp_config, runtime_context) = match thread { Some(thread) => { @@ -272,7 +270,7 @@ impl McpRequestProcessor { mcp_config, auth, runtime_context, - codex_apps_tools_cache, + mcp_manager, ) .await; }); @@ -286,7 +284,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, - codex_apps_tools_cache: ConnectorRuntimeManager, + mcp_manager: Arc, ) { let result = Self::list_mcp_server_status_response( request_id.request_id.to_string(), @@ -294,7 +292,7 @@ impl McpRequestProcessor { mcp_config, auth, runtime_context, - codex_apps_tools_cache, + mcp_manager, ) .await; outgoing.send_result(request_id, result).await; @@ -306,7 +304,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, - codex_apps_tools_cache: ConnectorRuntimeManager, + mcp_manager: Arc, ) -> Result { let detail = match params.detail.unwrap_or(McpServerStatusDetail::Full) { McpServerStatusDetail::Full => McpSnapshotDetail::Full, @@ -318,7 +316,8 @@ impl McpRequestProcessor { auth.as_ref(), request_id, runtime_context, - codex_apps_tools_cache, + mcp_manager.codex_apps_tools_cache(), + mcp_manager.tool_catalog_cache(), detail, ) .await; @@ -412,6 +411,7 @@ impl McpRequestProcessor { let mcp_manager = self.thread_manager.mcp_manager(); let mcp_config = mcp_manager.runtime_config(&config).await; let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); + let tool_catalog_cache = mcp_manager.tool_catalog_cache(); let auth = self.auth_manager.auth().await; let environment_manager = self.thread_manager.environment_manager(); // This threadless resource-read path has no turn cwd or turn-selected @@ -427,6 +427,7 @@ impl McpRequestProcessor { auth.as_ref(), runtime_context, codex_apps_tools_cache, + tool_catalog_cache, &server, &uri, ) diff --git a/codex-rs/codex-mcp/Cargo.toml b/codex-rs/codex-mcp/Cargo.toml index 428134b193..e1da7b22fc 100644 --- a/codex-rs/codex-mcp/Cargo.toml +++ b/codex-rs/codex-mcp/Cargo.toml @@ -29,6 +29,7 @@ codex-rmcp-client = { workspace = true } codex-utils-path-uri = { workspace = true } codex-utils-plugins = { workspace = true } futures = { workspace = true } +lru = { workspace = true } regex-lite = { workspace = true } rmcp = { workspace = true, default-features = false, features = ["base64", "macros", "schemars", "server"] } serde = { workspace = true, features = ["derive"] } diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index fbd5e59d87..eaf0719dfd 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -30,6 +30,7 @@ use crate::runtime::McpRuntimeContext; use crate::runtime::emit_duration; use crate::server::EffectiveMcpServer; use crate::server::McpServerMetadata; +use crate::tool_catalog_cache::McpToolCatalogCache; use crate::tools::ToolInfo; use crate::tools::filter_tools; use crate::tools::normalize_tools_for_model_with_prefix; @@ -136,6 +137,7 @@ impl McpConnectionManager { runtime_context: McpRuntimeContext, codex_home: PathBuf, codex_apps_tools_cache: ConnectorRuntimeManager, + tool_catalog_cache: McpToolCatalogCache, codex_apps_tools_cache_key: ConnectorRuntimeContextKey, prefix_mcp_tool_names: bool, client_elicitation_capability: ElicitationCapability, @@ -190,6 +192,10 @@ impl McpConnectionManager { ) .await; let configured_config = server.configured_config().cloned(); + let resolved_environment = configured_config.as_ref().map_or_else( + || Ok(None), + |config| runtime_context.resolve_server_environment(&server_name, config), + ); // For built-in Codex Apps, `CODEX_CONNECTORS_TOKEN` is a debug // override: it supplies runtime auth but bypasses the shared tools // cache. @@ -228,6 +234,22 @@ impl McpConnectionManager { } else { chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider) }; + let tool_catalog_cache_context = if server_name == CODEX_APPS_MCP_SERVER_NAME { + None + } else if let Some(config) = configured_config.as_ref() + && let Ok(environment) = resolved_environment.as_ref() + { + tool_catalog_cache.context( + &server_name, + config, + &runtime_context, + environment.as_ref(), + &client_elicitation_capability, + supports_openai_form_elicitation, + ) + } else { + None + }; let has_runtime_auth = runtime_auth_provider.is_some(); let async_managed_client = AsyncManagedClient::new( server_name.clone(), @@ -239,8 +261,10 @@ impl McpConnectionManager { tx_event.clone(), elicitation_requests.clone(), codex_apps_tools_cache_context, + tool_catalog_cache_context, Arc::clone(&tool_plugin_provenance), runtime_context.clone(), + resolved_environment, runtime_auth_provider, client_elicitation_capability.clone(), supports_openai_form_elicitation, @@ -586,6 +610,17 @@ impl McpConnectionManager { tools } + /// Returns one tool from the current live connection. + pub async fn tool_info(&self, server: &str, tool: &str) -> Option { + let client = self.clients.get(server)?; + let managed_client = client.client().await.ok()?; + let tool = client + .prepare_tools(managed_client.listed_tools()) + .into_iter() + .find(|tool_info| tool_info.tool.name == tool)?; + Some(self.with_server_metadata(tool)) + } + /// Force-refresh codex apps tools by bypassing the in-process cache. /// /// On success, the refreshed tools replace shared cache contents when the @@ -895,15 +930,16 @@ impl McpConnectionManager { .with_context(|| format!("resources/read failed for `{server}` ({uri})")) } - /// Returns presentation metadata without waiting for uncached clients still initializing. - /// Cached values will be used if available and the server is still starting up. + /// Returns presentation metadata from the current connection. + /// Codex Apps metadata may come from its existing cache; regular MCP server information is + /// connection-specific, so pending regular clients are awaited. pub(crate) async fn list_available_server_infos(&self) -> HashMap { let mut server_infos = HashMap::new(); for (server_name, client) in &self.clients { - if !client.startup_complete.load(Ordering::Acquire) { - if let Some(server_info) = client.cached_server_info.clone() { - server_infos.insert(server_name.clone(), server_info); - } + if !client.startup_complete.load(Ordering::Acquire) + && let Some(server_info) = client.cached_server_info.clone() + { + server_infos.insert(server_name.clone(), server_info); continue; } match client.client().await { diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index dece950259..12d50d4130 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -142,6 +142,7 @@ async fn create_ready_async_managed_client(tools: Vec) -> AsyncManaged is_codex_apps_mcp_server: false, cached_server_info: None, codex_apps_tools_cache_context: None, + tool_catalog_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), startup_reconnect: None, @@ -181,6 +182,7 @@ fn create_test_manager_with_failed_apps_startup( is_codex_apps_mcp_server: true, cached_server_info: None, codex_apps_tools_cache_context: Some(cache_context), + tool_catalog_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), startup_reconnect: Some(Arc::new(CodexAppsStartupReconnect::new(reconnect_factory))), @@ -722,6 +724,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { is_codex_apps_mcp_server: true, cached_server_info: None, codex_apps_tools_cache_context: Some(cache_context), + tool_catalog_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, @@ -742,6 +745,77 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { assert_eq!(tool.callable_name, "calendar_create_event"); } +#[tokio::test(start_paused = true)] +async fn tool_catalog_cache_sanitizes_tools_and_tracks_environment_generation() { + let cache = McpToolCatalogCache::default(); + let environment_manager = Arc::new(EnvironmentManager::without_environments()); + let replace_environment = |url: &str| { + environment_manager + .upsert_environment( + "remote".to_string(), + url.to_string(), + /*connect_timeout*/ None, + ) + .expect("replace environment"); + }; + replace_environment("ws://127.0.0.1:1"); + let runtime_context = + McpRuntimeContext::new(Arc::clone(&environment_manager), PathBuf::from("/tmp")); + let config: McpServerConfig = serde_json::from_value(serde_json::json!({ + "command": "docs-mcp", + "environment_id": "remote" + })) + .expect("MCP config"); + let resolve_environment = || { + runtime_context + .resolve_server_environment("docs", &config) + .expect("resolve environment") + .expect("remote environment") + }; + let cache_context = |environment: &Arc| { + cache + .context( + "docs", + &config, + &runtime_context, + Some(environment), + &ElicitationCapability::default(), + /*supports_openai_form_elicitation*/ false, + ) + .expect("cache context") + }; + let first_environment = resolve_environment(); + let first_environment_weak = Arc::downgrade(&first_environment); + let first_context = cache_context(&first_environment); + let mut tool = create_test_tool("docs", "search"); + tool.tool.annotations = Some(rmcp::model::ToolAnnotations::new().read_only(true)); + first_context.publish_if_newest(first_context.begin_fetch(), &[tool]); + assert_eq!( + first_context.current_tools().expect("cached tools")[0] + .tool + .annotations, + None + ); + + drop(first_environment); + replace_environment("ws://127.0.0.1:2"); + assert!(first_environment_weak.upgrade().is_none()); + let replacement_environment = resolve_environment(); + assert!(!cache_context(&replacement_environment).has_tools()); + + let older = first_context.begin_fetch(); + let newer = first_context.begin_fetch(); + first_context.publish_if_newest(newer, &[create_test_tool("docs", "new")]); + first_context.publish_if_newest(older, &[create_test_tool("docs", "old")]); + assert_eq!( + first_context.current_tools().expect("cached tools")[0].callable_name, + "new" + ); + + tokio::time::advance(Duration::from_secs(30 * 60 + 1)).await; + assert!(!first_context.has_tools()); +} + #[tokio::test] async fn list_available_server_infos_uses_cache_while_client_is_pending() { let pending_client = futures::future::pending::>() @@ -762,6 +836,7 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { is_codex_apps_mcp_server: true, cached_server_info: Some(server_info.clone()), codex_apps_tools_cache_context: None, + tool_catalog_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, @@ -863,6 +938,7 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { is_codex_apps_mcp_server: true, cached_server_info: None, codex_apps_tools_cache_context: None, + tool_catalog_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, @@ -918,6 +994,7 @@ async fn shutdown_cancels_pending_tool_listing() { is_codex_apps_mcp_server: true, cached_server_info: None, codex_apps_tools_cache_context: None, + tool_catalog_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, @@ -965,6 +1042,7 @@ async fn shutdown_continues_after_caller_is_aborted() { is_codex_apps_mcp_server: true, cached_server_info: None, codex_apps_tools_cache_context: None, + tool_catalog_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, @@ -1018,6 +1096,7 @@ async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() { is_codex_apps_mcp_server: true, cached_server_info: None, codex_apps_tools_cache_context: Some(cache_context), + tool_catalog_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, @@ -1071,6 +1150,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() is_codex_apps_mcp_server: true, cached_server_info: Some(server_info.clone()), codex_apps_tools_cache_context: Some(cache_context), + tool_catalog_cache_context: None, tool_filter: ToolFilter::default(), startup_complete, startup_reconnect: None, @@ -1503,6 +1583,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { ), codex_home.path().to_path_buf(), ConnectorRuntimeManager::::default(), + McpToolCatalogCache::default(), ConnectorRuntimeContextKey::personal( /*account_id*/ None, /*chatgpt_user_id*/ None, ), @@ -1520,6 +1601,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { assert!(manager.clients.contains_key("stdio")); assert!(manager.clients.contains_key("http")); + assert!(manager.clients["http"].tool_catalog_cache_context.is_none()); assert!( !manager .wait_for_server_ready("stdio", Duration::from_millis(10)) diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 079444070a..fa1d19023b 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -12,6 +12,7 @@ pub use resource_client::McpResourceReadResult; pub use rmcp_client::MCP_SANDBOX_STATE_META_CAPABILITY; pub use runtime::McpRuntimeContext; pub use runtime::SandboxState; +pub use tool_catalog_cache::McpToolCatalogCache; pub use tools::ToolInfo; /// Backward-compatible name for the shared Codex Apps tools runtime. @@ -90,4 +91,5 @@ mod resource_client; pub(crate) mod rmcp_client; pub(crate) mod runtime; pub(crate) mod server; +mod tool_catalog_cache; pub(crate) mod tools; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 345f425faa..ad42c5425c 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -305,6 +305,7 @@ pub async fn read_mcp_resource( auth: Option<&CodexAuth>, runtime_context: McpRuntimeContext, codex_apps_tools_cache: ConnectorRuntimeManager, + tool_catalog_cache: crate::McpToolCatalogCache, server: &str, uri: &str, ) -> anyhow::Result { @@ -325,6 +326,7 @@ pub async fn read_mcp_resource( runtime_context, config.codex_home.clone(), codex_apps_tools_cache, + tool_catalog_cache, connector_runtime_context_key(auth), config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), @@ -361,6 +363,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( submit_id: String, runtime_context: McpRuntimeContext, codex_apps_tools_cache: ConnectorRuntimeManager, + tool_catalog_cache: crate::McpToolCatalogCache, detail: McpSnapshotDetail, ) -> McpServerStatusSnapshot { let mcp_servers = effective_mcp_servers(config, auth); @@ -403,6 +406,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( runtime_context, config.codex_home.clone(), codex_apps_tools_cache, + tool_catalog_cache, connector_runtime_context_key(auth), config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), @@ -662,8 +666,12 @@ async fn collect_mcp_server_status_snapshot_from_manager( server_names: Vec, detail: McpSnapshotDetail, ) -> McpServerStatusSnapshot { - let (tools, resources, resource_templates) = tokio::join!( - mcp_connection_manager.list_all_tools(), + let ((server_infos, tools), resources, resource_templates) = tokio::join!( + async { + let server_infos = mcp_connection_manager.list_available_server_infos().await; + let tools = mcp_connection_manager.list_all_tools().await; + (server_infos, tools) + }, async { if detail.include_resources() { mcp_connection_manager.list_all_resources(|_| true).await @@ -681,7 +689,6 @@ async fn collect_mcp_server_status_snapshot_from_manager( } }, ); - let server_infos = mcp_connection_manager.list_available_server_infos().await; let mut tools_by_server = HashMap::>::new(); for tool_info in tools { diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index aefcac67ea..4ac96f867f 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -29,6 +29,8 @@ use crate::runtime::McpRuntimeContext; use crate::runtime::emit_duration; use crate::server::EffectiveMcpServer; use crate::server::McpServerLaunch; +use crate::tool_catalog_cache::McpToolCatalogCacheContext; +use crate::tool_catalog_cache::McpToolCatalogFetchTicket; use crate::tools::ToolFilter; use crate::tools::ToolInfo; use crate::tools::filter_tools; @@ -44,6 +46,7 @@ use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_connectors::ConnectorRuntimeContext; use codex_connectors::ConnectorRuntimeFetchSource; +use codex_exec_server::Environment; use codex_exec_server::HttpClient; use codex_exec_server::ReqwestHttpClient; use codex_protocol::mcp::McpServerInfo; @@ -109,7 +112,7 @@ pub(crate) struct ManagedClient { } impl ManagedClient { - fn listed_tools(&self) -> Vec { + pub(crate) fn listed_tools(&self) -> Vec { let total_start = Instant::now(); if let Some(tools) = self .codex_apps_tools_cache_context @@ -272,7 +275,9 @@ struct ManagedClientStartup { tx_event: Sender, elicitation_requests: ElicitationRequestManager, codex_apps_tools_cache_context: Option>, + tool_catalog_cache_context: Option, runtime_context: McpRuntimeContext, + resolved_environment: std::result::Result>, String>, runtime_auth_provider: Option, client_elicitation_capability: ElicitationCapability, supports_openai_form_elicitation: bool, @@ -290,7 +295,9 @@ impl ManagedClientStartup { tx_event, elicitation_requests, codex_apps_tools_cache_context, + tool_catalog_cache_context, runtime_context, + resolved_environment, runtime_auth_provider, client_elicitation_capability, supports_openai_form_elicitation, @@ -307,6 +314,9 @@ impl ManagedClientStartup { .and_then(|config| config.startup_timeout_sec) .unwrap_or(DEFAULT_STARTUP_TIMEOUT); let cancel_token_for_fut = cancel_token; + let tool_catalog_fetch_ticket = tool_catalog_cache_context + .as_ref() + .map(McpToolCatalogCacheContext::begin_fetch); async move { let refresh_start = is_codex_apps_mcp_server.then(Instant::now); let outcome = match async { @@ -322,6 +332,7 @@ impl ManagedClientStartup { store_mode, keyring_backend_kind, runtime_context, + resolved_environment, runtime_auth_provider, ), ) @@ -348,6 +359,8 @@ impl ManagedClientStartup { tx_event, elicitation_requests, codex_apps_tools_cache_context, + tool_catalog_cache_context, + tool_catalog_fetch_ticket, client_elicitation_capability, supports_openai_form_elicitation, }, @@ -385,6 +398,7 @@ pub(crate) struct AsyncManagedClient { pub(crate) is_codex_apps_mcp_server: bool, pub(crate) cached_server_info: Option, pub(crate) codex_apps_tools_cache_context: Option>, + pub(crate) tool_catalog_cache_context: Option, pub(crate) tool_filter: ToolFilter, pub(crate) startup_complete: Arc, pub(crate) startup_reconnect: Option>, @@ -407,8 +421,10 @@ impl AsyncManagedClient { tx_event: Sender, elicitation_requests: ElicitationRequestManager, codex_apps_tools_cache_context: Option>, + tool_catalog_cache_context: Option, tool_plugin_provenance: Arc, runtime_context: McpRuntimeContext, + resolved_environment: std::result::Result>, String>, runtime_auth_provider: Option, client_elicitation_capability: ElicitationCapability, supports_openai_form_elicitation: bool, @@ -436,7 +452,9 @@ impl AsyncManagedClient { tx_event, elicitation_requests, codex_apps_tools_cache_context: codex_apps_tools_cache_context.clone(), + tool_catalog_cache_context: tool_catalog_cache_context.clone(), runtime_context, + resolved_environment, runtime_auth_provider, client_elicitation_capability, supports_openai_form_elicitation, @@ -458,6 +476,9 @@ impl AsyncManagedClient { if codex_apps_tools_cache_context .as_ref() .is_some_and(ConnectorRuntimeContext::has_current_tools) + || tool_catalog_cache_context + .as_ref() + .is_some_and(McpToolCatalogCacheContext::has_tools) { let startup_task = client.clone(); tokio::spawn(async move { @@ -470,6 +491,7 @@ impl AsyncManagedClient { is_codex_apps_mcp_server, cached_server_info, codex_apps_tools_cache_context, + tool_catalog_cache_context, tool_filter, startup_complete, startup_reconnect, @@ -516,17 +538,34 @@ impl AsyncManagedClient { self.codex_apps_tools_cache_context .as_ref() .is_some_and(ConnectorRuntimeContext::has_current_tools) + || self + .tool_catalog_cache_context + .as_ref() + .is_some_and(McpToolCatalogCacheContext::has_tools) } fn cached_tools(&self) -> Option> { self.codex_apps_tools_cache_context .as_ref() .and_then(ConnectorRuntimeContext::current_tools) + .or_else(|| { + self.tool_catalog_cache_context + .as_ref() + .and_then(McpToolCatalogCacheContext::current_tools) + }) .map(|tools| filter_tools(tools, &self.tool_filter)) } + pub(crate) fn prepare_tools(&self, tools: Vec) -> Vec { + if self.is_codex_apps_mcp_server { + prepare_codex_apps_tools_for_model(tools, &self.tool_plugin_provenance) + } else { + prepare_regular_mcp_tools_for_model(tools, &self.tool_plugin_provenance) + } + } + pub(crate) async fn listed_tools(&self) -> Option> { - // Keep cache payloads raw; plugin provenance is resolved per-session at read time. + // Plugin provenance is resolved per-session rather than stored in shared cache payloads. let tools = if !self.startup_complete.load(Ordering::Acquire) && let Some(startup_tools) = self.cached_tools() { @@ -534,14 +573,11 @@ impl AsyncManagedClient { } else { match self.client().await { Ok(client) => Some(client.listed_tools()), - Err(_) => self.cached_tools(), + Err(_) if self.is_codex_apps_mcp_server => self.cached_tools(), + Err(_) => None, } }?; - Some(if self.is_codex_apps_mcp_server { - prepare_codex_apps_tools_for_model(tools, &self.tool_plugin_provenance) - } else { - prepare_regular_mcp_tools_for_model(tools, &self.tool_plugin_provenance) - }) + Some(self.prepare_tools(tools)) } } @@ -835,6 +871,8 @@ async fn start_server_task( tx_event, elicitation_requests, codex_apps_tools_cache_context, + tool_catalog_cache_context, + tool_catalog_fetch_ticket, client_elicitation_capability, supports_openai_form_elicitation, } = params; @@ -842,7 +880,6 @@ async fn start_server_task( client_elicitation_capability, supports_openai_form_elicitation, ); - let send_elicitation = elicitation_requests.make_sender(server_name.clone(), tx_event); let initialize_result = client @@ -878,7 +915,14 @@ async fn start_server_task( (None, None) => tools, _ => unreachable!("Codex Apps fetch ticket requires cache context"), }; - if is_codex_apps_mcp_server { + let has_shared_tool_catalog = is_codex_apps_mcp_server || tool_catalog_cache_context.is_some(); + if let (Some(cache_context), Some(fetch_ticket)) = ( + tool_catalog_cache_context.as_ref(), + tool_catalog_fetch_ticket, + ) { + cache_context.publish_if_newest(fetch_ticket, &tools); + } + if has_shared_tool_catalog { emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, list_start.elapsed(), @@ -944,6 +988,8 @@ struct StartServerTaskParams { tx_event: Sender, elicitation_requests: ElicitationRequestManager, codex_apps_tools_cache_context: Option>, + tool_catalog_cache_context: Option, + tool_catalog_fetch_ticket: Option, client_elicitation_capability: ElicitationCapability, supports_openai_form_elicitation: bool, } @@ -955,14 +1001,14 @@ async fn make_rmcp_client( store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, runtime_context: McpRuntimeContext, + resolved_environment: std::result::Result>, String>, runtime_auth_provider: Option, ) -> Result { let config = match server.launch() { McpServerLaunch::Configured(config) => config.as_ref().clone(), }; - let resolved_environment = runtime_context - .resolve_server_environment(server_name, &config) - .map_err(|err| StartupOutcomeError::from(anyhow!(err)))?; + let resolved_environment = + resolved_environment.map_err(|err| StartupOutcomeError::from(anyhow!(err)))?; let is_local_environment = config.is_local_environment(); let McpServerConfig { transport, .. } = config; diff --git a/codex-rs/codex-mcp/src/tool_catalog_cache.rs b/codex-rs/codex-mcp/src/tool_catalog_cache.rs new file mode 100644 index 0000000000..4f2f57140c --- /dev/null +++ b/codex-rs/codex-mcp/src/tool_catalog_cache.rs @@ -0,0 +1,276 @@ +use std::collections::BTreeMap; +use std::collections::hash_map::DefaultHasher; +use std::hash::Hash; +use std::hash::Hasher; +use std::num::NonZeroUsize; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::sync::Weak; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use codex_config::McpServerConfig; +use codex_config::McpServerTransportConfig; +use codex_exec_server::Environment; +use lru::LruCache; +use rmcp::model::ElicitationCapability; +use sha1::Digest; +use sha1::Sha1; +use tokio::time::Instant; + +use crate::McpRuntimeContext; +use crate::ToolInfo; + +const TOOL_CATALOG_CACHE_CAPACITY: usize = 32; +const TOOL_CATALOG_CACHE_TTL: Duration = Duration::from_secs(30 * 60); + +/// Process-scoped cache of recent reusable tool definitions for MCP servers. +#[derive(Clone)] +pub struct McpToolCatalogCache { + entries: Arc>>>, +} + +impl Default for McpToolCatalogCache { + fn default() -> Self { + Self { + entries: Arc::new(Mutex::new(LruCache::new( + NonZeroUsize::new(TOOL_CATALOG_CACHE_CAPACITY).unwrap_or(NonZeroUsize::MIN), + ))), + } + } +} + +struct ToolCatalogCacheEntry { + state: Mutex, + next_fetch_generation: AtomicU64, +} + +#[derive(Default)] +struct ToolCatalogCacheState { + snapshot: Option, + last_accepted_generation: u64, +} + +struct ToolCatalogSnapshot { + tools: Vec, + published_at: Instant, +} + +#[derive(Clone)] +pub(crate) struct McpToolCatalogCacheContext { + entry: Arc, +} + +pub(crate) struct McpToolCatalogFetchTicket { + generation: u64, +} + +impl McpToolCatalogCache { + pub(crate) fn context( + &self, + server_name: &str, + config: &McpServerConfig, + runtime_context: &McpRuntimeContext, + resolved_environment: Option<&Arc>, + client_elicitation_capability: &ElicitationCapability, + supports_openai_form_elicitation: bool, + ) -> Option { + let identity = ToolCatalogIdentity::new( + server_name, + config, + runtime_context, + resolved_environment, + client_elicitation_capability, + supports_openai_form_elicitation, + )?; + let entry = lock_unpoisoned(&self.entries) + .get_or_insert(identity, || Arc::new(ToolCatalogCacheEntry::default())) + .clone(); + Some(McpToolCatalogCacheContext { entry }) + } +} + +impl Default for ToolCatalogCacheEntry { + fn default() -> Self { + Self { + state: Mutex::new(ToolCatalogCacheState::default()), + next_fetch_generation: AtomicU64::new(0), + } + } +} + +impl McpToolCatalogCacheContext { + pub(crate) fn has_tools(&self) -> bool { + self.current_tools().is_some() + } + + pub(crate) fn current_tools(&self) -> Option> { + lock_unpoisoned(&self.entry.state) + .snapshot + .as_ref() + .filter(|snapshot| snapshot.published_at.elapsed() <= TOOL_CATALOG_CACHE_TTL) + .map(|snapshot| snapshot.tools.clone()) + } + + pub(crate) fn begin_fetch(&self) -> McpToolCatalogFetchTicket { + McpToolCatalogFetchTicket { + generation: self + .entry + .next_fetch_generation + .fetch_add(1, Ordering::Relaxed) + + 1, + } + } + + pub(crate) fn publish_if_newest(&self, ticket: McpToolCatalogFetchTicket, tools: &[ToolInfo]) { + let mut state = lock_unpoisoned(&self.entry.state); + if ticket.generation <= state.last_accepted_generation { + return; + } + + let mut tools = tools.to_vec(); + for tool in &mut tools { + // Initialize instructions belong to one live connection and must not cross sessions. + tool.namespace_description = None; + // Tool annotations affect approval and parallelism decisions, so only the live + // connection may supply them. + tool.tool.annotations = None; + } + state.last_accepted_generation = ticket.generation; + state.snapshot = Some(ToolCatalogSnapshot { + tools, + published_at: Instant::now(), + }); + } +} + +fn lock_unpoisoned(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +struct ToolCatalogIdentity { + server_name: String, + transport: ToolCatalogTransportIdentity, + environment: Option>, + local_stdio_fallback_cwd: Option, +} + +impl PartialEq for ToolCatalogIdentity { + fn eq(&self, other: &Self) -> bool { + self.server_name == other.server_name + && self.transport == other.transport + && self.local_stdio_fallback_cwd == other.local_stdio_fallback_cwd + && match (&self.environment, &other.environment) { + (Some(environment), Some(other)) => Weak::ptr_eq(environment, other), + (None, None) => true, + _ => false, + } + } +} + +impl Eq for ToolCatalogIdentity {} + +impl Hash for ToolCatalogIdentity { + fn hash(&self, state: &mut H) { + self.server_name.hash(state); + self.transport.hash(state); + self.local_stdio_fallback_cwd.hash(state); + self.environment + .as_ref() + .map(|environment| Weak::as_ptr(environment) as usize) + .hash(state); + } +} + +impl ToolCatalogIdentity { + fn new( + server_name: &str, + config: &McpServerConfig, + runtime_context: &McpRuntimeContext, + environment: Option<&Arc>, + client_elicitation_capability: &ElicitationCapability, + supports_openai_form_elicitation: bool, + ) -> Option { + let transport = ToolCatalogTransportIdentity::new( + config, + client_elicitation_capability, + supports_openai_form_elicitation, + )?; + Some(Self { + server_name: server_name.to_string(), + transport, + environment: environment.map(Arc::downgrade), + local_stdio_fallback_cwd: matches!( + &config.transport, + McpServerTransportConfig::Stdio { cwd: None, .. } + ) + .then(|| runtime_context.local_stdio_fallback_cwd()), + }) + } +} + +#[derive(PartialEq, Eq, Hash)] +enum ToolCatalogTransportIdentity { + Stdio { fingerprint: [u8; 20] }, +} + +impl ToolCatalogTransportIdentity { + fn new( + config: &McpServerConfig, + client_elicitation_capability: &ElicitationCapability, + supports_openai_form_elicitation: bool, + ) -> Option { + let McpServerTransportConfig::Stdio { + command, + args, + env, + env_vars, + cwd, + } = &config.transport + else { + // HTTP catalogs need a canonical resolved-auth identity before they can be shared. + return None; + }; + if env_vars + .iter() + .any(codex_config::McpServerEnvVar::is_remote_source) + { + return None; + } + + let mut hasher = Sha1::new(); + let env = env.as_ref().map(|env| { + env.iter() + .map(|(key, value)| (key.as_str(), value.as_str())) + .collect::>() + }); + hasher.update( + serde_json::to_vec(&( + command, + args, + env, + env_vars, + cwd, + &config.environment_id, + client_elicitation_capability, + supports_openai_form_elicitation, + )) + .ok()?, + ); + for env_var in env_vars { + hasher.update(env_var.name().as_bytes()); + let mut value_hasher = DefaultHasher::new(); + std::env::var_os(env_var.name()).hash(&mut value_hasher); + hasher.update(value_hasher.finish().to_le_bytes()); + } + + Some(Self::Stdio { + fingerprint: hasher.finalize().into(), + }) + } +} diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index de791d8bae..56b4623775 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -256,6 +256,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( runtime_context, config.codex_home.to_path_buf(), mcp_manager.codex_apps_tools_cache(), + mcp_manager.tool_catalog_cache(), connector_runtime_context_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability, diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 0930d49175..4d0e30f99f 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -18,6 +18,7 @@ use codex_mcp::EffectiveMcpServer; use codex_mcp::McpConfig; use codex_mcp::McpPluginAttribution; use codex_mcp::McpServerRegistration; +use codex_mcp::McpToolCatalogCache; use codex_mcp::ToolInfo; use codex_mcp::codex_apps_mcp_server_config; use codex_mcp::configured_mcp_servers; @@ -51,6 +52,7 @@ pub struct McpManager { plugins_manager: Arc, extensions: Arc>, codex_apps_tools_cache: ConnectorRuntimeManager, + tool_catalog_cache: McpToolCatalogCache, } impl McpManager { @@ -72,6 +74,7 @@ impl McpManager { plugins_manager, extensions, codex_apps_tools_cache, + tool_catalog_cache: McpToolCatalogCache::default(), } } @@ -79,6 +82,10 @@ impl McpManager { self.codex_apps_tools_cache.clone() } + pub fn tool_catalog_cache(&self) -> McpToolCatalogCache { + self.tool_catalog_cache.clone() + } + /// Returns the MCP config after applying compatibility built-ins and /// runtime-only extension overlays. pub async fn runtime_config(&self, config: &Config) -> McpConfig { diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index d268ea4b34..2d617a8a95 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -42,6 +42,7 @@ use codex_mcp::SandboxState; use codex_mcp::auth_elicitation_completed_result; use codex_mcp::build_auth_elicitation_plan; use codex_mcp::mcp_permission_prompt_is_auto_approved; +use codex_mcp::tool_is_model_visible; use codex_protocol::approvals::ElicitationRequest; use codex_protocol::items::McpToolCallError; use codex_protocol::items::McpToolCallItem; @@ -153,6 +154,23 @@ pub(crate) async fn handle_mcp_tool_call( ) .await; let item_metadata = McpToolCallItemMetadata::from_tool_metadata(&server, metadata.as_ref()); + if metadata.is_none() { + let result = notify_mcp_tool_call_skip( + sess.as_ref(), + turn_context.as_ref(), + &call_id, + invocation, + item_metadata, + format!("MCP tool `{server}/{tool_name}` is not available to the model"), + /*already_started*/ false, + ) + .await; + return HandledMcpToolCall { + result: CallToolResult::from_result(result), + tool_input: arguments_value + .unwrap_or_else(|| JsonValue::Object(serde_json::Map::new())), + }; + } let app_tool_policy = if server == CODEX_APPS_MCP_SERVER_NAME { let annotations = metadata .as_ref() @@ -1480,10 +1498,10 @@ pub(crate) async fn lookup_mcp_tool_metadata( let plugin_id = manager .plugin_id_for_mcp_server_name(server) .map(str::to_string); - let tools = manager.list_all_tools().await; - let tool_info = tools - .into_iter() - .find(|tool_info| tool_info.server_name == server && tool_info.tool.name == tool_name)?; + let tool_info = manager.tool_info(server, tool_name).await?; + if !tool_is_model_visible(&tool_info) { + return None; + } let connector_description = if server == CODEX_APPS_MCP_SERVER_NAME { let connectors = match connectors::list_cached_accessible_connectors_from_mcp_tools( turn_context.config.as_ref(), @@ -1598,17 +1616,10 @@ async fn lookup_mcp_app_usage_metadata( server: &str, tool_name: &str, ) -> Option { - 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 { - Some(McpAppUsageMetadata { - connector_id: tool_info.connector_id, - app_name: tool_info.connector_name, - }) - } else { - None - } + let tool_info = manager.tool_info(server, tool_name).await?; + Some(McpAppUsageMetadata { + connector_id: tool_info.connector_id, + app_name: tool_info.connector_name, }) } diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 6bd857b96f..0653578dde 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1454,6 +1454,7 @@ async fn host_owned_codex_apps_manager( ), turn_context.config.codex_home.to_path_buf(), session.services.mcp_manager.codex_apps_tools_cache(), + session.services.mcp_manager.tool_catalog_cache(), codex_connectors::connector_runtime_context_key(auth.as_ref()), turn_context.config.prefix_mcp_tool_names(), rmcp::model::ElicitationCapability::default(), diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index c215f89c7c..32bbf036f9 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -371,6 +371,7 @@ impl Session { mcp_runtime_context.clone(), mcp_config.codex_home.clone(), self.services.mcp_manager.codex_apps_tools_cache(), + self.services.mcp_manager.tool_catalog_cache(), connector_runtime_context_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability.clone(), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index ecabea1756..2b1943f932 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1206,6 +1206,7 @@ impl Session { mcp_runtime_context.clone(), config.codex_home.to_path_buf(), sess.services.mcp_manager.codex_apps_tools_cache(), + sess.services.mcp_manager.tool_catalog_cache(), connector_runtime_context_key(auth), config.prefix_mcp_tool_names(), mcp_projection diff --git a/codex-rs/core/tests/suite/mcp_auth_refresh.rs b/codex-rs/core/tests/suite/mcp_auth_refresh.rs index 0c05874ccd..a8bd10cfc6 100644 --- a/codex-rs/core/tests/suite/mcp_auth_refresh.rs +++ b/codex-rs/core/tests/suite/mcp_auth_refresh.rs @@ -16,6 +16,7 @@ use codex_mcp::EffectiveMcpServer; use codex_mcp::ElicitationRequestRouter; use codex_mcp::McpConnectionManager; use codex_mcp::McpRuntimeContext; +use codex_mcp::McpToolCatalogCache; use codex_mcp::ToolPluginProvenance; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; @@ -100,6 +101,7 @@ async fn hosted_plugin_runtime_ps_mcp_tool_calls_use_current_auth_manager_token( ), home.path().to_path_buf(), CodexAppsToolsCache::default(), + McpToolCatalogCache::default(), codex_mcp::codex_apps_tools_cache_key(Some(&expected_auth)), /*prefix_mcp_tool_names*/ true, ElicitationCapability::default(), diff --git a/codex-rs/core/tests/suite/mcp_tool_cache.rs b/codex-rs/core/tests/suite/mcp_tool_cache.rs new file mode 100644 index 0000000000..7a9e4126a9 --- /dev/null +++ b/codex-rs/core/tests/suite/mcp_tool_cache.rs @@ -0,0 +1,274 @@ +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context; +use codex_core::NewThread; +use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::RemoveOptions; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; +use codex_utils_path_uri::PathUri; +use core_test_support::responses; +use core_test_support::responses::ResponseMock; +use core_test_support::responses::mount_sse_once; +use core_test_support::skip_if_no_network; +use core_test_support::skip_if_wine_exec; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; + +use super::rmcp_client::remote_aware_environment_id; +use super::rmcp_client::remote_aware_stdio_server_bin; + +const SERVER_NAME: &str = "cached_rmcp"; +const NAMESPACE: &str = "mcp__cached_rmcp"; + +fn user_turn(prompt: &str) -> Op { + Op::UserInput { + items: vec![UserInput::Text { + text: prompt.to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { + approval_policy: Some(AskForApproval::Never), + permission_profile: Some(PermissionProfile::Disabled), + ..Default::default() + }, + } +} + +fn process_label(pid: &str) -> String { + format!("rmcp-test-process-{pid}") +} + +fn assert_definition(response: &ResponseMock, namespace_description: &str, tool_description: &str) { + let body = response.single_request().body_json(); + let namespace = body + .get("tools") + .and_then(Value::as_array) + .and_then(|tools| { + tools + .iter() + .find(|tool| tool.get("name").and_then(Value::as_str) == Some(NAMESPACE)) + }) + .expect("request should contain the MCP namespace"); + assert_eq!( + namespace.get("description").and_then(Value::as_str), + Some(namespace_description) + ); + assert_eq!( + responses::namespace_child_tool(&body, NAMESPACE, "echo") + .and_then(|tool| tool.get("description")) + .and_then(Value::as_str), + Some(tool_description) + ); +} + +async fn wait_for_new_pid( + fs: &dyn ExecutorFileSystem, + path: &PathUri, + previous_pid: Option<&str>, +) -> anyhow::Result { + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Ok(contents) = fs.read_file_text(path, /*sandbox*/ None).await { + let pid = contents.trim(); + if !pid.is_empty() && Some(pid) != previous_pid { + return pid.to_string(); + } + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("timed out waiting for a new MCP server process") +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow::Result<()> { + skip_if_wine_exec!( + Ok(()), + "requires a Windows test_stdio_server in the Wine-exec environment" + ); + skip_if_no_network!(Ok(())); + + let responses_server = responses::start_mock_server().await; + let command = remote_aware_stdio_server_bin()?; + let environment_id = remote_aware_environment_id(); + let fixture = test_codex() + .with_model_info_override("gpt-5.4", |model| model.supports_search_tool = false) + .with_config(move |config| { + let app_only_cwd_marker_file = config.cwd.join("cwd-app-only"); + let barrier_file = config.cwd.join("allow-initialize"); + let pid_file = config.cwd.join("mcp.pid"); + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + SERVER_NAME.to_string(), + serde_json::from_value(json!({ + "command": command, + "environment_id": environment_id, + "env": { + "MCP_TEST_APP_ONLY_CWD_MARKER_FILE": app_only_cwd_marker_file, + "MCP_TEST_INITIALIZE_BARRIER_FILE": barrier_file, + "MCP_TEST_DYNAMIC_SERVER_METADATA": "1", + "MCP_TEST_PID_FILE": pid_file, + }, + "enabled_tools": ["cwd", "echo"], + "startup_timeout_sec": 10, + })) + .expect("test MCP server configuration"), + ); + config + .mcp_servers + .set(servers) + .expect("test MCP server configuration"); + }) + .build_with_auto_env(&responses_server) + .await?; + let fs = fixture.fs(); + let app_only_cwd_marker_file = + PathUri::from_host_native_path(fixture.config.cwd.join("cwd-app-only"))?; + let barrier_file = PathUri::from_host_native_path(fixture.config.cwd.join("allow-initialize"))?; + let pid_file = PathUri::from_host_native_path(fixture.config.cwd.join("mcp.pid"))?; + + let cold_response = mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("cold"), + responses::ev_assistant_message("cold-message", "done"), + responses::ev_completed("cold"), + ]), + ) + .await; + fixture.codex.submit(user_turn("use the echo tool")).await?; + let first_pid = wait_for_new_pid(fs.as_ref(), &pid_file, /*previous_pid*/ None).await?; + fs.write_file(&barrier_file, b"ready".to_vec(), /*sandbox*/ None) + .await?; + wait_for_event(&fixture.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let first_process = process_label(&first_pid); + assert_definition( + &cold_response, + &format!("Use the tools from {first_process}."), + &format!("Echo from {first_process}."), + ); + + fs.remove( + &barrier_file, + RemoveOptions { + recursive: false, + force: false, + }, + /*sandbox*/ None, + ) + .await?; + fs.write_file( + &app_only_cwd_marker_file, + b"app-only".to_vec(), + /*sandbox*/ None, + ) + .await?; + let NewThread { + thread: second_thread, + .. + } = fixture + .thread_manager + .start_thread(fixture.config.clone()) + .await?; + let second_pid = wait_for_new_pid(fs.as_ref(), &pid_file, Some(&first_pid)).await?; + let second_process = process_label(&second_pid); + + let app_only_call_id = "cached-app-only-call"; + let cached_response = mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("cached-call"), + responses::ev_function_call_with_namespace( + "cached-call", + NAMESPACE, + "echo", + r#"{"message":"hello"}"#, + ), + responses::ev_function_call_with_namespace(app_only_call_id, NAMESPACE, "cwd", "{}"), + responses::ev_completed("cached-call"), + ]), + ) + .await; + let cached_done_response = mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("cached-done"), + responses::ev_assistant_message("cached-message", "done"), + responses::ev_completed("cached-done"), + ]), + ) + .await; + let second_for_turn = Arc::clone(&second_thread); + let cached_turn = tokio::spawn(async move { + second_for_turn + .submit(user_turn("call the echo and cwd tools")) + .await?; + let end = wait_for_event(&second_for_turn, |event| { + matches!( + event, + EventMsg::McpToolCallEnd(end) if end.call_id == "cached-call" + ) + }) + .await; + let EventMsg::McpToolCallEnd(end) = end else { + unreachable!("event predicate guarantees an MCP tool result"); + }; + let called_process = end + .result + .expect("echo call should succeed") + .structured_content + .and_then(|content| content.get("echo").cloned()) + .and_then(|echo| echo.as_str().map(ToString::to_string)) + .expect("echo result should identify its live server process"); + wait_for_event(&second_for_turn, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + anyhow::Ok(called_process) + }); + tokio::time::timeout(Duration::from_secs(2), async { + while cached_response.requests().is_empty() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("cached MCP definitions should reach inference before initialization")?; + assert_definition( + &cached_response, + &format!("Tools in the {NAMESPACE} namespace."), + &format!("Echo from {first_process}."), + ); + + fixture.codex.shutdown_and_wait().await?; + fs.write_file(&barrier_file, b"ready".to_vec(), /*sandbox*/ None) + .await?; + let expected_error = format!("MCP tool `{SERVER_NAME}/cwd` is not available to the model"); + assert_eq!(cached_turn.await??, second_process); + let output = cached_done_response + .single_request() + .function_call_output_text(app_only_call_id) + .expect("app-only tool error should be returned to the model"); + assert!( + output.contains(&expected_error), + "model-visible tool output should contain the live visibility error: {output}" + ); + + second_thread.shutdown_and_wait().await?; + responses_server.verify().await; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index f0b3fecb72..dd0a17e991 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -72,6 +72,7 @@ mod mcp_auth_elicitation; mod mcp_auth_refresh; #[cfg(unix)] mod mcp_refresh_cleanup; +mod mcp_tool_cache; mod mcp_tool_exposure; mod mcp_turn_metadata; mod model_overrides; diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 83c4e180a3..262da5d88c 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -169,7 +169,7 @@ enum McpCallEvent { const REMOTE_MCP_ENVIRONMENT: &str = "remote"; -fn remote_aware_environment_id() -> String { +pub(super) fn remote_aware_environment_id() -> String { if is_remote_test_environment() { REMOTE_MCP_ENVIRONMENT.to_string() } else { @@ -184,7 +184,7 @@ fn remote_aware_environment_id() -> String { /// would be meaningless to the process that actually launches the server. When /// the remote test environment is active, copy the binary into the executor /// container and return that in-container path instead. -fn remote_aware_stdio_server_bin() -> anyhow::Result { +pub(super) fn remote_aware_stdio_server_bin() -> anyhow::Result { let bin = stdio_server_bin()?; let Some(container_name) = test_docker_container_name() else { return Ok(bin); diff --git a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs index 722f5f7d78..3cc8e5d564 100644 --- a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs @@ -13,6 +13,7 @@ use rmcp::ServiceExt; use rmcp::handler::server::ServerHandler; use rmcp::model::CallToolRequestParams; use rmcp::model::CallToolResult; +use rmcp::model::Implementation; use rmcp::model::InitializeRequestParams; use rmcp::model::InitializeResult; use rmcp::model::JsonObject; @@ -50,6 +51,15 @@ const MEMO_URI: &str = "memo://codex/example-note"; const MEMO_CONTENT: &str = "This is a sample MCP resource served by the rmcp test server."; const SANDBOX_STATE_META_CAPABILITY: &str = "codex/sandbox-state-meta"; const SMALL_PNG_BASE64: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg=="; +const APP_ONLY_CWD_MARKER_FILE_ENV: &str = "MCP_TEST_APP_ONLY_CWD_MARKER_FILE"; +const DYNAMIC_SERVER_METADATA_ENV: &str = "MCP_TEST_DYNAMIC_SERVER_METADATA"; +const INITIALIZE_BARRIER_FILE_ENV: &str = "MCP_TEST_INITIALIZE_BARRIER_FILE"; + +fn dynamic_server_process_label() -> Option { + std::env::var_os(DYNAMIC_SERVER_METADATA_ENV) + .is_some() + .then(|| format!("rmcp-test-process-{}", std::process::id())) +} pub fn stdio() -> (tokio::io::Stdin, tokio::io::Stdout) { (tokio::io::stdin(), tokio::io::stdout()) @@ -88,7 +98,7 @@ impl TestToolServer { thread_hint_meta.insert("ui".to_string(), json!({ "visibility": [] })); thread_hint_tool.meta = Some(thread_hint_meta); - let tools = vec![ + let mut tools = vec![ Self::echo_tool(), Self::echo_dash_tool(), thread_hint_tool, @@ -100,6 +110,11 @@ impl TestToolServer { Self::image_scenario_tool(), sandbox_meta_tool, ]; + if let Some(process_label) = dynamic_server_process_label() + && let Some(echo) = tools.iter_mut().find(|tool| tool.name == "echo") + { + echo.description = Some(Cow::Owned(format!("Echo from {process_label}."))); + } let resources = vec![Self::memo_resource()]; let resource_templates = vec![Self::memo_template()]; Self { @@ -445,6 +460,11 @@ impl ServerHandler for TestToolServer { request: InitializeRequestParams, context: rmcp::service::RequestContext, ) -> Result { + if let Ok(barrier_file) = std::env::var(INITIALIZE_BARRIER_FILE_ENV) { + while !std::path::Path::new(&barrier_file).is_file() { + sleep(Duration::from_millis(10)).await; + } + } self.supports_openai_form_elicitation.store( request .capabilities @@ -468,8 +488,18 @@ impl ServerHandler for TestToolServer { JsonObject::new(), )])); - ServerInfo::new(capabilities) - .with_instructions("Use these tools to exercise the rmcp test server.") + let server_info = ServerInfo::new(capabilities); + match dynamic_server_process_label() { + Some(process_label) => server_info + .with_server_info( + Implementation::new("codex-rmcp-test-server", env!("CARGO_PKG_VERSION")) + .with_title(process_label.clone()), + ) + .with_instructions(format!("Use the tools from {process_label}.")), + None => { + server_info.with_instructions("Use these tools to exercise the rmcp test server.") + } + } } fn list_tools( @@ -479,8 +509,17 @@ impl ServerHandler for TestToolServer { ) -> impl std::future::Future> + Send + '_ { let tools = self.tools.clone(); async move { + let mut tools = (*tools).clone(); + if let Some(marker_file) = std::env::var_os(APP_ONLY_CWD_MARKER_FILE_ENV) + && std::path::Path::new(&marker_file).is_file() + && let Some(cwd) = tools.iter_mut().find(|tool| tool.name == "cwd") + { + cwd.meta + .get_or_insert_with(Meta::new) + .insert("ui".to_string(), json!({ "visibility": ["app"] })); + } Ok(ListToolsResult { - tools: (*tools).clone(), + tools, next_cursor: None, meta: None, }) @@ -588,8 +627,10 @@ impl ServerHandler for TestToolServer { let env_snapshot: HashMap = std::env::vars().collect(); let env_name = args.env_var.as_deref().unwrap_or("MCP_TEST_VALUE"); + let echo = dynamic_server_process_label() + .unwrap_or_else(|| format!("ECHOING: {}", args.message)); let structured_content = json!({ - "echo": format!("ECHOING: {}", args.message), + "echo": echo, "env": env_snapshot.get(env_name), });