agy-bridge 0.10.0

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

#[cfg(not(any(feature = "python", feature = "native")))]
compile_error!(
    "At least one backend feature must be enabled for `agy-bridge`. \
     Enable `python` (default) or `native` (pure-Rust standalone harness)."
);

// 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;
/// Protobuf definitions for the native harness backend.
#[cfg(feature = "native")]
pub mod proto;
/// 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, McpConfigError,
    McpConfigFile, McpServer, McpServerSpec, 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::{BackendLogLevel, RuntimeConfig};
pub use streaming::{ChatResponseHandle, ChatResult, ResponseEvent, StreamChunk};
pub use tools::{
    AvailableTool, RustTool, ToolContext, ToolDefinition, ToolError, ToolOutput, ToolRegistry,
    ToolSource,
};
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,
            McpConfigError, McpConfigFile, McpServer, McpServerSpec, McpSseServer, McpStdioServer,
            McpStreamableHttpServer, SystemInstructions,
        },
        content::{Audio, Content, ContentPrimitive, Document, Image, Video},
        error::Error,
        hooks::{HookPoint, HookResult, Hooks},
        policies::{AskUserHandler, PolicyDecision, PolicyRule, PolicySet},
        runtime::BackendLogLevel,
        streaming::{ChatResponseHandle, ChatResult, ResponseEvent, StreamChunk},
        tools::{
            AvailableTool, RustTool, ToolContext, ToolDefinition, ToolError, ToolOutput,
            ToolRegistry, ToolSource,
        },
        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.
/// // NOLINT: example code in documentation — `let _ =` demonstrates the return value exists
/// 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))
}

#[cfg(feature = "python")]
pub type DefaultRuntime = runtime::PythonRuntime;

#[cfg(all(not(feature = "python"), feature = "native"))]
pub type DefaultRuntime = runtime::NativeRuntime;

#[cfg(feature = "python")]
pub type PythonAgent = agent::AgentHandle<runtime::PythonRuntime>;

#[cfg(feature = "native")]
pub type NativeAgent = agent::AgentHandle<runtime::NativeRuntime>;

/// Convenience alias for an agent backed by the bridge's default runtime.
///
/// This hides the generic `Runtime` parameter so consumers can write `Agent`
/// with whichever runtime is default.
pub type Agent = agent::AgentHandle<DefaultRuntime>;

/// Primary entry point for the Antigravity bridge.
///
/// Wraps the runtime and provides a clean Rust API for creating agents.
pub struct AgyBridge<R: agent::Runtime + 'static = DefaultRuntime> {
    runtime: Arc<R>,
}

/// Builder for constructing an [`AgyBridge`] instance.
///
/// Created via [`AgyBridge::builder()`]. All settings have sensible defaults;
/// call [`.build()`](Self::build) to finalise.
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 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 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
    }

    /// Set the backend runtime log verbosity.
    ///
    /// Defaults to [`BackendLogLevel::Warn`]. Set to [`BackendLogLevel::Info`]
    /// or [`BackendLogLevel::Debug`] for verbose protocol-level diagnostics.
    #[must_use]
    pub fn backend_log_level(mut self, level: runtime::BackendLogLevel) -> Self {
        self.config.backend_log_level = level;
        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
    }

    /// Set the maximum number of consecutive model-quality errors before
    /// aborting a stream. Set to `0` to disable entirely (pure SDK pass-through).
    ///
    /// Defaults to `3` — enough to catch deterministically-bad model output
    /// without cutting off transient hiccups.
    #[must_use]
    pub fn max_consecutive_model_errors(mut self, limit: u32) -> Self {
        self.config.max_consecutive_model_errors = Some(limit);
        self
    }

    /// Set the maximum number of consecutive thinking-only/empty steps before
    /// aborting a stream. Set to `0` to disable entirely (pure SDK pass-through).
    ///
    /// Defaults to `500` — a generous ceiling to avoid false positives on
    /// legitimate long chains of thought.
    #[must_use]
    pub fn max_consecutive_empty_steps(mut self, limit: u32) -> Self {
        self.config.max_consecutive_empty_steps = Some(limit);
        self
    }

    /// Set the per-channel buffer size for streaming response channels.
    ///
    /// Each chat call creates ~7 channels of this size. Defaults to `256` —
    /// large enough to avoid backpressure under normal workloads.
    #[must_use]
    pub fn streaming_channel_buffer(mut self, size: usize) -> Self {
        self.config.streaming_channel_buffer = Some(size);
        self
    }

    /// Set an explicit path to the `localharness` binary (native backend).
    #[must_use]
    pub fn harness_path(mut self, path: impl Into<std::path::PathBuf>) -> Self {
        self.config.harness_binary_path = Some(path.into());
        self
    }

    #[cfg(feature = "python")]
    /// Build the [`AgyBridge`] with the Python runtime.
    ///
    /// # Errors
    ///
    /// Returns an error if the Python runtime fails to initialize.
    pub fn build_python(self) -> Result<AgyBridge<runtime::PythonRuntime>, error::Error> {
        Ok(AgyBridge {
            runtime: Arc::new(runtime::PythonRuntime::new(self.config)?),
        })
    }

    #[cfg(feature = "native")]
    /// Build the [`AgyBridge`] with the native local harness runtime.
    ///
    /// # Errors
    ///
    /// Returns an error if runtime initialization fails.
    pub fn build_native(self) -> Result<AgyBridge<runtime::NativeRuntime>, error::Error> {
        Ok(AgyBridge {
            runtime: Arc::new(runtime::NativeRuntime::new(self.config)),
        })
    }

    /// Build the [`AgyBridge`], starting the default runtime.
    ///
    /// # Errors
    ///
    /// Returns an error if the underlying runtime fails to initialize.
    pub fn build(self) -> Result<AgyBridge<DefaultRuntime>, error::Error> {
        #[cfg(feature = "python")]
        {
            self.build_python()
        }
        #[cfg(all(not(feature = "python"), feature = "native"))]
        {
            self.build_native()
        }
    }
}

