Skip to main content

a3s_code_core/mcp/
binding.rs

1//! Immutable, generation-exact MCP runtime bindings.
2//!
3//! Compatibility MCP APIs resolve calls through a mutable multi-server
4//! [`McpManager`](super::McpManager). Capability projections cannot use that
5//! boundary: a later registration under the same server name could otherwise
6//! change an already admitted Run. [`McpBinding`] instead owns one exact
7//! initialized client and one canonical frozen tool catalog.
8
9use std::fmt;
10use std::sync::Arc;
11
12use anyhow::{anyhow, Result};
13use async_trait::async_trait;
14use thiserror::Error;
15
16use super::client::McpClient;
17use super::protocol::McpTool;
18use super::result::project_tool_result;
19use super::tools::annotation_requires_confirmation;
20use crate::tools::{Tool, ToolContext, ToolOutput};
21
22pub const MAX_MCP_BINDING_TOOLS: usize = 1_024;
23pub const MAX_MCP_BINDING_DEFINITION_BYTES: usize = 16 * 1024 * 1024;
24const MAX_MCP_NAME_BYTES: usize = 256;
25const MAX_MCP_FULL_TOOL_NAME_BYTES: usize = 768;
26
27/// Invalid or no-longer-ready exact MCP binding.
28#[derive(Clone, Debug, Eq, Error, PartialEq)]
29pub enum McpBindingError {
30    #[error("MCP {field} is empty, padded, contains control characters, or exceeds its bound")]
31    InvalidName { field: &'static str },
32    #[error("MCP binding server name does not match its exact client identity")]
33    ClientNameMismatch,
34    #[error("MCP binding exceeds the {field} bound of {max}")]
35    BoundExceeded { field: &'static str, max: usize },
36    #[error("MCP binding repeats tool name '{name}'")]
37    DuplicateToolName { name: String },
38    #[error("MCP binding tool catalog is not serializable")]
39    InvalidToolCatalog,
40    #[error("MCP binding client is not initialized and connected")]
41    ClientNotReady,
42}
43
44/// One exact connected MCP server projected into an immutable Code catalog.
45///
46/// The binding deliberately contains no manager, package locator, Grant, or
47/// mutable discovery handle. Its client and canonical tool definitions are
48/// frozen together, while the owning capability transaction retains the
49/// client's asynchronous close effect. A Run separately retains the matching
50/// non-clone A3S Use generation lease.
51pub struct McpBinding {
52    server_name: Box<str>,
53    client: Arc<McpClient>,
54    tools: Arc<[McpTool]>,
55}
56
57impl McpBinding {
58    pub fn new(
59        server_name: impl Into<String>,
60        client: Arc<McpClient>,
61        tools: impl IntoIterator<Item = McpTool>,
62    ) -> std::result::Result<Self, McpBindingError> {
63        let server_name = server_name.into();
64        validate_name("server name", &server_name)?;
65        if client.name != server_name {
66            return Err(McpBindingError::ClientNameMismatch);
67        }
68        if !client.is_ready() {
69            return Err(McpBindingError::ClientNotReady);
70        }
71
72        let mut tools = tools.into_iter().collect::<Vec<_>>();
73        if tools.len() > MAX_MCP_BINDING_TOOLS {
74            return Err(McpBindingError::BoundExceeded {
75                field: "tool count",
76                max: MAX_MCP_BINDING_TOOLS,
77            });
78        }
79        tools.sort_by(|left, right| left.name.cmp(&right.name));
80
81        let mut definition_bytes = 0_usize;
82        let mut previous_name: Option<&str> = None;
83        for tool in &tools {
84            validate_name("tool name", &tool.name)?;
85            let full_name_len = "mcp__"
86                .len()
87                .saturating_add(server_name.len())
88                .saturating_add("__".len())
89                .saturating_add(tool.name.len());
90            if full_name_len > MAX_MCP_FULL_TOOL_NAME_BYTES {
91                return Err(McpBindingError::BoundExceeded {
92                    field: "fully qualified tool name bytes",
93                    max: MAX_MCP_FULL_TOOL_NAME_BYTES,
94                });
95            }
96            if previous_name == Some(tool.name.as_str()) {
97                return Err(McpBindingError::DuplicateToolName {
98                    name: tool.name.clone(),
99                });
100            }
101            previous_name = Some(&tool.name);
102            let encoded =
103                serde_json::to_vec(tool).map_err(|_| McpBindingError::InvalidToolCatalog)?;
104            definition_bytes = definition_bytes.saturating_add(encoded.len());
105            if definition_bytes > MAX_MCP_BINDING_DEFINITION_BYTES {
106                return Err(McpBindingError::BoundExceeded {
107                    field: "tool definition bytes",
108                    max: MAX_MCP_BINDING_DEFINITION_BYTES,
109                });
110            }
111        }
112
113        Ok(Self {
114            server_name: server_name.into_boxed_str(),
115            client,
116            tools: tools.into(),
117        })
118    }
119
120    pub fn server_name(&self) -> &str {
121        &self.server_name
122    }
123
124    pub fn tools(&self) -> &[McpTool] {
125        &self.tools
126    }
127
128    pub fn is_ready(&self) -> bool {
129        self.client.is_ready()
130    }
131
132    pub fn validate_run_scope(&self) -> std::result::Result<(), McpBindingError> {
133        if self.is_ready() {
134            Ok(())
135        } else {
136            Err(McpBindingError::ClientNotReady)
137        }
138    }
139
140    pub fn contains_public_tool_name(&self, full_name: &str) -> bool {
141        self.tools
142            .iter()
143            .any(|tool| full_name == format!("mcp__{}__{}", self.server_name, tool.name))
144    }
145
146    pub(crate) fn projected_tools(self: &Arc<Self>) -> Vec<Arc<dyn Tool>> {
147        self.tools
148            .iter()
149            .enumerate()
150            .map(|(tool_index, tool)| {
151                Arc::new(ProjectedMcpTool {
152                    full_name: format!("mcp__{}__{}", self.server_name, tool.name).into_boxed_str(),
153                    tool_index,
154                    binding: Arc::clone(self),
155                }) as Arc<dyn Tool>
156            })
157            .collect()
158    }
159
160    async fn call_tool(
161        &self,
162        tool_name: &str,
163        arguments: Option<serde_json::Value>,
164    ) -> Result<super::protocol::CallToolResult> {
165        if self
166            .tools
167            .binary_search_by(|tool| tool.name.as_str().cmp(tool_name))
168            .is_err()
169        {
170            return Err(anyhow!(
171                "MCP tool '{}' is not present in the frozen server binding",
172                tool_name
173            ));
174        }
175        self.client.call_tool(tool_name, arguments).await
176    }
177}
178
179impl fmt::Debug for McpBinding {
180    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
181        formatter
182            .debug_struct("McpBinding")
183            .field("server_name", &self.server_name)
184            .field("tool_count", &self.tools.len())
185            .field("ready", &self.is_ready())
186            .finish_non_exhaustive()
187    }
188}
189
190struct ProjectedMcpTool {
191    full_name: Box<str>,
192    tool_index: usize,
193    binding: Arc<McpBinding>,
194}
195
196impl ProjectedMcpTool {
197    fn tool(&self) -> &McpTool {
198        &self.binding.tools[self.tool_index]
199    }
200}
201
202#[async_trait]
203impl Tool for ProjectedMcpTool {
204    fn name(&self) -> &str {
205        &self.full_name
206    }
207
208    fn description(&self) -> &str {
209        self.tool().description.as_deref().unwrap_or("MCP tool")
210    }
211
212    fn parameters(&self) -> serde_json::Value {
213        self.tool().input_schema.clone()
214    }
215
216    fn requires_confirmation(&self, _args: &serde_json::Value) -> bool {
217        annotation_requires_confirmation(self.tool())
218    }
219
220    async fn execute(&self, args: &serde_json::Value, context: &ToolContext) -> Result<ToolOutput> {
221        let result = self
222            .binding
223            .call_tool(&self.tool().name, Some(args.clone()))
224            .await;
225        match result {
226            Ok(result) => project_tool_result(&self.full_name, &result, context).await,
227            Err(error) => Ok(ToolOutput::error(format!("MCP tool error: {error}"))),
228        }
229    }
230}
231
232fn validate_name(field: &'static str, value: &str) -> std::result::Result<(), McpBindingError> {
233    if value.is_empty()
234        || value.trim() != value
235        || value.len() > MAX_MCP_NAME_BYTES
236        || value.chars().any(char::is_control)
237    {
238        return Err(McpBindingError::InvalidName { field });
239    }
240    Ok(())
241}
242
243#[cfg(test)]
244mod tests {
245    use std::path::PathBuf;
246
247    use super::*;
248    use crate::mcp::test_support::{mcp_tool, ready_binding, RecordingMcpTransport};
249    use crate::mcp::transport::McpTransport;
250    use crate::mcp::McpProjectionAdapter;
251
252    #[tokio::test]
253    async fn binding_rejects_clients_that_are_not_ready_or_have_another_identity() {
254        let transport = RecordingMcpTransport::new("uninitialized", Vec::new());
255        let client = Arc::new(McpClient::new(
256            "catalog".to_string(),
257            Arc::clone(&transport) as Arc<dyn McpTransport>,
258        ));
259        assert_eq!(
260            McpBinding::new("catalog", Arc::clone(&client), Vec::new()).unwrap_err(),
261            McpBindingError::ClientNotReady
262        );
263
264        client.initialize().await.unwrap();
265        assert_eq!(
266            McpBinding::new("another", Arc::clone(&client), Vec::new()).unwrap_err(),
267            McpBindingError::ClientNameMismatch
268        );
269
270        transport.disconnect();
271        assert_eq!(
272            McpBinding::new("catalog", client, Vec::new()).unwrap_err(),
273            McpBindingError::ClientNotReady
274        );
275    }
276
277    #[tokio::test]
278    async fn binding_canonicalizes_tools_and_rejects_duplicate_or_oversized_catalogs() {
279        let (_, _transport, client) = ready_binding("catalog", "one", Vec::new()).await;
280        let binding = McpBinding::new(
281            "catalog",
282            Arc::clone(&client),
283            [mcp_tool("zeta", "last"), mcp_tool("alpha", "first")],
284        )
285        .unwrap();
286        assert_eq!(
287            binding
288                .tools()
289                .iter()
290                .map(|tool| tool.name.as_str())
291                .collect::<Vec<_>>(),
292            ["alpha", "zeta"]
293        );
294
295        assert_eq!(
296            McpBinding::new(
297                "catalog",
298                Arc::clone(&client),
299                [mcp_tool("same", "one"), mcp_tool("same", "two")],
300            )
301            .unwrap_err(),
302            McpBindingError::DuplicateToolName {
303                name: "same".to_string()
304            }
305        );
306
307        let too_many = (0..=MAX_MCP_BINDING_TOOLS)
308            .map(|index| mcp_tool(&format!("tool-{index:04}"), "bounded"))
309            .collect::<Vec<_>>();
310        assert_eq!(
311            McpBinding::new("catalog", Arc::clone(&client), too_many).unwrap_err(),
312            McpBindingError::BoundExceeded {
313                field: "tool count",
314                max: MAX_MCP_BINDING_TOOLS,
315            }
316        );
317
318        let oversized = mcp_tool("oversized", &"x".repeat(MAX_MCP_BINDING_DEFINITION_BYTES));
319        assert_eq!(
320            McpBinding::new("catalog", client, [oversized]).unwrap_err(),
321            McpBindingError::BoundExceeded {
322                field: "tool definition bytes",
323                max: MAX_MCP_BINDING_DEFINITION_BYTES,
324            }
325        );
326    }
327
328    #[tokio::test]
329    async fn projected_wrapper_calls_the_raw_tool_on_the_exact_client() {
330        let (binding, transport, _) = ready_binding(
331            "catalog",
332            "generation-one",
333            vec![mcp_tool("lookup", "generation-one")],
334        )
335        .await;
336        let wrappers = binding.projected_tools();
337        assert_eq!(wrappers.len(), 1);
338        assert_eq!(wrappers[0].name(), "mcp__catalog__lookup");
339        assert_eq!(wrappers[0].description(), "generation-one");
340
341        let arguments = serde_json::json!({"generation": "one"});
342        let output = wrappers[0]
343            .execute(&arguments, &ToolContext::new(PathBuf::from("/tmp")))
344            .await
345            .unwrap();
346        assert!(output.success);
347        assert_eq!(output.content, "generation-one");
348        assert_eq!(
349            transport.calls(),
350            [crate::mcp::test_support::RecordedMcpCall {
351                name: "lookup".to_string(),
352                arguments: Some(arguments),
353            }]
354        );
355    }
356
357    #[test]
358    fn projected_mcp_types_are_send_and_sync() {
359        fn assert_send_sync<T: Send + Sync>() {}
360
361        assert_send_sync::<McpBinding>();
362        assert_send_sync::<McpProjectionAdapter>();
363    }
364}