antigravity-sdk-rust 0.1.3

use crate::conversation::Conversation;
use crate::hooks::{Hook, HookRunner};
#[cfg(not(target_arch = "wasm32"))]
use crate::local::LocalConnectionStrategy;
use crate::policy::{self, Policy, PolicyEnforcer};
use crate::tools::{Tool, ToolRunner};
use crate::triggers::{Trigger, TriggerRunner};
use crate::types::{
    BuiltinTools, CapabilitiesConfig, ChatResponse, GeminiConfig, SystemInstructions,
};
use anyhow::anyhow;
use std::sync::Arc;

/// Configuration settings used to customize the behavior and capabilities of an [`Agent`].
#[derive(Default)]
pub struct AgentConfig {
    /// Optional path to the `localharness` binary. If not provided, it will be automatically
    /// resolved via standard paths or standard environments.
    pub binary_path: Option<String>,
    /// Gemini LLM configuration details (API key, default models, thinking settings, etc.).
    pub gemini_config: GeminiConfig,
    /// Capabilities config specifying enabled/disabled tools and threshold limits.
    pub capabilities: CapabilitiesConfig,
    /// Optional system instructions (either appended template sections or fully custom text).
    pub system_instructions: Option<SystemInstructions>,
    /// Optional directory to save session state logs.
    pub save_dir: Option<String>,
    /// Configured workspaces. If not provided, defaults to the current working directory.
    pub workspaces: Option<Vec<String>>,
    /// Paths to local folders containing custom skill modules.
    pub skills_paths: Vec<String>,
    /// Set of safety policies (e.g., workspace lock, run command approvals) to restrict tool execution.
    pub policies: Option<Vec<Policy>>,
    /// Handlers triggered during agent lifecycle hooks (pre/post tool calls, start session, etc.).
    pub hooks: Vec<Arc<dyn Hook>>,
    /// Custom triggers spawned when the agent starts.
    pub triggers: Vec<Arc<dyn Trigger>>,
    /// Custom Rust tools registered to be available for invocation.
    pub tools: Vec<Arc<dyn Tool>>,
    /// Specific conversation ID to assign or resume.
    pub conversation_id: Option<String>,
    /// Path to the application data directory where cache/configs are stored.
    pub app_data_dir: Option<String>,
    /// Optional JSON schema constraining the final structured tool output.
    pub response_schema: Option<String>,
}

impl std::fmt::Debug for AgentConfig {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AgentConfig")
            .field("binary_path", &self.binary_path)
            .field("gemini_config", &self.gemini_config)
            .field("capabilities", &self.capabilities)
            .field("system_instructions", &self.system_instructions)
            .field("save_dir", &self.save_dir)
            .field("workspaces", &self.workspaces)
            .field("skills_paths", &self.skills_paths)
            .field("policies", &self.policies)
            .field("hooks_count", &self.hooks.len())
            .field("triggers_count", &self.triggers.len())
            .field("tools_count", &self.tools.len())
            .field("conversation_id", &self.conversation_id)
            .field("app_data_dir", &self.app_data_dir)
            .field("response_schema", &self.response_schema)
            .finish()
    }
}

/// High-level orchestrator that manages an agentic execution session.
///
/// An `Agent` encapsulates binary discovery, WebSocket upgrades, tool wiring, safety policy enforcement,
/// and observer hook dispatch. It provides a simple `chat` API for sending prompts and retrieving responses.
///
/// # Examples
///
/// ```no_run
/// use antigravity_sdk_rust::agent::{Agent, AgentConfig};
/// use antigravity_sdk_rust::policy;
///
/// #[tokio::main]
/// async fn main() -> Result<(), anyhow::Error> {
///     let mut config = AgentConfig::default();
///     config.policies = Some(vec![policy::allow_all()]);
///
///     let mut agent = Agent::new(config);
///     agent.start().await?;
///
///     let response = agent.chat("What is 2+2?").await?;
///     println!("Agent: {}", response.text);
///
///     agent.stop().await?;
///     Ok(())
/// }
/// ```
pub struct Agent {
    config: AgentConfig,
    conversation: Option<Arc<Conversation>>,
    tool_runner: ToolRunner,
    hook_runner: HookRunner,
    trigger_runner: Option<TriggerRunner>,
}

impl std::fmt::Debug for Agent {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Agent")
            .field("config", &self.config)
            .field("conversation", &self.conversation)
            .field("tool_runner", &self.tool_runner)
            .field("hook_runner", &self.hook_runner)
            .field("trigger_runner", &self.trigger_runner)
            .finish()
    }
}

impl Agent {
    /// Creates a new `Agent` with the given configuration.
    pub fn new(config: AgentConfig) -> Self {
        Self {
            config,
            conversation: None,
            tool_runner: ToolRunner::new(),
            hook_runner: HookRunner::new(),
            trigger_runner: None,
        }
    }

