agy_bridge/config/
capabilities.rs

1//! Tool capability configuration.
2
3use serde::{Deserialize, Serialize};
4
5use super::DEFAULT_IMAGE_GENERATION_MODEL;
6
7#[non_exhaustive]
8#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
9#[serde(rename_all = "snake_case")]
10pub enum BuiltinTools {
11    /// List files and subdirectories.
12    ListDir,
13    /// Regex search within directory contents.
14    SearchDir,
15    /// Find files by name pattern.
16    FindFile,
17    /// Read file contents.
18    ViewFile,
19    /// Create a new file.
20    CreateFile,
21    /// Edit an existing file.
22    EditFile,
23    /// Execute a shell command.
24    RunCommand,
25    /// Ask the user a question.
26    AskQuestion,
27    /// Spawn a subagent.
28    StartSubagent,
29    /// Generate images from text prompts.
30    GenerateImage,
31    /// Signal task completion.
32    Finish,
33}
34
35impl BuiltinTools {
36    #[must_use]
37    /// Returns tools that only read (no writes, no command execution).
38    pub const fn read_only() -> &'static [Self] {
39        &[
40            Self::ListDir,
41            Self::SearchDir,
42            Self::FindFile,
43            Self::ViewFile,
44            Self::Finish,
45        ]
46    }
47
48    /// Returns tools that cannot delete content (all except `RunCommand`).
49    #[must_use]
50    pub const fn nondestructive() -> &'static [Self] {
51        &[
52            Self::ListDir,
53            Self::SearchDir,
54            Self::FindFile,
55            Self::ViewFile,
56            Self::CreateFile,
57            Self::EditFile,
58            Self::AskQuestion,
59            Self::StartSubagent,
60            Self::GenerateImage,
61            Self::Finish,
62        ]
63    }
64
65    /// Returns all builtin tools.
66    #[must_use]
67    pub const fn all_tools() -> &'static [Self] {
68        &[
69            Self::ListDir,
70            Self::SearchDir,
71            Self::FindFile,
72            Self::ViewFile,
73            Self::CreateFile,
74            Self::EditFile,
75            Self::RunCommand,
76            Self::AskQuestion,
77            Self::StartSubagent,
78            Self::GenerateImage,
79            Self::Finish,
80        ]
81    }
82
83    /// Returns tools that perform file read/write/create operations.
84    ///
85    /// These tools accept a file path argument and can be scoped to specific
86    /// workspace directories via `policy::workspace_only()`.
87    #[must_use]
88    pub const fn file_tools() -> &'static [Self] {
89        &[Self::ViewFile, Self::CreateFile, Self::EditFile]
90    }
91
92    /// Returns an empty tool list (no builtin tools).
93    #[must_use]
94    pub const fn none() -> &'static [Self] {
95        &[]
96    }
97
98    #[must_use]
99    /// Returns the Python SDK tool name string (e.g. `"list_directory"`).
100    pub const fn as_sdk_name(&self) -> &'static str {
101        match self {
102            Self::ListDir => "list_directory",
103            Self::SearchDir => "search_directory",
104            Self::FindFile => "find_file",
105            Self::ViewFile => "view_file",
106            Self::CreateFile => "create_file",
107            Self::EditFile => "edit_file",
108            Self::RunCommand => "run_command",
109            Self::AskQuestion => "ask_question",
110            Self::StartSubagent => "start_subagent",
111            Self::GenerateImage => "generate_image",
112            Self::Finish => "finish",
113        }
114    }
115}
116
117impl std::fmt::Display for BuiltinTools {
118    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119        f.write_str(self.as_sdk_name())
120    }
121}
122
123/// Agent capability toggles: tool allowlists, subagent support, and compaction.
124#[derive(Debug, Clone, Serialize, Deserialize)]
125pub struct CapabilitiesConfig {
126    /// Whether this agent can spawn subagents.
127    #[serde(default = "super::default_true")]
128    pub enable_subagents: bool,
129    /// If set, only these built-in tools are available (allowlist).
130    #[serde(default)]
131    pub enabled_tools: Option<Vec<BuiltinTools>>,
132    /// If set, these built-in tools are removed (denylist).
133    #[serde(default)]
134    pub disabled_tools: Option<Vec<BuiltinTools>>,
135    /// Token threshold that triggers conversation compaction.
136    pub compaction_threshold: Option<usize>,
137    /// The model to use for image generation.
138    ///
139    /// This setting is a shorthand for `GeminiConfig.models.image_generation.name`.
140    /// If both are specified, the value in [`GeminiConfig`](super::GeminiConfig) takes precedence and
141    /// this field is ignored.
142    #[serde(default = "super::default_image_model")]
143    pub image_model: String,
144    /// Optional JSON schema string for the finish tool's structured output.
145    #[serde(default)]
146    pub finish_tool_schema_json: Option<String>,
147}
148
149impl CapabilitiesConfig {
150    /// Create a capabilities config with only the specified tools enabled.
151    ///
152    /// Subagent support is enabled by default.
153    #[must_use]
154    pub fn with_tools(tools: Vec<BuiltinTools>) -> Self {
155        Self {
156            enabled_tools: Some(tools),
157            ..Self::default()
158        }
159    }
160
161    /// Create a capabilities config with all tools and subagent support.
162    #[must_use]
163    pub fn full() -> Self {
164        Self::default()
165    }
166
167    /// Create a capabilities config for read-only agents with subagent support.
168    #[must_use]
169    pub fn read_only() -> Self {
170        Self {
171            enabled_tools: Some(BuiltinTools::read_only().to_vec()),
172            ..Self::default()
173        }
174    }
175
176    /// Create a capabilities config with no builtin tools — only custom tools.
177    ///
178    /// Subagent support is still enabled.
179    #[must_use]
180    pub fn custom_tools_only() -> Self {
181        Self {
182            enabled_tools: Some(vec![]),
183            ..Self::default()
184        }
185    }
186
187    /// # Errors
188    ///
189    /// Returns an error if `enabled_tools` and `disabled_tools` are both provided.
190    pub const fn validate(&self) -> Result<(), &'static str> {
191        if self.enabled_tools.is_some() && self.disabled_tools.is_some() {
192            return Err("enabled_tools and disabled_tools are mutually exclusive");
193        }
194        Ok(())
195    }
196}
197
198impl Default for CapabilitiesConfig {
199    fn default() -> Self {
200        Self {
201            enable_subagents: true,
202            enabled_tools: None,
203            disabled_tools: None,
204            compaction_threshold: None,
205            image_model: DEFAULT_IMAGE_GENERATION_MODEL.to_owned(),
206            finish_tool_schema_json: None,
207        }
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use pyo3::types::PyAnyMethods;
214
215    use super::*;
216
217    #[test]
218    fn test_builtin_tools() {
219        let read_only = BuiltinTools::read_only();
220        assert_eq!(read_only.len(), 5);
221        assert!(read_only.contains(&BuiltinTools::ListDir));
222        assert!(read_only.contains(&BuiltinTools::Finish));
223        assert!(!read_only.contains(&BuiltinTools::CreateFile));
224
225        let all = BuiltinTools::all_tools();
226        assert_eq!(all.len(), 11);
227        assert!(all.contains(&BuiltinTools::CreateFile));
228        assert!(all.contains(&BuiltinTools::Finish));
229
230        assert_eq!(BuiltinTools::ListDir.as_sdk_name(), "list_directory");
231    }
232
233    #[test]
234    fn test_capabilities_validation() {
235        let mut caps = CapabilitiesConfig {
236            enable_subagents: true,
237            enabled_tools: Some(vec![BuiltinTools::ListDir]),
238            ..CapabilitiesConfig::default()
239        };
240        assert!(caps.validate().is_ok());
241
242        caps.disabled_tools = Some(vec![BuiltinTools::SearchDir]);
243        assert!(caps.validate().is_err());
244    }
245
246    #[test]
247
248    fn builtin_tools_serde_roundtrip_all_variants() {
249        let all = BuiltinTools::all_tools();
250        for tool in all {
251            let json = serde_json::to_string(tool).unwrap();
252            let parsed: BuiltinTools = serde_json::from_str(&json).unwrap();
253            assert_eq!(&parsed, tool, "Failed roundtrip for {tool:?}");
254        }
255    }
256
257    #[test]
258    fn builtin_tools_python_str_covers_all_variants() {
259        let expected = [
260            (BuiltinTools::ListDir, "list_directory"),
261            (BuiltinTools::SearchDir, "search_directory"),
262            (BuiltinTools::FindFile, "find_file"),
263            (BuiltinTools::ViewFile, "view_file"),
264            (BuiltinTools::CreateFile, "create_file"),
265            (BuiltinTools::EditFile, "edit_file"),
266            (BuiltinTools::RunCommand, "run_command"),
267            (BuiltinTools::AskQuestion, "ask_question"),
268            (BuiltinTools::StartSubagent, "start_subagent"),
269            (BuiltinTools::GenerateImage, "generate_image"),
270            (BuiltinTools::Finish, "finish"),
271        ];
272        for (variant, py_str) in expected {
273            assert_eq!(
274                variant.as_sdk_name(),
275                py_str,
276                "Python str mismatch for {variant:?}"
277            );
278        }
279    }
280
281    #[test]
282    fn builtin_tools_read_only_is_subset_of_all() {
283        let all = BuiltinTools::all_tools();
284        let read_only = BuiltinTools::read_only();
285        for tool in read_only {
286            assert!(
287                all.contains(tool),
288                "{tool:?} in read_only but not in all_tools"
289            );
290        }
291    }
292
293    #[test]
294    fn builtin_tools_read_only_excludes_write_tools() {
295        let read_only = BuiltinTools::read_only();
296        assert!(!read_only.contains(&BuiltinTools::CreateFile));
297        assert!(!read_only.contains(&BuiltinTools::EditFile));
298        assert!(!read_only.contains(&BuiltinTools::RunCommand));
299        assert!(!read_only.contains(&BuiltinTools::StartSubagent));
300        assert!(!read_only.contains(&BuiltinTools::GenerateImage));
301        assert!(!read_only.contains(&BuiltinTools::AskQuestion));
302    }
303
304    #[test]
305    fn capabilities_config_both_none_is_valid() {
306        let caps = CapabilitiesConfig::default();
307        assert!(caps.validate().is_ok());
308    }
309
310    #[test]
311    fn capabilities_config_only_disabled_is_valid() {
312        let caps = CapabilitiesConfig {
313            disabled_tools: Some(vec![BuiltinTools::RunCommand]),
314            compaction_threshold: Some(2000),
315            ..CapabilitiesConfig::default()
316        };
317        assert!(caps.validate().is_ok());
318    }
319
320    #[test]
321    fn capabilities_config_serde_roundtrip() {
322        let caps = CapabilitiesConfig {
323            enable_subagents: true,
324            enabled_tools: Some(vec![BuiltinTools::ViewFile, BuiltinTools::ListDir]),
325            compaction_threshold: Some(8000),
326            ..CapabilitiesConfig::default()
327        };
328        let json = serde_json::to_string(&caps).unwrap();
329        let parsed: CapabilitiesConfig = serde_json::from_str(&json).unwrap();
330        assert!(parsed.enable_subagents);
331        assert_eq!(parsed.enabled_tools.as_ref().unwrap().len(), 2);
332        assert_eq!(parsed.compaction_threshold, Some(8000));
333    }
334
335    #[test]
336    fn builtin_tools_snake_case_serde() {
337        // Verify that serde serializes with snake_case as specified by the attribute.
338        let tool = BuiltinTools::StartSubagent;
339        let json = serde_json::to_string(&tool).unwrap();
340        assert_eq!(json, "\"start_subagent\"");
341
342        let tool = BuiltinTools::GenerateImage;
343        let json = serde_json::to_string(&tool).unwrap();
344        assert_eq!(json, "\"generate_image\"");
345    }
346
347    #[test]
348    fn capabilities_config_empty_enabled_list_vs_none() {
349        // An explicitly empty enabled_tools list means "no tools enabled"
350        // whereas None means "use default set".
351        let caps_empty = CapabilitiesConfig {
352            enabled_tools: Some(vec![]),
353            ..CapabilitiesConfig::default()
354        };
355        assert!(caps_empty.validate().is_ok());
356        assert!(caps_empty.enabled_tools.as_ref().unwrap().is_empty());
357
358        let caps_none = CapabilitiesConfig::default();
359        assert!(caps_none.enabled_tools.is_none());
360    }
361
362    #[test]
363    fn capabilities_default_enables_subagents() {
364        // Matches the Python SDK default: enable_subagents=True
365        let caps = CapabilitiesConfig::default();
366        assert!(
367            caps.enable_subagents,
368            "enable_subagents should default to true, matching the SDK"
369        );
370    }
371
372    #[test]
373    fn capabilities_serde_missing_enable_subagents_defaults_true() {
374        // When enable_subagents is absent from JSON, it should default to true.
375        let json = r#"{"enabled_tools": ["view_file"]}"#;
376        let caps: CapabilitiesConfig = serde_json::from_str(json).unwrap();
377        assert!(
378            caps.enable_subagents,
379            "Missing enable_subagents in JSON should deserialize to true"
380        );
381    }
382
383    #[test]
384    fn capabilities_serde_explicit_false_is_respected() {
385        let json = r#"{"enable_subagents": false}"#;
386        let caps: CapabilitiesConfig = serde_json::from_str(json).unwrap();
387        assert!(!caps.enable_subagents, "Explicit false should be preserved");
388    }
389
390    #[test]
391    fn capabilities_with_tools_enables_subagents() {
392        let caps = CapabilitiesConfig::with_tools(vec![
393            BuiltinTools::ViewFile,
394            BuiltinTools::StartSubagent,
395        ]);
396        assert!(caps.enable_subagents);
397        assert_eq!(caps.enabled_tools.as_ref().unwrap().len(), 2);
398    }
399
400    #[test]
401    fn capabilities_full_enables_subagents() {
402        let caps = CapabilitiesConfig::full();
403        assert!(caps.enable_subagents);
404        assert!(caps.enabled_tools.is_none()); // None = SDK defaults (all tools)
405    }
406
407    #[test]
408    fn capabilities_read_only_enables_subagents_but_no_start_subagent() {
409        let caps = CapabilitiesConfig::read_only();
410        assert!(caps.enable_subagents);
411        let tools = caps.enabled_tools.as_ref().unwrap();
412        // read_only tools should NOT include StartSubagent
413        assert!(
414            !tools.contains(&BuiltinTools::StartSubagent),
415            "read_only should not include StartSubagent in enabled_tools"
416        );
417    }
418
419    #[test]
420    fn capabilities_custom_tools_only_enables_subagents() {
421        let caps = CapabilitiesConfig::custom_tools_only();
422        assert!(caps.enable_subagents);
423        assert!(caps.enabled_tools.as_ref().unwrap().is_empty());
424    }
425
426    #[test]
427    fn start_subagent_in_all_tools_and_nondestructive() {
428        let all = BuiltinTools::all_tools();
429        assert!(
430            all.contains(&BuiltinTools::StartSubagent),
431            "all_tools() must include StartSubagent"
432        );
433        let nondestructive = BuiltinTools::nondestructive();
434        assert!(
435            nondestructive.contains(&BuiltinTools::StartSubagent),
436            "nondestructive() must include StartSubagent"
437        );
438        let read_only = BuiltinTools::read_only();
439        assert!(
440            !read_only.contains(&BuiltinTools::StartSubagent),
441            "read_only() must NOT include StartSubagent"
442        );
443    }
444
445    /// Verify our `BuiltinTools` enum exactly matches the Python SDK's tool names.
446    #[test]
447    fn builtin_tools_match_python_sdk() {
448        pyo3::Python::initialize();
449        pyo3::Python::attach(|py| {
450            crate::runtime::venv::configure_python_sys_path(py)
451                .unwrap_or_else(|e| panic!("Failed to configure python sys.path: {e}"));
452            let types_mod = py
453                .import("google.antigravity.types")
454                .expect("Failed to import google.antigravity.types");
455            let bt = types_mod
456                .getattr("BuiltinTools")
457                .expect("Failed to get BuiltinTools");
458            // BuiltinTools is a (str, Enum) subclass — use `list(BuiltinTools)`
459            // to iterate members, then extract `.value` from each.
460            let builtins = py.import("builtins").expect("Failed to import builtins");
461            let members = builtins
462                .getattr("list")
463                .expect("Failed to get list")
464                .call1((bt,))
465                .expect("Failed to call list(BuiltinTools)");
466            let py_tools: Vec<String> = members
467                .try_iter()
468                .expect("Failed to iter members")
469                .map(|item| {
470                    item.and_then(|v| v.getattr("value"))
471                        .and_then(|v| v.extract::<String>())
472                })
473                .collect::<pyo3::PyResult<Vec<String>>>()
474                .expect("Failed to extract tool values");
475
476            let rust_tools: Vec<String> = BuiltinTools::all_tools()
477                .iter()
478                .map(|t| t.as_sdk_name().to_owned())
479                .collect();
480
481            assert_eq!(
482                rust_tools.len(),
483                py_tools.len(),
484                "Tool count mismatch: Rust has {}, Python has {}.\nRust: {rust_tools:?}\nPython: {py_tools:?}",
485                rust_tools.len(),
486                py_tools.len(),
487            );
488
489            for py_name in &py_tools {
490                assert!(
491                    rust_tools.contains(py_name),
492                    "Python SDK has tool '{py_name}' but Rust BuiltinTools does not"
493                );
494            }
495
496            for rust_name in &rust_tools {
497                assert!(
498                    py_tools.contains(rust_name),
499                    "Rust BuiltinTools has '{rust_name}' but Python SDK does not"
500                );
501            }
502        });
503    }
504
505    /// Verify the `BuiltinTools` enum maps correctly to the validate function.
506    #[test]
507    fn capabilities_validate_rejects_both_enabled_and_disabled() {
508        let caps = CapabilitiesConfig {
509            enabled_tools: Some(vec![BuiltinTools::ViewFile]),
510            disabled_tools: Some(vec![BuiltinTools::RunCommand]),
511            ..CapabilitiesConfig::default()
512        };
513        assert!(caps.validate().is_err());
514    }
515}
agy_bridge/config/capabilities.rs

agy_bridge/config/
capabilities.rs
