agy_bridge/

lib.rs

#![doc = include_str!("../README.md")]
//! agy-bridge: Standalone reusable `PyO3` bridge for the Google Antigravity SDK.

// Allow `::agy_bridge::` paths (generated by the `#[llm_tool]` proc macro) to
// resolve when compiling tests within this crate.
extern crate self as agy_bridge;

/// Agent lifecycle management: creation, chat, shutdown.
pub mod agent;
/// Configuration types for agents, models, capabilities, and MCP servers.
pub mod config;

/// Multimodal content types for chat input (text, image, document, audio, video).
pub mod content;
/// Error types for the bridge.
pub mod error;
/// Pre/post-turn and tool-call lifecycle hooks.
pub mod hooks;
/// Policy rules for tool-call filtering and workspace scoping.
pub mod policies;
/// Quota tracking and backoff state.
pub mod quota;
/// Python runtime bridge: command dispatch over a dedicated thread.
pub mod runtime;

/// Streaming response channels for text, thought, and tool-call events.
pub mod streaming;
/// Custom Rust tool dispatch and definition types.
pub mod tools;
/// Event-driven trigger definitions.
pub mod triggers;
/// Shared domain types (messages, steps, usage metadata).
pub mod types;

// ── Re-exports ──────────────────────────────────────────────────────────────
// Flat re-exports of the most commonly used types so callers can write
// `use agy_bridge::{AgentConfig, Error, ToolRegistry, Content};`
// without diving into sub-modules.

pub use config::{
    AgentConfig, BuiltinTools, CapabilitiesConfig, GeminiConfig, LocalAgentConfig, McpServer,
    McpSseServer, McpStdioServer, McpStreamableHttpServer, SystemInstructions,
};
pub use content::{Audio, Content, ContentPrimitive, Document, Image, Video};
pub use error::Error;
pub use hooks::{HookCallback, HookEntry, HookPoint, HookResult, HookSet, Hooks};
/// Re-export the `#[llm_tool]` proc-macro so users only need `agy_bridge` in
/// their dependency list.
pub use llm_tool_macros::llm_tool;
pub use policies::{AskUserHandler, PolicyDecision, PolicyRule, PolicySet};
pub use runtime::RuntimeConfig;
pub use streaming::{ChatResponseHandle, ChatResult, ResponseEvent, StreamChunk};
pub use tools::{RustTool, ToolContext, ToolDefinition, ToolError, ToolOutput, ToolRegistry};
pub use triggers::{TriggerConfig, TriggerEntry};
pub use types::{ConversationMessage, MessageRole, Step, UsageMetadata};

/// Convenience prelude — pull in everything you need with a single glob import.
///
/// ```
/// use agy_bridge::prelude::*;
/// ```
///
/// This re-exports the most commonly used types, traits, and macros from the
/// crate so you can get started quickly without hunting for individual paths.
pub mod prelude {
    pub use llm_tool_macros::llm_tool;

    pub use crate::{
        Agent, AgyBridge,
        config::{
            AgentConfig, BuiltinTools, CapabilitiesConfig, GeminiConfig, LocalAgentConfig,
            McpServer, McpSseServer, McpStdioServer, McpStreamableHttpServer, SystemInstructions,
        },
        content::{Audio, Content, ContentPrimitive, Document, Image, Video},
        error::Error,
        hooks::{HookPoint, HookResult, Hooks},
        policies::{AskUserHandler, PolicyDecision, PolicyRule, PolicySet},
        streaming::{ChatResponseHandle, ChatResult, ResponseEvent, StreamChunk},
        tools::{RustTool, ToolContext, ToolDefinition, ToolError, ToolOutput, ToolRegistry},
        triggers::{TriggerConfig, TriggerEntry},
        types::{ConversationMessage, MessageRole, Step, UsageMetadata},
    };
}

use std::sync::Arc;