    /// Registers a custom lifecycle hook. Hooks can observe or modify agent transitions.
    pub fn register_hook(&self, hook: Arc<dyn Hook>) {
        let hr = self.hook_runner.clone();
        crate::spawn_task(async move {
            hr.register(hook).await;
        });
    }

    /// Registers a custom background trigger. Triggers must be registered *before* the agent starts.
    ///
    /// # Errors
    ///
    /// Returns an error if the agent session has already been started.
    pub fn register_trigger(&mut self, trigger: Arc<dyn Trigger>) -> Result<(), anyhow::Error> {
        if self.conversation.is_some() {
            return Err(anyhow!(
                "Cannot register triggers after the agent has started."
            ));
        }
        self.config.triggers.push(trigger);
        Ok(())
    }

    /// Registers a custom tool available for execution by the agent.
    pub fn register_tool(&self, tool: Arc<dyn Tool>) {
        let tr = self.tool_runner.clone();
        crate::spawn_task(async move {
            tr.register(tool).await;
        });
    }

    /// Spawns the subprocess communication harness, initializes safety policies, registers tools/hooks,
    /// establishes the WebSocket session, and starts any configured triggers.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The `localharness` binary cannot be resolved.
    /// - Write tools are enabled but no safety policies are configured.
    /// - The WebSocket upgrade or subprocess connection fails.
    #[allow(clippy::too_many_lines)]
    pub async fn start(&mut self) -> Result<(), anyhow::Error> {
        if self.conversation.is_some() {
            return Ok(());
        }

        // 1. Resolve binary path
        #[cfg(not(target_arch = "wasm32"))]
        let binary_path = self.config.binary_path.clone()
            .or_else(get_default_binary_path)
            .ok_or_else(|| anyhow!("Could not find default localharness binary. Please specify binary_path explicitly."))?;

        // 2. Setup hook runner and register pending hooks
        for hook in &self.config.hooks {
            self.hook_runner.register(hook.clone()).await;
        }

        // 3. Process capabilities and active tools
        let enabled_tools = self.config.capabilities.enabled_tools.clone();
        let disabled_tools = self.config.capabilities.disabled_tools.clone();
        if enabled_tools.is_some() && disabled_tools.is_some() {
            return Err(anyhow!(
                "enabled_tools and disabled_tools are mutually exclusive"
            ));
        }

        let active_tools = enabled_tools.unwrap_or_else(|| {
            disabled_tools.map_or_else(
                || {
                    vec![
                        BuiltinTools::CreateFile,
                        BuiltinTools::EditFile,
                        BuiltinTools::FindFile,
                        BuiltinTools::ListDir,
                        BuiltinTools::RunCommand,
                        BuiltinTools::SearchDir,
                        BuiltinTools::ViewFile,
                        BuiltinTools::StartSubagent,
                        BuiltinTools::GenerateImage,
                        BuiltinTools::Finish,
                    ]
                },
                |disabled| {
                    let all = vec![
                        BuiltinTools::CreateFile,
                        BuiltinTools::EditFile,
                        BuiltinTools::FindFile,
                        BuiltinTools::ListDir,
                        BuiltinTools::RunCommand,
                        BuiltinTools::SearchDir,
                        BuiltinTools::ViewFile,
                        BuiltinTools::StartSubagent,
                        BuiltinTools::GenerateImage,
                        BuiltinTools::Finish,
                    ];
                    all.into_iter().filter(|t| !disabled.contains(t)).collect()
                },
            )
        });

        let read_only = BuiltinTools::read_only();
        let has_write_tools = active_tools.iter().any(|t| !read_only.contains(t));

        // 4. Set up policies
        let mut final_policies = self.config.policies.clone().unwrap_or_else(|| {
            // Default to confirm_run_command
            policy::confirm_run_command(None)
        });

        // Prepend workspace scoping policies if workspaces are configured
        let workspaces = self.config.workspaces.clone().unwrap_or_else(|| {
            std::env::current_dir().map_or_else(
                |_| Vec::new(),
                |cwd| vec![cwd.to_string_lossy().into_owned()],
            )
        });

        if !workspaces.is_empty() {
            let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
            let app_data_dir = self
                .config
                .app_data_dir
                .clone()
                .unwrap_or_else(|| format!("{home}/.gemini/antigravity"));
            let mut allowed_paths = workspaces;
            allowed_paths.push(app_data_dir);
            let mut ws_policies = policy::workspace_only(allowed_paths);
            ws_policies.append(&mut final_policies);
            final_policies = ws_policies;
        }

        // Safety policy check: if write tools are enabled, policies cannot be empty
        if has_write_tools && final_policies.is_empty() {
            return Err(anyhow!(
                "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."
            ));
        }

        if !final_policies.is_empty() {
            let enforcer = Arc::new(PolicyEnforcer::new(final_policies));
            self.hook_runner.register(enforcer).await;
        }

        // 5. Register configured tools
        for tool in &self.config.tools {
            self.tool_runner.register(tool.clone()).await;
        }

        // 6. Build and connect strategy
        #[cfg(target_arch = "wasm32")]
        {
            let mut cap = self.config.capabilities.clone();
            if let Some(ref schema) = self.config.response_schema {
                cap.finish_tool_schema_json = Some(schema.clone());
            }

            let strategy = crate::wasm::WasmConnectionStrategy {
                gemini_config: self.config.gemini_config.clone(),
                capabilities_config: cap,
                system_instructions: self.config.system_instructions.clone(),
                save_dir: self.config.save_dir.clone(),
                workspaces: self.config.workspaces.clone().unwrap_or_default(),
                skills_paths: self.config.skills_paths.clone(),
                tool_runner: Some(self.tool_runner.clone()),
                hook_runner: Some(self.hook_runner.clone()),
                conversation_id: self.config.conversation_id.clone().unwrap_or_default(),
            };

            let conn = strategy.connect().await?;
            let conversation = Arc::new(Conversation::new(Arc::new(conn), None));
            self.conversation = Some(conversation.clone());

            // 7. Start triggers
            if !self.config.triggers.is_empty() {
                let runner = TriggerRunner::new(self.config.triggers.clone());
                runner.start(&conversation.connection());
                self.trigger_runner = Some(runner);
            }

            Ok(())
        }

        #[cfg(not(target_arch = "wasm32"))]
        {
            let mut cap = self.config.capabilities.clone();
            if let Some(ref schema) = self.config.response_schema {
                cap.finish_tool_schema_json = Some(schema.clone());
            }

            let strategy = LocalConnectionStrategy::new(
                binary_path,
                self.config.gemini_config.clone(),
                cap,
                self.config.system_instructions.clone(),
                self.config.save_dir.clone(),
                self.config.workspaces.clone().unwrap_or_default(),
                self.config.skills_paths.clone(),
                Some(self.tool_runner.clone()),
                Some(self.hook_runner.clone()),
                self.config.conversation_id.clone().unwrap_or_default(),
            );

            let conn = strategy.connect().await?;
            let conversation = Arc::new(Conversation::new(Arc::new(conn), None));
            self.conversation = Some(conversation.clone());

            // 7. Start triggers
            if !self.config.triggers.is_empty() {
                let runner = TriggerRunner::new(self.config.triggers.clone());
                runner.start(&conversation.connection());
                self.trigger_runner = Some(runner);
            }

            Ok(())
        }
    }

