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