1use serde::{Deserialize, Serialize};
4use typed_builder::TypedBuilder;
5
6use super::DEFAULT_IMAGE_GENERATION_MODEL;
7
8#[non_exhaustive]
9#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
10pub enum BuiltinTools {
11 #[serde(rename = "list_directory")]
13 ListDir,
14 #[serde(rename = "search_directory")]
16 SearchDir,
17 #[serde(rename = "find_file")]
19 FindFile,
20 #[serde(rename = "view_file")]
22 ViewFile,
23 #[serde(rename = "create_file")]
25 CreateFile,
26 #[serde(rename = "edit_file")]
28 EditFile,
29 #[serde(rename = "run_command")]
31 RunCommand,
32 #[serde(rename = "ask_question")]
34 AskQuestion,
35 #[serde(rename = "start_subagent")]
37 StartSubagent,
38 #[serde(rename = "generate_image")]
40 GenerateImage,
41 #[serde(rename = "search_web")]
43 SearchWeb,
44 #[serde(rename = "read_url_content")]
46 ReadUrlContent,
47 #[serde(rename = "finish")]
49 Finish,
50}
51
52impl BuiltinTools {
53 #[must_use]
54 pub const fn read_only() -> &'static [Self] {
56 &[
57 Self::ListDir,
58 Self::SearchDir,
59 Self::FindFile,
60 Self::ViewFile,
61 Self::ReadUrlContent,
62 Self::Finish,
63 ]
64 }
65
66 #[must_use]
68 pub const fn nondestructive() -> &'static [Self] {
69 &[
70 Self::ListDir,
71 Self::SearchDir,
72 Self::FindFile,
73 Self::ViewFile,
74 Self::CreateFile,
75 Self::EditFile,
76 Self::AskQuestion,
77 Self::StartSubagent,
78 Self::GenerateImage,
79 Self::SearchWeb,
80 Self::ReadUrlContent,
81 Self::Finish,
82 ]
83 }
84
85 #[must_use]
87 pub const fn all_tools() -> &'static [Self] {
88 &[
89 Self::ListDir,
90 Self::SearchDir,
91 Self::FindFile,
92 Self::ViewFile,
93 Self::CreateFile,
94 Self::EditFile,
95 Self::RunCommand,
96 Self::AskQuestion,
97 Self::StartSubagent,
98 Self::GenerateImage,
99 Self::SearchWeb,
100 Self::ReadUrlContent,
101 Self::Finish,
102 ]
103 }
104
105 #[must_use]
110 pub const fn file_tools() -> &'static [Self] {
111 &[Self::ViewFile, Self::CreateFile, Self::EditFile]
112 }
113
114 #[must_use]
116 pub const fn none() -> &'static [Self] {
117 &[]
118 }
119
120 #[must_use]
121 pub const fn as_sdk_name(&self) -> &'static str {
123 match self {
124 Self::ListDir => "list_directory",
125 Self::SearchDir => "search_directory",
126 Self::FindFile => "find_file",
127 Self::ViewFile => "view_file",
128 Self::CreateFile => "create_file",
129 Self::EditFile => "edit_file",
130 Self::RunCommand => "run_command",
131 Self::AskQuestion => "ask_question",
132 Self::StartSubagent => "start_subagent",
133 Self::GenerateImage => "generate_image",
134 Self::SearchWeb => "search_web",
135 Self::ReadUrlContent => "read_url_content",
136 Self::Finish => "finish",
137 }
138 }
139
140 #[must_use]
142 pub const fn description(&self) -> &'static str {
143 match self {
144 Self::ListDir => "List files and subdirectories.",
145 Self::SearchDir => "Regex search within directory contents.",
146 Self::FindFile => "Find files by name pattern.",
147 Self::ViewFile => "Read file contents.",
148 Self::CreateFile => "Create a new file.",
149 Self::EditFile => "Edit an existing file.",
150 Self::RunCommand => "Execute a shell command.",
151 Self::AskQuestion => "Ask the user a question.",
152 Self::StartSubagent => "Spawn a subagent.",
153 Self::GenerateImage => "Generate images from text prompts.",
154 Self::SearchWeb => "Search the web for information.",
155 Self::ReadUrlContent => "Fetch content from a URL via HTTP request.",
156 Self::Finish => "Signal task completion.",
157 }
158 }
159}
160
161impl std::fmt::Display for BuiltinTools {
162 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
163 f.write_str(self.as_sdk_name())
164 }
165}
166
167#[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)]
185pub struct CapabilitiesConfig {
186 #[serde(default = "super::default_true")]
188 #[builder(default = true)]
189 pub enable_subagents: bool,
190 #[serde(default)]
192 #[builder(default, setter(strip_option))]
193 pub enabled_tools: Option<Vec<BuiltinTools>>,
194 #[serde(default)]
196 #[builder(default, setter(strip_option))]
197 pub disabled_tools: Option<Vec<BuiltinTools>>,
198 #[builder(default, setter(strip_option))]
200 pub compaction_threshold: Option<usize>,
201 #[serde(default = "super::default_image_model")]
207 #[builder(default = DEFAULT_IMAGE_GENERATION_MODEL.to_owned(), setter(into))]
208 pub image_model: String,
209 #[serde(default)]
211 #[builder(default, setter(into, strip_option))]
212 pub finish_tool_schema_json: Option<String>,
213}
214
215impl CapabilitiesConfig {
216 #[must_use]
220 pub fn with_tools(tools: Vec<BuiltinTools>) -> Self {
221 Self {
222 enabled_tools: Some(tools),
223 ..Self::default()
224 }
225 }
226
227 #[must_use]
229 pub fn full() -> Self {
230 Self::default()
231 }
232
233 #[must_use]
235 pub fn read_only() -> Self {
236 Self {
237 enabled_tools: Some(BuiltinTools::read_only().to_vec()),
238 ..Self::default()
239 }
240 }
241
242 #[must_use]
246 pub fn custom_tools_only() -> Self {
247 Self {
248 enabled_tools: Some(vec![]),
249 ..Self::default()
250 }
251 }
252
253 pub const fn validate(&self) -> Result<(), &'static str> {
257 if self.enabled_tools.is_some() && self.disabled_tools.is_some() {
258 return Err("enabled_tools and disabled_tools are mutually exclusive");
259 }
260 Ok(())
261 }
262}
263
264impl Default for CapabilitiesConfig {
265 fn default() -> Self {
266 Self::builder().build()
268 }
269}
270
271#[cfg(test)]
272mod tests {
273 use pyo3::types::PyAnyMethods;
274
275 use super::*;
276
277 #[test]
278 fn test_builtin_tools() {
279 let read_only = BuiltinTools::read_only();
280 assert_eq!(read_only.len(), 6);
281 assert!(read_only.contains(&BuiltinTools::ListDir));
282 assert!(!read_only.contains(&BuiltinTools::SearchWeb));
283 assert!(read_only.contains(&BuiltinTools::ReadUrlContent));
284 assert!(read_only.contains(&BuiltinTools::Finish));
285 assert!(!read_only.contains(&BuiltinTools::CreateFile));
286
287 let all = BuiltinTools::all_tools();
288 assert_eq!(all.len(), 13);
289 assert!(all.contains(&BuiltinTools::CreateFile));
290 assert!(all.contains(&BuiltinTools::Finish));
291 assert!(all.contains(&BuiltinTools::SearchWeb));
292 assert!(all.contains(&BuiltinTools::ReadUrlContent));
293
294 assert_eq!(BuiltinTools::ListDir.as_sdk_name(), "list_directory");
295 assert_eq!(BuiltinTools::SearchWeb.as_sdk_name(), "search_web");
296 assert_eq!(
297 BuiltinTools::ReadUrlContent.as_sdk_name(),
298 "read_url_content"
299 );
300 }
301
302 #[test]
303 fn test_capabilities_validation() {
304 let mut caps = CapabilitiesConfig {
305 enable_subagents: true,
306 enabled_tools: Some(vec![BuiltinTools::ListDir]),
307 ..CapabilitiesConfig::default()
308 };
309 assert!(caps.validate().is_ok());
310
311 caps.disabled_tools = Some(vec![BuiltinTools::SearchDir]);
312 assert!(caps.validate().is_err());
313 }
314
315 #[test]
316
317 fn builtin_tools_serde_roundtrip_all_variants() {
318 let all = BuiltinTools::all_tools();
319 for tool in all {
320 let json = serde_json::to_string(tool).unwrap();
321 let parsed: BuiltinTools = serde_json::from_str(&json).unwrap();
322 assert_eq!(&parsed, tool, "Failed roundtrip for {tool:?}");
323 }
324 }
325
326 #[test]
327 fn builtin_tools_python_str_covers_all_variants() {
328 let expected = [
329 (BuiltinTools::ListDir, "list_directory"),
330 (BuiltinTools::SearchDir, "search_directory"),
331 (BuiltinTools::FindFile, "find_file"),
332 (BuiltinTools::ViewFile, "view_file"),
333 (BuiltinTools::CreateFile, "create_file"),
334 (BuiltinTools::EditFile, "edit_file"),
335 (BuiltinTools::RunCommand, "run_command"),
336 (BuiltinTools::AskQuestion, "ask_question"),
337 (BuiltinTools::StartSubagent, "start_subagent"),
338 (BuiltinTools::GenerateImage, "generate_image"),
339 (BuiltinTools::SearchWeb, "search_web"),
340 (BuiltinTools::ReadUrlContent, "read_url_content"),
341 (BuiltinTools::Finish, "finish"),
342 ];
343 for (variant, py_str) in expected {
344 assert_eq!(
345 variant.as_sdk_name(),
346 py_str,
347 "Python str mismatch for {variant:?}"
348 );
349 }
350 }
351
352 #[test]
353 fn builtin_tools_read_only_is_subset_of_all() {
354 let all = BuiltinTools::all_tools();
355 let read_only = BuiltinTools::read_only();
356 for tool in read_only {
357 assert!(
358 all.contains(tool),
359 "{tool:?} in read_only but not in all_tools"
360 );
361 }
362 }
363
364 #[test]
365 fn builtin_tools_read_only_excludes_write_tools() {
366 let read_only = BuiltinTools::read_only();
367 assert!(!read_only.contains(&BuiltinTools::CreateFile));
368 assert!(!read_only.contains(&BuiltinTools::EditFile));
369 assert!(!read_only.contains(&BuiltinTools::RunCommand));
370 assert!(!read_only.contains(&BuiltinTools::StartSubagent));
371 assert!(!read_only.contains(&BuiltinTools::GenerateImage));
372 assert!(!read_only.contains(&BuiltinTools::AskQuestion));
373 }
374
375 #[test]
376 fn capabilities_config_both_none_is_valid() {
377 let caps = CapabilitiesConfig::default();
378 assert!(caps.validate().is_ok());
379 }
380
381 #[test]
382 fn capabilities_config_only_disabled_is_valid() {
383 let caps = CapabilitiesConfig {
384 disabled_tools: Some(vec![BuiltinTools::RunCommand]),
385 compaction_threshold: Some(2000),
386 ..CapabilitiesConfig::default()
387 };
388 assert!(caps.validate().is_ok());
389 }
390
391 #[test]
392 fn capabilities_config_serde_roundtrip() {
393 let caps = CapabilitiesConfig {
394 enable_subagents: true,
395 enabled_tools: Some(vec![BuiltinTools::ViewFile, BuiltinTools::ListDir]),
396 compaction_threshold: Some(8000),
397 ..CapabilitiesConfig::default()
398 };
399 let json = serde_json::to_string(&caps).unwrap();
400 let parsed: CapabilitiesConfig = serde_json::from_str(&json).unwrap();
401 assert!(parsed.enable_subagents);
402 assert_eq!(parsed.enabled_tools.as_ref().unwrap().len(), 2);
403 assert_eq!(parsed.compaction_threshold, Some(8000));
404 }
405
406 #[test]
407 fn builtin_tools_snake_case_serde() {
408 let tool = BuiltinTools::StartSubagent;
410 let json = serde_json::to_string(&tool).unwrap();
411 assert_eq!(json, "\"start_subagent\"");
412
413 let tool = BuiltinTools::GenerateImage;
414 let json = serde_json::to_string(&tool).unwrap();
415 assert_eq!(json, "\"generate_image\"");
416 }
417
418 #[test]
419 fn capabilities_config_empty_enabled_list_vs_none() {
420 let caps_empty = CapabilitiesConfig {
423 enabled_tools: Some(vec![]),
424 ..CapabilitiesConfig::default()
425 };
426 assert!(caps_empty.validate().is_ok());
427 assert!(caps_empty.enabled_tools.as_ref().unwrap().is_empty());
428
429 let caps_none = CapabilitiesConfig::default();
430 assert!(caps_none.enabled_tools.is_none());
431 }
432
433 #[test]
434 fn capabilities_default_enables_subagents() {
435 let caps = CapabilitiesConfig::default();
437 assert!(
438 caps.enable_subagents,
439 "enable_subagents should default to true, matching the SDK"
440 );
441 }
442
443 #[test]
444 fn capabilities_serde_missing_enable_subagents_defaults_true() {
445 let json = r#"{"enabled_tools": ["view_file"]}"#;
447 let caps: CapabilitiesConfig = serde_json::from_str(json).unwrap();
448 assert!(
449 caps.enable_subagents,
450 "Missing enable_subagents in JSON should deserialize to true"
451 );
452 }
453
454 #[test]
455 fn capabilities_serde_explicit_false_is_respected() {
456 let json = r#"{"enable_subagents": false}"#;
457 let caps: CapabilitiesConfig = serde_json::from_str(json).unwrap();
458 assert!(!caps.enable_subagents, "Explicit false should be preserved");
459 }
460
461 #[test]
462 fn capabilities_with_tools_enables_subagents() {
463 let caps = CapabilitiesConfig::with_tools(vec![
464 BuiltinTools::ViewFile,
465 BuiltinTools::StartSubagent,
466 ]);
467 assert!(caps.enable_subagents);
468 assert_eq!(caps.enabled_tools.as_ref().unwrap().len(), 2);
469 }
470
471 #[test]
472 fn capabilities_full_enables_subagents() {
473 let caps = CapabilitiesConfig::full();
474 assert!(caps.enable_subagents);
475 assert!(caps.enabled_tools.is_none()); }
477
478 #[test]
479 fn capabilities_read_only_enables_subagents_but_no_start_subagent() {
480 let caps = CapabilitiesConfig::read_only();
481 assert!(caps.enable_subagents);
482 let tools = caps.enabled_tools.as_ref().unwrap();
483 assert!(
485 !tools.contains(&BuiltinTools::StartSubagent),
486 "read_only should not include StartSubagent in enabled_tools"
487 );
488 }
489
490 #[test]
491 fn capabilities_custom_tools_only_enables_subagents() {
492 let caps = CapabilitiesConfig::custom_tools_only();
493 assert!(caps.enable_subagents);
494 assert!(caps.enabled_tools.as_ref().unwrap().is_empty());
495 }
496
497 #[test]
498 fn start_subagent_in_all_tools_and_nondestructive() {
499 let all = BuiltinTools::all_tools();
500 assert!(
501 all.contains(&BuiltinTools::StartSubagent),
502 "all_tools() must include StartSubagent"
503 );
504 let nondestructive = BuiltinTools::nondestructive();
505 assert!(
506 nondestructive.contains(&BuiltinTools::StartSubagent),
507 "nondestructive() must include StartSubagent"
508 );
509 let read_only = BuiltinTools::read_only();
510 assert!(
511 !read_only.contains(&BuiltinTools::StartSubagent),
512 "read_only() must NOT include StartSubagent"
513 );
514 }
515
516 #[test]
518 fn builtin_tools_match_python_sdk() {
519 pyo3::Python::initialize();
520 pyo3::Python::attach(|py| {
521 crate::runtime::venv::configure_python_sys_path(py)
522 .unwrap_or_else(|e| panic!("Failed to configure python sys.path: {e}"));
523 let types_mod =
524 crate::runtime::py_scripts::import_serialized(py, "google.antigravity.types")
525 .expect("Failed to import google.antigravity.types");
526 let bt = types_mod
527 .getattr("BuiltinTools")
528 .expect("Failed to get BuiltinTools");
529 let builtins = py.import("builtins").expect("Failed to import builtins");
532 let members = builtins
533 .getattr("list")
534 .expect("Failed to get list")
535 .call1((bt,))
536 .expect("Failed to call list(BuiltinTools)");
537 let py_tools: Vec<String> = members
538 .try_iter()
539 .expect("Failed to iter members")
540 .map(|item| {
541 item.and_then(|v| v.getattr("value"))
542 .and_then(|v| v.extract::<String>())
543 })
544 .collect::<pyo3::PyResult<Vec<String>>>()
545 .expect("Failed to extract tool values");
546
547 let rust_tools: Vec<String> = BuiltinTools::all_tools()
548 .iter()
549 .map(|t| t.as_sdk_name().to_owned())
550 .collect();
551
552 assert_eq!(
553 rust_tools.len(),
554 py_tools.len(),
555 "Tool count mismatch: Rust has {}, Python has {}.\nRust: {rust_tools:?}\nPython: {py_tools:?}",
556 rust_tools.len(),
557 py_tools.len(),
558 );
559
560 for py_name in &py_tools {
561 assert!(
562 rust_tools.contains(py_name),
563 "Python SDK has tool '{py_name}' but Rust BuiltinTools does not"
564 );
565 }
566
567 for rust_name in &rust_tools {
568 assert!(
569 py_tools.contains(rust_name),
570 "Rust BuiltinTools has '{rust_name}' but Python SDK does not"
571 );
572 }
573 });
574 }
575
576 #[test]
578 fn capabilities_validate_rejects_both_enabled_and_disabled() {
579 let caps = CapabilitiesConfig {
580 enabled_tools: Some(vec![BuiltinTools::ViewFile]),
581 disabled_tools: Some(vec![BuiltinTools::RunCommand]),
582 ..CapabilitiesConfig::default()
583 };
584 assert!(caps.validate().is_err());
585 }
586}