codex_mcp/
binding.rs

1//! Immutable MCP state bound to one model sampling request.
2
3use std::collections::HashMap;
4use std::fmt;
5use std::future::Future;
6use std::sync::Arc;
7
8use anyhow::Context;
9use anyhow::Result;
10use codex_config::AppToolApproval;
11use codex_protocol::mcp::CallToolResult;
12use rmcp::model::ListResourceTemplatesResult;
13use rmcp::model::ListResourcesResult;
14use rmcp::model::PaginatedRequestParams;
15use rmcp::model::ReadResourceRequestParams;
16use rmcp::model::ReadResourceResult;
17use rmcp::model::Resource;
18use rmcp::model::ResourceTemplate;
19use serde_json::Value as JsonValue;
20use tokio::sync::RwLock;
21
22use crate::McpConfig;
23use crate::binding_clients::McpBindingClients;
24use crate::connection_manager::McpConnectionManager;
25use crate::resource_client::McpResourceClient;
26use crate::rmcp_client::ManagedClient;
27use crate::server::McpServerMetadata;
28use crate::tools::ToolInfo;
29
30/// The exact tool catalog and execution handles for one model sampling request.
31pub struct McpBinding {
32    connections: Arc<McpConnectionManager>,
33    clients: Arc<McpBindingClients>,
34    config: Arc<McpConfig>,
35    plugins_available: bool,
36    tools: Vec<ToolInfo>,
37    calls: HashMap<(String, String), PreparedMcpCall>,
38}
39
40impl McpBinding {
41    /// Creates an empty binding for tests and callers without a materialized runtime.
42    pub fn empty(config: Arc<McpConfig>) -> Self {
43        Self::new(
44            Arc::new(McpConnectionManager::empty(config.prefix_mcp_tool_names)),
45            Arc::new(McpBindingClients::new(HashMap::new())),
46            config,
47            /*plugins_available*/ false,
48            Vec::new(),
49            HashMap::new(),
50        )
51    }
52
53    pub(crate) fn new(
54        connections: Arc<McpConnectionManager>,
55        clients: Arc<McpBindingClients>,
56        config: Arc<McpConfig>,
57        plugins_available: bool,
58        tools: Vec<ToolInfo>,
59        calls: HashMap<(String, String), PreparedMcpCall>,
60    ) -> Self {
61        Self {
62            connections,
63            clients,
64            config,
65            plugins_available,
66            tools,
67            calls,
68        }
69    }
70
71    pub fn config(&self) -> &Arc<McpConfig> {
72        &self.config
73    }
74
75    pub fn plugins_available(&self) -> bool {
76        self.plugins_available
77    }
78
79    /// Returns the frozen catalog advertised for this sampling request.
80    pub fn tools(&self) -> &[ToolInfo] {
81        &self.tools
82    }
83
84    /// Binds a call to the exact client and metadata advertised by this binding.
85    pub fn prepare_call(&self, server: &str, tool: &str) -> Option<PreparedMcpCall> {
86        self.calls
87            .get(&(server.to_string(), tool.to_string()))
88            .cloned()
89    }
90
91    pub fn has_servers(&self) -> bool {
92        self.connections.has_servers()
93    }
94
95    /// Returns resource access bound to this binding's exact connection set.
96    pub fn resource_client(&self) -> McpResourceClient {
97        McpResourceClient::for_binding(Arc::clone(&self.clients))
98    }
99
100    pub async fn list_resources(
101        &self,
102        server: &str,
103        params: Option<PaginatedRequestParams>,
104    ) -> Result<ListResourcesResult> {
105        self.clients.list_resources(server, params).await
106    }
107
108    pub async fn list_all_resources(
109        &self,
110        include_server: impl Fn(&str) -> bool,
111    ) -> HashMap<String, Vec<Resource>> {
112        self.clients.list_all_resources(include_server).await
113    }
114
115    pub async fn list_resource_templates(
116        &self,
117        server: &str,
118        params: Option<PaginatedRequestParams>,
119    ) -> Result<ListResourceTemplatesResult> {
120        self.clients.list_resource_templates(server, params).await
121    }
122
123    pub async fn list_all_resource_templates(
124        &self,
125        include_server: impl Fn(&str) -> bool,
126    ) -> HashMap<String, Vec<ResourceTemplate>> {
127        self.clients
128            .list_all_resource_templates(include_server)
129            .await
130    }
131
132    pub async fn read_resource(
133        &self,
134        server: &str,
135        params: ReadResourceRequestParams,
136    ) -> Result<ReadResourceResult> {
137        self.clients.read_resource(server, params).await
138    }
139}
140
141impl fmt::Debug for McpBinding {
142    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
143        formatter
144            .debug_struct("McpBinding")
145            .field("tools", &self.tools)
146            .field("prepared_call_count", &self.calls.len())
147            .finish_non_exhaustive()
148    }
149}
150
151/// A call bound to the exact client, tool, timeout, and server metadata seen by
152/// one [`McpBinding`].
153#[derive(Clone)]
154pub struct PreparedMcpCall {
155    _connections: Arc<McpConnectionManager>,
156    client: Arc<ManagedClient>,
157    catalog_revision: u64,
158    catalog_revision_source: Arc<RwLock<u64>>,
159    tool_info: ToolInfo,
160    server_name: String,
161    server_metadata: McpServerMetadata,
162    plugin_id: Option<String>,
163    selected_plugin_server: bool,
164}
165
166impl PreparedMcpCall {
167    #[expect(
168        clippy::too_many_arguments,
169        reason = "the exact call authority stays together"
170    )]
171    pub(crate) fn new(
172        connections: Arc<McpConnectionManager>,
173        client: Arc<ManagedClient>,
174        catalog_revision: u64,
175        catalog_revision_source: Arc<RwLock<u64>>,
176        tool_info: ToolInfo,
177        server_metadata: McpServerMetadata,
178        plugin_id: Option<String>,
179        selected_plugin_server: bool,
180    ) -> Self {
181        let server_name = tool_info.server_name.clone();
182        Self {
183            _connections: connections,
184            client,
185            catalog_revision,
186            catalog_revision_source,
187            tool_info,
188            server_name,
189            server_metadata,
190            plugin_id,
191            selected_plugin_server,
192        }
193    }
194
195    pub fn tool_info(&self) -> &ToolInfo {
196        &self.tool_info
197    }
198
199    pub fn server_name(&self) -> &str {
200        &self.server_name
201    }
202
203    pub fn server_origin(&self) -> Option<&str> {
204        self.server_metadata
205            .origin
206            .as_ref()
207            .map(super::server::McpServerOrigin::as_str)
208    }
209
210    pub fn server_environment_id(&self) -> &str {
211        &self.server_metadata.environment_id
212    }
213
214    pub fn server_pollutes_memory(&self) -> bool {
215        self.server_metadata.pollutes_memory
216    }
217
218    pub fn tool_approval_mode(&self) -> AppToolApproval {
219        self.server_metadata
220            .tool_approval_mode(&self.tool_info.tool.name)
221    }
222
223    pub fn plugin_id(&self) -> Option<&str> {
224        self.plugin_id.as_deref()
225    }
226
227    pub fn is_selected_plugin_server(&self) -> bool {
228        self.selected_plugin_server
229    }
230
231    pub async fn server_supports_sandbox_state_meta_capability(&self) -> Result<bool> {
232        Ok(self.client.server_supports_sandbox_state_meta_capability)
233    }
234
235    pub async fn call(
236        &self,
237        arguments: Option<JsonValue>,
238        meta: Option<JsonValue>,
239    ) -> Result<CallToolResult> {
240        self.call_with_preparation(|| async move { Ok((arguments, meta)) })
241            .await
242    }
243
244    /// Runs irreversible call preparation and execution under the authority of
245    /// this call's exact catalog revision.
246    #[expect(
247        clippy::await_holding_invalid_type,
248        reason = "catalog replacement must remain serialized with call preparation and execution"
249    )]
250    pub async fn call_with_preparation<F, Fut>(&self, prepare: F) -> Result<CallToolResult>
251    where
252        F: FnOnce() -> Fut,
253        Fut: Future<Output = Result<(Option<JsonValue>, Option<JsonValue>)>>,
254    {
255        let tool_name = self.tool_info.tool.name.to_string();
256        let current_revision = self.catalog_revision_source.read().await;
257        if *current_revision != self.catalog_revision {
258            return Err(anyhow::anyhow!(
259                "tool call rejected because the catalog changed after `{}/{tool_name}` was prepared",
260                self.server_name
261            ));
262        }
263        let (arguments, meta) = prepare().await?;
264        let result = self
265            .client
266            .client
267            .call_tool(tool_name.clone(), arguments, meta, self.client.tool_timeout)
268            .await
269            .with_context(|| format!("tool call failed for `{}/{tool_name}`", self.server_name))?;
270        drop(current_revision);
271        Ok(call_tool_result_from_rmcp(result))
272    }
273}
274
275fn call_tool_result_from_rmcp(result: rmcp::model::CallToolResult) -> CallToolResult {
276    let content = result
277        .content
278        .into_iter()
279        .map(|content| {
280            serde_json::to_value(content)
281                .unwrap_or_else(|_| JsonValue::String("<content>".to_string()))
282        })
283        .collect();
284    CallToolResult {
285        content,
286        structured_content: result.structured_content,
287        is_error: result.is_error,
288        meta: result.meta.and_then(|meta| serde_json::to_value(meta).ok()),
289    }
290}
291
292#[cfg(test)]
293#[path = "binding_tests.rs"]
294mod tests;
codex_mcp/binding.rs

codex_mcp/
binding.rs
