1use super::{ToolExposure, connection::Tool, naming::create_namespaced_tool_name, tool_filter::ToolFilter};
2use crate::status::{McpServerAuthCapability, McpServerStatus, McpServerStatusEntry};
3use llm::ToolDefinition;
4use std::collections::BTreeMap;
5
6pub const PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME: &str = "progressive-discovery";
7
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct ServerDescription {
10 pub name: String,
11 pub description: String,
12}
13
14#[derive(Debug, Clone, Default, PartialEq)]
15pub struct ToolCatalog {
16 servers: Vec<ServerCatalogEntry>,
17 progressive_discovery_instructions: Option<String>,
18}
19
20#[derive(Debug, Clone, PartialEq)]
21pub struct ServerCatalogEntry {
22 name: String,
23 description: String,
24 instructions: Option<String>,
25 status: McpServerStatus,
26 auth_capability: McpServerAuthCapability,
27 exposure: ToolExposure,
28 tools: Vec<CatalogTool>,
29}
30
31#[derive(Debug, Clone, Default, PartialEq)]
32pub struct CatalogTools<'a> {
33 pub model_visible: Vec<&'a CatalogTool>,
34 pub deferred: Vec<&'a CatalogTool>,
35}
36
37#[derive(Debug, Clone, PartialEq)]
38pub struct CatalogTool {
39 namespaced_name: String,
40 local_name: String,
41 definition: ToolDefinition,
42 exposure: ToolExposureKind,
43 allowed: bool,
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum ToolExposureKind {
48 ModelVisible,
49 Deferred,
50}
51
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub enum ToolRoute {
54 ModelVisible { namespaced_name: String },
55 Deferred { server: String, tool: String },
56}
57
58impl ToolCatalog {
59 pub fn new() -> Self {
60 Self::default()
61 }
62 pub fn servers(&self) -> &[ServerCatalogEntry] {
63 &self.servers
64 }
65
66 pub fn server(&self, name: &str) -> Option<&ServerCatalogEntry> {
67 self.servers.iter().find(|server| server.name == name)
68 }
69
70 pub fn tool(&self, namespaced_name: &str) -> Option<&CatalogTool> {
71 self.servers.iter().flat_map(|server| &server.tools).find(|tool| tool.namespaced_name == namespaced_name)
72 }
73
74 pub fn tools(&self) -> CatalogTools<'_> {
75 CatalogTools::from_tools(
76 self.servers.iter().filter(|server| server.is_connected()).flat_map(|server| server.tools.iter()),
77 )
78 }
79
80 pub fn tools_for(&self, server: &str) -> Option<CatalogTools<'_>> {
81 self.server(server).map(|entry| {
82 if entry.is_connected() { CatalogTools::from_tools(entry.tools.iter()) } else { CatalogTools::default() }
83 })
84 }
85
86 pub fn discoverable_deferred_servers(&self) -> Vec<ServerDescription> {
87 self.servers
88 .iter()
89 .filter(|server| {
90 server.is_connected()
91 && server.tools.iter().any(|tool| tool.allowed && tool.exposure == ToolExposureKind::Deferred)
92 })
93 .map(|server| ServerDescription { name: server.name.clone(), description: server.description.clone() })
94 .collect()
95 }
96
97 pub fn model_instructions(&self) -> BTreeMap<String, String> {
98 let mut instructions = self
99 .servers
100 .iter()
101 .filter(|server| server.is_connected())
102 .filter(|server| {
103 server.tools.iter().any(|tool| tool.allowed && tool.exposure == ToolExposureKind::ModelVisible)
104 })
105 .filter_map(|server| server.instructions.as_ref().map(|body| (server.name.clone(), body.clone())))
106 .collect::<BTreeMap<_, _>>();
107 if !self.discoverable_deferred_servers().is_empty()
108 && let Some(body) = &self.progressive_discovery_instructions
109 {
110 instructions.insert(PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME.to_string(), body.clone());
111 }
112 instructions
113 }
114
115 pub fn route_permitted(&self, route: &ToolRoute) -> bool {
116 let (namespaced_name, exposure) = match route {
117 ToolRoute::ModelVisible { namespaced_name } => (namespaced_name.clone(), ToolExposureKind::ModelVisible),
118 ToolRoute::Deferred { server, tool } => {
119 (create_namespaced_tool_name(server, tool), ToolExposureKind::Deferred)
120 }
121 };
122 let Some(server) = self.servers.iter().find(|server| {
123 server.is_connected() && server.tools.iter().any(|tool| tool.namespaced_name == namespaced_name)
124 }) else {
125 return false;
126 };
127 let Some(tool) = server.tools.iter().find(|tool| tool.namespaced_name == namespaced_name) else { return false };
128 tool.allowed && tool.exposure == exposure
129 }
130
131 pub fn server_statuses(&self) -> Vec<McpServerStatusEntry> {
132 self.servers.iter().map(ServerCatalogEntry::status_entry).collect()
133 }
134
135 pub fn upsert_server(&mut self, entry: ServerCatalogEntry) {
136 if let Some(existing) = self.servers.iter_mut().find(|server| server.name == entry.name) {
137 *existing = entry;
138 } else {
139 self.servers.push(entry);
140 }
141 }
142
143 pub fn remove_server(&mut self, name: &str) -> Option<ServerCatalogEntry> {
144 self.servers.iter().position(|server| server.name == name).map(|index| self.servers.remove(index))
145 }
146
147 pub fn set_progressive_discovery_instructions(&mut self, instructions: Option<String>) {
148 self.progressive_discovery_instructions = instructions;
149 }
150}
151
152impl ServerCatalogEntry {
153 #[allow(clippy::too_many_arguments)]
154 pub fn new(
155 name: impl Into<String>,
156 description: impl Into<String>,
157 instructions: Option<String>,
158 status: McpServerStatus,
159 auth_capability: McpServerAuthCapability,
160 exposure: ToolExposure,
161 tools: &[rmcp::model::Tool],
162 filter: &ToolFilter,
163 ) -> Self {
164 let tools = tools.iter().map(Tool::from).collect::<Vec<_>>();
165 Self::from_tools(
166 name.into(),
167 description.into(),
168 instructions,
169 status,
170 auth_capability,
171 exposure,
172 &tools,
173 filter,
174 )
175 }
176 pub fn name(&self) -> &str {
177 &self.name
178 }
179 pub fn description(&self) -> &str {
180 &self.description
181 }
182 pub fn instructions(&self) -> Option<&str> {
183 self.instructions.as_deref()
184 }
185 pub fn status(&self) -> &McpServerStatus {
186 &self.status
187 }
188 pub fn auth_capability(&self) -> McpServerAuthCapability {
189 self.auth_capability
190 }
191 pub fn exposure(&self) -> &ToolExposure {
192 &self.exposure
193 }
194 pub fn tools(&self) -> &[CatalogTool] {
195 &self.tools
196 }
197 pub fn status_entry(&self) -> McpServerStatusEntry {
198 McpServerStatusEntry::new(&self.name, self.status.clone())
199 .with_auth_capability(self.auth_capability)
200 .with_deferred_tools(self.exposure.has_deferred_tools())
201 }
202 pub(crate) fn pending(name: impl Into<String>, exposure: ToolExposure) -> Self {
203 let name = name.into();
204 Self {
205 description: name.clone(),
206 name,
207 instructions: None,
208 status: McpServerStatus::Connecting,
209 auth_capability: McpServerAuthCapability::Unavailable,
210 exposure,
211 tools: Vec::new(),
212 }
213 }
214 #[allow(clippy::too_many_arguments)]
215 pub(crate) fn from_tools(
216 name: String,
217 description: String,
218 instructions: Option<String>,
219 status: McpServerStatus,
220 auth_capability: McpServerAuthCapability,
221 exposure: ToolExposure,
222 tools: &[Tool],
223 filter: &ToolFilter,
224 ) -> Self {
225 let catalog_tools = tools
226 .iter()
227 .map(|tool| {
228 let definition = ToolDefinition::new(
229 create_namespaced_tool_name(&name, &tool.name),
230 tool.description.clone(),
231 tool.parameters.clone(),
232 )
233 .with_server(name.clone())
234 .with_annotations(tool.annotations.clone());
235 let exposure_kind = if exposure.is_model_visible_tool(&tool.name) {
236 ToolExposureKind::ModelVisible
237 } else {
238 ToolExposureKind::Deferred
239 };
240 CatalogTool {
241 namespaced_name: definition.name.clone(),
242 local_name: tool.name.clone(),
243 allowed: filter.is_tool_allowed(&definition),
244 definition,
245 exposure: exposure_kind,
246 }
247 })
248 .collect();
249 Self { name, description, instructions, status, auth_capability, exposure, tools: catalog_tools }
250 }
251 pub(crate) fn with_status(&self, status: McpServerStatus, auth_capability: McpServerAuthCapability) -> Self {
252 let mut next = self.clone();
253 next.status = status;
254 next.auth_capability = auth_capability;
255 if !next.is_connected() {
256 next.tools.clear();
257 next.instructions = None;
258 }
259 next
260 }
261 fn is_connected(&self) -> bool {
262 matches!(self.status, McpServerStatus::Connected { .. })
263 }
264}
265
266impl<'a> CatalogTools<'a> {
267 fn from_tools(tools: impl Iterator<Item = &'a CatalogTool>) -> Self {
268 let mut partitioned = Self::default();
269 for tool in tools.filter(|tool| tool.allowed) {
270 match tool.exposure {
271 ToolExposureKind::ModelVisible => partitioned.model_visible.push(tool),
272 ToolExposureKind::Deferred => partitioned.deferred.push(tool),
273 }
274 }
275 partitioned
276 }
277}
278
279impl CatalogTool {
280 pub fn namespaced_name(&self) -> &str {
281 &self.namespaced_name
282 }
283 pub fn local_name(&self) -> &str {
284 &self.local_name
285 }
286 pub fn definition(&self) -> &ToolDefinition {
287 &self.definition
288 }
289 pub fn exposure(&self) -> ToolExposureKind {
290 self.exposure
291 }
292 pub fn allowed(&self) -> bool {
293 self.allowed
294 }
295}