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:
Illia Polosukhin
2026-02-02 23:41:57 -08:00
co-authored by Claude Opus 4.5
parent d047c23b2d
commit 09032d69cb
26 changed files with 3944 additions and 272 deletions
+324
View File
@@ -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);
}
}
+236
View File
@@ -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);
}
}
+14
View File
@@ -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
View File
@@ -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.
}
}
+390
View File
@@ -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());
}
}
+203
View File
@@ -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());
}
}
+276
View File
@@ -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:"));
}
}
+252
View File
@@ -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
View File
@@ -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,
&params,
)
.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()
})
}
}
-124
View File
@@ -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();
}
+252
View File
@@ -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()
}
}
+318
View File
@@ -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");
}
}
+270
View File
@@ -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
}
}
}
+251
View File
@@ -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();
}
+145
View File
@@ -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);
}
}
+221
View File
@@ -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
View File
@@ -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
View File
@@ -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.
+8 -10
View File
@@ -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());
}
}
+4
View File
@@ -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
}
}
+5
View File
@@ -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)]
+46
View File
@@ -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.
+12
View File
@@ -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 {
+1 -3
View File
@@ -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"),
})