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        if context.is_cancelled() {
222            return Ok(ToolOutput::error(format!(
223                "MCP tool '{}' cancelled by caller",
224                self.full_name
225            )));
226        }
227
228        let cancellation = context.cancellation_token();
229        let call = self
230            .binding
231            .call_tool(&self.tool().name, Some(args.clone()));
232        let result = tokio::select! {
233            _ = cancellation.cancelled() => {
234                return Ok(ToolOutput::error(format!(
235                    "MCP tool '{}' cancelled by caller",
236                    self.full_name
237                )));
238            }
239            result = call => result,
240        };
241        match result {
242            Ok(result) => project_tool_result(&self.full_name, &result, context).await,
243            Err(error) => Ok(ToolOutput::error(format!("MCP tool error: {error}"))),
244        }
245    }
246}
247
248fn validate_name(field: &'static str, value: &str) -> std::result::Result<(), McpBindingError> {
249    if value.is_empty()
250        || value.trim() != value
251        || value.len() > MAX_MCP_NAME_BYTES
252        || value.chars().any(char::is_control)
253    {
254        return Err(McpBindingError::InvalidName { field });
255    }
256    Ok(())
257}
258
259#[cfg(test)]
260mod tests {
261    use std::path::PathBuf;
262
263    use super::*;
264    use crate::mcp::test_support::{mcp_tool, ready_binding, RecordingMcpTransport};
265    use crate::mcp::transport::McpTransport;
266    use crate::mcp::McpProjectionAdapter;
267    use async_trait::async_trait;
268
269    #[tokio::test]
270    async fn binding_rejects_clients_that_are_not_ready_or_have_another_identity() {
271        let transport = RecordingMcpTransport::new("uninitialized", Vec::new());
272        let client = Arc::new(McpClient::new(
273            "catalog".to_string(),
274            Arc::clone(&transport) as Arc<dyn McpTransport>,
275        ));
276        assert_eq!(
277            McpBinding::new("catalog", Arc::clone(&client), Vec::new()).unwrap_err(),
278            McpBindingError::ClientNotReady
279        );
280
281        client.initialize().await.unwrap();
282        assert_eq!(
283            McpBinding::new("another", Arc::clone(&client), Vec::new()).unwrap_err(),
284            McpBindingError::ClientNameMismatch
285        );
286
287        transport.disconnect();
288        assert_eq!(
289            McpBinding::new("catalog", client, Vec::new()).unwrap_err(),
290            McpBindingError::ClientNotReady
291        );
292    }
293
294    #[tokio::test]
295    async fn binding_canonicalizes_tools_and_rejects_duplicate_or_oversized_catalogs() {
296        let (_, _transport, client) = ready_binding("catalog", "one", Vec::new()).await;
297        let binding = McpBinding::new(
298            "catalog",
299            Arc::clone(&client),
300            [mcp_tool("zeta", "last"), mcp_tool("alpha", "first")],
301        )
302        .unwrap();
303        assert_eq!(
304            binding
305                .tools()
306                .iter()
307                .map(|tool| tool.name.as_str())
308                .collect::<Vec<_>>(),
309            ["alpha", "zeta"]
310        );
311
312        assert_eq!(
313            McpBinding::new(
314                "catalog",
315                Arc::clone(&client),
316                [mcp_tool("same", "one"), mcp_tool("same", "two")],
317            )
318            .unwrap_err(),
319            McpBindingError::DuplicateToolName {
320                name: "same".to_string()
321            }
322        );
323
324        let too_many = (0..=MAX_MCP_BINDING_TOOLS)
325            .map(|index| mcp_tool(&format!("tool-{index:04}"), "bounded"))
326            .collect::<Vec<_>>();
327        assert_eq!(
328            McpBinding::new("catalog", Arc::clone(&client), too_many).unwrap_err(),
329            McpBindingError::BoundExceeded {
330                field: "tool count",
331                max: MAX_MCP_BINDING_TOOLS,
332            }
333        );
334
335        let oversized = mcp_tool("oversized", &"x".repeat(MAX_MCP_BINDING_DEFINITION_BYTES));
336        assert_eq!(
337            McpBinding::new("catalog", client, [oversized]).unwrap_err(),
338            McpBindingError::BoundExceeded {
339                field: "tool definition bytes",
340                max: MAX_MCP_BINDING_DEFINITION_BYTES,
341            }
342        );
343    }
344
345    #[tokio::test]
346    async fn projected_wrapper_calls_the_raw_tool_on_the_exact_client() {
347        let (binding, transport, _) = ready_binding(
348            "catalog",
349            "generation-one",
350            vec![mcp_tool("lookup", "generation-one")],
351        )
352        .await;
353        let wrappers = binding.projected_tools();
354        assert_eq!(wrappers.len(), 1);
355        assert_eq!(wrappers[0].name(), "mcp__catalog__lookup");
356        assert_eq!(wrappers[0].description(), "generation-one");
357
358        let arguments = serde_json::json!({"generation": "one"});
359        let output = wrappers[0]
360            .execute(&arguments, &ToolContext::new(PathBuf::from("/tmp")))
361            .await
362            .unwrap();
363        assert!(output.success);
364        assert_eq!(output.content, "generation-one");
365        assert_eq!(
366            transport.calls(),
367            [crate::mcp::test_support::RecordedMcpCall {
368                name: "lookup".to_string(),
369                arguments: Some(arguments),
370            }]
371        );
372    }
373
374    #[tokio::test]
375    async fn projected_wrapper_returns_promptly_when_caller_cancels() {
376        let (binding, transport, _) = ready_binding(
377            "catalog",
378            "generation-one",
379            vec![mcp_tool("lookup", "generation-one")],
380        )
381        .await;
382        let wrappers = binding.projected_tools();
383        let cancel = tokio_util::sync::CancellationToken::new();
384        cancel.cancel();
385        let output = wrappers[0]
386            .execute(
387                &serde_json::json!({}),
388                &ToolContext::new(PathBuf::from("/tmp")).with_cancellation(cancel),
389            )
390            .await
391            .unwrap();
392        assert!(!output.success);
393        assert!(output.content.contains("cancelled by caller"));
394        assert!(
395            transport.calls().is_empty(),
396            "cancelled calls must not reach the MCP client"
397        );
398    }
399
400    #[test]
401    fn projected_mcp_types_are_send_and_sync() {
402        fn assert_send_sync<T: Send + Sync>() {}
403
404        assert_send_sync::<McpBinding>();
405        assert_send_sync::<McpProjectionAdapter>();
406    }
407
408    #[tokio::test]
409    async fn denied_governed_mcp_call_does_not_reach_the_server() {
410        let (binding, transport, _) = ready_binding(
411            "catalog",
412            "generation-one",
413            vec![mcp_tool("lookup", "lookup")],
414        )
415        .await;
416        let wrappers = binding.projected_tools();
417        let tool_name = wrappers[0].name().to_string();
418        assert_eq!(tool_name, "mcp__catalog__lookup");
419
420        let directory = tempfile::tempdir().unwrap();
421        let executor = Arc::new(crate::tools::ToolExecutor::new(
422            directory.path().to_string_lossy().to_string(),
423        ));
424        executor.register_dynamic_tool(Arc::clone(&wrappers[0]));
425        let session_id = "mcp-deny";
426        let context = ToolContext::new(directory.path().to_path_buf()).with_session_id(session_id);
427        let gate = Arc::new(DenyThenAllow {
428            allow: std::sync::atomic::AtomicBool::new(false),
429        });
430        let agent = crate::agent::AgentLoop::new(
431            Arc::new(IdleModel),
432            executor,
433            context.clone(),
434            crate::agent::AgentConfig {
435                permission_checker: Some(
436                    Arc::clone(&gate) as Arc<dyn crate::permissions::PermissionChecker>
437                ),
438                ..crate::agent::AgentConfig::default()
439            },
440        );
441
442        let denied = agent
443            .invoke_host_tool(
444                crate::tools::ToolInvocation::host_governed(
445                    "mcp-deny-1",
446                    tool_name.clone(),
447                    serde_json::json!({"generation": "secret"}),
448                ),
449                session_id,
450                &None,
451                &tokio_util::sync::CancellationToken::new(),
452                &context,
453            )
454            .await;
455        assert_ne!(denied.exit_code, 0, "{}", denied.output);
456        assert!(
457            denied.output.contains("Permission denied"),
458            "{}",
459            denied.output
460        );
461        assert!(
462            transport.calls().is_empty(),
463            "denied MCP call mutated the server: {:?}",
464            transport.calls()
465        );
466
467        gate.allow.store(true, std::sync::atomic::Ordering::SeqCst);
468        let allowed = agent
469            .invoke_host_tool(
470                crate::tools::ToolInvocation::host_governed(
471                    "mcp-allow-1",
472                    tool_name,
473                    serde_json::json!({"generation": "visible"}),
474                ),
475                session_id,
476                &None,
477                &tokio_util::sync::CancellationToken::new(),
478                &context,
479            )
480            .await;
481        assert_eq!(allowed.exit_code, 0, "{}", allowed.output);
482        assert_eq!(transport.calls().len(), 1);
483        assert_eq!(transport.calls()[0].name, "lookup");
484    }
485
486    struct DenyThenAllow {
487        allow: std::sync::atomic::AtomicBool,
488    }
489
490    impl crate::permissions::PermissionChecker for DenyThenAllow {
491        fn check(
492            &self,
493            tool_name: &str,
494            _args: &serde_json::Value,
495        ) -> crate::permissions::PermissionDecision {
496            assert_eq!(tool_name, "mcp__catalog__lookup");
497            if self.allow.load(std::sync::atomic::Ordering::SeqCst) {
498                crate::permissions::PermissionDecision::Allow
499            } else {
500                crate::permissions::PermissionDecision::Deny
501            }
502        }
503    }
504
505    struct IdleModel;
506
507    #[async_trait]
508    impl crate::llm::LlmClient for IdleModel {
509        async fn complete(
510            &self,
511            _messages: &[crate::llm::Message],
512            _system: Option<&str>,
513            _tools: &[crate::llm::ToolDefinition],
514        ) -> anyhow::Result<crate::llm::LlmResponse> {
515            anyhow::bail!("MCP policy test must not call the model")
516        }
517
518        async fn complete_streaming(
519            &self,
520            _messages: &[crate::llm::Message],
521            _system: Option<&str>,
522            _tools: &[crate::llm::ToolDefinition],
523            _cancel_token: tokio_util::sync::CancellationToken,
524        ) -> anyhow::Result<tokio::sync::mpsc::Receiver<crate::llm::StreamEvent>> {
525            anyhow::bail!("MCP policy test must not stream the model")
526        }
527    }
528}
a3s_code_core/mcp/binding.rs

a3s_code_core/mcp/
binding.rs
