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 std::sync::Arc;
13
14/// Configuration settings used to customize the behavior and capabilities of an [`Agent`].
15#[derive(Default)]
16pub struct AgentConfig {
17    /// Optional path to the `localharness` binary. If not provided, it will be automatically
18    /// resolved via standard paths or standard environments.
19    pub binary_path: Option<String>,
20    /// Gemini LLM configuration details (API key, default models, thinking settings, etc.).
21    pub gemini_config: GeminiConfig,
22    /// Capabilities config specifying enabled/disabled tools and threshold limits.
23    pub capabilities: CapabilitiesConfig,
24    /// Optional system instructions (either appended template sections or fully custom text).
25    pub system_instructions: Option<SystemInstructions>,
26    /// Optional directory to save session state logs.
27    pub save_dir: Option<String>,
28    /// Configured workspaces. If not provided, defaults to the current working directory.
29    pub workspaces: Option<Vec<String>>,
30    /// Paths to local folders containing custom skill modules.
31    pub skills_paths: Vec<String>,
32    /// Set of safety policies (e.g., workspace lock, run command approvals) to restrict tool execution.
33    pub policies: Option<Vec<Policy>>,
34    /// Handlers triggered during agent lifecycle hooks (pre/post tool calls, start session, etc.).
35    pub hooks: Vec<Arc<dyn DynHook>>,
36    /// Custom triggers spawned when the agent starts.
37    pub triggers: Vec<Arc<dyn DynTrigger>>,
38    /// Custom Rust tools registered to be available for invocation.
39    pub tools: Vec<Arc<dyn DynTool>>,
40    /// Specific conversation ID to assign or resume.
41    pub conversation_id: Option<String>,
42    /// Path to the application data directory where cache/configs are stored.
43    pub app_data_dir: Option<String>,
44    /// Optional JSON schema constraining the final structured tool output.
45    pub response_schema: Option<String>,
46}
47
48impl std::fmt::Debug for AgentConfig {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        f.debug_struct("AgentConfig")
51            .field("binary_path", &self.binary_path)
52            .field("gemini_config", &self.gemini_config)
53            .field("capabilities", &self.capabilities)
54            .field("system_instructions", &self.system_instructions)
55            .field("save_dir", &self.save_dir)
56            .field("workspaces", &self.workspaces)
57            .field("skills_paths", &self.skills_paths)
58            .field("policies", &self.policies)
59            .field("hooks_count", &self.hooks.len())
60            .field("triggers_count", &self.triggers.len())
61            .field("tools_count", &self.tools.len())
62            .field("conversation_id", &self.conversation_id)
63            .field("app_data_dir", &self.app_data_dir)
64            .field("response_schema", &self.response_schema)
65            .finish()
66    }
67}
68
69/// High-level orchestrator that manages an agentic execution session.
70///
71/// An `Agent` encapsulates binary discovery, WebSocket upgrades, tool wiring, safety policy enforcement,
72/// and observer hook dispatch. It provides a simple `chat` API for sending prompts and retrieving responses.
73///
74/// # Examples
75///
76/// ```no_run
77/// use antigravity_sdk_rust::agent::{Agent, AgentConfig};
78/// use antigravity_sdk_rust::policy;
79///
80/// #[tokio::main]
81/// async fn main() -> Result<(), anyhow::Error> {
82///     let mut config = AgentConfig::default();
83///     config.policies = Some(vec![policy::allow_all()]);
84///
85///     let mut agent = Agent::new(config);
86///     agent.start().await?;
87///
88///     let response = agent.chat("What is 2+2?").await?;
89///     println!("Agent: {}", response.text);
90///
91///     agent.stop().await?;
92///     Ok(())
93/// }
94/// ```
95pub struct Agent {
96    config: AgentConfig,
97    conversation: Option<Arc<Conversation>>,
98    tool_runner: ToolRunner,
99    hook_runner: HookRunner,
100    trigger_runner: Option<TriggerRunner>,
101}
102
103impl std::fmt::Debug for Agent {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        f.debug_struct("Agent")
106            .field("config", &self.config)
107            .field("conversation", &self.conversation)
108            .field("tool_runner", &self.tool_runner)
109            .field("hook_runner", &self.hook_runner)
110            .field("trigger_runner", &self.trigger_runner)
111            .finish()
112    }
113}
114
115impl Agent {
116    /// Creates a new `Agent` with the given configuration.
117    pub fn new(config: AgentConfig) -> Self {
118        Self {
119            config,
120            conversation: None,
121            tool_runner: ToolRunner::new(),
122            hook_runner: HookRunner::new(),
123            trigger_runner: None,
124        }
125    }
126
127    /// Registers a custom lifecycle hook. Hooks can observe or modify agent transitions.
128    pub fn register_hook(&self, hook: Arc<dyn DynHook>) {
129        let hr = self.hook_runner.clone();
130        crate::spawn_task(async move {
131            hr.register(hook).await;
132        });
133    }
134
135    /// Registers a custom background trigger. Triggers must be registered *before* the agent starts.
136    ///
137    /// # Errors
138    ///
139    /// Returns an error if the agent session has already been started.
140    pub fn register_trigger(&mut self, trigger: Arc<dyn DynTrigger>) -> Result<(), anyhow::Error> {
141        if self.conversation.is_some() {
142            return Err(anyhow!(
143                "Cannot register triggers after the agent has started."
144            ));
145        }
146        self.config.triggers.push(trigger);
147        Ok(())
148    }
149
150    /// Registers a custom tool available for execution by the agent.
151    pub fn register_tool(&self, tool: Arc<dyn DynTool>) {
152        let tr = self.tool_runner.clone();
153        crate::spawn_task(async move {
154            tr.register(tool).await;
155        });
156    }
157
158    /// Spawns the subprocess communication harness, initializes safety policies, registers tools/hooks,
159    /// establishes the WebSocket session, and starts any configured triggers.
160    ///
161    /// # Errors
162    ///
163    /// Returns an error if:
164    /// - The `localharness` binary cannot be resolved.
165    /// - Write tools are enabled but no safety policies are configured.
166    /// - The WebSocket upgrade or subprocess connection fails.
167    #[allow(clippy::too_many_lines)]
168    pub async fn start(&mut self) -> Result<(), anyhow::Error> {
169        if self.conversation.is_some() {
170            return Ok(());
171        }
172
173        // 1. Resolve binary path
174        #[cfg(not(target_arch = "wasm32"))]
175        let binary_path = self.config.binary_path.clone()
176            .or_else(get_default_binary_path)
177            .ok_or_else(|| anyhow!("Could not find default localharness binary. Please specify binary_path explicitly."))?;
178
179        // 2. Setup hook runner and register pending hooks
180        for hook in &self.config.hooks {
181            self.hook_runner.register(hook.clone()).await;
182        }
183
184        // 3. Process capabilities and active tools
185        let enabled_tools = self.config.capabilities.enabled_tools.clone();
186        let disabled_tools = self.config.capabilities.disabled_tools.clone();
187        if enabled_tools.is_some() && disabled_tools.is_some() {
188            return Err(anyhow!(
189                "enabled_tools and disabled_tools are mutually exclusive"
190            ));
191        }
192
193        let active_tools = enabled_tools.unwrap_or_else(|| {
194            disabled_tools.map_or_else(
195                || {
196                    vec![
197                        BuiltinTools::CreateFile,
198                        BuiltinTools::EditFile,
199                        BuiltinTools::FindFile,
200                        BuiltinTools::ListDir,
201                        BuiltinTools::RunCommand,
202                        BuiltinTools::SearchDir,
203                        BuiltinTools::ViewFile,
204                        BuiltinTools::StartSubagent,
205                        BuiltinTools::GenerateImage,
206                        BuiltinTools::Finish,
207                    ]
208                },
209                |disabled| {
210                    let all = vec![
211                        BuiltinTools::CreateFile,
212                        BuiltinTools::EditFile,
213                        BuiltinTools::FindFile,
214                        BuiltinTools::ListDir,
215                        BuiltinTools::RunCommand,
216                        BuiltinTools::SearchDir,
217                        BuiltinTools::ViewFile,
218                        BuiltinTools::StartSubagent,
219                        BuiltinTools::GenerateImage,
220                        BuiltinTools::Finish,
221                    ];
222                    all.into_iter().filter(|t| !disabled.contains(t)).collect()
223                },
224            )
225        });
226
227        let read_only = BuiltinTools::read_only();
228        let has_write_tools = active_tools.iter().any(|t| !read_only.contains(t));
229
230        // 4. Set up policies
231        let mut final_policies = self.config.policies.clone().unwrap_or_else(|| {
232            // Default to confirm_run_command
233            policy::confirm_run_command(None)
234        });
235
236        // Prepend workspace scoping policies if workspaces are configured
237        let workspaces = self.config.workspaces.clone().unwrap_or_else(|| {
238            std::env::current_dir().map_or_else(
239                |_| Vec::new(),
240                |cwd| vec![cwd.to_string_lossy().into_owned()],
241            )
242        });
243
244        if !workspaces.is_empty() {
245            let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
246            let app_data_dir = self
247                .config
248                .app_data_dir
249                .clone()
250                .unwrap_or_else(|| format!("{home}/.gemini/antigravity"));
251            let mut allowed_paths = workspaces;
252            allowed_paths.push(app_data_dir);
253            let mut ws_policies = policy::workspace_only(allowed_paths);
254            ws_policies.append(&mut final_policies);
255            final_policies = ws_policies;
256        }
257
258        // Safety policy check: if write tools are enabled, policies cannot be empty
259        if has_write_tools && final_policies.is_empty() {
260            return Err(anyhow!(
261                "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."
262            ));
263        }
264
265        if !final_policies.is_empty() {
266            let enforcer = Arc::new(PolicyEnforcer::new(final_policies));
267            self.hook_runner.register(enforcer).await;
268        }
269
270        // 5. Register configured tools
271        for tool in &self.config.tools {
272            self.tool_runner.register(tool.clone()).await;
273        }
274
275        // 6. Build and connect strategy
276        #[cfg(target_arch = "wasm32")]
277        {
278            let mut cap = self.config.capabilities.clone();
279            if let Some(ref schema) = self.config.response_schema {
280                cap.finish_tool_schema_json = Some(schema.clone());
281            }
282
283            let strategy = crate::wasm::WasmConnectionStrategy {
284                gemini_config: self.config.gemini_config.clone(),
285                capabilities_config: cap,
286                system_instructions: self.config.system_instructions.clone(),
287                save_dir: self.config.save_dir.clone(),
288                workspaces: self.config.workspaces.clone().unwrap_or_default(),
289                skills_paths: self.config.skills_paths.clone(),
290                tool_runner: Some(self.tool_runner.clone()),
291                hook_runner: Some(self.hook_runner.clone()),
292                conversation_id: self.config.conversation_id.clone().unwrap_or_default(),
293            };
294
295            let conn = strategy.connect().await?;
296            let conversation = Arc::new(Conversation::new(
297                crate::connection::AnyConnection::Wasm(Arc::new(conn)),
298                None,
299            ));
300            self.conversation = Some(conversation.clone());
301
302            // 7. Start triggers
303            if !self.config.triggers.is_empty() {
304                let runner = TriggerRunner::new(self.config.triggers.clone());
305                runner.start(&conversation.connection());
306                self.trigger_runner = Some(runner);
307            }
308
309            Ok(())
310        }
311
312        #[cfg(not(target_arch = "wasm32"))]
313        {
314            let mut cap = self.config.capabilities.clone();
315            if let Some(ref schema) = self.config.response_schema {
316                cap.finish_tool_schema_json = Some(schema.clone());
317            }
318
319            let strategy = LocalConnectionStrategy::new(
320                binary_path,
321                self.config.gemini_config.clone(),
322                cap,
323                self.config.system_instructions.clone(),
324                self.config.save_dir.clone(),
325                self.config.workspaces.clone().unwrap_or_default(),
326                self.config.skills_paths.clone(),
327                Some(self.tool_runner.clone()),
328                Some(self.hook_runner.clone()),
329                self.config.conversation_id.clone().unwrap_or_default(),
330            );
331
332            let conn = strategy.connect().await?;
333            let conversation = Arc::new(Conversation::new(
334                crate::connection::AnyConnection::Local(Arc::new(conn)),
335                None,
336            ));
337            self.conversation = Some(conversation.clone());
338
339            // 7. Start triggers
340            if !self.config.triggers.is_empty() {
341                let runner = TriggerRunner::new(self.config.triggers.clone());
342                runner.start(&conversation.connection());
343                self.trigger_runner = Some(runner);
344            }
345
346            Ok(())
347        }
348    }
349
350    /// Sends a prompt message to the active agent session and awaits the final completed response.
351    ///
352    /// # Errors
353    ///
354    /// Returns an error if the agent is not yet started or if the execution stream encounters a failure.
355    pub async fn chat(&self, prompt: &str) -> Result<ChatResponse, anyhow::Error> {
356        let conversation = self.conversation()?;
357        conversation.chat_to_completion(prompt).await
358    }
359
360    /// Returns the active [`Conversation`] session.
361    ///
362    /// # Errors
363    ///
364    /// Returns an error if the agent is not yet started.
365    pub fn conversation(&self) -> Result<Arc<Conversation>, anyhow::Error> {
366        self.conversation
367            .clone()
368            .ok_or_else(|| anyhow!("Agent session not started. Use start() first."))
369    }
370
371    /// Returns the active conversation ID if the session has started.
372    pub fn conversation_id(&self) -> Option<String> {
373        self.conversation
374            .as_ref()
375            .map(|c| c.conversation_id().to_string())
376    }
377
378    /// Gracefully stops the agent connection and disconnects the underlying harness.
379    ///
380    /// # Errors
381    ///
382    /// Returns an error if closing the connection fails.
383    pub async fn stop(&mut self) -> Result<(), anyhow::Error> {
384        if let Some(conversation) = self.conversation.take() {
385            conversation.disconnect().await?;
386        }
387        Ok(())
388    }
389}
390
391#[cfg(not(target_arch = "wasm32"))]
392fn get_default_binary_path() -> Option<String> {
393    if let Ok(path) = std::env::var("ANTIGRAVITY_HARNESS_PATH") {
394        return Some(path);
395    }
396    // Check if it is in standard PATH
397    if let Ok(paths) = std::env::var("PATH") {
398        for path in std::env::split_paths(&paths) {
399            let p = path.join("localharness");
400            if p.exists() {
401                return Some(p.to_string_lossy().into_owned());
402            }
403        }
404    }
405    // Check Python site-packages as a fallback since google-antigravity Python package installs it there
406    if let Some(output) = std::process::Command::new("python3")
407        .args([
408            "-c",
409            "import site; print('\\n'.join(site.getsitepackages()))",
410        ])
411        .output()
412        .ok()
413        .filter(|o| o.status.success())
414    {
415        let stdout = String::from_utf8_lossy(&output.stdout);
416        for line in stdout.lines() {
417            let p = std::path::Path::new(line.trim())
418                .join("google")
419                .join("antigravity")
420                .join("bin")
421                .join("localharness");
422            if p.exists() {
423                return Some(p.to_string_lossy().into_owned());
424            }
425        }
426    }
427    None
428}
antigravity_sdk_rust/agent.rs

antigravity_sdk_rust/
agent.rs
