From 1bc5c7f09947434a4c58df6beee3911ea9a9a056 Mon Sep 17 00:00:00 2001 From: codeitlikemiley Date: Fri, 29 May 2026 19:53:00 +0800 Subject: [PATCH 1/6] feat: add core types for Python SDK parity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add AntigravityError enum (Connection, Execution, Validation) in error.rs - Add HookContext hierarchical state store (Session→Turn→Operation) in context.rs - Add ToolContext with session-scoped state and agent communication in tool_context.rs - Add multimodal content pipeline: Media, Content, MimeType (32 extensions) - Add trigger types: TriggerDelivery, FileChange, FileChangeKind --- src/context.rs | 161 ++++++++++++++ src/error.rs | 35 +++ src/tool_context.rs | 112 ++++++++++ src/types.rs | 531 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 839 insertions(+) create mode 100644 src/context.rs create mode 100644 src/error.rs create mode 100644 src/tool_context.rs diff --git a/src/context.rs b/src/context.rs new file mode 100644 index 0000000..cd32449 --- /dev/null +++ b/src/context.rs @@ -0,0 +1,161 @@ +//! Hierarchical context system for hook state management. +//! +//! Provides a parent-chaining key-value store where `get()` walks the parent chain +//! and `set()` writes only to the local store. This enables state sharing across +//! hook lifecycle events (session → turn → operation scope). + +use serde::{Serialize, de::DeserializeOwned}; +use serde_json::Value; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +/// A hierarchical key-value store for sharing state across hook invocations. +/// +/// Contexts form a chain: `OperationContext` → `TurnContext` → `SessionContext`. +/// - `get()` searches the local store first, then walks up the parent chain. +/// - `set()` writes only to the local store (shadowing, not mutating parents). +#[derive(Debug, Clone)] +pub struct HookContext { + parent: Option>, + store: Arc>>, +} + +impl HookContext { + /// Creates a root context (no parent). Used as the session-level context. + pub fn new() -> Self { + Self { + parent: None, + store: Arc::new(Mutex::new(HashMap::new())), + } + } + + /// Creates a child context with the given parent. + /// Used to create turn-level (parent=session) or operation-level (parent=turn) contexts. + pub fn child(parent: Arc) -> Self { + Self { + parent: Some(parent), + store: Arc::new(Mutex::new(HashMap::new())), + } + } + + /// Retrieves a value by key, walking up the parent chain if not found locally. + /// Returns `None` if the key is not found in any context in the hierarchy. + #[allow(clippy::collapsible_if)] + pub fn get(&self, key: &str) -> Option { + // Check local store first + if let Ok(store) = self.store.lock() { + if let Some(value) = store.get(key) { + return serde_json::from_value(value.clone()).ok(); + } + } + // Walk up parent chain + self.parent.as_ref().and_then(|p| p.get(key)) + } + + /// Sets a value in the **local** store only (does not write to parents). + /// If the key already exists locally, it is overwritten. + #[allow(clippy::collapsible_if)] + pub fn set(&self, key: &str, value: T) { + if let Ok(mut store) = self.store.lock() { + if let Ok(v) = serde_json::to_value(value) { + store.insert(key.to_string(), v); + } + } + } + + /// Returns `true` if this context has a parent (i.e., is not a root/session context). + pub const fn has_parent(&self) -> bool { + self.parent.is_some() + } +} + +impl Default for HookContext { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + use super::*; + + #[test] + fn test_root_context_get_set() { + let ctx = HookContext::new(); + ctx.set("key1", "value1"); + assert_eq!(ctx.get::("key1"), Some("value1".to_string())); + assert_eq!(ctx.get::("missing"), None); + } + + #[test] + fn test_child_inherits_parent() { + let session = Arc::new(HookContext::new()); + session.set("session_key", "session_value"); + + let turn = HookContext::child(session); + // Child can read parent values + assert_eq!( + turn.get::("session_key"), + Some("session_value".to_string()) + ); + } + + #[test] + fn test_child_shadows_parent() { + let session = Arc::new(HookContext::new()); + session.set("key", "parent_value"); + + let turn = HookContext::child(session.clone()); + turn.set("key", "child_value"); + + // Child sees its own value + assert_eq!(turn.get::("key"), Some("child_value".to_string())); + // Parent is unchanged + assert_eq!( + session.get::("key"), + Some("parent_value".to_string()) + ); + } + + #[test] + fn test_three_level_hierarchy() { + let session = Arc::new(HookContext::new()); + session.set("level", "session"); + session.set("session_only", true); + + let turn = Arc::new(HookContext::child(session)); + turn.set("level", "turn"); + turn.set("turn_only", 42i32); + + let operation = HookContext::child(turn); + operation.set("level", "operation"); + + // Operation sees its own, then turn, then session + assert_eq!( + operation.get::("level"), + Some("operation".to_string()) + ); + assert_eq!(operation.get::("turn_only"), Some(42)); + assert_eq!(operation.get::("session_only"), Some(true)); + } + + #[test] + fn test_no_parent_for_root() { + let ctx = HookContext::new(); + assert!(!ctx.has_parent()); + } + + #[test] + fn test_child_has_parent() { + let parent = Arc::new(HookContext::new()); + let child = HookContext::child(parent); + assert!(child.has_parent()); + } + + #[test] + fn test_default_is_root() { + let ctx = HookContext::default(); + assert!(!ctx.has_parent()); + } +} diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..921f473 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,35 @@ +//! Typed error hierarchy for the Antigravity SDK. + +use serde::{Deserialize, Serialize}; + +/// Structured detail for a single validation failure. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ValidationDetail { + /// Location path segments indicating where the error occurred. + pub loc: Vec, + /// Human-readable error message. + pub msg: String, + /// Machine-readable error type tag. + pub error_type: String, +} + +/// Unified error type for the Antigravity SDK. +#[derive(Debug, Clone, thiserror::Error)] +pub enum AntigravityError { + /// A network or transport-level failure. + #[error("Connection error: {0}")] + Connection(String), + + /// A server-side or agent execution failure. + #[error("Execution error: {0}")] + Execution(String), + + /// One or more input validation failures. + #[error("Validation error: {message}")] + Validation { + /// Summary message for the validation error. + message: String, + /// Individual validation failures. + errors: Vec, + }, +} diff --git a/src/tool_context.rs b/src/tool_context.rs new file mode 100644 index 0000000..78f5bd1 --- /dev/null +++ b/src/tool_context.rs @@ -0,0 +1,112 @@ +//! Conversation-aware context injected into tools that request it. +//! +//! `ToolContext` provides access to session state, conversation metadata, +//! and the ability to send messages to the agent. State is scoped to the +//! session and is independent of `HookContext`. + +use crate::connection::AnyConnection; +use crate::connection::Connection; +use anyhow::Result; +use serde::{Serialize, de::DeserializeOwned}; +use serde_json::Value; +use std::collections::HashMap; +use std::sync::Mutex; + +/// Session-scoped context injected into tools that need conversation awareness. +/// +/// Provides: +/// - `conversation_id()` — current session identifier +/// - `is_idle()` — whether the agent is idle +/// - `send()` — push a trigger notification to the agent +/// - `get_state()` / `set_state()` — per-session key-value store +/// +/// State set by tools is **not** visible to hooks, and vice versa. +/// This separation is intentional (see hooks/README.md in the Python SDK). +#[derive(Debug)] +pub struct ToolContext { + connection: AnyConnection, + state: Mutex>, +} + +impl ToolContext { + /// Creates a new `ToolContext` wrapping the given connection. + pub fn new(connection: AnyConnection) -> Self { + Self { + connection, + state: Mutex::new(HashMap::new()), + } + } + + /// Returns the conversation ID for the current session. + pub fn conversation_id(&self) -> &str { + self.connection.conversation_id() + } + + /// Returns whether the agent is currently idle (not processing). + pub fn is_idle(&self) -> bool { + self.connection.is_idle() + } + + /// Sends a trigger notification message to the agent. + pub async fn send(&self, message: &str) -> Result<()> { + self.connection.send_trigger_notification(message).await + } + + /// Retrieves a previously stored value by key. + /// Returns `None` if the key doesn't exist or deserialization fails. + pub fn get_state(&self, key: &str) -> Option { + self.state + .lock() + .ok() + .and_then(|store| store.get(key).cloned()) + .and_then(|v| serde_json::from_value(v).ok()) + } + + /// Stores a value by key in the session-scoped state store. + #[allow(clippy::collapsible_if)] + pub fn set_state(&self, key: &str, value: T) { + if let Ok(mut store) = self.state.lock() { + if let Ok(v) = serde_json::to_value(value) { + store.insert(key.to_string(), v); + } + } + } +} + +#[cfg(test)] +mod tests { + #![allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::significant_drop_tightening + )] + use super::*; + + // ToolContext tests require a mock connection which is only available + // via the full test harness. Unit tests here validate the state store. + #[test] + fn test_state_set_and_get() { + let state: Mutex> = Mutex::new(HashMap::new()); + state + .lock() + .unwrap() + .insert("key".to_string(), serde_json::to_value("value").unwrap()); + let val: String = + serde_json::from_value(state.lock().unwrap().get("key").cloned().unwrap()).unwrap(); + assert_eq!(val, "value"); + } + + #[test] + fn test_state_overwrite() { + let state: Mutex> = Mutex::new(HashMap::new()); + { + let mut store = state.lock().unwrap(); + store.insert("key".to_string(), serde_json::to_value(1i32).unwrap()); + store.insert("key".to_string(), serde_json::to_value(2i32).unwrap()); + } + let val: i32 = + serde_json::from_value(state.lock().unwrap().get("key").cloned().unwrap()).unwrap(); + assert_eq!(val, 2); + } +} diff --git a/src/types.rs b/src/types.rs index 74b8a43..e34ee69 100644 --- a/src/types.rs +++ b/src/types.rs @@ -6,6 +6,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::HashMap; +use std::path::Path; /// The default model name used when none is specified. pub const DEFAULT_MODEL: &str = "gemini-3.5-flash"; @@ -630,6 +631,363 @@ pub enum StreamChunk { ToolCall(ToolCall), } +// ─── Multimodal Content Types ─────────────────────────────────────────────── + +/// Supported MIME types for image media. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum ImageMime { + /// BMP image. + #[serde(rename = "image/bmp")] + Bmp, + /// JPEG image. + #[serde(rename = "image/jpeg")] + Jpeg, + /// PNG image. + #[serde(rename = "image/png")] + Png, + /// WebP image. + #[serde(rename = "image/webp")] + Webp, +} + +/// Supported MIME types for document media. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum DocumentMime { + /// PDF document. + #[serde(rename = "application/pdf")] + Pdf, + /// JSON data. + #[serde(rename = "application/json")] + Json, + /// CSS stylesheet. + #[serde(rename = "text/css")] + Css, + /// CSV tabular data. + #[serde(rename = "text/csv")] + Csv, + /// HTML page. + #[serde(rename = "text/html")] + Html, + /// JavaScript source. + #[serde(rename = "application/javascript")] + Javascript, + /// Plain text. + #[serde(rename = "text/plain")] + PlainText, + /// RTF document. + #[serde(rename = "text/rtf")] + Rtf, + /// XML data. + #[serde(rename = "application/xml")] + Xml, +} + +/// Supported MIME types for audio media. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum AudioMime { + /// WAV audio. + #[serde(rename = "audio/wav")] + Wav, + /// MP3 audio. + #[serde(rename = "audio/mp3")] + Mp3, + /// AAC audio. + #[serde(rename = "audio/aac")] + Aac, + /// OGG audio. + #[serde(rename = "audio/ogg")] + Ogg, + /// FLAC audio. + #[serde(rename = "audio/flac")] + Flac, + /// Opus audio. + #[serde(rename = "audio/opus")] + Opus, + /// MPEG audio. + #[serde(rename = "audio/mpeg")] + Mpeg, + /// M4A audio. + #[serde(rename = "audio/m4a")] + M4a, + /// L16 raw audio. + #[serde(rename = "audio/l16")] + L16, +} + +/// Supported MIME types for video media. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum VideoMime { + /// 3GPP video. + #[serde(rename = "video/3gpp")] + Threegpp, + /// AVI video. + #[serde(rename = "video/x-msvideo")] + Avi, + /// MP4 video. + #[serde(rename = "video/mp4")] + Mp4, + /// MPEG video. + #[serde(rename = "video/mpeg")] + VideoMpeg, + /// MPG video. + #[serde(rename = "video/mpg")] + Mpg, + /// `QuickTime` video. + #[serde(rename = "video/quicktime")] + Quicktime, + /// `WebM` video. + #[serde(rename = "video/webm")] + Webm, + /// WMV video. + #[serde(rename = "video/x-ms-wmv")] + Wmv, + /// FLV video. + #[serde(rename = "video/x-flv")] + XFlv, +} + +/// Validated MIME type for any supported media category. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(untagged)] +pub enum MimeType { + /// Image MIME type. + Image(ImageMime), + /// Document MIME type. + Document(DocumentMime), + /// Audio MIME type. + Audio(AudioMime), + /// Video MIME type. + Video(VideoMime), +} + +impl MimeType { + /// Returns the MIME type string. + pub const fn as_str(&self) -> &'static str { + match self { + Self::Image(m) => match m { + ImageMime::Bmp => "image/bmp", + ImageMime::Jpeg => "image/jpeg", + ImageMime::Png => "image/png", + ImageMime::Webp => "image/webp", + }, + Self::Document(m) => match m { + DocumentMime::Pdf => "application/pdf", + DocumentMime::Json => "application/json", + DocumentMime::Css => "text/css", + DocumentMime::Csv => "text/csv", + DocumentMime::Html => "text/html", + DocumentMime::Javascript => "application/javascript", + DocumentMime::PlainText => "text/plain", + DocumentMime::Rtf => "text/rtf", + DocumentMime::Xml => "application/xml", + }, + Self::Audio(m) => match m { + AudioMime::Wav => "audio/wav", + AudioMime::Mp3 => "audio/mp3", + AudioMime::Aac => "audio/aac", + AudioMime::Ogg => "audio/ogg", + AudioMime::Flac => "audio/flac", + AudioMime::Opus => "audio/opus", + AudioMime::Mpeg => "audio/mpeg", + AudioMime::M4a => "audio/m4a", + AudioMime::L16 => "audio/l16", + }, + Self::Video(m) => match m { + VideoMime::Threegpp => "video/3gpp", + VideoMime::Avi => "video/x-msvideo", + VideoMime::Mp4 => "video/mp4", + VideoMime::VideoMpeg => "video/mpeg", + VideoMime::Mpg => "video/mpg", + VideoMime::Quicktime => "video/quicktime", + VideoMime::Webm => "video/webm", + VideoMime::Wmv => "video/x-ms-wmv", + VideoMime::XFlv => "video/x-flv", + }, + } + } +} + +impl std::fmt::Display for MimeType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +/// A validated media payload with raw bytes and MIME type. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Media { + /// Raw binary data of the media file. + pub data: Vec, + /// Validated MIME type. + pub mime_type: MimeType, + /// Optional human-readable description. + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, +} + +/// An image media payload. +pub type Image = Media; + +/// A document media payload. +pub type Document = Media; + +/// An audio media payload. +pub type Audio = Media; + +/// A video media payload. +pub type Video = Media; + +/// A single content primitive for agent prompts. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum ContentPrimitive { + /// Plain text content. + Text(String), + /// Binary media content (image, document, audio, or video). + Media(Media), +} + +/// Agent prompt content — a single primitive or a list of primitives. +/// +/// Use `Content::from("text")` or `"text".into()` for simple text prompts. +/// Use `Content::from_file("image.png", None)` for media files. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum Content { + /// A single content primitive. + Single(ContentPrimitive), + /// Multiple content primitives (e.g., text + image). + Multi(Vec), +} + +impl Content { + /// Creates a text-only content. + pub fn text(s: impl Into) -> Self { + Self::Single(ContentPrimitive::Text(s.into())) + } + + /// Creates content from a media payload. + pub const fn media(media: Media) -> Self { + Self::Single(ContentPrimitive::Media(media)) + } + + /// Reads a file and auto-detects the MIME type to construct the appropriate media content. + /// + /// # Errors + /// Returns an error if the file cannot be read or the MIME type is unsupported. + pub fn from_file(path: impl AsRef, description: Option<&str>) -> Result { + let path = path.as_ref(); + let data = std::fs::read(path).map_err(|e| format!("Failed to read file: {e}"))?; + + let ext = path + .extension() + .and_then(|e| e.to_str()) + .unwrap_or("") + .to_lowercase(); + + let mime_type = mime_from_extension(&ext) + .ok_or_else(|| format!("Unsupported file extension: .{ext}"))?; + + Ok(Self::Single(ContentPrimitive::Media(Media { + data, + mime_type, + description: description.map(String::from), + }))) + } + + /// Returns the text content if this is a single text primitive. + pub fn as_text(&self) -> Option<&str> { + match self { + Self::Single(ContentPrimitive::Text(s)) => Some(s), + _ => None, + } + } +} + +impl From<&str> for Content { + fn from(s: &str) -> Self { + Self::text(s) + } +} + +impl From for Content { + fn from(s: String) -> Self { + Self::text(s) + } +} + +/// Resolves a file extension to a validated `MimeType`. +fn mime_from_extension(ext: &str) -> Option { + match ext { + // Images + "bmp" => Some(MimeType::Image(ImageMime::Bmp)), + "jpg" | "jpeg" => Some(MimeType::Image(ImageMime::Jpeg)), + "png" => Some(MimeType::Image(ImageMime::Png)), + "webp" => Some(MimeType::Image(ImageMime::Webp)), + // Documents + "pdf" => Some(MimeType::Document(DocumentMime::Pdf)), + "json" => Some(MimeType::Document(DocumentMime::Json)), + "css" => Some(MimeType::Document(DocumentMime::Css)), + "csv" => Some(MimeType::Document(DocumentMime::Csv)), + "html" | "htm" => Some(MimeType::Document(DocumentMime::Html)), + "js" | "mjs" => Some(MimeType::Document(DocumentMime::Javascript)), + "txt" | "text" | "md" | "log" => Some(MimeType::Document(DocumentMime::PlainText)), + "rtf" => Some(MimeType::Document(DocumentMime::Rtf)), + "xml" => Some(MimeType::Document(DocumentMime::Xml)), + // Audio + "wav" => Some(MimeType::Audio(AudioMime::Wav)), + "mp3" => Some(MimeType::Audio(AudioMime::Mp3)), + "aac" => Some(MimeType::Audio(AudioMime::Aac)), + "ogg" | "oga" => Some(MimeType::Audio(AudioMime::Ogg)), + "flac" => Some(MimeType::Audio(AudioMime::Flac)), + "opus" => Some(MimeType::Audio(AudioMime::Opus)), + "m4a" => Some(MimeType::Audio(AudioMime::M4a)), + // Video + "3gp" | "3gpp" => Some(MimeType::Video(VideoMime::Threegpp)), + "avi" => Some(MimeType::Video(VideoMime::Avi)), + "mp4" | "m4v" => Some(MimeType::Video(VideoMime::Mp4)), + "mpeg" | "mpg" => Some(MimeType::Video(VideoMime::VideoMpeg)), + "mov" => Some(MimeType::Video(VideoMime::Quicktime)), + "webm" => Some(MimeType::Video(VideoMime::Webm)), + "wmv" => Some(MimeType::Video(VideoMime::Wmv)), + "flv" => Some(MimeType::Video(VideoMime::XFlv)), + _ => None, + } +} + +// ─── Trigger & File Change Types ──────────────────────────────────────────── + +/// Controls when trigger notifications are delivered to the agent. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum TriggerDelivery { + /// Deliver the notification immediately, even if the agent is busy. + SendImmediately, + /// Wait until the agent is idle before delivering. + WaitIdle, +} + +/// The kind of filesystem change detected by a file-watching trigger. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum FileChangeKind { + /// A new file was created. + Added, + /// An existing file was modified. + Modified, + /// A file was deleted. + Deleted, +} + +/// A single filesystem change event. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FileChange { + /// The type of change. + pub kind: FileChangeKind, + /// The path of the affected file. + pub path: String, +} + #[cfg(test)] mod tests { #![allow( @@ -781,4 +1139,177 @@ mod tests { assert!(config.disabled_tools.is_none()); assert!(config.compaction_threshold.is_none()); } + + #[test] + fn test_content_from_str() { + let content: Content = "hello".into(); + assert_eq!(content.as_text(), Some("hello")); + } + + #[test] + fn test_content_from_string() { + let content: Content = String::from("world").into(); + assert_eq!(content.as_text(), Some("world")); + } + + #[test] + fn test_content_text_constructor() { + let content = Content::text("test"); + assert_eq!(content.as_text(), Some("test")); + } + + #[test] + fn test_content_media_has_no_text() { + let media = Media { + data: vec![0xFF, 0xD8], + mime_type: MimeType::Image(ImageMime::Jpeg), + description: None, + }; + let content = Content::media(media); + assert_eq!(content.as_text(), None); + } + + #[test] + fn test_mime_from_extension_images() { + assert_eq!( + mime_from_extension("png"), + Some(MimeType::Image(ImageMime::Png)) + ); + assert_eq!( + mime_from_extension("jpg"), + Some(MimeType::Image(ImageMime::Jpeg)) + ); + assert_eq!( + mime_from_extension("jpeg"), + Some(MimeType::Image(ImageMime::Jpeg)) + ); + assert_eq!( + mime_from_extension("webp"), + Some(MimeType::Image(ImageMime::Webp)) + ); + assert_eq!( + mime_from_extension("bmp"), + Some(MimeType::Image(ImageMime::Bmp)) + ); + } + + #[test] + fn test_mime_from_extension_documents() { + assert_eq!( + mime_from_extension("pdf"), + Some(MimeType::Document(DocumentMime::Pdf)) + ); + assert_eq!( + mime_from_extension("json"), + Some(MimeType::Document(DocumentMime::Json)) + ); + assert_eq!( + mime_from_extension("txt"), + Some(MimeType::Document(DocumentMime::PlainText)) + ); + assert_eq!( + mime_from_extension("md"), + Some(MimeType::Document(DocumentMime::PlainText)) + ); + } + + #[test] + fn test_mime_from_extension_audio() { + assert_eq!( + mime_from_extension("mp3"), + Some(MimeType::Audio(AudioMime::Mp3)) + ); + assert_eq!( + mime_from_extension("wav"), + Some(MimeType::Audio(AudioMime::Wav)) + ); + assert_eq!( + mime_from_extension("flac"), + Some(MimeType::Audio(AudioMime::Flac)) + ); + } + + #[test] + fn test_mime_from_extension_video() { + assert_eq!( + mime_from_extension("mp4"), + Some(MimeType::Video(VideoMime::Mp4)) + ); + assert_eq!( + mime_from_extension("webm"), + Some(MimeType::Video(VideoMime::Webm)) + ); + assert_eq!( + mime_from_extension("mov"), + Some(MimeType::Video(VideoMime::Quicktime)) + ); + } + + #[test] + fn test_mime_from_extension_unsupported() { + assert_eq!(mime_from_extension("xyz"), None); + assert_eq!(mime_from_extension(""), None); + } + + #[test] + fn test_mime_type_display() { + let m = MimeType::Image(ImageMime::Png); + assert_eq!(m.to_string(), "image/png"); + assert_eq!(m.as_str(), "image/png"); + } + + #[test] + fn test_content_from_file_missing() { + let result = Content::from_file("/nonexistent/file.png", None); + assert!(result.is_err()); + } + + #[test] + fn test_content_from_file_unsupported_ext() { + // Create a temp file with unsupported extension + let dir = std::env::temp_dir(); + let path = dir.join("test_antigravity.xyz"); + std::fs::write(&path, b"data").unwrap(); + let result = Content::from_file(&path, None); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("Unsupported")); + std::fs::remove_file(&path).unwrap(); + } + + #[test] + fn test_content_from_file_success() { + let dir = std::env::temp_dir(); + let path = dir.join("test_antigravity.txt"); + std::fs::write(&path, b"hello world").unwrap(); + let result = Content::from_file(&path, Some("test doc")); + assert!(result.is_ok()); + let content = result.unwrap(); + // It should be a media content, not text + assert_eq!(content.as_text(), None); + std::fs::remove_file(&path).unwrap(); + } + + #[test] + fn test_trigger_delivery_serialization() { + let d = TriggerDelivery::SendImmediately; + let json_str = serde_json::to_string(&d).unwrap(); + assert_eq!(json_str, "\"send_immediately\""); + } + + #[test] + fn test_file_change_kind_serialization() { + let k = FileChangeKind::Modified; + let json_str = serde_json::to_string(&k).unwrap(); + assert_eq!(json_str, "\"modified\""); + } + + #[test] + fn test_file_change_construction() { + let fc = FileChange { + kind: FileChangeKind::Added, + path: "/tmp/new_file.txt".to_string(), + }; + assert!(matches!(fc.kind, FileChangeKind::Added)); + assert_eq!(fc.path, "/tmp/new_file.txt"); + } } From 67d712185ab8fab186898456b6507323786f52f9 Mon Sep 17 00:00:00 2001 From: codeitlikemiley Date: Fri, 29 May 2026 19:53:37 +0800 Subject: [PATCH 2/6] feat: extend hook and tool systems for Python SDK parity - Add on_session_end, post_turn, on_compaction hooks to Hook/DynHook traits - Add dispatch_session_end, dispatch_post_turn, dispatch_on_compaction to HookRunner - Add needs_context() and call_with_context() to Tool/DynTool traits - All additions are backwards compatible with default implementations --- src/hooks.rs | 173 ++++++++++++++++++++++++++++++++++++++++++++++++++- src/tools.rs | 41 ++++++++++++ 2 files changed, 212 insertions(+), 2 deletions(-) diff --git a/src/hooks.rs b/src/hooks.rs index e2dc35a..3b6cc82 100644 --- a/src/hooks.rs +++ b/src/hooks.rs @@ -3,7 +3,9 @@ //! This module defines the [`Hook`] trait, which allows implementing custom observers and middlewares //! to intercept session startup, pre/post tool invocations, execution errors, and user interactions. -use crate::types::{AskQuestionEntry, HookResult, QuestionHookResult, ToolCall, ToolResult}; +use crate::types::{ + AskQuestionEntry, ChatResponse, HookResult, QuestionHookResult, ToolCall, ToolResult, +}; use futures_util::future::BoxFuture; use std::sync::Arc; @@ -76,6 +78,26 @@ pub trait Hook: Send + Sync { { async { Ok(None) } } + /// Triggered when the session is ending (agent shutdown or disconnect). + fn on_session_end( + &self, + ) -> impl std::future::Future> + Send { + async { Ok(()) } + } + /// Triggered after a turn completes, receiving the full response. + fn post_turn<'a>( + &'a self, + _response: &'a ChatResponse, + ) -> impl std::future::Future> + Send { + async { Ok(()) } + } + /// Triggered when the conversation history is compacted/summarized. + fn on_compaction<'a>( + &'a self, + _summary: &'a str, + ) -> impl std::future::Future> + Send { + async { Ok(()) } + } } /// Object-safe version of the [`Hook`] trait, automatically implemented via a blanket impl. @@ -111,6 +133,18 @@ pub trait DynHook: Send + Sync { &'a self, questions: &'a [AskQuestionEntry], ) -> BoxFuture<'a, Result, anyhow::Error>>; + + /// Triggered when the session is ending. + fn on_session_end(&self) -> BoxFuture<'_, Result<(), anyhow::Error>>; + + /// Triggered after a turn completes. + fn post_turn<'a>( + &'a self, + response: &'a ChatResponse, + ) -> BoxFuture<'a, Result<(), anyhow::Error>>; + + /// Triggered when the conversation history is compacted. + fn on_compaction<'a>(&'a self, summary: &'a str) -> BoxFuture<'a, Result<(), anyhow::Error>>; } impl DynHook for T { @@ -149,6 +183,21 @@ impl DynHook for T { ) -> BoxFuture<'a, Result, anyhow::Error>> { Box::pin(async move { self.on_interaction(questions).await }) } + + fn on_session_end(&self) -> BoxFuture<'_, Result<(), anyhow::Error>> { + Box::pin(async move { self.on_session_end().await }) + } + + fn post_turn<'a>( + &'a self, + response: &'a ChatResponse, + ) -> BoxFuture<'a, Result<(), anyhow::Error>> { + Box::pin(async move { self.post_turn(response).await }) + } + + fn on_compaction<'a>(&'a self, summary: &'a str) -> BoxFuture<'a, Result<(), anyhow::Error>> { + Box::pin(async move { self.on_compaction(summary).await }) + } } /// Internal helper that manages a collection of registered [`Hook`]s and dispatches events sequentially. @@ -256,6 +305,33 @@ impl HookRunner { } Ok(None) } + + /// Dispatches `on_session_end` to all registered hooks. + pub async fn dispatch_session_end(&self) -> Result<(), anyhow::Error> { + let hooks = self.hooks.read().await.clone(); + for hook in &hooks { + hook.on_session_end().await?; + } + Ok(()) + } + + /// Dispatches `post_turn` to all registered hooks. + pub async fn dispatch_post_turn(&self, response: &ChatResponse) -> Result<(), anyhow::Error> { + let hooks = self.hooks.read().await.clone(); + for hook in &hooks { + hook.post_turn(response).await?; + } + Ok(()) + } + + /// Dispatches `on_compaction` to all registered hooks. + pub async fn dispatch_on_compaction(&self, summary: &str) -> Result<(), anyhow::Error> { + let hooks = self.hooks.read().await.clone(); + for hook in &hooks { + hook.on_compaction(summary).await?; + } + Ok(()) + } } #[cfg(test)] @@ -267,7 +343,7 @@ mod tests { clippy::field_reassign_with_default )] use super::*; - use crate::types::{HookResult, QuestionHookResult, ToolCall, ToolResult}; + use crate::types::{HookResult, QuestionHookResult, ToolCall, ToolResult, UsageMetadata}; use std::sync::Mutex; struct TrackerHook { @@ -372,6 +448,30 @@ mod tests { Ok(None) } } + + async fn on_session_end(&self) -> Result<(), anyhow::Error> { + self.calls + .lock() + .unwrap() + .push(format!("{}:session_end", self.name)); + Ok(()) + } + + async fn post_turn(&self, _response: &ChatResponse) -> Result<(), anyhow::Error> { + self.calls + .lock() + .unwrap() + .push(format!("{}:post_turn", self.name)); + Ok(()) + } + + async fn on_compaction(&self, _summary: &str) -> Result<(), anyhow::Error> { + self.calls + .lock() + .unwrap() + .push(format!("{}:on_compaction", self.name)); + Ok(()) + } } #[tokio::test] @@ -575,4 +675,73 @@ mod tests { let recorded = calls.lock().unwrap().clone(); assert_eq!(recorded, vec!["h1:on_interaction", "answer:on_interaction"]); } + + #[tokio::test] + async fn test_dispatch_session_end() { + let calls = Arc::new(Mutex::new(Vec::new())); + let runner = HookRunner::new(); + runner + .register(Arc::new(TrackerHook { + name: "h1".to_string(), + calls: calls.clone(), + })) + .await; + runner + .register(Arc::new(TrackerHook { + name: "h2".to_string(), + calls: calls.clone(), + })) + .await; + + runner.dispatch_session_end().await.unwrap(); + + let recorded = calls.lock().unwrap().clone(); + assert_eq!(recorded, vec!["h1:session_end", "h2:session_end"]); + } + + #[tokio::test] + async fn test_dispatch_post_turn() { + let calls = Arc::new(Mutex::new(Vec::new())); + let runner = HookRunner::new(); + runner + .register(Arc::new(TrackerHook { + name: "h1".to_string(), + calls: calls.clone(), + })) + .await; + + let response = ChatResponse { + text: "hello".to_string(), + thinking: String::new(), + steps: vec![], + usage_metadata: UsageMetadata::default(), + }; + runner.dispatch_post_turn(&response).await.unwrap(); + + let recorded = calls.lock().unwrap().clone(); + assert_eq!(recorded, vec!["h1:post_turn"]); + } + + #[tokio::test] + async fn test_dispatch_on_compaction() { + let calls = Arc::new(Mutex::new(Vec::new())); + let runner = HookRunner::new(); + runner + .register(Arc::new(TrackerHook { + name: "h1".to_string(), + calls: calls.clone(), + })) + .await; + runner + .register(Arc::new(TrackerHook { + name: "h2".to_string(), + calls: calls.clone(), + })) + .await; + + runner.dispatch_on_compaction("summary text").await.unwrap(); + + let recorded = calls.lock().unwrap().clone(); + assert_eq!(recorded, vec!["h1:on_compaction", "h2:on_compaction"]); + } } diff --git a/src/tools.rs b/src/tools.rs index 6ab401b..6d8f0af 100644 --- a/src/tools.rs +++ b/src/tools.rs @@ -4,6 +4,7 @@ //! API requests) to be registered with the agent and executed when requested by the model. //! Registration and concurrent execution is managed via [`ToolRunner`]. +use crate::tool_context::ToolContext; use futures_util::future::BoxFuture; use serde_json::Value; use std::collections::HashMap; @@ -29,6 +30,24 @@ pub trait Tool: Send + Sync { &self, args: Value, ) -> impl std::future::Future> + Send; + + /// Returns `true` if this tool requires a [`ToolContext`] to be injected. + /// Defaults to `false` for backwards compatibility. + fn needs_context(&self) -> bool { + false + } + + /// Executes the tool with access to the session-scoped [`ToolContext`]. + /// + /// Only called when [`needs_context()`](Tool::needs_context) returns `true`. + /// The default implementation ignores the context and delegates to [`call()`](Tool::call). + fn call_with_context( + &self, + args: Value, + _context: &ToolContext, + ) -> impl std::future::Future> + Send { + self.call(args) + } } /// Object-safe version of the [`Tool`] trait, automatically implemented via a blanket impl. @@ -46,6 +65,16 @@ pub trait DynTool: Send + Sync { /// Executes the tool with the given JSON arguments. fn call(&self, args: Value) -> BoxFuture<'_, Result>; + + /// Returns `true` if this tool requires a `ToolContext`. + fn needs_context(&self) -> bool; + + /// Executes the tool with a `ToolContext`. + fn call_with_context<'a>( + &'a self, + args: Value, + context: &'a ToolContext, + ) -> BoxFuture<'a, Result>; } impl DynTool for T { @@ -64,6 +93,18 @@ impl DynTool for T { fn call(&self, args: Value) -> BoxFuture<'_, Result> { Box::pin(async move { self.call(args).await }) } + + fn needs_context(&self) -> bool { + Tool::needs_context(self) + } + + fn call_with_context<'a>( + &'a self, + args: Value, + context: &'a ToolContext, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { self.call_with_context(args, context).await }) + } } /// Registry and concurrent runner for custom tool implementations. From 3f71823030d13a10d01c730a727dedfc75690de0 Mon Sep 17 00:00:00 2001 From: codeitlikemiley Date: Fri, 29 May 2026 19:53:59 +0800 Subject: [PATCH 3/6] feat: add trigger helpers, interactive loop, and wire new modules - Add trigger_helpers.rs with every() periodic trigger factory - Add interactive.rs with run_interactive_loop() REPL (native only) - Wire all new modules in lib.rs: context, error, tool_context, trigger_helpers, interactive --- src/interactive.rs | 90 ++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 7 ++++ src/trigger_helpers.rs | 66 +++++++++++++++++++++++++++++++ 3 files changed, 163 insertions(+) create mode 100644 src/interactive.rs create mode 100644 src/trigger_helpers.rs diff --git a/src/interactive.rs b/src/interactive.rs new file mode 100644 index 0000000..3f667db --- /dev/null +++ b/src/interactive.rs @@ -0,0 +1,90 @@ +//! Interactive CLI loop for conversational agent sessions. +//! +//! Provides [`run_interactive_loop()`] as a convenience for building REPL-style +//! agent interfaces. Matches Python SDK's `utils/interactive.py`. + +use crate::agent::{Agent, Started}; +use crate::types::ChatResponse; +use anyhow::Result; + +/// Runs an interactive read-eval-print loop with the given started agent. +/// +/// Reads user input from stdin, sends it to the agent via `chat()`, +/// and prints the response. Type `exit` or `quit` to end the session. +/// +/// # Example +/// ```no_run +/// use antigravity_sdk_rust::agent::Agent; +/// use antigravity_sdk_rust::interactive::run_interactive_loop; +/// +/// #[tokio::main] +/// async fn main() -> Result<(), anyhow::Error> { +/// let agent = Agent::builder().allow_all().build().start().await?; +/// run_interactive_loop(&agent).await?; +/// agent.stop().await?; +/// Ok(()) +/// } +/// ``` +pub async fn run_interactive_loop(agent: &Agent) -> Result<()> { + use std::io::{self, BufRead, Write}; + + let stdin = io::stdin(); + let mut stdout = io::stdout(); + + loop { + write!(stdout, "\n> ")?; + stdout.flush()?; + + let mut input = String::new(); + let bytes_read = stdin.lock().read_line(&mut input)?; + + // EOF (Ctrl-D) + if bytes_read == 0 { + println!("\nGoodbye!"); + break; + } + + let trimmed = input.trim(); + + // Exit commands + if matches!(trimmed.to_lowercase().as_str(), "exit" | "quit" | "q") { + println!("Goodbye!"); + break; + } + + // Skip empty input + if trimmed.is_empty() { + continue; + } + + // Send to agent + match agent.chat(trimmed).await { + Ok(response) => { + print_response(&response); + } + Err(e) => { + eprintln!("Error: {e}"); + } + } + } + + Ok(()) +} + +/// Formats and prints an agent response to stdout. +fn print_response(response: &ChatResponse) { + if !response.thinking.is_empty() { + println!("\n💭 {}", response.thinking); + } + if !response.text.is_empty() { + println!("\n{}", response.text); + } + // Print usage stats + let usage = &response.usage_metadata; + if usage.total_token_count > 0 { + println!( + "\n📊 Tokens: {} prompt, {} response, {} total", + usage.prompt_token_count, usage.candidates_token_count, usage.total_token_count + ); + } +} diff --git a/src/lib.rs b/src/lib.rs index 0a6134a..0f32cf2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -61,7 +61,9 @@ pub mod proto { pub mod agent; pub mod connection; +pub mod context; pub mod conversation; +pub mod error; pub mod hooks; #[cfg(not(target_arch = "wasm32"))] pub mod local; @@ -69,10 +71,15 @@ pub mod local; pub mod wasm; pub mod policy; +pub mod tool_context; pub mod tools; +pub mod trigger_helpers; pub mod triggers; pub mod types; +#[cfg(not(target_arch = "wasm32"))] +pub mod interactive; + /// Helper to spawn asynchronous tasks in a target-agnostic manner. pub fn spawn_task(future: F) where diff --git a/src/trigger_helpers.rs b/src/trigger_helpers.rs new file mode 100644 index 0000000..24fe6c0 --- /dev/null +++ b/src/trigger_helpers.rs @@ -0,0 +1,66 @@ +//! Convenience factory functions for common trigger patterns. +//! +//! These helpers create ready-to-use [`Trigger`] implementations for common scenarios +//! like periodic timers and filesystem watchers. + +use crate::connection::AnyConnection; +use crate::connection::Connection; +use crate::triggers::Trigger; +use std::time::Duration; + +// ─── Periodic (every) Trigger ─────────────────────────────────────────────── + +/// A trigger that fires at regular intervals. +/// +/// Created via [`every()`]. +#[derive(Debug)] +pub struct PeriodicTrigger { + interval: Duration, + message: String, +} + +impl Trigger for PeriodicTrigger { + async fn run(&self, connection: AnyConnection) -> Result<(), anyhow::Error> { + loop { + tokio::time::sleep(self.interval).await; + connection.send_trigger_notification(&self.message).await?; + } + } +} + +/// Creates a trigger that fires every `interval` duration, sending `message` to the agent. +/// +/// # Example +/// ```no_run +/// use antigravity_sdk_rust::trigger_helpers::every; +/// use std::time::Duration; +/// +/// let trigger = every(Duration::from_secs(30), "check_status"); +/// ``` +pub fn every(interval: Duration, message: impl Into) -> PeriodicTrigger { + PeriodicTrigger { + interval, + message: message.into(), + } +} + +#[cfg(test)] +mod tests { + #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + use super::*; + + #[test] + fn test_every_construction() { + let trigger = every(Duration::from_secs(10), "heartbeat"); + assert_eq!(trigger.interval, Duration::from_secs(10)); + assert_eq!(trigger.message, "heartbeat"); + } + + #[test] + fn test_every_with_string() { + let msg = String::from("custom message"); + let trigger = every(Duration::from_millis(500), msg); + assert_eq!(trigger.interval, Duration::from_millis(500)); + assert_eq!(trigger.message, "custom message"); + } +} From 6e8c7f22589683bdb6431fe8bbb5dcc3494d5c37 Mon Sep 17 00:00:00 2001 From: codeitlikemiley Date: Fri, 29 May 2026 19:54:14 +0800 Subject: [PATCH 4/6] docs: document all new features in README - Add Error Handling section with AntigravityError - Add Multimodal Content section with Content/Media/MimeType - Add Lifecycle Hooks section showing full hook API - Add Hook Context section with hierarchical state - Add Context-Aware Tools section with ToolContext - Add Trigger Helpers section with every() factory - Add Interactive Loop section with run_interactive_loop() --- README.md | 153 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 153 insertions(+) diff --git a/README.md b/README.md index ef9f252..93251cd 100644 --- a/README.md +++ b/README.md @@ -172,6 +172,159 @@ let policies = vec![ ]; ``` +### Error Handling + +The SDK provides a typed error hierarchy via `AntigravityError`: + +```rust,no_run +use antigravity_sdk_rust::error::AntigravityError; + +// Three error categories: +// - AntigravityError::Connection(msg) — network/transport failures +// - AntigravityError::Execution(msg) — agent execution failures +// - AntigravityError::Validation { message, errors } — input validation failures +``` + +### Multimodal Content + +Send images, documents, audio, and video to the agent: + +```rust,no_run +use antigravity_sdk_rust::types::{Content, Media, MimeType, ImageMime}; + +// Text-only (backwards compatible) +let response = agent.chat("What is 2+2?").await?; + +// From file (auto-detects MIME type from extension) +let content = Content::from_file("photo.jpg", Some("Describe this")).unwrap(); + +// From raw bytes +let media = Media { + data: vec![0xFF, 0xD8, 0xFF], + mime_type: MimeType::Image(ImageMime::Jpeg), + description: Some("A photo".to_string()), +}; +let content = Content::media(media); +``` + +Supported formats: BMP, JPEG, PNG, WebP (images), PDF, JSON, CSS, CSV, HTML, JS, TXT, RTF, XML (documents), WAV, MP3, AAC, OGG, FLAC, Opus, MPEG, M4A, L16 (audio), 3GPP, AVI, MP4, MPEG, MOV, WebM, WMV, FLV (video). + +### Lifecycle Hooks + +The full hook lifecycle is available: + +```rust,no_run +use antigravity_sdk_rust::hooks::Hook; +use antigravity_sdk_rust::types::{HookResult, ToolCall, ToolResult, ChatResponse}; + +struct MyHook; + +impl Hook for MyHook { + // Session lifecycle + async fn on_session_start(&self) -> Result<(), anyhow::Error> { Ok(()) } + async fn on_session_end(&self) -> Result<(), anyhow::Error> { Ok(()) } + + // Turn lifecycle + async fn pre_turn(&self) -> Result { + Ok(HookResult { allow: true, message: String::new() }) + } + async fn post_turn(&self, _response: &ChatResponse) -> Result<(), anyhow::Error> { Ok(()) } + + // Tool lifecycle + async fn pre_tool_call(&self, _call: &ToolCall) -> Result { + Ok(HookResult { allow: true, message: String::new() }) + } + async fn post_tool_call(&self, _result: &ToolResult) -> Result<(), anyhow::Error> { Ok(()) } + + // History compaction + async fn on_compaction(&self, _summary: &str) -> Result<(), anyhow::Error> { Ok(()) } +} +``` + +### Hook Context (Hierarchical State) + +Hooks can share state across lifecycle events via a parent-chaining key-value store: + +```rust,no_run +use antigravity_sdk_rust::context::HookContext; +use std::sync::Arc; + +// Session context (root) +let session = Arc::new(HookContext::new()); +session.set("user_id", "u-123"); + +// Turn context (inherits session) +let turn = Arc::new(HookContext::child(session)); +turn.set("turn_count", 1i32); + +// get() walks up the chain: +assert_eq!(turn.get::("user_id"), Some("u-123".to_string())); +``` + +### Context-Aware Tools + +Tools can opt-in to receiving a `ToolContext` for session state and agent communication: + +```rust,no_run +use antigravity_sdk_rust::tools::Tool; +use antigravity_sdk_rust::tool_context::ToolContext; +use serde_json::Value; + +struct StatefulTool; + +impl Tool for StatefulTool { + fn name(&self) -> &str { "stateful_tool" } + fn description(&self) -> &str { "A tool with session state" } + fn parameters_json_schema(&self) -> &str { r#"{"type":"object"}"# } + + fn needs_context(&self) -> bool { true } // Opt-in + + async fn call(&self, args: Value) -> Result { + Ok(Value::Null) // Fallback + } + + async fn call_with_context( + &self, + args: Value, + ctx: &ToolContext, + ) -> Result { + // Access session state + let count: i32 = ctx.get_state("call_count").unwrap_or(0); + ctx.set_state("call_count", count + 1); + Ok(serde_json::json!({ "calls": count + 1 })) + } +} +``` + +### Trigger Helpers + +Convenience factories for common trigger patterns: + +```rust,no_run +use antigravity_sdk_rust::trigger_helpers::every; +use std::time::Duration; + +// Periodic trigger — fires every 30 seconds +let heartbeat = every(Duration::from_secs(30), "check_status"); +``` + +### Interactive Loop + +Built-in REPL for conversational agents: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::interactive::run_interactive_loop; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let agent = Agent::builder().allow_all().build().start().await?; + run_interactive_loop(&agent).await?; + agent.stop().await?; + Ok(()) +} +``` + ### Google Search Grounding & Web Search Fallback The SDK supports server-side Google Search grounding and provides a client-side search fallback: From 80c546f200b4c1606469d6d69ad590e8ced50254 Mon Sep 17 00:00:00 2001 From: codeitlikemiley Date: Fri, 29 May 2026 20:10:26 +0800 Subject: [PATCH 5/6] feat: add MCP server builder methods and wiring - Add mcp_servers field to AgentConfig - Add .mcp_server() and .mcp_servers() builder methods - Wire MCP configs through Agent::start() to LocalConnectionStrategy - Add mcp_servers to LocalConnectionStrategy struct and constructor - Fix all call sites (tests, README example) --- src/agent.rs | 19 ++++++++++++++++++- src/local.rs | 8 ++++++-- tests/documentation_examples.rs | 1 + 3 files changed, 25 insertions(+), 3 deletions(-) diff --git a/src/agent.rs b/src/agent.rs index 4148246..ee8730d 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -6,7 +6,8 @@ use crate::policy::{self, Policy, PolicyEnforcer}; use crate::tools::{DynTool, ToolRunner}; use crate::triggers::{DynTrigger, TriggerRunner}; use crate::types::{ - BuiltinTools, CapabilitiesConfig, ChatResponse, GeminiConfig, SystemInstructions, + BuiltinTools, CapabilitiesConfig, ChatResponse, GeminiConfig, McpServerConfig, + SystemInstructions, }; use anyhow::anyhow; use futures_util::future::BoxFuture; @@ -44,6 +45,8 @@ pub struct AgentConfig { pub app_data_dir: Option, /// Optional JSON schema constraining the final structured tool output. pub response_schema: Option, + /// MCP server configurations to connect to external tool servers. + pub mcp_servers: Vec, } impl std::fmt::Debug for AgentConfig { @@ -63,6 +66,7 @@ impl std::fmt::Debug for AgentConfig { .field("conversation_id", &self.conversation_id) .field("app_data_dir", &self.app_data_dir) .field("response_schema", &self.response_schema) + .field("mcp_servers", &self.mcp_servers) .finish() } } @@ -358,6 +362,7 @@ impl Agent { Some(self.tool_runner.clone()), Some(self.hook_runner.clone()), self.config.conversation_id.clone().unwrap_or_default(), + self.config.mcp_servers.clone(), ); let conn = strategy.connect().await?; @@ -554,6 +559,18 @@ impl

