mcp_utils/client/
mcp_snapshot.rs

1use super::{
2    McpClient, McpError, ToolCatalog, ToolRoute,
3    naming::{create_namespaced_tool_name, split_on_server_name},
4};
5use llm::ToolDefinition;
6use rmcp::{RoleClient, model::CallToolRequestParams, service::RunningService};
7use serde_json::{Map, Value};
8use std::{collections::HashMap, fmt, sync::Arc};
9
10#[derive(Clone, Default)]
11pub struct McpSnapshot {
12    catalog: Arc<ToolCatalog>,
13    clients: Arc<HashMap<String, Arc<RunningService<RoleClient, McpClient>>>>,
14}
15
16impl fmt::Debug for McpSnapshot {
17    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
18        formatter
19            .debug_struct("McpSnapshot")
20            .field("catalog", &self.catalog)
21            .field("connected_servers", &self.clients.keys().collect::<Vec<_>>())
22            .finish()
23    }
24}
25
26impl McpSnapshot {
27    pub fn new(
28        catalog: Arc<ToolCatalog>,
29        clients: Arc<HashMap<String, Arc<RunningService<RoleClient, McpClient>>>>,
30    ) -> Self {
31        Self { catalog, clients }
32    }
33
34    pub fn catalog(&self) -> &Arc<ToolCatalog> {
35        &self.catalog
36    }
37
38    pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
39        self.catalog.tools().model_visible.into_iter().map(|tool| tool.definition().clone()).collect()
40    }
41
42    pub fn model_instructions(&self) -> std::collections::BTreeMap<String, String> {
43        self.catalog.model_instructions()
44    }
45
46    pub fn server_statuses(&self) -> Vec<crate::status::McpServerStatusEntry> {
47        self.catalog.server_statuses()
48    }
49
50    pub fn resolve(
51        &self,
52        route: ToolRoute,
53        arguments: Map<String, Value>,
54    ) -> super::Result<(Arc<RunningService<RoleClient, McpClient>>, CallToolRequestParams)> {
55        if let ToolRoute::ModelVisible { namespaced_name } = &route {
56            split_on_server_name(namespaced_name)
57                .ok_or_else(|| McpError::InvalidToolNameFormat(namespaced_name.clone()))?;
58        }
59        if !self.catalog.route_permitted(&route) {
60            let (tool_name, namespaced_name) = match &route {
61                ToolRoute::ModelVisible { namespaced_name } => (namespaced_name.clone(), namespaced_name.clone()),
62                ToolRoute::Deferred { server, tool } => (tool.clone(), create_namespaced_tool_name(server, tool)),
63            };
64            if matches!(route, ToolRoute::Deferred { .. })
65                && self.catalog.route_permitted(&ToolRoute::ModelVisible { namespaced_name: namespaced_name.clone() })
66            {
67                return Err(McpError::DirectToolRequiresDirectRoute { tool_name, direct_name: namespaced_name });
68            }
69            return Err(McpError::ToolNotFound(namespaced_name));
70        }
71        let (server, tool) = match route {
72            ToolRoute::ModelVisible { namespaced_name } => {
73                let (server, tool) =
74                    split_on_server_name(&namespaced_name).expect("model-visible route was validated above");
75                (server.to_string(), tool.to_string())
76            }
77            ToolRoute::Deferred { server, tool } => (server, tool),
78        };
79        let client = self.clients.get(&server).cloned().ok_or_else(|| McpError::ServerNotFound(server.clone()))?;
80        Ok((client, CallToolRequestParams::new(tool).with_arguments(arguments)))
81    }
82
83    pub fn clients_with_prompts(&self) -> Vec<(String, Arc<RunningService<RoleClient, McpClient>>)> {
84        self.clients
85            .iter()
86            .filter(|(_, client)| client.peer_info().is_some_and(|info| info.capabilities.prompts.is_some()))
87            .map(|(name, client)| (name.clone(), Arc::clone(client)))
88            .collect()
89    }
90
91    pub fn client_for_prompt(
92        &self,
93        namespaced_name: &str,
94    ) -> super::Result<(String, Arc<RunningService<RoleClient, McpClient>>)> {
95        let (server, prompt) = split_on_server_name(namespaced_name)
96            .ok_or_else(|| McpError::InvalidToolNameFormat(namespaced_name.to_string()))?;
97        let client = self.clients.get(server).cloned().ok_or_else(|| McpError::ServerNotFound(server.to_string()))?;
98        Ok((prompt.to_string(), client))
99    }
100}
mcp_utils/client/mcp_snapshot.rs

mcp_utils/client/
mcp_snapshot.rs
