openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 17:38:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

mcp resource read: add edgecases

Browse Source
This commit is contained in:
fbauer
2026-07-06 16:41:56 -07:00
parent 36d2e2dbdb
commit d4c3867155
12 changed files with 981 additions and 295 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/lib.rs
+1
View File
@@ -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;
codex-rs/app-server/src/mcp_resource_origin.rs
+102
View File
@@ -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<String>,
pub(crate) link_id: Option<String>,
pub(crate) resource_uri: Option<String>,
}
#[derive(Default)]
pub(crate) struct McpResourceOriginIndex(VecDeque<(String, McpResourceOrigin)>);
impl McpResourceOrigin {
pub(crate) fn find(items: &[ThreadItem], call_id: &str) -> Option<Self> {
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<McpResourceOrigin> {
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<Item = &'a ThreadItem>) {
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;
codex-rs/app-server/src/mcp_resource_origin_tests.rs
+79
View File
@@ -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<String>, 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()),
}
}
codex-rs/app-server/src/request_processors/mcp_resource_read.rs
+52 -78
View File
@@ -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<CodexThread>,
origin_call_id: &str,
) -> anyhow::Result<McpResourceOrigin> {
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<String>,
resource_uri: Option<String>,
}
fn find_mcp_resource_origin(
items: &[ThreadItem],
origin_call_id: &str,
) -> Option<anyhow::Result<McpResourceOrigin>> {
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::<McpResourceReadResponse>(structured_content)?;
Ok(serde_json::to_value(response)?)
}
#[cfg(test)]
#[path = "mcp_resource_read_tests.rs"]
mod tests;
codex-rs/app-server/src/request_processors/mcp_resource_read_tests.rs
+174
View File
@@ -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<T>(result: anyhow::Result<T>, expected: &str) {
let error = result.err().expect("expected validation failure");
assert!(
error.to_string().contains(expected),
"expected error containing {expected:?}, got {error:#}"
);
}
codex-rs/app-server/src/request_processors/thread_processor.rs
+10
View File
@@ -3558,6 +3558,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(
codex-rs/app-server/src/thread_state.rs
+19
View File
@@ -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<ThreadSettings>,
listener_command_tx: Option<mpsc::UnboundedSender<ThreadListenerCommand>>,
current_turn_history: ThreadHistoryBuilder,
pub(crate) mcp_resource_origins: McpResourceOriginIndex,
listener_thread: Option<Weak<CodexThread>>,
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<McpResourceOrigin> {
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()
codex-rs/app-server/tests/suite/v2/mcp_resource_origin.rs
+324 -153
View File
@@ -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: "<html>Catalog A</html>",
};
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: "<html>Catalog B</html>",
};
#[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<Mutex<Vec<Value>>>,
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<OriginTestScenario> {
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::<McpResourceReadResponse>(read_response)?,
McpResourceReadResponse {
contents: vec![McpResourceContent::Text {
uri: RESOURCE_URI.to_string(),
mime_type: Some("text/html".to_string()),
text: "<html>Walmart</html>".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<McpResourceReadResponse> {
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<JSONRPCError> {
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<i64> {
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<Mutex<Vec<Value>>>,
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::<Vec<_>>();
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::<Vec<_>>();
assert_eq!(actual, expected);
}
async fn start_test_app_server(
@@ -278,41 +446,11 @@ impl ServerHandler for PluginRuntimeMcpServer {
_request: Option<rmcp::model::PaginatedRequestParams>,
_context: RequestContext<RoleServer>,
) -> Result<ListToolsResult, rmcp::ErrorData> {
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<CallToolResult, rmcp::ErrorData> {
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": "<html>Walmart</html>",
"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<ItemCompletedNotification> {
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<CallToolResult, rmcp::ErrorData> {
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 },
}]
})))
}
codex-rs/codex-mcp/src/connection_manager.rs
+134 -45
View File
@@ -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<ToolInfo>,
}
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<serde_json::Value>,
meta: Option<serde_json::Value>,
) -> Result<CallToolResult> {
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<ReadResourceResult> {
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(
@@ -482,6 +530,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<McpServerConnection> {
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<ToolInfo> {
@@ -537,6 +598,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();
@@ -544,7 +610,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,
@@ -561,16 +627,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(),
@@ -743,33 +816,7 @@ impl McpConnectionManager {
meta: Option<serde_json::Value>,
) -> Result<CallToolResult> {
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("<content>".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(
@@ -821,14 +868,7 @@ impl McpConnectionManager {
params: ReadResourceRequestParams,
) -> Result<ReadResourceResult> {
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.
@@ -899,6 +939,55 @@ impl McpConnectionManager {
}
}
async fn call_tool_with_client(
client: &ManagedClient,
server: &str,
tool: &str,
arguments: Option<serde_json::Value>,
meta: Option<serde_json::Value>,
) -> Result<CallToolResult> {
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("<content>".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<ReadResourceResult> {
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();
codex-rs/codex-mcp/src/connection_manager_tests.rs
+60 -9
View File
@@ -115,7 +115,8 @@ async fn create_test_managed_client(tools: Vec<ToolInfo>) -> 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,
@@ -771,11 +772,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::<Result<ManagedClient, StartupOutcomeError>>()
.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(
@@ -788,7 +794,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)),
@@ -802,12 +808,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]
@@ -1145,6 +1189,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]
codex-rs/codex-mcp/src/lib.rs
+1
View File
@@ -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::ElicitationLifecycle;
pub use elicitation::ElicitationRequestRouter;
codex-rs/codex-mcp/src/rmcp_client.rs
+25 -10
View File
@@ -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<RmcpClient>,
pub(crate) server_info: McpServerInfo,
pub(crate) tools: Vec<ToolInfo>,
pub(crate) tools: Arc<ArcSwap<Vec<ToolInfo>>>,
pub(crate) tool_refresh_semaphore: Arc<tokio::sync::Semaphore>,
pub(crate) tool_filter: ToolFilter,
pub(crate) tool_timeout: Option<Duration>,
pub(crate) server_instructions: Option<String>,
@@ -109,7 +111,15 @@ pub(crate) struct ManagedClient {
}
impl ManagedClient {
fn listed_tools(&self) -> Vec<ToolInfo> {
pub(crate) fn current_tools(&self) -> Vec<ToolInfo> {
self.tools.load_full().as_ref().clone()
}
pub(crate) fn replace_tools(&self, tools: Vec<ToolInfo>) {
self.tools.store(Arc::new(tools));
}
pub(crate) fn listed_tools(&self) -> Vec<ToolInfo> {
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,