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    /// Signal task completion.
42    #[serde(rename = "finish")]
43    Finish,
44}
45
46impl BuiltinTools {
47    #[must_use]
48    /// Returns tools that only read (no writes, no command execution).
49    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    /// Returns tools that cannot delete content (all except `RunCommand`).
60    #[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    /// Returns all builtin tools.
77    #[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    /// Returns tools that perform file read/write/create operations.
95    ///
96    /// These tools accept a file path argument and can be scoped to specific
97    /// workspace directories via `policy::workspace_only()`.
98    #[must_use]
99    pub const fn file_tools() -> &'static [Self] {
100        &[Self::ViewFile, Self::CreateFile, Self::EditFile]
101    }
102
103    /// Returns an empty tool list (no builtin tools).
104    #[must_use]
105    pub const fn none() -> &'static [Self] {
106        &[]
107    }
108
109    #[must_use]
110    /// Returns the Python SDK tool name string (e.g. `"list_directory"`).
111    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    /// Human-readable description of this builtin tool.
128    #[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/// Agent capability toggles: tool allowlists, subagent support, and compaction.
153///
154/// Construct with the named constructors ([`with_tools`](Self::with_tools),
155/// [`read_only`](Self::read_only), …) for common presets, or with the
156/// [`TypedBuilder`] for fine-grained control. Both preserve the SDK-matching
157/// defaults (subagents enabled, default image model):
158///
159/// ```
160/// use agy_bridge::config::CapabilitiesConfig;
161///
162/// let caps = CapabilitiesConfig::builder()
163///     .enable_subagents(false)
164///     .compaction_threshold(8000)
165///     .build();
166/// assert!(!caps.enable_subagents);
167/// assert_eq!(caps.compaction_threshold, Some(8000));
168/// ```
169#[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)]
170pub struct CapabilitiesConfig {
171    /// Whether this agent can spawn subagents.
172    #[serde(default = "super::default_true")]
173    #[builder(default = true)]
174    pub enable_subagents: bool,
175    /// If set, only these built-in tools are available (allowlist).
176    #[serde(default)]
177    #[builder(default, setter(strip_option))]
178    pub enabled_tools: Option<Vec<BuiltinTools>>,
179    /// If set, these built-in tools are removed (denylist).
180    #[serde(default)]
181    #[builder(default, setter(strip_option))]
182    pub disabled_tools: Option<Vec<BuiltinTools>>,
183    /// Token threshold that triggers conversation compaction.
184    #[builder(default, setter(strip_option))]
185    pub compaction_threshold: Option<usize>,
186    /// The model to use for image generation.
187    ///
188    /// This setting is a shorthand for `GeminiConfig.models.image_generation.name`.
189    /// If both are specified, the value in [`GeminiConfig`](super::GeminiConfig) takes precedence and
190    /// this field is ignored.
191    #[serde(default = "super::default_image_model")]
192    #[builder(default = DEFAULT_IMAGE_GENERATION_MODEL.to_owned(), setter(into))]
193    pub image_model: String,
194    /// Optional JSON schema string for the finish tool's structured output.
195    #[serde(default)]
196    #[builder(default, setter(into, strip_option))]
197    pub finish_tool_schema_json: Option<String>,
198}
199
200impl CapabilitiesConfig {
201    /// Create a capabilities config with only the specified tools enabled.
202    ///
203    /// Subagent support is enabled by default.
204    #[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    /// Create a capabilities config with all tools and subagent support.
213    #[must_use]
214    pub fn full() -> Self {
215        Self::default()
216    }
217
218    /// Create a capabilities config for read-only agents with subagent support.
219    #[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    /// Create a capabilities config with no builtin tools — only custom tools.
228    ///
229    /// Subagent support is still enabled.
230    #[must_use]
231    pub fn custom_tools_only() -> Self {
232        Self {
233            enabled_tools: Some(vec![]),
234            ..Self::default()
235        }
236    }
237
238    /// # Errors
239    ///
240    /// Returns an error if `enabled_tools` and `disabled_tools` are both provided.
241    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        // Delegate to the builder so field defaults have a single source of truth.
252        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        // Verify that serde serializes with snake_case as specified by the attribute.
383        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        // An explicitly empty enabled_tools list means "no tools enabled"
395        // whereas None means "use default set".
396        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        // Matches the Python SDK default: enable_subagents=True
410        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        // When enable_subagents is absent from JSON, it should default to true.
420        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()); // None = SDK defaults (all tools)
450    }
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        // read_only tools should NOT include StartSubagent
458        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    /// Verify our `BuiltinTools` enum exactly matches the Python SDK's tool names.
491    #[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            // BuiltinTools is a (str, Enum) subclass — use `list(BuiltinTools)`
504            // to iterate members, then extract `.value` from each.
505            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    /// Verify the `BuiltinTools` enum maps correctly to the validate function.
551    #[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}