AgentBuilder

{ self } + /// Adds a single MCP server configuration. + pub fn mcp_server(mut self, server: McpServerConfig) -> Self { + self.config.mcp_servers.push(server); + self + } + + /// Sets the full list of MCP server configurations. + pub fn mcp_servers(mut self, servers: Vec) -> Self { + self.config.mcp_servers = servers; + self + } + pub fn policies(self, policies: Vec) -> AgentBuilder { let mut config = self.config; config.policies = Some(policies); diff --git a/src/local.rs b/src/local.rs index be59fc7..32a63d7 100644 --- a/src/local.rs +++ b/src/local.rs @@ -20,8 +20,8 @@ use crate::proto::localharness::{ use crate::tools::ToolRunner; use crate::types::{ AntigravityExecutionError, AskQuestionEntry, AskQuestionOption, BuiltinTools, - CapabilitiesConfig, GeminiConfig, QuestionHookResult, Step, StepSource, StepStatus, StepTarget, - StepType, SystemInstructions, ToolCall, ToolResult, UsageMetadata, + CapabilitiesConfig, GeminiConfig, McpServerConfig, QuestionHookResult, Step, StepSource, + StepStatus, StepTarget, StepType, SystemInstructions, ToolCall, ToolResult, UsageMetadata, }; use anyhow::anyhow; @@ -341,6 +341,8 @@ pub struct LocalConnectionStrategy { pub hook_runner: Option, /// Conversation ID for standard session resuming or tracking. pub conversation_id: String, + /// MCP server configurations. + pub mcp_servers: Vec, } impl LocalConnectionStrategy { @@ -357,6 +359,7 @@ impl LocalConnectionStrategy { tool_runner: Option, hook_runner: Option, conversation_id: String, + mcp_servers: Vec, ) -> Self { Self { binary_path, @@ -369,6 +372,7 @@ impl LocalConnectionStrategy { tool_runner, hook_runner, conversation_id, + mcp_servers, } } diff --git a/tests/documentation_examples.rs b/tests/documentation_examples.rs index 764bf79..f1a1254 100644 --- a/tests/documentation_examples.rs +++ b/tests/documentation_examples.rs @@ -41,6 +41,7 @@ async fn test_advanced_conversation_example() -> Result<(), anyhow::Error> { Some(tool_runner), None, "my_conversation_id".to_string(), + vec![], ); if false { From c945a43b3130ea4cf66d6beaeecd97ce050c2582 Mon Sep 17 00:00:00 2001 From: codeitlikemiley Date: Fri, 29 May 2026 20:10:44 +0800 Subject: [PATCH 6/6] docs: add per-component documentation and README sections MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add docs/agent.md — Agent builder, typestate, lifecycle - Add docs/connections.md — Connection trait, Local/WASM strategies - Add docs/conversation.md — Stateful sessions, streaming, chunks - Add docs/hooks.md — Hook trait, context, policy enforcer - Add docs/mcp.md — MCP server config, builder, policies - Add docs/tools.md — Tool trait, context-aware tools, ToolRunner - Add docs/triggers.md — Trigger trait, helpers, file watchers - Add MCP Integration section to README - Add Sugared Thoughts & Tool Call Streams section to README - Add Component Documentation links to README --- README.md | 85 ++++++ docs/agent.md | 696 +++++++++++++++++++++++++++++++++++++++++++ docs/connections.md | 386 ++++++++++++++++++++++++ docs/conversation.md | 408 +++++++++++++++++++++++++ docs/hooks.md | 608 +++++++++++++++++++++++++++++++++++++ docs/mcp.md | 289 ++++++++++++++++++ docs/tools.md | 240 +++++++++++++++ docs/triggers.md | 334 +++++++++++++++++++++ 8 files changed, 3046 insertions(+) create mode 100644 docs/agent.md create mode 100644 docs/connections.md create mode 100644 docs/conversation.md create mode 100644 docs/hooks.md create mode 100644 docs/mcp.md create mode 100644 docs/tools.md create mode 100644 docs/triggers.md diff --git a/README.md b/README.md index 93251cd..5aab294 100644 --- a/README.md +++ b/README.md @@ -106,6 +106,7 @@ async fn main() -> Result<(), anyhow::Error> { Some(tool_runner), None, "my_conversation_id".to_string(), + vec![], // MCP servers ); let connection = strategy.connect().await?; @@ -325,6 +326,78 @@ async fn main() -> Result<(), anyhow::Error> { } ``` +### MCP Integration + +Connect to external MCP servers and expose their tools to the agent: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::types::McpServerConfig; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let agent = Agent::builder() + .mcp_server(McpServerConfig::Stdio { + name: "my_server".to_string(), + command: "npx".to_string(), + args: vec!["my-mcp-server".to_string()], + enabled_tools: None, + disabled_tools: None, + }) + .allow_all() + .build(); + + let agent = agent.start().await?; + let response = agent.chat("Use the MCP tools to help me.").await?; + println!("{}", response.text); + agent.stop().await?; + Ok(()) +} +``` + +Three transport types are supported: +- **`McpServerConfig::Stdio`** — launch a local subprocess (e.g., `npx`, `uvx`) +- **`McpServerConfig::Sse`** — connect via Server-Sent Events +- **`McpServerConfig::Http`** — connect via standard HTTP with configurable timeouts + +Each variant supports `enabled_tools` / `disabled_tools` for fine-grained tool filtering. + +### Sugared Thoughts & Tool Call Streams (Advanced) + +For more complex use cases, stream internal model reasoning/thinking and intercept tool call dispatches in real-time using `StreamChunk`: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::types::StreamChunk; +use futures_util::StreamExt; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let agent = Agent::builder().allow_all().build().start().await?; + let conversation = agent.conversation(); + let mut stream = conversation.chat("Explain quantum computing").await?; + + while let Some(chunk_res) = stream.next().await { + match chunk_res? { + // 1. Stream reasoning/thinking deltas + StreamChunk::Thought { text, .. } => { + eprint!("💭 {}", text); // show thinking in grey/stderr + } + // 2. Stream response text tokens + StreamChunk::Text { text, .. } => { + print!("{}", text); + } + // 3. Stream strongly-typed ToolCall events + StreamChunk::ToolCall(call) => { + println!("\n🔧 Executing: {} (args: {})", call.name, call.args); + } + } + } + agent.stop().await?; + Ok(()) +} +``` + ### Google Search Grounding & Web Search Fallback The SDK supports server-side Google Search grounding and provides a client-side search fallback: @@ -447,6 +520,18 @@ This project uses [just](https://github.com/casey/just) to manage development ta just publish ``` +## Component Documentation + +For more detailed documentation on specific components, see: + +- **[Agent](docs/agent.md)** — High-level, batteries-included entry point. +- **[Connections](docs/connections.md)** — Transport and backend abstraction. +- **[Conversation](docs/conversation.md)** — Stateful session management. +- **[Hooks](docs/hooks.md)** — Agent lifecycle interception and policies. +- **[MCP](docs/mcp.md)** — Model Context Protocol integration. +- **[Tools](docs/tools.md)** — In-process tool execution. +- **[Triggers](docs/triggers.md)** — Background tasks and external events. + ## Architecture For more information, see [ARCHITECTURE.md](ARCHITECTURE.md). diff --git a/docs/agent.md b/docs/agent.md new file mode 100644 index 0000000..2518ce6 --- /dev/null +++ b/docs/agent.md @@ -0,0 +1,696 @@ +# Agent + +High-level, batteries-included entry point for building AI agents. + +## Overview + +The `Agent` struct manages the full lifecycle of an agentic AI session — binary discovery, tool wiring, hook registration, policy enforcement, and harness communication. It is the primary construct you'll interact with when building applications on the Antigravity SDK. + +```text +┌─────────────────────────────────────────────────────────────────┐ +│ Agent │ +│ ┌──────────────┐ ┌──────────┐ ┌─────────┐ ┌────────────┐ │ +│ │ AgentConfig │ │ToolRunner│ │HookRunner│ │PolicyEnforcer│ │ +│ │ (model, key, │ │(custom │ │(lifecycle│ │ (safety │ │ +│ │ policies…) │ │ tools) │ │ hooks) │ │ policies) │ │ +│ └──────────────┘ └──────────┘ └─────────┘ └────────────┘ │ +│ │ │ +│ .start().await? │ +│ ▼ │ +│ Agent │ +│ ┌──────────────────────┐ │ +│ │ Conversation (Arc) │ │ +│ │ ┌────────────────┐ │ │ +│ │ │ Connection │ │ │ +│ │ │ (WebSocket/IPC)│ │ │ +│ │ └────────────────┘ │ │ +│ └──────────────────────┘ │ +└─────────────────────────────────────────────────────────────────┘ +``` + +**Key responsibilities:** + +- **Binary discovery** — locates `localharness` automatically via env var, local install, `PATH`, or Python site-packages +- **Tool registration** — wires custom Rust `Tool` implementations into the model's callable toolset +- **Hook dispatch** — sequences lifecycle observer hooks (pre/post tool call, session start/end, etc.) +- **Policy enforcement** — compiles safety policies into a prioritized hook that gates tool execution +- **Connection management** — spawns the subprocess, upgrades to WebSocket, manages the conversation stream + +--- + +## Builder Pattern (Typestate) + +The agent uses a **compile-time typestate pattern** to guarantee that safety policies are always configured before an agent can be built. This eliminates an entire class of "forgot to set policies" runtime errors. + +```text +Agent::builder() + │ + ▼ +AgentBuilder ← .binary_path(), .api_key(), .tools(), .hooks(), etc. + │ + │── .allow_all() ──────────► AgentBuilder ──► .build() + │── .read_only() ──────────► AgentBuilder ──► .build() + │── .policies(…) ──────────► AgentBuilder ──► .build() + │── .policy(…) ──────────► AgentBuilder ──► .build() + │ + └── .build_unchecked() ────► Agent (escape hatch, skips check) +``` + +### How it works + +1. `Agent::builder()` returns `AgentBuilder` +2. Configuration methods (`.api_key()`, `.tools()`, `.hooks()`, etc.) return `Self`, preserving the current type-state +3. Policy-setting methods (`.allow_all()`, `.read_only()`, `.policies()`, `.policy()`) **consume** the builder and return `AgentBuilder` +4. `.build()` is only implemented on `AgentBuilder` — calling it on `NoPolicies` is a **compile error** +5. `.build_unchecked()` is available on any `AgentBuilder

