1use std::collections::BTreeMap;
8
9use serde::{Deserialize, Serialize};
10
11use super::{McpServerConfig, ModelRef, PermissionPolicy, SkillRef, ToolConfig};
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
34#[serde(rename_all = "snake_case")]
35#[non_exhaustive]
36pub struct ManagedAgentDef {
37 pub name: String,
39 pub model: ModelRef,
41 #[serde(skip_serializing_if = "Option::is_none")]
43 pub system: Option<String>,
44 #[serde(skip_serializing_if = "Option::is_none")]
46 pub description: Option<String>,
47 #[serde(default, skip_serializing_if = "Vec::is_empty")]
49 pub tools: Vec<ToolConfig>,
50 #[serde(default, skip_serializing_if = "Vec::is_empty")]
52 pub mcp_servers: Vec<McpServerConfig>,
53 #[serde(default, skip_serializing_if = "Vec::is_empty")]
55 pub skills: Vec<SkillRef>,
56 #[serde(skip_serializing_if = "Option::is_none")]
58 pub permission_policy: Option<PermissionPolicy>,
59 #[serde(skip_serializing_if = "Option::is_none")]
61 pub metadata: Option<BTreeMap<String, String>>,
62}
63
64impl ManagedAgentDef {
65 pub fn new(name: impl Into<String>, model: ModelRef) -> Self {
81 Self {
82 name: name.into(),
83 model,
84 system: None,
85 description: None,
86 tools: Vec::new(),
87 mcp_servers: Vec::new(),
88 skills: Vec::new(),
89 permission_policy: None,
90 metadata: None,
91 }
92 }
93
94 pub fn with_system(mut self, system: impl Into<String>) -> Self {
96 self.system = Some(system.into());
97 self
98 }
99
100 pub fn with_description(mut self, description: impl Into<String>) -> Self {
102 self.description = Some(description.into());
103 self
104 }
105
106 pub fn with_tools(mut self, tools: Vec<ToolConfig>) -> Self {
108 self.tools = tools;
109 self
110 }
111
112 pub fn with_mcp_servers(mut self, mcp_servers: Vec<McpServerConfig>) -> Self {
114 self.mcp_servers = mcp_servers;
115 self
116 }
117
118 pub fn with_skills(mut self, skills: Vec<SkillRef>) -> Self {
120 self.skills = skills;
121 self
122 }
123
124 pub fn with_permission_policy(mut self, policy: PermissionPolicy) -> Self {
126 self.permission_policy = Some(policy);
127 self
128 }
129
130 pub fn with_metadata(mut self, metadata: BTreeMap<String, String>) -> Self {
132 self.metadata = Some(metadata);
133 self
134 }
135}
136
137#[cfg(test)]
138mod tests {
139 use std::collections::HashMap;
140
141 use super::*;
142 use crate::types::{ModelConfig, PermissionMode, Provider};
143
144 #[test]
145 fn test_serialize_full_def_matches_canon() {
146 let def = ManagedAgentDef {
147 name: "research-agent".to_string(),
148 model: ModelRef::Structured {
149 provider: Provider::Openai,
150 model: ModelConfig::Name("gpt-4.1".to_string()),
151 speed: None,
152 },
153 system: Some("You are a research assistant.".to_string()),
154 description: Some("Researches topics using web search and custom tools".to_string()),
155 tools: vec![
156 ToolConfig::WebSearch {},
157 ToolConfig::Custom {
158 name: "get_papers".to_string(),
159 description: Some("Search academic papers".to_string()),
160 input_schema: serde_json::json!({
161 "type": "object",
162 "properties": {
163 "query": {"type": "string"},
164 "limit": {"type": "integer"}
165 },
166 "required": ["query"]
167 }),
168 },
169 ],
170 mcp_servers: vec![McpServerConfig {
171 name: "arxiv-server".to_string(),
172 transport: "stdio".to_string(),
173 command: Some("npx".to_string()),
174 args: vec!["arxiv-mcp-server".to_string()],
175 url: None,
176 env: HashMap::new(),
177 auto_approve: vec!["search".to_string()],
178 }],
179 skills: vec![SkillRef { skill_id: "web-research".to_string() }],
180 permission_policy: Some(PermissionPolicy {
181 default: PermissionMode::AutoApprove,
182 tools: {
183 let mut m = HashMap::new();
184 m.insert("delete_file".to_string(), PermissionMode::Prompt);
185 m
186 },
187 }),
188 metadata: Some({
189 let mut m = BTreeMap::new();
190 m.insert("team".to_string(), "platform".to_string());
191 m.insert("version".to_string(), "1.0".to_string());
192 m
193 }),
194 };
195
196 let json = serde_json::to_value(&def).unwrap();
197
198 assert_eq!(json["name"], "research-agent");
200 assert_eq!(json["system"], "You are a research assistant.");
201 assert_eq!(json["description"], "Researches topics using web search and custom tools");
202
203 let model = &json["model"];
205 assert_eq!(model["provider"], "openai");
206 assert_eq!(model["model"], "gpt-4.1");
207 assert!(model.get("speed").is_none());
208
209 let tools = json["tools"].as_array().unwrap();
211 assert_eq!(tools.len(), 2);
212 assert_eq!(tools[0]["type"], "web_search");
213 assert_eq!(tools[1]["type"], "custom");
214 assert_eq!(tools[1]["name"], "get_papers");
215 assert!(tools[1]["input_schema"]["properties"]["query"].is_object());
216
217 let mcp = json["mcp_servers"].as_array().unwrap();
219 assert_eq!(mcp.len(), 1);
220 assert_eq!(mcp[0]["name"], "arxiv-server");
221 assert_eq!(mcp[0]["transport"], "stdio");
222 assert_eq!(mcp[0]["command"], "npx");
223 assert_eq!(mcp[0]["args"][0], "arxiv-mcp-server");
224 assert_eq!(mcp[0]["auto_approve"][0], "search");
225
226 let skills = json["skills"].as_array().unwrap();
228 assert_eq!(skills.len(), 1);
229 assert_eq!(skills[0]["skill_id"], "web-research");
230
231 let policy = &json["permission_policy"];
233 assert_eq!(policy["default"], "auto_approve");
234 assert_eq!(policy["tools"]["delete_file"], "prompt");
235
236 let metadata = &json["metadata"];
238 assert_eq!(metadata["team"], "platform");
239 assert_eq!(metadata["version"], "1.0");
240 }
241
242 #[test]
243 fn test_deserialize_full_def() {
244 let json = serde_json::json!({
245 "name": "test-agent",
246 "model": "gemini-2.5-flash",
247 "system": "Be helpful.",
248 "tools": [
249 {"type": "bash"},
250 {"type": "filesystem"}
251 ],
252 "skills": [{"skill_id": "coding"}],
253 "metadata": {"env": "staging"}
254 });
255
256 let def: ManagedAgentDef = serde_json::from_value(json).unwrap();
257 assert_eq!(def.name, "test-agent");
258 assert_eq!(def.system, Some("Be helpful.".to_string()));
259 assert_eq!(def.tools.len(), 2);
260 assert_eq!(def.skills.len(), 1);
261 assert_eq!(def.mcp_servers.len(), 0);
262 assert_eq!(def.permission_policy, None);
263 assert_eq!(def.description, None);
264 assert_eq!(def.metadata.as_ref().unwrap().get("env"), Some(&"staging".to_string()));
265 }
266
267 #[test]
268 fn test_minimal_def_omits_optional_fields() {
269 let def = ManagedAgentDef {
270 name: "minimal".to_string(),
271 model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
272 system: None,
273 description: None,
274 tools: vec![],
275 mcp_servers: vec![],
276 skills: vec![],
277 permission_policy: None,
278 metadata: None,
279 };
280
281 let json = serde_json::to_value(&def).unwrap();
282 let obj = json.as_object().unwrap();
283
284 assert!(obj.contains_key("name"));
286 assert!(obj.contains_key("model"));
287
288 assert!(!obj.contains_key("system"));
290 assert!(!obj.contains_key("description"));
291 assert!(!obj.contains_key("tools"));
292 assert!(!obj.contains_key("mcp_servers"));
293 assert!(!obj.contains_key("skills"));
294 assert!(!obj.contains_key("permission_policy"));
295 assert!(!obj.contains_key("metadata"));
296 }
297
298 #[test]
299 fn test_round_trip_serialization() {
300 let def = ManagedAgentDef {
301 name: "roundtrip-agent".to_string(),
302 model: ModelRef::Shorthand("claude-3.5-sonnet".to_string()),
303 system: Some("System prompt".to_string()),
304 description: None,
305 tools: vec![ToolConfig::Bash {}],
306 mcp_servers: vec![],
307 skills: vec![],
308 permission_policy: None,
309 metadata: None,
310 };
311
312 let json_str = serde_json::to_string(&def).unwrap();
313 let deserialized: ManagedAgentDef = serde_json::from_str(&json_str).unwrap();
314
315 assert_eq!(deserialized.name, "roundtrip-agent");
316 assert_eq!(deserialized.system, Some("System prompt".to_string()));
317 assert_eq!(deserialized.tools.len(), 1);
318 }
319}