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    /// Optional initial conversation history to inject after agent creation.
237    ///
238    /// When set, these messages are inserted into the SDK's internal
239    /// `_history` list before the first chat turn, enabling warm-start
240    /// scenarios such as safety recovery (re-creating an agent with
241    /// curated prior context).
242    ///
243    /// Serialized as `"initial_history"` and consumed by the Python init
244    /// script. Skipped from serialization when empty.
245    #[serde(default, skip_serializing_if = "Vec::is_empty")]
246    #[builder(default, setter(transform = |v: impl IntoIterator<Item = impl Into<crate::types::ConversationMessage>>| v.into_iter().map(Into::into).collect()))]
247    pub initial_history: Vec<crate::types::ConversationMessage>,
248}
249
250impl Default for AgentConfig {
251    fn default() -> Self {
252        Self::builder().build()
253    }
254}
255
256impl AgentConfig {
257    /// Resolve the effective API key using the same priority chain as the
258    /// Python SDK's `LocalAgentConfig` → `_build_harness_config`:
259    ///
260    /// 1. Per-model key (`gemini.models.default.api_key`)
261    /// 2. Shared `GeminiConfig` key (`gemini.api_key`)
262    /// 3. Top-level shorthand (`api_key`)
263    /// 4. `$GEMINI_API_KEY` environment variable
264    #[must_use]
265    pub fn effective_api_key(&self) -> Option<String> {
266        self.gemini
267            .as_ref()
268            .and_then(|g| g.models.default.api_key.clone())
269            .or_else(|| self.gemini.as_ref().and_then(|g| g.api_key.clone()))
270            .or_else(|| self.api_key.clone())
271            // NOLINT: .ok() is intentional — env var not set returns None, which is the expected fallback
272            .or_else(|| std::env::var("GEMINI_API_KEY").ok())
273    }
274
275    /// Returns the names of all explicitly registered custom tools.
276    /// To get the full list of tools including built-ins, combine this
277    /// with `capabilities.enabled_tools` or examine `tools` + default semantics.
278    #[must_use]
279    pub fn custom_tool_names(&self) -> Vec<String> {
280        self.tools.iter().map(|t| t.name.clone()).collect()
281    }
282}
283
284// ─── LocalAgentConfig ────────────────────────────────────────────────────────
285
286/// Configuration for a local (on-device) agent, mirroring the Python SDK's
287/// `LocalAgentConfig`.
288///
289/// Wraps the standard [`AgentConfig`] via `#[serde(flatten)]`. The Python SDK's
290/// `LocalAgentConfig` extends the base `AgentConfig` with override defaults
291/// (e.g. `policies = confirm_run_command()`, `workspaces = [cwd]`), which
292/// our `AgentConfig` already matches.
293#[derive(Debug, Clone, Serialize, Deserialize, Default)]
294pub struct LocalAgentConfig {
295    /// The base agent configuration.
296    #[serde(flatten)]
297    pub agent: AgentConfig,
298}
299
300impl LocalAgentConfig {
301    /// Create a new `LocalAgentConfig` wrapping the given agent configuration.
302    #[must_use]
303    pub const fn new(agent: AgentConfig) -> Self {
304        Self { agent }
305    }
306}
307
308impl From<AgentConfig> for LocalAgentConfig {
309    fn from(agent: AgentConfig) -> Self {
310        Self::new(agent)
311    }
312}
313
314/// Matches the Python SDK's `LocalAgentConfig` default: block `run_command`,
315/// allow all other tools.
316fn default_policies() -> Vec<PolicyRule> {
317    vec![
318        PolicyRule::Deny("run_command".to_string()),
319        PolicyRule::AllowAll,
320    ]
321}
322
323#[cfg(test)]
324mod tests {
325    use pyo3::types::PyAnyMethods;
326
327    use super::{
328        super::{
329            DEFAULT_IMAGE_GENERATION_MODEL,
330            capabilities::BuiltinTools,
331            models::{
332                GenerationConfig, ModelConfig, ModelEntry, ThinkingLevel, default_image_model_entry,
333            },
334        },
335        *,
336    };
337
338    #[derive(schemars::JsonSchema)]
339    struct CustomToolParams {}
340
341    #[test]
342    fn test_roundtrip_serialization() {
343        let config = AgentConfig {
344            system_instructions: Some(SystemInstructions::Custom("Be helpful".to_string())),
345            capabilities: Some(CapabilitiesConfig {
346                enable_subagents: true,
347                enabled_tools: Some(vec![BuiltinTools::ListDir]),
348                compaction_threshold: Some(4000),
349                ..CapabilitiesConfig::default()
350            }),
351            workspaces: vec![PathBuf::from("/tmp")],
352            ..AgentConfig::default()
353        };
354
355        let json = serde_json::to_string(&config).unwrap();
356        let parsed: AgentConfig = serde_json::from_str(&json).unwrap();
357        assert_eq!(parsed.workspaces.len(), 1);
358        assert_eq!(
359            parsed.capabilities.unwrap().enabled_tools.unwrap()[0],
360            BuiltinTools::ListDir
361        );
362    }
363
364    #[test]
365    fn agent_config_builder_with_gemini() {
366        let gemini = GeminiConfig {
367            api_key: Some("test-key".to_string()),
368            base_url: None,
369            models: ModelConfig::default(),
370        };
371        let config = AgentConfig::builder().gemini(gemini).build();
372        let gemini_cfg = config.gemini.expect("gemini should be Some");
373        assert_eq!(gemini_cfg.api_key.as_deref(), Some("test-key"));
374        assert_eq!(gemini_cfg.models.default.name, DEFAULT_MODEL);
375    }
376
377    #[test]
378    fn agent_config_builder_gemini_with_thinking_level() {
379        let gemini = GeminiConfig {
380            api_key: None,
381            base_url: None,
382            models: ModelConfig {
383                default: ModelEntry {
384                    name: "gemini-3.5-flash".to_string(),
385                    api_key: None,
386                    generation: GenerationConfig {
387                        thinking_level: Some(ThinkingLevel::High),
388                    },
389                },
390                image_generation: default_image_model_entry(),
391            },
392        };
393        let config = AgentConfig::builder().gemini(gemini).build();
394        let gemini_cfg = config.gemini.expect("gemini should be Some");
395        assert_eq!(
396            gemini_cfg.models.default.generation.thinking_level,
397            Some(ThinkingLevel::High)
398        );
399        assert_eq!(gemini_cfg.models.default.name, "gemini-3.5-flash");
400    }
401
402    #[test]
403    fn agent_config_gemini_none_by_default() {
404        let config = AgentConfig::default();
405        assert!(config.gemini.is_none());
406    }
407
408    #[test]
409    fn agent_config_gemini_serde_roundtrip() {
410        let config = AgentConfig {
411            gemini: Some(GeminiConfig {
412                api_key: Some("roundtrip-key".to_string()),
413                base_url: None,
414                models: ModelConfig {
415                    default: ModelEntry {
416                        name: "gemini-3.5-flash".to_string(),
417                        api_key: Some("model-key".to_string()),
418                        generation: GenerationConfig {
419                            thinking_level: Some(ThinkingLevel::Medium),
420                        },
421                    },
422                    image_generation: default_image_model_entry(),
423                },
424            }),
425            ..AgentConfig::default()
426        };
427        let json = serde_json::to_string(&config).unwrap();
428        let parsed: AgentConfig = serde_json::from_str(&json).unwrap();
429        let gemini_cfg = parsed.gemini.expect("gemini should survive roundtrip");
430        assert_eq!(gemini_cfg.api_key.as_deref(), Some("roundtrip-key"));
431        assert_eq!(gemini_cfg.models.default.name, "gemini-3.5-flash");
432        assert_eq!(
433            gemini_cfg.models.default.api_key.as_deref(),
434            Some("model-key")
435        );
436        assert_eq!(
437            gemini_cfg.models.default.generation.thinking_level,
438            Some(ThinkingLevel::Medium)
439        );
440    }
441
442    #[test]
443    fn system_instructions_custom_serde() {
444        let instr = SystemInstructions::Custom("Be a helpful assistant".to_string());
445        let json = serde_json::to_string(&instr).unwrap();
446        let parsed: SystemInstructions = serde_json::from_str(&json).unwrap();
447        match parsed {
448            SystemInstructions::Custom(text) => assert_eq!(text, "Be a helpful assistant"),
449            SystemInstructions::Templated { .. } => {
450                panic!("Expected Custom, got Templated")
451            }
452        }
453    }
454
455    #[test]
456    fn system_instructions_templated_serde() {
457        let instr = SystemInstructions::Templated {
458            identity: Some("a security analyst".to_string()),
459            sections: vec![SystemInstructionSection {
460                content: "Always check permissions".to_string(),
461                title: "security".to_string(),
462            }],
463        };
464        let json = serde_json::to_string(&instr).unwrap();
465        let parsed: SystemInstructions = serde_json::from_str(&json).unwrap();
466        match parsed {
467            SystemInstructions::Templated { identity, sections } => {
468                assert_eq!(identity.as_deref(), Some("a security analyst"));
469                assert_eq!(sections.len(), 1);
470                assert_eq!(sections[0].content, "Always check permissions");
471            }
472            SystemInstructions::Custom(_) => {
473                panic!("Expected Templated, got Custom")
474            }
475        }
476    }
477
478    #[test]
479    fn agent_config_fully_populated_serde() {
480        let config = AgentConfig {
481            system_instructions: Some(SystemInstructions::Templated {
482                identity: Some("test-identity".to_string()),
483                sections: vec![],
484            }),
485            capabilities: Some(CapabilitiesConfig {
486                enable_subagents: true,
487                disabled_tools: Some(vec![BuiltinTools::RunCommand]),
488                compaction_threshold: Some(1000),
489                ..CapabilitiesConfig::default()
490            }),
491            workspaces: vec![PathBuf::from("/a"), PathBuf::from("/b")],
492            tools: vec![crate::tools::ToolDefinition {
493                name: "custom_tool".to_owned(),
494                description: "A custom tool".to_owned(),
495                parameter_schema: serde_json::to_value(schemars::schema_for!(CustomToolParams))
496                    .unwrap(),
497            }],
498            policies: vec![PolicyRule::DenyAll],
499            triggers: vec![TriggerEntry {
500                name: "poll".to_owned(),
501                config: crate::triggers::TriggerConfig::every_secs(30),
502                message_template: "time to poll".to_owned(),
503            }],
504            hooks: vec![HookEntry {
505                name: "pre_gate".to_owned(),
506                point: crate::hooks::HookPoint::PreTurn,
507                callback_id: "cb_pre".to_owned(),
508            }],
509            skills: vec![PathBuf::from("/skills/foo")],
510            ..AgentConfig::default()
511        };
512        let json = serde_json::to_string(&config).unwrap();
513        let parsed: AgentConfig = serde_json::from_str(&json).unwrap();
514        assert_eq!(parsed.workspaces.len(), 2);
515        assert_eq!(parsed.tools.len(), 1);
516        assert_eq!(parsed.policies.len(), 1);
517        assert_eq!(parsed.triggers.len(), 1);
518        assert_eq!(parsed.hooks.len(), 1);
519        assert_eq!(parsed.skills.len(), 1);
520    }
521
522    #[test]
523    fn agent_config_empty_defaults_serde() {
524        let json = r#"{"system_instructions":null}"#;
525        let parsed: AgentConfig = serde_json::from_str(json).unwrap();
526        assert!(parsed.system_instructions.is_none());
527        assert!(parsed.capabilities.is_none());
528        assert!(parsed.workspaces.is_empty());
529        assert!(parsed.tools.is_empty());
530        assert_eq!(
531            parsed.policies,
532            vec![
533                PolicyRule::Deny("run_command".to_string()),
534                PolicyRule::AllowAll,
535            ]
536        );
537        assert!(parsed.triggers.is_empty());
538        assert!(parsed.hooks.is_empty());
539        assert!(parsed.skills.is_empty());
540
541        assert!(parsed.gemini.is_none());
542    }
543
544    #[test]
545    fn agent_config_all_optional_fields_roundtrip() {
546        let config = AgentConfig {
547            workspaces: vec![PathBuf::from("/ws")],
548            skills: vec![PathBuf::from("/skills/test")],
549
550            conversation_id: Some("conv-123".to_string()),
551            save_dir: Some(PathBuf::from("/save")),
552            app_data_dir: Some(PathBuf::from("/app")),
553            response_schema: Some(JsonSchema::new(serde_json::json!({"type": "object"}))),
554            ..AgentConfig::default()
555        };
556        let json = serde_json::to_string(&config).unwrap();
557        let parsed: AgentConfig = serde_json::from_str(&json).unwrap();
558        assert_eq!(parsed.workspaces.len(), 1);
559        assert_eq!(parsed.conversation_id.as_deref(), Some("conv-123"));
560        assert_eq!(parsed.save_dir.as_ref().unwrap(), &PathBuf::from("/save"));
561        assert!(parsed.response_schema.is_some());
562    }
563
564    #[test]
565    fn agent_config_custom_tools_and_builtin_tools_coexist() {
566        // Audit 3: Verify that an AgentConfig can carry both custom tool
567        // definitions (for the ToolRegistry) AND SDK built-in tools
568        // (via CapabilitiesConfig.enabled_tools) at the same time.
569        let custom_tool = crate::tools::ToolDefinition {
570            name: "my_custom_tool".to_owned(),
571            description: "Does something custom".to_owned(),
572            parameter_schema: serde_json::json!({"type": "object", "properties": {}}),
573        };
574        let config = AgentConfig {
575            tools: vec![custom_tool],
576            capabilities: Some(CapabilitiesConfig {
577                enabled_tools: Some(vec![BuiltinTools::ViewFile, BuiltinTools::RunCommand]),
578                ..CapabilitiesConfig::default()
579            }),
580            ..AgentConfig::default()
581        };
582
583        // Serialize and deserialize to prove the combined config survives a roundtrip.
584        let json = serde_json::to_string(&config).unwrap();
585        let parsed: AgentConfig = serde_json::from_str(&json).unwrap();
586
587        // Custom tools are preserved.
588        assert_eq!(parsed.tools.len(), 1);
589        assert_eq!(parsed.tools[0].name, "my_custom_tool");
590
591        // Built-in tool selection is preserved.
592        let caps = parsed.capabilities.as_ref().unwrap();
593        let enabled = caps.enabled_tools.as_ref().unwrap();
594        assert_eq!(enabled.len(), 2);
595        assert!(enabled.contains(&BuiltinTools::ViewFile));
596        assert!(enabled.contains(&BuiltinTools::RunCommand));
597
598        // Validate the config is internally consistent.
599        assert!(caps.validate().is_ok());
600    }
601
602    #[test]
603    fn agent_config_custom_tools_only_no_builtins() {
604        // Verify custom_tools_only() + custom tools is valid.
605        let config = AgentConfig {
606            tools: vec![crate::tools::ToolDefinition {
607                name: "fetch_data".to_owned(),
608                description: "Fetches data".to_owned(),
609                parameter_schema: serde_json::json!({"type": "object"}),
610            }],
611            capabilities: Some(CapabilitiesConfig::custom_tools_only()),
612            ..AgentConfig::default()
613        };
614
615        let caps = config.capabilities.as_ref().unwrap();
616        assert!(caps.enabled_tools.as_ref().unwrap().is_empty());
617        assert!(caps.validate().is_ok());
618        assert_eq!(config.tools.len(), 1);
619    }
620
621    // ── LocalAgentConfig tests ───────────────────────────────────────
622
623    #[test]
624    fn local_agent_config_default() {
625        let config = LocalAgentConfig::default();
626        assert_eq!(config.agent.model, DEFAULT_MODEL);
627    }
628
629    #[test]
630    fn local_agent_config_from_agent_config() {
631        let agent_cfg = AgentConfig {
632            model: "gemini-3.5-flash".to_string(),
633            ..AgentConfig::default()
634        };
635        let local: LocalAgentConfig = agent_cfg.into();
636        assert_eq!(local.agent.model, "gemini-3.5-flash");
637    }
638
639    #[test]
640    fn local_agent_config_serde_roundtrip() {
641        let config = LocalAgentConfig::new(AgentConfig::default());
642        let json = serde_json::to_string(&config).unwrap();
643        let parsed: LocalAgentConfig = serde_json::from_str(&json).unwrap();
644        assert_eq!(parsed.agent.model, DEFAULT_MODEL);
645    }
646
647    // ── SDK field name alignment tests ────────────────────────────────
648    //
649    // Verify that serde serializes field names to match the Python SDK's
650    // expected JSON keys.
651
652    #[test]
653    fn skills_serializes_as_skills_paths() {
654        let config = AgentConfig::builder()
655            .skills(vec![PathBuf::from("/skill/a.md")])
656            .build();
657        let json = serde_json::to_string(&config).unwrap();
658        let v: serde_json::Value = serde_json::from_str(&json).unwrap();
659        assert!(
660            v.get("skills_paths").is_some(),
661            "Expected JSON key 'skills_paths', got: {json}"
662        );
663        assert!(
664            v.get("skills").is_none(),
665            "Should not have 'skills' key in JSON"
666        );
667    }
668
669    #[test]
670    fn skills_paths_deserializes_to_skills_field() {
671        let json = r#"{"skills_paths": ["/skill/a.md"]}"#;
672        let config: AgentConfig = serde_json::from_str(json).unwrap();
673        assert_eq!(config.skills.len(), 1);
674        assert_eq!(config.skills[0], PathBuf::from("/skill/a.md"));
675    }
676
677    #[test]
678    fn gemini_serializes_as_gemini_config() {
679        let config = AgentConfig::builder()
680            .gemini(super::super::GeminiConfig::default())
681            .build();
682        let json = serde_json::to_string(&config).unwrap();
683        let v: serde_json::Value = serde_json::from_str(&json).unwrap();
684        assert!(
685            v.get("gemini_config").is_some(),
686            "Expected JSON key 'gemini_config', got: {json}"
687        );
688        assert!(
689            v.get("gemini").is_none(),
690            "Should not have 'gemini' key in JSON"
691        );
692    }
693
694    #[test]
695    fn gemini_config_deserializes_to_gemini_field() {
696        let json = r#"{"gemini_config": {"api_key": "test-key"}}"#;
697        let config: AgentConfig = serde_json::from_str(json).unwrap();
698        assert_eq!(
699            config.gemini.as_ref().unwrap().api_key.as_deref(),
700            Some("test-key")
701        );
702    }
703
704    // ── skip_serializing_if tests ─────────────────────────────────────
705
706    #[test]
707    fn empty_vecs_omitted_from_json() {
708        let config = AgentConfig::default();
709        let json = serde_json::to_string(&config).unwrap();
710        let v: serde_json::Value = serde_json::from_str(&json).unwrap();
711        // These should all be absent when empty.
712        for key in &[
713            "workspaces",
714            "tools",
715            "triggers",
716            "hooks",
717            "skills_paths",
718            "mcp_servers",
719        ] {
720            assert!(
721                v.get(key).is_none(),
722                "Empty vec field '{key}' should be omitted from JSON, got: {json}"
723            );
724        }
725        // policies should always be present (non-empty default)
726        assert!(
727            v.get("policies").is_some(),
728            "policies should always be serialized"
729        );
730    }
731
732    #[test]
733    fn populated_vecs_included_in_json() {
734        let config = AgentConfig::builder()
735            .skills(vec![PathBuf::from("/skill.md")])
736            .workspaces(vec![PathBuf::from("/ws")])
737            .build();
738        let json = serde_json::to_string(&config).unwrap();
739        let v: serde_json::Value = serde_json::from_str(&json).unwrap();
740        assert!(
741            v.get("skills_paths").is_some(),
742            "Non-empty skills should be present"
743        );
744        assert!(
745            v.get("workspaces").is_some(),
746            "Non-empty workspaces should be present"
747        );
748    }
749
750    // ── Default policy tests ──────────────────────────────────────────
751
752    #[test]
753    fn default_policies_deny_run_command_allow_rest() {
754        let config = AgentConfig::default();
755        assert_eq!(config.policies.len(), 2);
756        assert_eq!(
757            config.policies[0],
758            PolicyRule::Deny("run_command".to_string())
759        );
760        assert_eq!(config.policies[1], PolicyRule::AllowAll);
761    }
762
763    // ── Python SDK mirror tests ────────────────────────────────────────
764    //
765    // These tests import the live Python SDK and verify that our Rust
766    // constants haven't drifted from the canonical Python values.  They
767    // require `pyo3::Python::initialize()` and a venv with the
768    // SDK installed.
769
770    /// Helper: extract a Python module-level attribute as a `String`.
771    fn py_str_attr(module: &str, attr: &str) -> String {
772        pyo3::Python::initialize();
773        pyo3::Python::attach(|py| {
774            crate::runtime::venv::configure_python_sys_path(py)
775                .unwrap_or_else(|e| panic!("Failed to configure python sys.path: {e}"));
776            let m = py
777                .import(module)
778                .unwrap_or_else(|e| panic!("Failed to import {module}: {e}"));
779            m.getattr(attr)
780                .unwrap_or_else(|e| panic!("Failed to get {module}.{attr}: {e}"))
781                .extract::<String>()
782                .unwrap_or_else(|e| panic!("Failed to extract {module}.{attr} as String: {e}"))
783        })
784    }
785
786    #[test]
787    fn default_model_matches_python_sdk() {
788        let py_val = py_str_attr("google.antigravity.types", "DEFAULT_MODEL");
789        assert_eq!(
790            DEFAULT_MODEL, py_val,
791            "Rust DEFAULT_MODEL ({DEFAULT_MODEL}) != Python SDK ({py_val})"
792        );
793    }
794
795    #[test]
796    fn default_image_model_matches_python_sdk() {
797        let py_val = py_str_attr("google.antigravity.types", "DEFAULT_IMAGE_GENERATION_MODEL");
798        assert_eq!(
799            DEFAULT_IMAGE_GENERATION_MODEL, py_val,
800            "Rust DEFAULT_IMAGE_GENERATION_MODEL ({DEFAULT_IMAGE_GENERATION_MODEL}) != Python SDK ({py_val})"
801        );
802    }
803
804    // ── effective_api_key tests ──────────────────────────────────────
805
806    #[test]
807    fn effective_api_key_prefers_per_model_key() {
808        let config = AgentConfig::builder()
809            .api_key("top-level-key")
810            .gemini(super::super::GeminiConfig {
811                api_key: Some("shared-key".into()),
812                base_url: None,
813                models: super::super::ModelConfig {
814                    default: super::super::ModelEntry {
815                        name: "gemini-3.5-flash".into(),
816                        api_key: Some("per-model-key".into()),
817                        generation: super::super::GenerationConfig::default(),
818                    },
819                    image_generation: super::super::ModelEntry {
820                        name: "imagen-4.0-generate-preview-06-03".into(),
821                        api_key: None,
822                        generation: super::super::GenerationConfig::default(),
823                    },
824                },
825            })
826            .build();
827        assert_eq!(config.effective_api_key().as_deref(), Some("per-model-key"));
828    }
829
830    #[test]
831    fn effective_api_key_falls_back_to_gemini_shared_key() {
832        let config = AgentConfig::builder()
833            .gemini(super::super::GeminiConfig {
834                api_key: Some("shared-key".into()),
835                ..Default::default()
836            })
837            .build();
838        assert_eq!(config.effective_api_key().as_deref(), Some("shared-key"));
839    }
840
841    #[test]
842    fn effective_api_key_falls_back_to_top_level() {
843        let config = AgentConfig::builder().api_key("top-level-key").build();
844        assert_eq!(config.effective_api_key().as_deref(), Some("top-level-key"));
845    }
846
847    #[test]
848    fn effective_api_key_none_without_any_key() {
849        // Build a config with no API key set at any level.
850        // We can't safely manipulate env vars in multi-threaded tests,
851        // so we test the chain up to the env-var fallback: if all config
852        // keys are None and the env var isn't set, the result is None.
853        // If GEMINI_API_KEY happens to be set, we verify it's returned.
854        let config = AgentConfig::builder().build();
855        let result = config.effective_api_key();
856        // NOLINT: .ok() is intentional — env var not set returns None, which is the expected fallback
857        match std::env::var("GEMINI_API_KEY").ok() {
858            Some(env_key) => assert_eq!(result.as_deref(), Some(env_key.as_str())),
859            None => assert!(result.is_none()),
860        }
861    }
862
863    #[test]
864    fn initial_history_empty_by_default() {
865        let config = AgentConfig::default();
866        assert!(config.initial_history.is_empty());
867
868        // Empty initial_history should be skipped in serialization.
869        let json = serde_json::to_string(&config).unwrap();
870        assert!(
871            !json.contains("initial_history"),
872            "empty initial_history should be skipped in JSON"
873        );
874    }
875
876    #[test]
877    fn initial_history_roundtrip() {
878        use crate::types::{ConversationMessage, MessageRole};
879
880        let config = AgentConfig::builder()
881            .initial_history(vec![
882                ConversationMessage {
883                    role: MessageRole::User,
884                    content: "Hello".to_string(),
885                },
886                ConversationMessage {
887                    role: MessageRole::Model,
888                    content: "Hi there!".to_string(),
889                },
890            ])
891            .build();
892
893        assert_eq!(config.initial_history.len(), 2);
894
895        let json = serde_json::to_string(&config).unwrap();
896        assert!(json.contains("initial_history"));
897
898        let parsed: AgentConfig = serde_json::from_str(&json).unwrap();
899        assert_eq!(parsed.initial_history.len(), 2);
900        assert_eq!(parsed.initial_history[0].role, MessageRole::User);
901        assert_eq!(parsed.initial_history[0].content, "Hello");
902        assert_eq!(parsed.initial_history[1].role, MessageRole::Model);
903        assert_eq!(parsed.initial_history[1].content, "Hi there!");
904    }
905}
agy_bridge/config/agent.rs

agy_bridge/config/
agent.rs
