mcp/
protocol.rs

1//! Tool protocol abstraction layer.
2//!
3//! This module provides a flexible abstraction for connecting agents to various tool protocols.
4//! It supports multiple standards including MCP (Model Context Protocol), custom function calling,
5//! Memory persistence, and allows users to implement their own tool communication mechanisms.
6//!
7//! # Architecture
8//!
9//! **Single Protocol** (traditional):
10//! ```text
11//! Agent → ToolRegistry → ToolProtocol → Single Tool Source
12//! ```
13//!
14//! **Multi-Protocol** (new in 0.5.0):
15//! ```text
16//! Agent → ToolRegistry → [Protocol1, Protocol2, Protocol3]
17//!         (routing map)     ↓          ↓          ↓
18//!                        Local      YouTube    GitHub
19//!                        Tools      Server     Server
20//! ```
21//!
22//! # Key Components
23//!
24//! - **ToolProtocol trait**: Define how tools are executed, discovered, and described
25//! - **ToolRegistry**: Single or multi-protocol tool aggregation with transparent routing
26//! - **ToolMetadata**: Tool identity, description, parameters
27//! - **ToolParameter**: Type-safe parameter definitions with validation
28//! - **ToolResult**: Structured tool execution results
29//! - **Tool**: Runtime tool instance bound to a protocol
30//!
31//! # Single Protocol Example
32//!
33//! ```rust,ignore
34//! use std::sync::Arc;
35//! use mcp::protocol::{ToolMetadata, ToolProtocol, ToolRegistry, ToolResult};
36//!
37//! struct MyProtocol;
38//!
39//! #[async_trait::async_trait]
40//! impl ToolProtocol for MyProtocol {
41//!     async fn execute(
42//!         &self,
43//!         _tool_name: &str,
44//!         _parameters: serde_json::Value,
45//!     ) -> Result<ToolResult, Box<dyn std::error::Error + Send + Sync>> {
46//!         Ok(ToolResult::success(serde_json::json!({"ok": true})))
47//!     }
48//!     async fn list_tools(&self) -> Result<Vec<ToolMetadata>, Box<dyn std::error::Error + Send + Sync>> {
49//!         Ok(vec![ToolMetadata::new("demo", "Demo tool")])
50//!     }
51//!     async fn get_tool_metadata(&self, _: &str) -> Result<ToolMetadata, Box<dyn std::error::Error + Send + Sync>> {
52//!         Ok(ToolMetadata::new("demo", "Demo tool"))
53//!     }
54//!     fn protocol_name(&self) -> &str { "demo" }
55//! }
56//!
57//! # #[tokio::main]
58//! # async fn main() {
59//! let protocol = Arc::new(MyProtocol);
60//! let mut registry = ToolRegistry::new(protocol);
61//! let _ = registry.discover_tools_from_primary().await;
62//! # }
63//! ```
64//!
65//! # Multi-Protocol Example
66//!
67//! ```rust,ignore
68//! use mcp::client::McpClientProtocol;
69//! use mcp::protocol::ToolRegistry;
70//!
71//! # #[tokio::main]
72//! # async fn main() {
73//! let mut registry = ToolRegistry::empty();
74//! let _ = registry
75//!     .add_protocol("remote", std::sync::Arc::new(McpClientProtocol::new("http://localhost:8080".to_string())))
76//!     .await;
77//! # }
78//! ```
79
80use crate::resources::ResourceMetadata;
81use async_trait::async_trait;
82use serde::{Deserialize, Serialize};
83use std::collections::HashMap;
84use std::error::Error;
85use std::fmt;
86use std::sync::Arc;
87
88/// Provider-agnostic tool definition suitable for LLM native tool calling.
89#[derive(Debug, Clone, Serialize, Deserialize)]
90pub struct ToolDefinition {
91    /// Tool name as it will appear in the API `tools` array.
92    pub name: String,
93    /// Human-readable description surfaced to the model.
94    pub description: String,
95    /// JSON Schema object describing accepted parameters.
96    pub parameters_schema: serde_json::Value,
97}
98
99/// Represents the result of a tool execution
100#[derive(Debug, Clone, Serialize, Deserialize)]
101pub struct ToolResult {
102    /// Whether the tool execution was successful
103    pub success: bool,
104    /// The output data from the tool
105    pub output: serde_json::Value,
106    /// Optional error message if execution failed
107    pub error: Option<String>,
108    /// Metadata about the execution (timing, cost, etc.)
109    pub metadata: HashMap<String, serde_json::Value>,
110}
111
112impl ToolResult {
113    /// Convenience constructor for successful tool execution.
114    pub fn success(output: serde_json::Value) -> Self {
115        Self {
116            success: true,
117            output,
118            error: None,
119            metadata: HashMap::new(),
120        }
121    }
122
123    /// Convenience constructor for failed tool execution.
124    pub fn failure(error: String) -> Self {
125        Self {
126            success: false,
127            output: serde_json::Value::Null,
128            error: Some(error),
129            metadata: HashMap::new(),
130        }
131    }
132
133    /// Attach protocol or application specific metadata to the result.
134    pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
135        self.metadata.insert(key.into(), value);
136        self
137    }
138}
139
140/// Defines the type of a tool parameter
141#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
142#[serde(rename_all = "lowercase")]
143pub enum ToolParameterType {
144    /// UTF-8 string value.
145    String,
146    /// Floating-point number.
147    Number,
148    /// Integer number.
149    Integer,
150    /// Boolean value.
151    Boolean,
152    /// JSON array.
153    Array,
154    /// JSON object.
155    Object,
156}
157
158/// Defines a parameter for a tool
159#[derive(Debug, Clone, Serialize, Deserialize)]
160pub struct ToolParameter {
161    /// Parameter name.
162    pub name: String,
163    /// Parameter value type.
164    #[serde(rename = "type")]
165    pub param_type: ToolParameterType,
166    /// Human-readable description.
167    pub description: Option<String>,
168    /// Whether the caller must provide this parameter.
169    pub required: bool,
170    /// Optional default value.
171    pub default: Option<serde_json::Value>,
172    /// For array types, specifies the type of items
173    pub items: Option<Box<ToolParameterType>>,
174    /// For object types, specifies nested properties
175    pub properties: Option<HashMap<String, ToolParameter>>,
176}
177
178impl ToolParameter {
179    /// Define a new tool parameter with the provided name and type.
180    pub fn new(name: impl Into<String>, param_type: ToolParameterType) -> Self {
181        Self {
182            name: name.into(),
183            param_type,
184            description: None,
185            required: false,
186            default: None,
187            items: None,
188            properties: None,
189        }
190    }
191
192    /// Add a human readable description that will surface in generated schemas.
193    pub fn with_description(mut self, description: impl Into<String>) -> Self {
194        self.description = Some(description.into());
195        self
196    }
197
198    /// Mark the argument as required.
199    pub fn required(mut self) -> Self {
200        self.required = true;
201        self
202    }
203
204    /// Provide a default value that will be used when the LLM omits the parameter.
205    pub fn with_default(mut self, default: serde_json::Value) -> Self {
206        self.default = Some(default);
207        self
208    }
209
210    /// For array parameters, declare the type of the contained items.
211    pub fn with_items(mut self, item_type: ToolParameterType) -> Self {
212        self.items = Some(Box::new(item_type));
213        self
214    }
215
216    /// For object parameters, describe the nested properties.
217    pub fn with_properties(mut self, properties: HashMap<String, ToolParameter>) -> Self {
218        self.properties = Some(properties);
219        self
220    }
221
222    /// Converts this parameter to a JSON Schema snippet suitable for native tool-calling.
223    ///
224    /// The generated schema follows [JSON Schema draft-07] conventions accepted by all major
225    /// providers (OpenAI, Anthropic, Grok, Gemini).
226    ///
227    /// | `ToolParameterType` | Schema produced |
228    /// |---|---|
229    /// | `String` | `{"type":"string"}` |
230    /// | `Number` | `{"type":"number"}` |
231    /// | `Integer` | `{"type":"integer"}` |
232    /// | `Boolean` | `{"type":"boolean"}` |
233    /// | `Array` | `{"type":"array","items":{...}}` |
234    /// | `Object` | `{"type":"object","properties":{...},"required":[...]}` |
235    ///
236    /// # Example
237    ///
238    /// ```rust
239    /// use cloudllm::tool_protocol::{ToolParameter, ToolParameterType};
240    ///
241    /// let param = ToolParameter::new("query", ToolParameterType::String)
242    ///     .with_description("Search query string")
243    ///     .required();
244    ///
245    /// let schema = param.to_json_schema();
246    /// assert_eq!(schema["type"], "string");
247    /// assert_eq!(schema["description"], "Search query string");
248    /// ```
249    pub fn to_json_schema(&self) -> serde_json::Value {
250        let mut schema = match &self.param_type {
251            ToolParameterType::String => serde_json::json!({"type": "string"}),
252            ToolParameterType::Number => serde_json::json!({"type": "number"}),
253            ToolParameterType::Integer => serde_json::json!({"type": "integer"}),
254            ToolParameterType::Boolean => serde_json::json!({"type": "boolean"}),
255            ToolParameterType::Array => {
256                let items_schema = self
257                    .items
258                    .as_ref()
259                    .map(|t| param_type_to_schema(t))
260                    .unwrap_or_else(|| serde_json::json!({}));
261                serde_json::json!({"type": "array", "items": items_schema})
262            }
263            ToolParameterType::Object => {
264                if let Some(props) = &self.properties {
265                    let properties: serde_json::Map<String, serde_json::Value> = props
266                        .iter()
267                        .map(|(k, v)| (k.clone(), v.to_json_schema()))
268                        .collect();
269                    let required: Vec<&str> = props
270                        .values()
271                        .filter(|p| p.required)
272                        .map(|p| p.name.as_str())
273                        .collect();
274                    serde_json::json!({
275                        "type": "object",
276                        "properties": properties,
277                        "required": required
278                    })
279                } else {
280                    serde_json::json!({"type": "object"})
281                }
282            }
283        };
284
285        // Attach description if present
286        if let Some(desc) = &self.description {
287            schema["description"] = serde_json::Value::String(desc.clone());
288        }
289
290        schema
291    }
292}
293
294/// Convert a bare [`ToolParameterType`] to a minimal JSON Schema value (no description).
295///
296/// Used internally when building array `items` schemas.
297fn param_type_to_schema(t: &ToolParameterType) -> serde_json::Value {
298    match t {
299        ToolParameterType::String => serde_json::json!({"type": "string"}),
300        ToolParameterType::Number => serde_json::json!({"type": "number"}),
301        ToolParameterType::Integer => serde_json::json!({"type": "integer"}),
302        ToolParameterType::Boolean => serde_json::json!({"type": "boolean"}),
303        ToolParameterType::Array => serde_json::json!({"type": "array"}),
304        ToolParameterType::Object => serde_json::json!({"type": "object"}),
305    }
306}
307
308/// Metadata about a tool
309#[derive(Debug, Clone, Serialize, Deserialize)]
310pub struct ToolMetadata {
311    /// Tool name.
312    pub name: String,
313    /// Human-readable description.
314    pub description: String,
315    /// Parameter schema entries.
316    pub parameters: Vec<ToolParameter>,
317    /// Additional metadata specific to the protocol
318    pub protocol_metadata: HashMap<String, serde_json::Value>,
319}
320
321impl ToolMetadata {
322    /// Create metadata with the supplied identifier and description.
323    pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
324        Self {
325            name: name.into(),
326            description: description.into(),
327            parameters: Vec::new(),
328            protocol_metadata: HashMap::new(),
329        }
330    }
331
332    /// Append a parameter definition to the tool metadata.
333    pub fn with_parameter(mut self, param: ToolParameter) -> Self {
334        self.parameters.push(param);
335        self
336    }
337
338    /// Add protocol specific metadata (e.g. MCP capability flags).
339    pub fn with_protocol_metadata(
340        mut self,
341        key: impl Into<String>,
342        value: serde_json::Value,
343    ) -> Self {
344        self.protocol_metadata.insert(key.into(), value);
345        self
346    }
347
348    /// Builds a [`ToolDefinition`] (JSON Schema) from this metadata.
349    ///
350    /// The resulting [`ToolDefinition`] is ready to be passed to
351    /// [`ClientWrapper::send_message`](crate::client_wrapper::ClientWrapper::send_message)
352    /// as part of the `tools` slice.  Parameters marked `required` are collected into the
353    /// JSON Schema `"required"` array automatically.
354    ///
355    /// # Example
356    ///
357    /// ```rust
358    /// use cloudllm::tool_protocol::{ToolMetadata, ToolParameter, ToolParameterType};
359    ///
360    /// let meta = ToolMetadata::new("calculator", "Evaluates a math expression")
361    ///     .with_parameter(
362    ///         ToolParameter::new("expression", ToolParameterType::String)
363    ///             .with_description("The expression to evaluate")
364    ///             .required(),
365    ///     );
366    ///
367    /// let def = meta.to_tool_definition();
368    /// assert_eq!(def.name, "calculator");
369    /// let schema = &def.parameters_schema;
370    /// assert_eq!(schema["type"], "object");
371    /// assert!(schema["properties"]["expression"].is_object());
372    /// ```
373    pub fn to_tool_definition(&self) -> ToolDefinition {
374        let mut properties = serde_json::Map::new();
375        let mut required: Vec<String> = Vec::new();
376
377        for param in &self.parameters {
378            properties.insert(param.name.clone(), param.to_json_schema());
379            if param.required {
380                required.push(param.name.clone());
381            }
382        }
383
384        let parameters_schema = serde_json::json!({
385            "type": "object",
386            "properties": properties,
387            "required": required
388        });
389
390        ToolDefinition {
391            name: self.name.clone(),
392            description: self.description.clone(),
393            parameters_schema,
394        }
395    }
396}
397
398/// Trait for implementing tool execution protocols
399#[async_trait]
400pub trait ToolProtocol: Send + Sync {
401    /// Execute a tool with the given parameters
402    async fn execute(
403        &self,
404        tool_name: &str,
405        parameters: serde_json::Value,
406    ) -> Result<ToolResult, Box<dyn Error + Send + Sync>>;
407
408    /// Get metadata about available tools
409    async fn list_tools(&self) -> Result<Vec<ToolMetadata>, Box<dyn Error + Send + Sync>>;
410
411    /// Get metadata about a specific tool
412    async fn get_tool_metadata(
413        &self,
414        tool_name: &str,
415    ) -> Result<ToolMetadata, Box<dyn Error + Send + Sync>>;
416
417    /// Protocol identifier (e.g., "mcp", "custom", "openai-functions")
418    fn protocol_name(&self) -> &str;
419
420    /// Initialize/connect to the tool protocol
421    async fn initialize(&mut self) -> Result<(), Box<dyn Error + Send + Sync>> {
422        Ok(())
423    }
424
425    /// Cleanup/disconnect from the tool protocol
426    async fn shutdown(&mut self) -> Result<(), Box<dyn Error + Send + Sync>> {
427        Ok(())
428    }
429
430    /// List available resources (MCP Resource support)
431    ///
432    /// Resources are application-provided contextual data that agents can read.
433    /// This method is optional and defaults to returning an empty list.
434    async fn list_resources(&self) -> Result<Vec<ResourceMetadata>, Box<dyn Error + Send + Sync>> {
435        Ok(Vec::new())
436    }
437
438    /// Read the content of a resource by URI (MCP Resource support)
439    ///
440    /// This method is optional and defaults to returning NotFound.
441    async fn read_resource(&self, uri: &str) -> Result<String, Box<dyn Error + Send + Sync>> {
442        Err(format!("Resource not found: {}", uri).into())
443    }
444
445    /// Check if this protocol supports resources
446    fn supports_resources(&self) -> bool {
447        false
448    }
449}
450
451/// Error types for tool operations
452#[derive(Debug, Clone)]
453pub enum ToolError {
454    /// Requested tool is not registered in the current registry/protocol.
455    NotFound(String),
456    /// Tool execution completed with an application level failure.
457    ExecutionFailed(String),
458    /// The provided JSON parameters failed validation or deserialization.
459    InvalidParameters(String),
460    /// A lower level protocol/transport error occurred.
461    ProtocolError(String),
462}
463
464impl fmt::Display for ToolError {
465    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
466        match self {
467            ToolError::NotFound(name) => write!(f, "Tool not found: {}", name),
468            ToolError::ExecutionFailed(msg) => write!(f, "Tool execution failed: {}", msg),
469            ToolError::InvalidParameters(msg) => write!(f, "Invalid parameters: {}", msg),
470            ToolError::ProtocolError(msg) => write!(f, "Protocol error: {}", msg),
471        }
472    }
473}
474
475impl Error for ToolError {}
476
477/// A tool that can be used by agents
478pub struct Tool {
479    /// Metadata describing the tool interface.
480    metadata: ToolMetadata,
481    /// Underlying protocol implementation that actually executes the tool.
482    protocol: Arc<dyn ToolProtocol>,
483}
484
485impl Tool {
486    /// Create a new tool bound to the supplied protocol implementation.
487    pub fn new(
488        name: impl Into<String>,
489        description: impl Into<String>,
490        protocol: Arc<dyn ToolProtocol>,
491    ) -> Self {
492        Self {
493            metadata: ToolMetadata::new(name, description),
494            protocol,
495        }
496    }
497
498    /// Add a parameter definition to the tool builder.
499    pub fn with_parameter(mut self, param: ToolParameter) -> Self {
500        self.metadata.parameters.push(param);
501        self
502    }
503
504    /// Attach protocol specific metadata to the tool builder.
505    pub fn with_protocol_metadata(
506        mut self,
507        key: impl Into<String>,
508        value: serde_json::Value,
509    ) -> Self {
510        self.metadata.protocol_metadata.insert(key.into(), value);
511        self
512    }
513
514    /// Borrow the static metadata for the tool.
515    pub fn metadata(&self) -> &ToolMetadata {
516        &self.metadata
517    }
518
519    /// Execute the tool using the configured protocol.
520    pub async fn execute(
521        &self,
522        parameters: serde_json::Value,
523    ) -> Result<ToolResult, Box<dyn Error + Send + Sync>> {
524        self.protocol.execute(&self.metadata.name, parameters).await
525    }
526}
527
528/// Registry for managing tools available to agents
529///
530/// Supports single or multiple tool protocols, enabling agents to transparently
531/// access tools from multiple sources (local functions, MCP servers, etc.)
532///
533/// # Single Protocol
534///
535/// ```rust,no_run
536/// use cloudllm::tool_protocol::ToolRegistry;
537/// use cloudllm::tool_protocols::CustomToolProtocol;
538/// use std::sync::Arc;
539///
540/// let protocol = Arc::new(CustomToolProtocol::new());
541/// let registry = ToolRegistry::new(protocol);
542/// ```
543///
544/// # Multiple Protocols
545///
546/// ```rust,no_run
547/// use cloudllm::tool_protocol::ToolRegistry;
548/// use cloudllm::tool_protocols::{CustomToolProtocol, McpClientProtocol};
549/// use std::sync::Arc;
550///
551/// # async {
552/// let mut registry = ToolRegistry::empty();
553///
554/// // Add local tools
555/// registry.add_protocol(
556///     "local",
557///     Arc::new(CustomToolProtocol::new())
558/// ).await.ok();
559///
560/// // Add remote MCP server
561/// registry.add_protocol(
562///     "youtube",
563///     Arc::new(McpClientProtocol::new("http://youtube-mcp:8081".to_string()))
564/// ).await.ok();
565///
566/// // Agent transparently accesses both
567/// # };
568/// ```
569pub struct ToolRegistry {
570    /// All discovered tools from all protocols
571    tools: HashMap<String, Tool>,
572    /// Mapping of tool_name -> protocol_name for routing
573    tool_to_protocol: HashMap<String, String>,
574    /// All registered protocols
575    protocols: HashMap<String, Arc<dyn ToolProtocol>>,
576    /// Primary protocol (for backwards compatibility with single-protocol code)
577    primary_protocol: Option<Arc<dyn ToolProtocol>>,
578}
579
580impl ToolRegistry {
581    /// Build a registry powered by a single protocol implementation.
582    ///
583    /// This is the traditional single-protocol mode. Use `empty()` and `add_protocol()`
584    /// for multi-protocol support.
585    pub fn new(protocol: Arc<dyn ToolProtocol>) -> Self {
586        Self {
587            tools: HashMap::new(),
588            tool_to_protocol: HashMap::new(),
589            protocols: {
590                let mut m = HashMap::new();
591                m.insert("primary".to_string(), protocol.clone());
592                m
593            },
594            primary_protocol: Some(protocol),
595        }
596    }
597
598    /// Create an empty registry ready to accept multiple protocols.
599    ///
600    /// Use `add_protocol()` to register protocols.
601    pub fn empty() -> Self {
602        Self {
603            tools: HashMap::new(),
604            tool_to_protocol: HashMap::new(),
605            protocols: HashMap::new(),
606            primary_protocol: None,
607        }
608    }
609
610    /// Register a protocol and discover its tools.
611    ///
612    /// # Arguments
613    ///
614    /// * `protocol_name` - Unique identifier for this protocol (e.g., "local", "youtube", "github")
615    /// * `protocol` - The ToolProtocol implementation
616    ///
617    /// # Tool Discovery
618    ///
619    /// This method calls `protocol.list_tools()` to discover available tools
620    /// and automatically registers them in the registry.
621    ///
622    /// # Conflicts
623    ///
624    /// If a tool with the same name already exists, it will be replaced.
625    /// The new protocol's tool takes precedence.
626    ///
627    /// # Example
628    ///
629    /// ```rust,no_run
630    /// use cloudllm::tool_protocol::ToolRegistry;
631    /// use cloudllm::tool_protocols::McpClientProtocol;
632    /// use std::sync::Arc;
633    ///
634    /// # async {
635    /// let mut registry = ToolRegistry::empty();
636    /// registry.add_protocol(
637    ///     "memory_server",
638    ///     Arc::new(McpClientProtocol::new("http://localhost:8080".to_string()))
639    /// ).await?;
640    /// # Ok::<(), Box<dyn std::error::Error + Send + Sync>>(())
641    /// # };
642    /// ```
643    pub async fn add_protocol(
644        &mut self,
645        protocol_name: &str,
646        protocol: Arc<dyn ToolProtocol>,
647    ) -> Result<(), Box<dyn Error + Send + Sync>> {
648        // Discover tools from this protocol
649        let discovered_tools = protocol.list_tools().await?;
650
651        // Register the protocol
652        self.protocols
653            .insert(protocol_name.to_string(), protocol.clone());
654
655        // Register discovered tools
656        for tool_meta in discovered_tools {
657            let tool_name = tool_meta.name.clone();
658
659            // Create a Tool that routes through this protocol
660            let tool = Tool::new(
661                tool_name.clone(),
662                tool_meta.description.clone(),
663                protocol.clone(),
664            );
665
666            // Copy over any additional parameters and metadata
667            let mut tool = tool;
668            for param in &tool_meta.parameters {
669                tool = tool.with_parameter(param.clone());
670            }
671            for (key, value) in &tool_meta.protocol_metadata {
672                tool = tool.with_protocol_metadata(key.clone(), value.clone());
673            }
674
675            // Register the tool and its routing
676            self.tools.insert(tool_name.clone(), tool);
677            self.tool_to_protocol
678                .insert(tool_name, protocol_name.to_string());
679        }
680
681        Ok(())
682    }
683
684    /// Remove a protocol and all its tools from the registry.
685    pub fn remove_protocol(&mut self, protocol_name: &str) {
686        self.protocols.remove(protocol_name);
687
688        // Collect tool names to remove
689        let tools_to_remove: Vec<String> = self
690            .tool_to_protocol
691            .iter()
692            .filter(|(_, pn)| *pn == protocol_name)
693            .map(|(tn, _)| tn.clone())
694            .collect();
695
696        // Remove the tools
697        for tool_name in tools_to_remove {
698            self.tools.remove(&tool_name);
699            self.tool_to_protocol.remove(&tool_name);
700        }
701    }
702
703    /// Insert or replace a tool definition (for manual tool registration).
704    pub fn add_tool(&mut self, tool: Tool) {
705        self.tools.insert(tool.metadata.name.clone(), tool);
706    }
707
708    /// Remove a tool by name returning the owned entry if present.
709    pub fn remove_tool(&mut self, name: &str) -> Option<Tool> {
710        self.tool_to_protocol.remove(name);
711        self.tools.remove(name)
712    }
713
714    /// Borrow a tool by name.
715    pub fn get_tool(&self, name: &str) -> Option<&Tool> {
716        self.tools.get(name)
717    }
718
719    /// List metadata for registered tools (iteration order follows the underlying map).
720    pub fn list_tools(&self) -> Vec<&ToolMetadata> {
721        self.tools.values().map(|t| &t.metadata).collect()
722    }
723
724    /// Discover tools from the primary protocol (for single-protocol registries).
725    ///
726    /// This is useful after registering tools with the protocol to populate the registry.
727    /// For multi-protocol registries, use `add_protocol()` instead.
728    pub async fn discover_tools_from_primary(
729        &mut self,
730    ) -> Result<(), Box<dyn Error + Send + Sync>> {
731        if let Some(protocol) = &self.primary_protocol {
732            let discovered_tools = protocol.list_tools().await?;
733            for tool_meta in discovered_tools {
734                let tool_name = tool_meta.name.clone();
735                let tool = Tool::new(
736                    tool_name.clone(),
737                    tool_meta.description.clone(),
738                    protocol.clone(),
739                );
740
741                // Copy over parameters and metadata
742                let mut tool = tool;
743                for param in &tool_meta.parameters {
744                    tool = tool.with_parameter(param.clone());
745                }
746                for (key, value) in &tool_meta.protocol_metadata {
747                    tool = tool.with_protocol_metadata(key.clone(), value.clone());
748                }
749
750                self.tools.insert(tool_name.clone(), tool);
751                self.tool_to_protocol
752                    .insert(tool_name, "primary".to_string());
753            }
754            Ok(())
755        } else {
756            Err("No primary protocol available".into())
757        }
758    }
759
760    /// Get which protocol handles a specific tool.
761    pub fn get_tool_protocol(&self, tool_name: &str) -> Option<&str> {
762        self.tool_to_protocol.get(tool_name).map(|s| s.as_str())
763    }
764
765    /// Get all registered protocol names.
766    pub fn list_protocols(&self) -> Vec<&str> {
767        self.protocols.keys().map(|s| s.as_str()).collect()
768    }
769
770    /// Execute a named tool with serialized parameters.
771    pub async fn execute_tool(
772        &self,
773        tool_name: &str,
774        parameters: serde_json::Value,
775    ) -> Result<ToolResult, Box<dyn Error + Send + Sync>> {
776        let tool = self
777            .tools
778            .get(tool_name)
779            .ok_or_else(|| ToolError::NotFound(tool_name.to_string()))?;
780
781        tool.execute(parameters).await
782    }
783
784    /// Borrow the primary protocol implementation (for single-protocol mode).
785    ///
786    /// Returns None if registry was created with `empty()` or has multiple protocols.
787    pub fn protocol(&self) -> Option<&Arc<dyn ToolProtocol>> {
788        self.primary_protocol.as_ref()
789    }
790
791    /// Returns all registered tools as [`ToolDefinition`]s ready to pass to the LLM.
792    ///
793    /// Iterates over every tool in the registry, calling
794    /// [`ToolMetadata::to_tool_definition`] on each, and returns the results as a `Vec`.
795    /// An empty `Vec` is returned when no tools have been registered.
796    ///
797    /// # Example
798    ///
799    /// ```rust
800    /// use cloudllm::tool_protocol::{ToolRegistry, Tool, ToolMetadata, ToolParameter, ToolParameterType};
801    /// use cloudllm::tool_protocols::CustomToolProtocol;
802    /// use std::sync::Arc;
803    ///
804    /// # #[tokio::main]
805    /// # async fn main() {
806    /// let protocol = Arc::new(CustomToolProtocol::new());
807    /// let mut registry = ToolRegistry::new(protocol.clone());
808    ///
809    /// let tool = Tool::new("calculator", "Evaluates math", protocol);
810    /// registry.add_tool(tool);
811    ///
812    /// let defs = registry.to_tool_definitions();
813    /// assert_eq!(defs.len(), 1);
814    /// assert_eq!(defs[0].name, "calculator");
815    /// # }
816    /// ```
817    pub fn to_tool_definitions(&self) -> Vec<ToolDefinition> {
818        self.tools
819            .values()
820            .map(|t| t.metadata.to_tool_definition())
821            .collect()
822    }
823}
824
825#[cfg(test)]
826mod tests {
827    use super::*;
828
829    struct MockProtocol;
830
831    #[async_trait]
832    impl ToolProtocol for MockProtocol {
833        async fn execute(
834            &self,
835            tool_name: &str,
836            _parameters: serde_json::Value,
837        ) -> Result<ToolResult, Box<dyn Error + Send + Sync>> {
838            Ok(ToolResult::success(serde_json::json!({
839                "tool": tool_name,
840                "result": "mock_result"
841            })))
842        }
843
844        async fn list_tools(&self) -> Result<Vec<ToolMetadata>, Box<dyn Error + Send + Sync>> {
845            Ok(vec![])
846        }
847
848        async fn get_tool_metadata(
849            &self,
850            _tool_name: &str,
851        ) -> Result<ToolMetadata, Box<dyn Error + Send + Sync>> {
852            Ok(ToolMetadata::new("mock_tool", "A mock tool"))
853        }
854
855        fn protocol_name(&self) -> &str {
856            "mock"
857        }
858    }
859
860    #[test]
861    fn test_tool_parameter_builder() {
862        let param = ToolParameter::new("test_param", ToolParameterType::String)
863            .with_description("A test parameter")
864            .required()
865            .with_default(serde_json::json!("default_value"));
866
867        assert_eq!(param.name, "test_param");
868        assert_eq!(param.param_type, ToolParameterType::String);
869        assert_eq!(param.description, Some("A test parameter".to_string()));
870        assert!(param.required);
871        assert_eq!(param.default, Some(serde_json::json!("default_value")));
872    }
873
874    #[tokio::test]
875    async fn test_tool_execution() {
876        let protocol = Arc::new(MockProtocol);
877        let tool = Tool::new("test_tool", "A test tool", protocol.clone());
878
879        let result = tool.execute(serde_json::json!({})).await.unwrap();
880        assert!(result.success);
881        assert_eq!(result.output["tool"], "test_tool");
882    }
883
884    #[tokio::test]
885    async fn test_tool_registry() {
886        let protocol = Arc::new(MockProtocol);
887        let mut registry = ToolRegistry::new(protocol.clone());
888
889        let tool = Tool::new("calculator", "Performs calculations", protocol.clone());
890        registry.add_tool(tool);
891
892        assert!(registry.get_tool("calculator").is_some());
893        assert_eq!(registry.list_tools().len(), 1);
894
895        let result = registry
896            .execute_tool("calculator", serde_json::json!({}))
897            .await
898            .unwrap();
899        assert!(result.success);
900    }
901
902    #[tokio::test]
903    async fn test_empty_registry_creation() {
904        let registry = ToolRegistry::empty();
905        assert_eq!(registry.list_tools().len(), 0);
906        assert_eq!(registry.list_protocols().len(), 0);
907        assert!(registry.protocol().is_none());
908    }
909
910    #[tokio::test]
911    async fn test_add_single_protocol_to_empty_registry() {
912        let protocol = Arc::new(MockProtocol);
913        let mut registry = ToolRegistry::empty();
914
915        // Add protocol
916        registry
917            .add_protocol("mock", protocol.clone())
918            .await
919            .unwrap();
920
921        // Verify protocol was added
922        assert_eq!(registry.list_protocols().len(), 1);
923        assert!(registry.list_protocols().contains(&"mock"));
924    }
925
926    #[tokio::test]
927    async fn test_add_multiple_protocols() {
928        let protocol1 = Arc::new(MockProtocol);
929        let protocol2 = Arc::new(MockProtocol);
930        let mut registry = ToolRegistry::empty();
931
932        // Add two protocols
933        registry
934            .add_protocol("protocol1", protocol1.clone())
935            .await
936            .unwrap();
937        registry
938            .add_protocol("protocol2", protocol2.clone())
939            .await
940            .unwrap();
941
942        // Verify both protocols are registered
943        assert_eq!(registry.list_protocols().len(), 2);
944        assert!(registry.list_protocols().contains(&"protocol1"));
945        assert!(registry.list_protocols().contains(&"protocol2"));
946    }
947
948    #[tokio::test]
949    async fn test_remove_protocol() {
950        let protocol = Arc::new(MockProtocol);
951        let mut registry = ToolRegistry::empty();
952
953        // Add protocol
954        registry
955            .add_protocol("protocol1", protocol.clone())
956            .await
957            .unwrap();
958        assert_eq!(registry.list_protocols().len(), 1);
959
960        // Remove protocol
961        registry.remove_protocol("protocol1");
962        assert_eq!(registry.list_protocols().len(), 0);
963    }
964
965    #[tokio::test]
966    async fn test_get_tool_protocol() {
967        let protocol = Arc::new(MockProtocol);
968        let mut registry = ToolRegistry::empty();
969
970        // Add protocol with a tool
971        registry
972            .add_protocol("local", protocol.clone())
973            .await
974            .unwrap();
975
976        // Add a tool manually for testing
977        let tool = Tool::new("calculator", "Performs calculations", protocol.clone());
978        registry.add_tool(tool);
979        registry
980            .tool_to_protocol
981            .insert("calculator".to_string(), "local".to_string());
982
983        // Verify tool-to-protocol mapping
984        assert_eq!(registry.get_tool_protocol("calculator"), Some("local"));
985        assert_eq!(registry.get_tool_protocol("nonexistent"), None);
986    }
987
988    #[tokio::test]
989    async fn test_remove_protocol_removes_tools() {
990        let protocol = Arc::new(MockProtocol);
991        let mut registry = ToolRegistry::empty();
992
993        // Add protocol and tools
994        registry
995            .add_protocol("protocol1", protocol.clone())
996            .await
997            .unwrap();
998
999        let tool1 = Tool::new("tool1", "First tool", protocol.clone());
1000        registry.add_tool(tool1);
1001        registry
1002            .tool_to_protocol
1003            .insert("tool1".to_string(), "protocol1".to_string());
1004
1005        let tool2 = Tool::new("tool2", "Second tool", protocol.clone());
1006        registry.add_tool(tool2);
1007        registry
1008            .tool_to_protocol
1009            .insert("tool2".to_string(), "protocol1".to_string());
1010
1011        assert_eq!(registry.list_tools().len(), 2);
1012
1013        // Remove protocol
1014        registry.remove_protocol("protocol1");
1015
1016        // Verify all tools from that protocol are removed
1017        assert_eq!(registry.list_tools().len(), 0);
1018        assert_eq!(registry.get_tool_protocol("tool1"), None);
1019        assert_eq!(registry.get_tool_protocol("tool2"), None);
1020    }
1021
1022    #[tokio::test]
1023    async fn test_execute_tool_through_registry() {
1024        let protocol = Arc::new(MockProtocol);
1025        let mut registry = ToolRegistry::empty();
1026
1027        // Add protocol
1028        registry
1029            .add_protocol("mock", protocol.clone())
1030            .await
1031            .unwrap();
1032
1033        // Add and execute tool
1034        let tool = Tool::new("test_tool", "A test tool", protocol.clone());
1035        registry.add_tool(tool);
1036
1037        let result = registry
1038            .execute_tool("test_tool", serde_json::json!({}))
1039            .await
1040            .unwrap();
1041
1042        assert!(result.success);
1043        assert_eq!(result.output["tool"], "test_tool");
1044    }
1045
1046    #[tokio::test]
1047    async fn test_backwards_compatibility_single_protocol() {
1048        let protocol = Arc::new(MockProtocol);
1049        let registry = ToolRegistry::new(protocol.clone());
1050
1051        // Single-protocol registry should have primary_protocol set
1052        assert!(registry.protocol().is_some());
1053        assert_eq!(registry.list_protocols().len(), 1);
1054        assert!(registry.list_protocols().contains(&"primary"));
1055    }
1056
1057    #[tokio::test]
1058    async fn test_discover_tools_from_primary() {
1059        struct TestProtocol {
1060            tools: Vec<ToolMetadata>,
1061        }
1062
1063        #[async_trait]
1064        impl ToolProtocol for TestProtocol {
1065            async fn execute(
1066                &self,
1067                tool_name: &str,
1068                _parameters: serde_json::Value,
1069            ) -> Result<ToolResult, Box<dyn Error + Send + Sync>> {
1070                Ok(ToolResult::success(serde_json::json!({
1071                    "tool": tool_name,
1072                })))
1073            }
1074
1075            async fn list_tools(&self) -> Result<Vec<ToolMetadata>, Box<dyn Error + Send + Sync>> {
1076                Ok(self.tools.clone())
1077            }
1078
1079            async fn get_tool_metadata(
1080                &self,
1081                tool_name: &str,
1082            ) -> Result<ToolMetadata, Box<dyn Error + Send + Sync>> {
1083                self.tools
1084                    .iter()
1085                    .find(|t| t.name == tool_name)
1086                    .cloned()
1087                    .ok_or_else(|| "Tool not found".into())
1088            }
1089
1090            fn protocol_name(&self) -> &str {
1091                "test"
1092            }
1093        }
1094
1095        // Create protocol with tools
1096        let protocol = Arc::new(TestProtocol {
1097            tools: vec![
1098                ToolMetadata::new("tool1", "First tool"),
1099                ToolMetadata::new("tool2", "Second tool"),
1100            ],
1101        });
1102
1103        let mut registry = ToolRegistry::new(protocol.clone());
1104
1105        // Initially, registry has no tools (they haven't been discovered)
1106        assert_eq!(registry.list_tools().len(), 0);
1107
1108        // Discover tools
1109        registry.discover_tools_from_primary().await.unwrap();
1110
1111        // Now registry should have the tools
1112        assert_eq!(registry.list_tools().len(), 2);
1113        assert!(registry.get_tool("tool1").is_some());
1114        assert!(registry.get_tool("tool2").is_some());
1115
1116        // Tools should be mapped to primary protocol
1117        assert_eq!(registry.get_tool_protocol("tool1"), Some("primary"));
1118        assert_eq!(registry.get_tool_protocol("tool2"), Some("primary"));
1119    }
1120
1121    // ──────────────────────────────────────────────────────────────────────
1122    // New tests for native tool-calling support (v0.11.1)
1123    // ──────────────────────────────────────────────────────────────────────
1124
1125    #[test]
1126    fn test_to_json_schema_string() {
1127        let param =
1128            ToolParameter::new("q", ToolParameterType::String).with_description("Search query");
1129        let schema = param.to_json_schema();
1130        assert_eq!(schema["type"], "string");
1131        assert_eq!(schema["description"], "Search query");
1132    }
1133
1134    #[test]
1135    fn test_to_json_schema_number() {
1136        let param = ToolParameter::new("value", ToolParameterType::Number);
1137        let schema = param.to_json_schema();
1138        assert_eq!(schema["type"], "number");
1139        // no description key when absent
1140        assert!(schema.get("description").is_none());
1141    }
1142
1143    #[test]
1144    fn test_to_json_schema_integer() {
1145        let schema = ToolParameter::new("n", ToolParameterType::Integer).to_json_schema();
1146        assert_eq!(schema["type"], "integer");
1147    }
1148
1149    #[test]
1150    fn test_to_json_schema_boolean() {
1151        let schema = ToolParameter::new("flag", ToolParameterType::Boolean).to_json_schema();
1152        assert_eq!(schema["type"], "boolean");
1153    }
1154
1155    #[test]
1156    fn test_to_json_schema_array_with_items() {
1157        let param = ToolParameter::new("ids", ToolParameterType::Array)
1158            .with_items(ToolParameterType::Integer);
1159        let schema = param.to_json_schema();
1160        assert_eq!(schema["type"], "array");
1161        assert_eq!(schema["items"]["type"], "integer");
1162    }
1163
1164    #[test]
1165    fn test_to_json_schema_array_without_items() {
1166        let schema = ToolParameter::new("items", ToolParameterType::Array).to_json_schema();
1167        assert_eq!(schema["type"], "array");
1168        // items key may be present but empty {}
1169        assert!(schema.get("items").is_some());
1170    }
1171
1172    #[test]
1173    fn test_to_json_schema_object_with_properties() {
1174        use std::collections::HashMap;
1175        let mut props = HashMap::new();
1176        props.insert(
1177            "name".to_string(),
1178            ToolParameter::new("name", ToolParameterType::String)
1179                .with_description("Person's name")
1180                .required(),
1181        );
1182        props.insert(
1183            "age".to_string(),
1184            ToolParameter::new("age", ToolParameterType::Integer),
1185        );
1186        let param = ToolParameter::new("person", ToolParameterType::Object).with_properties(props);
1187        let schema = param.to_json_schema();
1188        assert_eq!(schema["type"], "object");
1189        assert!(schema["properties"]["name"].is_object());
1190        assert!(schema["properties"]["age"].is_object());
1191        // "name" is required, "age" is not
1192        let required = schema["required"].as_array().unwrap();
1193        assert!(required.iter().any(|v| v.as_str() == Some("name")));
1194        assert!(!required.iter().any(|v| v.as_str() == Some("age")));
1195    }
1196
1197    #[test]
1198    fn test_to_tool_definition_roundtrip() {
1199        let meta = ToolMetadata::new("calculator", "Evaluates a math expression")
1200            .with_parameter(
1201                ToolParameter::new("expression", ToolParameterType::String)
1202                    .with_description("The expression")
1203                    .required(),
1204            )
1205            .with_parameter(
1206                ToolParameter::new("precision", ToolParameterType::Integer)
1207                    .with_description("Decimal places"),
1208            );
1209
1210        let def = meta.to_tool_definition();
1211        assert_eq!(def.name, "calculator");
1212        assert_eq!(def.description, "Evaluates a math expression");
1213        assert_eq!(def.parameters_schema["type"], "object");
1214        assert!(def.parameters_schema["properties"]["expression"].is_object());
1215        assert!(def.parameters_schema["properties"]["precision"].is_object());
1216
1217        let required = def.parameters_schema["required"].as_array().unwrap();
1218        assert!(required.iter().any(|v| v.as_str() == Some("expression")));
1219        assert!(!required.iter().any(|v| v.as_str() == Some("precision")));
1220    }
1221
1222    #[test]
1223    fn test_to_tool_definitions_collects_all() {
1224        let protocol = Arc::new(MockProtocol);
1225        let mut registry = ToolRegistry::empty();
1226
1227        let tool_a = Tool::new("tool_a", "First tool", protocol.clone());
1228        registry.add_tool(tool_a);
1229        let tool_b = Tool::new("tool_b", "Second tool", protocol.clone());
1230        registry.add_tool(tool_b);
1231
1232        let defs = registry.to_tool_definitions();
1233        assert_eq!(defs.len(), 2);
1234        let names: Vec<&str> = defs.iter().map(|d| d.name.as_str()).collect();
1235        assert!(names.contains(&"tool_a"));
1236        assert!(names.contains(&"tool_b"));
1237    }
1238
1239    #[test]
1240    fn test_to_tool_definitions_empty_registry() {
1241        let registry = ToolRegistry::empty();
1242        let defs = registry.to_tool_definitions();
1243        assert!(defs.is_empty());
1244    }
1245}
mcp/protocol.rs

mcp/
protocol.rs
