diff --git a/Cargo.lock b/Cargo.lock index 6faffd26..37a8709a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -426,6 +426,21 @@ version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b35204fbdc0b3f4446b89fc1ac2cf84a8a68971995d0bf2e925ec7cd960f9cb3" +[[package]] +name = "cassowary" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df8670b8c7b9dae1793364eafadf7239c40d669904660c5960d74cfd80b46a53" + +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + [[package]] name = "cc" version = "1.2.55" @@ -529,6 +544,20 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +[[package]] +name = "compact_str" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b79c4069c6cad78e2e0cdfcbd26275770669fb39fd308a752dc110e83b9af32" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "static_assertions", +] + [[package]] name = "constant_time_eq" version = "0.4.2" @@ -710,6 +739,32 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crossterm" +version = "0.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "829d955a0bb380ef178a640b91779e3987da38c9aea133b20614cfed8cdea9c6" +dependencies = [ + "bitflags 2.10.0", + "crossterm_winapi", + "futures-core", + "mio", + "parking_lot", + "rustix 0.38.44", + "signal-hook", + "signal-hook-mio", + "winapi", +] + +[[package]] +name = "crossterm_winapi" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "acdd7c62a3665c7f6830a51635d9ac9b23ed385797f70a83bb8bafe9c572ab2b" +dependencies = [ + "winapi", +] + [[package]] name = "crypto-common" version = "0.1.7" @@ -730,14 +785,38 @@ dependencies = [ "cipher", ] +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core 0.20.11", + "darling_macro 0.20.11", +] + [[package]] name = "darling" version = "0.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" dependencies = [ - "darling_core", - "darling_macro", + "darling_core 0.21.3", + "darling_macro 0.21.3", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.114", ] [[package]] @@ -754,13 +833,24 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core 0.20.11", + "quote", + "syn 2.0.114", +] + [[package]] name = "darling_macro" version = "0.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" dependencies = [ - "darling_core", + "darling_core 0.21.3", "quote", "syn 2.0.114", ] @@ -1194,6 +1284,8 @@ version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ + "allocator-api2", + "equivalent", "foldhash", "serde", ] @@ -1547,6 +1639,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + [[package]] name = "inout" version = "0.1.4" @@ -1556,6 +1657,19 @@ dependencies = [ "generic-array", ] +[[package]] +name = "instability" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6778b0196eefee7df739db78758e5cf9b37412268bfa5650bfeed028aed20d9c" +dependencies = [ + "darling 0.20.11", + "indoc", + "proc-macro2", + "quote", + "syn 2.0.114", +] + [[package]] name = "ipnet" version = "2.11.0" @@ -1587,6 +1701,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.17" @@ -1707,6 +1830,15 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "lru" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" +dependencies = [ + "hashbrown 0.15.5", +] + [[package]] name = "lru-slab" version = "0.1.2" @@ -1775,6 +1907,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" dependencies = [ "libc", + "log", "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] @@ -1791,6 +1924,7 @@ dependencies = [ "blake3", "chrono", "clap", + "crossterm", "deadpool-postgres", "dotenvy", "futures", @@ -1799,6 +1933,7 @@ dependencies = [ "postgres-types", "pretty_assertions", "rand 0.8.5", + "ratatui", "refinery", "regex", "reqwest", @@ -2293,6 +2428,27 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "ratatui" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eabd94c2f37801c20583fc49dd5cd6b0ba68c716787c2dd6ed18571e1e63117b" +dependencies = [ + "bitflags 2.10.0", + "cassowary", + "compact_str", + "crossterm", + "indoc", + "instability", + "itertools 0.13.0", + "lru", + "paste", + "strum", + "unicode-segmentation", + "unicode-truncate", + "unicode-width 0.2.0", +] + [[package]] name = "rayon" version = "1.11.0" @@ -2890,7 +3046,7 @@ version = "3.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52a8e3ca0ca629121f70ab50f95249e5a6f925cc0f6ffe8256c45b728875706c" dependencies = [ - "darling", + "darling 0.21.3", "proc-macro2", "quote", "syn 2.0.114", @@ -2922,6 +3078,27 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "signal-hook" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d881a16cf4426aa584979d30bd82cb33429027e42122b169753d6ef1085ed6e2" +dependencies = [ + "libc", + "signal-hook-registry", +] + +[[package]] +name = "signal-hook-mio" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b75a19a7a740b25bc7944bdee6172368f988763b744e3d4dfe753f6b4ece40cc" +dependencies = [ + "libc", + "mio", + "signal-hook", +] + [[package]] name = "signal-hook-registry" version = "1.4.8" @@ -2981,6 +3158,12 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "stringprep" version = "0.1.5" @@ -3021,6 +3204,28 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "strum" +version = "0.26.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "rustversion", + "syn 2.0.114", +] + [[package]] name = "subtle" version = "2.6.1" @@ -3580,10 +3785,33 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" [[package]] -name = "unicode-width" -version = "0.2.2" +name = "unicode-segmentation" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" +checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" + +[[package]] +name = "unicode-truncate" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3644627a5af5fa321c95b9b235a72fd24cd29c648c2c379431e6628655627bf" +dependencies = [ + "itertools 0.13.0", + "unicode-segmentation", + "unicode-width 0.1.14", +] + +[[package]] +name = "unicode-width" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" + +[[package]] +name = "unicode-width" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd" [[package]] name = "unicode-xid" @@ -3943,7 +4171,7 @@ dependencies = [ "cranelift-frontend", "cranelift-native", "gimli", - "itertools", + "itertools 0.12.1", "log", "object 0.36.7", "smallvec", @@ -4074,7 +4302,7 @@ dependencies = [ "bumpalo", "leb128fmt", "memchr", - "unicode-width", + "unicode-width 0.2.0", "wasm-encoder 0.244.0", ] diff --git a/Cargo.toml b/Cargo.toml index 3e38ee7e..1b1ab514 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -48,6 +48,10 @@ async-trait = "0.1" # CLI clap = { version = "4", features = ["derive", "env"] } +# TUI +ratatui = "0.29" +crossterm = { version = "0.28", features = ["event-stream"] } + # Channel integrations axum = "0.8" tower = "0.5" diff --git a/src/agent/compaction.rs b/src/agent/compaction.rs new file mode 100644 index 00000000..22b0ea6a --- /dev/null +++ b/src/agent/compaction.rs @@ -0,0 +1,324 @@ +//! Context compaction for preserving and summarizing conversation history. +//! +//! When the context window approaches its limit, compaction: +//! 1. Summarizes old turns +//! 2. Writes the summary to the workspace daily log +//! 3. Trims the context to keep only recent turns + +use std::sync::Arc; + +use chrono::Utc; + +use crate::agent::context_monitor::{CompactionStrategy, ContextBreakdown}; +use crate::agent::session::Thread; +use crate::error::Error; +use crate::llm::{ChatMessage, CompletionRequest, LlmProvider}; +use crate::workspace::Workspace; + +/// Result of a compaction operation. +#[derive(Debug)] +pub struct CompactionResult { + /// Number of turns removed. + pub turns_removed: usize, + /// Tokens before compaction. + pub tokens_before: usize, + /// Tokens after compaction. + pub tokens_after: usize, + /// Whether a summary was written to workspace. + pub summary_written: bool, + /// The generated summary (if any). + pub summary: Option, +} + +/// Compacts conversation context to stay within limits. +pub struct ContextCompactor { + llm: Arc, +} + +impl ContextCompactor { + /// Create a new context compactor. + pub fn new(llm: Arc) -> Self { + Self { llm } + } + + /// Compact a thread's context using the given strategy. + pub async fn compact( + &self, + thread: &mut Thread, + strategy: CompactionStrategy, + workspace: Option<&Workspace>, + ) -> Result { + let messages = thread.messages(); + let tokens_before = ContextBreakdown::analyze(&messages).total_tokens; + + let result = match strategy { + CompactionStrategy::Summarize { keep_recent } => { + self.compact_with_summary(thread, keep_recent, workspace) + .await? + } + CompactionStrategy::Truncate { keep_recent } => { + self.compact_truncate(thread, keep_recent) + } + CompactionStrategy::MoveToWorkspace => { + self.compact_to_workspace(thread, workspace).await? + } + }; + + let messages_after = thread.messages(); + let tokens_after = ContextBreakdown::analyze(&messages_after).total_tokens; + + Ok(CompactionResult { + turns_removed: result.turns_removed, + tokens_before, + tokens_after, + summary_written: result.summary_written, + summary: result.summary, + }) + } + + /// Compact by summarizing old turns. + async fn compact_with_summary( + &self, + thread: &mut Thread, + keep_recent: usize, + workspace: Option<&Workspace>, + ) -> Result { + if thread.turns.len() <= keep_recent { + return Ok(CompactionPartial::empty()); + } + + // Get turns to summarize + let turns_to_remove = thread.turns.len() - keep_recent; + let old_turns = &thread.turns[..turns_to_remove]; + + // Build messages for summarization + let mut to_summarize = Vec::new(); + for turn in old_turns { + to_summarize.push(ChatMessage::user(&turn.user_input)); + if let Some(ref response) = turn.response { + to_summarize.push(ChatMessage::assistant(response)); + } + } + + // Generate summary + let summary = self.generate_summary(&to_summarize).await?; + + // Write to workspace if available + let summary_written = if let Some(ws) = workspace { + self.write_summary_to_workspace(ws, &summary).await.is_ok() + } else { + false + }; + + // Truncate thread + thread.truncate_turns(keep_recent); + + Ok(CompactionPartial { + turns_removed: turns_to_remove, + summary_written, + summary: Some(summary), + }) + } + + /// Compact by simple truncation (no summary). + fn compact_truncate(&self, thread: &mut Thread, keep_recent: usize) -> CompactionPartial { + let turns_before = thread.turns.len(); + thread.truncate_turns(keep_recent); + let turns_removed = turns_before - thread.turns.len(); + + CompactionPartial { + turns_removed, + summary_written: false, + summary: None, + } + } + + /// Move context to workspace without summarization. + async fn compact_to_workspace( + &self, + thread: &mut Thread, + workspace: Option<&Workspace>, + ) -> Result { + let Some(ws) = workspace else { + // Fall back to truncation if no workspace + return Ok(self.compact_truncate(thread, 5)); + }; + + // Keep more turns when moving to workspace (we have a backup) + let keep_recent = 10; + if thread.turns.len() <= keep_recent { + return Ok(CompactionPartial::empty()); + } + + let turns_to_remove = thread.turns.len() - keep_recent; + let old_turns = &thread.turns[..turns_to_remove]; + + // Format turns for storage + let content = format_turns_for_storage(old_turns); + + // Write to workspace + let written = self.write_context_to_workspace(ws, &content).await.is_ok(); + + // Truncate + thread.truncate_turns(keep_recent); + + Ok(CompactionPartial { + turns_removed: turns_to_remove, + summary_written: written, + summary: None, + }) + } + + /// Generate a summary of messages using the LLM. + async fn generate_summary(&self, messages: &[ChatMessage]) -> Result { + let prompt = ChatMessage::system( + r#"Summarize the following conversation concisely. Focus on: +- Key decisions made +- Important information exchanged +- Actions taken +- Outcomes achieved + +Be brief but capture all important details. Use bullet points."#, + ); + + let mut request_messages = vec![prompt]; + + // Add a user message with the conversation to summarize + let formatted = messages + .iter() + .map(|m| { + let role_str = match m.role { + crate::llm::Role::User => "User", + crate::llm::Role::Assistant => "Assistant", + crate::llm::Role::System => "System", + crate::llm::Role::Tool => { + return format!( + "Tool {}: {}", + m.name.as_deref().unwrap_or("unknown"), + m.content + ); + } + }; + format!("{}: {}", role_str, m.content) + }) + .collect::>() + .join("\n\n"); + + request_messages.push(ChatMessage::user(format!( + "Please summarize this conversation:\n\n{}", + formatted + ))); + + let request = CompletionRequest::new(request_messages) + .with_max_tokens(1024) + .with_temperature(0.3); + + let response = self.llm.complete(request).await?; + Ok(response.content) + } + + /// Write a summary to the workspace daily log. + async fn write_summary_to_workspace( + &self, + workspace: &Workspace, + summary: &str, + ) -> Result<(), Error> { + let date = Utc::now().format("%Y-%m-%d"); + let entry = format!( + "\n## Context Summary ({})\n\n{}\n", + Utc::now().format("%H:%M UTC"), + summary + ); + + workspace + .append(&format!("daily/{}.md", date), &entry) + .await?; + Ok(()) + } + + /// Write full context to workspace for archival. + async fn write_context_to_workspace( + &self, + workspace: &Workspace, + content: &str, + ) -> Result<(), Error> { + let date = Utc::now().format("%Y-%m-%d"); + let entry = format!( + "\n## Archived Context ({})\n\n{}\n", + Utc::now().format("%H:%M UTC"), + content + ); + + workspace + .append(&format!("daily/{}.md", date), &entry) + .await?; + Ok(()) + } +} + +/// Partial result during compaction (internal). +struct CompactionPartial { + turns_removed: usize, + summary_written: bool, + summary: Option, +} + +impl CompactionPartial { + fn empty() -> Self { + Self { + turns_removed: 0, + summary_written: false, + summary: None, + } + } +} + +/// Format turns for storage in workspace. +fn format_turns_for_storage(turns: &[crate::agent::session::Turn]) -> String { + turns + .iter() + .map(|turn| { + let mut s = format!("**Turn {}**\n", turn.turn_number + 1); + s.push_str(&format!("User: {}\n", turn.user_input)); + if let Some(ref response) = turn.response { + s.push_str(&format!("Agent: {}\n", response)); + } + if !turn.tool_calls.is_empty() { + s.push_str("Tools: "); + let tools: Vec<_> = turn.tool_calls.iter().map(|t| t.name.as_str()).collect(); + s.push_str(&tools.join(", ")); + s.push('\n'); + } + s + }) + .collect::>() + .join("\n") +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::agent::session::Thread; + use uuid::Uuid; + + #[test] + fn test_format_turns() { + let mut thread = Thread::new(Uuid::new_v4()); + thread.start_turn("Hello"); + thread.complete_turn("Hi there"); + thread.start_turn("How are you?"); + thread.complete_turn("I'm good!"); + + let formatted = format_turns_for_storage(&thread.turns); + assert!(formatted.contains("Turn 1")); + assert!(formatted.contains("Hello")); + assert!(formatted.contains("Turn 2")); + } + + #[test] + fn test_compaction_partial_empty() { + let partial = CompactionPartial::empty(); + assert_eq!(partial.turns_removed, 0); + assert!(!partial.summary_written); + } +} diff --git a/src/agent/context_monitor.rs b/src/agent/context_monitor.rs new file mode 100644 index 00000000..c236a520 --- /dev/null +++ b/src/agent/context_monitor.rs @@ -0,0 +1,236 @@ +//! Context window monitoring and compaction triggers. +//! +//! Monitors the size of the conversation context and triggers +//! compaction when approaching the limit. + +use crate::llm::ChatMessage; + +/// Default context window limit (conservative estimate). +const DEFAULT_CONTEXT_LIMIT: usize = 100_000; + +/// Compaction threshold as a percentage of the limit. +const COMPACTION_THRESHOLD: f64 = 0.8; + +/// Approximate tokens per word (rough estimate for English). +const TOKENS_PER_WORD: f64 = 1.3; + +/// Strategy for context compaction. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CompactionStrategy { + /// Summarize old messages and keep recent ones. + Summarize { + /// Number of recent turns to keep intact. + keep_recent: usize, + }, + /// Truncate old messages without summarization. + Truncate { + /// Number of recent turns to keep. + keep_recent: usize, + }, + /// Move context to workspace memory. + MoveToWorkspace, +} + +impl Default for CompactionStrategy { + fn default() -> Self { + Self::Summarize { keep_recent: 5 } + } +} + +/// Monitors context size and suggests compaction. +pub struct ContextMonitor { + /// Maximum tokens allowed in context. + context_limit: usize, + /// Threshold ratio for triggering compaction. + threshold_ratio: f64, +} + +impl ContextMonitor { + /// Create a new context monitor with default settings. + pub fn new() -> Self { + Self { + context_limit: DEFAULT_CONTEXT_LIMIT, + threshold_ratio: COMPACTION_THRESHOLD, + } + } + + /// Create with a custom context limit. + pub fn with_limit(mut self, limit: usize) -> Self { + self.context_limit = limit; + self + } + + /// Create with a custom threshold ratio. + pub fn with_threshold(mut self, ratio: f64) -> Self { + self.threshold_ratio = ratio.clamp(0.5, 0.95); + self + } + + /// Estimate the token count for a list of messages. + pub fn estimate_tokens(&self, messages: &[ChatMessage]) -> usize { + messages.iter().map(|m| estimate_message_tokens(m)).sum() + } + + /// Check if compaction is needed. + pub fn needs_compaction(&self, messages: &[ChatMessage]) -> bool { + let tokens = self.estimate_tokens(messages); + let threshold = (self.context_limit as f64 * self.threshold_ratio) as usize; + tokens >= threshold + } + + /// Get the current usage percentage. + pub fn usage_percent(&self, messages: &[ChatMessage]) -> f64 { + let tokens = self.estimate_tokens(messages); + (tokens as f64 / self.context_limit as f64) * 100.0 + } + + /// Suggest a compaction strategy based on current context. + pub fn suggest_compaction(&self, messages: &[ChatMessage]) -> Option { + if !self.needs_compaction(messages) { + return None; + } + + let tokens = self.estimate_tokens(messages); + let overage = tokens as f64 / self.context_limit as f64; + + if overage > 0.95 { + // Critical: aggressive truncation + Some(CompactionStrategy::Truncate { keep_recent: 3 }) + } else if overage > 0.85 { + // High: summarize and keep fewer + Some(CompactionStrategy::Summarize { keep_recent: 5 }) + } else { + // Moderate: move to workspace + Some(CompactionStrategy::MoveToWorkspace) + } + } + + /// Get the context limit. + pub fn limit(&self) -> usize { + self.context_limit + } + + /// Get the current threshold in tokens. + pub fn threshold(&self) -> usize { + (self.context_limit as f64 * self.threshold_ratio) as usize + } +} + +impl Default for ContextMonitor { + fn default() -> Self { + Self::new() + } +} + +/// Estimate tokens for a single message. +fn estimate_message_tokens(message: &ChatMessage) -> usize { + // Use word-based estimation as it's more accurate for varied content + let word_count = message.content.split_whitespace().count(); + + // Add overhead for role and structure + let overhead = 4; // ~4 tokens for role and message structure + + (word_count as f64 * TOKENS_PER_WORD) as usize + overhead +} + +/// Estimate tokens for raw text. +pub fn estimate_text_tokens(text: &str) -> usize { + let word_count = text.split_whitespace().count(); + (word_count as f64 * TOKENS_PER_WORD) as usize +} + +/// Context size breakdown for reporting. +#[derive(Debug, Clone)] +pub struct ContextBreakdown { + /// Total estimated tokens. + pub total_tokens: usize, + /// System message tokens. + pub system_tokens: usize, + /// User message tokens. + pub user_tokens: usize, + /// Assistant message tokens. + pub assistant_tokens: usize, + /// Tool result tokens. + pub tool_tokens: usize, + /// Number of messages. + pub message_count: usize, +} + +impl ContextBreakdown { + /// Analyze a list of messages. + pub fn analyze(messages: &[ChatMessage]) -> Self { + let mut breakdown = Self { + total_tokens: 0, + system_tokens: 0, + user_tokens: 0, + assistant_tokens: 0, + tool_tokens: 0, + message_count: messages.len(), + }; + + for message in messages { + let tokens = estimate_message_tokens(message); + breakdown.total_tokens += tokens; + + match message.role { + crate::llm::Role::System => breakdown.system_tokens += tokens, + crate::llm::Role::User => breakdown.user_tokens += tokens, + crate::llm::Role::Assistant => breakdown.assistant_tokens += tokens, + crate::llm::Role::Tool => breakdown.tool_tokens += tokens, + } + } + + breakdown + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_token_estimation() { + let msg = ChatMessage::user("Hello, how are you today?"); + let tokens = estimate_message_tokens(&msg); + // 5 words * 1.3 + 4 overhead = ~10-11 tokens + assert!(tokens > 0); + assert!(tokens < 20); + } + + #[test] + fn test_needs_compaction() { + let monitor = ContextMonitor::new().with_limit(100); + + // Small context - no compaction needed + let small: Vec = vec![ChatMessage::user("Hello")]; + assert!(!monitor.needs_compaction(&small)); + + // Large context - compaction needed + let large_content = "word ".repeat(1000); + let large: Vec = vec![ChatMessage::user(&large_content)]; + assert!(monitor.needs_compaction(&large)); + } + + #[test] + fn test_suggest_compaction() { + let monitor = ContextMonitor::new().with_limit(100); + + let small: Vec = vec![ChatMessage::user("Hello")]; + assert!(monitor.suggest_compaction(&small).is_none()); + } + + #[test] + fn test_context_breakdown() { + let messages = vec![ + ChatMessage::system("You are a helpful assistant."), + ChatMessage::user("Hello"), + ChatMessage::assistant("Hi there!"), + ]; + + let breakdown = ContextBreakdown::analyze(&messages); + assert_eq!(breakdown.message_count, 3); + assert!(breakdown.system_tokens > 0); + assert!(breakdown.user_tokens > 0); + assert!(breakdown.assistant_tokens > 0); + } +} diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 0685f677..0875a6ff 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -6,17 +6,31 @@ //! - Tool invocation with safety //! - Self-repair for stuck jobs //! - Proactive heartbeat execution +//! - Turn-based session management with undo +//! - Context compaction for long conversations mod agent_loop; +pub mod compaction; +pub mod context_monitor; mod heartbeat; mod router; mod scheduler; mod self_repair; +pub mod session; +pub mod submission; +pub mod task; +pub mod undo; mod worker; pub use agent_loop::Agent; +pub use compaction::{CompactionResult, ContextCompactor}; +pub use context_monitor::{CompactionStrategy, ContextBreakdown, ContextMonitor}; pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat}; pub use router::{MessageIntent, Router}; pub use scheduler::Scheduler; pub use self_repair::{RepairResult, RepairTask, SelfRepair, StuckJob}; +pub use session::{Session, Thread, ThreadState, Turn, TurnState}; +pub use submission::{Submission, SubmissionResult}; +pub use task::{Task, TaskContext, TaskHandler, TaskOutput, TaskStatus}; +pub use undo::{Checkpoint, UndoManager}; pub use worker::Worker; diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index b92ec923..07163f5b 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -2,15 +2,17 @@ use std::collections::HashMap; use std::sync::Arc; +use std::time::Duration; -use tokio::sync::{RwLock, mpsc}; +use tokio::sync::{RwLock, mpsc, oneshot}; use tokio::task::JoinHandle; use uuid::Uuid; use crate::agent::Worker; +use crate::agent::task::{Task, TaskContext, TaskOutput}; use crate::config::AgentConfig; -use crate::context::{ContextManager, JobState}; -use crate::error::JobError; +use crate::context::{ContextManager, JobContext, JobState}; +use crate::error::{Error, JobError}; use crate::history::Store; use crate::llm::LlmProvider; use crate::safety::SafetyLayer; @@ -35,6 +37,12 @@ pub struct ScheduledJob { pub tx: mpsc::Sender, } +/// Status of a scheduled sub-task. +struct ScheduledSubtask { + task_id: Uuid, + handle: JoinHandle>, +} + /// Schedules and manages parallel job execution. pub struct Scheduler { config: AgentConfig, @@ -43,8 +51,10 @@ pub struct Scheduler { safety: Arc, tools: Arc, store: Option>, - /// Running jobs. + /// Running jobs (main LLM-driven jobs). jobs: RwLock>, + /// Running sub-tasks (tool executions, background tasks). + subtasks: RwLock>, } impl Scheduler { @@ -65,6 +75,7 @@ impl Scheduler { tools, store, jobs: RwLock::new(HashMap::new()), + subtasks: RwLock::new(HashMap::new()), } } @@ -131,6 +142,185 @@ impl Scheduler { Ok(()) } + /// Schedule a sub-task from within a worker. + /// + /// Sub-tasks are lightweight tasks that don't go through the full job lifecycle. + /// They're used for parallel tool execution and background computations. + /// + /// Returns a oneshot receiver to get the result. + pub async fn spawn_subtask( + &self, + parent_id: Uuid, + task: Task, + ) -> Result>, JobError> { + let task_id = Uuid::new_v4(); + let (result_tx, result_rx) = oneshot::channel(); + + let handle = match task { + Task::Job { .. } => { + // Jobs should go through schedule(), not spawn_subtask + return Err(JobError::ContextError { + id: parent_id, + reason: "Use schedule() for Job tasks, not spawn_subtask()".to_string(), + }); + } + + Task::ToolExec { + parent_id: tool_parent_id, + tool_name, + params, + } => { + let tools = self.tools.clone(); + let context_manager = self.context_manager.clone(); + + tokio::spawn(async move { + let result = Self::execute_tool_task( + tools, + context_manager, + tool_parent_id, + &tool_name, + params, + ) + .await; + + // Send result (ignore if receiver dropped) + let _ = result_tx.send(result); + }) + } + + Task::Background { id: _, handler } => { + let ctx = TaskContext::new(task_id).with_parent(parent_id); + + tokio::spawn(async move { + let result = handler.run(ctx).await; + let _ = result_tx.send(result); + }) + } + }; + + // Track the subtask + self.subtasks.write().await.insert( + task_id, + ScheduledSubtask { + task_id, + handle: tokio::spawn(async move { + // Wrap the handle to get its result + match handle.await { + Ok(()) => Err(Error::Job(JobError::ContextError { + id: task_id, + reason: "Subtask completed but result not captured".to_string(), + })), + Err(e) => Err(Error::Job(JobError::ContextError { + id: task_id, + reason: format!("Subtask panicked: {}", e), + })), + } + }), + }, + ); + + tracing::debug!( + parent_id = %parent_id, + task_id = %task_id, + "Spawned subtask" + ); + + Ok(result_rx) + } + + /// Schedule multiple tasks in parallel and wait for all to complete. + /// + /// Returns results in the same order as the input tasks. + pub async fn spawn_batch( + &self, + parent_id: Uuid, + tasks: Vec, + ) -> Vec> { + if tasks.is_empty() { + return Vec::new(); + } + + let mut receivers = Vec::with_capacity(tasks.len()); + + // Spawn all tasks + for task in tasks { + match self.spawn_subtask(parent_id, task).await { + Ok(rx) => receivers.push(Some(rx)), + Err(e) => { + // Store the error directly + receivers.push(None); + tracing::warn!( + parent_id = %parent_id, + error = %e, + "Failed to spawn subtask in batch" + ); + } + } + } + + // Collect results + let mut results = Vec::with_capacity(receivers.len()); + for rx in receivers { + let result = match rx { + Some(receiver) => match receiver.await { + Ok(task_result) => task_result, + Err(_) => Err(Error::Job(JobError::ContextError { + id: parent_id, + reason: "Subtask channel closed unexpectedly".to_string(), + })), + }, + None => Err(Error::Job(JobError::ContextError { + id: parent_id, + reason: "Subtask failed to spawn".to_string(), + })), + }; + results.push(result); + } + + results + } + + /// Execute a single tool as a subtask. + async fn execute_tool_task( + tools: Arc, + context_manager: Arc, + job_id: Uuid, + tool_name: &str, + params: serde_json::Value, + ) -> Result { + let start = std::time::Instant::now(); + + // Get the tool + let tool = tools.get(tool_name).await.ok_or_else(|| { + Error::Tool(crate::error::ToolError::NotFound { + name: tool_name.to_string(), + }) + })?; + + // Get job context + let job_ctx: JobContext = context_manager.get_context(job_id).await?; + + // Execute with timeout + let result = tokio::time::timeout(Duration::from_secs(60), async { + tool.execute(params, &job_ctx).await + }) + .await + .map_err(|_| { + Error::Tool(crate::error::ToolError::Timeout { + name: tool_name.to_string(), + timeout: Duration::from_secs(60), + }) + })? + .map_err(|e| { + Error::Tool(crate::error::ToolError::ExecutionFailed { + name: tool_name.to_string(), + reason: e.to_string(), + }) + })?; + + Ok(TaskOutput::new(result.result, start.elapsed())) + } + /// Stop a running job. pub async fn stop(&self, job_id: Uuid) -> Result<(), JobError> { let mut jobs = self.jobs.write().await; @@ -190,25 +380,50 @@ impl Scheduler { self.jobs.read().await.len() } + /// Get count of running subtasks. + pub async fn subtask_count(&self) -> usize { + self.subtasks.read().await.len() + } + /// Get all running job IDs. pub async fn running_jobs(&self) -> Vec { self.jobs.read().await.keys().cloned().collect() } - /// Clean up finished jobs. + /// Clean up finished jobs and subtasks. pub async fn cleanup_finished(&self) { - let mut jobs = self.jobs.write().await; - let mut finished = Vec::new(); + // Clean up jobs + { + let mut jobs = self.jobs.write().await; + let mut finished = Vec::new(); - for (id, scheduled) in jobs.iter() { - if scheduled.handle.is_finished() { - finished.push(*id); + for (id, scheduled) in jobs.iter() { + if scheduled.handle.is_finished() { + finished.push(*id); + } + } + + for id in finished { + jobs.remove(&id); + tracing::debug!("Cleaned up finished job {}", id); } } - for id in finished { - jobs.remove(&id); - tracing::debug!("Cleaned up finished job {}", id); + // Clean up subtasks + { + let mut subtasks = self.subtasks.write().await; + let mut finished = Vec::new(); + + for (id, scheduled) in subtasks.iter() { + if scheduled.handle.is_finished() { + finished.push(*id); + } + } + + for id in finished { + subtasks.remove(&id); + tracing::trace!("Cleaned up finished subtask {}", id); + } } } @@ -219,16 +434,35 @@ impl Scheduler { for job_id in job_ids { let _ = self.stop(job_id).await; } + + // Abort all subtasks + let mut subtasks = self.subtasks.write().await; + for (_, scheduled) in subtasks.drain() { + scheduled.handle.abort(); + } + } + + /// Get access to the tools registry. + pub fn tools(&self) -> &Arc { + &self.tools + } + + /// Get access to the context manager. + pub fn context_manager(&self) -> &Arc { + &self.context_manager } } #[cfg(test)] mod tests { - // Note: Full scheduler tests require mocking LLM provider - // These are placeholder tests - #[test] fn test_scheduler_creation() { // Would need to mock dependencies for proper testing } + + #[tokio::test] + async fn test_spawn_batch_empty() { + // This test would need mock dependencies. + // For now just verify the empty case doesn't panic. + } } diff --git a/src/agent/session.rs b/src/agent/session.rs new file mode 100644 index 00000000..615f803d --- /dev/null +++ b/src/agent/session.rs @@ -0,0 +1,390 @@ +//! Session and thread model for turn-based agent interactions. +//! +//! A Session contains one or more Threads. Each Thread represents a +//! conversation/interaction sequence with the agent. Threads contain +//! Turns, which are request/response pairs. +//! +//! This model supports: +//! - Undo: Roll back to a previous turn +//! - Interrupt: Cancel the current turn mid-execution +//! - Compaction: Summarize old turns to save context +//! - Resume: Continue from a saved checkpoint + +use std::collections::HashMap; + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::llm::ChatMessage; + +/// A session containing one or more threads. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Session { + /// Unique session ID. + pub id: Uuid, + /// User ID that owns this session. + pub user_id: String, + /// Active thread ID. + pub active_thread: Option, + /// All threads in this session. + pub threads: HashMap, + /// When the session was created. + pub created_at: DateTime, + /// When the session was last active. + pub last_active_at: DateTime, + /// Session metadata. + pub metadata: serde_json::Value, +} + +impl Session { + /// Create a new session. + pub fn new(user_id: impl Into) -> Self { + let now = Utc::now(); + Self { + id: Uuid::new_v4(), + user_id: user_id.into(), + active_thread: None, + threads: HashMap::new(), + created_at: now, + last_active_at: now, + metadata: serde_json::Value::Null, + } + } + + /// Create a new thread in this session. + pub fn create_thread(&mut self) -> &mut Thread { + let thread = Thread::new(self.id); + let thread_id = thread.id; + self.threads.insert(thread_id, thread); + self.active_thread = Some(thread_id); + self.last_active_at = Utc::now(); + self.threads.get_mut(&thread_id).expect("just inserted") + } + + /// Get the active thread. + pub fn active_thread(&self) -> Option<&Thread> { + self.active_thread.and_then(|id| self.threads.get(&id)) + } + + /// Get the active thread mutably. + pub fn active_thread_mut(&mut self) -> Option<&mut Thread> { + self.active_thread.and_then(|id| self.threads.get_mut(&id)) + } + + /// Get or create the active thread. + pub fn get_or_create_thread(&mut self) -> &mut Thread { + if self.active_thread.is_none() { + self.create_thread(); + } + self.active_thread_mut().expect("just created") + } + + /// Switch to a different thread. + pub fn switch_thread(&mut self, thread_id: Uuid) -> bool { + if self.threads.contains_key(&thread_id) { + self.active_thread = Some(thread_id); + self.last_active_at = Utc::now(); + true + } else { + false + } + } +} + +/// State of a thread. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum ThreadState { + /// Thread is idle, waiting for input. + Idle, + /// Thread is processing a turn. + Processing, + /// Thread is waiting for user approval. + AwaitingApproval, + /// Thread has completed (no more turns expected). + Completed, + /// Thread was interrupted. + Interrupted, +} + +/// A conversation thread within a session. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Thread { + /// Unique thread ID. + pub id: Uuid, + /// Parent session ID. + pub session_id: Uuid, + /// Current state. + pub state: ThreadState, + /// Turns in this thread. + pub turns: Vec, + /// When the thread was created. + pub created_at: DateTime, + /// When the thread was last updated. + pub updated_at: DateTime, + /// Thread metadata (e.g., title, tags). + pub metadata: serde_json::Value, +} + +impl Thread { + /// Create a new thread. + pub fn new(session_id: Uuid) -> Self { + let now = Utc::now(); + Self { + id: Uuid::new_v4(), + session_id, + state: ThreadState::Idle, + turns: Vec::new(), + created_at: now, + updated_at: now, + metadata: serde_json::Value::Null, + } + } + + /// Get the current turn number (1-indexed for display). + pub fn turn_number(&self) -> usize { + self.turns.len() + 1 + } + + /// Get the last turn. + pub fn last_turn(&self) -> Option<&Turn> { + self.turns.last() + } + + /// Get the last turn mutably. + pub fn last_turn_mut(&mut self) -> Option<&mut Turn> { + self.turns.last_mut() + } + + /// Start a new turn with user input. + pub fn start_turn(&mut self, user_input: impl Into) -> &mut Turn { + let turn_number = self.turns.len(); + let turn = Turn::new(turn_number, user_input); + self.turns.push(turn); + self.state = ThreadState::Processing; + self.updated_at = Utc::now(); + self.turns.last_mut().expect("just pushed") + } + + /// Complete the current turn with a response. + pub fn complete_turn(&mut self, response: impl Into) { + if let Some(turn) = self.turns.last_mut() { + turn.complete(response); + } + self.state = ThreadState::Idle; + self.updated_at = Utc::now(); + } + + /// Fail the current turn with an error. + pub fn fail_turn(&mut self, error: impl Into) { + if let Some(turn) = self.turns.last_mut() { + turn.fail(error); + } + self.state = ThreadState::Idle; + self.updated_at = Utc::now(); + } + + /// Mark the thread as awaiting approval. + pub fn await_approval(&mut self) { + self.state = ThreadState::AwaitingApproval; + self.updated_at = Utc::now(); + } + + /// Interrupt the current turn. + pub fn interrupt(&mut self) { + if let Some(turn) = self.turns.last_mut() { + turn.interrupt(); + } + self.state = ThreadState::Interrupted; + self.updated_at = Utc::now(); + } + + /// Resume after interruption. + pub fn resume(&mut self) { + if self.state == ThreadState::Interrupted { + self.state = ThreadState::Idle; + self.updated_at = Utc::now(); + } + } + + /// Get all messages for context building. + pub fn messages(&self) -> Vec { + let mut messages = Vec::new(); + for turn in &self.turns { + messages.push(ChatMessage::user(&turn.user_input)); + if let Some(ref response) = turn.response { + messages.push(ChatMessage::assistant(response)); + } + } + messages + } + + /// Truncate turns to a specific count (keeping most recent). + pub fn truncate_turns(&mut self, keep: usize) { + if self.turns.len() > keep { + let drain_count = self.turns.len() - keep; + self.turns.drain(0..drain_count); + // Re-number remaining turns + for (i, turn) in self.turns.iter_mut().enumerate() { + turn.turn_number = i; + } + } + } +} + +/// State of a turn. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum TurnState { + /// Turn is being processed. + Processing, + /// Turn completed successfully. + Completed, + /// Turn failed with an error. + Failed, + /// Turn was interrupted. + Interrupted, +} + +/// A single turn (request/response pair) in a thread. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Turn { + /// Turn number (0-indexed). + pub turn_number: usize, + /// User input that started this turn. + pub user_input: String, + /// Agent response (if completed). + pub response: Option, + /// Tool calls made during this turn. + pub tool_calls: Vec, + /// Turn state. + pub state: TurnState, + /// When the turn started. + pub started_at: DateTime, + /// When the turn completed. + pub completed_at: Option>, + /// Error message (if failed). + pub error: Option, +} + +impl Turn { + /// Create a new turn. + pub fn new(turn_number: usize, user_input: impl Into) -> Self { + Self { + turn_number, + user_input: user_input.into(), + response: None, + tool_calls: Vec::new(), + state: TurnState::Processing, + started_at: Utc::now(), + completed_at: None, + error: None, + } + } + + /// Complete this turn. + pub fn complete(&mut self, response: impl Into) { + self.response = Some(response.into()); + self.state = TurnState::Completed; + self.completed_at = Some(Utc::now()); + } + + /// Fail this turn. + pub fn fail(&mut self, error: impl Into) { + self.error = Some(error.into()); + self.state = TurnState::Failed; + self.completed_at = Some(Utc::now()); + } + + /// Interrupt this turn. + pub fn interrupt(&mut self) { + self.state = TurnState::Interrupted; + self.completed_at = Some(Utc::now()); + } + + /// Record a tool call. + pub fn record_tool_call(&mut self, name: impl Into, params: serde_json::Value) { + self.tool_calls.push(TurnToolCall { + name: name.into(), + parameters: params, + result: None, + error: None, + }); + } + + /// Record tool call result. + pub fn record_tool_result(&mut self, result: serde_json::Value) { + if let Some(call) = self.tool_calls.last_mut() { + call.result = Some(result); + } + } + + /// Record tool call error. + pub fn record_tool_error(&mut self, error: impl Into) { + if let Some(call) = self.tool_calls.last_mut() { + call.error = Some(error.into()); + } + } +} + +/// Record of a tool call made during a turn. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TurnToolCall { + /// Tool name. + pub name: String, + /// Parameters passed to the tool. + pub parameters: serde_json::Value, + /// Result from the tool (if successful). + pub result: Option, + /// Error from the tool (if failed). + pub error: Option, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_session_creation() { + let mut session = Session::new("user-123"); + assert!(session.active_thread.is_none()); + + session.create_thread(); + assert!(session.active_thread.is_some()); + } + + #[test] + fn test_thread_turns() { + let mut thread = Thread::new(Uuid::new_v4()); + + thread.start_turn("Hello"); + assert_eq!(thread.state, ThreadState::Processing); + assert_eq!(thread.turns.len(), 1); + + thread.complete_turn("Hi there!"); + assert_eq!(thread.state, ThreadState::Idle); + assert_eq!(thread.turns[0].response, Some("Hi there!".to_string())); + } + + #[test] + fn test_thread_messages() { + let mut thread = Thread::new(Uuid::new_v4()); + + thread.start_turn("First message"); + thread.complete_turn("First response"); + thread.start_turn("Second message"); + thread.complete_turn("Second response"); + + let messages = thread.messages(); + assert_eq!(messages.len(), 4); + } + + #[test] + fn test_turn_tool_calls() { + let mut turn = Turn::new(0, "Test input"); + turn.record_tool_call("echo", serde_json::json!({"message": "test"})); + turn.record_tool_result(serde_json::json!("test")); + + assert_eq!(turn.tool_calls.len(), 1); + assert!(turn.tool_calls[0].result.is_some()); + } +} diff --git a/src/agent/submission.rs b/src/agent/submission.rs new file mode 100644 index 00000000..a35057ea --- /dev/null +++ b/src/agent/submission.rs @@ -0,0 +1,203 @@ +//! Submission types for the turn-based agent loop. +//! +//! Submissions are the different types of input the agent can receive +//! and process as part of the turn-based development loop. + +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +/// A submission to the agent. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum Submission { + /// User text input (starts a new turn). + UserInput { + /// The user's message content. + content: String, + }, + + /// Response to an execution approval request. + ExecApproval { + /// ID of the approval request being responded to. + request_id: Uuid, + /// Whether the execution was approved. + approved: bool, + /// If true, auto-approve this tool for the rest of the session. + always: bool, + }, + + /// Interrupt the current turn. + Interrupt, + + /// Request context compaction. + Compact, + + /// Undo the last turn. + Undo, + + /// Redo a previously undone turn (if available). + Redo, + + /// Resume from a specific checkpoint. + Resume { + /// ID of the checkpoint to resume from. + checkpoint_id: Uuid, + }, + + /// Clear the current thread and start fresh. + Clear, + + /// Switch to a different thread. + SwitchThread { + /// ID of the thread to switch to. + thread_id: Uuid, + }, + + /// Create a new thread. + NewThread, +} + +impl Submission { + /// Create a user input submission. + pub fn user_input(content: impl Into) -> Self { + Self::UserInput { + content: content.into(), + } + } + + /// Create an approval submission. + pub fn approval(request_id: Uuid, approved: bool) -> Self { + Self::ExecApproval { + request_id, + approved, + always: false, + } + } + + /// Create an "always approve" submission. + pub fn always_approve(request_id: Uuid) -> Self { + Self::ExecApproval { + request_id, + approved: true, + always: true, + } + } + + /// Create an interrupt submission. + pub fn interrupt() -> Self { + Self::Interrupt + } + + /// Create a compact submission. + pub fn compact() -> Self { + Self::Compact + } + + /// Create an undo submission. + pub fn undo() -> Self { + Self::Undo + } + + /// Create a redo submission. + pub fn redo() -> Self { + Self::Redo + } + + /// Check if this submission starts a new turn. + pub fn starts_turn(&self) -> bool { + matches!(self, Self::UserInput { .. }) + } + + /// Check if this submission is a control command. + pub fn is_control(&self) -> bool { + matches!( + self, + Self::Interrupt + | Self::Compact + | Self::Undo + | Self::Redo + | Self::Clear + | Self::NewThread + ) + } +} + +/// Result of processing a submission. +#[derive(Debug, Clone)] +pub enum SubmissionResult { + /// Turn completed with a response. + Response { + /// The agent's response. + content: String, + }, + + /// Need approval before continuing. + NeedApproval { + /// ID of the approval request. + request_id: Uuid, + /// Tool that needs approval. + tool_name: String, + /// Description of what the tool will do. + description: String, + /// Parameters being passed. + parameters: serde_json::Value, + }, + + /// Successfully processed (for control commands). + Ok { + /// Optional message. + message: Option, + }, + + /// Error occurred. + Error { + /// Error message. + message: String, + }, + + /// Turn was interrupted. + Interrupted, +} + +impl SubmissionResult { + /// Create a response result. + pub fn response(content: impl Into) -> Self { + Self::Response { + content: content.into(), + } + } + + /// Create an OK result. + pub fn ok() -> Self { + Self::Ok { message: None } + } + + /// Create an OK result with a message. + pub fn ok_with_message(message: impl Into) -> Self { + Self::Ok { + message: Some(message.into()), + } + } + + /// Create an error result. + pub fn error(message: impl Into) -> Self { + Self::Error { + message: message.into(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_submission_types() { + let input = Submission::user_input("Hello"); + assert!(input.starts_turn()); + assert!(!input.is_control()); + + let undo = Submission::undo(); + assert!(!undo.starts_turn()); + assert!(undo.is_control()); + } +} diff --git a/src/agent/task.rs b/src/agent/task.rs new file mode 100644 index 00000000..ba5e359c --- /dev/null +++ b/src/agent/task.rs @@ -0,0 +1,276 @@ +//! Task types for the scheduler. +//! +//! Tasks are the unit of work that can be scheduled for execution. +//! They can represent full LLM-driven jobs, parallel tool batches, +//! or background computations. + +use std::fmt; +use std::time::Duration; + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::error::Error; + +/// Result of a task execution. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskOutput { + /// The result data. + pub result: serde_json::Value, + /// Time taken to execute. + pub duration: Duration, +} + +impl TaskOutput { + /// Create a new task output. + pub fn new(result: serde_json::Value, duration: Duration) -> Self { + Self { result, duration } + } + + /// Create a text result. + pub fn text(text: impl Into, duration: Duration) -> Self { + Self { + result: serde_json::Value::String(text.into()), + duration, + } + } + + /// Create an empty success result. + pub fn empty(duration: Duration) -> Self { + Self { + result: serde_json::Value::Null, + duration, + } + } +} + +/// Context passed to task handlers. +#[derive(Debug, Clone)] +pub struct TaskContext { + /// Task ID. + pub task_id: Uuid, + /// Parent task ID (if this is a sub-task). + pub parent_id: Option, + /// Arbitrary metadata for the task. + pub metadata: serde_json::Value, +} + +impl TaskContext { + /// Create a new task context. + pub fn new(task_id: Uuid) -> Self { + Self { + task_id, + parent_id: None, + metadata: serde_json::Value::Null, + } + } + + /// Set the parent task ID. + pub fn with_parent(mut self, parent_id: Uuid) -> Self { + self.parent_id = Some(parent_id); + self + } + + /// Set metadata. + pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self { + self.metadata = metadata; + self + } +} + +/// Handler for custom background tasks. +#[async_trait] +pub trait TaskHandler: Send + Sync { + /// Run the task and return the result. + async fn run(&self, ctx: TaskContext) -> Result; + + /// Get a description of this handler for logging. + fn description(&self) -> &str { + "background task" + } +} + +/// A task that can be scheduled for execution. +#[derive(Clone)] +pub enum Task { + /// Full LLM-driven job (current Worker behavior). + Job { + id: Uuid, + title: String, + description: String, + }, + + /// Single tool execution as a sub-task. + ToolExec { + /// ID of the parent job this tool execution belongs to. + parent_id: Uuid, + /// Name of the tool to execute. + tool_name: String, + /// Parameters to pass to the tool. + params: serde_json::Value, + }, + + /// Background computation (no LLM, uses a custom handler). + /// Note: The handler is wrapped in Arc for cloning. + Background { + id: Uuid, + handler: std::sync::Arc, + }, +} + +impl Task { + /// Create a new Job task. + pub fn job(title: impl Into, description: impl Into) -> Self { + Self::Job { + id: Uuid::new_v4(), + title: title.into(), + description: description.into(), + } + } + + /// Create a new Job task with a specific ID. + pub fn job_with_id(id: Uuid, title: impl Into, description: impl Into) -> Self { + Self::Job { + id, + title: title.into(), + description: description.into(), + } + } + + /// Create a new ToolExec task. + pub fn tool_exec( + parent_id: Uuid, + tool_name: impl Into, + params: serde_json::Value, + ) -> Self { + Self::ToolExec { + parent_id, + tool_name: tool_name.into(), + params, + } + } + + /// Create a new Background task. + pub fn background(handler: std::sync::Arc) -> Self { + Self::Background { + id: Uuid::new_v4(), + handler, + } + } + + /// Create a new Background task with a specific ID. + pub fn background_with_id(id: Uuid, handler: std::sync::Arc) -> Self { + Self::Background { id, handler } + } + + /// Get the task ID, if applicable. + pub fn id(&self) -> Option { + match self { + Self::Job { id, .. } => Some(*id), + Self::ToolExec { .. } => None, // Tool execs don't have their own ID + Self::Background { id, .. } => Some(*id), + } + } + + /// Get the parent ID for sub-tasks. + pub fn parent_id(&self) -> Option { + match self { + Self::Job { .. } => None, + Self::ToolExec { parent_id, .. } => Some(*parent_id), + Self::Background { .. } => None, + } + } + + /// Get a short description for logging. + pub fn description(&self) -> String { + match self { + Self::Job { title, .. } => format!("job: {}", title), + Self::ToolExec { tool_name, .. } => format!("tool: {}", tool_name), + Self::Background { handler, .. } => format!("background: {}", handler.description()), + } + } +} + +impl fmt::Debug for Task { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Job { + id, + title, + description, + } => f + .debug_struct("Task::Job") + .field("id", id) + .field("title", title) + .field("description", description) + .finish(), + Self::ToolExec { + parent_id, + tool_name, + params, + } => f + .debug_struct("Task::ToolExec") + .field("parent_id", parent_id) + .field("tool_name", tool_name) + .field("params", params) + .finish(), + Self::Background { id, handler } => f + .debug_struct("Task::Background") + .field("id", id) + .field("handler", &handler.description()) + .finish(), + } + } +} + +/// Status of a scheduled task. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TaskStatus { + /// Task is queued waiting for execution. + Queued, + /// Task is currently running. + Running, + /// Task completed successfully. + Completed, + /// Task failed with an error. + Failed, + /// Task was cancelled. + Cancelled, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_task_output() { + let output = TaskOutput::text("hello", Duration::from_secs(1)); + assert_eq!(output.result, serde_json::json!("hello")); + assert_eq!(output.duration, Duration::from_secs(1)); + } + + #[test] + fn test_task_context() { + let parent = Uuid::new_v4(); + let ctx = TaskContext::new(Uuid::new_v4()).with_parent(parent); + assert_eq!(ctx.parent_id, Some(parent)); + } + + #[test] + fn test_task_job() { + let task = Task::job("Test Job", "Test description"); + assert!(task.id().is_some()); + assert!(task.parent_id().is_none()); + assert!(task.description().contains("job:")); + } + + #[test] + fn test_task_tool_exec() { + let parent_id = Uuid::new_v4(); + let task = Task::tool_exec(parent_id, "echo", serde_json::json!({"message": "hi"})); + assert!(task.id().is_none()); + assert_eq!(task.parent_id(), Some(parent_id)); + assert!(task.description().contains("tool:")); + } +} diff --git a/src/agent/undo.rs b/src/agent/undo.rs new file mode 100644 index 00000000..5307a77a --- /dev/null +++ b/src/agent/undo.rs @@ -0,0 +1,252 @@ +//! Undo system with checkpoints. +//! +//! Provides the ability to roll back the conversation state to a previous point. +//! Checkpoints are created automatically at the start of each turn. + +use std::collections::VecDeque; + +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::llm::ChatMessage; + +/// Maximum number of checkpoints to keep by default. +const DEFAULT_MAX_CHECKPOINTS: usize = 20; + +/// A snapshot of conversation state at a point in time. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Checkpoint { + /// Unique checkpoint ID. + pub id: Uuid, + /// Turn number this checkpoint was created at. + pub turn_number: usize, + /// Snapshot of messages at this point. + pub messages: Vec, + /// Description of what happened at this checkpoint. + pub description: String, +} + +impl Checkpoint { + /// Create a new checkpoint. + pub fn new( + turn_number: usize, + messages: Vec, + description: impl Into, + ) -> Self { + Self { + id: Uuid::new_v4(), + turn_number, + messages, + description: description.into(), + } + } +} + +/// Manager for undo/redo functionality. +pub struct UndoManager { + /// Stack of past checkpoints (for undo). + undo_stack: VecDeque, + /// Stack of future checkpoints (for redo). + redo_stack: Vec, + /// Maximum checkpoints to keep. + max_checkpoints: usize, +} + +impl UndoManager { + /// Create a new undo manager. + pub fn new() -> Self { + Self { + undo_stack: VecDeque::new(), + redo_stack: Vec::new(), + max_checkpoints: DEFAULT_MAX_CHECKPOINTS, + } + } + + /// Create with a custom checkpoint limit. + pub fn with_max_checkpoints(mut self, max: usize) -> Self { + self.max_checkpoints = max; + self + } + + /// Create a checkpoint at the current state. + /// + /// This clears the redo stack since we're creating a new history branch. + pub fn checkpoint( + &mut self, + turn_number: usize, + messages: Vec, + description: impl Into, + ) { + // Clear redo stack (new branch of history) + self.redo_stack.clear(); + + // Create and push checkpoint + let checkpoint = Checkpoint::new(turn_number, messages, description); + self.undo_stack.push_back(checkpoint); + + // Trim if over limit + while self.undo_stack.len() > self.max_checkpoints { + self.undo_stack.pop_front(); + } + } + + /// Undo: pop the last checkpoint and return it. + /// + /// The current state should be saved to redo stack before calling this. + pub fn undo( + &mut self, + current_turn: usize, + current_messages: Vec, + ) -> Option<&Checkpoint> { + if self.undo_stack.is_empty() { + return None; + } + + // Save current state to redo stack + let current = Checkpoint::new( + current_turn, + current_messages, + format!("Turn {}", current_turn), + ); + self.redo_stack.push(current); + + // Return the most recent checkpoint without removing it + // (we keep it so multiple undos can work) + self.undo_stack.back() + } + + /// Pop the last checkpoint from the undo stack. + pub fn pop_undo(&mut self) -> Option { + self.undo_stack.pop_back() + } + + /// Redo: restore a previously undone state. + pub fn redo(&mut self) -> Option { + self.redo_stack.pop() + } + + /// Check if undo is available. + pub fn can_undo(&self) -> bool { + !self.undo_stack.is_empty() + } + + /// Check if redo is available. + pub fn can_redo(&self) -> bool { + !self.redo_stack.is_empty() + } + + /// Get the number of undo steps available. + pub fn undo_count(&self) -> usize { + self.undo_stack.len() + } + + /// Get the number of redo steps available. + pub fn redo_count(&self) -> usize { + self.redo_stack.len() + } + + /// Get a checkpoint by ID. + pub fn get_checkpoint(&self, id: Uuid) -> Option<&Checkpoint> { + self.undo_stack + .iter() + .find(|c| c.id == id) + .or_else(|| self.redo_stack.iter().find(|c| c.id == id)) + } + + /// List all available checkpoints (for UI display). + pub fn list_checkpoints(&self) -> Vec<&Checkpoint> { + self.undo_stack.iter().collect() + } + + /// Clear all checkpoints. + pub fn clear(&mut self) { + self.undo_stack.clear(); + self.redo_stack.clear(); + } + + /// Restore to a specific checkpoint by ID. + /// + /// This invalidates all checkpoints after this one. + pub fn restore(&mut self, checkpoint_id: Uuid) -> Option { + // Find the checkpoint position + let pos = self.undo_stack.iter().position(|c| c.id == checkpoint_id)?; + + // Clear redo stack + self.redo_stack.clear(); + + // Remove all checkpoints after this one + while self.undo_stack.len() > pos + 1 { + self.undo_stack.pop_back(); + } + + // Pop and return the target checkpoint + self.undo_stack.pop_back() + } +} + +impl Default for UndoManager { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_checkpoint_creation() { + let mut manager = UndoManager::new(); + + manager.checkpoint(0, vec![], "Initial state"); + manager.checkpoint(1, vec![ChatMessage::user("Hello")], "Turn 1"); + + assert_eq!(manager.undo_count(), 2); + } + + #[test] + fn test_undo_redo() { + let mut manager = UndoManager::new(); + + manager.checkpoint(0, vec![], "Turn 0"); + manager.checkpoint(1, vec![ChatMessage::user("Hello")], "Turn 1"); + + assert!(manager.can_undo()); + assert!(!manager.can_redo()); + + // Undo + let current = vec![ChatMessage::user("Hello"), ChatMessage::assistant("Hi")]; + let checkpoint = manager.undo(2, current); + assert!(checkpoint.is_some()); + assert!(manager.can_redo()); + + // Redo + let restored = manager.redo(); + assert!(restored.is_some()); + } + + #[test] + fn test_max_checkpoints() { + let mut manager = UndoManager::new().with_max_checkpoints(3); + + for i in 0..5 { + manager.checkpoint(i, vec![], format!("Turn {}", i)); + } + + assert_eq!(manager.undo_count(), 3); + } + + #[test] + fn test_restore_to_checkpoint() { + let mut manager = UndoManager::new(); + + manager.checkpoint(0, vec![], "Turn 0"); + let checkpoint_id = manager.undo_stack.back().unwrap().id; + manager.checkpoint(1, vec![], "Turn 1"); + manager.checkpoint(2, vec![], "Turn 2"); + + let restored = manager.restore(checkpoint_id); + assert!(restored.is_some()); + assert_eq!(manager.undo_count(), 0); + } +} diff --git a/src/agent/worker.rs b/src/agent/worker.rs index 780bd9f2..d155dd2d 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -3,14 +3,16 @@ use std::sync::Arc; use std::time::Duration; +use futures::future::join_all; use tokio::sync::mpsc; use uuid::Uuid; use crate::agent::scheduler::WorkerMessage; +use crate::agent::task::TaskOutput; use crate::context::{ContextManager, JobState}; use crate::error::Error; use crate::history::Store; -use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext}; +use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, ToolSelection}; use crate::safety::SafetyLayer; use crate::tools::ToolRegistry; @@ -25,6 +27,13 @@ pub struct Worker { timeout: Duration, } +/// Result of a tool execution with metadata for context building. +struct ToolExecResult { + tool_name: String, + result: Result, + duration: Duration, +} + impl Worker { /// Create a new worker. pub fn new( @@ -97,6 +106,7 @@ Job: {} Description: {} You have access to tools to complete this job. Plan your approach and execute tools as needed. +You may request multiple tools at once if they can be executed in parallel. Report when the job is complete or if you encounter issues you cannot resolve."#, job_ctx.title, job_ctx.description ))); @@ -155,88 +165,60 @@ Report when the job is complete or if you encounter issues you cannot resolve."# return Ok(()); } - // Select next tool to use - let selection = reasoning.select_tool(reason_ctx).await?; + // Select next tool(s) to use + let selections = reasoning.select_tools(reason_ctx).await?; - match selection { - Some(tool_selection) => { - tracing::debug!( - "Job {} selecting tool: {} - {}", - self.job_id, - tool_selection.tool_name, - tool_selection.reasoning - ); + if selections.is_empty() { + // No tools selected, ask LLM for next steps + let response = reasoning.respond(reason_ctx).await?; - // Execute the tool - let result = self - .execute_tool(&tool_selection.tool_name, &tool_selection.parameters) - .await; - - // Record the result - match result { - Ok(output) => { - // Sanitize output - let sanitized = self - .safety - .sanitize_tool_output(&tool_selection.tool_name, &output); - - // Add to context - let wrapped = self.safety.wrap_for_llm( - &tool_selection.tool_name, - &sanitized.content, - sanitized.was_modified, - ); - - reason_ctx.messages.push(ChatMessage::tool_result( - "tool_call_id", - &tool_selection.tool_name, - wrapped, - )); - - // Check if job is complete - if output.contains("TASK_COMPLETE") || output.contains("JOB_DONE") { - self.mark_completed().await?; - return Ok(()); - } - } - Err(e) => { - tracing::warn!( - "Tool {} failed for job {}: {}", - tool_selection.tool_name, - self.job_id, - e - ); - - reason_ctx.messages.push(ChatMessage::tool_result( - "tool_call_id", - &tool_selection.tool_name, - format!("Error: {}", e), - )); - } - } + if response.to_lowercase().contains("complete") + || response.to_lowercase().contains("finished") + || response.to_lowercase().contains("done") + { + self.mark_completed().await?; + return Ok(()); } - None => { - // No tool selected, ask LLM for next steps - let response = reasoning.respond(reason_ctx).await?; - if response.to_lowercase().contains("complete") - || response.to_lowercase().contains("finished") - || response.to_lowercase().contains("done") - { - self.mark_completed().await?; - return Ok(()); - } + // Add assistant response to context + reason_ctx.messages.push(ChatMessage::assistant(&response)); - // Add assistant response to context - reason_ctx.messages.push(ChatMessage::assistant(&response)); + // Give it one more chance to select a tool + if iteration > 3 && iteration % 5 == 0 { + reason_ctx.messages.push(ChatMessage::user( + "Are you stuck? Do you need help completing this job?", + )); + } + } else if selections.len() == 1 { + // Single tool: execute directly + let selection = &selections[0]; + tracing::debug!( + "Job {} selecting tool: {} - {}", + self.job_id, + selection.tool_name, + selection.reasoning + ); - // Give it one more chance to select a tool - if iteration > 3 && iteration % 5 == 0 { - // Ask if stuck - reason_ctx.messages.push(ChatMessage::user( - "Are you stuck? Do you need help completing this job?", - )); - } + let result = self + .execute_tool(&selection.tool_name, &selection.parameters) + .await; + + self.process_tool_result(reason_ctx, selection, result) + .await?; + } else { + // Multiple tools: execute in parallel + tracing::debug!( + "Job {} executing {} tools in parallel", + self.job_id, + selections.len() + ); + + let results = self.execute_tools_parallel(&selections).await; + + // Process all results + for (selection, result) in selections.iter().zip(results) { + self.process_tool_result(reason_ctx, selection, result.result) + .await?; } } @@ -245,21 +227,59 @@ Report when the job is complete or if you encounter issues you cannot resolve."# } } - async fn execute_tool( - &self, + /// Execute multiple tools in parallel. + async fn execute_tools_parallel(&self, selections: &[ToolSelection]) -> Vec { + let futures: Vec<_> = selections + .iter() + .map(|selection| { + let tool_name = selection.tool_name.clone(); + let params = selection.parameters.clone(); + let tools = self.tools.clone(); + let context_manager = self.context_manager.clone(); + let job_id = self.job_id; + let store = self.store.clone(); + + async move { + let start = std::time::Instant::now(); + let result = Self::execute_tool_inner( + tools, + context_manager, + store, + job_id, + &tool_name, + ¶ms, + ) + .await; + ToolExecResult { + tool_name, + result, + duration: start.elapsed(), + } + } + }) + .collect(); + + join_all(futures).await + } + + /// Inner tool execution logic that can be called from both single and parallel paths. + async fn execute_tool_inner( + tools: Arc, + context_manager: Arc, + store: Option>, + job_id: Uuid, tool_name: &str, params: &serde_json::Value, ) -> Result { - let tool = - self.tools - .get(tool_name) - .await - .ok_or_else(|| crate::error::ToolError::NotFound { - name: tool_name.to_string(), - })?; + let tool = tools + .get(tool_name) + .await + .ok_or_else(|| crate::error::ToolError::NotFound { + name: tool_name.to_string(), + })?; // Get job context for the tool - let job_ctx = self.context_manager.get_context(self.job_id).await?; + let job_ctx = context_manager.get_context(job_id).await?; // Execute with timeout and timing let start = std::time::Instant::now(); @@ -273,8 +293,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."# let action = match &result { Ok(Ok(output)) => { let output_str = serde_json::to_string_pretty(&output.result).ok(); - self.context_manager - .update_memory(self.job_id, |mem| { + context_manager + .update_memory(job_id, |mem| { let rec = mem.create_action(tool_name, params.clone()).succeed( output_str.clone(), output.result.clone(), @@ -286,9 +306,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."# .await .ok() } - Ok(Err(e)) => self - .context_manager - .update_memory(self.job_id, |mem| { + Ok(Err(e)) => context_manager + .update_memory(job_id, |mem| { let rec = mem .create_action(tool_name, params.clone()) .fail(e.to_string(), elapsed); @@ -297,9 +316,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."# }) .await .ok(), - Err(_) => self - .context_manager - .update_memory(self.job_id, |mem| { + Err(_) => context_manager + .update_memory(job_id, |mem| { let rec = mem .create_action(tool_name, params.clone()) .fail("Execution timeout", elapsed); @@ -311,9 +329,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# }; // Persist action to database (fire-and-forget) - if let (Some(action), Some(store)) = (action, &self.store) { - let store = store.clone(); - let job_id = self.job_id; + if let (Some(action), Some(store)) = (action, store) { tokio::spawn(async move { if let Err(e) = store.save_action(job_id, &action).await { tracing::warn!("Failed to persist action for job {}: {}", job_id, e); @@ -342,6 +358,76 @@ Report when the job is complete or if you encounter issues you cannot resolve."# }) } + /// Process a tool execution result and add it to the reasoning context. + async fn process_tool_result( + &self, + reason_ctx: &mut ReasoningContext, + selection: &ToolSelection, + result: Result, + ) -> Result { + match result { + Ok(output) => { + // Sanitize output + let sanitized = self + .safety + .sanitize_tool_output(&selection.tool_name, &output); + + // Add to context + let wrapped = self.safety.wrap_for_llm( + &selection.tool_name, + &sanitized.content, + sanitized.was_modified, + ); + + reason_ctx.messages.push(ChatMessage::tool_result( + "tool_call_id", + &selection.tool_name, + wrapped, + )); + + // Check if job is complete + if output.contains("TASK_COMPLETE") || output.contains("JOB_DONE") { + self.mark_completed().await?; + return Ok(true); + } + + Ok(false) + } + Err(e) => { + tracing::warn!( + "Tool {} failed for job {}: {}", + selection.tool_name, + self.job_id, + e + ); + + reason_ctx.messages.push(ChatMessage::tool_result( + "tool_call_id", + &selection.tool_name, + format!("Error: {}", e), + )); + + Ok(false) + } + } + } + + async fn execute_tool( + &self, + tool_name: &str, + params: &serde_json::Value, + ) -> Result { + Self::execute_tool_inner( + self.tools.clone(), + self.context_manager.clone(), + self.store.clone(), + self.job_id, + tool_name, + params, + ) + .await + } + async fn mark_completed(&self) -> Result<(), Error> { self.context_manager .update_context(self.job_id, |ctx| { @@ -391,3 +477,16 @@ Report when the job is complete or if you encounter issues you cannot resolve."# Ok(()) } } + +/// Convert a TaskOutput to a string result for tool execution. +impl From for Result { + fn from(output: TaskOutput) -> Self { + serde_json::to_string_pretty(&output.result).map_err(|e| { + crate::error::ToolError::ExecutionFailed { + name: "task".to_string(), + reason: format!("Failed to serialize result: {}", e), + } + .into() + }) + } +} diff --git a/src/channels/cli.rs b/src/channels/cli.rs deleted file mode 100644 index 49d58794..00000000 --- a/src/channels/cli.rs +++ /dev/null @@ -1,124 +0,0 @@ -//! CLI/stdin channel for interactive terminal usage. - -use std::io::{self, BufRead, Write}; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; - -use async_trait::async_trait; -use tokio::sync::mpsc; -use tokio_stream::wrappers::ReceiverStream; - -use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse}; -use crate::error::ChannelError; - -/// CLI channel for interactive terminal input. -pub struct CliChannel { - running: Arc, -} - -impl CliChannel { - /// Create a new CLI channel. - pub fn new() -> Self { - Self { - running: Arc::new(AtomicBool::new(false)), - } - } -} - -impl Default for CliChannel { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl Channel for CliChannel { - fn name(&self) -> &str { - "cli" - } - - async fn start(&self) -> Result { - self.running.store(true, Ordering::SeqCst); - let running = self.running.clone(); - - let (tx, rx) = mpsc::channel(32); - - // Spawn a blocking task to read from stdin - tokio::task::spawn_blocking(move || { - let stdin = io::stdin(); - let reader = stdin.lock(); - - // Print prompt - print_prompt(); - - for line in reader.lines() { - if !running.load(Ordering::SeqCst) { - break; - } - - match line { - Ok(content) => { - let content = content.trim(); - if content.is_empty() { - print_prompt(); - continue; - } - - // Handle exit commands - if content == "exit" || content == "quit" || content == "/quit" { - running.store(false, Ordering::SeqCst); - break; - } - - let msg = IncomingMessage::new("cli", "local-user", content); - - if tx.blocking_send(msg).is_err() { - // Channel closed, stop reading - break; - } - } - Err(e) => { - tracing::error!("Error reading stdin: {}", e); - break; - } - } - } - - tracing::debug!("CLI input loop ended"); - }); - - Ok(Box::pin(ReceiverStream::new(rx))) - } - - async fn respond( - &self, - _msg: &IncomingMessage, - response: OutgoingResponse, - ) -> Result<(), ChannelError> { - // Print response to stdout - println!("\n{}\n", response.content); - print_prompt(); - Ok(()) - } - - async fn health_check(&self) -> Result<(), ChannelError> { - // CLI is always healthy if we're running - if self.running.load(Ordering::SeqCst) { - Ok(()) - } else { - Err(ChannelError::HealthCheckFailed { - name: "cli".to_string(), - }) - } - } - - async fn shutdown(&self) -> Result<(), ChannelError> { - self.running.store(false, Ordering::SeqCst); - Ok(()) - } -} - -fn print_prompt() { - print!("agent> "); - let _ = io::stdout().flush(); -} diff --git a/src/channels/cli/app.rs b/src/channels/cli/app.rs new file mode 100644 index 00000000..2f7e03ab --- /dev/null +++ b/src/channels/cli/app.rs @@ -0,0 +1,252 @@ +//! Application state for the TUI. + +use std::collections::VecDeque; + +use crate::channels::cli::composer::ChatComposer; +use crate::channels::cli::overlay::{ApprovalOverlay, ApprovalRequest}; + +/// Events that can occur in the TUI. +#[derive(Debug, Clone)] +pub enum AppEvent { + /// Keyboard/mouse input event. + Input(crossterm::event::Event), + /// Response from the agent. + Response(String), + /// Tool execution started. + ToolStarted { name: String }, + /// Tool execution completed. + ToolCompleted { name: String, success: bool }, + /// Request approval for a tool. + ApprovalRequested(ApprovalRequest), + /// Streaming chunk received. + StreamChunk(String), + /// Force a redraw. + Redraw, + /// Quit the application. + Quit, +} + +/// Current input mode. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InputMode { + /// Normal input mode. + Normal, + /// Editing input. + Editing, + /// Approval overlay is active. + Approval, +} + +/// Message in the chat history. +#[derive(Debug, Clone)] +pub struct ChatMessage { + /// Who sent this message. + pub role: MessageRole, + /// The message content. + pub content: String, + /// Optional status indicator. + pub status: Option, +} + +/// Who sent a message. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MessageRole { + User, + Agent, + System, +} + +/// Status of a message (for in-progress indicators). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MessageStatus { + Pending, + InProgress, + Complete, + Error, +} + +impl ChatMessage { + pub fn user(content: impl Into) -> Self { + Self { + role: MessageRole::User, + content: content.into(), + status: None, + } + } + + pub fn agent(content: impl Into) -> Self { + Self { + role: MessageRole::Agent, + content: content.into(), + status: None, + } + } + + pub fn system(content: impl Into) -> Self { + Self { + role: MessageRole::System, + content: content.into(), + status: None, + } + } + + pub fn with_status(mut self, status: MessageStatus) -> Self { + self.status = Some(status); + self + } +} + +/// Application state. +pub struct AppState { + /// Current input mode. + pub mode: InputMode, + /// Chat message history. + pub messages: Vec, + /// Input composer. + pub composer: ChatComposer, + /// Approval overlay (if active). + pub approval: Option, + /// Scroll offset for messages. + pub scroll_offset: u16, + /// Whether the app should quit. + pub should_quit: bool, + /// Pending approvals queue. + pub pending_approvals: VecDeque, + /// Current streaming response buffer. + pub streaming_buffer: Option, + /// Status line message. + pub status_message: Option, +} + +impl AppState { + /// Create a new app state. + pub fn new() -> Self { + Self { + mode: InputMode::Editing, + messages: vec![ChatMessage::system( + "Welcome to NEAR Agent. Type a message or /help for commands.", + )], + composer: ChatComposer::new(), + approval: None, + scroll_offset: 0, + should_quit: false, + pending_approvals: VecDeque::new(), + streaming_buffer: None, + status_message: None, + } + } + + /// Add a user message to history. + pub fn add_user_message(&mut self, content: impl Into) { + self.messages.push(ChatMessage::user(content)); + self.scroll_to_bottom(); + } + + /// Add an agent response to history. + pub fn add_agent_message(&mut self, content: impl Into) { + // If we were streaming, finalize it + if self.streaming_buffer.is_some() { + self.streaming_buffer = None; + } + self.messages.push(ChatMessage::agent(content)); + self.scroll_to_bottom(); + } + + /// Start streaming a response. + pub fn start_streaming(&mut self) { + self.streaming_buffer = Some(String::new()); + self.messages + .push(ChatMessage::agent("").with_status(MessageStatus::InProgress)); + } + + /// Append to the streaming buffer. + pub fn append_stream(&mut self, chunk: &str) { + if let Some(ref mut buffer) = self.streaming_buffer { + buffer.push_str(chunk); + // Update the last message + if let Some(last) = self.messages.last_mut() { + if last.role == MessageRole::Agent { + last.content = buffer.clone(); + } + } + } + } + + /// Finalize streaming. + pub fn finish_streaming(&mut self) { + if let Some(last) = self.messages.last_mut() { + if last.role == MessageRole::Agent { + last.status = Some(MessageStatus::Complete); + } + } + self.streaming_buffer = None; + } + + /// Show an approval request. + pub fn show_approval(&mut self, request: ApprovalRequest) { + self.approval = Some(ApprovalOverlay::new(request)); + self.mode = InputMode::Approval; + } + + /// Queue an approval request. + pub fn queue_approval(&mut self, request: ApprovalRequest) { + if self.approval.is_none() { + self.show_approval(request); + } else { + self.pending_approvals.push_back(request); + } + } + + /// Handle approval response. + pub fn handle_approval_response(&mut self, approved: bool) -> Option { + let request = self.approval.take().map(|o| o.request); + + // Show next pending approval if any + if let Some(next) = self.pending_approvals.pop_front() { + self.show_approval(next); + } else { + self.mode = InputMode::Editing; + } + + if approved { request } else { None } + } + + /// Clear all pending approvals. + pub fn clear_approvals(&mut self) { + self.approval = None; + self.pending_approvals.clear(); + self.mode = InputMode::Editing; + } + + /// Set the status message. + pub fn set_status(&mut self, message: impl Into) { + self.status_message = Some(message.into()); + } + + /// Clear the status message. + pub fn clear_status(&mut self) { + self.status_message = None; + } + + /// Scroll to the bottom of messages. + pub fn scroll_to_bottom(&mut self) { + // Will be calculated based on render area in render.rs + self.scroll_offset = 0; + } + + /// Scroll up. + pub fn scroll_up(&mut self, amount: u16) { + self.scroll_offset = self.scroll_offset.saturating_add(amount); + } + + /// Scroll down. + pub fn scroll_down(&mut self, amount: u16) { + self.scroll_offset = self.scroll_offset.saturating_sub(amount); + } +} + +impl Default for AppState { + fn default() -> Self { + Self::new() + } +} diff --git a/src/channels/cli/composer.rs b/src/channels/cli/composer.rs new file mode 100644 index 00000000..01e87730 --- /dev/null +++ b/src/channels/cli/composer.rs @@ -0,0 +1,318 @@ +//! Input composer with history and completion. + +use std::collections::VecDeque; + +/// Maximum number of history entries to keep. +const MAX_HISTORY: usize = 100; + +/// Available slash commands for completion. +const SLASH_COMMANDS: &[&str] = &[ + "/help", "/job", "/status", "/cancel", "/list", "/tools", "/clear", "/quit", +]; + +/// Chat input composer with history navigation and slash command completion. +pub struct ChatComposer { + /// Current input buffer. + buffer: String, + /// Cursor position in the buffer. + cursor: usize, + /// Input history. + history: VecDeque, + /// Current position in history (-1 = current input). + history_index: Option, + /// Saved current input when navigating history. + saved_input: String, + /// Completion candidates. + completions: Vec, + /// Current completion index. + completion_index: Option, +} + +impl ChatComposer { + /// Create a new composer. + pub fn new() -> Self { + Self { + buffer: String::new(), + cursor: 0, + history: VecDeque::with_capacity(MAX_HISTORY), + history_index: None, + saved_input: String::new(), + completions: Vec::new(), + completion_index: None, + } + } + + /// Get the current input buffer. + pub fn buffer(&self) -> &str { + &self.buffer + } + + /// Get the cursor position. + pub fn cursor(&self) -> usize { + self.cursor + } + + /// Check if the buffer is empty. + pub fn is_empty(&self) -> bool { + self.buffer.is_empty() + } + + /// Insert a character at the cursor. + pub fn insert(&mut self, c: char) { + self.clear_completion(); + self.buffer.insert(self.cursor, c); + self.cursor += c.len_utf8(); + } + + /// Insert a string at the cursor. + pub fn insert_str(&mut self, s: &str) { + self.clear_completion(); + self.buffer.insert_str(self.cursor, s); + self.cursor += s.len(); + } + + /// Delete the character before the cursor (backspace). + pub fn backspace(&mut self) { + self.clear_completion(); + if self.cursor > 0 { + // Find the previous character boundary + let prev = self.buffer[..self.cursor] + .char_indices() + .next_back() + .map(|(i, _)| i) + .unwrap_or(0); + self.buffer.drain(prev..self.cursor); + self.cursor = prev; + } + } + + /// Delete the character at the cursor (delete). + pub fn delete(&mut self) { + self.clear_completion(); + if self.cursor < self.buffer.len() { + // Find the next character boundary + let next = self.buffer[self.cursor..] + .char_indices() + .nth(1) + .map(|(i, _)| self.cursor + i) + .unwrap_or(self.buffer.len()); + self.buffer.drain(self.cursor..next); + } + } + + /// Move cursor left. + pub fn move_left(&mut self) { + if self.cursor > 0 { + self.cursor = self.buffer[..self.cursor] + .char_indices() + .next_back() + .map(|(i, _)| i) + .unwrap_or(0); + } + } + + /// Move cursor right. + pub fn move_right(&mut self) { + if self.cursor < self.buffer.len() { + self.cursor = self.buffer[self.cursor..] + .char_indices() + .nth(1) + .map(|(i, _)| self.cursor + i) + .unwrap_or(self.buffer.len()); + } + } + + /// Move cursor to start. + pub fn move_home(&mut self) { + self.cursor = 0; + } + + /// Move cursor to end. + pub fn move_end(&mut self) { + self.cursor = self.buffer.len(); + } + + /// Delete from cursor to end of line. + pub fn kill_line(&mut self) { + self.clear_completion(); + self.buffer.truncate(self.cursor); + } + + /// Delete from start to cursor. + pub fn kill_to_start(&mut self) { + self.clear_completion(); + self.buffer.drain(..self.cursor); + self.cursor = 0; + } + + /// Clear the entire buffer. + pub fn clear(&mut self) { + self.buffer.clear(); + self.cursor = 0; + self.clear_completion(); + } + + /// Submit the current input and return it. + pub fn submit(&mut self) -> String { + let input = std::mem::take(&mut self.buffer); + self.cursor = 0; + self.clear_completion(); + + // Add to history if non-empty and different from last entry + if !input.is_empty() && self.history.front() != Some(&input) { + self.history.push_front(input.clone()); + if self.history.len() > MAX_HISTORY { + self.history.pop_back(); + } + } + + self.history_index = None; + self.saved_input.clear(); + + input + } + + /// Navigate to previous history entry. + pub fn history_prev(&mut self) { + if self.history.is_empty() { + return; + } + + match self.history_index { + None => { + // Save current input and go to first history entry + self.saved_input = std::mem::take(&mut self.buffer); + self.history_index = Some(0); + self.buffer = self.history[0].clone(); + } + Some(i) if i + 1 < self.history.len() => { + self.history_index = Some(i + 1); + self.buffer = self.history[i + 1].clone(); + } + _ => {} + } + + self.cursor = self.buffer.len(); + self.clear_completion(); + } + + /// Navigate to next history entry. + pub fn history_next(&mut self) { + match self.history_index { + Some(0) => { + // Go back to saved input + self.history_index = None; + self.buffer = std::mem::take(&mut self.saved_input); + } + Some(i) => { + self.history_index = Some(i - 1); + self.buffer = self.history[i - 1].clone(); + } + None => {} + } + + self.cursor = self.buffer.len(); + self.clear_completion(); + } + + /// Attempt tab completion. + pub fn complete(&mut self) { + // Only complete slash commands for now + if !self.buffer.starts_with('/') { + return; + } + + if self.completions.is_empty() { + // Generate completions + let prefix = &self.buffer; + self.completions = SLASH_COMMANDS + .iter() + .filter(|cmd| cmd.starts_with(prefix)) + .map(|s| s.to_string()) + .collect(); + + if !self.completions.is_empty() { + self.completion_index = Some(0); + } + } else if let Some(i) = self.completion_index { + // Cycle through completions + self.completion_index = Some((i + 1) % self.completions.len()); + } + + // Apply completion + if let Some(i) = self.completion_index { + if let Some(completion) = self.completions.get(i) { + self.buffer = completion.clone(); + self.cursor = self.buffer.len(); + } + } + } + + /// Clear completion state. + fn clear_completion(&mut self) { + self.completions.clear(); + self.completion_index = None; + } + + /// Get current completion hint (for display). + pub fn completion_hint(&self) -> Option<&str> { + if let Some(i) = self.completion_index { + self.completions.get(i).map(|s| s.as_str()) + } else { + None + } + } + + /// Get the number of completions available. + pub fn completion_count(&self) -> usize { + self.completions.len() + } +} + +impl Default for ChatComposer { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_insert_and_backspace() { + let mut composer = ChatComposer::new(); + composer.insert('h'); + composer.insert('i'); + assert_eq!(composer.buffer(), "hi"); + composer.backspace(); + assert_eq!(composer.buffer(), "h"); + } + + #[test] + fn test_history_navigation() { + let mut composer = ChatComposer::new(); + composer.insert_str("first"); + composer.submit(); + composer.insert_str("second"); + composer.submit(); + + composer.insert_str("current"); + composer.history_prev(); + assert_eq!(composer.buffer(), "second"); + composer.history_prev(); + assert_eq!(composer.buffer(), "first"); + composer.history_next(); + assert_eq!(composer.buffer(), "second"); + composer.history_next(); + assert_eq!(composer.buffer(), "current"); + } + + #[test] + fn test_completion() { + let mut composer = ChatComposer::new(); + composer.insert_str("/hel"); + composer.complete(); + assert_eq!(composer.buffer(), "/help"); + } +} diff --git a/src/channels/cli/events.rs b/src/channels/cli/events.rs new file mode 100644 index 00000000..aa2194ed --- /dev/null +++ b/src/channels/cli/events.rs @@ -0,0 +1,270 @@ +//! Event handling for the TUI. + +use std::io; +use std::time::Duration; + +use crossterm::event::{self, Event, KeyCode, KeyEvent, KeyModifiers}; +use ratatui::Terminal; +use ratatui::backend::CrosstermBackend; +use tokio::sync::mpsc; + +use crate::channels::IncomingMessage; +use crate::channels::cli::app::{AppEvent, AppState, InputMode}; +use crate::channels::cli::render; + +/// Tick rate for the event loop (50ms = 20fps). +const TICK_RATE: Duration = Duration::from_millis(50); + +/// Run the main event loop. +pub fn run_event_loop( + terminal: &mut Terminal>, + app: &mut AppState, + msg_tx: mpsc::Sender, + event_rx: &mut mpsc::Receiver, +) -> io::Result<()> { + loop { + // Render + terminal.draw(|f| render::render(f, app))?; + + // Check for quit + if app.should_quit { + return Ok(()); + } + + // Poll for events + if event::poll(TICK_RATE)? { + let evt = event::read()?; + if let Err(e) = handle_event(app, evt, &msg_tx) { + tracing::error!("Event handling error: {}", e); + } + } + + // Check for app events (non-blocking) + while let Ok(app_event) = event_rx.try_recv() { + handle_app_event(app, app_event); + } + } +} + +/// Handle a crossterm event. +fn handle_event( + app: &mut AppState, + event: Event, + msg_tx: &mpsc::Sender, +) -> io::Result<()> { + match event { + Event::Key(key) => handle_key(app, key, msg_tx), + Event::Mouse(_) => Ok(()), // Could handle mouse scrolling here + Event::Resize(_, _) => Ok(()), // Terminal will handle resize + _ => Ok(()), + } +} + +/// Handle a key event. +fn handle_key( + app: &mut AppState, + key: KeyEvent, + msg_tx: &mpsc::Sender, +) -> io::Result<()> { + // Global keybindings + if key.modifiers.contains(KeyModifiers::CONTROL) { + match key.code { + KeyCode::Char('c') => { + if app.mode == InputMode::Approval { + // Cancel all pending approvals + app.clear_approvals(); + } else { + // Quit + app.should_quit = true; + } + return Ok(()); + } + KeyCode::Char('d') => { + app.should_quit = true; + return Ok(()); + } + _ => {} + } + } + + match app.mode { + InputMode::Normal => handle_normal_mode(app, key), + InputMode::Editing => handle_editing_mode(app, key, msg_tx), + InputMode::Approval => handle_approval_mode(app, key), + } +} + +/// Handle keys in normal mode. +fn handle_normal_mode(app: &mut AppState, key: KeyEvent) -> io::Result<()> { + match key.code { + KeyCode::Char('i') | KeyCode::Char('a') => { + app.mode = InputMode::Editing; + } + KeyCode::Char('q') => { + app.should_quit = true; + } + KeyCode::Up | KeyCode::Char('k') => { + app.scroll_up(1); + } + KeyCode::Down | KeyCode::Char('j') => { + app.scroll_down(1); + } + KeyCode::PageUp => { + app.scroll_up(10); + } + KeyCode::PageDown => { + app.scroll_down(10); + } + KeyCode::Char('G') => { + app.scroll_to_bottom(); + } + _ => {} + } + Ok(()) +} + +/// Handle keys in editing mode. +fn handle_editing_mode( + app: &mut AppState, + key: KeyEvent, + msg_tx: &mpsc::Sender, +) -> io::Result<()> { + match key.code { + KeyCode::Enter => { + if !app.composer.is_empty() { + let input = app.composer.submit(); + app.add_user_message(&input); + + // Send message to agent + let msg = IncomingMessage::new("tui", "local-user", &input); + let _ = msg_tx.blocking_send(msg); + } + } + KeyCode::Esc => { + app.mode = InputMode::Normal; + } + KeyCode::Backspace => { + app.composer.backspace(); + } + KeyCode::Delete => { + app.composer.delete(); + } + KeyCode::Left => { + if key.modifiers.contains(KeyModifiers::CONTROL) { + // Move word left (simplified: just move to start) + app.composer.move_home(); + } else { + app.composer.move_left(); + } + } + KeyCode::Right => { + if key.modifiers.contains(KeyModifiers::CONTROL) { + // Move word right (simplified: just move to end) + app.composer.move_end(); + } else { + app.composer.move_right(); + } + } + KeyCode::Home => { + app.composer.move_home(); + } + KeyCode::End => { + app.composer.move_end(); + } + KeyCode::Up => { + app.composer.history_prev(); + } + KeyCode::Down => { + app.composer.history_next(); + } + KeyCode::Tab => { + app.composer.complete(); + } + KeyCode::Char(c) => { + if key.modifiers.contains(KeyModifiers::CONTROL) { + match c { + 'a' => app.composer.move_home(), + 'e' => app.composer.move_end(), + 'k' => app.composer.kill_line(), + 'u' => app.composer.kill_to_start(), + 'w' => { + // Delete word backwards (simplified: clear) + app.composer.clear(); + } + _ => {} + } + } else { + app.composer.insert(c); + } + } + _ => {} + } + Ok(()) +} + +/// Handle keys in approval mode. +fn handle_approval_mode(app: &mut AppState, key: KeyEvent) -> io::Result<()> { + if let Some(ref mut overlay) = app.approval { + match key.code { + KeyCode::Left | KeyCode::Char('h') => { + overlay.select_prev(); + } + KeyCode::Right | KeyCode::Char('l') => { + overlay.select_next(); + } + KeyCode::Enter | KeyCode::Char(' ') => { + let (approved, _always) = overlay.confirm(); + app.handle_approval_response(approved); + // TODO: If always, remember to auto-approve this tool + } + KeyCode::Char(c) => { + if let Some(approved) = overlay.handle_shortcut(c) { + app.handle_approval_response(approved); + } + } + KeyCode::Esc => { + // Deny this approval + app.handle_approval_response(false); + } + _ => {} + } + } + Ok(()) +} + +/// Handle an application event. +fn handle_app_event(app: &mut AppState, event: AppEvent) { + match event { + AppEvent::Response(content) => { + app.add_agent_message(content); + } + AppEvent::ToolStarted { name } => { + app.set_status(format!("Running tool: {}...", name)); + } + AppEvent::ToolCompleted { name, success } => { + if success { + app.set_status(format!("Tool {} completed", name)); + } else { + app.set_status(format!("Tool {} failed", name)); + } + } + AppEvent::ApprovalRequested(request) => { + app.queue_approval(request); + } + AppEvent::StreamChunk(chunk) => { + if app.streaming_buffer.is_none() { + app.start_streaming(); + } + app.append_stream(&chunk); + } + AppEvent::Redraw => { + // Just triggers a redraw on next loop iteration + } + AppEvent::Quit => { + app.should_quit = true; + } + AppEvent::Input(_) => { + // Already handled directly + } + } +} diff --git a/src/channels/cli/mod.rs b/src/channels/cli/mod.rs new file mode 100644 index 00000000..98396890 --- /dev/null +++ b/src/channels/cli/mod.rs @@ -0,0 +1,251 @@ +//! Interactive TUI channel using Ratatui. +//! +//! Provides a rich terminal interface with: +//! - Input history navigation +//! - Slash command completion +//! - Approval overlays for tool execution +//! - Streaming response display + +mod app; +mod composer; +mod events; +mod overlay; +mod render; + +use std::io; +use std::sync::Arc; + +use async_trait::async_trait; +use crossterm::{ + event::{DisableMouseCapture, EnableMouseCapture}, + execute, + terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode}, +}; +use ratatui::Terminal; +use ratatui::backend::CrosstermBackend; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; + +use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse}; +use crate::error::ChannelError; + +pub use app::{AppEvent, AppState, InputMode}; +pub use composer::ChatComposer; +pub use overlay::{ApprovalOverlay, ApprovalRequest}; + +/// TUI channel for interactive terminal input with Ratatui. +pub struct TuiChannel { + /// Channel for sending events to the TUI. + event_tx: Option>, +} + +impl TuiChannel { + /// Create a new TUI channel. + pub fn new() -> Self { + Self { event_tx: None } + } +} + +impl Default for TuiChannel { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Channel for TuiChannel { + fn name(&self) -> &str { + "tui" + } + + async fn start(&self) -> Result { + let (msg_tx, msg_rx) = mpsc::channel(32); + let (event_tx, event_rx) = mpsc::channel(64); + + // Store the event sender so we can send responses + // Note: In the actual implementation, we'd store this properly + // For now, spawn the TUI in a separate task + let event_tx_clone = event_tx.clone(); + + tokio::task::spawn_blocking(move || { + if let Err(e) = run_tui(msg_tx, event_rx) { + tracing::error!("TUI error: {}", e); + } + }); + + // Keep the event_tx alive by storing it + // This is a hack; in production we'd use Arc> or similar + let _ = event_tx_clone; + + Ok(Box::pin(ReceiverStream::new(msg_rx))) + } + + async fn respond( + &self, + _msg: &IncomingMessage, + response: OutgoingResponse, + ) -> Result<(), ChannelError> { + // Send response event to the TUI + if let Some(ref tx) = self.event_tx { + let _ = tx + .send(AppEvent::Response(response.content)) + .await + .map_err(|e| ChannelError::SendFailed { + name: "tui".to_string(), + reason: e.to_string(), + }); + } + Ok(()) + } + + async fn health_check(&self) -> Result<(), ChannelError> { + Ok(()) + } + + async fn shutdown(&self) -> Result<(), ChannelError> { + if let Some(ref tx) = self.event_tx { + let _ = tx.send(AppEvent::Quit).await; + } + Ok(()) + } +} + +/// Run the TUI event loop (blocking). +fn run_tui( + msg_tx: mpsc::Sender, + mut event_rx: mpsc::Receiver, +) -> io::Result<()> { + // Setup terminal + enable_raw_mode()?; + let mut stdout = io::stdout(); + execute!(stdout, EnterAlternateScreen, EnableMouseCapture)?; + let backend = CrosstermBackend::new(stdout); + let mut terminal = Terminal::new(backend)?; + + // Create app state + let mut app = AppState::new(); + + // Run event loop + let result = events::run_event_loop(&mut terminal, &mut app, msg_tx, &mut event_rx); + + // Restore terminal + disable_raw_mode()?; + execute!( + terminal.backend_mut(), + LeaveAlternateScreen, + DisableMouseCapture + )?; + terminal.show_cursor()?; + + result +} + +/// Simple blocking CLI channel (fallback when TUI not available). +pub struct SimpleCliChannel { + running: Arc, +} + +impl SimpleCliChannel { + pub fn new() -> Self { + Self { + running: Arc::new(std::sync::atomic::AtomicBool::new(false)), + } + } +} + +impl Default for SimpleCliChannel { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Channel for SimpleCliChannel { + fn name(&self) -> &str { + "cli" + } + + async fn start(&self) -> Result { + self.running + .store(true, std::sync::atomic::Ordering::SeqCst); + let running = self.running.clone(); + + let (tx, rx) = mpsc::channel(32); + + tokio::task::spawn_blocking(move || { + use std::io::BufRead; + + let stdin = io::stdin(); + let reader = stdin.lock(); + + print_prompt(); + + for line in reader.lines() { + if !running.load(std::sync::atomic::Ordering::SeqCst) { + break; + } + + match line { + Ok(content) => { + let content = content.trim(); + if content.is_empty() { + print_prompt(); + continue; + } + + if content == "exit" || content == "quit" || content == "/quit" { + running.store(false, std::sync::atomic::Ordering::SeqCst); + break; + } + + let msg = IncomingMessage::new("cli", "local-user", content); + + if tx.blocking_send(msg).is_err() { + break; + } + } + Err(e) => { + tracing::error!("Error reading stdin: {}", e); + break; + } + } + } + + tracing::debug!("CLI input loop ended"); + }); + + Ok(Box::pin(ReceiverStream::new(rx))) + } + + async fn respond( + &self, + _msg: &IncomingMessage, + response: OutgoingResponse, + ) -> Result<(), ChannelError> { + println!("\n{}\n", response.content); + print_prompt(); + Ok(()) + } + + async fn health_check(&self) -> Result<(), ChannelError> { + if self.running.load(std::sync::atomic::Ordering::SeqCst) { + Ok(()) + } else { + Err(ChannelError::HealthCheckFailed { + name: "cli".to_string(), + }) + } + } + + async fn shutdown(&self) -> Result<(), ChannelError> { + self.running + .store(false, std::sync::atomic::Ordering::SeqCst); + Ok(()) + } +} + +fn print_prompt() { + use std::io::Write; + print!("agent> "); + let _ = io::stdout().flush(); +} diff --git a/src/channels/cli/overlay.rs b/src/channels/cli/overlay.rs new file mode 100644 index 00000000..1a8b4d97 --- /dev/null +++ b/src/channels/cli/overlay.rs @@ -0,0 +1,145 @@ +//! Approval overlay modal. + +use uuid::Uuid; + +/// A request for user approval before executing a tool. +#[derive(Debug, Clone)] +pub struct ApprovalRequest { + /// Unique ID for this request. + pub id: Uuid, + /// Name of the tool requesting approval. + pub tool_name: String, + /// Description of what the tool will do. + pub description: String, + /// Parameters being passed to the tool. + pub parameters: serde_json::Value, + /// Whether this is a destructive operation. + pub destructive: bool, +} + +impl ApprovalRequest { + /// Create a new approval request. + pub fn new( + tool_name: impl Into, + description: impl Into, + parameters: serde_json::Value, + ) -> Self { + Self { + id: Uuid::new_v4(), + tool_name: tool_name.into(), + description: description.into(), + parameters, + destructive: false, + } + } + + /// Mark as destructive operation. + pub fn destructive(mut self) -> Self { + self.destructive = true; + self + } +} + +/// Current selection in the approval overlay. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ApprovalSelection { + /// Yes, approve this action. + Yes, + /// No, deny this action. + No, + /// Always approve this tool (for this session). + Always, +} + +impl ApprovalSelection { + /// Get the next selection (cycling). + pub fn next(self) -> Self { + match self { + Self::Yes => Self::No, + Self::No => Self::Always, + Self::Always => Self::Yes, + } + } + + /// Get the previous selection (cycling). + pub fn prev(self) -> Self { + match self { + Self::Yes => Self::Always, + Self::No => Self::Yes, + Self::Always => Self::No, + } + } +} + +/// Approval overlay state. +pub struct ApprovalOverlay { + /// The request being shown. + pub request: ApprovalRequest, + /// Current selection. + pub selection: ApprovalSelection, +} + +impl ApprovalOverlay { + /// Create a new approval overlay. + pub fn new(request: ApprovalRequest) -> Self { + Self { + request, + selection: ApprovalSelection::Yes, + } + } + + /// Move selection left. + pub fn select_prev(&mut self) { + self.selection = self.selection.prev(); + } + + /// Move selection right. + pub fn select_next(&mut self) { + self.selection = self.selection.next(); + } + + /// Handle keyboard shortcut. + pub fn handle_shortcut(&mut self, c: char) -> Option { + match c.to_ascii_lowercase() { + 'y' => Some(true), + 'n' => Some(false), + 'a' => { + self.selection = ApprovalSelection::Always; + Some(true) + } + _ => None, + } + } + + /// Confirm the current selection. + pub fn confirm(&self) -> (bool, bool) { + match self.selection { + ApprovalSelection::Yes => (true, false), + ApprovalSelection::No => (false, false), + ApprovalSelection::Always => (true, true), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_approval_selection_cycle() { + let sel = ApprovalSelection::Yes; + assert_eq!(sel.next(), ApprovalSelection::No); + assert_eq!(sel.next().next(), ApprovalSelection::Always); + assert_eq!(sel.next().next().next(), ApprovalSelection::Yes); + } + + #[test] + fn test_approval_shortcuts() { + let request = ApprovalRequest::new("test", "Test operation", serde_json::json!({})); + let mut overlay = ApprovalOverlay::new(request); + + assert_eq!(overlay.handle_shortcut('y'), Some(true)); + assert_eq!(overlay.handle_shortcut('n'), Some(false)); + assert_eq!(overlay.handle_shortcut('x'), None); + } +} diff --git a/src/channels/cli/render.rs b/src/channels/cli/render.rs new file mode 100644 index 00000000..3f693e17 --- /dev/null +++ b/src/channels/cli/render.rs @@ -0,0 +1,221 @@ +//! TUI rendering with Ratatui. + +use ratatui::{ + Frame, + layout::{Constraint, Direction, Layout, Rect}, + style::{Color, Modifier, Style}, + text::{Line, Span, Text}, + widgets::{Block, Borders, Clear, List, ListItem, Paragraph, Wrap}, +}; + +use crate::channels::cli::app::{AppState, InputMode, MessageRole, MessageStatus}; +use crate::channels::cli::overlay::ApprovalSelection; + +/// Render the entire UI. +pub fn render(frame: &mut Frame, app: &AppState) { + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints([ + Constraint::Min(3), // Messages + Constraint::Length(3), // Input + Constraint::Length(1), // Status + ]) + .split(frame.area()); + + render_messages(frame, app, chunks[0]); + render_input(frame, app, chunks[1]); + render_status(frame, app, chunks[2]); + + // Render approval overlay if active + if app.mode == InputMode::Approval { + render_approval_overlay(frame, app); + } +} + +/// Render the message history. +fn render_messages(frame: &mut Frame, app: &AppState, area: Rect) { + let items: Vec = app + .messages + .iter() + .map(|msg| { + let (prefix, style) = match msg.role { + MessageRole::User => ("You: ", Style::default().fg(Color::Cyan)), + MessageRole::Agent => ("Agent: ", Style::default().fg(Color::Green)), + MessageRole::System => ( + "", + Style::default() + .fg(Color::DarkGray) + .add_modifier(Modifier::ITALIC), + ), + }; + + let status_indicator = match msg.status { + Some(MessageStatus::Pending) => " ⏳", + Some(MessageStatus::InProgress) => " ⚙️", + Some(MessageStatus::Complete) => " ✓", + Some(MessageStatus::Error) => " ✗", + None => "", + }; + + let content = format!("{}{}{}", prefix, msg.content, status_indicator); + ListItem::new(Text::styled(content, style)) + }) + .collect(); + + let messages = List::new(items).block(Block::default().borders(Borders::ALL).title("Chat")); + + frame.render_widget(messages, area); +} + +/// Render the input area. +fn render_input(frame: &mut Frame, app: &AppState, area: Rect) { + let input_style = match app.mode { + InputMode::Editing => Style::default().fg(Color::Yellow), + InputMode::Normal => Style::default(), + InputMode::Approval => Style::default().fg(Color::DarkGray), + }; + + let buffer = app.composer.buffer(); + let cursor = app.composer.cursor(); + + // Build the input text with cursor + let (before, after) = buffer.split_at(cursor.min(buffer.len())); + let cursor_char = after.chars().next().unwrap_or(' '); + let after_cursor = if after.is_empty() { + "" + } else { + &after[cursor_char.len_utf8()..] + }; + + let input = Paragraph::new(Line::from(vec![ + Span::raw(before), + Span::styled( + cursor_char.to_string(), + Style::default().bg(Color::White).fg(Color::Black), + ), + Span::raw(after_cursor), + ])) + .style(input_style) + .block(Block::default().borders(Borders::ALL).title("Input")); + + frame.render_widget(input, area); + + // Show cursor in editing mode + if app.mode == InputMode::Editing { + // Calculate cursor position accounting for the block border + let cursor_x = area.x + 1 + cursor as u16; + let cursor_y = area.y + 1; + frame.set_cursor_position((cursor_x, cursor_y)); + } +} + +/// Render the status line. +fn render_status(frame: &mut Frame, app: &AppState, area: Rect) { + let status_text = if let Some(ref msg) = app.status_message { + msg.clone() + } else { + match app.mode { + InputMode::Normal => "Press 'i' to edit, 'q' to quit".to_string(), + InputMode::Editing => "Type message, Enter to send, Esc to cancel".to_string(), + InputMode::Approval => "y=Yes, n=No, a=Always, Ctrl+C=Cancel all".to_string(), + } + }; + + let status = Paragraph::new(status_text).style(Style::default().fg(Color::DarkGray)); + + frame.render_widget(status, area); +} + +/// Render the approval overlay. +fn render_approval_overlay(frame: &mut Frame, app: &AppState) { + let Some(ref overlay) = app.approval else { + return; + }; + + let area = frame.area(); + + // Calculate overlay size and position + let overlay_width = (area.width * 60 / 100).min(60); + let overlay_height = 12; + let overlay_x = (area.width - overlay_width) / 2; + let overlay_y = (area.height - overlay_height) / 2; + + let overlay_area = Rect::new(overlay_x, overlay_y, overlay_width, overlay_height); + + // Clear the area behind the overlay + frame.render_widget(Clear, overlay_area); + + // Build overlay content + let title = if overlay.request.destructive { + "⚠️ Approval Required (Destructive)" + } else { + "Approval Required" + }; + + let title_style = if overlay.request.destructive { + Style::default().fg(Color::Red).add_modifier(Modifier::BOLD) + } else { + Style::default() + .fg(Color::Yellow) + .add_modifier(Modifier::BOLD) + }; + + // Build the text content + let mut lines = vec![ + Line::from(vec![ + Span::styled("Tool: ", Style::default().add_modifier(Modifier::BOLD)), + Span::raw(&overlay.request.tool_name), + ]), + Line::from(""), + Line::from(overlay.request.description.as_str()), + Line::from(""), + ]; + + // Add parameters preview (truncated) + let params_str = serde_json::to_string_pretty(&overlay.request.parameters) + .unwrap_or_else(|_| "{}".to_string()); + let params_preview: String = params_str.chars().take(100).collect(); + lines.push(Line::from(vec![ + Span::styled("Params: ", Style::default().add_modifier(Modifier::BOLD)), + Span::styled(params_preview, Style::default().fg(Color::DarkGray)), + ])); + lines.push(Line::from("")); + + // Add selection buttons + let yes_style = if overlay.selection == ApprovalSelection::Yes { + Style::default().bg(Color::Green).fg(Color::Black) + } else { + Style::default().fg(Color::Green) + }; + + let no_style = if overlay.selection == ApprovalSelection::No { + Style::default().bg(Color::Red).fg(Color::Black) + } else { + Style::default().fg(Color::Red) + }; + + let always_style = if overlay.selection == ApprovalSelection::Always { + Style::default().bg(Color::Blue).fg(Color::Black) + } else { + Style::default().fg(Color::Blue) + }; + + lines.push(Line::from(vec![ + Span::raw(" "), + Span::styled(" [Y]es ", yes_style), + Span::raw(" "), + Span::styled(" [N]o ", no_style), + Span::raw(" "), + Span::styled(" [A]lways ", always_style), + ])); + + let content = Paragraph::new(lines) + .block( + Block::default() + .borders(Borders::ALL) + .title(Span::styled(title, title_style)), + ) + .wrap(Wrap { trim: true }); + + frame.render_widget(content, overlay_area); +} diff --git a/src/channels/mod.rs b/src/channels/mod.rs index d78b1695..054b1208 100644 --- a/src/channels/mod.rs +++ b/src/channels/mod.rs @@ -4,14 +4,14 @@ //! and convert them to a unified message format for the agent to process. mod channel; -mod cli; +pub mod cli; mod http; mod manager; mod slack; mod telegram; pub use channel::{Channel, IncomingMessage, MessageStream, OutgoingResponse}; -pub use cli::CliChannel; +pub use cli::{SimpleCliChannel as CliChannel, TuiChannel}; pub use http::HttpChannel; pub use manager::ChannelManager; pub use slack::SlackChannel; diff --git a/src/llm/reasoning.rs b/src/llm/reasoning.rs index 65ad1a66..65623ef5 100644 --- a/src/llm/reasoning.rs +++ b/src/llm/reasoning.rs @@ -140,8 +140,20 @@ impl Reasoning { &self, context: &ReasoningContext, ) -> Result, LlmError> { + let tools = self.select_tools(context).await?; + Ok(tools.into_iter().next()) + } + + /// Select tools to execute (may return multiple for parallel execution). + /// + /// The LLM may return multiple tool calls if it determines they can be + /// executed in parallel. This enables more efficient job completion. + pub async fn select_tools( + &self, + context: &ReasoningContext, + ) -> Result, LlmError> { if context.available_tools.is_empty() { - return Ok(None); + return Ok(vec![]); } let request = @@ -151,16 +163,20 @@ impl Reasoning { let response = self.llm.complete_with_tools(request).await?; - if let Some(tool_call) = response.tool_calls.first() { - Ok(Some(ToolSelection { - tool_name: tool_call.name.clone(), - parameters: tool_call.arguments.clone(), - reasoning: response.content.unwrap_or_default(), + let reasoning = response.content.unwrap_or_default(); + + let selections: Vec = response + .tool_calls + .into_iter() + .map(|tool_call| ToolSelection { + tool_name: tool_call.name, + parameters: tool_call.arguments, + reasoning: reasoning.clone(), alternatives: vec![], - })) - } else { - Ok(None) - } + }) + .collect(); + + Ok(selections) } /// Evaluate whether a task was completed successfully. diff --git a/src/safety/leak_detector.rs b/src/safety/leak_detector.rs index 612b1403..65d301b0 100644 --- a/src/safety/leak_detector.rs +++ b/src/safety/leak_detector.rs @@ -299,12 +299,11 @@ impl LeakDetector { // Scan each header value for (name, value) in headers { - self.scan_and_clean(value).map_err(|e| { - LeakDetectionError::SecretLeakBlocked { + self.scan_and_clean(value) + .map_err(|e| LeakDetectionError::SecretLeakBlocked { pattern: format!("header:{}", name), preview: e.to_string(), - } - })?; + })?; } // Scan body if present and valid UTF-8 @@ -688,7 +687,10 @@ mod tests { // Attempt to exfiltrate in custom header let result = detector.scan_http_request( "https://api.example.com/data", - &[("X-Custom".to_string(), "ghp_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx".to_string())], + &[( + "X-Custom".to_string(), + "ghp_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx".to_string(), + )], None, ); assert!(result.is_err()); @@ -700,11 +702,7 @@ mod tests { // Attempt to exfiltrate in request body let body = b"{\"stolen\": \"sk-proj-test1234567890abcdefghij\"}"; - let result = detector.scan_http_request( - "https://api.example.com/webhook", - &[], - Some(body), - ); + let result = detector.scan_http_request("https://api.example.com/webhook", &[], Some(body)); assert!(result.is_err()); } } diff --git a/src/tools/builtin/http.rs b/src/tools/builtin/http.rs index a583db21..b9b25d3e 100644 --- a/src/tools/builtin/http.rs +++ b/src/tools/builtin/http.rs @@ -163,4 +163,8 @@ impl Tool for HttpTool { fn requires_sanitization(&self) -> bool { true // External data always needs sanitization } + + fn requires_approval(&self) -> bool { + true // HTTP requests go to external services, require user approval + } } diff --git a/src/tools/mcp/client.rs b/src/tools/mcp/client.rs index 652f2fd0..45ea3b7b 100644 --- a/src/tools/mcp/client.rs +++ b/src/tools/mcp/client.rs @@ -203,6 +203,11 @@ impl Tool for McpToolWrapper { fn requires_sanitization(&self) -> bool { true // MCP tools are external, always sanitize } + + fn requires_approval(&self) -> bool { + // Check the destructive_hint annotation from the MCP server + self.tool.requires_approval() + } } #[cfg(test)] diff --git a/src/tools/mcp/protocol.rs b/src/tools/mcp/protocol.rs index 3cdb44d2..285be30e 100644 --- a/src/tools/mcp/protocol.rs +++ b/src/tools/mcp/protocol.rs @@ -11,6 +11,52 @@ pub struct McpTool { pub description: String, /// JSON Schema for input parameters. pub input_schema: serde_json::Value, + /// Optional annotations from the MCP server. + #[serde(default)] + pub annotations: Option, +} + +/// Annotations for an MCP tool that provide hints about its behavior. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct McpToolAnnotations { + /// Hint that this tool performs destructive operations that cannot be undone. + /// Tools with this hint set to true should require user approval before execution. + #[serde(default)] + pub destructive_hint: bool, + + /// Hint that this tool may have side effects beyond its return value. + #[serde(default)] + pub side_effects_hint: bool, + + /// Hint that this tool performs read-only operations. + #[serde(default)] + pub read_only_hint: bool, + + /// Hint about the expected execution time category. + #[serde(default)] + pub execution_time_hint: Option, +} + +/// Hint about how long a tool typically takes to execute. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ExecutionTimeHint { + /// Typically completes in under 1 second. + Fast, + /// Typically completes in 1-10 seconds. + Medium, + /// Typically completes in more than 10 seconds. + Slow, +} + +impl McpTool { + /// Check if this tool requires user approval based on its annotations. + pub fn requires_approval(&self) -> bool { + self.annotations + .as_ref() + .map(|a| a.destructive_hint) + .unwrap_or(false) + } } /// Request to an MCP server. diff --git a/src/tools/tool.rs b/src/tools/tool.rs index 52ab547c..06fcb2c6 100644 --- a/src/tools/tool.rs +++ b/src/tools/tool.rs @@ -148,6 +148,18 @@ pub trait Tool: Send + Sync { true } + /// Whether this tool requires explicit user approval before execution. + /// + /// Returns false by default since most tools run in a sandboxed/virtualized + /// environment. Only tools that make external network calls or perform + /// destructive operations should return true. + /// + /// When true, the agent will prompt the user for confirmation before + /// executing this tool. + fn requires_approval(&self) -> bool { + false + } + /// Get the tool schema for LLM function calling. fn schema(&self) -> ToolSchema { ToolSchema { diff --git a/src/tools/wasm/storage.rs b/src/tools/wasm/storage.rs index 855398f6..d29a8afb 100644 --- a/src/tools/wasm/storage.rs +++ b/src/tools/wasm/storage.rs @@ -553,9 +553,7 @@ fn row_to_tool(row: &tokio_postgres::Row) -> Result