` as an escape hatch + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; + +// ✅ Compiles — policies are set via .allow_all() +let agent = Agent::builder() + .api_key("my-key") + .allow_all() + .build(); + +// ❌ Compile error — .build() is not available on AgentBuilder +// let agent = Agent::builder() +// .api_key("my-key") +// .build(); // error[E0599]: no method named `build` found + +// ⚠️ Escape hatch — bypasses compile-time policy check +let agent = Agent::builder() + .api_key("my-key") + .build_unchecked(); +``` + +### Type-state markers + +```rust,no_run +// Marker types — you won't construct these directly +pub struct NoPolicies; // Initial state, .build() is NOT available +pub struct HasPolicies; // After setting policies, .build() IS available +``` + +> **Python SDK comparison:** The Python SDK performs this check at runtime during `agent.start()`, raising a `ValueError` if write tools are enabled without policies. The Rust SDK catches it at compile time, shifting the error left. + +--- + +## AgentConfig Fields + +`AgentConfig` is the resolved configuration struct that powers the `Agent`. While you'll typically use the builder, understanding the underlying fields is useful for advanced use cases. + +```rust,no_run +use antigravity_sdk_rust::agent::AgentConfig; + +// AgentConfig is Default — all fields start as None/empty +let config = AgentConfig::default(); +``` + +| Field | Type | Description | +|---|---|---| +| `binary_path` | `Option` | Path to the `localharness` binary. If `None`, auto-discovered via the [binary discovery](#binary-discovery) algorithm. | +| `gemini_config` | `GeminiConfig` | Model selection, API key, Vertex AI project/location, search grounding, URL context. | +| `capabilities` | `CapabilitiesConfig` | Tool enable/disable lists, compaction threshold, image model override, finish tool schema. | +| `system_instructions` | `Option` | Either `Custom` (full text override) or `Appended` (sections added to default identity). | +| `save_dir` | `Option` | Directory for persisting session state/logs. | +| `workspaces` | `Option>` | Working directories the agent may access. Defaults to `cwd` if `None`. Used by `workspace_only` policies. | +| `skills_paths` | `Vec` | Paths to folders containing custom skill modules. | +| `policies` | `Option>` | Safety policies controlling tool execution (approve, deny, ask_user). | +| `hooks` | `Vec>` | Lifecycle hooks — observe/intercept session start, tool calls, turns, errors, etc. | +| `triggers` | `Vec>` | Background async workers spawned when the agent starts. | +| `tools` | `Vec>` | Custom Rust tools registered for model invocation. | +| `mcp_servers` | `Vec` | Model Context Protocol server configurations (stdio, SSE, or HTTP transports). | +| `conversation_id` | `Option` | Assign or resume a specific conversation ID. | +| `app_data_dir` | `Option` | Application data directory for cache/configs. Defaults to `$HOME/.gemini/antigravity`. | +| `response_schema` | `Option` | JSON Schema string constraining the agent's final structured output. | + +--- + +## Builder Methods + +All builder methods follow the fluent chaining pattern. Methods that don't affect policy state preserve the current typestate (`Self`). Methods that set policies transition from `NoPolicies` → `HasPolicies`. + +### Configuration Methods (available on any `AgentBuilder

