From 3c074bb44a31f474cc60ccbd7a7230414992465b Mon Sep 17 00:00:00 2001 From: yitaoh-oai Date: Tue, 3 Nov 2026 01:44:09 +0000 Subject: [PATCH] Add model Search discovery to MCP bindings (#59034) ## What changed Add `McpBinding` APIs to discover the model's Search tool profile through an already authorized, host-owned Apps connection. Discovery sends turn metadata, paginates `tools/list`, and returns a scoped binding with the refreshed declaration and connector-provided invocation context. Require a visible Search tool and an object-valued invocation context. Preserve the context only for calls using the same client, while allowing normal dispatch to recheck visibility, account selection, and approval policy. Bound discovery to five seconds and distinguish critical timeout, transport, and server errors from unavailability. Expose the existing retryable transport error classifier for this distinction. ## Testing Add coverage for opaque invocation context extraction and extend transport error classification coverage to HTTP 503 and 403 responses, exercising the exported classifier through a contextual error. GitOrigin-RevId: 9848d4adb91e3fb2b9406349c31d1fd5c5c47a33 --- codex-rs/codex-mcp/src/binding.rs | 6 + codex-rs/codex-mcp/src/lib.rs | 3 + codex-rs/codex-mcp/src/model_search.rs | 143 ++++++++++++++++++ codex-rs/codex-mcp/src/model_search_tests.rs | 20 +++ .../src/rmcp_client/managed_client.rs | 5 +- codex-rs/rmcp-client/src/lib.rs | 1 + codex-rs/rmcp-client/src/rmcp_client.rs | 9 +- codex-rs/rmcp-client/src/service_error.rs | 8 + .../src/streamable_http_retry_tests.rs | 15 +- 9 files changed, 202 insertions(+), 8 deletions(-) create mode 100644 codex-rs/codex-mcp/src/model_search.rs create mode 100644 codex-rs/codex-mcp/src/model_search_tests.rs diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index 6b07723b25..36809350f2 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -1,5 +1,11 @@ //! Immutable MCP catalog and execution handles. +#[path = "model_search.rs"] +mod model_search; +pub use model_search::SEARCH_CONNECTOR; +pub use model_search::SearchDiscovery; +pub use model_search::is_model_search_tool; + use std::collections::HashMap; use std::fmt; use std::future::Future; diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 7a40003261..afc5cfb1da 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -1,6 +1,9 @@ pub use binding::McpBinding; pub use binding::McpPermissionSource; pub use binding::PreparedMcpCall; +pub use binding::SEARCH_CONNECTOR; +pub use binding::SearchDiscovery; +pub use binding::is_model_search_tool; pub use client_capabilities::client_mcp_extensions; pub use client_tool_catalog::CodexAppsToolSnapshot; pub use codex_rmcp_client::McpProtocolMode; diff --git a/codex-rs/codex-mcp/src/model_search.rs b/codex-rs/codex-mcp/src/model_search.rs new file mode 100644 index 0000000000..b492c25aac --- /dev/null +++ b/codex-rs/codex-mcp/src/model_search.rs @@ -0,0 +1,143 @@ +//! Resolve the model's Search profile on an already authorized Apps connection. +//! Keep the projected declaration and invocation pin private to its sampling step. + +use super::McpBinding; +use crate::CODEX_APPS_MCP_SERVER_NAME; +use crate::MCP_TOOL_CODEX_APPS_META_KEY; +use crate::pagination::MAX_CODEX_APPS_TOOL_CATALOG_ITEMS; +use crate::pagination::collect_paginated_with_limit; +use rmcp::model::RequestMetaObject; +use serde_json::Value as JsonValue; +use std::sync::Arc; +use std::time::Duration; + +pub const SEARCH_CONNECTOR: &str = "connector_openai_search_service"; + +pub fn is_model_search_tool(tool: &crate::ToolInfo) -> bool { + tool.server_name == CODEX_APPS_MCP_SERVER_NAME + && tool.connector_id.as_deref() == Some(SEARCH_CONNECTOR) + && tool.tool.name.rsplit('.').next() == Some("web_run") +} + +pub enum SearchDiscovery { + Ready(Arc), + Critical(&'static str), + Unavailable, +} + +impl McpBinding { + /// Preserve the discovery pin while letting normal dispatch recheck current tool visibility, + /// account selection and approval policy. Never transfer the pin to a different transport. + pub fn search_invocation_context(&self, live: &crate::PreparedMcpCall) -> Option<&JsonValue> { + if !live.is_host_owned_apps() { + return None; + } + let tool = live.tool_info(); + let scoped = self + .calls + .get(&(tool.server_name.clone(), tool.tool.name.to_string()))?; + if !Arc::ptr_eq(&scoped.client.client, &live.client.client) { + return None; + } + invocation_context(&scoped.tool_info().tool) + } + + /// A catalog omission, policy denial, or untrusted registration never authorizes fallback. + pub async fn discover_model_search(&self, turn_metadata: JsonValue) -> SearchDiscovery { + let Some((index, original)) = self + .tools + .iter() + .enumerate() + .find(|(_, tool)| is_model_search_tool(tool)) + else { + return SearchDiscovery::Unavailable; + }; + let key = (original.server_name.clone(), original.tool.name.to_string()); + let Some(call) = self + .calls + .get(&key) + .filter(|call| call.is_host_owned_apps()) + else { + return SearchDiscovery::Unavailable; + }; + let meta = RequestMetaObject::from(serde_json::Map::from_iter([( + "x-codex-turn-metadata".to_string(), + turn_metadata, + )])); + let client = &call.client; + let limit = Duration::from_secs(5); + let listed = collect_paginated_with_limit( + "tools/list", + /*overall_timeout*/ None, + MAX_CODEX_APPS_TOOL_CATALOG_ITEMS, + |params| { + let mut params = params.unwrap_or_default(); + params.meta = Some(meta.clone()); + async move { + let page = client + .client + .list_tools_with_connector_ids(Some(params), /*timeout*/ None) + .await?; + Ok((page.tools, page.next_cursor)) + } + }, + ); + let listed = tokio::time::timeout(limit, client.with_lease_request(listed)).await; + let tools = match listed { + Err(_) => return SearchDiscovery::Critical("discovery_timeout"), + Ok(Err(err)) => { + if matches!( + codex_rmcp_client::request_error_code(&err), + Some("transport_closed" | "request_timeout") + ) || codex_rmcp_client::is_retryable_transport_error(&err) + || codex_rmcp_client::mcp_error(&err).is_some_and(|e| e.code.0 == -32603) + { + return SearchDiscovery::Critical("discovery_error"); + } + return SearchDiscovery::Unavailable; + } + Ok(Ok(tools)) => tools, + }; + let Some(listed) = tools.into_iter().find(|tool| { + tool.connector_id.as_deref() == Some(SEARCH_CONNECTOR) + && tool.tool.name == original.tool.name + }) else { + return SearchDiscovery::Unavailable; + }; + let mut projected = original.clone(); + projected.namespace_description = listed.connector_description; + projected.tool = listed.tool; + if !crate::tool_is_model_visible(&projected) + || invocation_context(&projected.tool).is_none() + { + return SearchDiscovery::Unavailable; + } + let mut tools = self.tools.to_vec(); + tools[index] = projected.clone(); + let mut calls = self.calls.clone(); + let mut call = call.clone(); + call.tool_info = projected.into(); + calls.insert(key, call); + SearchDiscovery::Ready(Arc::new(Self::new( + Arc::clone(&self.connections), + Arc::clone(&self.clients), + Arc::clone(&self.config), + self.plugins_available, + tools, + calls, + ))) + } +} + +// The connector owns the contents of its invocation context. +fn invocation_context(tool: &rmcp::model::Tool) -> Option<&JsonValue> { + tool.meta + .as_ref()? + .get(MCP_TOOL_CODEX_APPS_META_KEY)? + .get("invocation_context") + .filter(|context| context.is_object()) +} + +#[cfg(test)] +#[path = "model_search_tests.rs"] +mod tests; diff --git a/codex-rs/codex-mcp/src/model_search_tests.rs b/codex-rs/codex-mcp/src/model_search_tests.rs new file mode 100644 index 0000000000..942d9da977 --- /dev/null +++ b/codex-rs/codex-mcp/src/model_search_tests.rs @@ -0,0 +1,20 @@ +#![cfg(test)] + +use super::invocation_context; +use rmcp::model::Tool; +use serde_json::json; + +#[test] +fn discovery_preserves_the_connectors_opaque_invocation_context() { + let mut declaration = json!({ + "name": "search_service.web_run", + "inputSchema": {"type": "object"}, + }); + let tool: Tool = serde_json::from_value(declaration.clone()).unwrap(); + assert_eq!(invocation_context(&tool), None); + + let pin = json!({"revision": "connector-selected"}); + declaration["_meta"] = json!({"_codex_apps": {"invocation_context": pin}}); + let tool: Tool = serde_json::from_value(declaration).unwrap(); + assert_eq!(invocation_context(&tool), Some(&pin)); +} diff --git a/codex-rs/codex-mcp/src/rmcp_client/managed_client.rs b/codex-rs/codex-mcp/src/rmcp_client/managed_client.rs index 20da4e82a4..7c3d830dc2 100644 --- a/codex-rs/codex-mcp/src/rmcp_client/managed_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client/managed_client.rs @@ -59,7 +59,10 @@ impl ManagedClient { } } - async fn with_lease_request(&self, operation: impl Future>) -> Result { + pub(crate) async fn with_lease_request( + &self, + operation: impl Future>, + ) -> Result { // Keep one copy of the request future as it passes through both lease scopes. let operation = std::pin::pin!(operation); self.with_active_shared_lease(async { diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index c78b788d08..50a4f531d6 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -106,6 +106,7 @@ pub use rmcp_client::RmcpClient; pub use rmcp_client::SendElicitation; pub use rmcp_client::StreamableHttpBearerToken; pub use rmcp_client::ToolWithConnectorId; +pub use service_error::is_retryable_transport_error; pub use service_error::mcp_error; pub use service_error::request_error_code; pub use service_error::request_timeout_error; diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 46b86cac6b..d2766f0155 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -1820,7 +1820,9 @@ impl RmcpClient { .await { Ok(result) => return Ok(result), - Err(error) if Self::is_retryable_tools_list_error(label, &error) => { + Err(error) + if label == "tools/list" && Self::is_retryable_transport_error(&error) => + { let Some(retry_delay_ms) = retry_delay_ms else { return Err(error); }; @@ -1871,10 +1873,7 @@ impl RmcpClient { } } - fn is_retryable_tools_list_error(label: &str, error: &ClientOperationError) -> bool { - if label != "tools/list" { - return false; - } + pub(crate) fn is_retryable_transport_error(error: &ClientOperationError) -> bool { let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) = error else { diff --git a/codex-rs/rmcp-client/src/service_error.rs b/codex-rs/rmcp-client/src/service_error.rs index 60dd3d2990..ee411f2161 100644 --- a/codex-rs/rmcp-client/src/service_error.rs +++ b/codex-rs/rmcp-client/src/service_error.rs @@ -6,6 +6,14 @@ use rmcp::service::ServiceError; use std::time::Duration; use crate::rmcp_client::ClientOperationError; +use crate::rmcp_client::RmcpClient; + +/// Whether tools/list would retry this transport error before returning it to the caller. +pub fn is_retryable_transport_error(error: &Error) -> bool { + error + .downcast_ref() + .is_some_and(RmcpClient::is_retryable_transport_error) +} /// Returns the server's structured protocol error, including its original data. pub fn mcp_error(error: &Error) -> Option<&ErrorData> { diff --git a/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs b/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs index a1e4dc717f..3cda85b432 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs @@ -55,6 +55,8 @@ fn streamable_http_retry_classifies_transport_and_auth_errors() { ), )), StreamableHttpError::UnexpectedServerResponse("HTTP 502: upstream failure".into()), + StreamableHttpError::UnexpectedServerResponse("HTTP 503: service unavailable".into()), + StreamableHttpError::UnexpectedServerResponse("HTTP 403: forbidden".into()), StreamableHttpError::UnexpectedServerResponse("HTTP 400: bad request".into()), StreamableHttpError::Client(StreamableHttpClientAdapterError::Auth( codex_api::AuthError::Transient("credential refresh unavailable".to_string()), @@ -68,9 +70,18 @@ fn streamable_http_retry_classifies_transport_and_auth_errors() { ]; assert_eq!( - errors.map(|error| RmcpClient::is_retryable_streamable_http_error(&error)), + errors.map(|error| { + let transport = DynamicTransportError::from_parts( + "streamable_http", + TypeId::of::<()>(), + Box::new(error), + ); + let error = + ClientOperationError::from(rmcp::service::ServiceError::TransportSend(transport)); + crate::is_retryable_transport_error(&anyhow::Error::new(error).context("listing tools")) + }), [ - false, true, true, true, false, true, false, true, false, false + false, true, true, true, false, true, true, false, false, true, false, false ], ); }