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

agy_bridge/config/
agent.rs