/// Load environment variables from a `.env` file into the process environment.
///
/// 1. Walks upward from `CARGO_MANIFEST_DIR` (if set) or the current working
///    directory to find the nearest `.env` file.
/// 2. Parses each `KEY=VALUE` line (skipping blanks and `#`-comments).
/// 3. For every key that is **not** already present in the process
///    environment, calls [`std::env::set_var`] to inject it.
/// 4. Returns a [`HashMap`](std::collections::HashMap) of the newly-set
///    key/value pairs (keys that were already set are omitted).
///
/// Results are cached via [`OnceLock`](std::sync::OnceLock) — the file is
/// read and environment variables are set at most once. Subsequent calls
/// return a clone of the cached map without re-reading the file or
/// modifying the environment.
///
/// # Safety
///
/// This function calls [`std::env::set_var`], which is **not** thread-safe.
/// It **must** be called during single-threaded startup, before any
/// additional threads are spawned (including the Tokio runtime). Calling it
/// after threads exist is undefined behaviour.
///
/// # Example
///
/// ```
/// let new_vars = agy_bridge::load_dotenv();
/// // OnceLock-cached: safe to call multiple times, only loads .env once.
/// let _ = new_vars.len();
/// ```
pub fn load_dotenv() -> &'static std::collections::HashMap<String, String> {
    use std::sync::OnceLock;

    static CACHED: OnceLock<std::collections::HashMap<String, String>> = OnceLock::new();

    CACHED.get_or_init(|| {
        let start = std::env::var_os("CARGO_MANIFEST_DIR").map_or_else(
            || {
                std::env::current_dir().unwrap_or_else(|e| {
                    tracing::debug!("load_dotenv: current_dir() failed: {e}, using fallback \".\"");
                    std::path::PathBuf::from(".")
                })
            },
            std::path::PathBuf::from,
        );

        let mut dir = start.as_path();
        loop {
            let candidate = dir.join(".env");
            if candidate.is_file() {
                let mut env_map = std::collections::HashMap::new();
                match std::fs::read_to_string(&candidate) {
                    Ok(contents) => {
                        for line in contents.lines() {
                            if let Some((k, v)) = parse_dotenv_line(line)
                                && std::env::var_os(k).is_none()
                            {
                                // SAFETY: Called inside the OnceLock closure during
                                // single-threaded initialization, before any threads
                                // are spawned. set_var is not thread-safe, but here
                                // we are the only thread.
                                unsafe {
                                    std::env::set_var(k, v);
                                }
                                env_map.insert(k.to_owned(), v.to_owned());
                            }
                        }
                    }
                    Err(e) => {
                        tracing::warn!(error = %e, "Failed to read .env file at {}", candidate.display());
                    }
                }
                return env_map;
            }
            match dir.parent() {
                Some(parent) => dir = parent,
                None => return std::collections::HashMap::new(),
            }
        }
    })
}

/// Parse a single line from a `.env` file.
///
/// Returns `Some((key, value))` for valid `KEY=VALUE` lines, stripping
/// surrounding whitespace and quotes (single or double) from the value.
/// Returns `None` for blank lines, comments, or lines without `=`.
///
/// This is factored out of [`load_dotenv`] for testability.
pub(crate) fn parse_dotenv_line(line: &str) -> Option<(&str, &str)> {
    let line = line.trim();
    if line.is_empty() || line.starts_with('#') {
        return None;
    }
    let (k, v) = line.split_once('=')?;
    let k = k.trim();
    if k.is_empty() {
        return None;
    }
    let v = v.trim();
    // Strip surrounding quotes (single or double)
    let v = v
        .strip_prefix('"')
        .and_then(|s| s.strip_suffix('"'))
        .or_else(|| v.strip_prefix('\'').and_then(|s| s.strip_suffix('\'')))
        .unwrap_or(v);
    Some((k, v))
}

/// Convenience alias for an agent backed by the bridge's runtime.
///
/// This hides the generic `Runtime` parameter so consumers never see the
/// underlying Python bridge type.
pub type Agent = agent::AgentHandle<runtime::PythonRuntime>;