`) + +These methods return `Self` and can be called in any order, regardless of policy state. + +| Method | Signature | Description | +|---|---|---| +| `.binary_path()` | `fn binary_path(self, path: impl Into) -> Self` | Sets the path to the `localharness` binary. | +| `.gemini_config()` | `fn gemini_config(self, config: GeminiConfig) -> Self` | Sets the full Gemini configuration (model, API key, Vertex AI, search). | +| `.api_key()` | `fn api_key(self, key: impl Into) -> Self` | Shorthand to set the API key on `gemini_config`. | +| `.default_model()` | `fn default_model(self, model: impl Into) -> Self` | Shorthand to set the default model name (e.g. `"gemini-3.5-flash"`). | +| `.capabilities()` | `fn capabilities(self, caps: CapabilitiesConfig) -> Self` | Sets tool enable/disable lists and thresholds. | +| `.system_instructions()` | `fn system_instructions(self, si: SystemInstructions) -> Self` | Sets custom or appended system instructions. | +| `.save_dir()` | `fn save_dir(self, dir: impl Into) -> Self` | Sets the session state log directory. | +| `.workspaces()` | `fn workspaces(self, ws: Vec) -> Self` | Sets the allowed workspace directories. | +| `.skills_paths()` | `fn skills_paths(self, paths: Vec) -> Self` | Sets custom skill module folder paths. | +| `.hooks()` | `fn hooks(self, hooks: Vec>) -> Self` | Sets the full list of lifecycle hooks (replaces any existing). | +| `.hook()` | `fn hook(self, hook: Arc) -> Self` | Appends a single lifecycle hook. | +| `.triggers()` | `fn triggers(self, triggers: Vec>) -> Self` | Sets the full list of background triggers (replaces any existing). | +| `.trigger()` | `fn trigger(self, trigger: Arc) -> Self` | Appends a single background trigger. | +| `.tools()` | `fn tools(self, tools: Vec>) -> Self` | Sets the full list of custom tools (replaces any existing). | +| `.tool()` | `fn tool(self, tool: Arc) -> Self` | Appends a single custom tool. | +| `.conversation_id()` | `fn conversation_id(self, id: impl Into) -> Self` | Assigns or resumes a conversation ID. | +| `.app_data_dir()` | `fn app_data_dir(self, dir: impl Into) -> Self` | Sets the application data directory. | +| `.response_schema()` | `fn response_schema(self, schema: impl Into) -> Self` | Sets a JSON Schema for structured output. | +| `.mcp_server()` | `fn mcp_server(self, server: McpServerConfig) -> Self` | Appends a single MCP server configuration. | +| `.mcp_servers()` | `fn mcp_servers(self, servers: Vec) -> Self` | Sets the full list of MCP server configurations (replaces any existing). | + +### Policy Methods (transition `NoPolicies` → `HasPolicies`) + +These methods **consume** the builder and return `AgentBuilder`, enabling `.build()`. + +| Method | Signature | Description | +|---|---|---| +| `.policy()` | `fn policy(self, policy: Policy) -> AgentBuilder` | Appends a single policy and transitions to `HasPolicies`. | +| `.policies()` | `fn policies(self, policies: Vec) -> AgentBuilder` | Sets the full policy list and transitions to `HasPolicies`. | +| `.allow_all()` | `fn allow_all(self) -> AgentBuilder` | Convenience: sets `policy::allow_all()` — approves all tool calls unconditionally. | +| `.read_only()` | `fn read_only(self) -> AgentBuilder` | Convenience: denies all tools except read-only ones (`FindFile`, `ListDir`, `ViewFile`, `SearchDir`). | + +### Build Methods + +| Method | Signature | Available On | Description | +|---|---|---|---| +| `.build()` | `fn build(self) -> Agent` | `AgentBuilder` only | Constructs the agent. Compile error if policies are not set. | +| `.build_unchecked()` | `fn build_unchecked(self) -> Agent` | Any `AgentBuilder

