From 6d7db13219a67708c68ac90f7f93aef3cd631751 Mon Sep 17 00:00:00 2001 From: fbauer Date: Mon, 6 Jul 2026 14:46:07 -0700 Subject: [PATCH] mcp resource read: add edgecases --- codex-rs/app-server/src/lib.rs | 1 + .../app-server/src/mcp_resource_origin.rs | 102 ++++ .../src/mcp_resource_origin_tests.rs | 79 +++ .../request_processors/mcp_resource_read.rs | 130 ++--- .../mcp_resource_read_tests.rs | 174 +++++++ .../request_processors/thread_processor.rs | 10 + codex-rs/app-server/src/thread_state.rs | 19 + .../tests/suite/v2/mcp_resource_origin.rs | 477 ++++++++++++------ codex-rs/codex-mcp/src/connection_manager.rs | 179 +++++-- .../codex-mcp/src/connection_manager_tests.rs | 69 ++- codex-rs/codex-mcp/src/lib.rs | 1 + codex-rs/codex-mcp/src/rmcp_client.rs | 35 +- 12 files changed, 981 insertions(+), 295 deletions(-) create mode 100644 codex-rs/app-server/src/mcp_resource_origin.rs create mode 100644 codex-rs/app-server/src/mcp_resource_origin_tests.rs create mode 100644 codex-rs/app-server/src/request_processors/mcp_resource_read_tests.rs diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 0edb57d8a2..784d57dfb9 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -102,6 +102,7 @@ mod fuzzy_file_search; mod image_url; pub mod in_process; mod mcp_refresh; +mod mcp_resource_origin; mod message_processor; mod models; mod models_refresh_worker; diff --git a/codex-rs/app-server/src/mcp_resource_origin.rs b/codex-rs/app-server/src/mcp_resource_origin.rs new file mode 100644 index 0000000000..2438496b82 --- /dev/null +++ b/codex-rs/app-server/src/mcp_resource_origin.rs @@ -0,0 +1,102 @@ +use codex_app_server_protocol::ThreadItem; +use std::collections::VecDeque; + +const MAX_MCP_RESOURCE_ORIGINS: usize = 1024; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct McpResourceOrigin { + pub(crate) server: String, + pub(crate) tool: String, + pub(crate) connector_id: Option, + pub(crate) link_id: Option, + pub(crate) resource_uri: Option, +} + +#[derive(Default)] +pub(crate) struct McpResourceOriginIndex(VecDeque<(String, McpResourceOrigin)>); + +impl McpResourceOrigin { + pub(crate) fn find(items: &[ThreadItem], call_id: &str) -> Option { + items + .iter() + .find(|item| matches!(item, ThreadItem::McpToolCall { id, .. } if id == call_id)) + .and_then(|item| Self::from_item(item).map(|(_, origin)| origin)) + } + + fn from_item(item: &ThreadItem) -> Option<(&str, Self)> { + let ThreadItem::McpToolCall { + id, + server, + tool, + app_context, + .. + } = item + else { + return None; + }; + let (connector_id, link_id, resource_uri) = app_context + .as_ref() + .map(|context| { + ( + Some(context.connector_id.clone()), + context.link_id.clone(), + context.resource_uri.clone(), + ) + }) + .unwrap_or_default(); + Some(( + id, + Self { + server: server.clone(), + tool: tool.clone(), + connector_id, + link_id, + resource_uri, + }, + )) + } +} + +impl McpResourceOriginIndex { + pub(crate) fn get(&self, call_id: &str) -> Option { + self.0.iter().find_map(|(existing_call_id, origin)| { + (existing_call_id == call_id).then(|| origin.clone()) + }) + } + + pub(crate) fn insert_if_absent( + &mut self, + call_id: &str, + origin: McpResourceOrigin, + ) -> McpResourceOrigin { + if let Some(origin) = self.get(call_id) { + return origin; + } + self.insert(call_id.to_string(), origin.clone()); + origin + } + + pub(crate) fn seed<'a>(&mut self, items: impl IntoIterator) { + for (call_id, origin) in items.into_iter().filter_map(McpResourceOrigin::from_item) { + self.insert(call_id.to_string(), origin); + } + } + + fn insert(&mut self, call_id: String, origin: McpResourceOrigin) { + if let Some(index) = self + .0 + .iter() + .position(|(existing_call_id, _)| existing_call_id == &call_id) + { + self.0.remove(index); + } + self.0.push_back((call_id, origin)); + if self.0.len() > MAX_MCP_RESOURCE_ORIGINS { + self.0.pop_front(); + } + } +} + +#[cfg(test)] +#[path = "mcp_resource_origin_tests.rs"] +mod tests; diff --git a/codex-rs/app-server/src/mcp_resource_origin_tests.rs b/codex-rs/app-server/src/mcp_resource_origin_tests.rs new file mode 100644 index 0000000000..5d943475a3 --- /dev/null +++ b/codex-rs/app-server/src/mcp_resource_origin_tests.rs @@ -0,0 +1,79 @@ +use super::*; +use codex_app_server_protocol::McpToolCallAppContext; +use codex_app_server_protocol::McpToolCallStatus; +use pretty_assertions::assert_eq; + +#[test] +fn bounds_entries_and_keeps_updated_origins() { + let mut index = McpResourceOriginIndex::default(); + insert(&mut index, "call-0", "old-connector"); + for index_value in 1..MAX_MCP_RESOURCE_ORIGINS { + insert(&mut index, format!("call-{index_value}"), "connector"); + } + + insert(&mut index, "call-0", "new-connector"); + insert(&mut index, "overflow", "connector"); + + assert_eq!( + index.get("call-0"), + Some(origin_with_app_context("new-connector")) + ); + assert_eq!(index.get("call-1"), None); +} + +#[test] +fn history_does_not_replace_live_origin() { + let mut index = McpResourceOriginIndex::default(); + insert(&mut index, "call-1", "live-connector"); + + let resolved = index.insert_if_absent("call-1", origin_with_app_context("history-connector")); + + assert_eq!(resolved, origin_with_app_context("live-connector")); + assert_eq!(index.get("call-1"), Some(resolved)); +} + +#[test] +fn seeds_origins_from_materialized_items() { + let mut index = McpResourceOriginIndex::default(); + let items = [ThreadItem::McpToolCall { + id: "call-1".to_string(), + server: "plugin-runtime".to_string(), + tool: "search".to_string(), + status: McpToolCallStatus::Completed, + arguments: serde_json::Value::Null, + app_context: Some(McpToolCallAppContext { + connector_id: "connector".to_string(), + link_id: Some("link".to_string()), + resource_uri: Some("ui://shared/widget.html".to_string()), + app_name: Some("App".to_string()), + template_id: None, + action_name: Some("Search".to_string()), + }), + mcp_app_resource_uri: Some("ui://shared/widget.html".to_string()), + plugin_id: None, + result: None, + error: None, + duration_ms: Some(1), + }]; + + index.seed(&items); + + assert_eq!( + index.get("call-1"), + Some(origin_with_app_context("connector")) + ); +} + +fn insert(index: &mut McpResourceOriginIndex, call_id: impl Into, connector_id: &str) { + index.insert(call_id.into(), origin_with_app_context(connector_id)); +} + +fn origin_with_app_context(connector_id: &str) -> McpResourceOrigin { + McpResourceOrigin { + server: "plugin-runtime".to_string(), + tool: "search".to_string(), + connector_id: Some(connector_id.to_string()), + link_id: Some("link".to_string()), + resource_uri: Some("ui://shared/widget.html".to_string()), + } +} diff --git a/codex-rs/app-server/src/request_processors/mcp_resource_read.rs b/codex-rs/app-server/src/request_processors/mcp_resource_read.rs index 7f6c407629..d330433d95 100644 --- a/codex-rs/app-server/src/request_processors/mcp_resource_read.rs +++ b/codex-rs/app-server/src/request_processors/mcp_resource_read.rs @@ -1,5 +1,6 @@ use super::mcp_processor::with_mcp_tool_call_thread_id_meta; use super::*; +use crate::mcp_resource_origin::McpResourceOrigin; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; use codex_mcp::ToolInfo; @@ -24,12 +25,13 @@ pub(super) async fn read_thread_mcp_resource( return thread.read_mcp_resource(server, uri).await; } - let server_info = thread - .mcp_server_info(server) - .await - .ok_or_else(|| anyhow::anyhow!("unable to identify codex_apps MCP server"))?; - match server_info.name.as_str() { - LEGACY_CODEX_APPS_MCP_SERVER_NAME => return thread.read_mcp_resource(server, uri).await, + let runtime = thread.current_mcp_runtime().await; + let connection = runtime.manager().server_connection(server).await?; + match connection.server_info().name.as_str() { + LEGACY_CODEX_APPS_MCP_SERVER_NAME => { + let result = connection.read_resource(uri).await?; + return Ok(serde_json::to_value(result)?); + } PLUGIN_RUNTIME_MCP_SERVER_NAME => {} implementation => { anyhow::bail!("unsupported codex_apps MCP server implementation `{implementation}`") @@ -39,6 +41,9 @@ pub(super) async fn read_thread_mcp_resource( let origin = resolve_mcp_resource_origin(thread_state_manager, thread_id, thread, origin_call_id) .await?; + if origin.connector_id.is_none() { + anyhow::bail!("originating MCP tool call `{origin_call_id}` has no app context"); + } if origin.server != server { anyhow::bail!( "originating MCP tool call server `{}` does not match resource server `{server}`", @@ -52,20 +57,16 @@ pub(super) async fn read_thread_mcp_resource( ); } - let tool_info = thread - .mcp_tool_info(server, &origin.tool) - .await - .ok_or_else(|| { - anyhow::anyhow!( - "originating MCP tool `{}` is not available on server `{server}`", - origin.tool - ) - })?; - let meta = build_plugin_runtime_fetch_resource_meta(&origin, uri, &tool_info)?; + let tool_info = connection.tool_info(&origin.tool).ok_or_else(|| { + anyhow::anyhow!( + "originating MCP tool `{}` is not available on server `{server}`", + origin.tool + ) + })?; + let meta = build_plugin_runtime_fetch_resource_meta(&origin, uri, tool_info)?; let meta = with_mcp_tool_call_thread_id_meta(Some(meta), thread_id_string); - let result = thread - .call_mcp_tool( - server, + let result = connection + .call_tool( FETCH_RESOURCE_TOOL_NAME, Some(serde_json::json!({ "uri": uri })), meta, @@ -80,70 +81,39 @@ async fn resolve_mcp_resource_origin( thread: &Arc, origin_call_id: &str, ) -> anyhow::Result { - let active_turn = { - let thread_state = thread_state_manager.thread_state(thread_id).await; - thread_state.lock().await.active_turn_snapshot() - }; - if let Some(origin) = active_turn - .as_ref() - .and_then(|turn| find_mcp_resource_origin(&turn.items, origin_call_id)) + let thread_state = thread_state_manager.thread_state(thread_id).await; + if let Some(origin) = thread_state + .lock() + .await + .mcp_resource_origin(origin_call_id) { - return origin; + return Ok(origin); } - let history = thread.load_history(/*include_archived*/ true).await?; - let turns = build_api_turns_from_rollout_items(&history.items); - turns + let history = thread.load_history(/*include_archived*/ true).await; + if let Some(origin) = thread_state + .lock() + .await + .mcp_resource_origin(origin_call_id) + { + return Ok(origin); + } + let history = history?; + let history_origin = build_api_turns_from_rollout_items(&history.items) .iter() .rev() - .find_map(|turn| find_mcp_resource_origin(&turn.items, origin_call_id)) - .unwrap_or_else(|| { - Err(anyhow::anyhow!( - "originating MCP tool call `{origin_call_id}` was not found" - )) + .find_map(|turn| McpResourceOrigin::find(&turn.items, origin_call_id)); + let mut thread_state = thread_state.lock().await; + history_origin + .map(|origin| { + thread_state + .mcp_resource_origins + .insert_if_absent(origin_call_id, origin) }) -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct McpResourceOrigin { - server: String, - tool: String, - connector_id: String, - link_id: Option, - resource_uri: Option, -} - -fn find_mcp_resource_origin( - items: &[ThreadItem], - origin_call_id: &str, -) -> Option> { - items.iter().find_map(|item| { - let ThreadItem::McpToolCall { - id, - server, - tool, - app_context, - .. - } = item - else { - return None; - }; - if id != origin_call_id { - return None; - } - Some(match app_context { - Some(app_context) => Ok(McpResourceOrigin { - server: server.clone(), - tool: tool.clone(), - connector_id: app_context.connector_id.clone(), - link_id: app_context.link_id.clone(), - resource_uri: app_context.resource_uri.clone(), - }), - None => Err(anyhow::anyhow!( - "originating MCP tool call `{origin_call_id}` has no app context" - )), + .or_else(|| thread_state.mcp_resource_origin(origin_call_id)) + .ok_or_else(|| { + anyhow::anyhow!("originating MCP tool call `{origin_call_id}` was not found") }) - }) } fn build_plugin_runtime_fetch_resource_meta( @@ -177,9 +147,9 @@ fn build_plugin_runtime_fetch_resource_meta( "originating MCP tool resource URI `{resource_uri}` does not match requested URI `{uri}`" ); } - if connector_id != origin.connector_id { + if Some(connector_id) != origin.connector_id.as_deref() { anyhow::bail!( - "originating MCP tool connector `{connector_id}` does not match app context connector `{}`", + "originating MCP tool connector `{connector_id}` does not match app context connector {:?}", origin.connector_id ); } @@ -298,3 +268,7 @@ fn plugin_runtime_fetch_resource_response( let response = serde_json::from_value::(structured_content)?; Ok(serde_json::to_value(response)?) } + +#[cfg(test)] +#[path = "mcp_resource_read_tests.rs"] +mod tests; diff --git a/codex-rs/app-server/src/request_processors/mcp_resource_read_tests.rs b/codex-rs/app-server/src/request_processors/mcp_resource_read_tests.rs new file mode 100644 index 0000000000..934d99fba2 --- /dev/null +++ b/codex-rs/app-server/src/request_processors/mcp_resource_read_tests.rs @@ -0,0 +1,174 @@ +use super::*; +use pretty_assertions::assert_eq; +use rmcp::model::JsonObject; +use rmcp::model::Meta; +use rmcp::model::Tool; +use serde_json::json; +use std::borrow::Cow; +use std::sync::Arc; + +const CONNECTOR_ID: &str = "catalog-a"; +const LINK_ID: &str = "link-catalog-a"; +const RESOURCE_URI: &str = "ui://widget/product-search.html"; +const TOOL_NAME: &str = "catalog_a_search"; + +#[test] +fn rejects_connector_link_uri_and_route_identity_mismatches() { + let mut mismatched_connector = origin(Some(LINK_ID)); + mismatched_connector.connector_id = Some("catalog-b".to_string()); + assert_error_contains( + build_plugin_runtime_fetch_resource_meta(&mismatched_connector, RESOURCE_URI, &tool_info()), + "does not match app context connector", + ); + + let mismatched_link = origin(Some("link-catalog-b")); + assert_error_contains( + build_plugin_runtime_fetch_resource_meta(&mismatched_link, RESOURCE_URI, &tool_info()), + "link does not match app context link", + ); + + assert_error_contains( + build_plugin_runtime_fetch_resource_meta( + &origin(Some(LINK_ID)), + "ui://widget/other.html", + &tool_info(), + ), + "does not match requested URI", + ); + + let mut mismatched_route = tool_info(); + set_codex_apps_meta( + &mut mismatched_route, + json!({ + "resource_uri": "/catalog-b/link-catalog-b/catalog_b_search", + "contains_mcp_source": true, + }), + ); + assert_error_contains( + build_plugin_runtime_fetch_resource_meta( + &origin(Some(LINK_ID)), + RESOURCE_URI, + &mismatched_route, + ), + "resource route does not match connector and link", + ); +} + +#[test] +fn accepts_missing_origin_link_only_for_synthetic_tools() { + let mut synthetic_tool = tool_info(); + set_codex_apps_meta( + &mut synthetic_tool, + json!({ + "resource_uri": format!("/{CONNECTOR_ID}/{LINK_ID}/{TOOL_NAME}"), + "contains_mcp_source": true, + "synthetic_link": true, + }), + ); + + assert_eq!( + build_plugin_runtime_fetch_resource_meta( + &origin(/*link_id*/ None), + RESOURCE_URI, + &synthetic_tool, + ) + .expect("synthetic tool accepts missing origin link") + .pointer("/_codex_apps/synthetic_link"), + Some(&json!(true)) + ); + assert_error_contains( + build_plugin_runtime_fetch_resource_meta( + &origin(/*link_id*/ None), + RESOURCE_URI, + &tool_info(), + ), + "link does not match app context link", + ); +} + +#[test] +fn rejects_missing_or_malformed_tool_metadata() { + let mut missing_metadata = tool_info(); + missing_metadata.tool.meta = None; + assert_error_contains( + build_plugin_runtime_fetch_resource_meta( + &origin(Some(LINK_ID)), + RESOURCE_URI, + &missing_metadata, + ), + "missing metadata", + ); + + let mut malformed_route = tool_info(); + set_codex_apps_meta( + &mut malformed_route, + json!({ "resource_uri": "/catalog-a/link-catalog-a" }), + ); + assert_error_contains( + build_plugin_runtime_fetch_resource_meta( + &origin(Some(LINK_ID)), + RESOURCE_URI, + &malformed_route, + ), + "invalid resource route", + ); +} + +fn origin(link_id: Option<&str>) -> McpResourceOrigin { + McpResourceOrigin { + server: CODEX_APPS_MCP_SERVER_NAME.to_string(), + tool: TOOL_NAME.to_string(), + connector_id: Some(CONNECTOR_ID.to_string()), + link_id: link_id.map(str::to_string), + resource_uri: Some(RESOURCE_URI.to_string()), + } +} + +fn set_codex_apps_meta(tool_info: &mut ToolInfo, value: serde_json::Value) { + tool_info + .tool + .meta + .as_mut() + .expect("tool metadata") + .0 + .insert(MCP_TOOL_CODEX_APPS_META_KEY.to_string(), value); +} + +fn tool_info() -> ToolInfo { + let mut tool = Tool::new( + Cow::Borrowed(TOOL_NAME), + Cow::Borrowed("Search catalog A."), + Arc::new(JsonObject::new()), + ); + tool.meta = Some(Meta(serde_json::Map::from_iter([ + ("link_id".to_string(), json!(LINK_ID)), + ("openai/outputTemplate".to_string(), json!(RESOURCE_URI)), + ( + MCP_TOOL_CODEX_APPS_META_KEY.to_string(), + json!({ + "resource_uri": format!("/{CONNECTOR_ID}/{LINK_ID}/{TOOL_NAME}"), + "contains_mcp_source": true, + }), + ), + ]))); + ToolInfo { + server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, + callable_name: "_product_search".to_string(), + callable_namespace: "mcp__codex_apps__catalog_a".to_string(), + namespace_description: Some("Search catalog A.".to_string()), + tool, + connector_id: Some(CONNECTOR_ID.to_string()), + connector_name: Some("Catalog A".to_string()), + plugin_display_names: Vec::new(), + } +} + +fn assert_error_contains(result: anyhow::Result, expected: &str) { + let error = result.err().expect("expected validation failure"); + assert!( + error.to_string().contains(expected), + "expected error containing {expected:?}, got {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 b78c53a383..6dfaa4c593 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -3537,6 +3537,16 @@ impl ThreadRequestProcessor { let instruction_sources = forked_thread.legacy_instruction_sources().await; + if session_configured.rollout_path.is_none() { + let turns = build_api_turns_from_rollout_items(&history_items); + let thread_state = self.thread_state_manager.thread_state(thread_id).await; + thread_state + .lock() + .await + .mcp_resource_origins + .seed(turns.iter().flat_map(|turn| &turn.items)); + } + // Auto-attach a conversation listener when forking a thread. log_listener_attach_result( self.ensure_conversation_listener( diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index fec786373c..6a325cfa2f 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -1,3 +1,5 @@ +use crate::mcp_resource_origin::McpResourceOrigin; +use crate::mcp_resource_origin::McpResourceOriginIndex; use crate::outgoing_message::ConnectionId; use crate::outgoing_message::ConnectionRequestId; use codex_app_server_protocol::RequestId; @@ -86,6 +88,7 @@ pub(crate) struct ThreadState { last_thread_settings: Option, listener_command_tx: Option>, current_turn_history: ThreadHistoryBuilder, + pub(crate) mcp_resource_origins: McpResourceOriginIndex, listener_thread: Option>, watch_registration: WatchRegistration, } @@ -141,10 +144,26 @@ impl ThreadState { self.current_turn_history.active_turn_snapshot() } + pub(crate) fn mcp_resource_origin(&self, call_id: &str) -> Option { + self.active_turn_snapshot() + .as_ref() + .and_then(|turn| McpResourceOrigin::find(&turn.items, call_id)) + .or_else(|| self.mcp_resource_origins.get(call_id)) + } + pub(crate) fn track_current_turn_event(&mut self, event_turn_id: &str, event: &EventMsg) { if let EventMsg::TurnStarted(payload) = event { self.turn_summary.started_at = payload.started_at; } + match event { + EventMsg::ThreadRolledBack(_) => self.mcp_resource_origins = Default::default(), + EventMsg::TurnAborted(_) | EventMsg::TurnComplete(_) => { + if let Some(turn) = self.active_turn_snapshot() { + self.mcp_resource_origins.seed(&turn.items); + } + } + _ => {} + } self.current_turn_history.handle_event(event); if matches!(event, EventMsg::TurnAborted(_) | EventMsg::TurnComplete(_)) && !self.current_turn_history.has_active_turn() diff --git a/codex-rs/app-server/tests/suite/v2/mcp_resource_origin.rs b/codex-rs/app-server/tests/suite/v2/mcp_resource_origin.rs index 0c117a86ce..b270b76278 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_resource_origin.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_resource_origin.rs @@ -9,13 +9,13 @@ use app_test_support::TestAppServer; use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use axum::Router; -use codex_app_server_protocol::ItemCompletedNotification; +use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::McpResourceContent; use codex_app_server_protocol::McpResourceReadParams; use codex_app_server_protocol::McpResourceReadResponse; use codex_app_server_protocol::RequestId; -use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; @@ -26,7 +26,6 @@ use pretty_assertions::assert_eq; use rmcp::handler::server::ServerHandler; use rmcp::model::CallToolRequestParams; use rmcp::model::CallToolResult; -use rmcp::model::Content; use rmcp::model::Implementation; use rmcp::model::JsonObject; use rmcp::model::ListToolsResult; @@ -49,36 +48,144 @@ use tokio::task::JoinHandle; use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); -const CONNECTOR_ID: &str = "walmart"; -const CONNECTOR_NAME: &str = "Walmart"; -const LINK_ID: &str = "link_walmart"; -const ORIGIN_CALL_ID: &str = "walmart-product-search-call"; -const PRODUCT_SEARCH_TOOL: &str = "walmart_product_search"; -const PRODUCT_SEARCH_MODEL_NAMESPACE: &str = "mcp__codex_apps__walmart"; +const CONNECTOR_DESCRIPTION: &str = "Search catalog."; const PRODUCT_SEARCH_MODEL_TOOL: &str = "_product_search"; const RESOURCE_URI: &str = "ui://widget/product-search.html"; +const WRONG_RESOURCE_URI: &str = "ui://widget/wrong-product-search.html"; + +#[derive(Clone, Copy)] +struct AppFixture { + connector_id: &'static str, + link_id: &'static str, + origin_call_id: &'static str, + product_search_tool: &'static str, + model_namespace: &'static str, + html: &'static str, +} + +const CATALOG_A: AppFixture = AppFixture { + connector_id: "catalog-a", + link_id: "link-catalog-a", + origin_call_id: "catalog-a-search-call", + product_search_tool: "catalog_a_product_search", + model_namespace: "mcp__codex_apps__catalog_a", + html: "Catalog A", +}; + +const CATALOG_B: AppFixture = AppFixture { + connector_id: "catalog-b", + link_id: "link-catalog-b", + origin_call_id: "catalog-b-search-call", + product_search_tool: "catalog_b_product_search", + model_namespace: "mcp__codex_apps__catalog_b", + html: "Catalog B", +}; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn mcp_resource_read_with_origin_calls_plugin_runtime_fetch_resource() -> Result<()> { +async fn reopened_mcp_resource_reads_are_scoped_to_the_originating_app() -> Result<()> { + let mut scenario = start_origin_test_scenario(/*ephemeral*/ false).await?; + + let catalog_a_response = read_mcp_resource( + &mut scenario.mcp, + &scenario.thread_id, + RESOURCE_URI, + CATALOG_A.origin_call_id, + ) + .await?; + assert_eq!( + catalog_a_response, + expected_resource_response(CATALOG_A), + "catalog A's call must render catalog A's HTML" + ); + let catalog_b_response = read_mcp_resource( + &mut scenario.mcp, + &scenario.thread_id, + RESOURCE_URI, + CATALOG_B.origin_call_id, + ) + .await?; + assert_eq!( + catalog_b_response, + expected_resource_response(CATALOG_B), + "catalog B's call must render catalog B's HTML" + ); + + assert_fetch_resource_calls( + &scenario.recorded_calls, + &[CATALOG_A, CATALOG_B], + &scenario.thread_id, + ); + + scenario.stop().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn completed_ephemeral_origin_is_retained_and_identity_mismatches_fail_closed() -> Result<()> +{ + let mut scenario = start_origin_test_scenario(/*ephemeral*/ true).await?; + + let response = read_mcp_resource( + &mut scenario.mcp, + &scenario.thread_id, + RESOURCE_URI, + CATALOG_A.origin_call_id, + ) + .await?; + assert_eq!(response, expected_resource_response(CATALOG_A)); + + let wrong_uri_error = read_mcp_resource_error( + &mut scenario.mcp, + &scenario.thread_id, + WRONG_RESOURCE_URI, + CATALOG_A.origin_call_id, + ) + .await?; + assert!( + wrong_uri_error.error.message.contains("does not match"), + "unexpected wrong-URI error: {wrong_uri_error:?}" + ); + + read_mcp_resource_error( + &mut scenario.mcp, + &scenario.thread_id, + RESOURCE_URI, + "unknown-product-search-call", + ) + .await?; + assert_fetch_resource_calls(&scenario.recorded_calls, &[CATALOG_A], &scenario.thread_id); + + scenario.stop().await; + Ok(()) +} + +struct OriginTestScenario { + _codex_home: TempDir, + mcp: TestAppServer, + thread_id: String, + recorded_calls: Arc>>, + apps_server_handle: JoinHandle<()>, +} + +impl OriginTestScenario { + async fn stop(self) { + self.apps_server_handle.abort(); + let _ = self.apps_server_handle.await; + } +} + +async fn start_origin_test_scenario(ephemeral: bool) -> Result { let responses_server = responses::start_mock_server().await; let (apps_server_url, recorded_calls, apps_server_handle) = start_plugin_runtime_mcp_server().await?; - let (_codex_home, mut mcp) = + let (codex_home, mut mcp) = start_test_app_server(&apps_server_url, &responses_server.uri()).await?; let response_mock = responses::mount_sse_sequence( &responses_server, vec![ - responses::sse(vec![ - responses::ev_response_created("resp-product-search"), - responses::ev_function_call_with_namespace( - ORIGIN_CALL_ID, - PRODUCT_SEARCH_MODEL_NAMESPACE, - PRODUCT_SEARCH_MODEL_TOOL, - &json!({ "query": "bed lamps" }).to_string(), - ), - responses::ev_completed("resp-product-search"), - ]), + product_search_response("resp-catalog-a-search", CATALOG_A), + product_search_response("resp-catalog-b-search", CATALOG_B), responses::sse(vec![ responses::ev_response_created("resp-done"), responses::ev_assistant_message("msg-done", "Done"), @@ -91,6 +198,7 @@ async fn mcp_resource_read_with_origin_calls_plugin_runtime_fetch_resource() -> let thread_start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { model: Some("mock-model".to_string()), + ephemeral: Some(ephemeral), ..Default::default() }) .await?; @@ -105,7 +213,7 @@ async fn mcp_resource_read_with_origin_calls_plugin_runtime_fetch_resource() -> .send_turn_start_request(TurnStartParams { thread_id: thread.id.clone(), input: vec![UserInput::Text { - text: "Find bed lamps".to_string(), + text: "Compare product results".to_string(), text_elements: Vec::new(), }], ..Default::default() @@ -117,81 +225,141 @@ async fn mcp_resource_read_with_origin_calls_plugin_runtime_fetch_resource() -> ) .await??; - let completed = wait_for_mcp_tool_call_completed(&mut mcp, ORIGIN_CALL_ID).await?; - let ThreadItem::McpToolCall { - app_context: Some(app_context), - .. - } = completed.item - else { - anyhow::bail!("originating MCP tool call should include app context"); - }; - assert_eq!(app_context.connector_id, CONNECTOR_ID); - assert_eq!(app_context.link_id.as_deref(), Some(LINK_ID)); - assert_eq!(app_context.resource_uri.as_deref(), Some(RESOURCE_URI)); - timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), ) .await??; - let read_request_id = mcp - .send_mcp_resource_read_request(McpResourceReadParams { - thread_id: Some(thread.id.clone()), - server: "codex_apps".to_string(), - uri: RESOURCE_URI.to_string(), - origin_call_id: Some(ORIGIN_CALL_ID.to_string()), - }) - .await?; - let read_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_request_id)), - ) - .await??; - - assert_eq!( - to_response::(read_response)?, - McpResourceReadResponse { - contents: vec![McpResourceContent::Text { - uri: RESOURCE_URI.to_string(), - mime_type: Some("text/html".to_string()), - text: "Walmart".to_string(), - meta: Some(json!({ "app": CONNECTOR_ID })), - }], - } - ); - - { - let calls = recorded_calls - .lock() - .expect("recorded tool calls lock poisoned"); - let fetch_resource_call = calls - .iter() - .find(|call| call["name"] == "fetch_resource") - .expect("resource read should call fetch_resource"); - assert_eq!( - fetch_resource_call["arguments"], - json!({ "uri": RESOURCE_URI }) - ); - assert_eq!( - fetch_resource_call.pointer("/meta/_codex_apps"), - Some(&json!({ - "resource_uri": format!("/{CONNECTOR_ID}/{LINK_ID}/fetch_resource"), - "contains_mcp_source": true, - })) - ); - assert_eq!( - fetch_resource_call - .pointer("/meta/x-codex-turn-metadata/mcp_request_meta/selected_connector_ids"), - Some(&json!([CONNECTOR_ID])) - ); - assert_eq!(fetch_resource_call["meta"]["threadId"], thread.id); + if !ephemeral { + drop(mcp); + mcp = TestAppServer::new_with_auto_env(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let resume_id = mcp + .send_thread_resume_request(ThreadResumeParams { + thread_id: thread.id.clone(), + ..Default::default() + }) + .await?; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), + ) + .await??; } - assert_eq!(response_mock.requests().len(), 2); - apps_server_handle.abort(); - let _ = apps_server_handle.await; - Ok(()) + assert_eq!(response_mock.requests().len(), 3); + Ok(OriginTestScenario { + _codex_home: codex_home, + mcp, + thread_id: thread.id, + recorded_calls, + apps_server_handle, + }) +} + +fn product_search_response(response_id: &str, app: AppFixture) -> String { + responses::sse(vec![ + responses::ev_response_created(response_id), + responses::ev_function_call_with_namespace( + app.origin_call_id, + app.model_namespace, + PRODUCT_SEARCH_MODEL_TOOL, + &json!({ "query": "bed lamps" }).to_string(), + ), + responses::ev_completed(response_id), + ]) +} + +async fn read_mcp_resource( + mcp: &mut TestAppServer, + thread_id: &str, + uri: &str, + origin_call_id: &str, +) -> Result { + let request_id = send_read_request(mcp, thread_id, uri, origin_call_id).await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + to_response(response) +} + +async fn read_mcp_resource_error( + mcp: &mut TestAppServer, + thread_id: &str, + uri: &str, + origin_call_id: &str, +) -> Result { + let request_id = send_read_request(mcp, thread_id, uri, origin_call_id).await?; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await? +} + +async fn send_read_request( + mcp: &mut TestAppServer, + thread_id: &str, + uri: &str, + origin_call_id: &str, +) -> Result { + mcp.send_mcp_resource_read_request(McpResourceReadParams { + thread_id: Some(thread_id.to_string()), + server: "codex_apps".to_string(), + uri: uri.to_string(), + origin_call_id: Some(origin_call_id.to_string()), + }) + .await +} + +fn expected_resource_response(app: AppFixture) -> McpResourceReadResponse { + McpResourceReadResponse { + contents: vec![McpResourceContent::Text { + uri: RESOURCE_URI.to_string(), + mime_type: Some("text/html".to_string()), + text: app.html.to_string(), + meta: Some(json!({ "app": app.connector_id })), + }], + } +} + +fn assert_fetch_resource_calls( + recorded_calls: &Arc>>, + apps: &[AppFixture], + thread_id: &str, +) { + let actual = recorded_calls + .lock() + .expect("recorded tool calls lock poisoned") + .iter() + .filter(|call| call["name"] == "fetch_resource") + .cloned() + .collect::>(); + let expected = apps + .iter() + .map(|app| { + json!({ + "name": "fetch_resource", + "arguments": { "uri": RESOURCE_URI }, + "meta": { + "_codex_apps": { + "resource_uri": format!("/{}/{}/fetch_resource", app.connector_id, app.link_id), + "contains_mcp_source": true, + }, + "connector_name": app.connector_id, + "connector_description": CONNECTOR_DESCRIPTION, + "x-codex-turn-metadata": { + "mcp_request_meta": { "selected_connector_ids": [app.connector_id] }, + }, + "threadId": thread_id, + }, + }) + }) + .collect::>(); + assert_eq!(actual, expected); } async fn start_test_app_server( @@ -278,41 +446,11 @@ impl ServerHandler for PluginRuntimeMcpServer { _request: Option, _context: RequestContext, ) -> Result { - let input_schema: JsonObject = serde_json::from_value(json!({ - "type": "object", - "properties": { "query": { "type": "string" } }, - "required": ["query"], - "additionalProperties": false - })) - .map_err(|error| rmcp::ErrorData::internal_error(error.to_string(), None))?; - let mut tool = Tool::new( - Cow::Borrowed(PRODUCT_SEARCH_TOOL), - Cow::Borrowed("Search Walmart products."), - Arc::new(input_schema), - ); - tool.annotations = Some(ToolAnnotations::new().read_only(true)); - tool.meta = Some(Meta(serde_json::Map::from_iter([ - ("connector_id".to_string(), json!(CONNECTOR_ID)), - ("connector_name".to_string(), json!(CONNECTOR_NAME)), - ( - "connector_description".to_string(), - json!("Search Walmart products."), - ), - ("link_id".to_string(), json!(LINK_ID)), - ("openai/outputTemplate".to_string(), json!(RESOURCE_URI)), - ( - "_codex_apps".to_string(), - json!({ - "resource_uri": format!( - "/{CONNECTOR_ID}/{LINK_ID}/{PRODUCT_SEARCH_TOOL}" - ), - "contains_mcp_source": true, - }), - ), - ]))); - Ok(ListToolsResult { - tools: vec![tool], + tools: [CATALOG_A, CATALOG_B] + .into_iter() + .map(product_search_tool) + .collect(), next_cursor: None, meta: None, }) @@ -325,6 +463,13 @@ impl ServerHandler for PluginRuntimeMcpServer { ) -> Result { let mut meta = context.meta; meta.0.remove("progressToken"); + let resource_route = meta + .0 + .get("_codex_apps") + .and_then(Value::as_object) + .and_then(|codex_apps| codex_apps.get("resource_uri")) + .and_then(Value::as_str) + .map(str::to_string); self.calls .lock() .expect("recorded tool calls lock poisoned") @@ -335,19 +480,12 @@ impl ServerHandler for PluginRuntimeMcpServer { })); match request.name.as_ref() { - PRODUCT_SEARCH_TOOL => Ok(CallToolResult::structured(json!({ "products": [] }))), - "fetch_resource" => { - let mut result = CallToolResult::structured(json!({ - "contents": [{ - "uri": RESOURCE_URI, - "mimeType": "text/html", - "text": "Walmart", - "meta": { "app": CONNECTOR_ID }, - }] - })); - result.content = vec![Content::text("Fetched Walmart resource")]; - Ok(result) + name if [CATALOG_A.product_search_tool, CATALOG_B.product_search_tool] + .contains(&name) => + { + Ok(CallToolResult::structured(json!({ "products": [] }))) } + "fetch_resource" => fetch_resource_result(resource_route.as_deref()), name => Err(rmcp::ErrorData::invalid_params( format!("unknown tool: {name}"), None, @@ -356,22 +494,55 @@ impl ServerHandler for PluginRuntimeMcpServer { } } -async fn wait_for_mcp_tool_call_completed( - mcp: &mut TestAppServer, - call_id: &str, -) -> Result { - loop { - let notification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("item/completed"), - ) - .await??; - let Some(params) = notification.params else { - continue; - }; - let completed: ItemCompletedNotification = serde_json::from_value(params)?; - if matches!(&completed.item, ThreadItem::McpToolCall { id, .. } if id == call_id) { - return Ok(completed); - } - } +fn product_search_tool(app: AppFixture) -> Tool { + let mut tool = Tool::new( + Cow::Borrowed(app.product_search_tool), + Cow::Borrowed(CONNECTOR_DESCRIPTION), + Arc::new(JsonObject::new()), + ); + tool.annotations = Some(ToolAnnotations::new().read_only(true)); + tool.meta = Some(Meta(serde_json::Map::from_iter([ + ("connector_id".to_string(), json!(app.connector_id)), + ("connector_name".to_string(), json!(app.connector_id)), + ( + "connector_description".to_string(), + json!(CONNECTOR_DESCRIPTION), + ), + ("link_id".to_string(), json!(app.link_id)), + ("openai/outputTemplate".to_string(), json!(RESOURCE_URI)), + ( + "_codex_apps".to_string(), + json!({ + "resource_uri": format!( + "/{}/{}/{}", + app.connector_id, app.link_id, app.product_search_tool + ), + "contains_mcp_source": true, + }), + ), + ]))); + tool +} + +fn fetch_resource_result(resource_route: Option<&str>) -> Result { + let app = [CATALOG_A, CATALOG_B] + .into_iter() + .find(|app| { + resource_route + == Some(format!("/{}/{}/fetch_resource", app.connector_id, app.link_id).as_str()) + }) + .ok_or_else(|| { + rmcp::ErrorData::invalid_params( + format!("unknown fetch_resource route: {resource_route:?}"), + None, + ) + })?; + Ok(CallToolResult::structured(json!({ + "contents": [{ + "uri": RESOURCE_URI, + "mimeType": "text/html", + "text": app.html, + "meta": { "app": app.connector_id }, + }] + }))) } diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 960baf3423..5fb2c68c83 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -120,6 +120,54 @@ pub struct McpConnectionManager { startup_cancellation_token: CancellationToken, } +/// A live connection with server identity, tools, and calls pinned to one client. +pub struct McpServerConnection { + server_name: String, + managed_client: ManagedClient, + tools: Vec, +} + +impl McpServerConnection { + /// Returns authoritative server implementation metadata from initialization. + pub fn server_info(&self) -> &McpServerInfo { + &self.managed_client.server_info + } + + /// Returns the filtered metadata for a currently available raw MCP tool. + pub fn tool_info(&self, tool: &str) -> Option<&ToolInfo> { + self.tools + .iter() + .find(|tool_info| tool_info.tool.name.as_ref() == tool) + } + + /// Invoke a tool through this pinned server connection. + pub async fn call_tool( + &self, + tool: &str, + arguments: Option, + meta: Option, + ) -> Result { + call_tool_with_client( + &self.managed_client, + &self.server_name, + tool, + arguments, + meta, + ) + .await + } + + /// Read a resource through this pinned server connection. + pub async fn read_resource(&self, uri: &str) -> Result { + read_resource_with_client( + &self.managed_client, + &self.server_name, + ReadResourceRequestParams::new(uri), + ) + .await + } +} + impl McpConnectionManager { #[allow(clippy::too_many_arguments)] pub async fn new( @@ -479,6 +527,19 @@ impl McpConnectionManager { } } + /// Acquires a live connection to `server` and snapshots its filtered tools. + /// + /// Unlike presentation-oriented APIs, this waits for authoritative server + /// initialization and never falls back to cached startup metadata. + pub async fn server_connection(&self, server: &str) -> Result { + let managed_client = self.client_by_name(server).await?; + Ok(McpServerConnection { + server_name: server.to_string(), + tools: managed_client.current_tools(), + managed_client, + }) + } + /// Returns all tools with model-visible names normalized. #[instrument(level = "trace", skip_all, fields(mcp_server_count = self.clients.len()))] pub async fn list_all_tools(&self) -> Vec { @@ -534,6 +595,11 @@ impl McpConnectionManager { .client() .await .context("failed to get client")?; + let _tool_refresh_permit = managed_client + .tool_refresh_semaphore + .acquire() + .await + .context("tool refresh semaphore closed")?; let list_start = Instant::now(); let fetch_start = Instant::now(); @@ -541,7 +607,7 @@ impl McpConnectionManager { .codex_apps_tools_cache_context .as_ref() .map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh)); - let tools = list_tools_for_client_uncached( + let fetched_tools = list_tools_for_client_uncached( CODEX_APPS_MCP_SERVER_NAME, /*is_codex_apps_mcp_server*/ true, &managed_client.client, @@ -558,16 +624,23 @@ impl McpConnectionManager { &[], ); - let tools = - match ( - managed_client.codex_apps_tools_cache_context.as_ref(), + managed_client.replace_tools(filter_tools( + fetched_tools.clone(), + &managed_client.tool_filter, + )); + + let tools = match ( + managed_client.codex_apps_tools_cache_context.as_ref(), + fetch_ticket, + ) { + (Some(cache_context), Some(fetch_ticket)) => cache_context.publish_if_newest_accepted( fetch_ticket, - ) { - (Some(cache_context), Some(fetch_ticket)) => cache_context - .publish_if_newest_accepted(fetch_ticket, &managed_client.server_info, tools), - (None, None) => tools, - _ => unreachable!("Codex Apps fetch ticket requires cache context"), - }; + &managed_client.server_info, + fetched_tools, + ), + (None, None) => fetched_tools, + _ => unreachable!("Codex Apps fetch ticket requires cache context"), + }; emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, list_start.elapsed(), @@ -740,33 +813,7 @@ impl McpConnectionManager { meta: Option, ) -> Result { let client = self.client_by_name(server).await?; - if !client.tool_filter.allows(tool) { - return Err(anyhow!( - "tool '{tool}' is disabled for MCP server '{server}'" - )); - } - - let result: rmcp::model::CallToolResult = client - .client - .call_tool(tool.to_string(), arguments, meta, client.tool_timeout) - .await - .with_context(|| format!("tool call failed for `{server}/{tool}`"))?; - - let content = result - .content - .into_iter() - .map(|content| { - serde_json::to_value(content) - .unwrap_or_else(|_| serde_json::Value::String("".to_string())) - }) - .collect(); - - Ok(CallToolResult { - content, - structured_content: result.structured_content, - is_error: result.is_error, - meta: result.meta.and_then(|meta| serde_json::to_value(meta).ok()), - }) + call_tool_with_client(&client, server, tool, arguments, meta).await } pub async fn server_supports_sandbox_state_meta_capability( @@ -818,14 +865,7 @@ impl McpConnectionManager { params: ReadResourceRequestParams, ) -> Result { let managed = self.client_by_name(server).await?; - let client = managed.client.clone(); - let timeout = managed.tool_timeout; - let uri = params.uri.clone(); - - client - .read_resource(params, timeout) - .await - .with_context(|| format!("resources/read failed for `{server}` ({uri})")) + read_resource_with_client(&managed, server, params).await } /// Returns presentation metadata without waiting for uncached clients still initializing. @@ -896,6 +936,55 @@ impl McpConnectionManager { } } +async fn call_tool_with_client( + client: &ManagedClient, + server: &str, + tool: &str, + arguments: Option, + meta: Option, +) -> Result { + if !client.tool_filter.allows(tool) { + return Err(anyhow!( + "tool '{tool}' is disabled for MCP server '{server}'" + )); + } + + let result: rmcp::model::CallToolResult = client + .client + .call_tool(tool.to_string(), arguments, meta, client.tool_timeout) + .await + .with_context(|| format!("tool call failed for `{server}/{tool}`"))?; + + let content = result + .content + .into_iter() + .map(|content| { + serde_json::to_value(content) + .unwrap_or_else(|_| serde_json::Value::String("".to_string())) + }) + .collect(); + + Ok(CallToolResult { + content, + structured_content: result.structured_content, + is_error: result.is_error, + meta: result.meta.and_then(|meta| serde_json::to_value(meta).ok()), + }) +} + +async fn read_resource_with_client( + managed: &ManagedClient, + server: &str, + params: ReadResourceRequestParams, +) -> Result { + let uri = params.uri.clone(); + managed + .client + .read_resource(params, managed.tool_timeout) + .await + .with_context(|| format!("resources/read failed for `{server}` ({uri})")) +} + impl Drop for McpConnectionManager { fn drop(&mut self) { self.startup_cancellation_token.cancel(); diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 1f61ab700e..69492d0432 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -114,7 +114,8 @@ async fn create_test_managed_client(tools: Vec) -> ManagedClient { .expect("create in-process RMCP client"), ), server_info: create_test_server_info("Ready"), - tools, + tools: Arc::new(arc_swap::ArcSwap::from_pointee(tools)), + tool_refresh_semaphore: Arc::new(tokio::sync::Semaphore::new(1)), tool_filter: ToolFilter::default(), tool_timeout: None, server_instructions: None, @@ -748,11 +749,16 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { ); cache_context.store_current_tools_for_test(vec![create_test_tool( CODEX_APPS_MCP_SERVER_NAME, - "calendar_create_event", + "cached_tool", )]); - let pending_client = futures::future::pending::>() - .boxed() - .shared(); + let (live_tx, live_rx) = tokio::sync::oneshot::channel(); + let pending_client = async move { + Ok(live_rx + .await + .expect("live client sender should remain available")) + } + .boxed() + .shared(); let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionManager::new_uninitialized( @@ -765,7 +771,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_server_info: None, + cached_server_info: Some(create_test_server_info("Cached")), codex_apps_tools_cache_context: Some(cache_context), tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), @@ -779,12 +785,50 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { let tool = tools .iter() .find(|tool| { - tool.canonical_tool_name() - == ToolName::namespaced("mcp__codex_apps", "calendar_create_event") + tool.canonical_tool_name() == ToolName::namespaced("mcp__codex_apps", "cached_tool") }) .expect("tool from shared cache"); assert_eq!(tool.server_name, CODEX_APPS_MCP_SERVER_NAME); - assert_eq!(tool.callable_name, "calendar_create_event"); + assert_eq!(tool.callable_name, "cached_tool"); + + assert!( + tokio::time::timeout( + Duration::from_millis(10), + manager.server_connection(CODEX_APPS_MCP_SERVER_NAME), + ) + .await + .is_err() + ); + + let mut live_client = create_test_managed_client(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "live_tool", + )]) + .await; + live_client.server_info = create_test_server_info("Live"); + let live_client_for_refresh = live_client.clone(); + assert!(live_tx.send(live_client).is_ok()); + + let connection = manager + .server_connection(CODEX_APPS_MCP_SERVER_NAME) + .await + .expect("live connection"); + assert_eq!(connection.server_info().title.as_deref(), Some("Live")); + assert!(connection.tool_info("live_tool").is_some()); + assert!(connection.tool_info("cached_tool").is_none()); + + live_client_for_refresh.replace_tools(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "refreshed_tool", + )]); + let refreshed = manager + .server_connection(CODEX_APPS_MCP_SERVER_NAME) + .await + .expect("refreshed connection"); + assert!(connection.tool_info("live_tool").is_some()); + assert!(connection.tool_info("refreshed_tool").is_none()); + assert!(refreshed.tool_info("live_tool").is_none()); + assert!(refreshed.tool_info("refreshed_tool").is_some()); } #[tokio::test] @@ -1122,6 +1166,13 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() .get(CODEX_APPS_MCP_SERVER_NAME), Some(&server_info) ); + assert!( + manager + .server_connection(CODEX_APPS_MCP_SERVER_NAME) + .await + .is_err(), + "authoritative connections must not fall back to cached startup data" + ); } #[tokio::test] diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index eda7ac2399..abd407b8bd 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -1,4 +1,5 @@ pub use connection_manager::McpConnectionManager; +pub use connection_manager::McpServerConnection; pub use connection_manager::tool_is_model_visible; pub use elicitation::ElicitationRequestRouter; pub use elicitation::ElicitationReviewRequest; diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 37819203f6..ff693c2611 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -37,6 +37,7 @@ use crate::tools::filter_tools; use crate::tools::tool_with_model_visible_input_schema; use anyhow::Result; use anyhow::anyhow; +use arc_swap::ArcSwap; use async_channel::Sender; use codex_api::SharedAuthProvider; use codex_async_utils::CancelErr; @@ -100,7 +101,8 @@ const UNTRUSTED_CONNECTOR_META_KEYS: &[&str] = &[ pub(crate) struct ManagedClient { pub(crate) client: Arc, pub(crate) server_info: McpServerInfo, - pub(crate) tools: Vec, + pub(crate) tools: Arc>>, + pub(crate) tool_refresh_semaphore: Arc, pub(crate) tool_filter: ToolFilter, pub(crate) tool_timeout: Option, pub(crate) server_instructions: Option, @@ -109,7 +111,15 @@ pub(crate) struct ManagedClient { } impl ManagedClient { - fn listed_tools(&self) -> Vec { + pub(crate) fn current_tools(&self) -> Vec { + self.tools.load_full().as_ref().clone() + } + + pub(crate) fn replace_tools(&self, tools: Vec) { + self.tools.store(Arc::new(tools)); + } + + pub(crate) fn listed_tools(&self) -> Vec { let total_start = Instant::now(); if let Some(tools) = self .codex_apps_tools_cache_context @@ -132,7 +142,7 @@ impl ManagedClient { ); } - self.tools.clone() + self.current_tools() } } @@ -824,7 +834,7 @@ async fn start_server_task( let fetch_ticket = codex_apps_tools_cache_context .as_ref() .map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::Startup)); - let tools = list_tools_for_client_uncached( + let fetched_tools = list_tools_for_client_uncached( &server_name, is_codex_apps_mcp_server, &client, @@ -839,13 +849,17 @@ async fn start_server_task( &[], ); let server_info = mcp_server_info_from_implementation(initialize_result.server_info); - let tools = match (codex_apps_tools_cache_context.as_ref(), fetch_ticket) { + match (codex_apps_tools_cache_context.as_ref(), fetch_ticket) { (Some(cache_context), Some(fetch_ticket)) => { - cache_context.publish_if_newest_accepted(fetch_ticket, &server_info, tools) + cache_context.publish_if_newest_accepted( + fetch_ticket, + &server_info, + fetched_tools.clone(), + ); } - (None, None) => tools, + (None, None) => {} _ => unreachable!("Codex Apps fetch ticket requires cache context"), - }; + } if is_codex_apps_mcp_server { emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, @@ -853,12 +867,13 @@ async fn start_server_task( &[("cache", "miss")], ); } - let tools = filter_tools(tools, &tool_filter); + let tools = filter_tools(fetched_tools, &tool_filter); let managed = ManagedClient { client: Arc::clone(&client), server_info, - tools, + tools: Arc::new(ArcSwap::from_pointee(tools)), + tool_refresh_semaphore: Arc::new(tokio::sync::Semaphore::new(1)), tool_timeout: Some(tool_timeout), tool_filter, server_instructions: initialize_result.instructions,