    /// Sends a prompt message to the active agent session and awaits the final completed response.
    ///
    /// # Errors
    ///
    /// Returns an error if the agent is not yet started or if the execution stream encounters a failure.
    pub async fn chat(&self, prompt: &str) -> Result<ChatResponse, anyhow::Error> {
        let conversation = self.conversation()?;
        conversation.chat_to_completion(prompt).await
    }

    /// Returns the active [`Conversation`] session.
    ///
    /// # Errors
    ///
    /// Returns an error if the agent is not yet started.
    pub fn conversation(&self) -> Result<Arc<Conversation>, anyhow::Error> {
        self.conversation
            .clone()
            .ok_or_else(|| anyhow!("Agent session not started. Use start() first."))
    }

    /// Returns the active conversation ID if the session has started.
    pub fn conversation_id(&self) -> Option<String> {
        self.conversation
            .as_ref()
            .map(|c| c.conversation_id().to_string())
    }

    /// Gracefully stops the agent connection and disconnects the underlying harness.
    ///
    /// # Errors
    ///
    /// Returns an error if closing the connection fails.
    pub async fn stop(&mut self) -> Result<(), anyhow::Error> {
        if let Some(conversation) = self.conversation.take() {
            conversation.disconnect().await?;
        }
        Ok(())
    }
}

#[cfg(not(target_arch = "wasm32"))]
fn get_default_binary_path() -> Option<String> {
    if let Ok(path) = std::env::var("ANTIGRAVITY_HARNESS_PATH") {
        return Some(path);
    }
    // Check if it is in standard PATH
    if let Ok(paths) = std::env::var("PATH") {
        for path in std::env::split_paths(&paths) {
            let p = path.join("localharness");
            if p.exists() {
                return Some(p.to_string_lossy().into_owned());
            }
        }
    }
    // Check Python site-packages as a fallback since google-antigravity Python package installs it there
    if let Some(output) = std::process::Command::new("python3")
        .args([
            "-c",
            "import site; print('\\n'.join(site.getsitepackages()))",
        ])
        .output()
        .ok()
        .filter(|o| o.status.success())
    {
        let stdout = String::from_utf8_lossy(&output.stdout);
        for line in stdout.lines() {
            let p = std::path::Path::new(line.trim())
                .join("google")
                .join("antigravity")
                .join("bin")
                .join("localharness");
            if p.exists() {
                return Some(p.to_string_lossy().into_owned());
            }
        }
    }
    None
}