P0 deep-dive: core, providers, agent, memory, TUI, setup wizard
DEEP DIVE — Agent Loop (gyandoot-agent): - Real conversation loop with interactive (run()) and non-interactive (run_conversation()) modes - Retry logic with exponential backoff and error classification (RateLimit, Auth, Timeout, ServerError, ContextOverflow) - Context window management — token estimation, 75% threshold compression - 7 built-in tools: Read, Write, Edit, Bash, Grep, Find, Ls - BashTool with destructive-command security blocks - ContextCompressor — prune tool results, summarize middle messages DEEP DIVE — Memory (gyandoot-memory): - Real MemPalace with ~/.gyandoot/palace/wings/<wing>/ directory tree - Daily transcript files (YYYY-MM-DD_transcript.md) with append-only format - L0/L1 identity persistence (< 157 tokens for L0) - Session digest to drawer.md - File persistence support for transcripts and identity store DEEP DIVE — Core (gyandoot-core): - GyandootConfig with TOML deserialization for all sections - ConfigLoader with env > profile > file > default chain - Rich types: Message, ToolCall, ToolResult, SessionId - Comprehensive error types with thiserror derive DEEP DIVE — Providers (gyandoot-providers): - Provider trait with chat/chat_stream/models - OpenAI provider with streaming SSE and function calling - Anthropic provider with Messages API + tool use - Ollama provider with merge_system/no_tools auto-detect - SseParser, StreamEvent, TokenStream for SSE parsing - ProviderFactory for creation from config NEW — TUI crate (gyandoot-tui): - Banner: Kapil-approved robot head ASCII art with ANSI colors - Setup wizard: 7 providers, interactive config, connection test - Terminal helpers: spinner, confirm, text input, selection menu - Ratatui TUI: keyboard input, commands, status display All 12 crates + CLI binary compile with 0 errors.
This commit is contained in:
@@ -1,47 +1,134 @@
|
||||
//! The core agent inference loop.
|
||||
//! Production-grade agent inference loop.
|
||||
//!
|
||||
//! The `AgentLoop` manages the cycle: receive message → build context →
|
||||
//! call LLM → process tool calls → repeat or respond.
|
||||
//! `AgentLoop` manages the full cycle: receive message → build context →
|
||||
//! call LLM → execute tools → respond → compress context. Supports
|
||||
//! interactive (stdin/stdout) and non-interactive (gateway/sub-agent) modes,
|
||||
//! retry logic with exponential backoff, error classification, and streaming
|
||||
//! response delivery.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use gyandoot_core::{
|
||||
AgentConfig, AgentStatus, ConversationContext, Message, MessageRole, ToolCall, ToolResult,
|
||||
AgentConfig, AgentError, AgentResult, AgentStatus, ConversationContext, Message,
|
||||
MessageRole,
|
||||
};
|
||||
use gyandoot_memory::palace::MemPalace;
|
||||
use gyandoot_providers::Provider;
|
||||
use gyandoot_providers::{Provider, ProviderResponse};
|
||||
use gyandoot_security::sanitizer::Sanitizer;
|
||||
use gyandoot_skills::SkillsManager;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{debug, error, info, span, Level};
|
||||
use tracing::{debug, error, info, span, warn, Level};
|
||||
|
||||
use crate::context::ContextCompressor;
|
||||
use crate::tools::ToolExecutor;
|
||||
|
||||
/// The main agent execution loop.
|
||||
// ─── Error classification ───────────────────────────────────────────────
|
||||
|
||||
/// Categories of errors the agent loop can encounter during LLM calls.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AgentErrorKind {
|
||||
/// Provider rejected authentication credentials.
|
||||
Auth,
|
||||
/// API rate-limit exceeded.
|
||||
RateLimit,
|
||||
/// Request timed out (network or provider-side).
|
||||
Timeout,
|
||||
/// Server returned a 5xx or similar transient error.
|
||||
ServerError,
|
||||
/// Context window exceeded the model's maximum capacity.
|
||||
ContextOverflow,
|
||||
/// All other / unclassified errors.
|
||||
Other,
|
||||
}
|
||||
|
||||
impl AgentErrorKind {
|
||||
/// Classify a provider- or agent-error string into a structured kind.
|
||||
pub fn classify(error_text: &str) -> Self {
|
||||
let lower = error_text.to_lowercase();
|
||||
if lower.contains("rate limit")
|
||||
|| lower.contains("rate_limit")
|
||||
|| lower.contains("429")
|
||||
{
|
||||
Self::RateLimit
|
||||
} else if lower.contains("auth")
|
||||
|| lower.contains("unauthorized")
|
||||
|| lower.contains("401")
|
||||
|| lower.contains("api key")
|
||||
{
|
||||
Self::Auth
|
||||
} else if lower.contains("timeout")
|
||||
|| lower.contains("timed out")
|
||||
|| lower.contains("connection reset")
|
||||
{
|
||||
Self::Timeout
|
||||
} else if lower.contains("5")
|
||||
|| lower.contains("server error")
|
||||
|| lower.contains("internal server error")
|
||||
{
|
||||
Self::ServerError
|
||||
} else if lower.contains("context length")
|
||||
|| lower.contains("maximum context")
|
||||
|| lower.contains("token limit")
|
||||
|| lower.contains("context overflow")
|
||||
{
|
||||
Self::ContextOverflow
|
||||
} else {
|
||||
Self::Other
|
||||
}
|
||||
}
|
||||
|
||||
/// True for errors that are safe to retry.
|
||||
pub fn is_retryable(&self) -> bool {
|
||||
matches!(self, Self::RateLimit | Self::Timeout | Self::ServerError | Self::Other)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── AgentLoop ──────────────────────────────────────────────────────────
|
||||
|
||||
/// Configuration for the agent-loop retry behaviour.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RetryConfig {
|
||||
/// Maximum number of retry attempts (including the initial call).
|
||||
pub max_attempts: u32,
|
||||
/// Base delay in milliseconds for exponential back-off.
|
||||
pub base_delay_ms: u64,
|
||||
/// Maximum delay cap in milliseconds.
|
||||
pub max_delay_ms: u64,
|
||||
}
|
||||
|
||||
impl Default for RetryConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_attempts: 3,
|
||||
base_delay_ms: 1000,
|
||||
max_delay_ms: 30_000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The main agent execution loop, usable in both interactive and
|
||||
/// non-interactive contexts.
|
||||
pub struct AgentLoop {
|
||||
config: AgentConfig,
|
||||
provider: Box<dyn Provider + Send + Sync>,
|
||||
tool_executor: ToolExecutor,
|
||||
context_compressor: ContextCompressor,
|
||||
memory: Arc<Mutex<MemPalace>>,
|
||||
skills_manager: SkillsManager,
|
||||
_skills_manager: SkillsManager,
|
||||
sanitizer: Sanitizer,
|
||||
status: Arc<Mutex<AgentStatus>>,
|
||||
retry_cfg: RetryConfig,
|
||||
}
|
||||
|
||||
impl AgentLoop {
|
||||
/// Create a new agent loop with the given configuration.
|
||||
///
|
||||
/// This initializes the LLM provider, tool executor, memory system, and
|
||||
/// skills manager based on the provided config.
|
||||
/// Create a new agent loop from a fully-resolved configuration.
|
||||
pub fn new(config: AgentConfig) -> Self {
|
||||
let provider = create_provider(&config);
|
||||
let tool_executor = ToolExecutor::new();
|
||||
let context_compressor = ContextCompressor::new();
|
||||
let memory = Arc::new(Mutex::new(MemPalace::new()));
|
||||
let tool_executor = ToolExecutor::with_defaults();
|
||||
let context_compressor = ContextCompressor::new(config.max_tokens);
|
||||
let memory = Arc::new(Mutex::new(MemPalace::new("default")));
|
||||
let skills_manager = SkillsManager::new();
|
||||
let sanitizer = Sanitizer::new();
|
||||
|
||||
Self {
|
||||
config,
|
||||
@@ -49,96 +136,81 @@ impl AgentLoop {
|
||||
tool_executor,
|
||||
context_compressor,
|
||||
memory,
|
||||
skills_manager,
|
||||
sanitizer,
|
||||
_skills_manager: skills_manager,
|
||||
sanitizer: Sanitizer::new(),
|
||||
status: Arc::new(Mutex::new(AgentStatus::Idle)),
|
||||
retry_cfg: RetryConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Run the agent loop interactively, reading from stdin and writing to stdout.
|
||||
///
|
||||
/// This is the primary entry point for CLI usage. The loop reads user input,
|
||||
/// calls the LLM, executes any requested tools, and continues until the
|
||||
/// user exits or an error occurs.
|
||||
/// Create an agent loop with a custom retry policy.
|
||||
pub fn with_retry(mut self, cfg: RetryConfig) -> Self {
|
||||
self.retry_cfg = cfg;
|
||||
self
|
||||
}
|
||||
|
||||
// ── Interactive loop (stdin → stdout) ──────────────────────────────
|
||||
|
||||
/// Run the agent loop interactively, reading from stdin and writing to
|
||||
/// stdout until the user types `exit` / `quit` or an unrecoverable error
|
||||
/// occurs.
|
||||
pub async fn run(&self) -> anyhow::Result<()> {
|
||||
let span = span!(Level::INFO, "agent_loop", agent = %self.config.name);
|
||||
let _guard = span.enter();
|
||||
info!("Agent loop started (interactive)");
|
||||
|
||||
info!("Agent loop started");
|
||||
|
||||
let mut messages: Vec<Message> = vec![
|
||||
Message::new(MessageRole::System, &self.config.system_prompt),
|
||||
];
|
||||
let mut messages: Vec<Message> =
|
||||
vec![Message::new(MessageRole::System, &self.config.system_prompt)];
|
||||
|
||||
loop {
|
||||
// Read user input
|
||||
// 1. Read user input
|
||||
let user_input = self.read_input().await?;
|
||||
|
||||
// Check for exit conditions
|
||||
// 2. Check exit conditions
|
||||
if self.should_exit(&user_input) {
|
||||
info!("User requested exit");
|
||||
break;
|
||||
}
|
||||
|
||||
// Sanitize input
|
||||
// 3. Sanitize
|
||||
let clean_input = self.sanitizer.sanitize(&user_input);
|
||||
|
||||
// Add user message
|
||||
messages.push(Message::new(MessageRole::User, clean_input));
|
||||
|
||||
// Build conversation context
|
||||
let context = ConversationContext {
|
||||
messages: messages.clone(),
|
||||
system_prompt: self.config.system_prompt.clone(),
|
||||
tools: self.tool_executor.definitions(),
|
||||
max_tokens: self.config.max_tokens,
|
||||
};
|
||||
// 4. Run the core turn (LLM + tool execution)
|
||||
let new_messages = self.execute_turn(messages).await?;
|
||||
messages = new_messages;
|
||||
|
||||
// Call the LLM
|
||||
*self.status.lock().await = AgentStatus::Thinking;
|
||||
let response = self.provider.chat(context).await?;
|
||||
|
||||
// Process assistant response
|
||||
messages.push(Message::new(MessageRole::Assistant, &response.content));
|
||||
|
||||
// Handle tool calls
|
||||
if let Some(tool_calls) = response.tool_calls {
|
||||
for tool_call in tool_calls {
|
||||
debug!("Executing tool: {}", tool_call.name);
|
||||
let result = self.tool_executor.execute(&tool_call).await?;
|
||||
messages.push(
|
||||
Message::new(MessageRole::Tool, &result.output)
|
||||
.with_tool_result(result.output.clone()),
|
||||
);
|
||||
}
|
||||
|
||||
// Call LLM again with tool results
|
||||
let context = ConversationContext {
|
||||
messages: messages.clone(),
|
||||
system_prompt: self.config.system_prompt.clone(),
|
||||
tools: self.tool_executor.definitions(),
|
||||
max_tokens: self.config.max_tokens,
|
||||
};
|
||||
|
||||
let follow_up = self.provider.chat(context).await?;
|
||||
messages.push(Message::new(MessageRole::Assistant, &follow_up.content));
|
||||
self.display_output(&follow_up.content);
|
||||
} else {
|
||||
self.display_output(&response.content);
|
||||
}
|
||||
|
||||
// Store in memory
|
||||
// 5. Store conversation turn in memory
|
||||
{
|
||||
let mut memory = self.memory.lock().await;
|
||||
memory.store_conversation_turn(&messages).await?;
|
||||
let memory = self.memory.lock().await;
|
||||
memory
|
||||
.transcript_append("user", &user_input)
|
||||
.await
|
||||
.unwrap_or_else(|e| warn!("Failed to store user message: {e}"));
|
||||
// Store last assistant response
|
||||
if let Some(last) = messages.last() {
|
||||
if last.role == MessageRole::Assistant {
|
||||
memory
|
||||
.transcript_append("assistant", &last.content)
|
||||
.await
|
||||
.unwrap_or_else(|e| warn!("Failed to store assistant message: {e}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Compress context if too large
|
||||
if self.should_compress(&messages) {
|
||||
debug!("Context size threshold reached, compressing");
|
||||
messages = self.context_compressor.compress(messages).await;
|
||||
// 6. Compress context if token budget is tight
|
||||
let total_tokens: usize =
|
||||
messages.iter().map(|m| estimate_tokens(&m.content)).sum();
|
||||
let limit = self.config.max_tokens;
|
||||
if total_tokens > (limit * 3 / 4) {
|
||||
debug!(
|
||||
"Context at ~{:.0}% of limit ({total_tokens}/{limit}), compressing",
|
||||
(total_tokens as f64 / limit as f64) * 100.0
|
||||
);
|
||||
messages = self.context_compressor.compress(messages, limit).await;
|
||||
}
|
||||
|
||||
// 7. Update status
|
||||
*self.status.lock().await = AgentStatus::Idle;
|
||||
}
|
||||
|
||||
@@ -146,11 +218,263 @@ impl AgentLoop {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read input from stdin (CLI mode).
|
||||
// ── Non-interactive (gateway / sub-agent mode) ─────────────────────
|
||||
|
||||
/// Run a single conversation turn given an incoming user message.
|
||||
/// Returns the assistant response(s) and any tool outputs.
|
||||
///
|
||||
/// This is the entry point used by the gateway (Discord, Telegram, etc.)
|
||||
/// and by sub-agents that receive delegated tasks.
|
||||
pub async fn run_conversation(
|
||||
&self,
|
||||
message: &str,
|
||||
) -> AgentResult<ConversationOutput> {
|
||||
let span = span!(Level::INFO, "conversation", agent = %self.config.name);
|
||||
let _guard = span.enter();
|
||||
debug!("run_conversation: processing message");
|
||||
|
||||
let mut messages: Vec<Message> =
|
||||
vec![Message::new(MessageRole::System, &self.config.system_prompt)];
|
||||
|
||||
// Load recent context from memory for continuity
|
||||
{
|
||||
let memory = self.memory.lock().await;
|
||||
if let Ok(transcript) = memory.get_transcript(4096, 10).await {
|
||||
if !transcript.is_empty() {
|
||||
messages.push(Message::new(
|
||||
MessageRole::System,
|
||||
format!("[Recent conversation context]\n{transcript}"),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
messages.push(Message::new(MessageRole::User, message));
|
||||
|
||||
// Execute turn with retry
|
||||
let (final_messages, output) = self
|
||||
.execute_turn_with_output(messages)
|
||||
.await?;
|
||||
|
||||
// Store in memory
|
||||
{
|
||||
let memory = self.memory.lock().await;
|
||||
memory
|
||||
.transcript_append("user", message)
|
||||
.await
|
||||
.unwrap_or_else(|e| warn!("Failed to store message: {e}"));
|
||||
if let Some(last) = final_messages.last() {
|
||||
if last.role == MessageRole::Assistant {
|
||||
memory
|
||||
.transcript_append("assistant", &last.content)
|
||||
.await
|
||||
.unwrap_or_else(|e| warn!("Failed to store response: {e}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
/// Get the current agent status.
|
||||
pub async fn get_status(&self) -> AgentStatus {
|
||||
self.status.lock().await.clone()
|
||||
}
|
||||
|
||||
/// Provide access to tool definitions (e.g. for OpenAI function-calling).
|
||||
pub fn tool_definitions(&self) -> Vec<gyandoot_core::ToolDefinition> {
|
||||
self.tool_executor.definitions()
|
||||
}
|
||||
|
||||
/// Provide a reference to the tool executor.
|
||||
pub fn tool_executor(&self) -> &ToolExecutor {
|
||||
&self.tool_executor
|
||||
}
|
||||
|
||||
// ── Internal helpers ────────────────────────────────────────────────
|
||||
|
||||
/// Execute a single turn: LLM call → tool calls → LLM follow-up.
|
||||
/// Returns the updated message list.
|
||||
async fn execute_turn(&self, mut messages: Vec<Message>) -> AgentResult<Vec<Message>> {
|
||||
let _limit = self.config.max_tokens;
|
||||
|
||||
// Build context and call LLM (with retry)
|
||||
let context = self
|
||||
.context_compressor
|
||||
.build_context(&messages, &self.config.system_prompt, &self.tool_executor.definitions());
|
||||
|
||||
*self.status.lock().await = AgentStatus::Thinking;
|
||||
let response = self.call_with_retry(context).await?;
|
||||
|
||||
messages.push(Message::new(MessageRole::Assistant, &response.content)
|
||||
.with_tool_calls(response.tool_calls.clone().unwrap_or_default()));
|
||||
|
||||
// If there are tool calls, execute them and loop back
|
||||
if let Some(calls) = response.tool_calls {
|
||||
if !calls.is_empty() {
|
||||
*self.status.lock().await = AgentStatus::WaitingForTool;
|
||||
for call in calls {
|
||||
debug!("Executing tool: {}", call.name);
|
||||
let result = self.tool_executor.execute(&call).await.map_err(|e| {
|
||||
AgentError::ToolExecution(format!("{}: {}", call.name, e))
|
||||
})?;
|
||||
let output = if result.success {
|
||||
result.output
|
||||
} else {
|
||||
format!("ERROR: {}", result.error.unwrap_or_default())
|
||||
};
|
||||
messages.push(
|
||||
Message::new(MessageRole::Tool, &output)
|
||||
.with_tool_result(output),
|
||||
);
|
||||
}
|
||||
|
||||
// Follow-up call with tool results
|
||||
let follow_context = self
|
||||
.context_compressor
|
||||
.build_context(&messages, &self.config.system_prompt, &self.tool_executor.definitions());
|
||||
|
||||
*self.status.lock().await = AgentStatus::Thinking;
|
||||
let follow_up = self.call_with_retry(follow_context).await?;
|
||||
messages.push(
|
||||
Message::new(MessageRole::Assistant, &follow_up.content)
|
||||
.with_tool_calls(follow_up.tool_calls.unwrap_or_default()),
|
||||
);
|
||||
self.display_output(&follow_up.content);
|
||||
} else {
|
||||
self.display_output(&response.content);
|
||||
}
|
||||
} else {
|
||||
self.display_output(&response.content);
|
||||
}
|
||||
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
/// Execute a single turn and return the structured output (for
|
||||
/// non-interactive callers).
|
||||
async fn execute_turn_with_output(
|
||||
&self,
|
||||
mut messages: Vec<Message>,
|
||||
) -> AgentResult<(Vec<Message>, ConversationOutput)> {
|
||||
let _limit = self.config.max_tokens;
|
||||
|
||||
let context = self
|
||||
.context_compressor
|
||||
.build_context(&messages, &self.config.system_prompt, &self.tool_executor.definitions());
|
||||
|
||||
*self.status.lock().await = AgentStatus::Thinking;
|
||||
let response = self.call_with_retry(context).await?;
|
||||
|
||||
let mut assistant_text = response.content.clone();
|
||||
let mut tool_results: Vec<String> = Vec::new();
|
||||
|
||||
messages.push(Message::new(MessageRole::Assistant, &response.content)
|
||||
.with_tool_calls(response.tool_calls.clone().unwrap_or_default()));
|
||||
|
||||
if let Some(calls) = response.tool_calls {
|
||||
if !calls.is_empty() {
|
||||
*self.status.lock().await = AgentStatus::WaitingForTool;
|
||||
for call in calls {
|
||||
let result = self.tool_executor.execute(&call).await.map_err(|e| {
|
||||
AgentError::ToolExecution(format!("{}: {}", call.name, e))
|
||||
})?;
|
||||
let output = if result.success {
|
||||
result.output.clone()
|
||||
} else {
|
||||
format!("ERROR: {}", result.error.unwrap_or_default())
|
||||
};
|
||||
tool_results.push(output.clone());
|
||||
messages.push(
|
||||
Message::new(MessageRole::Tool, &output)
|
||||
.with_tool_result(output),
|
||||
);
|
||||
}
|
||||
|
||||
let follow_context = self
|
||||
.context_compressor
|
||||
.build_context(&messages, &self.config.system_prompt, &self.tool_executor.definitions());
|
||||
|
||||
*self.status.lock().await = AgentStatus::Thinking;
|
||||
let follow_up = self.call_with_retry(follow_context).await?;
|
||||
assistant_text = follow_up.content.clone();
|
||||
messages.push(
|
||||
Message::new(MessageRole::Assistant, &follow_up.content)
|
||||
.with_tool_calls(follow_up.tool_calls.unwrap_or_default()),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Ok((
|
||||
messages,
|
||||
ConversationOutput {
|
||||
text: assistant_text,
|
||||
tool_results,
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
/// Call the LLM provider with automatic retry and exponential back-off.
|
||||
async fn call_with_retry(
|
||||
&self,
|
||||
context: ConversationContext,
|
||||
) -> AgentResult<ProviderResponse> {
|
||||
let mut last_error = None;
|
||||
|
||||
for attempt in 1..=self.retry_cfg.max_attempts {
|
||||
match self.provider.chat(context.clone()).await {
|
||||
Ok(response) => return Ok(response),
|
||||
Err(e) => {
|
||||
let error_text = e.to_string();
|
||||
let kind = AgentErrorKind::classify(&error_text);
|
||||
|
||||
// Non-retryable errors bubble up immediately.
|
||||
if !kind.is_retryable() {
|
||||
return Err(AgentError::Provider(format!(
|
||||
"Non-retryable provider error: {error_text}"
|
||||
)));
|
||||
}
|
||||
|
||||
if attempt < self.retry_cfg.max_attempts {
|
||||
let delay = calculate_backoff(
|
||||
attempt,
|
||||
self.retry_cfg.base_delay_ms,
|
||||
self.retry_cfg.max_delay_ms,
|
||||
);
|
||||
warn!(
|
||||
"Provider call failed (attempt {}/{}, kind={kind:?}), \
|
||||
retrying in {delay}ms: {error_text}",
|
||||
attempt,
|
||||
self.retry_cfg.max_attempts,
|
||||
);
|
||||
tokio::time::sleep(Duration::from_millis(delay)).await;
|
||||
last_error = Some(e);
|
||||
} else {
|
||||
error!(
|
||||
"Provider call failed after {} attempts: {error_text}",
|
||||
self.retry_cfg.max_attempts
|
||||
);
|
||||
return Err(AgentError::Provider(format!(
|
||||
"Failed after {} retries: {error_text}",
|
||||
self.retry_cfg.max_attempts
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(AgentError::Provider(format!(
|
||||
"Exhausted retries: {:?}",
|
||||
last_error.map(|e| e.to_string())
|
||||
)))
|
||||
}
|
||||
|
||||
/// Read a single line from stdin with a prompt.
|
||||
async fn read_input(&self) -> anyhow::Result<String> {
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
let mut stdout = tokio::io::stdout();
|
||||
use tokio::io::AsyncWriteExt;
|
||||
|
||||
let mut stdout = tokio::io::stdout();
|
||||
stdout.write_all(b"\n> ").await?;
|
||||
stdout.flush().await?;
|
||||
|
||||
@@ -158,76 +482,79 @@ impl AgentLoop {
|
||||
let reader = tokio::io::BufReader::new(stdin);
|
||||
let mut lines = reader.lines();
|
||||
let mut input = String::new();
|
||||
|
||||
if let Some(line) = lines.next_line().await? {
|
||||
input = line;
|
||||
}
|
||||
|
||||
Ok(input)
|
||||
}
|
||||
|
||||
/// Check if the input indicates the user wants to exit.
|
||||
fn should_exit(&self, input: &str) -> bool {
|
||||
matches!(input.trim().to_lowercase().as_str(), "exit" | "quit" | "/exit" | "/quit")
|
||||
matches!(
|
||||
input.trim().to_lowercase().as_str(),
|
||||
"exit" | "quit" | "/exit" | "/quit"
|
||||
)
|
||||
}
|
||||
|
||||
/// Display output to the user.
|
||||
fn display_output(&self, output: &str) {
|
||||
println!("\n{}", output);
|
||||
}
|
||||
|
||||
/// Determine if context compression is needed based on message count and token estimate.
|
||||
fn should_compress(&self, messages: &[Message]) -> bool {
|
||||
// Rough estimate: ~4 chars per token, compress when over 75% of max
|
||||
let estimated_tokens: usize = messages.iter().map(|m| m.content.len() / 4).sum();
|
||||
estimated_tokens > (self.config.max_tokens * 3 / 4)
|
||||
}
|
||||
|
||||
/// Get the current agent status.
|
||||
pub async fn get_status(&self) -> AgentStatus {
|
||||
self.status.lock().await.clone()
|
||||
}
|
||||
}
|
||||
|
||||
/// Create the appropriate LLM provider based on configuration.
|
||||
// ─── ConversationOutput (for non-interactive callers) ───────────────────
|
||||
|
||||
/// The result of a single `run_conversation` call.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ConversationOutput {
|
||||
/// The final assistant response text.
|
||||
pub text: String,
|
||||
/// Any tool outputs generated during the turn.
|
||||
pub tool_results: Vec<String>,
|
||||
}
|
||||
|
||||
// ─── Backoff calculation ───────────────────────────────────────────────
|
||||
|
||||
/// Calculate exponential back-off with jitter.
|
||||
fn calculate_backoff(attempt: u32, base_ms: u64, max_ms: u64) -> u64 {
|
||||
let delay = base_ms.saturating_mul(2u64.saturating_pow(attempt.saturating_sub(1)));
|
||||
let clamped = delay.min(max_ms);
|
||||
// Add small jitter (±10%)
|
||||
let jitter = rand::random::<f64>() * 0.2 - 0.1; // -10% to +10%
|
||||
(clamped as f64 * (1.0 + jitter)) as u64
|
||||
}
|
||||
|
||||
/// Rough token estimator (~4 chars per token).
|
||||
fn estimate_tokens(text: &str) -> usize {
|
||||
(text.len() / 4).max(1)
|
||||
}
|
||||
|
||||
// ─── Provider factory ──────────────────────────────────────────────────
|
||||
|
||||
fn create_provider(config: &AgentConfig) -> Box<dyn Provider + Send + Sync> {
|
||||
match config.provider.default.as_str() {
|
||||
"openai" => {
|
||||
let openai_config = config
|
||||
.provider
|
||||
.openai
|
||||
.clone()
|
||||
.unwrap_or_default();
|
||||
let cfg = config.provider.openai.clone().unwrap_or_default();
|
||||
Box::new(gyandoot_providers::openai::OpenAIProvider::new(
|
||||
openai_config.api_key.clone(),
|
||||
openai_config.model.clone(),
|
||||
openai_config.base_url.clone(),
|
||||
cfg.api_key,
|
||||
cfg.model,
|
||||
cfg.base_url,
|
||||
))
|
||||
}
|
||||
"anthropic" => {
|
||||
let anthropic_config = config
|
||||
.provider
|
||||
.anthropic
|
||||
.clone()
|
||||
.unwrap_or_default();
|
||||
let cfg = config.provider.anthropic.clone().unwrap_or_default();
|
||||
Box::new(gyandoot_providers::anthropic::AnthropicProvider::new(
|
||||
anthropic_config.api_key.clone(),
|
||||
anthropic_config.model.clone(),
|
||||
cfg.api_key,
|
||||
cfg.model,
|
||||
))
|
||||
}
|
||||
"ollama" => {
|
||||
let ollama_config = config
|
||||
.provider
|
||||
.ollama
|
||||
.clone()
|
||||
.unwrap_or_default();
|
||||
let cfg = config.provider.ollama.clone().unwrap_or_default();
|
||||
Box::new(gyandoot_providers::ollama::OllamaProvider::new(
|
||||
ollama_config.base_url.clone(),
|
||||
ollama_config.model.clone(),
|
||||
cfg.base_url,
|
||||
cfg.model,
|
||||
))
|
||||
}
|
||||
other => {
|
||||
tracing::warn!("Unknown provider '{}', falling back to OpenAI", other);
|
||||
tracing::warn!("Unknown provider '{other}', falling back to OpenAI");
|
||||
Box::new(gyandoot_providers::openai::OpenAIProvider::new(
|
||||
String::new(),
|
||||
config.model.clone(),
|
||||
@@ -237,10 +564,48 @@ fn create_provider(config: &AgentConfig) -> Box<dyn Provider + Send + Sync> {
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Tests ──────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_error_classification() {
|
||||
assert_eq!(AgentErrorKind::classify("rate limit exceeded"), AgentErrorKind::RateLimit);
|
||||
assert_eq!(AgentErrorKind::classify("401 unauthorized"), AgentErrorKind::Auth);
|
||||
assert_eq!(AgentErrorKind::classify("timeout waiting"), AgentErrorKind::Timeout);
|
||||
assert_eq!(AgentErrorKind::classify("500 server error"), AgentErrorKind::ServerError);
|
||||
assert_eq!(AgentErrorKind::classify("context length exceeded"), AgentErrorKind::ContextOverflow);
|
||||
assert_eq!(AgentErrorKind::classify("unknown error"), AgentErrorKind::Other);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_retryable_classification() {
|
||||
assert!(AgentErrorKind::RateLimit.is_retryable());
|
||||
assert!(AgentErrorKind::Timeout.is_retryable());
|
||||
assert!(!AgentErrorKind::Auth.is_retryable());
|
||||
assert!(!AgentErrorKind::ContextOverflow.is_retryable());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_backoff_is_monotonic() {
|
||||
let d1 = calculate_backoff(1, 1000, 30_000);
|
||||
let d2 = calculate_backoff(2, 1000, 30_000);
|
||||
let d3 = calculate_backoff(3, 1000, 30_000);
|
||||
// With jitter these might not be strictly monotonic, but should be roughly so.
|
||||
assert!(d1 >= 900 && d1 <= 1100); // ~1000ms ±10%
|
||||
assert!(d2 >= 1800 && d2 <= 2200); // ~2000ms ±10%
|
||||
assert!(d3 >= 3600 && d3 <= 4400); // ~4000ms ±10%
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_tokens() {
|
||||
assert_eq!(estimate_tokens("abcd"), 1);
|
||||
assert_eq!(estimate_tokens(""), 1); // at least 1
|
||||
assert!(estimate_tokens("hello world this is a test") > 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_should_exit() {
|
||||
let config = AgentConfig::default();
|
||||
@@ -249,15 +614,24 @@ mod tests {
|
||||
"gpt-4o".to_string(),
|
||||
None,
|
||||
);
|
||||
let _agent = AgentLoop {
|
||||
let tool_executor = ToolExecutor::with_defaults();
|
||||
let context_compressor = ContextCompressor::new(4096);
|
||||
|
||||
let loop_ = AgentLoop {
|
||||
config,
|
||||
provider: Box::new(provider),
|
||||
tool_executor: ToolExecutor::new(),
|
||||
context_compressor: ContextCompressor::new(),
|
||||
memory: Arc::new(Mutex::new(MemPalace::new())),
|
||||
skills_manager: SkillsManager::new(&gyandoot_core::SkillsConfig::default()),
|
||||
tool_executor,
|
||||
context_compressor,
|
||||
memory: Arc::new(Mutex::new(MemPalace::new("default"))),
|
||||
_skills_manager: SkillsManager::new(),
|
||||
sanitizer: Sanitizer::new(),
|
||||
status: Arc::new(Mutex::new(AgentStatus::Idle)),
|
||||
retry_cfg: RetryConfig::default(),
|
||||
};
|
||||
|
||||
assert!(loop_.should_exit("exit"));
|
||||
assert!(loop_.should_exit("quit"));
|
||||
assert!(loop_.should_exit(" /exit "));
|
||||
assert!(!loop_.should_exit("hello"));
|
||||
}
|
||||
}
|
||||
@@ -1,53 +1,152 @@
|
||||
//! Context compression logic.
|
||||
//! Production-grade context compression engine.
|
||||
//!
|
||||
//! When the agent's conversation context grows too large, this module
|
||||
//! intelligently compresses or summarizes older messages while preserving
|
||||
//! critical information.
|
||||
//! Manages LLM conversation context by estimating token usage, pruning verbose
|
||||
//! tool outputs, summarizing older messages, and assembling the final context
|
||||
//! that gets sent to the provider.
|
||||
|
||||
use gyandoot_core::Message;
|
||||
use gyandoot_core::{ConversationContext, Message, MessageRole, ToolDefinition};
|
||||
use tracing::debug;
|
||||
|
||||
/// Compresses conversation context when it exceeds size thresholds.
|
||||
// ─── ContextCompressor ──────────────────────────────────────────────────
|
||||
|
||||
/// Token-aware context window manager.
|
||||
///
|
||||
/// Strategies:
|
||||
/// - Rough token estimation (chars/4) for fast checks.
|
||||
/// - Summarization of older messages into a compact system message.
|
||||
/// - Pruning of verbose tool outputs from mid-history.
|
||||
/// - Keeping the most recent N turns intact.
|
||||
pub struct ContextCompressor {
|
||||
/// Hard limit on context messages before summarization kicks in.
|
||||
max_messages_before_compress: usize,
|
||||
/// Minimum recent messages to keep intact.
|
||||
recent_context_count: usize,
|
||||
}
|
||||
|
||||
impl ContextCompressor {
|
||||
/// Create a new context compressor.
|
||||
/// Create a new compressor with sensible defaults.
|
||||
///
|
||||
/// By default, compression triggers when there are more than 30 messages
|
||||
/// (roughly corresponding to 15 turns of conversation).
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
max_messages_before_compress: 30,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a compressor with a custom message threshold.
|
||||
pub fn with_threshold(max_messages: usize) -> Self {
|
||||
/// `max_tokens` is the model's context window size; compress triggers
|
||||
/// at 75% of this value.
|
||||
pub fn new(max_tokens: usize) -> Self {
|
||||
let max_messages = (max_tokens / 200).max(20);
|
||||
Self {
|
||||
max_messages_before_compress: max_messages,
|
||||
recent_context_count: (max_messages / 2).max(4),
|
||||
}
|
||||
}
|
||||
|
||||
/// Compress the message list by summarizing older messages.
|
||||
/// Create a compressor with explicit message thresholds.
|
||||
pub fn with_thresholds(max_messages: usize, recent_count: usize) -> Self {
|
||||
Self {
|
||||
max_messages_before_compress: max_messages,
|
||||
recent_context_count: recent_count.max(2),
|
||||
}
|
||||
}
|
||||
|
||||
/// Estimate the number of tokens in a text string.
|
||||
///
|
||||
/// Uses a simple char/4 heuristic. In production, swap this with a
|
||||
/// tiktoken-rs call for per-model accuracy.
|
||||
pub fn estimate_tokens(text: &str) -> usize {
|
||||
if text.is_empty() {
|
||||
return 1;
|
||||
}
|
||||
// ~4 chars per token on average for English text
|
||||
(text.len() / 4).max(1)
|
||||
}
|
||||
|
||||
/// Estimate tokens for a slice of messages.
|
||||
pub fn estimate_messages_tokens(messages: &[Message]) -> usize {
|
||||
messages
|
||||
.iter()
|
||||
.map(|m| Self::estimate_tokens(&m.content) + 4) // +4 for role overhead
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// Compress a message list to fit within `max_tokens`.
|
||||
///
|
||||
/// Strategy:
|
||||
/// - Keep the system prompt (first message).
|
||||
/// - Keep the most recent N messages (recent context).
|
||||
/// - Summarize the messages in between into a single condensed message.
|
||||
/// - Preserve all tool call/result pairs.
|
||||
pub async fn compress(&self, messages: Vec<Message>) -> Vec<Message> {
|
||||
if messages.len() <= self.max_messages_before_compress {
|
||||
/// 1. Always keep the system prompt (first message).
|
||||
/// 2. Prune verbose tool result bodies from middle messages.
|
||||
/// 3. Summarize older non-recent messages into a single condensed entry.
|
||||
/// 4. Keep the most recent N messages intact.
|
||||
pub async fn compress(
|
||||
&self,
|
||||
messages: Vec<Message>,
|
||||
max_tokens: usize,
|
||||
) -> Vec<Message> {
|
||||
let total_tokens = Self::estimate_messages_tokens(&messages);
|
||||
let threshold = max_tokens * 3 / 4;
|
||||
|
||||
if total_tokens <= threshold && messages.len() <= self.max_messages_before_compress {
|
||||
return messages;
|
||||
}
|
||||
|
||||
debug!(
|
||||
"Compressing {} messages (threshold: {})",
|
||||
"Compressing {} messages (estimated {} tokens, threshold {} tokens)",
|
||||
messages.len(),
|
||||
self.max_messages_before_compress
|
||||
total_tokens,
|
||||
threshold,
|
||||
);
|
||||
|
||||
// Step 1: Prune verbose tool results from middle messages
|
||||
let mut pruned = self.prune_tool_results(messages);
|
||||
|
||||
// Step 2: If still over threshold, summarize
|
||||
let pruned_tokens = Self::estimate_messages_tokens(&pruned);
|
||||
if pruned_tokens > threshold || pruned.len() > self.max_messages_before_compress {
|
||||
pruned = self.summarize_middle(pruned);
|
||||
}
|
||||
|
||||
debug!("Compressed to {} messages", pruned.len());
|
||||
pruned
|
||||
}
|
||||
|
||||
/// Prune verbose tool outputs from message history, keeping only a
|
||||
/// compact summary for non-recent tool results.
|
||||
pub fn prune_tool_results(&self, messages: Vec<Message>) -> Vec<Message> {
|
||||
let keep_recent = self.recent_context_count;
|
||||
let mut result: Vec<Message> = Vec::with_capacity(messages.len());
|
||||
|
||||
for (i, msg) in messages.iter().enumerate() {
|
||||
// Always keep the first message (system prompt)
|
||||
if i == 0 {
|
||||
result.push(msg.clone());
|
||||
continue;
|
||||
}
|
||||
|
||||
// Keep recent messages intact
|
||||
if i >= messages.len().saturating_sub(keep_recent) {
|
||||
result.push(msg.clone());
|
||||
continue;
|
||||
}
|
||||
|
||||
// Prune tool result content in middle messages
|
||||
if msg.role == MessageRole::Tool {
|
||||
let tool_result = msg.tool_result.as_deref().unwrap_or("");
|
||||
let truncated = if tool_result.len() > 200 {
|
||||
format!(
|
||||
"[Tool result: {} chars — truncated]\n...{}",
|
||||
tool_result.len(),
|
||||
&tool_result[..tool_result.len().min(200)]
|
||||
)
|
||||
} else {
|
||||
msg.content.clone()
|
||||
};
|
||||
result.push(Message::new(MessageRole::Tool, truncated));
|
||||
} else {
|
||||
result.push(msg.clone());
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Summarize non-recent, non-system messages into a single condensed
|
||||
/// system message.
|
||||
fn summarize_middle(&self, messages: Vec<Message>) -> Vec<Message> {
|
||||
let keep_recent = self.recent_context_count;
|
||||
let mut compressed: Vec<Message> = Vec::with_capacity(self.max_messages_before_compress);
|
||||
|
||||
// Always keep the system prompt
|
||||
@@ -55,108 +154,240 @@ impl ContextCompressor {
|
||||
compressed.push(system.clone());
|
||||
}
|
||||
|
||||
// Keep the most recent N-2 messages (reserve space for summary)
|
||||
let recent_count = (self.max_messages_before_compress / 2).max(2);
|
||||
let keep_start = messages.len().saturating_sub(recent_count);
|
||||
|
||||
// Summarize middle messages
|
||||
let middle: Vec<&Message> = messages[1..keep_start]
|
||||
// Messages to summarize: everything between system prompt and recent tail
|
||||
let summarize_end = messages.len().saturating_sub(keep_recent);
|
||||
let middle: Vec<&Message> = messages[1..summarize_end]
|
||||
.iter()
|
||||
.filter(|m| m.role != gyandoot_core::MessageRole::Tool)
|
||||
.filter(|m| m.role != MessageRole::Tool)
|
||||
.collect();
|
||||
|
||||
if !middle.is_empty() {
|
||||
let summary = self.summarize_messages(&middle);
|
||||
let summary = self.build_summary(&middle);
|
||||
compressed.push(Message::new(
|
||||
gyandoot_core::MessageRole::System,
|
||||
format!("[Summarized context from previous conversation: {}]", summary),
|
||||
MessageRole::System,
|
||||
format!(
|
||||
"[Summarized conversation history — {}]\n\
|
||||
(Context from earlier turns in this session)",
|
||||
summary
|
||||
),
|
||||
));
|
||||
}
|
||||
|
||||
// Append the most recent messages
|
||||
let keep_start = messages.len().saturating_sub(keep_recent).max(1);
|
||||
for msg in messages.iter().skip(keep_start) {
|
||||
compressed.push(msg.clone());
|
||||
}
|
||||
|
||||
debug!("Compressed to {} messages", compressed.len());
|
||||
compressed
|
||||
}
|
||||
|
||||
/// Create a text summary from a slice of messages.
|
||||
fn summarize_messages(&self, messages: &[&Message]) -> String {
|
||||
/// Build a text summary from a slice of middle messages.
|
||||
fn build_summary(&self, messages: &[&Message]) -> String {
|
||||
let count = messages.len();
|
||||
let user_msgs: Vec<&str> = messages
|
||||
.iter()
|
||||
.filter(|m| m.role == gyandoot_core::MessageRole::User)
|
||||
.filter(|m| m.role == MessageRole::User)
|
||||
.map(|m| m.content.as_str())
|
||||
.collect();
|
||||
|
||||
if user_msgs.is_empty() {
|
||||
format!("{} previous assistant messages", count)
|
||||
format!("{} assistant messages", count)
|
||||
} else if user_msgs.len() <= 3 {
|
||||
format!(
|
||||
"Earlier discussion: {}",
|
||||
user_msgs.join(" | ")
|
||||
"Previous discussion: {}",
|
||||
user_msgs
|
||||
.iter()
|
||||
.map(|s| truncate_str(s, 80))
|
||||
.collect::<Vec<_>>()
|
||||
.join(" | ")
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"{} previous turns covering topics: {}, ... (and {} more)",
|
||||
"{} turns covering: {}, ... (and {} more)",
|
||||
count,
|
||||
user_msgs[..3].join(", "),
|
||||
user_msgs[..3]
|
||||
.iter()
|
||||
.map(|s| truncate_str(s, 80))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", "),
|
||||
user_msgs.len().saturating_sub(3)
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ContextCompressor {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
/// Assemble the final `ConversationContext` from messages, system
|
||||
/// prompt, and tool definitions.
|
||||
pub fn build_context(
|
||||
&self,
|
||||
messages: &[Message],
|
||||
system_prompt: &str,
|
||||
tool_definitions: &[ToolDefinition],
|
||||
) -> ConversationContext {
|
||||
ConversationContext {
|
||||
messages: messages.to_vec(),
|
||||
system_prompt: system_prompt.to_string(),
|
||||
tools: tool_definitions.to_vec(),
|
||||
max_tokens: 0, // caller sets this or provider defaults
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Truncate a string to `max_chars`, appending "…" if shortened.
|
||||
fn truncate_str(s: &str, max_chars: usize) -> String {
|
||||
if s.len() <= max_chars {
|
||||
s.to_string()
|
||||
} else {
|
||||
format!("{}…", &s[..max_chars.saturating_sub(1)])
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Tests ──────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use gyandoot_core::MessageRole;
|
||||
|
||||
fn make_msg(role: MessageRole, content: impl Into<String>) -> Message {
|
||||
Message::new(role, content)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_tokens() {
|
||||
assert_eq!(ContextCompressor::estimate_tokens("abcd"), 1);
|
||||
assert_eq!(ContextCompressor::estimate_tokens(""), 1);
|
||||
let long = "a".repeat(400);
|
||||
assert_eq!(ContextCompressor::estimate_tokens(&long), 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_messages_tokens() {
|
||||
let msgs = vec![
|
||||
make_msg(MessageRole::User, "hello"),
|
||||
make_msg(MessageRole::Assistant, "world"),
|
||||
];
|
||||
let tokens = ContextCompressor::estimate_messages_tokens(&msgs);
|
||||
// Each content ~1 token, each +4 overhead = ~10 total
|
||||
assert!(tokens >= 10);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_no_compress_below_threshold() {
|
||||
let compressor = ContextCompressor::with_threshold(10);
|
||||
let compressor = ContextCompressor::with_thresholds(20, 5);
|
||||
let messages: Vec<Message> = (0..5)
|
||||
.map(|i| Message::new(MessageRole::User, format!("message {}", i)))
|
||||
.map(|i| make_msg(MessageRole::User, format!("msg {i}")))
|
||||
.collect();
|
||||
|
||||
let result = compressor.compress(messages.clone()).await;
|
||||
let result = compressor.compress(messages.clone(), 9999).await;
|
||||
assert_eq!(result.len(), 5);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_compress_above_threshold() {
|
||||
let compressor = ContextCompressor::with_threshold(10);
|
||||
let mut messages = vec![Message::new(MessageRole::System, "You are an AI.")];
|
||||
let compressor = ContextCompressor::with_thresholds(10, 4);
|
||||
let mut messages = vec![make_msg(MessageRole::System, "You are an AI.")];
|
||||
for i in 0..20 {
|
||||
messages.push(Message::new(MessageRole::User, format!("user message {}", i)));
|
||||
messages.push(Message::new(MessageRole::Assistant, format!("response {}", i)));
|
||||
messages.push(make_msg(MessageRole::User, format!("user msg {i}")));
|
||||
messages.push(make_msg(MessageRole::Assistant, format!("response {i}")));
|
||||
}
|
||||
|
||||
let result = compressor.compress(messages).await;
|
||||
let result = compressor.compress(messages, 9999).await;
|
||||
assert!(result.len() < 20);
|
||||
assert_eq!(result[0].content, "You are an AI.");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_compress_token_budget() {
|
||||
let compressor = ContextCompressor::with_thresholds(50, 8);
|
||||
let mut messages = vec![make_msg(MessageRole::System, "sys")];
|
||||
for i in 0..30 {
|
||||
messages.push(make_msg(MessageRole::User, format!("user msg {i}")));
|
||||
messages.push(make_msg(MessageRole::Assistant, format!("response {i}")));
|
||||
}
|
||||
// total messages: 1 + 60 = 61
|
||||
// estimated tokens: ~ 61 * (3/4 + 4) ≈ 290
|
||||
let result = compressor.compress(messages, 500).await;
|
||||
assert!(!result.is_empty());
|
||||
assert_eq!(result[0].content, "sys");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_summarize_empty() {
|
||||
let compressor = ContextCompressor::new();
|
||||
let summary = compressor.summarize_messages(&[]);
|
||||
fn test_prune_tool_results_short() {
|
||||
let compressor = ContextCompressor::with_thresholds(20, 5);
|
||||
let messages = vec![
|
||||
make_msg(MessageRole::System, "sys"),
|
||||
make_msg(MessageRole::User, "hello"),
|
||||
Message::new(MessageRole::Tool, "short result").with_tool_result("short result".into()),
|
||||
];
|
||||
let pruned = compressor.prune_tool_results(messages);
|
||||
assert_eq!(pruned.len(), 3);
|
||||
assert_eq!(pruned[2].content, "short result");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_prune_tool_results_long() {
|
||||
let compressor = ContextCompressor::with_thresholds(20, 3);
|
||||
let long_result = "x".repeat(500);
|
||||
let mut messages = vec![
|
||||
make_msg(MessageRole::System, "sys"),
|
||||
make_msg(MessageRole::User, "hello"),
|
||||
];
|
||||
// Add middle messages so tool result falls outside recent zone
|
||||
// total = 10, recent = 3, recent_start = 7
|
||||
for i in 0..5 {
|
||||
messages.push(make_msg(MessageRole::User, format!("middle {i}")));
|
||||
}
|
||||
messages.push(
|
||||
Message::new(MessageRole::Tool, &long_result)
|
||||
.with_tool_result(long_result.clone()),
|
||||
);
|
||||
messages.push(make_msg(MessageRole::Assistant, "response 1"));
|
||||
messages.push(make_msg(MessageRole::Assistant, "response 2"));
|
||||
messages.push(make_msg(MessageRole::Assistant, "recent response"));
|
||||
|
||||
let pruned = compressor.prune_tool_results(messages);
|
||||
// Find the tool message
|
||||
let tool_msg = pruned.iter().find(|m| m.role == MessageRole::Tool).unwrap();
|
||||
// Long tool content should be truncated
|
||||
assert!(
|
||||
tool_msg.content.contains("[Tool result:"),
|
||||
"Expected truncated tool result, got: {}...",
|
||||
&tool_msg.content[..tool_msg.content.len().min(100)]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_str() {
|
||||
assert_eq!(truncate_str("hello", 10), "hello");
|
||||
let long = "a".repeat(100);
|
||||
let truncated = truncate_str(&long, 10);
|
||||
// 9 chars + '…' (multi-byte) = 10 characters
|
||||
assert_eq!(truncated.chars().count(), 10);
|
||||
assert!(truncated.ends_with('…'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_summary_empty() {
|
||||
let compressor = ContextCompressor::with_thresholds(20, 5);
|
||||
let summary = compressor.build_summary(&[]);
|
||||
assert!(summary.contains("0"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_summarize_few() {
|
||||
let compressor = ContextCompressor::new();
|
||||
let msg = Message::new(MessageRole::User, "hello");
|
||||
let summary = compressor.summarize_messages(&[&msg]);
|
||||
assert!(summary.contains("hello"));
|
||||
fn test_build_summary_few_messages() {
|
||||
let compressor = ContextCompressor::with_thresholds(20, 5);
|
||||
let msg = make_msg(MessageRole::User, "hello there");
|
||||
let msgs = vec![&msg];
|
||||
let summary = compressor.build_summary(&msgs);
|
||||
assert!(summary.contains("hello there"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_context() {
|
||||
let compressor = ContextCompressor::with_thresholds(20, 5);
|
||||
let messages = vec![make_msg(MessageRole::User, "hi")];
|
||||
let ctx = compressor.build_context(&messages, "system prompt", &[]);
|
||||
assert_eq!(ctx.messages.len(), 1);
|
||||
assert_eq!(ctx.system_prompt, "system prompt");
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,6 @@ pub mod browser_tools;
|
||||
pub mod tools;
|
||||
pub mod context;
|
||||
|
||||
pub use agent_loop::AgentLoop;
|
||||
pub use tools::ToolExecutor;
|
||||
pub use agent_loop::{AgentLoop, AgentErrorKind, ConversationOutput, RetryConfig};
|
||||
pub use tools::{BuiltinTool, BashTool, BashToolConfig, Tool, ToolExecutor, ReadTool, WriteTool, EditTool, GrepTool, FindTool, LsTool};
|
||||
pub use context::ContextCompressor;
|
||||
@@ -1,42 +1,119 @@
|
||||
//! Tool execution subsystem.
|
||||
//! Production-grade tool execution subsystem.
|
||||
//!
|
||||
//! Discovers available tools, executes tool calls, and returns results.
|
||||
//! Defines the `Tool` trait and `ToolExecutor` which discovers, registers,
|
||||
//! and executes tools on behalf of the LLM. Includes built-in tools for
|
||||
//! filesystem operations and shell command execution, all with appropriate
|
||||
//! security checks and resource limits.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
|
||||
use gyandoot_core::{ToolCall, ToolDefinition, ToolResult};
|
||||
use serde_json;
|
||||
use tracing::{debug, error, warn};
|
||||
|
||||
/// Executes tool calls made by the LLM.
|
||||
// ─── Tool trait ─────────────────────────────────────────────────────────
|
||||
|
||||
/// Trait that all tools (built-in or skill-based) must implement.
|
||||
#[async_trait::async_trait]
|
||||
pub trait Tool: Send + Sync {
|
||||
/// Short, unique name for the tool (e.g. `"read"`, `"bash"`).
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// Human-readable description for the LLM.
|
||||
fn description(&self) -> &str;
|
||||
|
||||
/// JSON Schema for the tool's parameters.
|
||||
fn parameters(&self) -> serde_json::Value;
|
||||
|
||||
/// Execute the tool with the given arguments. Returns the output text.
|
||||
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<String>;
|
||||
}
|
||||
|
||||
// ─── BuiltinTool enum ───────────────────────────────────────────────────
|
||||
|
||||
/// Enumeration of all built-in tools available to the agent.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum BuiltinTool {
|
||||
Read,
|
||||
Write,
|
||||
Edit,
|
||||
Bash,
|
||||
Grep,
|
||||
Find,
|
||||
Ls,
|
||||
}
|
||||
|
||||
impl BuiltinTool {
|
||||
/// All built-in tool variants.
|
||||
pub fn all() -> &'static [BuiltinTool] {
|
||||
&[
|
||||
BuiltinTool::Read,
|
||||
BuiltinTool::Write,
|
||||
BuiltinTool::Edit,
|
||||
BuiltinTool::Bash,
|
||||
BuiltinTool::Grep,
|
||||
BuiltinTool::Find,
|
||||
BuiltinTool::Ls,
|
||||
]
|
||||
}
|
||||
|
||||
fn name(&self) -> &'static str {
|
||||
match self {
|
||||
BuiltinTool::Read => "read",
|
||||
BuiltinTool::Write => "write",
|
||||
BuiltinTool::Edit => "edit",
|
||||
BuiltinTool::Bash => "bash",
|
||||
BuiltinTool::Grep => "grep",
|
||||
BuiltinTool::Find => "find",
|
||||
BuiltinTool::Ls => "ls",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── ToolExecutor ───────────────────────────────────────────────────────
|
||||
|
||||
/// Manages tool registration, definition generation, and execution.
|
||||
pub struct ToolExecutor {
|
||||
tools: HashMap<String, Box<dyn Tool + Send + Sync>>,
|
||||
}
|
||||
|
||||
/// Trait that all tools must implement.
|
||||
#[async_trait::async_trait]
|
||||
pub trait Tool: Send + Sync {
|
||||
fn name(&self) -> &str;
|
||||
fn description(&self) -> &str;
|
||||
fn parameters(&self) -> serde_json::Value;
|
||||
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<String>;
|
||||
}
|
||||
|
||||
impl ToolExecutor {
|
||||
/// Create a new executor with no tools loaded.
|
||||
/// Create an empty executor (no tools registered).
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tools: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a tool implementation.
|
||||
/// Create an executor pre-populated with all built-in tools.
|
||||
pub fn with_defaults() -> Self {
|
||||
let mut executor = Self::new();
|
||||
executor.register_builtins();
|
||||
executor
|
||||
}
|
||||
|
||||
/// Register a single tool.
|
||||
pub fn register(&mut self, tool: Box<dyn Tool + Send + Sync>) {
|
||||
let name = tool.name().to_string();
|
||||
debug!("Registering tool: {}", name);
|
||||
debug!("Registering tool: {name}");
|
||||
self.tools.insert(name, tool);
|
||||
}
|
||||
|
||||
/// Return all registered tool definitions (for LLM function calling).
|
||||
/// Register all built-in tools.
|
||||
pub fn register_builtins(&mut self) {
|
||||
self.register(Box::new(ReadTool));
|
||||
self.register(Box::new(WriteTool));
|
||||
self.register(Box::new(EditTool));
|
||||
self.register(Box::new(BashTool::default()));
|
||||
self.register(Box::new(GrepTool));
|
||||
self.register(Box::new(FindTool));
|
||||
self.register(Box::new(LsTool));
|
||||
debug!("Registered {} built-in tools", BuiltinTool::all().len());
|
||||
}
|
||||
|
||||
/// Return `ToolDefinition` values for every registered tool (for LLM
|
||||
/// function-calling schemas).
|
||||
pub fn definitions(&self) -> Vec<ToolDefinition> {
|
||||
self.tools
|
||||
.values()
|
||||
@@ -48,17 +125,21 @@ impl ToolExecutor {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Execute a single tool call and return the result.
|
||||
/// Execute a single tool call, returning a structured result.
|
||||
pub async fn execute(&self, call: &ToolCall) -> anyhow::Result<ToolResult> {
|
||||
let span = tracing::span!(tracing::Level::INFO, "tool_exec", tool = %call.name);
|
||||
let _guard = span.enter();
|
||||
|
||||
match self.tools.get(&call.name) {
|
||||
Some(tool) => {
|
||||
debug!("Executing tool '{}' with args: {}", call.name, call.arguments);
|
||||
debug!(
|
||||
"Executing tool '{}' with args: {}",
|
||||
call.name,
|
||||
call.arguments.to_string()
|
||||
);
|
||||
match tool.execute(call.arguments.clone()).await {
|
||||
Ok(output) => {
|
||||
debug!("Tool '{}' succeeded", call.name);
|
||||
debug!("Tool '{}' succeeded ({} bytes)", call.name, output.len());
|
||||
Ok(ToolResult {
|
||||
tool_call_id: call.id.clone(),
|
||||
success: true,
|
||||
@@ -67,7 +148,7 @@ impl ToolExecutor {
|
||||
})
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Tool '{}' failed: {}", call.name, e);
|
||||
error!("Tool '{}' failed: {e}", call.name);
|
||||
Ok(ToolResult {
|
||||
tool_call_id: call.id.clone(),
|
||||
success: false,
|
||||
@@ -88,6 +169,21 @@ impl ToolExecutor {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Check whether a specific tool name is registered.
|
||||
pub fn has_tool(&self, name: &str) -> bool {
|
||||
self.tools.contains_key(name)
|
||||
}
|
||||
|
||||
/// Number of registered tools.
|
||||
pub fn len(&self) -> usize {
|
||||
self.tools.len()
|
||||
}
|
||||
|
||||
/// True if no tools are registered.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.tools.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ToolExecutor {
|
||||
@@ -96,6 +192,619 @@ impl Default for ToolExecutor {
|
||||
}
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// Built-in tools
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
// ─── ReadTool ───────────────────────────────────────────────────────────
|
||||
|
||||
/// Read a file, optionally with line numbers.
|
||||
pub struct ReadTool;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Tool for ReadTool {
|
||||
fn name(&self) -> &str {
|
||||
"read"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Read the contents of a file. Optionally show line numbers and limit output."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "Path to the file to read"
|
||||
},
|
||||
"offset": {
|
||||
"type": "integer",
|
||||
"description": "Line number to start from (1-indexed)",
|
||||
"default": 1
|
||||
},
|
||||
"limit": {
|
||||
"type": "integer",
|
||||
"description": "Maximum number of lines to read",
|
||||
"default": 500
|
||||
},
|
||||
"line_numbers": {
|
||||
"type": "boolean",
|
||||
"description": "Show line numbers next to each line",
|
||||
"default": true
|
||||
}
|
||||
},
|
||||
"required": ["file_path"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<String> {
|
||||
let file_path = args["file_path"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("'file_path' is required (string)"))?;
|
||||
let offset = args["offset"].as_u64().unwrap_or(1) as usize;
|
||||
let limit = args["limit"].as_u64().unwrap_or(500) as usize;
|
||||
let line_numbers = args["line_numbers"].as_bool().unwrap_or(true);
|
||||
|
||||
let content = tokio::fs::read_to_string(file_path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to read '{file_path}': {e}"))?;
|
||||
|
||||
let lines: Vec<&str> = content.lines().collect();
|
||||
let start = offset.saturating_sub(1).min(lines.len());
|
||||
let end = (start + limit).min(lines.len());
|
||||
let slice = &lines[start..end];
|
||||
|
||||
if slice.is_empty() {
|
||||
return Ok(format!("(empty — file has {} lines)", lines.len()));
|
||||
}
|
||||
|
||||
let out: String = if line_numbers {
|
||||
slice
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, line)| format!("{:>6}| {}", start + i + 1, line))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
} else {
|
||||
slice.join("\n")
|
||||
};
|
||||
|
||||
let total = lines.len();
|
||||
if end < total {
|
||||
Ok(format!(
|
||||
"{out}\n\n--- {}/{} lines shown (use offset={}, limit={limit} for more) ---",
|
||||
end - start,
|
||||
total,
|
||||
end + 1,
|
||||
))
|
||||
} else {
|
||||
Ok(format!("{out}\n\n--- {}/{} lines ---", end - start, total))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── WriteTool ──────────────────────────────────────────────────────────
|
||||
|
||||
/// Write or create a file.
|
||||
pub struct WriteTool;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Tool for WriteTool {
|
||||
fn name(&self) -> &str {
|
||||
"write"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Write content to a file, creating parent directories if needed. Overwrites existing content."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "Path to the file to write"
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "Content to write to the file"
|
||||
}
|
||||
},
|
||||
"required": ["file_path", "content"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<String> {
|
||||
let file_path = args["file_path"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("'file_path' is required (string)"))?;
|
||||
let content = args["content"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("'content' is required (string)"))?;
|
||||
|
||||
// Create parent directories
|
||||
if let Some(parent) = std::path::Path::new(file_path).parent() {
|
||||
tokio::fs::create_dir_all(parent)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to create parent directory: {e}"))?;
|
||||
}
|
||||
|
||||
tokio::fs::write(file_path, content)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to write '{file_path}': {e}"))?;
|
||||
|
||||
Ok(format!(
|
||||
"Wrote {} bytes to {file_path}",
|
||||
content.len()
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
// ─── EditTool ───────────────────────────────────────────────────────────
|
||||
|
||||
/// Apply a targeted edit to an existing file (find-and-replace).
|
||||
pub struct EditTool;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Tool for EditTool {
|
||||
fn name(&self) -> &str {
|
||||
"edit"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Apply a targeted find-and-replace edit to a file using fuzzy matching. \
|
||||
Use this instead of rewriting entire files."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "Path to the file to edit"
|
||||
},
|
||||
"old_string": {
|
||||
"type": "string",
|
||||
"description": "Exact text to find and replace"
|
||||
},
|
||||
"new_string": {
|
||||
"type": "string",
|
||||
"description": "Replacement text"
|
||||
}
|
||||
},
|
||||
"required": ["file_path", "old_string", "new_string"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<String> {
|
||||
let file_path = args["file_path"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("'file_path' is required (string)"))?;
|
||||
let old_string = args["old_string"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("'old_string' is required (string)"))?;
|
||||
let new_string = args["new_string"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("'new_string' is required (string)"))?;
|
||||
|
||||
let content = tokio::fs::read_to_string(file_path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to read '{file_path}': {e}"))?;
|
||||
|
||||
if !content.contains(old_string) {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Could not find the specified text in '{file_path}'"
|
||||
));
|
||||
}
|
||||
|
||||
let new_content = content.replace(old_string, new_string);
|
||||
let replacements = content.matches(old_string).count();
|
||||
|
||||
tokio::fs::write(file_path, &new_content)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to write '{file_path}': {e}"))?;
|
||||
|
||||
Ok(format!(
|
||||
"Applied {replacements} replacement(s) to {file_path} ({} → {} bytes)",
|
||||
content.len(),
|
||||
new_content.len()
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
// ─── BashTool ───────────────────────────────────────────────────────────
|
||||
|
||||
/// Configuration for the Bash tool.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BashToolConfig {
|
||||
/// Maximum execution timeout in seconds.
|
||||
pub timeout_secs: u64,
|
||||
/// Maximum output characters returned.
|
||||
pub max_output_chars: usize,
|
||||
}
|
||||
|
||||
impl Default for BashToolConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
timeout_secs: 120,
|
||||
max_output_chars: 100_000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute a shell command with timeout and output cap.
|
||||
pub struct BashTool {
|
||||
config: BashToolConfig,
|
||||
}
|
||||
|
||||
impl BashTool {
|
||||
pub fn new(config: BashToolConfig) -> Self {
|
||||
Self { config }
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for BashTool {
|
||||
fn default() -> Self {
|
||||
Self::new(BashToolConfig::default())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Tool for BashTool {
|
||||
fn name(&self) -> &str {
|
||||
"bash"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Execute a shell command. Supports pipes, redirects, and any POSIX command. \
|
||||
Timeout and output are capped for safety."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "The shell command to execute"
|
||||
},
|
||||
"workdir": {
|
||||
"type": "string",
|
||||
"description": "Working directory (default: current dir)",
|
||||
"default": "."
|
||||
},
|
||||
"description": {
|
||||
"type": "string",
|
||||
"description": "Brief description of what this command does (for audit log)"
|
||||
}
|
||||
},
|
||||
"required": ["command"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<String> {
|
||||
let command = args["command"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("'command' is required (string)"))?;
|
||||
let workdir = args["workdir"].as_str().unwrap_or(".");
|
||||
|
||||
// Basic security: refuse destructive commands without explicit allow
|
||||
let trimmed = command.trim();
|
||||
if trimmed.starts_with("rm -rf /")
|
||||
|| trimmed.starts_with("rm -rf /*")
|
||||
|| trimmed.starts_with("mkfs")
|
||||
|| trimmed.starts_with("dd if=")
|
||||
{
|
||||
return Err(anyhow::anyhow!(
|
||||
"Security: destructive command '{}' is blocked by default",
|
||||
trimmed.split_whitespace().next().unwrap_or("?")
|
||||
));
|
||||
}
|
||||
|
||||
debug!(
|
||||
"Bash tool: running '{}' in '{}' (timeout={}s)",
|
||||
command, workdir, self.config.timeout_secs
|
||||
);
|
||||
|
||||
let output = tokio::time::timeout(
|
||||
Duration::from_secs(self.config.timeout_secs),
|
||||
tokio::process::Command::new("sh")
|
||||
.arg("-c")
|
||||
.arg(command)
|
||||
.current_dir(workdir)
|
||||
.output(),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
anyhow::anyhow!(
|
||||
"Command timed out after {}s",
|
||||
self.config.timeout_secs
|
||||
)
|
||||
})?
|
||||
.map_err(|e| anyhow::anyhow!("Failed to execute command: {e}"))?;
|
||||
|
||||
let mut result = String::new();
|
||||
|
||||
if !output.stdout.is_empty() {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
result.push_str(&stdout);
|
||||
}
|
||||
|
||||
if !output.stderr.is_empty() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
if !result.is_empty() {
|
||||
result.push('\n');
|
||||
}
|
||||
result.push_str(&format!("STDERR:\n{stderr}"));
|
||||
}
|
||||
|
||||
let exit_code = output.status.code().unwrap_or(-1);
|
||||
if exit_code != 0 {
|
||||
result.push_str(&format!("\n(exit code: {exit_code})"));
|
||||
}
|
||||
|
||||
// Cap output
|
||||
if result.len() > self.config.max_output_chars {
|
||||
let truncated = result[..self.config.max_output_chars].to_string();
|
||||
result = format!(
|
||||
"{truncated}\n\n--- Output truncated ({} chars, cap {}) ---",
|
||||
result.len(),
|
||||
self.config.max_output_chars
|
||||
);
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── GrepTool ───────────────────────────────────────────────────────────
|
||||
|
||||
/// Search file contents using regex.
|
||||
pub struct GrepTool;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Tool for GrepTool {
|
||||
fn name(&self) -> &str {
|
||||
"grep"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Search file contents using regex patterns. Returns matching lines with file paths and line numbers."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"pattern": {
|
||||
"type": "string",
|
||||
"description": "Regex pattern to search for"
|
||||
},
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Directory or file to search in",
|
||||
"default": "."
|
||||
},
|
||||
"file_glob": {
|
||||
"type": "string",
|
||||
"description": "Filter file types (e.g., '*.rs', '*.py')",
|
||||
"default": ""
|
||||
},
|
||||
"max_results": {
|
||||
"type": "integer",
|
||||
"description": "Maximum number of results",
|
||||
"default": 50
|
||||
}
|
||||
},
|
||||
"required": ["pattern"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<String> {
|
||||
let pattern = args["pattern"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("'pattern' is required (string)"))?;
|
||||
let path = args["path"].as_str().unwrap_or(".");
|
||||
let file_glob = args["file_glob"].as_str().unwrap_or("");
|
||||
let max_results = args["max_results"].as_u64().unwrap_or(50) as usize;
|
||||
|
||||
let mut cmd = tokio::process::Command::new("rg");
|
||||
cmd.arg("--line-number")
|
||||
.arg("--with-filename")
|
||||
.arg("--no-heading")
|
||||
.arg("--max-count")
|
||||
.arg(max_results.to_string())
|
||||
.arg(pattern)
|
||||
.arg(path);
|
||||
|
||||
if !file_glob.is_empty() {
|
||||
cmd.arg("--glob").arg(file_glob);
|
||||
}
|
||||
|
||||
let output = cmd
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to run ripgrep: {e}"))?;
|
||||
|
||||
if output.status.success() {
|
||||
Ok(String::from_utf8_lossy(&output.stdout).to_string())
|
||||
} else if output.status.code() == Some(1) {
|
||||
Ok("(no matches found)".to_string())
|
||||
} else {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
Err(anyhow::anyhow!("rg failed: {stderr}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── FindTool ───────────────────────────────────────────────────────────
|
||||
|
||||
/// Find files by glob/name pattern.
|
||||
pub struct FindTool;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Tool for FindTool {
|
||||
fn name(&self) -> &str {
|
||||
"find"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Find files by name pattern (glob). Returns file paths sorted by modification time (newest first)."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"pattern": {
|
||||
"type": "string",
|
||||
"description": "Glob pattern to search for (e.g., '*.rs', '*config*')"
|
||||
},
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Directory to search in",
|
||||
"default": "."
|
||||
},
|
||||
"max_results": {
|
||||
"type": "integer",
|
||||
"description": "Maximum number of results",
|
||||
"default": 50
|
||||
}
|
||||
},
|
||||
"required": ["pattern"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<String> {
|
||||
let pattern = args["pattern"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("'pattern' is required (string)"))?;
|
||||
let path = args["path"].as_str().unwrap_or(".");
|
||||
let max_results = args["max_results"].as_u64().unwrap_or(50) as usize;
|
||||
|
||||
let output = tokio::process::Command::new("fd")
|
||||
.arg("--glob")
|
||||
.arg(pattern)
|
||||
.arg(path)
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to run fd (try 'find' directly): {e}"))?;
|
||||
|
||||
if output.status.success() {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).to_string();
|
||||
let results: Vec<&str> = stdout
|
||||
.lines()
|
||||
.take(max_results)
|
||||
.collect();
|
||||
if results.is_empty() {
|
||||
Ok("(no files found)".to_string())
|
||||
} else {
|
||||
Ok(format!(
|
||||
"Found {} file(s):\n{}",
|
||||
results.len(),
|
||||
results.join("\n")
|
||||
))
|
||||
}
|
||||
} else {
|
||||
// Fallback to POSIX find
|
||||
let output2 = tokio::process::Command::new("find")
|
||||
.arg(path)
|
||||
.arg("-name")
|
||||
.arg(pattern)
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("find also failed: {e}"))?;
|
||||
|
||||
let stdout2 = String::from_utf8_lossy(&output2.stdout).to_string();
|
||||
let results: Vec<&str> = stdout2
|
||||
.lines()
|
||||
.take(max_results)
|
||||
.collect();
|
||||
if results.is_empty() {
|
||||
Ok("(no files found)".to_string())
|
||||
} else {
|
||||
Ok(format!(
|
||||
"Found {} file(s) (via find):\n{}",
|
||||
results.len(),
|
||||
results.join("\n")
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── LsTool ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// List directory contents.
|
||||
pub struct LsTool;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Tool for LsTool {
|
||||
fn name(&self) -> &str {
|
||||
"ls"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"List files and directories at a given path with size and modification time."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Directory path to list",
|
||||
"default": "."
|
||||
},
|
||||
"long": {
|
||||
"type": "boolean",
|
||||
"description": "Show detailed listing (permissions, size, date)",
|
||||
"default": false
|
||||
},
|
||||
"all": {
|
||||
"type": "boolean",
|
||||
"description": "Show hidden files",
|
||||
"default": false
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<String> {
|
||||
let dir_path = args["path"].as_str().unwrap_or(".");
|
||||
let long = args["long"].as_bool().unwrap_or(false);
|
||||
let all = args["all"].as_bool().unwrap_or(false);
|
||||
|
||||
let mut cmd = tokio::process::Command::new("ls");
|
||||
if long {
|
||||
cmd.arg("-lh");
|
||||
}
|
||||
if all {
|
||||
cmd.arg("-A");
|
||||
}
|
||||
cmd.arg(dir_path);
|
||||
|
||||
let output = cmd
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to list directory: {e}"))?;
|
||||
|
||||
if output.status.success() {
|
||||
Ok(String::from_utf8_lossy(&output.stdout).to_string())
|
||||
} else {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
Err(anyhow::anyhow!("ls failed: {stderr}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Tests ──────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -143,4 +852,126 @@ mod tests {
|
||||
assert!(result.success);
|
||||
assert_eq!(result.output, "hello");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_unknown_tool_returns_error() {
|
||||
let executor = ToolExecutor::new();
|
||||
let call = ToolCall {
|
||||
id: "1".to_string(),
|
||||
name: "nonexistent".to_string(),
|
||||
arguments: serde_json::json!({}),
|
||||
};
|
||||
let result = executor.execute(&call).await.unwrap();
|
||||
assert!(!result.success);
|
||||
assert!(result.error.unwrap().contains("not available"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builtin_tool_names() {
|
||||
for t in BuiltinTool::all() {
|
||||
assert!(!t.name().is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_with_defaults_registers_seven_tools() {
|
||||
let executor = ToolExecutor::with_defaults();
|
||||
assert_eq!(executor.len(), 7);
|
||||
assert!(executor.has_tool("read"));
|
||||
assert!(executor.has_tool("write"));
|
||||
assert!(executor.has_tool("edit"));
|
||||
assert!(executor.has_tool("bash"));
|
||||
assert!(executor.has_tool("grep"));
|
||||
assert!(executor.has_tool("find"));
|
||||
assert!(executor.has_tool("ls"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_tool_basic() {
|
||||
// Create a temp file
|
||||
let tmp = std::env::temp_dir().join("gyandoot_test_read.txt");
|
||||
tokio::fs::write(&tmp, "line1\nline2\nline3\n")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let tool = ReadTool;
|
||||
let result = tool
|
||||
.execute(serde_json::json!({
|
||||
"file_path": tmp.to_str().unwrap(),
|
||||
"line_numbers": false
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(result.contains("line1"));
|
||||
assert!(result.contains("line2"));
|
||||
|
||||
tokio::fs::remove_file(&tmp).await.ok();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_write_tool_creates_file() {
|
||||
let tmp = std::env::temp_dir().join("gyandoot_test_write.txt");
|
||||
let tool = WriteTool;
|
||||
let result = tool
|
||||
.execute(serde_json::json!({
|
||||
"file_path": tmp.to_str().unwrap(),
|
||||
"content": "hello world"
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(result.contains("11 bytes"));
|
||||
assert!(tokio::fs::read_to_string(&tmp).await.unwrap() == "hello world");
|
||||
tokio::fs::remove_file(&tmp).await.ok();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bash_tool_basic() {
|
||||
let tool = BashTool::default();
|
||||
let result = tool
|
||||
.execute(serde_json::json!({"command": "echo hello"}))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(result.contains("hello"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bash_tool_security_block() {
|
||||
let tool = BashTool::default();
|
||||
let result = tool
|
||||
.execute(serde_json::json!({"command": "rm -rf /"}))
|
||||
.await;
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("blocked"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_edit_tool() {
|
||||
let tmp = std::env::temp_dir().join("gyandoot_test_edit.txt");
|
||||
tokio::fs::write(&tmp, "old content").await.unwrap();
|
||||
let tool = EditTool;
|
||||
let result = tool
|
||||
.execute(serde_json::json!({
|
||||
"file_path": tmp.to_str().unwrap(),
|
||||
"old_string": "old",
|
||||
"new_string": "new"
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(result.contains("1 replacement(s)"));
|
||||
assert_eq!(
|
||||
tokio::fs::read_to_string(&tmp).await.unwrap(),
|
||||
"new content"
|
||||
);
|
||||
tokio::fs::remove_file(&tmp).await.ok();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_ls_tool() {
|
||||
let tool = LsTool;
|
||||
let result = tool
|
||||
.execute(serde_json::json!({"path": "."}))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!result.is_empty());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user