/// Primary entry point for the Antigravity bridge.
///
/// Wraps the runtime and provides a clean Rust API for creating agents.
///
/// # Example
///
/// ```rust
/// # use agy_bridge::AgyBridge;
/// # use agy_bridge::config::AgentConfig;
/// # #[tokio::main]
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// # agy_bridge::load_dotenv();
/// // Zero-config:
/// // let bridge = AgyBridge::builder().build()?;
///
/// // With custom timeouts:
/// let bridge = AgyBridge::builder()
///     .chat_timeout(std::time::Duration::from_secs(120))
///     .build()?;
///
/// // Create an agent (simple):
/// // let agent = bridge.agent(AgentConfig::default()).await?;
/// # let manifest_dir = std::path::PathBuf::from(std::env::var("CARGO_MANIFEST_DIR")?);
/// # let project_root = manifest_dir.parent().unwrap().parent().unwrap();
/// # let agent = bridge.agent(
/// #     AgentConfig::builder()
/// #         .system_instructions("Reply with 'Hello!' and nothing else. Never use tools.")
/// #         .capabilities(agy_bridge::config::CapabilitiesConfig::custom_tools_only())
/// #         .workspaces(vec![project_root])
/// #         .build()
/// # ).await?;
///
/// // Create an agent with tools and hooks:
/// // let agent = bridge.agent(config)
/// //     .tools(registry)
/// //     .hooks(hooks)
/// //     .await?;
///
/// let answer = agent.chat("Hello!").await?.text().await?;
/// # Ok(())
/// # }
/// ```
pub struct AgyBridge {
    runtime: Arc<runtime::PythonRuntime>,
}

/// Builder for constructing an [`AgyBridge`] instance.
///
/// Created via [`AgyBridge::builder()`]. All settings have sensible defaults;
/// call [`.build()`](Self::build) to finalise.
///
/// # Example
///
/// ```
/// # use agy_bridge::AgyBridge;
/// let bridge = AgyBridge::builder()
///     .chat_timeout(std::time::Duration::from_secs(120))
///     .channel_capacity(128)
///     .build()?;
/// # Ok::<(), agy_bridge::error::Error>(())
/// ```
pub struct AgyBridgeBuilder {
    config: runtime::RuntimeConfig,
}

impl AgyBridgeBuilder {
    /// Set the mpsc channel buffer size for the command channel.
    #[must_use]
    pub fn channel_capacity(mut self, capacity: usize) -> Self {
        self.config.channel_capacity = capacity;
        self
    }

    /// Set the timeout for individual runtime operations.
    #[must_use]
    pub fn operation_timeout(mut self, timeout: std::time::Duration) -> Self {
        self.config.operation_timeout = timeout;
        self
    }

    /// Set the timeout for joining the Python thread on shutdown.
    #[must_use]
    pub fn shutdown_timeout(mut self, timeout: std::time::Duration) -> Self {
        self.config.shutdown_timeout = timeout;
        self
    }

    /// Set the timeout for a single `agent.chat()` round-trip.
    #[must_use]
    pub fn chat_timeout(mut self, timeout: std::time::Duration) -> Self {
        self.config.chat_timeout = timeout;
        self
    }

    /// Set the delay between successive chat commands to prevent burst requests.
    #[must_use]
    pub fn inter_agent_delay(mut self, delay: std::time::Duration) -> Self {
        self.config.inter_agent_delay = delay;
        self
    }

    /// Replace the entire runtime configuration at once.
    ///
    /// Useful when you already have a [`RuntimeConfig`] struct. Individual
    /// setters called *after* this will override the corresponding fields.
    #[must_use]
    pub fn runtime_config(mut self, config: runtime::RuntimeConfig) -> Self {
        self.config = config;
        self
    }

    /// Build the [`AgyBridge`], starting the Python runtime.
    ///
    /// # Errors
    ///
    /// Returns [`error::Error`] if the Python runtime cannot be started
    /// (e.g. missing Antigravity SDK installation).
    pub fn build(self) -> Result<AgyBridge, error::Error> {
        Ok(AgyBridge {
            runtime: Arc::new(runtime::PythonRuntime::new(self.config)?),
        })
    }
}