` | Escape hatch — builds without compile-time policy check. Runtime errors may still occur at `.start()` if write tools are enabled without policies. | + +--- + +## Lifecycle: Unstarted → Started + +The `Agent` is generic over its lifecycle state, using the `AgentLifecycle` trait: + +```text +Agent ── .start().await? ──► Agent ── .stop().await? +``` + +### `Agent` + +An agent that has been configured but not yet connected. Available methods: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::hooks::DynHook; +use antigravity_sdk_rust::tools::DynTool; +use antigravity_sdk_rust::triggers::DynTrigger; +use std::sync::Arc; + +let mut agent = Agent::builder().allow_all().build(); + +// Additional registrations before starting +// agent.register_hook(hook); +// agent.register_tool(tool); +// agent.register_trigger(trigger)?; +``` + +| Method | Signature | Description | +|---|---|---| +| `Agent::new()` | `fn new(config: AgentConfig) -> Self` | Direct construction from config (prefer the builder). | +| `Agent::builder()` | `fn builder() -> AgentBuilder` | Returns a new builder. | +| `.register_hook()` | `fn register_hook(&mut self, hook: Arc)` | Adds a hook after construction but before starting. | +| `.register_tool()` | `fn register_tool(&mut self, tool: Arc)` | Adds a tool after construction but before starting. | +| `.register_trigger()` | `fn register_trigger(&mut self, trigger: Arc) -> Result<()>` | Adds a trigger after construction but before starting. | +| `.start()` | `fn start(self) -> BoxFuture<'static, Result>>` | Resolves the binary, connects, registers tools/hooks/policies, starts triggers. Consumes `self`. | + +### `Agent` + +A running agent with an active connection. Available methods: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; + +# async fn example() -> Result<(), anyhow::Error> { +let agent = Agent::builder() + .allow_all() + .build() + .start() + .await?; + +// Send a prompt and wait for the full response +let response = agent.chat("What is 2+2?").await?; +println!("Text: {}", response.text); +println!("Thinking: {}", response.thinking); +println!("Steps: {}", response.steps.len()); + +// Access the conversation for streaming or advanced use +let conversation = agent.conversation(); + +// Get the conversation ID +let id = agent.conversation_id(); + +// Shut down gracefully +agent.stop().await?; +# Ok(()) +# } +``` + +| Method | Signature | Description | +|---|---|---| +| `.chat()` | `async fn chat(&self, prompt: &str) -> Result` | Sends a prompt and awaits the complete response. Returns text, thinking, steps, and usage metadata. | +| `.conversation()` | `fn conversation(&self) -> Arc` | Returns the active `Conversation` for streaming or direct access. | +| `.conversation_id()` | `fn conversation_id(&self) -> String` | Returns the conversation session ID. | +| `.stop()` | `async fn stop(&self) -> Result<()>` | Gracefully disconnects the harness and stops the session. | + +### `ChatResponse` + +The response from `.chat()` contains: + +```rust,no_run +use antigravity_sdk_rust::types::ChatResponse; + +// ChatResponse { +// text: String, // Combined model text output +// thinking: String, // Combined reasoning/thinking text +// steps: Vec, // All intermediate execution steps +// usage_metadata: UsageMetadata, // Token consumption stats +// } +``` + +--- + +## `.start()` Internals + +When `.start()` is called, the following happens in order: + +1. **Resolve binary path** — Uses the [binary discovery](#binary-discovery) algorithm to find `localharness` +2. **Register hooks** — All configured hooks are added to the `HookRunner` +3. **Process capabilities** — Resolves enabled/disabled tool lists (mutually exclusive) +4. **Compile policies** — Builds `PolicyEnforcer` from configured policies; prepends `workspace_only` policies unless `allow_all()` was used +5. **Safety check** — Fails if write tools are enabled without any policies +6. **Register tools** — All custom tools are added to the `ToolRunner` +7. **Connect** — Spawns `localharness` subprocess, establishes WebSocket connection +8. **Start triggers** — Spawns background trigger tasks + +### Errors + +`.start()` returns `Err` if: + +- The `localharness` binary cannot be found +- `enabled_tools` and `disabled_tools` are both set (mutually exclusive) +- Write tools are enabled but no policies are configured +- The subprocess or WebSocket connection fails + +--- + +## Binary Discovery + +When `binary_path` is not explicitly set, the SDK searches for `localharness` in the following order: + +| Priority | Location | Description | +|---|---|---| +| 1 | `$ANTIGRAVITY_HARNESS_PATH` | Environment variable override | +| 2 | `./bin/localharness` | Local install relative to `cwd` (where `just install` places it) | +| 3 | `$PATH` lookup | Standard PATH search (e.g. via `pip install google-antigravity`) | +| 4 | Python site-packages | Fallback: queries `python3 -c "import site; ..."` and checks `google/antigravity/bin/localharness` | + +If none are found, `.start()` returns an error with a message to specify `binary_path` explicitly. + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; + +# async fn example() -> Result<(), anyhow::Error> { +// Explicit path — skips discovery +let agent = Agent::builder() + .binary_path("/usr/local/bin/localharness") + .allow_all() + .build() + .start() + .await?; + +// Auto-discovery — checks env var, ./bin, PATH, site-packages +let agent = Agent::builder() + .allow_all() + .build() + .start() + .await?; +# Ok(()) +# } +``` + +--- + +## Examples + +### 1. Minimal Hello World + +The simplest possible agent — auto-discovers the binary, uses default model, permits all tool calls: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let agent = Agent::builder() + .allow_all() + .build() + .start() + .await?; + + let response = agent.chat("Say 'Hello World!'").await?; + println!("Agent: {}", response.text); + + agent.stop().await?; + Ok(()) +} +``` + +### 2. Custom Model + API Key + +Configure a specific model and API key, with an explicit binary path: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + dotenvy::dotenv().ok(); + + let agent = Agent::builder() + .binary_path(std::env::var("ANTIGRAVITY_HARNESS_PATH").unwrap()) + .api_key(std::env::var("GEMINI_API_KEY").unwrap()) + .default_model("gemini-3.5-flash") + .allow_all() + .build() + .start() + .await?; + + let response = agent.chat("Explain Rust's ownership model in 3 sentences.").await?; + println!("{}", response.text); + println!("Tokens used: {}", response.usage_metadata.total_token_count); + + agent.stop().await?; + Ok(()) +} +``` + +### 3. Custom Tools + Hooks + Policies + +Register custom tools with fine-grained safety policies — deny all built-in tools, allow only your custom ones: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::hooks::Hook; +use antigravity_sdk_rust::policy; +use antigravity_sdk_rust::tools::Tool; +use antigravity_sdk_rust::types::{ + ChatResponse, CustomSystemInstructions, HookResult, SystemInstructions, ToolCall, +}; +use serde_json::Value; +use std::sync::Arc; + +// ── Custom Tool ───────────────────────────────────────────────────── + +struct WeatherTool; + +impl Tool for WeatherTool { + fn name(&self) -> &'static str { + "get_weather" + } + + fn description(&self) -> &'static str { + "Returns the current weather for a given city." + } + + fn parameters_json_schema(&self) -> &'static str { + r#"{ + "type": "object", + "properties": { + "city": { "type": "string", "description": "City name" } + }, + "required": ["city"] + }"# + } + + async fn call(&self, args: Value) -> Result { + let city = args.get("city").and_then(Value::as_str).unwrap_or("unknown"); + Ok(Value::String(format!("Weather in {city}: 22°C, partly cloudy"))) + } +} + +// ── Custom Hook ───────────────────────────────────────────────────── + +struct AuditHook; + +impl Hook for AuditHook { + async fn pre_tool_call(&self, tool_call: &ToolCall) -> Result { + println!("[AUDIT] Tool called: {} with args: {}", tool_call.name, tool_call.args); + Ok(HookResult { allow: true, message: String::new() }) + } + + async fn post_turn(&self, response: &ChatResponse) -> Result<(), anyhow::Error> { + println!("[AUDIT] Turn complete. Tokens: {}", response.usage_metadata.total_token_count); + Ok(()) + } +} + +// ── Main ──────────────────────────────────────────────────────────── + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let agent = Agent::builder() + .api_key("your-api-key") + .default_model("gemini-3.5-flash") + .system_instructions(SystemInstructions::Custom(CustomSystemInstructions { + text: "You are a helpful weather assistant. Use get_weather to answer weather questions.".to_string(), + })) + .tool(Arc::new(WeatherTool)) + .hook(Arc::new(AuditHook)) + .policies(vec![ + policy::deny_all(), // Deny everything by default + policy::allow("get_weather"), // Allow only our custom tool + ]) + .build() + .start() + .await?; + + let response = agent.chat("What's the weather like in Tokyo?").await?; + println!("Agent: {}", response.text); + + agent.stop().await?; + Ok(()) +} +``` + +### 4. Structured Output + +Constrain the agent to return JSON matching a schema: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let schema = r#"{ + "type": "object", + "properties": { + "name": { "type": "string" }, + "capital": { "type": "string" }, + "population": { "type": "integer" } + }, + "required": ["name", "capital", "population"] + }"#; + + let agent = Agent::builder() + .api_key("your-api-key") + .default_model("gemini-3.5-flash") + .response_schema(schema) + .allow_all() + .build() + .start() + .await?; + + let response = agent.chat("Tell me about Japan.").await?; + println!("Raw text: {}", response.text); + + // The last Finish step contains the parsed structured output + for step in &response.steps { + if let Some(ref output) = step.structured_output { + println!("Structured: {}", serde_json::to_string_pretty(output)?); + } + } + + agent.stop().await?; + Ok(()) +} +``` + +### 5. Streaming via `conversation()` + +For token-by-token streaming, access the `Conversation` directly: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::types::StreamChunk; +use futures_util::StreamExt; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let agent = Agent::builder() + .allow_all() + .build() + .start() + .await?; + + let conversation = agent.conversation(); + + // Send a prompt and get a streaming chunk iterator + let mut stream = conversation.chat("Write a haiku about Rust.").await?; + + // Process chunks as they arrive + while let Some(chunk_result) = stream.next().await { + match chunk_result? { + StreamChunk::Thought { text, .. } => { + eprint!("[thinking] {}", text); + } + StreamChunk::Text { text, .. } => { + print!("{}", text); + } + StreamChunk::ToolCall(call) => { + println!("\n[tool] {} called with: {}", call.name, call.args); + } + } + } + println!(); + + // Access conversation metadata + println!("Total turns: {}", conversation.turn_count().await); + println!("Total tokens: {}", conversation.total_usage().await.total_token_count); + println!("History steps: {}", conversation.history().await.len()); + + agent.stop().await?; + Ok(()) +} +``` + +### 6. Read-Only Agent + +Restrict the agent to only read files — no writes, no commands: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let agent = Agent::builder() + .workspaces(vec!["/home/user/project".to_string()]) + .read_only() // Only FindFile, ListDir, ViewFile, SearchDir + .build() + .start() + .await?; + + let response = agent.chat("Summarize the README.md").await?; + println!("{}", response.text); + + agent.stop().await?; + Ok(()) +} +``` + +### 7. Multi-Turn Conversation + +The agent maintains state across turns within a single session: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let agent = Agent::builder() + .allow_all() + .build() + .start() + .await?; + + let r1 = agent.chat("My name is Alice.").await?; + println!("Agent: {}", r1.text); + + let r2 = agent.chat("What's my name?").await?; + println!("Agent: {}", r2.text); // Should reference "Alice" + + let r3 = agent.chat("Summarize our conversation.").await?; + println!("Agent: {}", r3.text); + + agent.stop().await?; + Ok(()) +} +``` + +### 8. Advanced: Vertex AI + Full GeminiConfig + +Use Vertex AI backend with custom model configuration: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::types::{ + GeminiConfig, GenerationConfig, ModelConfig, ModelEntry, ThinkingLevel, +}; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let gemini_config = GeminiConfig { + vertex: true, + project: Some("my-gcp-project".to_string()), + location: Some("us-central1".to_string()), + models: ModelConfig { + default: ModelEntry { + name: "gemini-3.5-flash".to_string(), + api_key: None, + generation: GenerationConfig { + thinking_level: Some(ThinkingLevel::High), + }, + }, + ..Default::default() + }, + enable_google_search: Some(true), + enable_url_context: Some(true), + ..Default::default() + }; + + let agent = Agent::builder() + .gemini_config(gemini_config) + .allow_all() + .build() + .start() + .await?; + + let response = agent.chat("Explain quantum computing.").await?; + println!("{}", response.text); + if !response.thinking.is_empty() { + println!("\n--- Thinking ---\n{}", response.thinking); + } + + agent.stop().await?; + Ok(()) +} +``` + +--- + +## Comparison with Python SDK + +| Feature | Python SDK | Rust SDK | +|---|---|---| +| Policy enforcement | Runtime `ValueError` during `agent.connect()` | Compile-time typestate — `.build()` unavailable without policies | +| Tool trait | `Tool` base class with `call(args)` | `Tool` trait with `name()`, `description()`, `parameters_json_schema()`, `call(args)` | +| Hook trait | `Hook` base class, overridable methods | `Hook` trait with default async no-ops | +| Builder pattern | `Agent(model=..., tools=[...], ...)` constructor kwargs | Typestate builder: `Agent::builder().model().tools().allow_all().build()` | +| Streaming | `async for chunk in conversation.chat(prompt)` | `conversation.chat(prompt).await?` → `StreamExt::next()` on `BoxStream` | +| Lifecycle states | Implicit (`agent.connect()` / `agent.close()`) | Typestate: `Agent` / `Agent` — method availability enforced at compile time | +| Async runtime | `asyncio` | `tokio` | +| Object safety | Not applicable (duck typing) | `DynTool` / `DynHook` / `DynTrigger` object-safe wrappers via blanket impls | + +--- + +## Related Types + +For deeper dives into the types referenced here, see: + +- **[`Conversation`](conversation.md)** — Stateful session wrapper with streaming and history +- **[`Tool`](tools.md)** — Custom tool trait and registration +- **[`Hook`](hooks.md)** — Lifecycle event hooks +- **[`Policy`](policy.md)** — Safety policy system +- **[`Trigger`](triggers.md)** — Background async workers +- **[`GeminiConfig`](types.md)** — Model and API configuration +- **[`McpServerConfig`](types.md)** — MCP server configuration (stdio, SSE, HTTP) diff --git a/docs/connections.md b/docs/connections.md new file mode 100644 index 0000000..0037287 --- /dev/null +++ b/docs/connections.md @@ -0,0 +1,386 @@ +# Connections + +Transport and backend abstraction layer for the Antigravity Rust SDK. + +## Overview + +The connection layer abstracts communication between the SDK and the `localharness` backend. +Every agent session is backed by a **connection** — an active transport channel that sends +prompts, receives step updates, dispatches tool handshakes, and manages idle/disconnect lifecycle. + +Most users never interact with connections directly; the [`Agent`](./agent.md) builder handles +connection setup automatically. This module is relevant when you need to: + +- Understand the transport architecture +- Implement a custom connection strategy +- Use `Conversation` directly without the `Agent` wrapper + +## Connection Trait + +The `Connection` trait defines the transport-agnostic interface that all connection +implementations must satisfy: + +```rust,no_run +use futures_util::stream::BoxStream; +use antigravity_sdk_rust::types::{QuestionHookResult, Step, ToolResult}; + +pub trait Connection: Send + Sync { + /// Returns the unique conversation ID for this session. + fn conversation_id(&self) -> &str; + + /// Returns whether the connection is currently idle (no active turns). + fn is_idle(&self) -> bool; + + /// Subscribes to the stream of step updates from the harness. + fn receive_steps(&self) -> BoxStream<'static, Result>; + + /// Sends a user text prompt to the harness. + fn send( + &self, + content: &str, + ) -> impl std::future::Future> + Send; + + /// Sends a trigger notification message (from automated triggers). + fn send_trigger_notification( + &self, + content: &str, + ) -> impl std::future::Future> + Send; + + /// Sends a halt/cancellation request. + fn send_halt_request( + &self, + ) -> impl std::future::Future> + Send; + + /// Sends approval or rejection for a tool execution confirmation. + fn send_tool_confirmation( + &self, + trajectory_id: &str, + step_index: u32, + accepted: bool, + ) -> impl std::future::Future> + Send; + + /// Sends the result of a client-side tool execution back to the harness. + fn send_tool_response( + &self, + id: &str, + result: ToolResult, + ) -> impl std::future::Future> + Send; + + /// Sends answers to interactive user questions back to the harness. + fn send_question_response( + &self, + trajectory_id: &str, + step_index: u32, + answers: QuestionHookResult, + ) -> impl std::future::Future> + Send; + + /// Gracefully closes the connection and releases resources. + fn disconnect( + &self, + ) -> impl std::future::Future> + Send; +} +``` + +### Key Design Points + +| Method | Purpose | +|--------|---------| +| `conversation_id()` | Session identifier — learned from the harness during the first `StepUpdate` | +| `is_idle()` | Tracks idle state via `AtomicBool`; transitions on `TrajectoryStateUpdate` events | +| `receive_steps()` | Returns a `BoxStream` that yields `Step` events until the connection goes idle | +| `send()` | Serializes the prompt into an `InputEvent::UserInput` and sends over WebSocket | +| `send_halt_request()` | Cancels an in-progress agent turn | +| `send_tool_confirmation()` | Accepts or rejects a built-in tool call (policy enforcement) | +| `send_tool_response()` | Returns custom tool results to the harness for continued execution | +| `disconnect()` | Kills the subprocess (local) or closes the WebSocket (WASM) | + +## AnyConnection + +`AnyConnection` is a target-aware enum that dispatches to the concrete connection for the +current compilation target. It implements `Connection` by delegating every method to the +inner variant: + +```rust,no_run +#[derive(Clone)] +pub enum AnyConnection { + // Native platforms (not wasm32) + #[cfg(not(target_arch = "wasm32"))] + Local(std::sync::Arc), + + // WebAssembly targets + #[cfg(target_arch = "wasm32")] + Wasm(std::sync::Arc), + + // Testing only + #[cfg(test)] + Mock(std::sync::Arc), +} +``` + +| Variant | Target | Transport | +|---------|--------|-----------| +| `Local(Arc)` | Native (macOS, Linux, Windows) | Subprocess → WebSocket | +| `Wasm(Arc)` | `wasm32` | TCP → WebSocket to remote harness | +| `Mock(Arc)` | `#[cfg(test)]` only | In-memory channel for unit tests | + +The `AnyConnection` enum is `Clone` and `Debug`, and is the concrete type stored inside +`Conversation`. Users typically never construct it directly — `Agent::start()` does this +internally. + +## LocalConnectionStrategy + +`LocalConnectionStrategy` is the connection factory for **native platforms**. It orchestrates +a multi-phase startup sequence: + +### Connection Lifecycle + +```text +┌──────────────────────────────────────────────────────────────┐ +│ 1. Spawn localharness subprocess (stdin/stdout/stderr) │ +│ 2. Length-prefixed protobuf handshake over stdin/stdout │ +│ ├─ Send: InputConfig { storage_dir, client_info } │ +│ └─ Recv: OutputConfig { port, api_key } │ +│ 3. WebSocket upgrade to ws://localhost:{port}/ │ +│ └─ Header: x-goog-api-key = {harness_api_key} │ +│ 4. Send InitializeConversationEvent with HarnessConfig │ +│ ├─ GeminiConfig (model, API key, thinking level) │ +│ ├─ SystemInstructions (custom or appended) │ +│ ├─ Custom tools (name, description, JSON schema) │ +│ ├─ Harness-side tools (find, edit, run_command, etc.) │ +│ ├─ Workspaces (filesystem directories) │ +│ └─ Skills paths │ +│ 5. Spawn background tasks: │ +│ ├─ WS Sender loop (mpsc channel → ws_write) │ +│ ├─ WS Reader loop (ws_read → step_tx channel) │ +│ └─ Stderr reader (logs harness stderr as tracing::info) │ +└──────────────────────────────────────────────────────────────┘ +``` + +### Configuration + +```rust,no_run +use antigravity_sdk_rust::local::LocalConnectionStrategy; +use antigravity_sdk_rust::types::{ + GeminiConfig, CapabilitiesConfig, SystemInstructions, +}; +use antigravity_sdk_rust::tools::ToolRunner; +use antigravity_sdk_rust::hooks::HookRunner; + +let strategy = LocalConnectionStrategy::new( + "/path/to/localharness".to_string(), // binary_path + GeminiConfig::default(), // gemini_config + CapabilitiesConfig::default(), // capabilities_config + None, // system_instructions + Some("/tmp/agent-state".to_string()), // save_dir + vec!["/my/workspace".to_string()], // workspaces + vec![], // skills_paths + Some(ToolRunner::new()), // tool_runner + Some(HookRunner::new()), // hook_runner + "conv-123".to_string(), // conversation_id +); +``` + +### API Key Resolution + +The strategy resolves the Gemini API key in this priority order: + +1. `gemini_config.models.default.api_key` (model-level override) +2. `gemini_config.api_key` (global config) +3. `GEMINI_API_KEY` environment variable + +For **Vertex AI** mode (`gemini_config.vertex = true`), either an API key (Express Mode) +or both `project` and `location` must be set. + +### WebSocket Retry + +The initial WebSocket connection retries up to **5 times** with exponential backoff +starting at 100ms. If all attempts fail, the subprocess is killed and an error is returned. + +## WasmConnectionStrategy + +`WasmConnectionStrategy` is the connection factory for **WebAssembly** targets. Instead of +spawning a subprocess, it connects to a **remote** `localharness` instance over TCP → WebSocket: + +```rust,no_run +use antigravity_sdk_rust::wasm::WasmConnectionStrategy; +use antigravity_sdk_rust::types::{GeminiConfig, CapabilitiesConfig}; + +let strategy = WasmConnectionStrategy { + gemini_config: GeminiConfig::default(), + capabilities_config: CapabilitiesConfig::default(), + system_instructions: None, + save_dir: None, + workspaces: vec![], + skills_paths: vec![], + tool_runner: None, + hook_runner: None, + conversation_id: "wasm-conv-1".to_string(), +}; +``` + +### Environment Variables + +| Variable | Default | Purpose | +|----------|---------|---------| +| `ANTIGRAVITY_HARNESS_HOST` | `127.0.0.1` | Remote harness host | +| `ANTIGRAVITY_HARNESS_PORT` | `8000` | Remote harness port | +| `ANTIGRAVITY_API_KEY` | — | API key (fallback after `GEMINI_API_KEY`) | + +### Key Differences from LocalConnectionStrategy + +| Aspect | Local | WASM | +|--------|-------|------| +| Subprocess | Spawns `localharness` | Connects to pre-running instance | +| Transport | `tokio-tungstenite` async WebSocket | `tungstenite` sync WebSocket with non-blocking TCP | +| Task spawning | `tokio::spawn` | `any_spawner::Executor::spawn_local` | +| API key env var | `GEMINI_API_KEY` | `ANTIGRAVITY_API_KEY` (then `GEMINI_API_KEY`) | +| Reader loop | Async `ws_read.next().await` | Polling `ws_lock.read()` with yield | + +## MockConnection (Testing) + +Available only in `#[cfg(test)]` builds, `MockConnection` provides an in-memory +connection for unit testing: + +```rust,no_run +// Only available in test builds +use antigravity_sdk_rust::connection::MockConnection; +use std::sync::Arc; + +let mock = Arc::new(MockConnection::new("test-conv-1")); + +// Pre-load steps for the test +mock.steps_to_yield.lock().unwrap().push(/* Step { ... } */); + +// After use, inspect what was sent +let sent = mock.sent_prompts.lock().unwrap(); +assert_eq!(sent[0], "Hello"); +``` + +### MockConnection Fields + +| Field | Type | Purpose | +|-------|------|---------| +| `id` | `String` | Conversation ID returned by `conversation_id()` | +| `is_idle` | `AtomicBool` | Controllable idle state | +| `steps_to_yield` | `Mutex>` | Pre-loaded steps for `receive_steps()` | +| `sent_prompts` | `Mutex>` | Records all prompts passed to `send()` | + +## Direct Usage Example + +For advanced users who need direct control over the connection layer without the +`Agent` builder: + +```rust,no_run +use antigravity_sdk_rust::local::LocalConnectionStrategy; +use antigravity_sdk_rust::connection::{AnyConnection, Connection}; +use antigravity_sdk_rust::conversation::Conversation; +use antigravity_sdk_rust::types::{ + GeminiConfig, CapabilitiesConfig, StreamChunk, +}; +use futures_util::StreamExt; +use std::sync::Arc; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + // 1. Create the connection strategy + let strategy = LocalConnectionStrategy::new( + std::env::var("ANTIGRAVITY_HARNESS_PATH") + .unwrap_or_else(|_| "./bin/localharness".to_string()), + GeminiConfig { + api_key: std::env::var("GEMINI_API_KEY").ok(), + ..Default::default() + }, + CapabilitiesConfig::default(), + None, // system_instructions + None, // save_dir + vec![".".to_string()], // workspaces + vec![], // skills_paths + None, // tool_runner + None, // hook_runner + String::new(), // conversation_id (auto-assigned) + ); + + // 2. Connect — spawns subprocess, handshakes, upgrades to WebSocket + let conn = strategy.connect().await?; + let any_conn = AnyConnection::Local(Arc::new(conn)); + + // 3. Wrap in a Conversation for state tracking + let conversation = Conversation::new(any_conn, Some(5000)); + + // 4. Chat with streaming + let mut stream = conversation.chat("Explain Rust's ownership model").await?; + while let Some(chunk_res) = stream.next().await { + match chunk_res? { + StreamChunk::Thought { text, .. } => { + eprint!("{}", text); // thoughts to stderr + } + StreamChunk::Text { text, .. } => { + print!("{}", text); // response to stdout + } + StreamChunk::ToolCall(call) => { + eprintln!("[Tool: {} args: {}]", call.name, call.args); + } + } + } + + // 5. Inspect usage + let usage = conversation.total_usage().await; + eprintln!( + "\nTokens — prompt: {}, candidates: {}, total: {}", + usage.prompt_token_count, + usage.candidates_token_count, + usage.total_token_count, + ); + + // 6. Disconnect + conversation.disconnect().await?; + Ok(()) +} +``` + +> **Note:** When using `LocalConnectionStrategy` directly, you bypass the `Agent`'s +> automatic binary resolution, policy enforcement, hook registration, and trigger startup. +> This is useful for embedding the SDK in custom orchestration frameworks. + +## Python SDK Comparison + +| Concept | Python SDK | Rust SDK | +|---------|-----------|----------| +| Connection trait | `Connection` ABC | `trait Connection: Send + Sync` | +| Dispatch enum | Runtime dispatch | `AnyConnection` compile-time `#[cfg]` dispatch | +| Local strategy | `LocalConnectionStrategy` class | `LocalConnectionStrategy` struct | +| WASM strategy | N/A | `WasmConnectionStrategy` (wasm32 target) | +| Subprocess init | Length-prefixed protobuf | Same protocol, `prost` + `tokio::process` | +| WebSocket lib | `websockets` (Python) | `tokio-tungstenite` / `tungstenite` | +| Idle detection | Event-based | `AtomicBool` + `TrajectoryStateUpdate` sentinel pattern | +| Mock for tests | `unittest.mock` | `MockConnection` with `#[cfg(test)]` | + +## Architecture Diagram + +```text +┌─────────────────────────┐ +│ Agent::start() │ +│ (resolves binary, │ +│ wires policies/tools) │ +└──────────┬──────────────┘ + │ + ▼ +┌─────────────────────────┐ +│ ConnectionStrategy │ +│ ├─ Local (native) │ +│ └─ Wasm (wasm32) │ +└──────────┬──────────────┘ + │ .connect() + ▼ +┌─────────────────────────┐ +│ AnyConnection │ +│ (enum dispatch) │ +└──────────┬──────────────┘ + │ + ▼ +┌─────────────────────────┐ +│ Conversation │ +│ (state tracking, │ +│ streaming, usage) │ +└─────────────────────────┘ +``` diff --git a/docs/conversation.md b/docs/conversation.md new file mode 100644 index 0000000..a0f3ad0 --- /dev/null +++ b/docs/conversation.md @@ -0,0 +1,408 @@ +# Conversation + +Stateful session management for the Antigravity Rust SDK. + +## Overview + +`Conversation` is a stateful wrapper around a [`Connection`](./connections.md) that tracks +step history, accumulates token usage metadata, and processes stream chunks. It is the +primary interface for interacting with an active agent session — whether you want streaming +chunk-by-chunk output or a simple blocking call-and-response. + +Most users access `Conversation` through `Agent`: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let agent = Agent::builder() + .allow_all() + .build() + .start().await?; + + // Get the active Conversation (Arc) + let conversation = agent.conversation(); + + // Use it for streaming or blocking chat + let response = conversation.chat_to_completion("Hello!").await?; + println!("{}", response.text); + + agent.stop().await +} +``` + +## Creating a Conversation + +`Conversation` wraps an `AnyConnection` and an optional history size limit: + +```rust,no_run +use antigravity_sdk_rust::conversation::Conversation; +use antigravity_sdk_rust::connection::AnyConnection; + +// Created internally by Agent::start(), but can be constructed directly: +let conversation = Conversation::new( + any_connection, // AnyConnection (Local, Wasm, or Mock) + Some(5000), // max_history_size (None → default 10,000; Some(0) → unlimited) +); +``` + +### Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `conn` | `AnyConnection` | — | The underlying connection to the harness | +| `max_history_size` | `Option` | `Some(10_000)` | Max steps retained in memory. `None` uses 10,000. `Some(0)` disables trimming. | + +## Key Methods + +### Chatting + +#### `chat(prompt)` — Streaming + +Sends a prompt and returns a `BoxStream` for real-time processing: + +```rust,no_run +use antigravity_sdk_rust::types::StreamChunk; +use futures_util::StreamExt; + +let mut stream = conversation.chat("Explain monads").await?; + +while let Some(chunk_res) = stream.next().await { + match chunk_res? { + StreamChunk::Thought { text, step_index } => { + // Model's internal reasoning (thinking tokens) + eprint!("{}", text); + } + StreamChunk::Text { text, step_index } => { + // Response text fragments + print!("{}", text); + } + StreamChunk::ToolCall(call) => { + // The model requested a tool execution + println!("[Tool: {} | Args: {}]", call.name, call.args); + } + } +} +``` + +**Signature:** +```rust,no_run +pub async fn chat( + &self, + prompt: &str, +) -> Result>, anyhow::Error> +``` + +#### `chat_to_completion(prompt)` — Blocking + +Sends a prompt and waits for the complete response, accumulating all chunks internally: + +```rust,no_run +let response = conversation.chat_to_completion("What is 2 + 2?").await?; + +println!("Response: {}", response.text); +println!("Thinking: {}", response.thinking); +println!("Steps: {}", response.steps.len()); +println!("Total tokens: {}", response.usage_metadata.total_token_count); +``` + +**Signature:** +```rust,no_run +pub async fn chat_to_completion( + &self, + prompt: &str, +) -> Result +``` + +### Sending & Receiving (Low-Level) + +| Method | Signature | Description | +|--------|-----------|-------------| +| `send(prompt)` | `async fn send(&self, prompt: &str) -> Result<()>` | Sends a raw prompt and registers a turn boundary. Does **not** return a stream. | +| `receive_steps()` | `fn receive_steps(&self) -> BoxStream<'static, Result>` | Raw step-level stream. Inserts steps into history, tracks compaction, enforces history limits. | +| `receive_chunks()` | `fn receive_chunks(&self) -> BoxStream<'static, Result>` | Filters `receive_steps()` into high-level `StreamChunk` events. Only emits model→user content. | + +#### `receive_chunks()` Filtering Logic + +The chunk stream applies these rules: + +1. **Thought chunks** — emitted when `source == Model`, `target == User`, and `thinking_delta` is non-empty +2. **Text chunks** — emitted when `source == Model`, `target == User`, and `content_delta` is non-empty +3. **ToolCall chunks** — emitted for each `ToolCall` in the step, **deduplicated by `call.id`** (empty IDs are never deduplicated) +4. **Environment-targeted steps** are silently filtered out (e.g., harness internal tool dispatches) + +### Connection Access + +| Method | Return Type | Description | +|--------|-------------|-------------| +| `connection()` | `AnyConnection` | Returns the underlying connection (clone of `Arc`) | +| `conversation_id()` | `&str` | Session ID (delegates to `Connection::conversation_id()`) | +| `is_idle()` | `bool` | Whether the connection is idle | +| `disconnect()` | `async -> Result<()>` | Gracefully closes the connection | + +## StreamChunk + +The streaming fragment enum used by `chat()` and `receive_chunks()`: + +```rust,no_run +use antigravity_sdk_rust::types::ToolCall; + +#[derive(Debug, Clone)] +pub enum StreamChunk { + /// Model's internal reasoning/thinking fragment. + Thought { + /// Index of the step that produced this chunk. + step_index: u32, + /// Thinking text delta. + text: String, + }, + + /// Response text fragment directed at the user. + Text { + /// Index of the step that produced this chunk. + step_index: u32, + /// Text content delta. + text: String, + }, + + /// A complete tool call requested by the model. + ToolCall(ToolCall), +} +``` + +### ToolCall Structure + +```rust,no_run +#[derive(Debug, Clone)] +pub struct ToolCall { + /// Unique correlation ID for this call. + pub id: String, + /// Name of the tool to invoke. + pub name: String, + /// Arguments as a JSON value. + pub args: serde_json::Value, + /// Optional canonical filesystem path (for file-targeting tools). + pub canonical_path: Option, +} +``` + +## Streaming Example + +The complete streaming pattern from `examples/streaming.rs`: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::types::StreamChunk; +use futures_util::StreamExt; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + dotenvy::dotenv().ok(); + + let agent = Agent::builder() + .default_model("gemini-3.5-flash") + .allow_all() + .build() + .start().await?; + + let conversation = agent.conversation(); + + let prompt = "Solve this riddle: I speak without a mouth \ + and hear without ears. What am I?"; + println!("User: {}\n", prompt); + + let mut stream = conversation.chat(prompt).await?; + + while let Some(chunk_res) = stream.next().await { + match chunk_res? { + StreamChunk::Thought { text, .. } => { + // Print reasoning tokens as they arrive + print!("{}", text); + std::io::Write::flush(&mut std::io::stdout())?; + } + StreamChunk::Text { text, .. } => { + // Print response text tokens as they arrive + print!("{}", text); + std::io::Write::flush(&mut std::io::stdout())?; + } + StreamChunk::ToolCall(call) => { + println!( + "\n[Tool Call: {} with args: {}]", + call.name, call.args + ); + } + } + } + + println!(); // Final newline + agent.stop().await?; + Ok(()) +} +``` + +## State Management + +`Conversation` maintains a `ConversationState` protected by an async `Mutex`: + +```rust,no_run +pub struct ConversationState { + /// All executed steps (prompts, tool calls, responses, compaction markers). + pub steps: Vec, + /// Step indices marking each user prompt turn boundary. + pub turn_start_indices: Vec, + /// Step indices where history compaction occurred. + pub compaction_indices: Vec, + /// Cumulative token usage across all turns. + pub cumulative_usage: UsageMetadata, + /// Token usage for the current active turn. + pub turn_usage: Option, +} +``` + +### State Query Methods + +All state methods are `async` because they acquire the internal `Mutex`: + +| Method | Return Type | Description | +|--------|-------------|-------------| +| `history()` | `Vec` | Clone of the full step history | +| `turn_count()` | `usize` | Number of user-initiated turns | +| `compaction_indices()` | `Vec` | Where the harness compacted history | +| `last_response()` | `String` | Text content of the last complete model response | +| `total_usage()` | `UsageMetadata` | Cumulative token usage for the entire session | +| `last_turn_usage()` | `Option` | Token usage for the most recent turn | +| `clear_history()` | `()` | Resets all steps, turns, compactions, and usage stats | + +### Usage Example + +```rust,no_run +// After a chat interaction +let history = conversation.history().await; +println!("Total steps: {}", history.len()); +println!("Turns completed: {}", conversation.turn_count().await); + +let total = conversation.total_usage().await; +println!("Session tokens: {} prompt, {} generated, {} total", + total.prompt_token_count, + total.candidates_token_count, + total.total_token_count, +); + +if let Some(turn) = conversation.last_turn_usage().await { + println!("Last turn: {} thinking tokens", turn.thoughts_token_count); +} + +// Check if compaction happened +let compactions = conversation.compaction_indices().await; +if !compactions.is_empty() { + println!("History was compacted at step indices: {:?}", compactions); +} +``` + +### History Auto-Trimming + +When `max_history_size > 0` and the step count exceeds the limit, older steps are +drained from the front. Turn and compaction indices are adjusted accordingly: + +```rust,no_run +// Only keep the last 100 steps in memory +let conversation = Conversation::new(any_connection, Some(100)); + +// Disable auto-trimming entirely +let conversation = Conversation::new(any_connection, Some(0)); + +// Use the default (10,000 steps) +let conversation = Conversation::new(any_connection, None); +``` + +## ChatResponse + +The complete result returned by `chat_to_completion()`: + +```rust,no_run +#[derive(Debug, Clone)] +pub struct ChatResponse { + /// Combined final response text. + pub text: String, + /// Combined model reasoning/thinking text. + pub thinking: String, + /// All steps executed during this turn. + pub steps: Vec, + /// Cumulative token usage metrics. + pub usage_metadata: UsageMetadata, +} +``` + +### UsageMetadata + +```rust,no_run +#[derive(Debug, Clone, Default)] +pub struct UsageMetadata { + /// Tokens included in the request prompt. + pub prompt_token_count: i32, + /// Tokens generated in model candidates. + pub candidates_token_count: i32, + /// Total combined tokens. + pub total_token_count: i32, + /// Cache-hit content tokens. + pub cached_content_token_count: i32, + /// Tokens consumed during reasoning/thinking. + pub thoughts_token_count: i32, +} +``` + +## Step Model + +Each event in the trajectory is represented as a `Step`: + +```rust,no_run +#[derive(Debug, Clone)] +pub struct Step { + pub id: String, // Unique step ID + pub step_index: u32, // Position in trajectory + pub r#type: StepType, // TextResponse, ToolCall, Compaction, Finish, etc. + pub source: StepSource, // System, User, Model + pub target: StepTarget, // User, Environment + pub status: StepStatus, // Active, Done, Error, TerminalError, etc. + pub content: String, // Full accumulated text + pub content_delta: String, // Incremental text delta + pub thinking: String, // Full accumulated thinking + pub thinking_delta: String, // Incremental thinking delta + pub tool_calls: Vec, // Tool calls in this step + pub error: String, // Error message (if any) + pub is_complete_response: Option, // True for final model response + pub structured_output: Option, // Parsed structured output + pub usage_metadata: Option, // Per-step token usage + pub cascade_id: String, // Execution cascade group + pub trajectory_id: String, // Sub-agent trajectory group + pub http_code: u32, // HTTP status code (if applicable) +} +``` + +### Step Enums + +| Enum | Variants | +|------|----------| +| `StepType` | `TextResponse`, `ToolCall`, `SystemMessage`, `Compaction`, `Finish`, `Unknown` | +| `StepSource` | `System`, `User`, `Model`, `Unknown` | +| `StepTarget` | `User`, `Environment`, `Unspecified`, `Unknown` | +| `StepStatus` | `Active`, `Done`, `WaitingForUser`, `Error`, `Canceled`, `TerminalError`, `Unknown` | + +## Python SDK Comparison + +| Concept | Python SDK | Rust SDK | +|---------|-----------|----------| +| Conversation class | `Conversation` | `Conversation` (same name) | +| Streaming | `async for chunk in conversation.chat(prompt)` | `conversation.chat(prompt).await?.next().await` | +| Blocking | `conversation.chat_to_completion(prompt)` | `conversation.chat_to_completion(prompt).await?` | +| History | `conversation.history` (property) | `conversation.history().await` (async method) | +| Turn count | `conversation.turn_count` | `conversation.turn_count().await` | +| Token usage | `conversation.total_usage` | `conversation.total_usage().await` | +| State mutex | GIL + threading lock | `tokio::sync::Mutex` | +| Stream type | `AsyncIterator[StreamChunk]` | `BoxStream<'static, Result>` | +| Max history | `max_history_size` param | `max_history_size: Option` (default 10,000) | + +> **Key difference:** In the Rust SDK, all state-querying methods are `async` because the +> internal state is protected by a `tokio::sync::Mutex`. In Python, these are synchronous +> properties protected by the GIL. diff --git a/docs/hooks.md b/docs/hooks.md new file mode 100644 index 0000000..6c3dbe3 --- /dev/null +++ b/docs/hooks.md @@ -0,0 +1,608 @@ +# Hooks + +Agent lifecycle interception and policies. + +## Overview + +Hooks observe or modify agent behavior at key lifecycle points. They let you +log events, enforce safety rules, rate-limit turns, audit tool calls, and +recover from errors — all without modifying the core agent loop. + +### Python SDK vs Rust SDK + +The Python SDK splits hooks into separate base classes by category: + +| Python Category | Python Base Classes | Rust Equivalent | +|---|---|---| +| **Inspect** (read-only) | `OnSessionStartHook`, `PostToolCallHook`, `OnSessionEndHook`, `PostTurnHook`, `OnCompactionHook` | Default no-op methods on `Hook` | +| **Decide** (blocking) | `PreTurnHook`, `PreToolCallDecideHook` | `pre_turn()`, `pre_tool_call()` return `HookResult` | +| **Transform** (modifying) | `OnToolErrorHook`, `OnInteractionHook` | `on_tool_error()`, `on_interaction()` return recovery data | + +The Rust SDK merges all of these into a **single `Hook` trait** with 9 async +methods. Every method has a default no-op implementation, so you only override +what you need. + +--- + +## Hook Trait + +All 9 lifecycle methods are defined on the `Hook` trait. Each has a default +no-op implementation so you only need to override the methods relevant to your +use case: + +```rust,no_run +use antigravity_sdk_rust::hooks::Hook; +use antigravity_sdk_rust::types::{ + AskQuestionEntry, ChatResponse, HookResult, QuestionHookResult, + ToolCall, ToolResult, +}; + +pub trait Hook: Send + Sync { + // ── Session lifecycle ────────────────────────────────────────── + + /// Called when the agent establishes a connection and starts a session. + async fn on_session_start(&self) -> Result<(), anyhow::Error> { + Ok(()) + } + + /// Called when the session is ending (agent shutdown or disconnect). + async fn on_session_end(&self) -> Result<(), anyhow::Error> { + Ok(()) + } + + // ── Turn lifecycle ───────────────────────────────────────────── + + /// Intercepts the start of a user turn before the LLM processes the prompt. + /// Return `allow: false` to halt execution. + async fn pre_turn(&self) -> Result { + Ok(HookResult { allow: true, message: String::new() }) + } + + /// Called after a turn completes, receiving the full response. + async fn post_turn(&self, _response: &ChatResponse) -> Result<(), anyhow::Error> { + Ok(()) + } + + // ── Tool lifecycle ───────────────────────────────────────────── + + /// Intercepts a tool call before execution. + /// Return `allow: false` to prevent the tool from running. + async fn pre_tool_call(&self, _tool_call: &ToolCall) -> Result { + Ok(HookResult { allow: true, message: String::new() }) + } + + /// Called after a tool successfully returns a result. + async fn post_tool_call(&self, _result: &ToolResult) -> Result<(), anyhow::Error> { + Ok(()) + } + + // ── Error recovery ───────────────────────────────────────────── + + /// Called when a tool execution encounters an error. + /// Return `(HookResult { allow: true, .. }, Some(value))` to provide a + /// recovery payload instead of propagating the error. + async fn on_tool_error( + &self, + error: &anyhow::Error, + ) -> Result<(HookResult, Option), anyhow::Error> { + Ok(( + HookResult { allow: false, message: error.to_string() }, + None, + )) + } + + // ── User interaction ─────────────────────────────────────────── + + /// Intercepts interactive user questions. + /// Return `Some(result)` to answer programmatically. + async fn on_interaction( + &self, + _questions: &[AskQuestionEntry], + ) -> Result, anyhow::Error> { + Ok(None) + } + + // ── History compaction ───────────────────────────────────────── + + /// Called when the conversation history is compacted/summarized. + async fn on_compaction(&self, _summary: &str) -> Result<(), anyhow::Error> { + Ok(()) + } +} +``` + +### Method Reference + +| Method | Category | Return | Short-Circuit? | +|---|---|---|---| +| `on_session_start` | Inspect | `()` | No | +| `on_session_end` | Inspect | `()` | No | +| `pre_turn` | Decide | `HookResult` | Yes — first `allow: false` | +| `post_turn` | Inspect | `()` | No | +| `pre_tool_call` | Decide | `HookResult` | Yes — first `allow: false` | +| `post_tool_call` | Inspect | `()` | No | +| `on_tool_error` | Transform | `(HookResult, Option)` | Yes — first `allow: true` | +| `on_interaction` | Transform | `Option` | Yes — first `Some(result)` | +| `on_compaction` | Inspect | `()` | No | + +--- + +## DynHook (Object-Safe Wrapper) + +Rust's `async fn` in traits produces `impl Future` return types, which are not +object-safe — you cannot use `dyn Hook` directly. The SDK solves this with a +companion `DynHook` trait that wraps every method in a `BoxFuture`: + +```rust,no_run +use futures_util::future::BoxFuture; +use antigravity_sdk_rust::types::{HookResult, ToolCall, ToolResult, ChatResponse, AskQuestionEntry, QuestionHookResult}; + +/// Object-safe version of `Hook`, used internally for dynamic dispatch. +pub trait DynHook: Send + Sync { + fn on_session_start(&self) -> BoxFuture<'_, Result<(), anyhow::Error>>; + fn pre_turn(&self) -> BoxFuture<'_, Result>; + fn pre_tool_call<'a>(&'a self, tool_call: &'a ToolCall) -> BoxFuture<'a, Result>; + fn post_tool_call<'a>(&'a self, result: &'a ToolResult) -> BoxFuture<'a, Result<(), anyhow::Error>>; + fn on_tool_error<'a>(&'a self, error: &'a anyhow::Error) -> BoxFuture<'a, Result<(HookResult, Option), anyhow::Error>>; + fn on_interaction<'a>(&'a self, questions: &'a [AskQuestionEntry]) -> BoxFuture<'a, Result, anyhow::Error>>; + fn on_session_end(&self) -> BoxFuture<'_, Result<(), anyhow::Error>>; + fn post_turn<'a>(&'a self, response: &'a ChatResponse) -> BoxFuture<'a, Result<(), anyhow::Error>>; + fn on_compaction<'a>(&'a self, summary: &'a str) -> BoxFuture<'a, Result<(), anyhow::Error>>; +} +``` + +A **blanket implementation** automatically bridges the two traits: + +```rust,no_run +// Any type that implements Hook also implements DynHook — for free. +// impl DynHook for T { ... } +``` + +This means you always implement `Hook` (the ergonomic trait) and the SDK +converts it to `Arc` for storage and dispatch. You never need to +implement `DynHook` directly. + +--- + +## HookRunner + +The `HookRunner` is the internal dispatch engine that manages registered hooks +and dispatches lifecycle events sequentially. It stores hooks as +`Vec>` behind a `tokio::sync::RwLock`. + +### Dispatch Behavior + +| Dispatch Method | Behavior | +|---|---| +| `dispatch_session_start` | Calls all hooks sequentially | +| `dispatch_session_end` | Calls all hooks sequentially | +| `dispatch_pre_turn` | **Short-circuits** at first `allow: false` | +| `dispatch_post_turn` | Calls all hooks sequentially | +| `dispatch_pre_tool_call` | **Short-circuits** at first `allow: false` | +| `dispatch_post_tool_call` | Calls all hooks sequentially | +| `dispatch_on_tool_error` | **Short-circuits** at first `allow: true` (recovery) | +| `dispatch_interaction` | **Short-circuits** at first `Some(result)` | +| `dispatch_on_compaction` | Calls all hooks sequentially | + +Short-circuiting means that once a decisive result is found, remaining hooks +are skipped for that event. This allows safety hooks to block execution +without later hooks overriding the decision. + +### Registration + +```rust,no_run +use antigravity_sdk_rust::hooks::{HookRunner, DynHook}; +use std::sync::Arc; + +let runner = HookRunner::new(); + +// Register a hook +// runner.register(Arc::new(my_hook)).await; +``` + +--- + +## HookContext (Hierarchical State) + +`HookContext` provides a parent-chaining key-value store for sharing state +across hook invocations at different scopes. Values are stored as +`serde_json::Value` internally. + +### 3-Level Hierarchy + +```text +┌─────────────────────────────────┐ +│ Session Context (root) │ ← HookContext::new() +│ "session_id" = "abc-123" │ +│ │ +│ ┌───────────────────────────┐ │ +│ │ Turn Context │ │ ← HookContext::child(session) +│ │ "turn_count" = 5 │ │ +│ │ │ │ +│ │ ┌─────────────────────┐ │ │ +│ │ │ Operation Context │ │ │ ← HookContext::child(turn) +│ │ │ "tool_name" = "X" │ │ │ +│ │ └─────────────────────┘ │ │ +│ └───────────────────────────┘ │ +└─────────────────────────────────┘ +``` + +### Key Semantics + +- **`get(key)`** — searches the local store first, then walks up the parent + chain. Returns `None` if the key is not found at any level. +- **`set(key, value)`** — writes only to the **local** store. This shadows + (but does not mutate) any parent value with the same key. +- **`has_parent()`** — returns `true` if this context is not the root. + +### Usage + +```rust,no_run +use antigravity_sdk_rust::context::HookContext; +use std::sync::Arc; + +// Session-level context (root) +let session = Arc::new(HookContext::new()); +session.set("user_id", "user-42"); + +// Turn-level context +let turn = Arc::new(HookContext::child(session.clone())); +turn.set("turn_number", 1u32); + +// Operation-level context +let operation = HookContext::child(turn.clone()); +operation.set("tool_name", "read_file"); + +// get() walks up the chain: +assert_eq!(operation.get::("tool_name"), Some("read_file".to_string())); +assert_eq!(operation.get::("turn_number"), Some(1)); +assert_eq!(operation.get::("user_id"), Some("user-42".to_string())); + +// set() is local-only — parent is unchanged: +operation.set("user_id", "override"); +assert_eq!(operation.get::("user_id"), Some("override".to_string())); +assert_eq!(session.get::("user_id"), Some("user-42".to_string())); +``` + +--- + +## PolicyEnforcer + +`PolicyEnforcer` is a built-in `Hook` implementation that enforces safety +policies on tool calls. It intercepts `pre_tool_call` and evaluates registered +policies against a 9-bucket priority system. + +### Priority Bucket System + +Policies are sorted into 9 buckets organized by **specificity** (3 levels) × +**decision type** (3 types). Lower bucket index = higher priority: + +| Bucket | Specificity | Decision | Index | +|---|---|---|---| +| Specific Deny | Exact tool name | `Deny` | 0 (highest) | +| Specific Ask | Exact tool name | `AskUser` | 1 | +| Specific Allow | Exact tool name | `Approve` | 2 | +| Prefix Deny | Server wildcard (`server/*`) | `Deny` | 3 | +| Prefix Ask | Server wildcard | `AskUser` | 4 | +| Prefix Allow | Server wildcard | `Approve` | 5 | +| Global Deny | Global wildcard (`*`) | `Deny` | 6 | +| Global Ask | Global wildcard | `AskUser` | 7 | +| Global Allow | Global wildcard | `Approve` | 8 (lowest) | + +**Key rules:** +- Specific rules always override prefix and global rules. +- Prefix rules (`server/*`) override global wildcards (`*`). +- At the same specificity, **Deny beats Ask beats Allow**. +- If no policy matches, the tool call is **allowed** (open by default). + +### Policy Helper Functions + +```rust,no_run +use antigravity_sdk_rust::policy; + +// ── Single-tool policies ─────────────────────────────────────────── +let _ = policy::allow("read_file"); // Approve a specific tool +let _ = policy::deny("run_command"); // Deny a specific tool +let _ = policy::ask_user("run_command", |_tc| { + // Return true = user approved, false = user denied + true +}); + +// ── Wildcard policies ────────────────────────────────────────────── +let _ = policy::allow_all(); // Approve all tools (wildcard) +let _ = policy::deny_all(); // Deny all tools (wildcard) + +// ── Composite policies ──────────────────────────────────────────── +// Require user confirmation for RUN_COMMAND, allow everything else: +let _ = policy::confirm_run_command(Some(std::sync::Arc::new(|_| true))); + +// Restrict filesystem tools to specific workspace directories: +let _ = policy::workspace_only(vec!["/my/project".to_string()]); +``` + +### MCP Server Policies + +For controlling tools exposed by MCP (Model Context Protocol) servers: + +```rust,no_run +use antigravity_sdk_rust::policy; +use antigravity_sdk_rust::types::McpServerConfig; + +let server = McpServerConfig::Stdio { + name: "math".to_string(), + command: "math-server".to_string(), + args: vec![], + enabled_tools: None, + disabled_tools: None, +}; + +// Allow all tools from the server +let _ = policy::allow_mcp(&server, None); + +// Deny specific tools from the server +let _ = policy::deny_mcp(&server, Some(&["dangerous_calc"])); + +// Ask user for specific MCP tools +let _ = policy::ask_user_mcp(&server, Some(&["execute"]), |_| true); +``` + +### Conditional Policies with `when()` + +Policies can include a predicate that narrows when they apply: + +```rust,no_run +use antigravity_sdk_rust::policy; + +// Only deny run_command when the command contains "rm" +let _ = policy::deny("run_command").when(|tc| { + tc.args + .get("CommandLine") + .and_then(|v| v.as_str()) + .map_or(false, |s| s.contains("rm")) +}); +``` + +### Creating a PolicyEnforcer + +Use `policy::enforce()` to validate and compile policies: + +```rust,no_run +use antigravity_sdk_rust::policy; + +// Without MCP servers +let enforcer = policy::enforce( + vec![ + policy::deny("run_command"), + policy::allow_all(), + ], + None, // no MCP servers +).expect("policy validation failed"); + +// enforce() will return an error if: +// - An AskUser policy is missing its handler callback +// - MCP policies are present but no MCP servers are registered (fail-closed) +``` + +--- + +## Examples + +### 1. Logging Hook + +A simple hook that logs all lifecycle events: + +```rust,no_run +use antigravity_sdk_rust::hooks::Hook; +use antigravity_sdk_rust::types::{ChatResponse, HookResult, ToolCall, ToolResult}; + +struct LoggingHook; + +impl Hook for LoggingHook { + async fn on_session_start(&self) -> Result<(), anyhow::Error> { + println!("🟢 Session started"); + Ok(()) + } + + async fn on_session_end(&self) -> Result<(), anyhow::Error> { + println!("🔴 Session ended"); + Ok(()) + } + + async fn pre_tool_call(&self, tool_call: &ToolCall) -> Result { + println!("🔧 Calling tool: {}", tool_call.name); + Ok(HookResult { allow: true, message: String::new() }) + } + + async fn post_tool_call(&self, result: &ToolResult) -> Result<(), anyhow::Error> { + println!("✅ Tool {} completed", result.name); + Ok(()) + } + + async fn post_turn(&self, response: &ChatResponse) -> Result<(), anyhow::Error> { + println!("💬 Response length: {} chars", response.text.len()); + Ok(()) + } +} +``` + +### 2. Rate-Limiting Hook + +A `pre_turn` hook that enforces a maximum number of turns per session: + +```rust,no_run +use antigravity_sdk_rust::hooks::Hook; +use antigravity_sdk_rust::types::HookResult; +use std::sync::atomic::{AtomicU32, Ordering}; + +struct RateLimitHook { + max_turns: u32, + turn_count: AtomicU32, +} + +impl RateLimitHook { + fn new(max_turns: u32) -> Self { + Self { + max_turns, + turn_count: AtomicU32::new(0), + } + } +} + +impl Hook for RateLimitHook { + async fn pre_turn(&self) -> Result { + let count = self.turn_count.fetch_add(1, Ordering::SeqCst); + if count >= self.max_turns { + Ok(HookResult { + allow: false, + message: format!( + "Rate limit exceeded: {} of {} turns used", + count, self.max_turns + ), + }) + } else { + Ok(HookResult { allow: true, message: String::new() }) + } + } +} +``` + +### 3. Tool Audit Hook + +A `post_tool_call` hook that records all tool invocations for auditing: + +```rust,no_run +use antigravity_sdk_rust::hooks::Hook; +use antigravity_sdk_rust::types::{HookResult, ToolCall, ToolResult}; +use std::sync::Mutex; + +struct AuditRecord { + tool_name: String, + success: bool, + timestamp: std::time::Instant, +} + +struct AuditHook { + records: Mutex>, +} + +impl AuditHook { + fn new() -> Self { + Self { + records: Mutex::new(Vec::new()), + } + } +} + +impl Hook for AuditHook { + async fn pre_tool_call(&self, tool_call: &ToolCall) -> Result { + println!("📝 Audit: tool '{}' invoked", tool_call.name); + Ok(HookResult { allow: true, message: String::new() }) + } + + async fn post_tool_call(&self, result: &ToolResult) -> Result<(), anyhow::Error> { + let record = AuditRecord { + tool_name: result.name.clone(), + success: result.error.is_none(), + timestamp: std::time::Instant::now(), + }; + if let Ok(mut records) = self.records.lock() { + records.push(record); + } + Ok(()) + } +} +``` + +### 4. Error Recovery Hook + +An `on_tool_error` hook that provides fallback values when specific tools fail: + +```rust,no_run +use antigravity_sdk_rust::hooks::Hook; +use antigravity_sdk_rust::types::HookResult; + +struct ErrorRecoveryHook; + +impl Hook for ErrorRecoveryHook { + async fn on_tool_error( + &self, + error: &anyhow::Error, + ) -> Result<(HookResult, Option), anyhow::Error> { + let msg = error.to_string(); + + // Provide a fallback value for network-related errors + if msg.contains("timeout") || msg.contains("connection refused") { + println!("⚠️ Network error detected, providing fallback"); + Ok(( + HookResult { + allow: true, + message: "Recovered from network error with fallback".to_string(), + }, + Some(serde_json::json!({ + "error": "network_unavailable", + "fallback": true, + "message": "Service temporarily unavailable, using cached data" + })), + )) + } else { + // Let other hooks handle it, or propagate the error + Ok(( + HookResult { + allow: false, + message: msg, + }, + None, + )) + } + } +} +``` + +### 5. Registering Hooks with the Agent Builder + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::hooks::DynHook; +use std::sync::Arc; + +// Assuming LoggingHook and RateLimitHook are defined as above + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let agent = Agent::builder() + // Register hooks via the builder + .hook(Arc::new(LoggingHook) as Arc) + // .hook(Arc::new(RateLimitHook::new(100)) as Arc) + .allow_all() + .build(); + + let agent = agent.start().await?; + let response = agent.chat("Hello!").await?; + println!("{}", response.text); + agent.stop().await?; + Ok(()) +} + +struct LoggingHook; +impl antigravity_sdk_rust::hooks::Hook for LoggingHook { + async fn on_session_start(&self) -> Result<(), anyhow::Error> { + println!("Session started"); + Ok(()) + } +} +``` + +You can also register hooks after construction but before starting: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::hooks::DynHook; +use std::sync::Arc; + +struct MyHook; +impl antigravity_sdk_rust::hooks::Hook for MyHook {} + +let mut agent = Agent::builder().allow_all().build(); +agent.register_hook(Arc::new(MyHook) as Arc); +// agent.start().await?; +``` diff --git a/docs/mcp.md b/docs/mcp.md new file mode 100644 index 0000000..225698c --- /dev/null +++ b/docs/mcp.md @@ -0,0 +1,289 @@ +# MCP + +Model Context Protocol integration for the Antigravity Rust SDK. + +## Overview + +MCP (Model Context Protocol) allows the agent to connect to external tool servers. The SDK configures MCP servers which are managed by the underlying localharness runtime — the SDK passes configuration, and the harness handles client sessions, tool discovery, and execution. + +## McpServerConfig + +Three transport variants are supported: + +### Stdio + +Launch a local MCP server as a subprocess: + +```rust,no_run +use antigravity_sdk_rust::types::McpServerConfig; + +let server = McpServerConfig::Stdio { + name: "my_server".to_string(), + command: "npx".to_string(), + args: vec!["my-mcp-server".to_string()], + enabled_tools: None, + disabled_tools: None, +}; +``` + +Use Stdio for local MCP servers distributed as npm packages, Python packages, or standalone binaries. + +### SSE (Server-Sent Events) + +Connect to a remote MCP server via SSE: + +```rust,no_run +use antigravity_sdk_rust::types::McpServerConfig; + +let server = McpServerConfig::Sse { + name: "remote_server".to_string(), + url: "https://my-server.example.com/sse".to_string(), + headers: None, + enabled_tools: None, + disabled_tools: None, +}; +``` + +### HTTP + +Connect via standard HTTP with configurable timeouts: + +```rust,no_run +use antigravity_sdk_rust::types::McpServerConfig; +use std::collections::HashMap; + +let server = McpServerConfig::Http { + name: "http_server".to_string(), + url: "https://my-server.example.com/mcp".to_string(), + headers: Some(HashMap::from([ + ("Authorization".to_string(), "Bearer my-token".to_string()), + ])), + timeout: 30.0, // Connection timeout (seconds) + sse_read_timeout: 300.0, // SSE read timeout (seconds) + terminate_on_close: true, // Terminate channel on close + enabled_tools: None, + disabled_tools: None, +}; +``` + +## Tool Filtering + +Each transport variant supports fine-grained tool control: + +- **`enabled_tools`**: Allowlist — only these tools are exposed to the model +- **`disabled_tools`**: Denylist — these tools are hidden from the model + +These are mutually exclusive. When both are `None`, all tools from the server are available. + +```rust,no_run +use antigravity_sdk_rust::types::McpServerConfig; + +// Only expose specific tools +let server = McpServerConfig::Stdio { + name: "fs_server".to_string(), + command: "npx".to_string(), + args: vec!["@anthropic/mcp-fs-server".to_string()], + enabled_tools: Some(vec!["read_file".to_string(), "list_dir".to_string()]), + disabled_tools: None, +}; + +// Or disable specific tools +let server = McpServerConfig::Stdio { + name: "fs_server".to_string(), + command: "npx".to_string(), + args: vec!["@anthropic/mcp-fs-server".to_string()], + enabled_tools: None, + disabled_tools: Some(vec!["delete_file".to_string()]), +}; +``` + +## Agent Builder Integration + +### Single Server + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::types::McpServerConfig; + +let agent = Agent::builder() + .mcp_server(McpServerConfig::Stdio { + name: "my_server".to_string(), + command: "npx".to_string(), + args: vec!["my-mcp-server".to_string()], + enabled_tools: None, + disabled_tools: None, + }) + .allow_all() + .build(); +``` + +### Multiple Servers + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::types::McpServerConfig; + +let agent = Agent::builder() + .mcp_servers(vec![ + McpServerConfig::Stdio { + name: "fs".to_string(), + command: "npx".to_string(), + args: vec!["@anthropic/mcp-fs-server".to_string()], + enabled_tools: None, + disabled_tools: None, + }, + McpServerConfig::Sse { + name: "api".to_string(), + url: "https://api.example.com/mcp/sse".to_string(), + headers: None, + enabled_tools: None, + disabled_tools: None, + }, + ]) + .allow_all() + .build(); +``` + +### Chaining + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::types::McpServerConfig; + +let agent = Agent::builder() + .mcp_server(McpServerConfig::Stdio { + name: "server_a".to_string(), + command: "npx".to_string(), + args: vec!["server-a".to_string()], + enabled_tools: None, + disabled_tools: None, + }) + .mcp_server(McpServerConfig::Stdio { + name: "server_b".to_string(), + command: "npx".to_string(), + args: vec!["server-b".to_string()], + enabled_tools: None, + disabled_tools: None, + }) + .allow_all() + .build(); +``` + +## Policy Integration + +MCP tools are named `{server_name}_{tool_name}` in the policy system. The SDK provides helpers: + +```rust,no_run +use antigravity_sdk_rust::policy; +use antigravity_sdk_rust::types::McpServerConfig; + +let server = McpServerConfig::Stdio { + name: "fs".to_string(), + command: "npx".to_string(), + args: vec!["mcp-fs-server".to_string()], + enabled_tools: None, + disabled_tools: None, +}; + +// Allow all tools from this server +let policies = policy::allow_mcp(&server, None); + +// Allow only specific tools +let policies = policy::allow_mcp(&server, Some(&["read_file", "list_dir"])); + +// Deny specific tools +let policies = policy::deny_mcp(&server, Some(&["delete_file"])); + +// Require user confirmation for all tools +let policies = policy::ask_user_mcp(&server, None); +``` + +### Combined with Agent Builder + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::policy; +use antigravity_sdk_rust::types::McpServerConfig; + +let fs_server = McpServerConfig::Stdio { + name: "fs".to_string(), + command: "npx".to_string(), + args: vec!["mcp-fs-server".to_string()], + enabled_tools: None, + disabled_tools: None, +}; + +let agent = Agent::builder() + .mcp_server(fs_server.clone()) + .policies(vec![ + policy::deny_all(), // Block everything by default + policy::allow("VIEW_FILE"), // Allow built-in VIEW_FILE + policy::allow_mcp(&fs_server, Some(&["read"])), // Allow MCP fs.read + ]) + .build(); +``` + +## Architecture + +``` +┌─────────────────────────────────────────────┐ +│ Your Rust Application │ +│ │ +│ Agent::builder() │ +│ .mcp_server(McpServerConfig::Stdio {...}) │ +│ .build() │ +│ .start().await │ +└──────────────────┬──────────────────────────┘ + │ Config passed via WebSocket + ▼ +┌──────────────────────────────────────────────┐ +│ localharness (subprocess) │ +│ │ +│ ┌──────────┐ ┌──────────┐ ┌───────────┐ │ +│ │ MCP │ │ MCP │ │ MCP │ │ +│ │ Client 1 │ │ Client 2 │ │ Client N │ │ +│ └────┬─────┘ └────┬─────┘ └─────┬─────┘ │ +│ │ │ │ │ +│ ▼ ▼ ▼ │ +│ ┌──────────┐ ┌──────────┐ ┌───────────┐ │ +│ │ External │ │ External │ │ External │ │ +│ │ MCP │ │ MCP │ │ MCP │ │ +│ │ Server │ │ Server │ │ Server │ │ +│ └──────────┘ └──────────┘ └───────────┘ │ +└──────────────────────────────────────────────┘ +``` + +The SDK configures MCP servers; the localharness manages their lifecycle (connecting, tool discovery, execution, disconnection). + +## Common MCP Servers + +| Package | Name | Transport | Description | +|---------|------|-----------|-------------| +| `@anthropic/mcp-fs-server` | `fs` | Stdio | File system access | +| `@anthropic/mcp-memory` | `memory` | Stdio | Persistent memory | +| `@anthropic/mcp-github` | `github` | Stdio | GitHub API | +| `@anthropic/mcp-slack` | `slack` | Stdio | Slack messaging | + +## McpServerConfig API + +```rust,no_run +impl McpServerConfig { + /// Returns the unique name identifier of this MCP server. + pub fn name(&self) -> &str; +} +``` + +The name is used for: +- Tool call routing (tools are prefixed with the server name) +- Policy matching (`{server_name}/{tool_name}` or `{server_name}/*`) +- Logging and diagnostics + +## Python SDK Comparison + +| Python | Rust | +|--------|------| +| `McpStdioServer(name, command, args)` | `McpServerConfig::Stdio { name, command, args, ... }` | +| `McpSseServer(name, url, headers)` | `McpServerConfig::Sse { name, url, headers, ... }` | +| N/A | `McpServerConfig::Http { ... }` (Rust-only) | +| `LocalAgentConfig(mcp_servers=[...])` | `Agent::builder().mcp_servers(vec![...])` | +| `McpBridge` (runtime client) | Handled by localharness (not in SDK) | diff --git a/docs/tools.md b/docs/tools.md new file mode 100644 index 0000000..dad2401 --- /dev/null +++ b/docs/tools.md @@ -0,0 +1,240 @@ +# Tools + +In-process tool execution for the Antigravity Rust SDK. + +## Overview + +Tools are Rust functions exposed to the Gemini model for invocation. The SDK provides a trait-based system with JSON schema validation, automatic registration, and optional session-scoped context. + +## Tool Trait + +Define a tool by implementing the `Tool` trait: + +```rust,no_run +use antigravity_sdk_rust::tools::Tool; +use serde_json::Value; + +struct WeatherTool; + +impl Tool for WeatherTool { + fn name(&self) -> &str { + "get_weather" + } + + fn description(&self) -> &str { + "Get the current weather for a city." + } + + fn parameters_json_schema(&self) -> &str { + r#"{ + "type": "object", + "properties": { + "city": { "type": "string", "description": "City name" } + }, + "required": ["city"] + }"# + } + + async fn call(&self, args: Value) -> Result { + let city = args.get("city").and_then(|c| c.as_str()).unwrap_or("Unknown"); + Ok(serde_json::json!({ + "temperature": 22, + "condition": "sunny", + "city": city + })) + } +} +``` + +### Full Trait Signature + +```rust,no_run +pub trait Tool: Send + Sync { + /// Unique name used in tool calls from the model. + fn name(&self) -> &str; + + /// Human-readable description of what the tool does. + fn description(&self) -> &str; + + /// JSON Schema string describing the expected arguments. + fn parameters_json_schema(&self) -> &str; + + /// Execute the tool with the given arguments. + async fn call(&self, args: Value) -> Result; + + /// Whether this tool requires a `ToolContext`. Default: `false`. + fn needs_context(&self) -> bool { false } + + /// Execute with session context. Default: delegates to `call()`. + async fn call_with_context( + &self, + args: Value, + ctx: &ToolContext, + ) -> Result { + self.call(args).await + } +} +``` + +## DynTool (Object-Safe) + +The `Tool` trait uses `async fn`, which isn't object-safe. The SDK provides `DynTool` — an object-safe wrapper trait that uses `BoxFuture` — with an automatic blanket implementation: + +```rust,no_run +// Any T: Tool automatically implements DynTool +// You never need to implement DynTool manually +let tool: Arc = Arc::new(WeatherTool); +``` + +## ToolRunner + +`ToolRunner` manages tool registration and dispatch: + +```rust,no_run +use antigravity_sdk_rust::tools::ToolRunner; +use std::sync::Arc; + +let runner = ToolRunner::new(); + +// Register tools +runner.register(Arc::new(WeatherTool)).await; + +// Execute a single tool by name +let result = runner.execute("get_weather", serde_json::json!({"city": "Tokyo"})).await?; + +// Batch-execute multiple tool calls +let results = runner.process_tool_calls(&tool_calls).await; +``` + +Tools are stored behind `Arc>>>` for concurrent access. + +## Context-Aware Tools + +Tools can opt-in to receiving a `ToolContext` for session state and agent communication: + +```rust,no_run +use antigravity_sdk_rust::tools::Tool; +use antigravity_sdk_rust::tool_context::ToolContext; +use serde_json::Value; + +struct CounterTool; + +impl Tool for CounterTool { + fn name(&self) -> &str { "counter" } + fn description(&self) -> &str { "Increment and return a counter" } + fn parameters_json_schema(&self) -> &str { r#"{"type":"object"}"# } + + fn needs_context(&self) -> bool { true } // Opt-in + + async fn call(&self, _args: Value) -> Result { + Ok(Value::Null) // Fallback when no context + } + + async fn call_with_context( + &self, + _args: Value, + ctx: &ToolContext, + ) -> Result { + let count: i32 = ctx.get_state("count").unwrap_or(0); + ctx.set_state("count", count + 1); + Ok(serde_json::json!({ "count": count + 1 })) + } +} +``` + +### ToolContext API + +```rust,no_run +pub struct ToolContext { + // Methods: + fn conversation_id(&self) -> &str; + fn is_idle(&self) -> bool; + async fn send(&self, message: &str) -> Result<()>; + fn get_state(&self, key: &str) -> Option; + fn set_state(&self, key: &str, value: T); +} +``` + +> **Note**: Tool state is independent of Hook state. They use separate stores. + +## Built-in Tools + +The SDK provides these built-in tools (managed by the harness): + +| Tool | Description | +|------|-------------| +| `CreateFile` | Create a new file | +| `EditFile` | Edit an existing file | +| `FindFile` | Search for files by name | +| `ListDir` | List directory contents | +| `RunCommand` | Execute shell commands | +| `SearchDir` | Search file contents | +| `ViewFile` | Read file contents | +| `StartSubagent` | Launch sub-agents | +| `GenerateImage` | Generate images | +| `Finish` | Signal task completion | +| `GrepSearch` | Grep-based search | + +### Read-Only Tools + +`BuiltinTools::read_only()` returns: `FindFile`, `ListDir`, `SearchDir`, `ViewFile`, `GrepSearch`. + +## Agent Builder Integration + +Register tools via the builder: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use std::sync::Arc; + +let agent = Agent::builder() + .tool(Arc::new(WeatherTool)) // single tool + .tools(vec![ // multiple tools + Arc::new(WeatherTool), + Arc::new(CounterTool), + ]) + .allow_all() + .build(); +``` + +## Shared State Pattern + +For tools that need shared mutable state, use `Arc>`: + +```rust,no_run +use std::sync::{Arc, Mutex}; +use std::collections::HashMap; + +struct InventoryTool { + db: Arc>>, +} + +impl Tool for InventoryTool { + fn name(&self) -> &str { "check_inventory" } + fn description(&self) -> &str { "Check item inventory" } + fn parameters_json_schema(&self) -> &str { r#"{"type":"object","properties":{"item":{"type":"string"}}}"# } + + async fn call(&self, args: Value) -> Result { + let item = args["item"].as_str().unwrap_or(""); + let db = self.db.lock().unwrap(); + let count = db.get(item).copied().unwrap_or(0); + Ok(serde_json::json!({ "item": item, "count": count })) + } +} +``` + +## Google Search Fallback + +If the model invokes `google_search` or `web_search` and no custom tool handles it, the SDK runs a built-in DuckDuckGo scraper: +- **Native**: Spawns `python3` subprocess to scrape results +- **WASM**: Returns empty result (not available) + +## Python SDK Comparison + +| Python | Rust | +|--------|------| +| `@tool` decorator or `ToolWithSchema` | `impl Tool for T` trait | +| `ToolRunner.register()` | `ToolRunner::register()` | +| `ToolRunner.execute()` | `ToolRunner::execute()` | +| `ToolContext` with `get_state`/`set_state` | `ToolContext` with `get_state`/`set_state` | +| Sync/async auto-detection | All tools are async | diff --git a/docs/triggers.md b/docs/triggers.md new file mode 100644 index 0000000..d3c7654 --- /dev/null +++ b/docs/triggers.md @@ -0,0 +1,334 @@ +# Triggers + +Background tasks and external events. + +## Overview + +Triggers are long-running async tasks that run concurrently with the agent +loop. They are spawned when the agent starts and can push notifications to the +agent via the connection. Common use cases include periodic health checks, +heartbeat monitors, file watchers, and external event listeners. + +### Python SDK Comparison + +The trigger concept is consistent across both SDKs. In Python, triggers are +registered via `Agent(triggers=[...])`. In Rust, triggers are registered +through the builder's `.trigger()` or `.triggers()` methods and implement the +`Trigger` trait. + +--- + +## Trigger Trait + +The `Trigger` trait defines a single `run` method that receives the active +connection and executes for the lifetime of the agent: + +```rust,no_run +use antigravity_sdk_rust::connection::AnyConnection; + +/// A trait for defining asynchronous background tasks that execute +/// during a connection lifecycle. +pub trait Trigger: Send + Sync { + /// Launches the trigger task with the active connection. + /// + /// This method runs for the lifetime of the agent. Use the connection + /// to send notifications back to the agent. + async fn run(&self, connection: AnyConnection) -> Result<(), anyhow::Error>; +} +``` + +The connection parameter gives triggers access to +`send_trigger_notification()`, which pushes a message string into the agent's +event stream. + +--- + +## DynTrigger (Object-Safe Wrapper) + +Like the `Hook`/`DynHook` pattern, triggers have an object-safe counterpart +that wraps the async `run` method in a `BoxFuture`: + +```rust,no_run +use futures_util::future::BoxFuture; +use antigravity_sdk_rust::connection::AnyConnection; + +/// Object-safe version of `Trigger`, automatically implemented +/// via a blanket impl for any `T: Trigger`. +pub trait DynTrigger: Send + Sync { + fn run(&self, connection: AnyConnection) -> BoxFuture<'_, Result<(), anyhow::Error>>; +} + +// Blanket impl: any type implementing Trigger automatically implements DynTrigger. +// impl DynTrigger for T { ... } +``` + +You always implement `Trigger` — the SDK handles the `DynTrigger` conversion +automatically. + +--- + +## TriggerRunner + +The `TriggerRunner` orchestrates the spawning and lifecycle of all registered +triggers: + +```rust,no_run +use antigravity_sdk_rust::triggers::TriggerRunner; +use std::sync::Arc; + +// TriggerRunner stores triggers as Vec> +``` + +### API + +| Method | Description | +|---|---| +| `TriggerRunner::new(triggers)` | Creates a runner wrapping a `Vec>` | +| `runner.start(connection)` | Spawns each trigger as an independent tokio task | + +When `start()` is called, each trigger is cloned (via `Arc`) and spawned into +its own `tokio::spawn` block. If a trigger's `run` method returns an error, +it is logged via `tracing::error!` but does not crash the agent. + +> [!NOTE] +> Triggers run independently — one trigger failing does not affect others. +> Errors are logged but silently swallowed to maintain agent stability. + +--- + +## Trigger Helpers + +### `every()` — Periodic Timer + +The `every()` factory function creates a trigger that fires at regular +intervals, sending a message to the agent each time: + +```rust,no_run +use antigravity_sdk_rust::trigger_helpers::every; +use std::time::Duration; + +// Send "check_status" to the agent every 30 seconds +let heartbeat = every(Duration::from_secs(30), "check_status"); + +// With a custom message +let monitor = every(Duration::from_millis(500), "fast_poll"); +``` + +The returned `PeriodicTrigger` loops indefinitely: + +```text +loop { + sleep(interval) + connection.send_trigger_notification(message) +} +``` + +### `PeriodicTrigger` Internals + +```rust,no_run +use std::time::Duration; + +/// A trigger that fires at regular intervals. +/// Created via `every()`. +pub struct PeriodicTrigger { + /// How long to wait between notifications. + interval: Duration, + /// The message sent to the agent on each tick. + message: String, +} +``` + +--- + +## Types + +### TriggerDelivery + +Controls when trigger notifications are delivered to the agent: + +```rust,no_run +/// Controls when trigger notifications are delivered to the agent. +#[derive(Debug, Clone, Copy)] +pub enum TriggerDelivery { + /// Deliver the notification immediately, even if the agent is busy. + SendImmediately, + /// Wait until the agent is idle before delivering. + WaitIdle, +} +``` + +### FileChange + +Represents a single filesystem change event, useful for file-watching triggers: + +```rust,no_run +/// The kind of filesystem change detected by a file-watching trigger. +#[derive(Debug, Clone, Copy)] +pub enum FileChangeKind { + /// A new file was created. + Added, + /// An existing file was modified. + Modified, + /// A file was deleted. + Deleted, +} + +/// A single filesystem change event. +#[derive(Debug, Clone)] +pub struct FileChange { + /// The type of change. + pub kind: FileChangeKind, + /// The path of the affected file. + pub path: String, +} +``` + +--- + +## Examples + +### 1. Heartbeat Trigger + +The simplest trigger — periodic status checks using the built-in `every()` helper: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::trigger_helpers::every; +use antigravity_sdk_rust::triggers::DynTrigger; +use std::sync::Arc; +use std::time::Duration; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let heartbeat = every(Duration::from_secs(60), "heartbeat_check"); + + let agent = Agent::builder() + .trigger(Arc::new(heartbeat) as Arc) + .allow_all() + .build(); + + let agent = agent.start().await?; + // The heartbeat trigger is now running in the background, + // sending "heartbeat_check" every 60 seconds. + + let response = agent.chat("Monitor the system").await?; + println!("{}", response.text); + + agent.stop().await?; + Ok(()) +} +``` + +### 2. Custom Monitoring Trigger + +A trigger that monitors a resource and sends notifications when conditions change: + +```rust,no_run +use antigravity_sdk_rust::connection::AnyConnection; +use antigravity_sdk_rust::connection::Connection; +use antigravity_sdk_rust::triggers::Trigger; +use std::time::Duration; + +struct DiskSpaceMonitor { + path: String, + threshold_mb: u64, + check_interval: Duration, +} + +impl DiskSpaceMonitor { + fn new(path: impl Into, threshold_mb: u64) -> Self { + Self { + path: path.into(), + threshold_mb, + check_interval: Duration::from_secs(300), // every 5 minutes + } + } +} + +impl Trigger for DiskSpaceMonitor { + async fn run(&self, connection: AnyConnection) -> Result<(), anyhow::Error> { + loop { + tokio::time::sleep(self.check_interval).await; + + // Simulate checking disk space + // In production, use a filesystem API + let available_mb = 500u64; // placeholder + + if available_mb < self.threshold_mb { + let msg = format!( + "⚠️ Low disk space on {}: {}MB remaining (threshold: {}MB)", + self.path, available_mb, self.threshold_mb + ); + connection.send_trigger_notification(&msg).await?; + } + } + } +} +``` + +### 3. Builder Integration + +Registering multiple triggers with the agent builder: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::trigger_helpers::every; +use antigravity_sdk_rust::triggers::DynTrigger; +use std::sync::Arc; +use std::time::Duration; + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let agent = Agent::builder() + // Add individual triggers + .trigger(Arc::new(every(Duration::from_secs(30), "health_check")) as Arc) + .trigger(Arc::new(every(Duration::from_secs(300), "metrics_report")) as Arc) + .allow_all() + .build(); + + let agent = agent.start().await?; + // Both triggers are now running concurrently in the background. + + agent.stop().await?; + Ok(()) +} +``` + +Or register after construction but before starting: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::trigger_helpers::every; +use antigravity_sdk_rust::triggers::DynTrigger; +use std::sync::Arc; +use std::time::Duration; + +let mut agent = Agent::builder().allow_all().build(); +agent.register_trigger( + Arc::new(every(Duration::from_secs(60), "late_trigger")) as Arc +).expect("trigger registration failed"); +// agent.start().await?; +``` + +### 4. Batch Triggers via `.triggers()` + +You can set all triggers at once using the plural form: + +```rust,no_run +use antigravity_sdk_rust::agent::Agent; +use antigravity_sdk_rust::trigger_helpers::every; +use antigravity_sdk_rust::triggers::DynTrigger; +use std::sync::Arc; +use std::time::Duration; + +let all_triggers: Vec> = vec![ + Arc::new(every(Duration::from_secs(30), "check_a")), + Arc::new(every(Duration::from_secs(60), "check_b")), + Arc::new(every(Duration::from_secs(120), "check_c")), +]; + +let agent = Agent::builder() + .triggers(all_triggers) + .allow_all() + .build(); +```