impl<R: agent::Runtime + 'static> AgyBridge<R> {
    /// Create a new bridge instance from an existing runtime Arc.
    #[must_use]
    pub fn new(runtime: Arc<R>) -> Self {
        Self { runtime }
    }

    /// Begin building a new agent on this bridge.
    #[must_use]
    pub fn agent(&self, config: config::AgentConfig) -> AgentBuilder<'_, R> {
        AgentBuilder {
            bridge: self,
            config,
            registry: None,
            hooks: None,
            policy_handler: None,
        }
    }

    /// Convenience shorthand for `self.agent(AgentConfig::default())`.
    #[must_use]
    pub fn default_agent(&self) -> AgentBuilder<'_, R> {
        self.agent(config::AgentConfig::default())
    }

    /// Return a reference to the underlying runtime Arc.
    #[must_use]
    pub fn runtime(&self) -> &Arc<R> {
        &self.runtime
    }
}

#[cfg(feature = "python")]
impl AgyBridge<runtime::PythonRuntime> {
    /// Return the number of agents currently live on this bridge.
    ///
    /// # Errors
    ///
    /// Returns an error if querying the active agent count fails.
    pub async fn active_agent_count(&self) -> Result<usize, error::Error> {
        self.runtime.active_agent_count().await
    }
}

impl AgyBridge<DefaultRuntime> {
    /// Create a new builder for configuring and constructing an [`AgyBridge`].
    #[must_use]
    pub fn builder() -> AgyBridgeBuilder {
        AgyBridgeBuilder {
            config: runtime::RuntimeConfig::default(),
        }
    }
}

#[cfg(feature = "native")]
impl AgyBridge<runtime::NativeRuntime> {
    /// Create a new builder for configuring and constructing a native [`AgyBridge`].
    #[must_use]
    pub fn native_builder() -> AgyBridgeBuilder {
        AgyBridgeBuilder {
            config: runtime::RuntimeConfig::default(),
        }
    }

    /// Return the number of agents currently live on this native bridge.
    ///
    /// # Errors
    ///
    /// Returns an error if querying the active agent count fails.
    pub async fn active_agent_count(&self) -> Result<usize, error::Error> {
        self.runtime.active_agent_count().await
    }
}

/// Builder for creating an [`Agent`] on an [`AgyBridge`].
pub struct AgentBuilder<'a, R: agent::Runtime + 'static = DefaultRuntime> {
    bridge: &'a AgyBridge<R>,
    config: config::AgentConfig,
    registry: Option<tools::ToolRegistry>,
    hooks: Option<hooks::Hooks>,
    policy_handler: Option<Arc<dyn policies::AskUserHandler>>,
}

impl<R: agent::Runtime + 'static> AgentBuilder<'_, R> {
    /// Attach a [`ToolRegistry`] containing custom
    /// Rust tools for the agent.
    #[must_use]
    pub fn tools(mut self, registry: tools::ToolRegistry) -> Self {
        self.registry = Some(registry);
        self
    }

    /// Attach [`Hooks`] for lifecycle event callbacks.
    #[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
    }

    /// Resume an existing conversation by its SDK id.
    #[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 the API key for backend requests.
    #[must_use]
    pub fn api_key(mut self, key: impl Into<String>) -> Self {
        self.config.api_key = Some(key.into());
        self
    }

    /// Replace the 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
    }

    /// Set the capabilities configuration for the agent.
    #[must_use]
    pub fn capabilities(mut self, capabilities: config::CapabilitiesConfig) -> Self {
        self.config.capabilities = Some(capabilities);
        self
    }

    /// Append workspace directories to the agent's permitted filesystem scope.
    #[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-call 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
    }

    /// Validate configuration and create the agent.
    ///
    /// # Errors
    ///
    /// Returns [`Error::InvalidConfig`] if configuration
    /// validation fails, or a backend error if agent creation fails.
    pub async fn build(mut self) -> Result<agent::AgentHandle<R>, error::Error> {
        // Validate capabilities.
        if let Some(ref caps) = self.config.capabilities {
            caps.validate().map_err(|msg| error::Error::InvalidConfig {
                message: msg.to_string(),
            })?;
        }

        // Validate response_schema.
        if let Some(ref schema) = self.config.response_schema {
            schema
                .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, R: agent::Runtime + 'static> std::future::IntoFuture for AgentBuilder<'a, R> {
    type Output = Result<agent::AgentHandle<R>, 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));
    }
}