impl AgyBridge {
    /// Create a new builder for configuring and constructing an [`AgyBridge`].
    ///
    /// # Example
    ///
    /// ```
    /// # use agy_bridge::AgyBridge;
    /// let bridge = AgyBridge::builder().build()?;
    /// # Ok::<(), agy_bridge::error::Error>(())
    /// ```
    #[must_use]
    pub fn builder() -> AgyBridgeBuilder {
        AgyBridgeBuilder {
            config: runtime::RuntimeConfig::default(),
        }
    }

    /// Begin building a new agent on this bridge.
    ///
    /// Returns an [`AgentBuilder`] that can be directly `.await`ed for the
    /// simple case, or chained with [`.tools()`](AgentBuilder::tools) and
    /// [`.hooks()`](AgentBuilder::hooks) before awaiting.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use agy_bridge::{AgyBridge, config::AgentConfig};
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), agy_bridge::error::Error> {
    /// # agy_bridge::load_dotenv();
    /// # let bridge = AgyBridge::builder().build()?;
    /// // Simple — no tools or hooks:
    /// let agent = bridge.agent(AgentConfig::default()).await?;
    ///
    /// // With tools:
    /// // let agent = bridge.agent(config).tools(registry).await?;
    ///
    /// // With tools and hooks:
    /// // let agent = bridge.agent(config).tools(registry).hooks(hooks).await?;
    /// # Ok(())
    /// # }
    /// ```
    #[must_use]
    pub fn agent(&self, config: config::AgentConfig) -> AgentBuilder<'_> {
        AgentBuilder {
            bridge: self,
            config,
            registry: None,
            hooks: None,
            policy_handler: None,
        }
    }

    /// Convenience shorthand for `self.agent(AgentConfig::default())`.
    ///
    /// Creates an agent builder with default configuration. Chain
    /// [`.tools()`](AgentBuilder::tools) or [`.hooks()`](AgentBuilder::hooks)
    /// before awaiting, or `.await` directly for a bare agent.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use agy_bridge::AgyBridge;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), agy_bridge::error::Error> {
    /// # agy_bridge::load_dotenv();
    /// # let bridge = AgyBridge::builder().build()?;
    /// let agent = bridge.default_agent().await?;
    /// # Ok(())
    /// # }
    /// ```
    #[must_use]
    pub fn default_agent(&self) -> AgentBuilder<'_> {
        self.agent(config::AgentConfig::default())
    }
}

/// Builder for creating an [`Agent`] on an [`AgyBridge`].
///
/// Obtained from [`AgyBridge::agent()`]. Implements [`IntoFuture`] so you can
/// `.await` it directly, or chain optional [`.tools()`](Self::tools) /
/// [`.hooks()`](Self::hooks) calls before awaiting.
pub struct AgentBuilder<'a> {
    bridge: &'a AgyBridge,
    config: config::AgentConfig,
    registry: Option<tools::ToolRegistry>,
    hooks: Option<hooks::Hooks>,
    policy_handler: Option<Arc<dyn policies::AskUserHandler>>,
}

