Skip to main content

agy_bridge/config/
capabilities.rs

1//! Tool capability configuration.
2
3use 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    /// List files and subdirectories.
12    #[serde(rename = "list_directory")]
13    ListDir,
14    /// Regex search within directory contents.
15    #[serde(rename = "search_directory")]
16    SearchDir,
17    /// Find files by name pattern.
18    #[serde(rename = "find_file")]
19    FindFile,
20    /// Read file contents.
21    #[serde(rename = "view_file")]
22    ViewFile,
23    /// Create a new file.
24    #[serde(rename = "create_file")]
25    CreateFile,
26    /// Edit an existing file.
27    #[serde(rename = "edit_file")]
28    EditFile,
29    /// Execute a shell command.
30    #[serde(rename = "run_command")]
31    RunCommand,
32    /// Ask the user a question.
33    #[serde(rename = "ask_question")]
34    AskQuestion,
35    /// Spawn a subagent.
36    #[serde(rename = "start_subagent")]
37    StartSubagent,
38    /// Generate images from text prompts.
39    #[serde(rename = "generate_image")]
40    GenerateImage,
41    /// Search the web for information.
42    #[serde(rename = "search_web")]
43    SearchWeb,
44    /// Fetch content from a URL via HTTP request.
45    #[serde(rename = "read_url_content")]
46    ReadUrlContent,
47    /// Signal task completion.
48    #[serde(rename = "finish")]
49    Finish,
50}
51
52impl BuiltinTools {
53    #[must_use]
54    /// Returns tools that only read (no writes, no command execution).
55    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    /// Returns tools that cannot delete content (all except `RunCommand`).
67    #[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    /// Returns all builtin tools.
86    #[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    /// Returns tools that perform file read/write/create operations.
106    ///
107    /// These tools accept a file path argument and can be scoped to specific
108    /// workspace directories via `policy::workspace_only()`.
109    #[must_use]
110    pub const fn file_tools() -> &'static [Self] {
111        &[Self::ViewFile, Self::CreateFile, Self::EditFile]
112    }
113
114    /// Returns an empty tool list (no builtin tools).
115    #[must_use]
116    pub const fn none() -> &'static [Self] {
117        &[]
118    }
119
120    #[must_use]
121    /// Returns the Python SDK tool name string (e.g. `"list_directory"`).
122    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    /// Human-readable description of this builtin tool.
141    #[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/// Agent capability toggles: tool allowlists, subagent support, and compaction.
168///
169/// Construct with the named constructors ([`with_tools`](Self::with_tools),
170/// [`read_only`](Self::read_only), …) for common presets, or with the
171/// [`TypedBuilder`] for fine-grained control. Both preserve the SDK-matching
172/// defaults (subagents enabled, default image model):
173///
174/// ```
175/// use agy_bridge::config::CapabilitiesConfig;
176///
177/// let caps = CapabilitiesConfig::builder()
178///     .enable_subagents(false)
179///     .compaction_threshold(8000)
180///     .build();
181/// assert!(!caps.enable_subagents);
182/// assert_eq!(caps.compaction_threshold, Some(8000));
183/// ```
184#[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)]
185pub struct CapabilitiesConfig {
186    /// Whether this agent can spawn subagents.
187    #[serde(default = "super::default_true")]
188    #[builder(default = true)]
189    pub enable_subagents: bool,
190    /// If set, only these built-in tools are available (allowlist).
191    #[serde(default)]
192    #[builder(default, setter(strip_option))]
193    pub enabled_tools: Option<Vec<BuiltinTools>>,
194    /// If set, these built-in tools are removed (denylist).
195    #[serde(default)]
196    #[builder(default, setter(strip_option))]
197    pub disabled_tools: Option<Vec<BuiltinTools>>,
198    /// Token threshold that triggers conversation compaction.
199    #[builder(default, setter(strip_option))]
200    pub compaction_threshold: Option<usize>,
201    /// The model to use for image generation.
202    ///
203    /// This setting is a shorthand for `GeminiConfig.models.image_generation.name`.
204    /// If both are specified, the value in [`GeminiConfig`](super::GeminiConfig) takes precedence and
205    /// this field is ignored.
206    #[serde(default = "super::default_image_model")]
207    #[builder(default = DEFAULT_IMAGE_GENERATION_MODEL.to_owned(), setter(into))]
208    pub image_model: String,
209    /// Optional JSON schema string for the finish tool's structured output.
210    #[serde(default)]
211    #[builder(default, setter(into, strip_option))]
212    pub finish_tool_schema_json: Option<String>,
213}
214
215impl CapabilitiesConfig {
216    /// Create a capabilities config with only the specified tools enabled.
217    ///
218    /// Subagent support is enabled by default.
219    #[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    /// Create a capabilities config with all tools and subagent support.
228    #[must_use]
229    pub fn full() -> Self {
230        Self::default()
231    }
232
233    /// Create a capabilities config for read-only agents with subagent support.
234    #[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    /// Create a capabilities config with no builtin tools — only custom tools.
243    ///
244    /// Subagent support is still enabled.
245    #[must_use]
246    pub fn custom_tools_only() -> Self {
247        Self {
248            enabled_tools: Some(vec![]),
249            ..Self::default()
250        }
251    }
252
253    /// # Errors
254    ///
255    /// Returns an error if `enabled_tools` and `disabled_tools` are both provided.
256    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        // Delegate to the builder so field defaults have a single source of truth.
267        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        // Verify that serde serializes with snake_case as specified by the attribute.
409        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        // An explicitly empty enabled_tools list means "no tools enabled"
421        // whereas None means "use default set".
422        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        // Matches the Python SDK default: enable_subagents=True
436        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        // When enable_subagents is absent from JSON, it should default to true.
446        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()); // None = SDK defaults (all tools)
476    }
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        // read_only tools should NOT include StartSubagent
484        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    /// Verify our `BuiltinTools` enum exactly matches the Python SDK's tool names.
517    #[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            // BuiltinTools is a (str, Enum) subclass — use `list(BuiltinTools)`
530            // to iterate members, then extract `.value` from each.
531            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    /// Verify the `BuiltinTools` enum maps correctly to the validate function.
577    #[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}