antigravity_sdk_rust/
agent.rs

1use crate::conversation::Conversation;
2use crate::hooks::{DynHook, HookRunner};
3#[cfg(not(target_arch = "wasm32"))]
4use crate::local::LocalConnectionStrategy;
5use crate::policy::{self, Policy, PolicyEnforcer};
6use crate::tools::{DynTool, ToolRunner};
7use crate::triggers::{DynTrigger, TriggerRunner};
8use crate::types::{
9    BuiltinTools, CapabilitiesConfig, ChatResponse, GeminiConfig, SystemInstructions,
10};
11use anyhow::anyhow;
12use futures_util::future::BoxFuture;
13use std::sync::Arc;
14
15/// Configuration settings used to customize the behavior and capabilities of an [`Agent`].
16#[derive(Default)]
17pub struct AgentConfig {
18    /// Optional path to the `localharness` binary. If not provided, it will be automatically
19    /// resolved via standard paths or standard environments.
20    pub binary_path: Option<String>,
21    /// Gemini LLM configuration details (API key, default models, thinking settings, etc.).
22    pub gemini_config: GeminiConfig,
23    /// Capabilities config specifying enabled/disabled tools and threshold limits.
24    pub capabilities: CapabilitiesConfig,
25    /// Optional system instructions (either appended template sections or fully custom text).
26    pub system_instructions: Option<SystemInstructions>,
27    /// Optional directory to save session state logs.
28    pub save_dir: Option<String>,
29    /// Configured workspaces. If not provided, defaults to the current working directory.
30    pub workspaces: Option<Vec<String>>,
31    /// Paths to local folders containing custom skill modules.
32    pub skills_paths: Vec<String>,
33    /// Set of safety policies (e.g., workspace lock, run command approvals) to restrict tool execution.
34    pub policies: Option<Vec<Policy>>,
35    /// Handlers triggered during agent lifecycle hooks (pre/post tool calls, start session, etc.).
36    pub hooks: Vec<Arc<dyn DynHook>>,
37    /// Custom triggers spawned when the agent starts.
38    pub triggers: Vec<Arc<dyn DynTrigger>>,
39    /// Custom Rust tools registered to be available for invocation.
40    pub tools: Vec<Arc<dyn DynTool>>,
41    /// Specific conversation ID to assign or resume.
42    pub conversation_id: Option<String>,
43    /// Path to the application data directory where cache/configs are stored.
44    pub app_data_dir: Option<String>,
45    /// Optional JSON schema constraining the final structured tool output.
46    pub response_schema: Option<String>,
47}
48
49impl std::fmt::Debug for AgentConfig {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        f.debug_struct("AgentConfig")
52            .field("binary_path", &self.binary_path)
53            .field("gemini_config", &self.gemini_config)
54            .field("capabilities", &self.capabilities)
55            .field("system_instructions", &self.system_instructions)
56            .field("save_dir", &self.save_dir)
57            .field("workspaces", &self.workspaces)
58            .field("skills_paths", &self.skills_paths)
59            .field("policies", &self.policies)
60            .field("hooks_count", &self.hooks.len())
61            .field("triggers_count", &self.triggers.len())
62            .field("tools_count", &self.tools.len())
63            .field("conversation_id", &self.conversation_id)
64            .field("app_data_dir", &self.app_data_dir)
65            .field("response_schema", &self.response_schema)
66            .finish()
67    }
68}
69
70/// High-level orchestrator that manages an agentic execution session.
71///
72/// An `Agent` encapsulates binary discovery, WebSocket upgrades, tool wiring, safety policy enforcement,
73/// and observer hook dispatch. It provides a simple `chat` API for sending prompts and retrieving responses.
74///
75/// # Examples
76///
77/// ```no_run
78/// use antigravity_sdk_rust::agent::Agent;
79///
80/// #[tokio::main]
81/// async fn main() -> Result<(), anyhow::Error> {
82///     let agent = Agent::builder()
83///         .allow_all()
84///         .build();
85///     let agent = agent.start().await?;
86///
87///     let response = agent.chat("What is 2+2?").await?;
88///     println!("Agent: {}", response.text);
89///
90///     let _ = agent.stop().await;
91///     Ok(())
92/// }
93/// ```
94/// Marker trait for all valid agent lifecycles.
95pub trait AgentLifecycle: Send + Sync + std::fmt::Debug {}
96
97/// Represents an agent that has been configured but not yet started.
98#[derive(Debug)]
99pub struct Unstarted;
100impl AgentLifecycle for Unstarted {}
101
102/// Represents an active, running agent session.
103pub struct Started {
104    pub(crate) conversation: Arc<Conversation>,
105    pub(crate) trigger_runner: Option<TriggerRunner>,
106}
107
108impl AgentLifecycle for Started {}
109
110impl std::fmt::Debug for Started {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        f.debug_struct("Started")
113            .field("conversation", &self.conversation)
114            .field("trigger_runner", &self.trigger_runner)
115            .finish()
116    }
117}
118
119pub struct Agent<S: AgentLifecycle = Unstarted> {
120    config: AgentConfig,
121    tool_runner: ToolRunner,
122    hook_runner: HookRunner,
123    state: S,
124}
125
126impl<S: AgentLifecycle> std::fmt::Debug for Agent<S> {
127    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128        f.debug_struct("Agent")
129            .field("config", &self.config)
130            .field("tool_runner", &self.tool_runner)
131            .field("hook_runner", &self.hook_runner)
132            .field("state", &self.state)
133            .finish()
134    }
135}
136
137impl Agent<Unstarted> {
138    /// Creates a new `Agent` with the given configuration.
139    pub fn new(config: AgentConfig) -> Self {
140        Self {
141            config,
142            tool_runner: ToolRunner::new(),
143            hook_runner: HookRunner::new(),
144            state: Unstarted,
145        }
146    }
147
148    /// Returns an `AgentBuilder` to configure and construct an `Agent`.
149    pub fn builder() -> AgentBuilder<NoPolicies> {
150        AgentBuilder::new()
151    }
152
153    /// Registers a custom lifecycle hook during configuration.
154    pub fn register_hook(&mut self, hook: Arc<dyn DynHook>) {
155        self.config.hooks.push(hook);
156    }
157
158    /// Registers a custom background trigger.
159    pub fn register_trigger(&mut self, trigger: Arc<dyn DynTrigger>) -> Result<(), anyhow::Error> {
160        self.config.triggers.push(trigger);
161        Ok(())
162    }
163
164    /// Registers a custom tool during configuration.
165    pub fn register_tool(&mut self, tool: Arc<dyn DynTool>) {
166        self.config.tools.push(tool);
167    }
168
169    /// Spawns the subprocess communication harness, initializes safety policies, registers tools/hooks,
170    /// establishes the WebSocket session, and starts any configured triggers.
171    ///
172    /// # Errors
173    ///
174    /// Returns an error if:
175    /// - The `localharness` binary cannot be resolved.
176    /// - Write tools are enabled but no safety policies are configured.
177    /// - The WebSocket upgrade or subprocess connection fails.
178    #[allow(clippy::too_many_lines)]
179    pub fn start(self) -> BoxFuture<'static, Result<Agent<Started>, anyhow::Error>> {
180        Box::pin(async move {
181            // 1. Resolve binary path
182            #[cfg(not(target_arch = "wasm32"))]
183        let binary_path = self.config.binary_path.clone()
184            .or_else(get_default_binary_path)
185            .ok_or_else(|| anyhow!("Could not find default localharness binary. Please specify binary_path explicitly."))?;
186
187            // 2. Setup hook runner and register pending hooks
188            for hook in &self.config.hooks {
189                self.hook_runner.register(hook.clone()).await;
190            }
191
192            // 3. Process capabilities and active tools
193            let enabled_tools = self.config.capabilities.enabled_tools.clone();
194            let disabled_tools = self.config.capabilities.disabled_tools.clone();
195            if enabled_tools.is_some() && disabled_tools.is_some() {
196                return Err(anyhow!(
197                    "enabled_tools and disabled_tools are mutually exclusive"
198                ));
199            }
200
201            let active_tools = enabled_tools.unwrap_or_else(|| {
202                disabled_tools.map_or_else(
203                    || {
204                        vec![
205                            BuiltinTools::CreateFile,
206                            BuiltinTools::EditFile,
207                            BuiltinTools::FindFile,
208                            BuiltinTools::ListDir,
209                            BuiltinTools::RunCommand,
210                            BuiltinTools::SearchDir,
211                            BuiltinTools::ViewFile,
212                            BuiltinTools::StartSubagent,
213                            BuiltinTools::GenerateImage,
214                            BuiltinTools::Finish,
215                        ]
216                    },
217                    |disabled| {
218                        let all = vec![
219                            BuiltinTools::CreateFile,
220                            BuiltinTools::EditFile,
221                            BuiltinTools::FindFile,
222                            BuiltinTools::ListDir,
223                            BuiltinTools::RunCommand,
224                            BuiltinTools::SearchDir,
225                            BuiltinTools::ViewFile,
226                            BuiltinTools::StartSubagent,
227                            BuiltinTools::GenerateImage,
228                            BuiltinTools::Finish,
229                        ];
230                        all.into_iter().filter(|t| !disabled.contains(t)).collect()
231                    },
232                )
233            });
234
235            let read_only = BuiltinTools::read_only();
236            let has_write_tools = active_tools.iter().any(|t| !read_only.contains(t));
237
238            // 4. Set up policies
239            let mut final_policies = self.config.policies.clone().unwrap_or_else(|| {
240                // Default to confirm_run_command
241                policy::confirm_run_command(None)
242            });
243
244            // Prepend workspace scoping policies ONLY if the caller has not explicitly opted
245            // into allow_all(). When allow_all() is in the policy set, the intent is to
246            // approve every tool call (typically managed by a ConfirmHook instead). In that
247            // case, prepending workspace_only would silently override the user's approval
248            // because Deny policies have higher bucket priority than wildcard Approve policies.
249            //
250            // We detect allow_all by looking for a wildcard Approve policy named "allow_all".
251            // If found, skip the workspace gate entirely.
252            let has_allow_all = final_policies.iter().any(|p| {
253                p.tool == "*"
254                    && p.decision == crate::policy::Decision::Approve
255                    && p.name == "allow_all"
256            });
257
258            if !has_allow_all {
259                let workspaces = self.config.workspaces.clone().unwrap_or_else(|| {
260                    std::env::current_dir().map_or_else(
261                        |_| Vec::new(),
262                        |cwd| vec![cwd.to_string_lossy().into_owned()],
263                    )
264                });
265
266                if !workspaces.is_empty() {
267                    let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
268                    let app_data_dir = self
269                        .config
270                        .app_data_dir
271                        .clone()
272                        .unwrap_or_else(|| format!("{home}/.gemini/antigravity"));
273                    let mut allowed_paths = workspaces;
274                    allowed_paths.push(app_data_dir);
275                    let mut ws_policies = policy::workspace_only(allowed_paths);
276                    ws_policies.append(&mut final_policies);
277                    final_policies = ws_policies;
278                }
279            }
280
281            // Safety policy check: if write tools are enabled, policies cannot be empty
282            if has_write_tools && final_policies.is_empty() {
283                return Err(anyhow!(
284                    "Write tools are enabled without a safety policy. Add policies=[policy.allow_all()] to approve all tool calls, or policies=[policy.deny_all(), policy.allow(\"tool_name\")] to selectively allow specific tools."
285                ));
286            }
287
288            if !final_policies.is_empty() {
289                let enforcer = Arc::new(PolicyEnforcer::new(final_policies));
290                self.hook_runner.register(enforcer).await;
291            }
292
293            // 5. Register configured tools
294            for tool in &self.config.tools {
295                self.tool_runner.register(tool.clone()).await;
296            }
297
298            // 6. Build and connect strategy
299            #[cfg(target_arch = "wasm32")]
300            {
301                let mut cap = self.config.capabilities.clone();
302                if let Some(ref schema) = self.config.response_schema {
303                    cap.finish_tool_schema_json = Some(schema.clone());
304                }
305
306                let strategy = crate::wasm::WasmConnectionStrategy {
307                    gemini_config: self.config.gemini_config.clone(),
308                    capabilities_config: cap,
309                    system_instructions: self.config.system_instructions.clone(),
310                    save_dir: self.config.save_dir.clone(),
311                    workspaces: self.config.workspaces.clone().unwrap_or_default(),
312                    skills_paths: self.config.skills_paths.clone(),
313                    tool_runner: Some(self.tool_runner.clone()),
314                    hook_runner: Some(self.hook_runner.clone()),
315                    conversation_id: self.config.conversation_id.clone().unwrap_or_default(),
316                };
317
318                let conn = strategy.connect().await?;
319                let conversation = Arc::new(Conversation::new(
320                    crate::connection::AnyConnection::Wasm(Arc::new(conn)),
321                    None,
322                ));
323
324                // 7. Start triggers
325                let mut trigger_runner = None;
326                if !self.config.triggers.is_empty() {
327                    let runner = TriggerRunner::new(self.config.triggers.clone());
328                    runner.start(&conversation.connection());
329                    trigger_runner = Some(runner);
330                }
331
332                Ok(Agent {
333                    config: self.config,
334                    tool_runner: self.tool_runner,
335                    hook_runner: self.hook_runner,
336                    state: Started {
337                        conversation,
338                        trigger_runner,
339                    },
340                })
341            }
342
343            #[cfg(not(target_arch = "wasm32"))]
344            {
345                let mut cap = self.config.capabilities.clone();
346                if let Some(ref schema) = self.config.response_schema {
347                    cap.finish_tool_schema_json = Some(schema.clone());
348                }
349
350                let strategy = LocalConnectionStrategy::new(
351                    binary_path,
352                    self.config.gemini_config.clone(),
353                    cap,
354                    self.config.system_instructions.clone(),
355                    self.config.save_dir.clone(),
356                    self.config.workspaces.clone().unwrap_or_default(),
357                    self.config.skills_paths.clone(),
358                    Some(self.tool_runner.clone()),
359                    Some(self.hook_runner.clone()),
360                    self.config.conversation_id.clone().unwrap_or_default(),
361                );
362
363                let conn = strategy.connect().await?;
364                let conversation = Arc::new(Conversation::new(
365                    crate::connection::AnyConnection::Local(Arc::new(conn)),
366                    None,
367                ));
368
369                // 7. Start triggers
370                let trigger_runner = if self.config.triggers.is_empty() {
371                    None
372                } else {
373                    let runner = TriggerRunner::new(self.config.triggers.clone());
374                    runner.start(&conversation.connection());
375                    Some(runner)
376                };
377
378                Ok(Agent {
379                    config: self.config,
380                    tool_runner: self.tool_runner,
381                    hook_runner: self.hook_runner,
382                    state: Started {
383                        conversation,
384                        trigger_runner,
385                    },
386                })
387            }
388        }) // end Box::pin
389    }
390}
391
392impl Agent<Started> {
393    /// Sends a prompt message to the active agent session and awaits the final completed response.
394    ///
395    /// # Errors
396    ///
397    /// Returns an error if the execution stream encounters a failure.
398    pub async fn chat(&self, prompt: &str) -> Result<ChatResponse, anyhow::Error> {
399        self.state.conversation.chat_to_completion(prompt).await
400    }
401
402    /// Returns the active [`Conversation`] session.
403    pub fn conversation(&self) -> Arc<Conversation> {
404        self.state.conversation.clone()
405    }
406
407    /// Returns the active conversation ID.
408    pub fn conversation_id(&self) -> String {
409        self.state.conversation.conversation_id().to_string()
410    }
411
412    /// Gracefully stops the agent connection and disconnects the underlying harness.
413    ///
414    /// # Errors
415    ///
416    /// Returns an error if closing the connection fails.
417    pub async fn stop(&self) -> Result<(), anyhow::Error> {
418        self.state.conversation.disconnect().await?;
419        Ok(())
420    }
421}
422
423#[derive(Debug)]
424pub struct NoPolicies;
425#[derive(Debug)]
426pub struct HasPolicies;
427
428pub struct AgentBuilder<P = NoPolicies> {
429    config: AgentConfig,
430    _policy_marker: std::marker::PhantomData<P>,
431}
432
433impl<P> std::fmt::Debug for AgentBuilder<P> {
434    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
435        f.debug_struct("AgentBuilder")
436            .field("config", &self.config)
437            .finish()
438    }
439}
440
441impl AgentBuilder<NoPolicies> {
442    pub fn new() -> Self {
443        Self {
444            config: AgentConfig::default(),
445            _policy_marker: std::marker::PhantomData,
446        }
447    }
448}
449
450impl Default for AgentBuilder<NoPolicies> {
451    fn default() -> Self {
452        Self::new()
453    }
454}
455
456impl<P> AgentBuilder<P> {
457    pub fn binary_path(mut self, path: impl Into<String>) -> Self {
458        self.config.binary_path = Some(path.into());
459        self
460    }
461
462    pub fn gemini_config(mut self, gemini_config: GeminiConfig) -> Self {
463        self.config.gemini_config = gemini_config;
464        self
465    }
466
467    pub fn api_key(mut self, key: impl Into<String>) -> Self {
468        self.config.gemini_config.api_key = Some(key.into());
469        self
470    }
471
472    pub fn default_model(mut self, model: impl Into<String>) -> Self {
473        self.config.gemini_config.models.default.name = model.into();
474        self
475    }
476
477    pub fn capabilities(mut self, capabilities: CapabilitiesConfig) -> Self {
478        self.config.capabilities = capabilities;
479        self
480    }
481
482    pub fn system_instructions(mut self, system_instructions: SystemInstructions) -> Self {
483        self.config.system_instructions = Some(system_instructions);
484        self
485    }
486
487    pub fn save_dir(mut self, save_dir: impl Into<String>) -> Self {
488        self.config.save_dir = Some(save_dir.into());
489        self
490    }
491
492    pub fn workspaces(mut self, workspaces: Vec<String>) -> Self {
493        self.config.workspaces = Some(workspaces);
494        self
495    }
496
497    pub fn skills_paths(mut self, skills_paths: Vec<String>) -> Self {
498        self.config.skills_paths = skills_paths;
499        self
500    }
501
502    pub fn hooks(mut self, hooks: Vec<Arc<dyn DynHook>>) -> Self {
503        self.config.hooks = hooks;
504        self
505    }
506
507    pub fn triggers(mut self, triggers: Vec<Arc<dyn DynTrigger>>) -> Self {
508        self.config.triggers = triggers;
509        self
510    }
511
512    pub fn tools(mut self, tools: Vec<Arc<dyn DynTool>>) -> Self {
513        self.config.tools = tools;
514        self
515    }
516
517    pub fn tool(mut self, tool: Arc<dyn DynTool>) -> Self {
518        self.config.tools.push(tool);
519        self
520    }
521
522    pub fn hook(mut self, hook: Arc<dyn DynHook>) -> Self {
523        self.config.hooks.push(hook);
524        self
525    }
526
527    pub fn trigger(mut self, trigger: Arc<dyn DynTrigger>) -> Self {
528        self.config.triggers.push(trigger);
529        self
530    }
531
532    pub fn policy(mut self, policy: Policy) -> AgentBuilder<HasPolicies> {
533        let mut policies = self.config.policies.take().unwrap_or_default();
534        policies.push(policy);
535        self.config.policies = Some(policies);
536        AgentBuilder {
537            config: self.config,
538            _policy_marker: std::marker::PhantomData,
539        }
540    }
541
542    pub fn conversation_id(mut self, conversation_id: impl Into<String>) -> Self {
543        self.config.conversation_id = Some(conversation_id.into());
544        self
545    }
546
547    pub fn app_data_dir(mut self, app_data_dir: impl Into<String>) -> Self {
548        self.config.app_data_dir = Some(app_data_dir.into());
549        self
550    }
551
552    pub fn response_schema(mut self, response_schema: impl Into<String>) -> Self {
553        self.config.response_schema = Some(response_schema.into());
554        self
555    }
556
557    pub fn policies(self, policies: Vec<Policy>) -> AgentBuilder<HasPolicies> {
558        let mut config = self.config;
559        config.policies = Some(policies);
560        AgentBuilder {
561            config,
562            _policy_marker: std::marker::PhantomData,
563        }
564    }
565
566    pub fn allow_all(self) -> AgentBuilder<HasPolicies> {
567        let mut config = self.config;
568        config.policies = Some(vec![policy::allow_all()]);
569        AgentBuilder {
570            config,
571            _policy_marker: std::marker::PhantomData,
572        }
573    }
574
575    pub fn read_only(self) -> AgentBuilder<HasPolicies> {
576        let mut config = self.config;
577        let read_only_tools = BuiltinTools::read_only();
578        let mut policies = vec![policy::deny_all()];
579        for tool in read_only_tools {
580            policies.push(policy::allow(tool.as_str()));
581        }
582        config.policies = Some(policies);
583        AgentBuilder {
584            config,
585            _policy_marker: std::marker::PhantomData,
586        }
587    }
588
589    /// Builder escape hatch to construct `Agent<Unstarted>` without compile-time check for policies.
590    pub fn build_unchecked(self) -> Agent<Unstarted> {
591        Agent::new(self.config)
592    }
593}
594
595impl AgentBuilder<HasPolicies> {
596    pub fn build(self) -> Agent<Unstarted> {
597        Agent::new(self.config)
598    }
599}
600
601#[cfg(not(target_arch = "wasm32"))]
602fn get_default_binary_path() -> Option<String> {
603    if let Ok(path) = std::env::var("ANTIGRAVITY_HARNESS_PATH") {
604        return Some(path);
605    }
606    // Check if it is in standard PATH
607    if let Ok(paths) = std::env::var("PATH") {
608        for path in std::env::split_paths(&paths) {
609            let p = path.join("localharness");
610            if p.exists() {
611                return Some(p.to_string_lossy().into_owned());
612            }
613        }
614    }
615    // Check Python site-packages as a fallback since google-antigravity Python package installs it there
616    if let Some(output) = std::process::Command::new("python3")
617        .args([
618            "-c",
619            "import site; print('\\n'.join(site.getsitepackages()))",
620        ])
621        .output()
622        .ok()
623        .filter(|o| o.status.success())
624    {
625        let stdout = String::from_utf8_lossy(&output.stdout);
626        for line in stdout.lines() {
627            let p = std::path::Path::new(line.trim())
628                .join("google")
629                .join("antigravity")
630                .join("bin")
631                .join("localharness");
632            if p.exists() {
633                return Some(p.to_string_lossy().into_owned());
634            }
635        }
636    }
637    None
638}
antigravity_sdk_rust/agent.rs

antigravity_sdk_rust/
agent.rs