impl AgentBuilder<'_> {
    /// Attach a [`ToolRegistry`] containing custom
    /// Rust tools for the agent.
    ///
    /// The registry's tool definitions are automatically merged into the
    /// agent configuration.
    ///
    /// # Errors (at build time)
    ///
    /// Returns [`error::Error::InvalidConfig`] if `config.tools` is
    /// already non-empty — pass tools via the registry **or** via
    /// `config.tools`, not both.
    #[must_use]
    pub fn tools(mut self, registry: tools::ToolRegistry) -> Self {
        self.registry = Some(registry);
        self
    }

    /// Attach [`Hooks`] for lifecycle event
    /// callbacks (pre/post turn, tool-call gating, etc.).
    #[must_use]
    pub fn hooks(mut self, hooks: hooks::Hooks) -> Self {
        self.hooks = Some(hooks);
        self
    }

    /// Attach a custom [`AskUserHandler`] to manage interactive tool-call confirmations.
    #[must_use]
    pub fn policy_handler(mut self, handler: impl policies::AskUserHandler + 'static) -> Self {
        self.policy_handler = Some(Arc::new(handler));
        self
    }

    /// Set a pre-existing conversation ID to resume.
    #[must_use]
    pub fn conversation_id(mut self, id: impl Into<String>) -> Self {
        self.config.conversation_id = Some(id.into());
        self
    }

    /// Set the model backend (e.g. `"gemini-3.5-flash"`).
    #[must_use]
    pub fn model(mut self, model: impl Into<String>) -> Self {
        self.config.model = model.into();
        self
    }

    /// Set system instructions for the agent.
    #[must_use]
    pub fn system_instructions(
        mut self,
        instructions: impl Into<config::SystemInstructions>,
    ) -> Self {
        self.config.system_instructions = Some(instructions.into());
        self
    }

    /// Append workspace directories the agent is allowed to access and modify.
    #[must_use]
    pub fn workspaces(
        mut self,
        workspaces: impl IntoIterator<Item = impl Into<std::path::PathBuf>>,
    ) -> Self {
        self.config
            .workspaces
            .extend(workspaces.into_iter().map(Into::into));
        self
    }

    /// Append policy rules to govern tool execution.
    #[must_use]
    pub fn policies(
        mut self,
        policies: impl IntoIterator<Item = impl Into<policies::PolicyRule>>,
    ) -> Self {
        self.config
            .policies
            .extend(policies.into_iter().map(Into::into));
        self
    }

    /// Append triggers that autonomously wake the agent.
    #[must_use]
    pub fn triggers(
        mut self,
        triggers: impl IntoIterator<Item = impl Into<triggers::TriggerEntry>>,
    ) -> Self {
        self.config
            .triggers
            .extend(triggers.into_iter().map(Into::into));
        self
    }

    /// Append MCP servers for the agent.
    #[must_use]
    pub fn mcp_servers(
        mut self,
        servers: impl IntoIterator<Item = impl Into<config::McpServer>>,
    ) -> Self {
        self.config
            .mcp_servers
            .extend(servers.into_iter().map(Into::into));
        self
    }

    /// Append paths for agent skills.
    #[must_use]
    pub fn skills(
        mut self,
        skills: impl IntoIterator<Item = impl Into<std::path::PathBuf>>,
    ) -> Self {
        self.config
            .skills
            .extend(skills.into_iter().map(Into::into));
        self
    }

    /// Set the maximum number of quota retry attempts before giving up.
    #[must_use]
    pub fn max_quota_retries(mut self, retries: u32) -> Self {
        self.config.max_quota_retries = Some(retries);
        self
    }

    /// Validate configuration and create the agent.
    ///
    /// Prefer using `.await` directly on the builder (via [`IntoFuture`])
    /// instead of calling this method explicitly.
    ///
    /// # Errors
    ///
    /// Returns [`error::Error::InvalidConfig`] if:
    /// - `config.tools` is non-empty **and** a `ToolRegistry` was provided.
    /// - The capabilities configuration is self-contradictory (e.g. both
    ///   `enabled_tools` and `disabled_tools` specified).
    ///
    /// Returns other [`error::Error`] variants if agent creation fails.
    pub async fn build(mut self) -> Result<Agent, error::Error> {
        // Validate capabilities.
        if let Some(ref caps) = self.config.capabilities {
            caps.validate().map_err(|msg| error::Error::InvalidConfig {
                message: msg.to_string(),
            })?;
        }

        // Handle tool registry.
        let arc_registry = if let Some(registry) = self.registry {
            if !self.config.tools.is_empty() {
                return Err(error::Error::InvalidConfig {
                    message: "config.tools is non-empty and a ToolRegistry was also provided; \
                              pass tools via the registry or via config.tools, not both"
                        .to_string(),
                });
            }
            self.config.tools = registry.definitions();
            Some(Arc::new(registry))
        } else {
            None
        };

        // Handle hooks.
        let arc_hooks = if let Some(hooks) = self.hooks {
            if !self.config.hooks.is_empty() {
                return Err(error::Error::InvalidConfig {
                    message: "config.hooks is non-empty and a Hooks instance was also provided; \
                              configure hooks via Hooks or config.hooks, not both"
                        .to_string(),
                });
            }
            self.config.hooks = hooks.entries();
            Some(Arc::new(hooks))
        } else {
            None
        };

        // Handle policy handler.
        let arc_policy = self.policy_handler;

        agent::AgentHandle::new(
            Arc::clone(&self.bridge.runtime),
            self.config,
            arc_registry,
            arc_hooks,
            arc_policy,
        )
        .await
    }
}

