mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 23:10:11 +00:00
Integrate Codex patterns: task scheduler, TUI, sessions, compaction
Add Codex-inspired patterns for improved agent architecture: **Task Scheduler (Phase 1 & 4)** - New Task enum with Job, ToolExec, Background variants - TaskHandler trait for custom background tasks - Scheduler.spawn_subtask() and spawn_batch() for parallel execution - Worker executes multiple tools in parallel via futures::join_all **Tool Approval System (Phase 6)** - Tool.requires_approval() method (default false for sandboxed tools) - HttpTool marked as requiring approval (external network) - MCP protocol annotations: destructive_hint, side_effects_hint **Ratatui TUI CLI (Phase 2)** - Replace blocking stdin with event-driven TUI - ChatComposer with history navigation and tab completion - ApprovalOverlay modal with y/n/a keyboard shortcuts - Raw mode with proper terminal cleanup **Session/Turn Model (Phase 3)** - Session, Thread, Turn structs for conversation tracking - Submission enum for user input, approvals, undo, interrupt - UndoManager with checkpoint-based undo/redo **Context Compaction (Phase 5)** - ContextMonitor with token estimation and threshold checks - CompactionStrategy: Summarize, Truncate, MoveToWorkspace - ContextCompactor writes summaries to workspace daily logs Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
co-authored by
Claude Opus 4.5
parent
d047c23b2d
commit
09032d69cb
@@ -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<String>,
|
||||
}
|
||||
|
||||
/// Compacts conversation context to stay within limits.
|
||||
pub struct ContextCompactor {
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
}
|
||||
|
||||
impl ContextCompactor {
|
||||
/// Create a new context compactor.
|
||||
pub fn new(llm: Arc<dyn LlmProvider>) -> 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<CompactionResult, Error> {
|
||||
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<CompactionPartial, Error> {
|
||||
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<CompactionPartial, Error> {
|
||||
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<String, Error> {
|
||||
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::<Vec<_>>()
|
||||
.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<String>,
|
||||
}
|
||||
|
||||
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::<Vec<_>>()
|
||||
.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);
|
||||
}
|
||||
}
|
||||
@@ -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<CompactionStrategy> {
|
||||
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<ChatMessage> = vec![ChatMessage::user("Hello")];
|
||||
assert!(!monitor.needs_compaction(&small));
|
||||
|
||||
// Large context - compaction needed
|
||||
let large_content = "word ".repeat(1000);
|
||||
let large: Vec<ChatMessage> = 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<ChatMessage> = 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);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
+250
-16
@@ -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<WorkerMessage>,
|
||||
}
|
||||
|
||||
/// Status of a scheduled sub-task.
|
||||
struct ScheduledSubtask {
|
||||
task_id: Uuid,
|
||||
handle: JoinHandle<Result<TaskOutput, Error>>,
|
||||
}
|
||||
|
||||
/// Schedules and manages parallel job execution.
|
||||
pub struct Scheduler {
|
||||
config: AgentConfig,
|
||||
@@ -43,8 +51,10 @@ pub struct Scheduler {
|
||||
safety: Arc<SafetyLayer>,
|
||||
tools: Arc<ToolRegistry>,
|
||||
store: Option<Arc<Store>>,
|
||||
/// Running jobs.
|
||||
/// Running jobs (main LLM-driven jobs).
|
||||
jobs: RwLock<HashMap<Uuid, ScheduledJob>>,
|
||||
/// Running sub-tasks (tool executions, background tasks).
|
||||
subtasks: RwLock<HashMap<Uuid, ScheduledSubtask>>,
|
||||
}
|
||||
|
||||
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<oneshot::Receiver<Result<TaskOutput, Error>>, 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<Task>,
|
||||
) -> Vec<Result<TaskOutput, Error>> {
|
||||
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<ToolRegistry>,
|
||||
context_manager: Arc<ContextManager>,
|
||||
job_id: Uuid,
|
||||
tool_name: &str,
|
||||
params: serde_json::Value,
|
||||
) -> Result<TaskOutput, Error> {
|
||||
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<Uuid> {
|
||||
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<ToolRegistry> {
|
||||
&self.tools
|
||||
}
|
||||
|
||||
/// Get access to the context manager.
|
||||
pub fn context_manager(&self) -> &Arc<ContextManager> {
|
||||
&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.
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<Uuid>,
|
||||
/// All threads in this session.
|
||||
pub threads: HashMap<Uuid, Thread>,
|
||||
/// When the session was created.
|
||||
pub created_at: DateTime<Utc>,
|
||||
/// When the session was last active.
|
||||
pub last_active_at: DateTime<Utc>,
|
||||
/// Session metadata.
|
||||
pub metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
impl Session {
|
||||
/// Create a new session.
|
||||
pub fn new(user_id: impl Into<String>) -> 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<Turn>,
|
||||
/// When the thread was created.
|
||||
pub created_at: DateTime<Utc>,
|
||||
/// When the thread was last updated.
|
||||
pub updated_at: DateTime<Utc>,
|
||||
/// 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<String>) -> &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<String>) {
|
||||
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<String>) {
|
||||
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<ChatMessage> {
|
||||
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<String>,
|
||||
/// Tool calls made during this turn.
|
||||
pub tool_calls: Vec<TurnToolCall>,
|
||||
/// Turn state.
|
||||
pub state: TurnState,
|
||||
/// When the turn started.
|
||||
pub started_at: DateTime<Utc>,
|
||||
/// When the turn completed.
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
/// Error message (if failed).
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
impl Turn {
|
||||
/// Create a new turn.
|
||||
pub fn new(turn_number: usize, user_input: impl Into<String>) -> 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<String>) {
|
||||
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<String>) {
|
||||
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<String>, 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<String>) {
|
||||
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<serde_json::Value>,
|
||||
/// Error from the tool (if failed).
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
#[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());
|
||||
}
|
||||
}
|
||||
@@ -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<String>) -> 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<String>,
|
||||
},
|
||||
|
||||
/// Error occurred.
|
||||
Error {
|
||||
/// Error message.
|
||||
message: String,
|
||||
},
|
||||
|
||||
/// Turn was interrupted.
|
||||
Interrupted,
|
||||
}
|
||||
|
||||
impl SubmissionResult {
|
||||
/// Create a response result.
|
||||
pub fn response(content: impl Into<String>) -> 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<String>) -> Self {
|
||||
Self::Ok {
|
||||
message: Some(message.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an error result.
|
||||
pub fn error(message: impl Into<String>) -> 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());
|
||||
}
|
||||
}
|
||||
@@ -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<String>, 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<Uuid>,
|
||||
/// 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<TaskOutput, Error>;
|
||||
|
||||
/// 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<dyn TaskHandler>,
|
||||
},
|
||||
}
|
||||
|
||||
impl Task {
|
||||
/// Create a new Job task.
|
||||
pub fn job(title: impl Into<String>, description: impl Into<String>) -> 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<String>, description: impl Into<String>) -> 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<String>,
|
||||
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<dyn TaskHandler>) -> 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<dyn TaskHandler>) -> Self {
|
||||
Self::Background { id, handler }
|
||||
}
|
||||
|
||||
/// Get the task ID, if applicable.
|
||||
pub fn id(&self) -> Option<Uuid> {
|
||||
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<Uuid> {
|
||||
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:"));
|
||||
}
|
||||
}
|
||||
@@ -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<ChatMessage>,
|
||||
/// Description of what happened at this checkpoint.
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
impl Checkpoint {
|
||||
/// Create a new checkpoint.
|
||||
pub fn new(
|
||||
turn_number: usize,
|
||||
messages: Vec<ChatMessage>,
|
||||
description: impl Into<String>,
|
||||
) -> 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<Checkpoint>,
|
||||
/// Stack of future checkpoints (for redo).
|
||||
redo_stack: Vec<Checkpoint>,
|
||||
/// 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<ChatMessage>,
|
||||
description: impl Into<String>,
|
||||
) {
|
||||
// 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<ChatMessage>,
|
||||
) -> 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<Checkpoint> {
|
||||
self.undo_stack.pop_back()
|
||||
}
|
||||
|
||||
/// Redo: restore a previously undone state.
|
||||
pub fn redo(&mut self) -> Option<Checkpoint> {
|
||||
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<Checkpoint> {
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
+197
-98
@@ -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<String, Error>,
|
||||
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<ToolExecResult> {
|
||||
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<ToolRegistry>,
|
||||
context_manager: Arc<ContextManager>,
|
||||
store: Option<Arc<Store>>,
|
||||
job_id: Uuid,
|
||||
tool_name: &str,
|
||||
params: &serde_json::Value,
|
||||
) -> Result<String, Error> {
|
||||
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<String, Error>,
|
||||
) -> Result<bool, Error> {
|
||||
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<String, Error> {
|
||||
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<TaskOutput> for Result<String, Error> {
|
||||
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()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<AtomicBool>,
|
||||
}
|
||||
|
||||
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<MessageStream, ChannelError> {
|
||||
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();
|
||||
}
|
||||
@@ -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<MessageStatus>,
|
||||
}
|
||||
|
||||
/// 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<String>) -> Self {
|
||||
Self {
|
||||
role: MessageRole::User,
|
||||
content: content.into(),
|
||||
status: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn agent(content: impl Into<String>) -> Self {
|
||||
Self {
|
||||
role: MessageRole::Agent,
|
||||
content: content.into(),
|
||||
status: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn system(content: impl Into<String>) -> 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<ChatMessage>,
|
||||
/// Input composer.
|
||||
pub composer: ChatComposer,
|
||||
/// Approval overlay (if active).
|
||||
pub approval: Option<ApprovalOverlay>,
|
||||
/// Scroll offset for messages.
|
||||
pub scroll_offset: u16,
|
||||
/// Whether the app should quit.
|
||||
pub should_quit: bool,
|
||||
/// Pending approvals queue.
|
||||
pub pending_approvals: VecDeque<ApprovalRequest>,
|
||||
/// Current streaming response buffer.
|
||||
pub streaming_buffer: Option<String>,
|
||||
/// Status line message.
|
||||
pub status_message: Option<String>,
|
||||
}
|
||||
|
||||
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<String>) {
|
||||
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<String>) {
|
||||
// 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<ApprovalRequest> {
|
||||
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<String>) {
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -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<String>,
|
||||
/// Current position in history (-1 = current input).
|
||||
history_index: Option<usize>,
|
||||
/// Saved current input when navigating history.
|
||||
saved_input: String,
|
||||
/// Completion candidates.
|
||||
completions: Vec<String>,
|
||||
/// Current completion index.
|
||||
completion_index: Option<usize>,
|
||||
}
|
||||
|
||||
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");
|
||||
}
|
||||
}
|
||||
@@ -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<CrosstermBackend<io::Stdout>>,
|
||||
app: &mut AppState,
|
||||
msg_tx: mpsc::Sender<IncomingMessage>,
|
||||
event_rx: &mut mpsc::Receiver<AppEvent>,
|
||||
) -> 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<IncomingMessage>,
|
||||
) -> 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<IncomingMessage>,
|
||||
) -> 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<IncomingMessage>,
|
||||
) -> 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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<mpsc::Sender<AppEvent>>,
|
||||
}
|
||||
|
||||
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<MessageStream, ChannelError> {
|
||||
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<Mutex<>> 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<IncomingMessage>,
|
||||
mut event_rx: mpsc::Receiver<AppEvent>,
|
||||
) -> 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<std::sync::atomic::AtomicBool>,
|
||||
}
|
||||
|
||||
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<MessageStream, ChannelError> {
|
||||
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();
|
||||
}
|
||||
@@ -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<String>,
|
||||
description: impl Into<String>,
|
||||
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<bool> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -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<ListItem> = 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);
|
||||
}
|
||||
+2
-2
@@ -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;
|
||||
|
||||
+26
-10
@@ -140,8 +140,20 @@ impl Reasoning {
|
||||
&self,
|
||||
context: &ReasoningContext,
|
||||
) -> Result<Option<ToolSelection>, 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<Vec<ToolSelection>, 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<ToolSelection> = 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.
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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<McpToolAnnotations>,
|
||||
}
|
||||
|
||||
/// 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<ExecutionTimeHint>,
|
||||
}
|
||||
|
||||
/// 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.
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -553,9 +553,7 @@ fn row_to_tool(row: &tokio_postgres::Row) -> Result<StoredWasmTool, WasmStorageE
|
||||
trust_level: trust_level_str
|
||||
.parse()
|
||||
.map_err(WasmStorageError::InvalidData)?,
|
||||
status: status_str
|
||||
.parse()
|
||||
.map_err(WasmStorageError::InvalidData)?,
|
||||
status: status_str.parse().map_err(WasmStorageError::InvalidData)?,
|
||||
created_at: row.get("created_at"),
|
||||
updated_at: row.get("updated_at"),
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user