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 ], ); }