impl<'a> std::future::IntoFuture for AgentBuilder<'a> {
    type Output = Result<Agent, error::Error>;
    type IntoFuture =
        std::pin::Pin<Box<dyn std::future::Future<Output = Self::Output> + Send + 'a>>;

    fn into_future(self) -> Self::IntoFuture {
        Box::pin(self.build())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── parse_dotenv_line regression tests ────────────────────────────

    #[test]
    fn dotenv_strips_double_quotes() {
        let (k, v) = parse_dotenv_line(r#"API_KEY="my-secret""#).unwrap();
        assert_eq!(k, "API_KEY");
        assert_eq!(v, "my-secret");
    }

    #[test]
    fn dotenv_strips_single_quotes() {
        let (k, v) = parse_dotenv_line("TOKEN='abc123'").unwrap();
        assert_eq!(k, "TOKEN");
        assert_eq!(v, "abc123");
    }

    #[test]
    fn dotenv_unquoted_value_unchanged() {
        let (k, v) = parse_dotenv_line("FOO=bar").unwrap();
        assert_eq!(k, "FOO");
        assert_eq!(v, "bar");
    }

    #[test]
    fn dotenv_mismatched_quotes_preserved() {
        // Opening double quote but closing single quote → not stripped.
        let (k, v) = parse_dotenv_line(r#"KEY="value'"#).unwrap();
        assert_eq!(k, "KEY");
        assert_eq!(v, r#""value'"#);
    }

    #[test]
    fn dotenv_empty_quoted_value() {
        let (k, v) = parse_dotenv_line(r#"EMPTY="""#).unwrap();
        assert_eq!(k, "EMPTY");
        assert_eq!(v, "");
    }

    #[test]
    fn dotenv_whitespace_around_key_value() {
        let (k, v) = parse_dotenv_line("  MY_VAR  =  \"hello world\"  ").unwrap();
        assert_eq!(k, "MY_VAR");
        assert_eq!(v, "hello world");
    }

    #[test]
    fn dotenv_comment_line_is_none() {
        assert!(parse_dotenv_line("# this is a comment").is_none());
    }

    #[test]
    fn dotenv_blank_line_is_none() {
        assert!(parse_dotenv_line("   ").is_none());
    }

    #[test]
    fn dotenv_empty_key_is_none() {
        assert!(parse_dotenv_line("=value").is_none());
    }

    #[test]
    fn dotenv_no_equals_is_none() {
        assert!(parse_dotenv_line("JUSTKEY").is_none());
    }

    #[test]
    fn dotenv_value_with_internal_equals() {
        let (k, v) = parse_dotenv_line("DSN=postgres://host:5432/db?opt=1").unwrap();
        assert_eq!(k, "DSN");
        assert_eq!(v, "postgres://host:5432/db?opt=1");
    }

    #[test]
    fn dotenv_value_with_embedded_quotes_not_stripped() {
        // Quotes in the middle are not stripped — only surrounding ones.
        let (k, v) = parse_dotenv_line(r#"MSG=say "hello""#).unwrap();
        assert_eq!(k, "MSG");
        assert_eq!(v, r#"say "hello""#);
    }

    #[test]
    fn test_load_dotenv_returns_static_reference_identity() {
        let map1 = load_dotenv();
        let map2 = load_dotenv();
        assert!(std::ptr::eq(map1, map2));
    }
}