agy_bridge/config/
agent.rs

1//! Agent configuration, system instructions, and local agent config.
2
3use std::path::PathBuf;
4
5use serde::{Deserialize, Serialize};
6use typed_builder::TypedBuilder;
7
8use super::{
9    DEFAULT_MODEL, capabilities::CapabilitiesConfig, mcp::McpServer, models::GeminiConfig,
10};
11use crate::{
12    hooks::HookEntry, policies::PolicyRule, tools::ToolDefinition, triggers::TriggerEntry,
13};
14
15/// A section within a system instruction, with a label and body text.
16#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
17pub struct SystemInstructionSection {
18    /// The content of the section.
19    pub content: String,
20    /// The title/label for this section.
21    #[serde(default = "default_section_title")]
22    pub title: String,
23}
24
25fn default_section_title() -> String {
26    "user_system_instructions".to_owned()
27}
28
29fn default_model_name() -> String {
30    DEFAULT_MODEL.to_owned()
31}
32
33/// System instruction configuration, mirroring the Python SDK's union type.
34///
35/// Uses internal tagging via `#[serde(untagged)]` so each variant is
36/// distinguishable by its `"mode"` field in JSON (e.g. `{"mode": "Custom", "text": "..."}`).
37#[non_exhaustive]
38#[derive(Debug, Clone, Serialize, Deserialize)]
39#[serde(untagged)]
40pub enum SystemInstructions {
41    /// Completely replace the default system instructions (advanced usage).
42    Custom(String),
43    /// Override identity and/or append sections to the defaults (recommended).
44    Templated {
45        /// Optional identity string that replaces the agent's default persona.
46        #[serde(default)]
47        identity: Option<String>,
48        /// Sections appended to the default system instructions.
49        #[serde(default)]
50        sections: Vec<SystemInstructionSection>,
51    },
52}
53
54impl SystemInstructions {
55    /// Create custom system instructions from a plain text string.
56    #[must_use]
57    pub fn custom(text: impl Into<String>) -> Self {
58        Self::Custom(text.into())
59    }
60}
61
62impl From<&str> for SystemInstructions {
63    fn from(s: &str) -> Self {
64        Self::custom(s)
65    }
66}
67
68impl From<String> for SystemInstructions {
69    fn from(s: String) -> Self {
70        Self::custom(s)
71    }
72}
73
74// ─── JSON Schema newtype ──────────────────────────────────────────────────────────────────
75
76/// A JSON Schema definition for structured output.
77#[derive(Debug, Clone, Serialize, Deserialize)]
78#[serde(transparent)]
79pub struct JsonSchema(serde_json::Value);
80
81impl JsonSchema {
82    #[must_use]
83    /// Wrap a raw `serde_json::Value` as a `JsonSchema`.
84    pub const fn new(value: serde_json::Value) -> Self {
85        Self(value)
86    }
87
88    #[must_use]
89    /// Return a reference to the inner JSON value.
90    pub const fn as_value(&self) -> &serde_json::Value {
91        &self.0
92    }
93
94    /// Validate that the schema is structurally sound.
95    ///
96    /// Currently checks that the top-level value is a JSON object (i.e.
97    /// `serde_json::Value::Object`), which is the minimum requirement for a
98    /// valid JSON Schema.
99    ///
100    /// # Errors
101    ///
102    /// Returns a static error message if the schema is not an object.
103    pub fn validate(&self) -> Result<(), &'static str> {
104        if self.0.is_object() {
105            Ok(())
106        } else {
107            Err("JSON Schema must be a JSON object at the top level")
108        }
109    }
110}
111
112// ─── AgentConfig ─────────────────────────────────────────────────────────────
113
114/// Full configuration for creating an agent.
115///
116/// Covers model selection, system instructions, capabilities, tools,
117/// policies, hooks, MCP servers, structured output, and Gemini backend
118/// settings. All fields have sensible defaults.
119///
120/// # Construction patterns
121///
122/// `AgentConfig` deliberately supports **two** construction paths:
123///
124/// 1. **[`TypedBuilder`]** — ergonomic chained construction with `impl
125///    IntoIterator` setters for collection fields. Preferred for
126///    programmatic use:
127///    ```
128///    # use agy_bridge::config::AgentConfig;
129///    let config = AgentConfig::builder()
130///        .model("gemini-3.5-flash")
131///        .build();
132///    ```
133///
134/// 2. **Struct literal with `..Default::default()`** — convenient for
135///    deserialization (`serde`), config files, and framework code that
136///    already has fully-formed values:
137///    ```
138///    # use agy_bridge::config::AgentConfig;
139///    let config = AgentConfig {
140///        model: "gemini-3.5-flash".into(),
141///        ..AgentConfig::default()
142///    };
143///    ```
144///
145/// Both paths are supported intentionally. The builder provides ergonomic
146/// setters (e.g. accepting `impl IntoIterator` for collection fields),
147/// while struct literals enable direct field access for serialization
148/// roundtrips and downstream framework integration.
149#[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)]
150#[builder(field_defaults(default))]
151pub struct AgentConfig {
152    /// The model name (e.g. `"gemini-3.5-flash"`).
153    #[serde(default = "default_model_name")]
154    #[builder(default = DEFAULT_MODEL.to_owned(), setter(into))]
155    pub model: String,
156    /// API key. Falls back to `GEMINI_API_KEY` env var if `None`.
157    #[serde(default)]
158    #[builder(setter(into, strip_option))]
159    pub api_key: Option<String>,
160    /// Optional system instructions (custom text or templated sections).
161    #[builder(setter(into, strip_option))]
162    pub system_instructions: Option<SystemInstructions>,
163    #[serde(default)]
164    /// Agent capability toggles (tool lists, subagents, compaction).
165    #[builder(setter(strip_option))]
166    pub capabilities: Option<CapabilitiesConfig>,
167    #[serde(default, skip_serializing_if = "Vec::is_empty")]
168    /// Workspace directories the agent can access.
169    ///
170    /// When empty (the default), the Python SDK's own default of
171    /// `[os.getcwd()]` applies. Set explicitly to override.
172    #[builder(setter(transform = |v: impl IntoIterator<Item = impl Into<PathBuf>>| v.into_iter().map(Into::into).collect()))]
173    pub workspaces: Vec<PathBuf>,
174    #[serde(default, skip_serializing_if = "Vec::is_empty")]
175    /// Custom tool definitions exposed to the agent.
176    #[builder(setter(transform = |v: impl IntoIterator<Item = impl Into<ToolDefinition>>| v.into_iter().map(Into::into).collect()))]
177    pub tools: Vec<ToolDefinition>,
178    #[serde(default = "default_policies")]
179    /// Policy rules evaluated before each tool call.
180    #[builder(default = default_policies(), setter(transform = |v: impl IntoIterator<Item = impl Into<PolicyRule>>| v.into_iter().map(Into::into).collect()))]
181    pub policies: Vec<PolicyRule>,
182    #[serde(default, skip_serializing_if = "Vec::is_empty")]
183    /// Event-driven triggers attached to this agent.
184    #[builder(setter(transform = |v: impl IntoIterator<Item = impl Into<TriggerEntry>>| v.into_iter().map(Into::into).collect()))]
185    pub triggers: Vec<TriggerEntry>,
186    #[serde(default, skip_serializing_if = "Vec::is_empty")]
187    /// Lifecycle hooks (pre-turn, post-turn, pre-tool, etc.).
188    #[builder(setter(transform = |v: impl IntoIterator<Item = impl Into<HookEntry>>| v.into_iter().map(Into::into).collect()))]
189    pub hooks: Vec<HookEntry>,
190    #[serde(
191        default,
192        skip_serializing_if = "Vec::is_empty",
193        rename = "skills_paths"
194    )]
195    /// Paths to skill instruction files loaded into the agent.
196    ///
197    /// Serializes as `"skills_paths"` to match the Python SDK field name.
198    #[builder(setter(transform = |v: impl IntoIterator<Item = impl Into<PathBuf>>| v.into_iter().map(Into::into).collect()))]
199    pub skills: Vec<PathBuf>,
200
201    /// MCP server configurations.
202    #[serde(default, skip_serializing_if = "Vec::is_empty")]
203    #[builder(setter(transform = |v: impl IntoIterator<Item = impl Into<McpServer>>| v.into_iter().map(Into::into).collect()))]
204    pub mcp_servers: Vec<McpServer>,
205    /// Pre-existing conversation ID to resume.
206    #[serde(default)]
207    #[builder(setter(into, strip_option))]
208    pub conversation_id: Option<String>,
209    /// Directory where conversation state is saved.
210    #[serde(default)]
211    #[builder(setter(into, strip_option))]
212    pub save_dir: Option<PathBuf>,
213    /// Application data directory.
214    #[serde(default)]
215    #[builder(setter(into, strip_option))]
216    pub app_data_dir: Option<PathBuf>,
217    /// Optional JSON schema for structured responses.
218    #[serde(default)]
219    #[builder(setter(strip_option))]
220    pub response_schema: Option<JsonSchema>,
221    /// Gemini model backend configuration.
222    ///
223    /// Controls per-model API keys, model selection per capability,
224    /// and generation parameters such as `thinking_level`.
225    ///
226    /// Serializes as `"gemini_config"` to match the Python SDK field name.
227    #[serde(default, rename = "gemini_config")]
228    #[builder(setter(strip_option))]
229    pub gemini: Option<GeminiConfig>,
230    /// Maximum number of quota retry attempts before giving up.
231    ///
232    /// If `None`, defaults to 0 (no retries).
233    #[serde(default)]
234    #[builder(setter(into, strip_option))]
235    pub max_quota_retries: Option<u32>,
236}
237
238impl Default for AgentConfig {
239    fn default() -> Self {
240        Self::builder().build()
241    }
242}
243
244impl AgentConfig {
245    /// Resolve the effective API key using the same priority chain as the
246    /// Python SDK's `LocalAgentConfig` → `_build_harness_config`:
247    ///
248    /// 1. Per-model key (`gemini.models.default.api_key`)
249    /// 2. Shared `GeminiConfig` key (`gemini.api_key`)
250    /// 3. Top-level shorthand (`api_key`)
251    /// 4. `$GEMINI_API_KEY` environment variable
252    #[must_use]
253    pub fn effective_api_key(&self) -> Option<String> {
254        self.gemini
255            .as_ref()
256            .and_then(|g| g.models.default.api_key.clone())
257            .or_else(|| self.gemini.as_ref().and_then(|g| g.api_key.clone()))
258            .or_else(|| self.api_key.clone())
259            .or_else(|| std::env::var("GEMINI_API_KEY").ok())
260    }
261
262    /// Returns the names of all explicitly registered custom tools.
263    /// To get the full list of tools including built-ins, combine this
264    /// with `capabilities.enabled_tools` or examine `tools` + default semantics.
265    #[must_use]
266    pub fn custom_tool_names(&self) -> Vec<String> {
267        self.tools.iter().map(|t| t.name.clone()).collect()
268    }
269}
270
271// ─── LocalAgentConfig ────────────────────────────────────────────────────────
272
273/// Configuration for a local (on-device) agent, mirroring the Python SDK's
274/// `LocalAgentConfig`.
275///
276/// Wraps the standard [`AgentConfig`] via `#[serde(flatten)]`. The Python SDK's
277/// `LocalAgentConfig` extends the base `AgentConfig` with override defaults
278/// (e.g. `policies = confirm_run_command()`, `workspaces = [cwd]`), which
279/// our `AgentConfig` already matches.
280#[derive(Debug, Clone, Serialize, Deserialize, Default)]
281pub struct LocalAgentConfig {
282    /// The base agent configuration.
283    #[serde(flatten)]
284    pub agent: AgentConfig,
285}
286
287impl LocalAgentConfig {
288    /// Create a new `LocalAgentConfig` wrapping the given agent configuration.
289    #[must_use]
290    pub const fn new(agent: AgentConfig) -> Self {
291        Self { agent }
292    }
293}
294
295impl From<AgentConfig> for LocalAgentConfig {
296    fn from(agent: AgentConfig) -> Self {
297        Self::new(agent)
298    }
299}
300
301/// Matches the Python SDK's `LocalAgentConfig` default: block `run_command`,
302/// allow all other tools.
303fn default_policies() -> Vec<PolicyRule> {
304    vec![
305        PolicyRule::Deny("run_command".to_string()),
306        PolicyRule::AllowAll,
307    ]
308}
309
310#[cfg(test)]
311mod tests {
312    use pyo3::types::PyAnyMethods;
313
314    use super::{
315        super::{
316            DEFAULT_IMAGE_GENERATION_MODEL,
317            capabilities::BuiltinTools,
318            models::{
319                GenerationConfig, ModelConfig, ModelEntry, ThinkingLevel, default_image_model_entry,
320            },
321        },
322        *,
323    };
324
325    #[derive(schemars::JsonSchema)]
326    struct CustomToolParams {}
327
328    #[test]
329    fn test_roundtrip_serialization() {
330        let config = AgentConfig {
331            system_instructions: Some(SystemInstructions::Custom("Be helpful".to_string())),
332            capabilities: Some(CapabilitiesConfig {
333                enable_subagents: true,
334                enabled_tools: Some(vec![BuiltinTools::ListDir]),
335                compaction_threshold: Some(4000),
336                ..CapabilitiesConfig::default()
337            }),
338            workspaces: vec![PathBuf::from("/tmp")],
339            ..AgentConfig::default()
340        };
341
342        let json = serde_json::to_string(&config).unwrap();
343        let parsed: AgentConfig = serde_json::from_str(&json).unwrap();
344        assert_eq!(parsed.workspaces.len(), 1);
345        assert_eq!(
346            parsed.capabilities.unwrap().enabled_tools.unwrap()[0],
347            BuiltinTools::ListDir
348        );
349    }
350
351    #[test]
352    fn agent_config_builder_with_gemini() {
353        let gemini = GeminiConfig {
354            api_key: Some("test-key".to_string()),
355            base_url: None,
356            models: ModelConfig::default(),
357        };
358        let config = AgentConfig::builder().gemini(gemini).build();
359        let gemini_cfg = config.gemini.expect("gemini should be Some");
360        assert_eq!(gemini_cfg.api_key.as_deref(), Some("test-key"));
361        assert_eq!(gemini_cfg.models.default.name, DEFAULT_MODEL);
362    }
363
364    #[test]
365    fn agent_config_builder_gemini_with_thinking_level() {
366        let gemini = GeminiConfig {
367            api_key: None,
368            base_url: None,
369            models: ModelConfig {
370                default: ModelEntry {
371                    name: "gemini-3.5-flash".to_string(),
372                    api_key: None,
373                    generation: GenerationConfig {
374                        thinking_level: Some(ThinkingLevel::High),
375                    },
376                },
377                image_generation: default_image_model_entry(),
378            },
379        };
380        let config = AgentConfig::builder().gemini(gemini).build();
381        let gemini_cfg = config.gemini.expect("gemini should be Some");
382        assert_eq!(
383            gemini_cfg.models.default.generation.thinking_level,
384            Some(ThinkingLevel::High)
385        );
386        assert_eq!(gemini_cfg.models.default.name, "gemini-3.5-flash");
387    }
388
389    #[test]
390    fn agent_config_gemini_none_by_default() {
391        let config = AgentConfig::default();
392        assert!(config.gemini.is_none());
393    }
394
395    #[test]
396    fn agent_config_gemini_serde_roundtrip() {
397        let config = AgentConfig {
398            gemini: Some(GeminiConfig {
399                api_key: Some("roundtrip-key".to_string()),
400                base_url: None,
401                models: ModelConfig {
402                    default: ModelEntry {
403                        name: "gemini-3.5-flash".to_string(),
404                        api_key: Some("model-key".to_string()),
405                        generation: GenerationConfig {
406                            thinking_level: Some(ThinkingLevel::Medium),
407                        },
408                    },
409                    image_generation: default_image_model_entry(),
410                },
411            }),
412            ..AgentConfig::default()
413        };
414        let json = serde_json::to_string(&config).unwrap();
415        let parsed: AgentConfig = serde_json::from_str(&json).unwrap();
416        let gemini_cfg = parsed.gemini.expect("gemini should survive roundtrip");
417        assert_eq!(gemini_cfg.api_key.as_deref(), Some("roundtrip-key"));
418        assert_eq!(gemini_cfg.models.default.name, "gemini-3.5-flash");
419        assert_eq!(
420            gemini_cfg.models.default.api_key.as_deref(),
421            Some("model-key")
422        );
423        assert_eq!(
424            gemini_cfg.models.default.generation.thinking_level,
425            Some(ThinkingLevel::Medium)
426        );
427    }
428
429    #[test]
430    fn system_instructions_custom_serde() {
431        let instr = SystemInstructions::Custom("Be a helpful assistant".to_string());
432        let json = serde_json::to_string(&instr).unwrap();
433        let parsed: SystemInstructions = serde_json::from_str(&json).unwrap();
434        match parsed {
435            SystemInstructions::Custom(text) => assert_eq!(text, "Be a helpful assistant"),
436            SystemInstructions::Templated { .. } => {
437                panic!("Expected Custom, got Templated")
438            }
439        }
440    }
441
442    #[test]
443    fn system_instructions_templated_serde() {
444        let instr = SystemInstructions::Templated {
445            identity: Some("a security analyst".to_string()),
446            sections: vec![SystemInstructionSection {
447                content: "Always check permissions".to_string(),
448                title: "security".to_string(),
449            }],
450        };
451        let json = serde_json::to_string(&instr).unwrap();
452        let parsed: SystemInstructions = serde_json::from_str(&json).unwrap();
453        match parsed {
454            SystemInstructions::Templated { identity, sections } => {
455                assert_eq!(identity.as_deref(), Some("a security analyst"));
456                assert_eq!(sections.len(), 1);
457                assert_eq!(sections[0].content, "Always check permissions");
458            }
459            SystemInstructions::Custom(_) => {
460                panic!("Expected Templated, got Custom")
461            }
462        }
463    }
464
465    #[test]
466    fn agent_config_fully_populated_serde() {
467        let config = AgentConfig {
468            system_instructions: Some(SystemInstructions::Templated {
469                identity: Some("test-identity".to_string()),
470                sections: vec![],
471            }),
472            capabilities: Some(CapabilitiesConfig {
473                enable_subagents: true,
474                disabled_tools: Some(vec![BuiltinTools::RunCommand]),
475                compaction_threshold: Some(1000),
476                ..CapabilitiesConfig::default()
477            }),
478            workspaces: vec![PathBuf::from("/a"), PathBuf::from("/b")],
479            tools: vec![crate::tools::ToolDefinition {
480                name: "custom_tool".to_owned(),
481                description: "A custom tool".to_owned(),
482                parameter_schema: serde_json::to_value(schemars::schema_for!(CustomToolParams))
483                    .unwrap(),
484            }],
485            policies: vec![PolicyRule::DenyAll],
486            triggers: vec![TriggerEntry {
487                name: "poll".to_owned(),
488                config: crate::triggers::TriggerConfig::every_secs(30),
489                message_template: "time to poll".to_owned(),
490            }],
491            hooks: vec![HookEntry {
492                name: "pre_gate".to_owned(),
493                point: crate::hooks::HookPoint::PreTurn,
494                callback_id: "cb_pre".to_owned(),
495            }],
496            skills: vec![PathBuf::from("/skills/foo")],
497            ..AgentConfig::default()
498        };
499        let json = serde_json::to_string(&config).unwrap();
500        let parsed: AgentConfig = serde_json::from_str(&json).unwrap();
501        assert_eq!(parsed.workspaces.len(), 2);
502        assert_eq!(parsed.tools.len(), 1);
503        assert_eq!(parsed.policies.len(), 1);
504        assert_eq!(parsed.triggers.len(), 1);
505        assert_eq!(parsed.hooks.len(), 1);
506        assert_eq!(parsed.skills.len(), 1);
507    }
508
509    #[test]
510    fn agent_config_empty_defaults_serde() {
511        let json = r#"{"system_instructions":null}"#;
512        let parsed: AgentConfig = serde_json::from_str(json).unwrap();
513        assert!(parsed.system_instructions.is_none());
514        assert!(parsed.capabilities.is_none());
515        assert!(parsed.workspaces.is_empty());
516        assert!(parsed.tools.is_empty());
517        assert_eq!(
518            parsed.policies,
519            vec![
520                PolicyRule::Deny("run_command".to_string()),
521                PolicyRule::AllowAll,
522            ]
523        );
524        assert!(parsed.triggers.is_empty());
525        assert!(parsed.hooks.is_empty());
526        assert!(parsed.skills.is_empty());
527
528        assert!(parsed.gemini.is_none());
529    }
530
531    #[test]
532    fn thinking_level_all_variants_python_str() {
533        assert_eq!(ThinkingLevel::Minimal.as_str(), "minimal");
534        assert_eq!(ThinkingLevel::Low.as_str(), "low");
535        assert_eq!(ThinkingLevel::Medium.as_str(), "medium");
536        assert_eq!(ThinkingLevel::High.as_str(), "high");
537    }
538
539    #[test]
540    fn thinking_level_all_variants_serde() {
541        for (variant, expected) in [
542            (ThinkingLevel::Minimal, "\"minimal\""),
543            (ThinkingLevel::Low, "\"low\""),
544            (ThinkingLevel::Medium, "\"medium\""),
545            (ThinkingLevel::High, "\"high\""),
546        ] {
547            let json = serde_json::to_string(&variant).unwrap();
548            assert_eq!(json, expected);
549            let parsed: ThinkingLevel = serde_json::from_str(&json).unwrap();
550            assert_eq!(parsed, variant);
551        }
552    }
553
554    #[test]
555    fn system_instructions_custom_variant() {
556        let instr = SystemInstructions::Custom("Be helpful".to_string());
557        let json = serde_json::to_string(&instr).unwrap();
558        let parsed: SystemInstructions = serde_json::from_str(&json).unwrap();
559        match parsed {
560            SystemInstructions::Custom(text) => assert_eq!(text, "Be helpful"),
561            SystemInstructions::Templated { .. } => {
562                panic!("Expected Custom, got Templated")
563            }
564        }
565    }
566
567    #[test]
568    fn agent_config_all_optional_fields_roundtrip() {
569        let config = AgentConfig {
570            workspaces: vec![PathBuf::from("/ws")],
571            skills: vec![PathBuf::from("/skills/test")],
572
573            conversation_id: Some("conv-123".to_string()),
574            save_dir: Some(PathBuf::from("/save")),
575            app_data_dir: Some(PathBuf::from("/app")),
576            response_schema: Some(JsonSchema::new(serde_json::json!({"type": "object"}))),
577            ..AgentConfig::default()
578        };
579        let json = serde_json::to_string(&config).unwrap();
580        let parsed: AgentConfig = serde_json::from_str(&json).unwrap();
581        assert_eq!(parsed.workspaces.len(), 1);
582        assert_eq!(parsed.conversation_id.as_deref(), Some("conv-123"));
583        assert_eq!(parsed.save_dir.as_ref().unwrap(), &PathBuf::from("/save"));
584        assert!(parsed.response_schema.is_some());
585    }
586
587    #[test]
588
589    fn agent_config_custom_tools_and_builtin_tools_coexist() {
590        // Audit 3: Verify that an AgentConfig can carry both custom tool
591        // definitions (for the ToolRegistry) AND SDK built-in tools
592        // (via CapabilitiesConfig.enabled_tools) at the same time.
593        let custom_tool = crate::tools::ToolDefinition {
594            name: "my_custom_tool".to_owned(),
595            description: "Does something custom".to_owned(),
596            parameter_schema: serde_json::json!({"type": "object", "properties": {}}),
597        };
598        let config = AgentConfig {
599            tools: vec![custom_tool],
600            capabilities: Some(CapabilitiesConfig {
601                enabled_tools: Some(vec![BuiltinTools::ViewFile, BuiltinTools::RunCommand]),
602                ..CapabilitiesConfig::default()
603            }),
604            ..AgentConfig::default()
605        };
606
607        // Serialize and deserialize to prove the combined config survives a roundtrip.
608        let json = serde_json::to_string(&config).unwrap();
609        let parsed: AgentConfig = serde_json::from_str(&json).unwrap();
610
611        // Custom tools are preserved.
612        assert_eq!(parsed.tools.len(), 1);
613        assert_eq!(parsed.tools[0].name, "my_custom_tool");
614
615        // Built-in tool selection is preserved.
616        let caps = parsed.capabilities.as_ref().unwrap();
617        let enabled = caps.enabled_tools.as_ref().unwrap();
618        assert_eq!(enabled.len(), 2);
619        assert!(enabled.contains(&BuiltinTools::ViewFile));
620        assert!(enabled.contains(&BuiltinTools::RunCommand));
621
622        // Validate the config is internally consistent.
623        assert!(caps.validate().is_ok());
624    }
625
626    #[test]
627    fn agent_config_custom_tools_only_no_builtins() {
628        // Verify custom_tools_only() + custom tools is valid.
629        let config = AgentConfig {
630            tools: vec![crate::tools::ToolDefinition {
631                name: "fetch_data".to_owned(),
632                description: "Fetches data".to_owned(),
633                parameter_schema: serde_json::json!({"type": "object"}),
634            }],
635            capabilities: Some(CapabilitiesConfig::custom_tools_only()),
636            ..AgentConfig::default()
637        };
638
639        let caps = config.capabilities.as_ref().unwrap();
640        assert!(caps.enabled_tools.as_ref().unwrap().is_empty());
641        assert!(caps.validate().is_ok());
642        assert_eq!(config.tools.len(), 1);
643    }
644
645    // ── LocalAgentConfig tests ───────────────────────────────────────
646
647    #[test]
648    fn local_agent_config_default() {
649        let config = LocalAgentConfig::default();
650        assert_eq!(config.agent.model, DEFAULT_MODEL);
651    }
652
653    #[test]
654    fn local_agent_config_from_agent_config() {
655        let agent_cfg = AgentConfig {
656            model: "gemini-3.5-flash".to_string(),
657            ..AgentConfig::default()
658        };
659        let local: LocalAgentConfig = agent_cfg.into();
660        assert_eq!(local.agent.model, "gemini-3.5-flash");
661    }
662
663    #[test]
664    fn local_agent_config_serde_roundtrip() {
665        let config = LocalAgentConfig::new(AgentConfig::default());
666        let json = serde_json::to_string(&config).unwrap();
667        let parsed: LocalAgentConfig = serde_json::from_str(&json).unwrap();
668        assert_eq!(parsed.agent.model, DEFAULT_MODEL);
669    }
670
671    // ── SDK field name alignment tests ────────────────────────────────
672    //
673    // Verify that serde serializes field names to match the Python SDK's
674    // expected JSON keys.
675
676    #[test]
677    fn skills_serializes_as_skills_paths() {
678        let config = AgentConfig::builder()
679            .skills(vec![PathBuf::from("/skill/a.md")])
680            .build();
681        let json = serde_json::to_string(&config).unwrap();
682        let v: serde_json::Value = serde_json::from_str(&json).unwrap();
683        assert!(
684            v.get("skills_paths").is_some(),
685            "Expected JSON key 'skills_paths', got: {json}"
686        );
687        assert!(
688            v.get("skills").is_none(),
689            "Should not have 'skills' key in JSON"
690        );
691    }
692
693    #[test]
694    fn skills_paths_deserializes_to_skills_field() {
695        let json = r#"{"skills_paths": ["/skill/a.md"]}"#;
696        let config: AgentConfig = serde_json::from_str(json).unwrap();
697        assert_eq!(config.skills.len(), 1);
698        assert_eq!(config.skills[0], PathBuf::from("/skill/a.md"));
699    }
700
701    #[test]
702    fn gemini_serializes_as_gemini_config() {
703        let config = AgentConfig::builder()
704            .gemini(super::super::GeminiConfig::default())
705            .build();
706        let json = serde_json::to_string(&config).unwrap();
707        let v: serde_json::Value = serde_json::from_str(&json).unwrap();
708        assert!(
709            v.get("gemini_config").is_some(),
710            "Expected JSON key 'gemini_config', got: {json}"
711        );
712        assert!(
713            v.get("gemini").is_none(),
714            "Should not have 'gemini' key in JSON"
715        );
716    }
717
718    #[test]
719    fn gemini_config_deserializes_to_gemini_field() {
720        let json = r#"{"gemini_config": {"api_key": "test-key"}}"#;
721        let config: AgentConfig = serde_json::from_str(json).unwrap();
722        assert_eq!(
723            config.gemini.as_ref().unwrap().api_key.as_deref(),
724            Some("test-key")
725        );
726    }
727
728    // ── skip_serializing_if tests ─────────────────────────────────────
729
730    #[test]
731    fn empty_vecs_omitted_from_json() {
732        let config = AgentConfig::default();
733        let json = serde_json::to_string(&config).unwrap();
734        let v: serde_json::Value = serde_json::from_str(&json).unwrap();
735        // These should all be absent when empty.
736        for key in &[
737            "workspaces",
738            "tools",
739            "triggers",
740            "hooks",
741            "skills_paths",
742            "mcp_servers",
743        ] {
744            assert!(
745                v.get(key).is_none(),
746                "Empty vec field '{key}' should be omitted from JSON, got: {json}"
747            );
748        }
749        // policies should always be present (non-empty default)
750        assert!(
751            v.get("policies").is_some(),
752            "policies should always be serialized"
753        );
754    }
755
756    #[test]
757    fn populated_vecs_included_in_json() {
758        let config = AgentConfig::builder()
759            .skills(vec![PathBuf::from("/skill.md")])
760            .workspaces(vec![PathBuf::from("/ws")])
761            .build();
762        let json = serde_json::to_string(&config).unwrap();
763        let v: serde_json::Value = serde_json::from_str(&json).unwrap();
764        assert!(
765            v.get("skills_paths").is_some(),
766            "Non-empty skills should be present"
767        );
768        assert!(
769            v.get("workspaces").is_some(),
770            "Non-empty workspaces should be present"
771        );
772    }
773
774    // ── Default policy tests ──────────────────────────────────────────
775
776    #[test]
777    fn default_policies_deny_run_command_allow_rest() {
778        let config = AgentConfig::default();
779        assert_eq!(config.policies.len(), 2);
780        assert_eq!(
781            config.policies[0],
782            PolicyRule::Deny("run_command".to_string())
783        );
784        assert_eq!(config.policies[1], PolicyRule::AllowAll);
785    }
786
787    // ── Python SDK mirror tests ────────────────────────────────────────
788    //
789    // These tests import the live Python SDK and verify that our Rust
790    // constants haven't drifted from the canonical Python values.  They
791    // require `pyo3::prepare_freethreaded_python()` and a venv with the
792    // SDK installed.
793
794    /// Helper: extract a Python module-level attribute as a `String`.
795    fn py_str_attr(module: &str, attr: &str) -> String {
796        pyo3::prepare_freethreaded_python();
797        pyo3::Python::with_gil(|py| {
798            crate::runtime::venv::configure_python_sys_path(py)
799                .unwrap_or_else(|e| panic!("Failed to configure python sys.path: {e}"));
800            let m = py
801                .import_bound(module)
802                .unwrap_or_else(|e| panic!("Failed to import {module}: {e}"));
803            m.getattr(attr)
804                .unwrap_or_else(|e| panic!("Failed to get {module}.{attr}: {e}"))
805                .extract::<String>()
806                .unwrap_or_else(|e| panic!("Failed to extract {module}.{attr} as String: {e}"))
807        })
808    }
809
810    #[test]
811    fn default_model_matches_python_sdk() {
812        let py_val = py_str_attr("google.antigravity.types", "DEFAULT_MODEL");
813        assert_eq!(
814            DEFAULT_MODEL, py_val,
815            "Rust DEFAULT_MODEL ({DEFAULT_MODEL}) != Python SDK ({py_val})"
816        );
817    }
818
819    #[test]
820    fn default_image_model_matches_python_sdk() {
821        let py_val = py_str_attr("google.antigravity.types", "DEFAULT_IMAGE_GENERATION_MODEL");
822        assert_eq!(
823            DEFAULT_IMAGE_GENERATION_MODEL, py_val,
824            "Rust DEFAULT_IMAGE_GENERATION_MODEL ({DEFAULT_IMAGE_GENERATION_MODEL}) != Python SDK ({py_val})"
825        );
826    }
827
828    // ── effective_api_key tests ──────────────────────────────────────
829
830    #[test]
831    fn effective_api_key_prefers_per_model_key() {
832        let config = AgentConfig::builder()
833            .api_key("top-level-key")
834            .gemini(super::super::GeminiConfig {
835                api_key: Some("shared-key".into()),
836                base_url: None,
837                models: super::super::ModelConfig {
838                    default: super::super::ModelEntry {
839                        name: "gemini-3.5-flash".into(),
840                        api_key: Some("per-model-key".into()),
841                        generation: super::super::GenerationConfig::default(),
842                    },
843                    image_generation: super::super::ModelEntry {
844                        name: "imagen-4.0-generate-preview-06-03".into(),
845                        api_key: None,
846                        generation: super::super::GenerationConfig::default(),
847                    },
848                },
849            })
850            .build();
851        assert_eq!(config.effective_api_key().as_deref(), Some("per-model-key"));
852    }
853
854    #[test]
855    fn effective_api_key_falls_back_to_gemini_shared_key() {
856        let config = AgentConfig::builder()
857            .gemini(super::super::GeminiConfig {
858                api_key: Some("shared-key".into()),
859                ..Default::default()
860            })
861            .build();
862        assert_eq!(config.effective_api_key().as_deref(), Some("shared-key"));
863    }
864
865    #[test]
866    fn effective_api_key_falls_back_to_top_level() {
867        let config = AgentConfig::builder().api_key("top-level-key").build();
868        assert_eq!(config.effective_api_key().as_deref(), Some("top-level-key"));
869    }
870
871    #[test]
872    fn effective_api_key_none_without_any_key() {
873        // Temporarily remove the env var if set so the test is deterministic.
874        let saved = std::env::var("GEMINI_API_KEY").ok();
875        unsafe { std::env::remove_var("GEMINI_API_KEY") };
876        let config = AgentConfig::builder().build();
877        let result = config.effective_api_key();
878        // Restore
879        if let Some(v) = saved {
880            unsafe { std::env::set_var("GEMINI_API_KEY", v) };
881        }
882        assert!(result.is_none());
883    }
884}
agy_bridge/config/agent.rs

agy_bridge/config/
agent.rs
