Skip to main content

ai_agents_runtime/spec/
tool.rs

1//! Tool configuration types for agent specification.
2//!
3//! `ToolEntry` supports plain strings, structured builtins, and MCP wrapper entries.
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7
8use ai_agents_tools::mcp::wrapper::MCPWrapperConfig;
9
10/// A single entry in the agent-level `tools:` list.
11///
12/// Supports three YAML forms:
13/// 1. Plain string:         `- datetime`
14/// 2. Builtin with config:  `- name: http`
15/// 3. MCP wrapper:          `- name: github\n  type: mcp\n  transport: stdio\n  ...`
16#[derive(Debug, Clone, Serialize, Deserialize)]
17#[serde(untagged)]
18pub enum ToolEntry {
19    /// Plain string tool reference (e.g., `- datetime`).
20    Simple(String),
21    /// Structured entry — builtin with extra config or MCP wrapper.
22    Structured(StructuredToolEntry),
23}
24
25/// A structured tool entry with a name, optional type, and extra configuration.
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct StructuredToolEntry {
28    pub name: String,
29
30    /// `"mcp"` for MCP wrapper tools, absent/null for builtin tools.
31    #[serde(rename = "type", default)]
32    pub tool_type: Option<String>,
33
34    /// Extra fields — flattened from YAML. For builtins, holds additional config.
35    /// For MCP tools, holds transport/env/security/etc.
36    #[serde(flatten)]
37    pub extra: Value,
38}
39
40impl ToolEntry {
41    /// Get the tool name regardless of entry form.
42    pub fn name(&self) -> &str {
43        match self {
44            ToolEntry::Simple(name) => name,
45            ToolEntry::Structured(s) => &s.name,
46        }
47    }
48
49    /// Check if this entry is an MCP wrapper tool.
50    pub fn is_mcp(&self) -> bool {
51        match self {
52            ToolEntry::Simple(_) => false,
53            ToolEntry::Structured(s) => s.tool_type.as_deref() == Some("mcp"),
54        }
55    }
56
57    /// Extract MCP wrapper config from a structured MCP entry.
58    ///
59    /// Re-serializes the structured entry to JSON, then deserializes as
60    /// `MCPWrapperConfig` to pick up all flattened MCP fields.
61    pub fn to_mcp_config(&self) -> Option<MCPWrapperConfig> {
62        if !self.is_mcp() {
63            return None;
64        }
65        match self {
66            ToolEntry::Structured(s) => {
67                let value = serde_json::to_value(s).ok()?;
68                serde_json::from_value(value).ok()
69            }
70            _ => None,
71        }
72    }
73}
74
75/// Backward compatibility alias — existing code using `ToolConfig` keeps working.
76pub type ToolConfig = ToolEntry;
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    #[test]
83    fn test_tool_entry_plain_string() {
84        let yaml = "datetime";
85        let entry: ToolEntry = serde_yaml::from_str(yaml).unwrap();
86        assert_eq!(entry.name(), "datetime");
87        assert!(!entry.is_mcp());
88    }
89
90    #[test]
91    fn test_tool_entry_structured_builtin() {
92        let yaml = "name: http";
93        let entry: ToolEntry = serde_yaml::from_str(yaml).unwrap();
94        assert_eq!(entry.name(), "http");
95        assert!(!entry.is_mcp());
96    }
97
98    #[test]
99    fn test_tool_entry_mcp() {
100        let yaml = r#"
101name: github
102type: mcp
103transport: stdio
104command: npx
105args: ["-y", "@modelcontextprotocol/server-github"]
106env:
107  GITHUB_TOKEN: "test"
108"#;
109        let entry: ToolEntry = serde_yaml::from_str(yaml).unwrap();
110        assert_eq!(entry.name(), "github");
111        assert!(entry.is_mcp());
112        let config = entry.to_mcp_config().unwrap();
113        assert_eq!(config.name, "github");
114    }
115
116    #[test]
117    fn test_tool_entry_mixed_list() {
118        let yaml = r#"
119- datetime
120- name: http
121- name: github
122  type: mcp
123  transport: stdio
124  command: npx
125  args: ["-y", "@modelcontextprotocol/server-github"]
126  env:
127    GITHUB_TOKEN: "test"
128"#;
129        let entries: Vec<ToolEntry> = serde_yaml::from_str(yaml).unwrap();
130        assert_eq!(entries.len(), 3);
131        assert_eq!(entries[0].name(), "datetime");
132        assert!(!entries[0].is_mcp());
133        assert_eq!(entries[1].name(), "http");
134        assert!(!entries[1].is_mcp());
135        assert_eq!(entries[2].name(), "github");
136        assert!(entries[2].is_mcp());
137    }
138
139    #[test]
140    fn test_tool_config_backward_compat() {
141        let config = ToolConfig::Simple("echo".to_string());
142        assert_eq!(config.name(), "echo");
143    }
144
145    #[test]
146    fn test_tool_entry_mcp_with_views() {
147        let yaml = r#"
148name: github
149type: mcp
150transport: stdio
151command: npx
152args: ["-y", "@modelcontextprotocol/server-github"]
153env:
154  GITHUB_TOKEN: "test"
155views:
156  github_issues:
157    functions: [create_issue, list_issues]
158  github_code:
159    functions: [search_code]
160    description: "Code search"
161"#;
162        let entry: ToolEntry = serde_yaml::from_str(yaml).unwrap();
163        assert_eq!(entry.name(), "github");
164        assert!(entry.is_mcp());
165        let config = entry.to_mcp_config().unwrap();
166        assert_eq!(config.views.len(), 2);
167        assert_eq!(
168            config.views["github_issues"].functions,
169            vec!["create_issue", "list_issues"]
170        );
171        assert_eq!(
172            config.views["github_code"].description.as_deref(),
173            Some("Code search")
174        );
175    }
176
177    #[test]
178    fn test_tool_entry_name_method() {
179        let simple = ToolEntry::Simple("calculator".to_string());
180        assert_eq!(simple.name(), "calculator");
181
182        let structured = ToolEntry::Structured(StructuredToolEntry {
183            name: "custom_tool".to_string(),
184            tool_type: None,
185            extra: serde_json::json!({}),
186        });
187        assert_eq!(structured.name(), "custom_tool");
188    }
189}