Implement coding-agent-loop crate: tools, profiles, subagents, parallel execution

Add the core infrastructure for a programmable agentic coding loop:

- Core tool executors (read_file, write_file, edit_file, shell, grep, glob)
- Project doc discovery (AGENTS.md, provider-specific files, 32KB budget)
- Provider profiles: Anthropic (200K ctx), OpenAI (128K ctx, v4a apply_patch),
  Gemini (1M ctx) with provider-specific system prompts
- Subagent system with spawn/wait/close, depth limiting
- Parallel tool execution via futures::join_all when provider supports it
- Context window awareness with 80% threshold warning events

163 tests passing.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Bryan Helmkamp 2026-02-20 12:57:49 -04:00
parent 255a93c2e8
commit 6a53634b05
21 changed files with 6930 additions and 0 deletions

View file

@ -0,0 +1,26 @@
[package]
name = "coding-agent-loop"
edition.workspace = true
version.workspace = true
license.workspace = true
description = "A programmable agentic loop for coding agents"
repository = "https://github.com/brynary/attractor-rust"
readme = "README.md"
keywords = ["llm", "ai", "agent", "coding"]
categories = ["api-bindings"]
[dependencies]
unified-llm = { path = "../unified-llm" }
thiserror.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio.workspace = true
uuid.workspace = true
futures.workspace = true
async-trait.workspace = true
[dev-dependencies]
tokio = { workspace = true, features = ["test-util", "macros"] }
[lints]
workspace = true

View file

@ -0,0 +1,64 @@
use std::collections::HashMap;
#[derive(Debug, Clone)]
pub struct SessionConfig {
pub max_turns: usize,
pub max_tool_rounds_per_input: usize,
pub default_command_timeout_ms: u64,
pub max_command_timeout_ms: u64,
pub reasoning_effort: Option<String>,
pub tool_output_limits: HashMap<String, usize>,
pub tool_line_limits: HashMap<String, usize>,
pub enable_loop_detection: bool,
pub loop_detection_window: usize,
pub max_subagent_depth: usize,
}
impl Default for SessionConfig {
fn default() -> Self {
Self {
max_turns: 0,
max_tool_rounds_per_input: 200,
default_command_timeout_ms: 10_000,
max_command_timeout_ms: 600_000,
reasoning_effort: None,
tool_output_limits: HashMap::new(),
tool_line_limits: HashMap::new(),
enable_loop_detection: true,
loop_detection_window: 10,
max_subagent_depth: 1,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_config_values() {
let config = SessionConfig::default();
assert_eq!(config.max_turns, 0);
assert_eq!(config.max_tool_rounds_per_input, 200);
assert_eq!(config.default_command_timeout_ms, 10_000);
assert_eq!(config.max_command_timeout_ms, 600_000);
assert!(config.reasoning_effort.is_none());
assert!(config.tool_output_limits.is_empty());
assert!(config.tool_line_limits.is_empty());
assert!(config.enable_loop_detection);
assert_eq!(config.loop_detection_window, 10);
assert_eq!(config.max_subagent_depth, 1);
}
#[test]
fn config_with_custom_values() {
let config = SessionConfig {
max_turns: 50,
reasoning_effort: Some("high".into()),
..Default::default()
};
assert_eq!(config.max_turns, 50);
assert_eq!(config.reasoning_effort, Some("high".into()));
assert_eq!(config.max_tool_rounds_per_input, 200);
}
}

View file

@ -0,0 +1,67 @@
use unified_llm::error::SdkError;
#[derive(Debug, thiserror::Error)]
pub enum AgentError {
#[error("LLM error: {0}")]
Llm(#[from] SdkError),
#[error("Session is closed")]
SessionClosed,
#[error("Invalid state: {0}")]
InvalidState(String),
#[error("Tool execution error: {0}")]
ToolExecution(String),
#[error("IO error: {0}")]
Io(String),
#[error("Aborted")]
Aborted,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn agent_error_from_sdk_error() {
let sdk_err = SdkError::Network {
message: "connection refused".into(),
};
let agent_err = AgentError::from(sdk_err);
assert!(matches!(agent_err, AgentError::Llm(_)));
assert!(agent_err.to_string().contains("connection refused"));
}
#[test]
fn session_closed_display() {
let err = AgentError::SessionClosed;
assert_eq!(err.to_string(), "Session is closed");
}
#[test]
fn invalid_state_display() {
let err = AgentError::InvalidState("bad state".into());
assert_eq!(err.to_string(), "Invalid state: bad state");
}
#[test]
fn tool_execution_display() {
let err = AgentError::ToolExecution("command failed".into());
assert_eq!(err.to_string(), "Tool execution error: command failed");
}
#[test]
fn io_error_display() {
let err = AgentError::Io("file not found".into());
assert_eq!(err.to_string(), "IO error: file not found");
}
#[test]
fn aborted_display() {
let err = AgentError::Aborted;
assert_eq!(err.to_string(), "Aborted");
}
}

View file

@ -0,0 +1,109 @@
use crate::types::{EventKind, SessionEvent};
use std::collections::HashMap;
use std::time::SystemTime;
use tokio::sync::broadcast;
#[derive(Clone)]
pub struct EventEmitter {
sender: broadcast::Sender<SessionEvent>,
}
impl EventEmitter {
#[must_use]
pub fn new() -> Self {
let (sender, _) = broadcast::channel(1024);
Self { sender }
}
pub fn emit(
&self,
kind: EventKind,
session_id: String,
data: HashMap<String, serde_json::Value>,
) {
let event = SessionEvent {
kind,
timestamp: SystemTime::now(),
session_id,
data,
};
// Ignore send error (no receivers)
let _ = self.sender.send(event);
}
#[must_use]
pub fn subscribe(&self) -> broadcast::Receiver<SessionEvent> {
self.sender.subscribe()
}
}
impl Default for EventEmitter {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn emit_and_receive_event() {
let emitter = EventEmitter::new();
let mut receiver = emitter.subscribe();
emitter.emit(
EventKind::SessionStart,
"sess-1".into(),
HashMap::new(),
);
let event = receiver.recv().await.unwrap();
assert_eq!(event.kind, EventKind::SessionStart);
assert_eq!(event.session_id, "sess-1");
assert!(event.data.is_empty());
}
#[tokio::test]
async fn emit_with_data() {
let emitter = EventEmitter::new();
let mut receiver = emitter.subscribe();
let mut data = HashMap::new();
data.insert("text".into(), serde_json::json!("hello world"));
emitter.emit(EventKind::AssistantTextDelta, "sess-2".into(), data);
let event = receiver.recv().await.unwrap();
assert_eq!(event.kind, EventKind::AssistantTextDelta);
assert_eq!(event.data["text"], serde_json::json!("hello world"));
}
#[tokio::test]
async fn multiple_subscribers() {
let emitter = EventEmitter::new();
let mut rx1 = emitter.subscribe();
let mut rx2 = emitter.subscribe();
emitter.emit(EventKind::SessionEnd, "sess-3".into(), HashMap::new());
let e1 = rx1.recv().await.unwrap();
let e2 = rx2.recv().await.unwrap();
assert_eq!(e1.kind, EventKind::SessionEnd);
assert_eq!(e2.kind, EventKind::SessionEnd);
assert_eq!(e1.session_id, "sess-3");
assert_eq!(e2.session_id, "sess-3");
}
#[test]
fn emit_without_subscribers_does_not_panic() {
let emitter = EventEmitter::new();
emitter.emit(EventKind::Error, "sess-4".into(), HashMap::new());
}
#[test]
fn default_creates_emitter() {
let emitter = EventEmitter::default();
let _rx = emitter.subscribe();
}
}

View file

@ -0,0 +1,184 @@
use async_trait::async_trait;
#[derive(Debug, Clone)]
pub struct ExecResult {
pub stdout: String,
pub stderr: String,
pub exit_code: i32,
pub timed_out: bool,
pub duration_ms: u64,
}
#[derive(Debug, Clone)]
pub struct DirEntry {
pub name: String,
pub is_dir: bool,
pub size: Option<u64>,
}
#[derive(Debug, Clone, Default)]
pub struct GrepOptions {
pub glob_filter: Option<String>,
pub case_insensitive: bool,
pub max_results: Option<usize>,
}
#[async_trait]
pub trait ExecutionEnvironment: Send + Sync {
async fn read_file(&self, path: &str) -> Result<String, String>;
async fn write_file(&self, path: &str, content: &str) -> Result<(), String>;
async fn file_exists(&self, path: &str) -> Result<bool, String>;
async fn list_directory(&self, path: &str) -> Result<Vec<DirEntry>, String>;
async fn exec_command(
&self,
command: &str,
args: &[String],
timeout_ms: u64,
) -> Result<ExecResult, String>;
async fn grep(
&self,
pattern: &str,
path: &str,
options: &GrepOptions,
) -> Result<Vec<String>, String>;
async fn glob(&self, pattern: &str) -> Result<Vec<String>, String>;
async fn initialize(&self) -> Result<(), String>;
async fn cleanup(&self) -> Result<(), String>;
fn working_directory(&self) -> &str;
fn platform(&self) -> &str;
fn os_version(&self) -> String;
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
struct MockEnv;
#[async_trait]
impl ExecutionEnvironment for MockEnv {
async fn read_file(&self, _path: &str) -> Result<String, String> {
Ok("hello".into())
}
async fn write_file(&self, _path: &str, _content: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _path: &str) -> Result<bool, String> {
Ok(true)
}
async fn list_directory(&self, _path: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![DirEntry {
name: "test.rs".into(),
is_dir: false,
size: Some(100),
}])
}
async fn exec_command(
&self,
_command: &str,
_args: &[String],
_timeout_ms: u64,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: "output".into(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 10,
})
}
async fn grep(
&self,
_pattern: &str,
_path: &str,
_options: &GrepOptions,
) -> Result<Vec<String>, String> {
Ok(vec!["match".into()])
}
async fn glob(&self, _pattern: &str) -> Result<Vec<String>, String> {
Ok(vec!["file.rs".into()])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp"
}
fn platform(&self) -> &str {
"darwin"
}
fn os_version(&self) -> String {
"Darwin 24.0.0".into()
}
}
#[tokio::test]
async fn mock_env_read_file() {
let env: Arc<dyn ExecutionEnvironment> = Arc::new(MockEnv);
let result = env.read_file("test.rs").await.unwrap();
assert_eq!(result, "hello");
}
#[tokio::test]
async fn mock_env_exec_command() {
let env: Arc<dyn ExecutionEnvironment> = Arc::new(MockEnv);
let result = env.exec_command("echo", &[], 5000).await.unwrap();
assert_eq!(result.exit_code, 0);
assert!(!result.timed_out);
}
#[tokio::test]
async fn mock_env_list_directory() {
let env: Arc<dyn ExecutionEnvironment> = Arc::new(MockEnv);
let entries = env.list_directory("/tmp").await.unwrap();
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].name, "test.rs");
assert!(!entries[0].is_dir);
}
#[test]
fn exec_result_fields() {
let result = ExecResult {
stdout: "out".into(),
stderr: "err".into(),
exit_code: 1,
timed_out: true,
duration_ms: 5000,
};
assert_eq!(result.exit_code, 1);
assert!(result.timed_out);
assert_eq!(result.duration_ms, 5000);
}
#[test]
fn dir_entry_fields() {
let entry = DirEntry {
name: "src".into(),
is_dir: true,
size: None,
};
assert_eq!(entry.name, "src");
assert!(entry.is_dir);
assert!(entry.size.is_none());
}
#[test]
fn grep_options_defaults() {
let opts = GrepOptions::default();
assert!(opts.glob_filter.is_none());
assert!(!opts.case_insensitive);
assert!(opts.max_results.is_none());
}
#[test]
fn mock_env_platform() {
let env = MockEnv;
assert_eq!(env.platform(), "darwin");
assert_eq!(env.working_directory(), "/tmp");
assert_eq!(env.os_version(), "Darwin 24.0.0");
}
}

View file

@ -0,0 +1,278 @@
use crate::types::Turn;
use unified_llm::types::{ContentPart, Message, Role};
#[derive(Debug, Clone, Default)]
pub struct History {
turns: Vec<Turn>,
}
impl History {
pub fn new() -> Self {
Self { turns: Vec::new() }
}
pub fn push(&mut self, turn: Turn) {
self.turns.push(turn);
}
pub fn turns(&self) -> &[Turn] {
&self.turns
}
pub fn count_turns(&self) -> usize {
self.turns.len()
}
pub fn convert_to_messages(&self) -> Vec<Message> {
self.turns
.iter()
.map(|turn| match turn {
Turn::User { content, .. } => Message::user(content),
Turn::Assistant {
content,
tool_calls,
reasoning,
..
} => {
let mut parts: Vec<ContentPart> = Vec::new();
if let Some(reasoning_text) = reasoning {
parts.push(ContentPart::Thinking(
unified_llm::types::ThinkingData {
text: reasoning_text.clone(),
signature: None,
redacted: false,
},
));
}
if !content.is_empty() {
parts.push(ContentPart::text(content));
}
for tc in tool_calls {
parts.push(ContentPart::ToolCall(tc.clone()));
}
Message {
role: Role::Assistant,
content: parts,
name: None,
tool_call_id: None,
}
}
Turn::ToolResults { results, .. } => {
let content: Vec<ContentPart> = results
.iter()
.map(|r| ContentPart::ToolResult(r.clone()))
.collect();
// Use the first result's tool_call_id if available
let tool_call_id = results.first().map(|r| r.tool_call_id.clone());
Message {
role: Role::Tool,
content,
name: None,
tool_call_id,
}
}
Turn::System { content, .. } => Message::system(content),
Turn::Steering { content, .. } => Message {
role: Role::Developer,
content: vec![ContentPart::text(content)],
name: None,
tool_call_id: None,
},
})
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::SystemTime;
use unified_llm::types::{ToolCall, ToolResult, Usage};
#[test]
fn empty_history_produces_empty_messages() {
let history = History::new();
assert!(history.convert_to_messages().is_empty());
assert_eq!(history.count_turns(), 0);
}
#[test]
fn user_turn_maps_to_user_message() {
let mut history = History::new();
history.push(Turn::User {
content: "Hello".into(),
timestamp: SystemTime::now(),
});
let messages = history.convert_to_messages();
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].role, Role::User);
assert_eq!(messages[0].text(), "Hello");
}
#[test]
fn assistant_turn_maps_to_assistant_message() {
let mut history = History::new();
history.push(Turn::Assistant {
content: "Hi there".into(),
tool_calls: vec![],
reasoning: None,
usage: Usage::default(),
response_id: "resp_1".into(),
timestamp: SystemTime::now(),
});
let messages = history.convert_to_messages();
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].role, Role::Assistant);
assert_eq!(messages[0].text(), "Hi there");
}
#[test]
fn assistant_turn_with_tool_calls() {
let mut history = History::new();
let tc = ToolCall::new("call_1", "read_file", serde_json::json!({"path": "foo.rs"}));
history.push(Turn::Assistant {
content: "Let me read that".into(),
tool_calls: vec![tc],
reasoning: None,
usage: Usage::default(),
response_id: "resp_2".into(),
timestamp: SystemTime::now(),
});
let messages = history.convert_to_messages();
assert_eq!(messages[0].role, Role::Assistant);
let tool_call_parts: Vec<_> = messages[0]
.content
.iter()
.filter(|p| matches!(p, ContentPart::ToolCall(_)))
.collect();
assert_eq!(tool_call_parts.len(), 1);
}
#[test]
fn assistant_turn_with_reasoning() {
let mut history = History::new();
history.push(Turn::Assistant {
content: "The answer is 42".into(),
tool_calls: vec![],
reasoning: Some("Let me think about this...".into()),
usage: Usage::default(),
response_id: "resp_3".into(),
timestamp: SystemTime::now(),
});
let messages = history.convert_to_messages();
let thinking_parts: Vec<_> = messages[0]
.content
.iter()
.filter(|p| matches!(p, ContentPart::Thinking(_)))
.collect();
assert_eq!(thinking_parts.len(), 1);
}
#[test]
fn tool_results_turn_maps_to_tool_message() {
let mut history = History::new();
let result = ToolResult {
tool_call_id: "call_1".into(),
content: serde_json::json!("file contents here"),
is_error: false,
image_data: None,
image_media_type: None,
};
history.push(Turn::ToolResults {
results: vec![result],
timestamp: SystemTime::now(),
});
let messages = history.convert_to_messages();
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].role, Role::Tool);
assert_eq!(messages[0].tool_call_id, Some("call_1".into()));
}
#[test]
fn system_turn_maps_to_system_message() {
let mut history = History::new();
history.push(Turn::System {
content: "You are a coding assistant".into(),
timestamp: SystemTime::now(),
});
let messages = history.convert_to_messages();
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].role, Role::System);
assert_eq!(messages[0].text(), "You are a coding assistant");
}
#[test]
fn steering_turn_maps_to_developer_message() {
let mut history = History::new();
history.push(Turn::Steering {
content: "Focus on the main task".into(),
timestamp: SystemTime::now(),
});
let messages = history.convert_to_messages();
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].role, Role::Developer);
assert_eq!(messages[0].text(), "Focus on the main task");
}
#[test]
fn count_turns_matches_push_count() {
let mut history = History::new();
assert_eq!(history.count_turns(), 0);
history.push(Turn::User {
content: "First".into(),
timestamp: SystemTime::now(),
});
assert_eq!(history.count_turns(), 1);
history.push(Turn::Assistant {
content: "Second".into(),
tool_calls: vec![],
reasoning: None,
usage: Usage::default(),
response_id: "resp_1".into(),
timestamp: SystemTime::now(),
});
assert_eq!(history.count_turns(), 2);
}
#[test]
fn round_trip_preserves_content() {
let mut history = History::new();
history.push(Turn::User {
content: "Hello".into(),
timestamp: SystemTime::now(),
});
history.push(Turn::Assistant {
content: "Hi".into(),
tool_calls: vec![ToolCall::new(
"c1",
"shell",
serde_json::json!({"cmd": "ls"}),
)],
reasoning: Some("thinking...".into()),
usage: Usage {
input_tokens: 10,
output_tokens: 5,
total_tokens: 15,
..Default::default()
},
response_id: "resp_1".into(),
timestamp: SystemTime::now(),
});
history.push(Turn::ToolResults {
results: vec![ToolResult {
tool_call_id: "c1".into(),
content: serde_json::json!("file1.rs\nfile2.rs"),
is_error: false,
image_data: None,
image_media_type: None,
}],
timestamp: SystemTime::now(),
});
let messages = history.convert_to_messages();
assert_eq!(messages.len(), 3);
assert_eq!(messages[0].role, Role::User);
assert_eq!(messages[1].role, Role::Assistant);
assert_eq!(messages[2].role, Role::Tool);
}
}

View file

@ -0,0 +1,36 @@
pub mod config;
pub mod error;
pub mod event;
pub mod execution_env;
pub mod history;
pub mod local_env;
pub mod loop_detection;
pub mod profiles;
pub mod project_docs;
pub mod provider_profile;
pub mod session;
pub mod subagent;
pub mod tool_registry;
pub mod tools;
pub mod truncation;
pub mod types;
pub use config::SessionConfig;
pub use error::AgentError;
pub use event::EventEmitter;
pub use execution_env::{DirEntry, ExecResult, ExecutionEnvironment, GrepOptions};
pub use history::History;
pub use local_env::LocalExecutionEnvironment;
pub use loop_detection::detect_loop;
pub use project_docs::discover_project_docs;
pub use profiles::{AnthropicProfile, GeminiProfile, OpenAiProfile};
pub use provider_profile::ProviderProfile;
pub use session::Session;
pub use subagent::{SubAgent, SubAgentManager};
pub use tool_registry::ToolRegistry;
pub use tools::{
make_edit_file_tool, make_glob_tool, make_grep_tool, make_read_file_tool, make_shell_tool,
make_write_file_tool,
};
pub use truncation::{truncate_lines, truncate_output, truncate_tool_output, TruncationMode};
pub use types::{EventKind, SessionEvent, SessionState, Turn};

View file

@ -0,0 +1,572 @@
use crate::execution_env::{DirEntry, ExecResult, ExecutionEnvironment, GrepOptions};
use async_trait::async_trait;
use std::path::{Path, PathBuf};
use std::time::Instant;
use tokio::io::AsyncReadExt;
use tokio::process::Command;
pub struct LocalExecutionEnvironment {
working_directory: PathBuf,
}
impl LocalExecutionEnvironment {
#[must_use]
pub const fn new(working_directory: PathBuf) -> Self {
Self { working_directory }
}
fn format_line_numbered(content: &str) -> String {
use std::fmt::Write;
let lines: Vec<&str> = content.lines().collect();
let width = lines.len().to_string().len().max(1);
let mut result = String::new();
let mut line_num = 1;
for line in &lines {
let _ = writeln!(result, "{line_num:>width$} | {line}");
line_num += 1;
}
result
}
fn should_filter_env_var(key: &str) -> bool {
let lower = key.to_lowercase();
lower.ends_with("_api_key")
|| lower.ends_with("_secret")
|| lower.ends_with("_token")
|| lower.ends_with("_password")
|| lower.ends_with("_credential")
}
fn resolve_path(&self, path: &str) -> PathBuf {
let p = Path::new(path);
if p.is_absolute() {
p.to_path_buf()
} else {
self.working_directory.join(p)
}
}
}
#[async_trait]
impl ExecutionEnvironment for LocalExecutionEnvironment {
async fn read_file(&self, path: &str) -> Result<String, String> {
let full_path = self.resolve_path(path);
let content = tokio::fs::read_to_string(&full_path)
.await
.map_err(|e| format!("Failed to read {}: {e}", full_path.display()))?;
Ok(Self::format_line_numbered(&content))
}
async fn write_file(&self, path: &str, content: &str) -> Result<(), String> {
let full_path = self.resolve_path(path);
if let Some(parent) = full_path.parent() {
tokio::fs::create_dir_all(parent)
.await
.map_err(|e| format!("Failed to create parent dirs: {e}"))?;
}
tokio::fs::write(&full_path, content)
.await
.map_err(|e| format!("Failed to write {}: {e}", full_path.display()))
}
async fn file_exists(&self, path: &str) -> Result<bool, String> {
let full_path = self.resolve_path(path);
Ok(full_path.exists())
}
async fn list_directory(&self, path: &str) -> Result<Vec<DirEntry>, String> {
let full_path = self.resolve_path(path);
let mut entries = Vec::new();
let mut read_dir = tokio::fs::read_dir(&full_path)
.await
.map_err(|e| format!("Failed to read directory {}: {e}", full_path.display()))?;
while let Some(entry) = read_dir
.next_entry()
.await
.map_err(|e| format!("Failed to read entry: {e}"))?
{
let metadata = entry
.metadata()
.await
.map_err(|e| format!("Failed to read metadata: {e}"))?;
entries.push(DirEntry {
name: entry.file_name().to_string_lossy().into_owned(),
is_dir: metadata.is_dir(),
size: if metadata.is_file() {
Some(metadata.len())
} else {
None
},
});
}
entries.sort_by(|a, b| a.name.cmp(&b.name));
Ok(entries)
}
async fn exec_command(
&self,
command: &str,
args: &[String],
timeout_ms: u64,
) -> Result<ExecResult, String> {
let start = Instant::now();
let filtered_env: Vec<(String, String)> = std::env::vars()
.filter(|(key, _)| !Self::should_filter_env_var(key))
.collect();
let mut cmd = Command::new(command);
cmd.args(args)
.current_dir(&self.working_directory)
.env_clear()
.envs(filtered_env)
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped());
let mut child = cmd
.spawn()
.map_err(|e| format!("Failed to spawn {command}: {e}"))?;
let timeout_duration = std::time::Duration::from_millis(timeout_ms);
let (timed_out, exit_code) =
if let Ok(status_result) = tokio::time::timeout(timeout_duration, child.wait()).await {
let status =
status_result.map_err(|e| format!("Failed to wait for process: {e}"))?;
(false, status.code().unwrap_or(-1))
} else {
let _ = child.kill().await;
let _ = child.wait().await;
(true, -1)
};
let duration_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
let mut stdout_str = String::new();
if let Some(mut stdout) = child.stdout.take() {
let _ = stdout.read_to_string(&mut stdout_str).await;
}
let mut stderr_str = String::new();
if let Some(mut stderr) = child.stderr.take() {
let _ = stderr.read_to_string(&mut stderr_str).await;
}
Ok(ExecResult {
stdout: stdout_str,
stderr: stderr_str,
exit_code,
timed_out,
duration_ms,
})
}
async fn grep(
&self,
pattern: &str,
path: &str,
options: &GrepOptions,
) -> Result<Vec<String>, String> {
let full_path = self.resolve_path(path);
let mut args = vec!["-rn".to_string()];
if options.case_insensitive {
args.push("-i".into());
}
if let Some(ref glob_filter) = options.glob_filter {
args.push("--include".into());
args.push(glob_filter.clone());
}
if let Some(max) = options.max_results {
args.push("-m".into());
args.push(max.to_string());
}
args.push(pattern.into());
args.push(full_path.to_string_lossy().into_owned());
let output = std::process::Command::new("grep")
.args(&args)
.output()
.map_err(|e| format!("Failed to run grep: {e}"))?;
let stdout = String::from_utf8_lossy(&output.stdout);
let results: Vec<String> = stdout.lines().map(String::from).filter(|l| !l.is_empty()).collect();
Ok(results)
}
async fn glob(&self, pattern: &str) -> Result<Vec<String>, String> {
// Use find + fnmatch-style pattern via shell glob expansion
let full_pattern = if Path::new(pattern).is_absolute() {
pattern.to_string()
} else {
format!("{}/{pattern}", self.working_directory.display())
};
// Use shell globbing via ls
let output = std::process::Command::new("sh")
.args(["-c", &format!("ls -d {full_pattern} 2>/dev/null")])
.output()
.map_err(|e| format!("Failed to run glob: {e}"))?;
let stdout = String::from_utf8_lossy(&output.stdout);
let results: Vec<String> = stdout.lines().map(String::from).filter(|l| !l.is_empty()).collect();
Ok(results)
}
async fn initialize(&self) -> Result<(), String> {
tokio::fs::create_dir_all(&self.working_directory)
.await
.map_err(|e| format!("Failed to create working directory: {e}"))
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
self.working_directory.to_str().unwrap_or(".")
}
fn platform(&self) -> &str {
if cfg!(target_os = "macos") {
"darwin"
} else if cfg!(target_os = "linux") {
"linux"
} else if cfg!(target_os = "windows") {
"windows"
} else {
"unknown"
}
}
fn os_version(&self) -> String {
#[cfg(unix)]
{
let output = std::process::Command::new("uname")
.arg("-r")
.output();
match output {
Ok(out) => {
let version = String::from_utf8_lossy(&out.stdout).trim().to_string();
format!("{} {version}", self.platform())
}
Err(_) => self.platform().to_string(),
}
}
#[cfg(not(unix))]
{
self.platform().to_string()
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
fn temp_dir() -> PathBuf {
let dir = std::env::temp_dir().join(format!("local_env_test_{}", uuid::Uuid::new_v4()));
std::fs::create_dir_all(&dir).unwrap();
dir
}
#[tokio::test]
async fn read_file_with_line_numbers() {
let dir = temp_dir();
std::fs::write(dir.join("test.txt"), "hello\nworld\nfoo").unwrap();
let env = LocalExecutionEnvironment::new(dir.clone());
let result = env.read_file("test.txt").await.unwrap();
assert_eq!(result, "1 | hello\n2 | world\n3 | foo\n");
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn read_file_line_number_padding() {
let dir = temp_dir();
let content: String = (1..=12).map(|i| format!("line {i}\n")).collect();
std::fs::write(dir.join("padded.txt"), content.trim_end()).unwrap();
let env = LocalExecutionEnvironment::new(dir.clone());
let result = env.read_file("padded.txt").await.unwrap();
assert!(result.starts_with(" 1 | line 1\n"));
assert!(result.contains("12 | line 12\n"));
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn read_file_not_found() {
let dir = temp_dir();
let env = LocalExecutionEnvironment::new(dir.clone());
let result = env.read_file("nonexistent.txt").await;
assert!(result.is_err());
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn write_file_creates_parent_dirs() {
let dir = temp_dir();
let env = LocalExecutionEnvironment::new(dir.clone());
env.write_file("sub/dir/test.txt", "content").await.unwrap();
let written = std::fs::read_to_string(dir.join("sub/dir/test.txt")).unwrap();
assert_eq!(written, "content");
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn file_exists_true() {
let dir = temp_dir();
std::fs::write(dir.join("exists.txt"), "data").unwrap();
let env = LocalExecutionEnvironment::new(dir.clone());
assert!(env.file_exists("exists.txt").await.unwrap());
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn file_exists_false() {
let dir = temp_dir();
let env = LocalExecutionEnvironment::new(dir.clone());
assert!(!env.file_exists("nope.txt").await.unwrap());
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn list_directory_sorted() {
let dir = temp_dir();
std::fs::write(dir.join("b.txt"), "b").unwrap();
std::fs::write(dir.join("a.txt"), "a").unwrap();
std::fs::create_dir(dir.join("c_dir")).unwrap();
let env = LocalExecutionEnvironment::new(dir.clone());
let entries = env.list_directory(".").await.unwrap();
assert_eq!(entries.len(), 3);
assert_eq!(entries[0].name, "a.txt");
assert!(!entries[0].is_dir);
assert!(entries[0].size.is_some());
assert_eq!(entries[1].name, "b.txt");
assert_eq!(entries[2].name, "c_dir");
assert!(entries[2].is_dir);
assert!(entries[2].size.is_none());
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn exec_command_echo() {
let dir = temp_dir();
let env = LocalExecutionEnvironment::new(dir.clone());
let result = env
.exec_command("echo", &["hello".into()], 5000)
.await
.unwrap();
assert_eq!(result.stdout.trim(), "hello");
assert_eq!(result.exit_code, 0);
assert!(!result.timed_out);
assert!(result.duration_ms < 5000);
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn exec_command_exit_code() {
let dir = temp_dir();
let env = LocalExecutionEnvironment::new(dir.clone());
let result = env
.exec_command("sh", &["-c".into(), "exit 42".into()], 5000)
.await
.unwrap();
assert_eq!(result.exit_code, 42);
assert!(!result.timed_out);
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn exec_command_timeout() {
let dir = temp_dir();
let env = LocalExecutionEnvironment::new(dir.clone());
let result = env
.exec_command("sleep", &["10".into()], 200)
.await
.unwrap();
assert!(result.timed_out);
assert_eq!(result.exit_code, -1);
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn exec_command_stderr() {
let dir = temp_dir();
let env = LocalExecutionEnvironment::new(dir.clone());
let result = env
.exec_command("sh", &["-c".into(), "echo err >&2".into()], 5000)
.await
.unwrap();
assert_eq!(result.stderr.trim(), "err");
std::fs::remove_dir_all(&dir).unwrap();
}
#[test]
fn env_var_filtering() {
assert!(LocalExecutionEnvironment::should_filter_env_var(
"OPENAI_API_KEY"
));
assert!(LocalExecutionEnvironment::should_filter_env_var(
"ANTHROPIC_API_KEY"
));
assert!(LocalExecutionEnvironment::should_filter_env_var(
"DB_PASSWORD"
));
assert!(LocalExecutionEnvironment::should_filter_env_var(
"AWS_SECRET"
));
assert!(LocalExecutionEnvironment::should_filter_env_var(
"AUTH_TOKEN"
));
assert!(LocalExecutionEnvironment::should_filter_env_var(
"MY_CREDENTIAL"
));
// Case insensitive
assert!(LocalExecutionEnvironment::should_filter_env_var(
"my_api_key"
));
assert!(LocalExecutionEnvironment::should_filter_env_var(
"Some_Secret"
));
// Should not filter
assert!(!LocalExecutionEnvironment::should_filter_env_var("PATH"));
assert!(!LocalExecutionEnvironment::should_filter_env_var("HOME"));
assert!(!LocalExecutionEnvironment::should_filter_env_var("EDITOR"));
assert!(!LocalExecutionEnvironment::should_filter_env_var(
"SECRET_PATH"
));
}
#[test]
fn platform_is_known() {
let env = LocalExecutionEnvironment::new(PathBuf::from("/tmp"));
let platform = env.platform();
assert!(
platform == "darwin" || platform == "linux" || platform == "windows",
"Unknown platform: {platform}"
);
}
#[test]
fn os_version_contains_platform() {
let env = LocalExecutionEnvironment::new(PathBuf::from("/tmp"));
let version = env.os_version();
assert!(
version.contains(env.platform()),
"OS version should contain platform: {version}"
);
}
#[test]
fn working_directory_accessor() {
let env = LocalExecutionEnvironment::new(PathBuf::from("/tmp/test_dir"));
assert_eq!(env.working_directory(), "/tmp/test_dir");
}
#[tokio::test]
async fn initialize_creates_directory() {
let dir = std::env::temp_dir().join(format!("init_test_{}", uuid::Uuid::new_v4()));
let env = LocalExecutionEnvironment::new(dir.clone());
env.initialize().await.unwrap();
assert!(dir.exists());
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn grep_finds_matches() {
let dir = temp_dir();
std::fs::write(dir.join("test.rs"), "fn main() {\n println!(\"hello\");\n}\n").unwrap();
let env = LocalExecutionEnvironment::new(dir.clone());
let results = env
.grep("println", "test.rs", &GrepOptions::default())
.await
.unwrap();
assert_eq!(results.len(), 1);
assert!(results[0].contains("println"));
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn grep_case_insensitive() {
let dir = temp_dir();
std::fs::write(dir.join("test.txt"), "Hello\nhello\nHELLO\n").unwrap();
let env = LocalExecutionEnvironment::new(dir.clone());
let results = env
.grep(
"hello",
"test.txt",
&GrepOptions {
case_insensitive: true,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results.len(), 3);
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn grep_max_results() {
let dir = temp_dir();
std::fs::write(dir.join("test.txt"), "match1\nmatch2\nmatch3\nmatch4\n").unwrap();
let env = LocalExecutionEnvironment::new(dir.clone());
let results = env
.grep(
"match",
"test.txt",
&GrepOptions {
max_results: Some(2),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results.len(), 2);
std::fs::remove_dir_all(&dir).unwrap();
}
#[tokio::test]
async fn glob_finds_files() {
let dir = temp_dir();
std::fs::write(dir.join("a.rs"), "").unwrap();
std::fs::write(dir.join("b.rs"), "").unwrap();
std::fs::write(dir.join("c.txt"), "").unwrap();
let env = LocalExecutionEnvironment::new(dir.clone());
let results = env.glob("*.rs").await.unwrap();
assert_eq!(results.len(), 2);
std::fs::remove_dir_all(&dir).unwrap();
}
#[test]
fn format_line_numbered_empty() {
let result = LocalExecutionEnvironment::format_line_numbered("");
assert_eq!(result, "");
}
#[test]
fn format_line_numbered_single_line() {
let result = LocalExecutionEnvironment::format_line_numbered("hello");
assert_eq!(result, "1 | hello\n");
}
}

View file

@ -0,0 +1,215 @@
use crate::history::History;
use crate::types::Turn;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
fn tool_call_signature(name: &str, arguments: &serde_json::Value) -> u64 {
let mut hasher = DefaultHasher::new();
name.hash(&mut hasher);
let args_str = arguments.to_string();
args_str.hash(&mut hasher);
hasher.finish()
}
fn extract_signatures_from_assistant(turn: &Turn) -> Vec<u64> {
match turn {
Turn::Assistant { tool_calls, .. } => tool_calls
.iter()
.map(|tc| tool_call_signature(&tc.name, &tc.arguments))
.collect(),
_ => vec![],
}
}
pub fn detect_loop(history: &History, window_size: usize) -> bool {
// Extract tool call signatures from the last N assistant turns that have tool calls
let turns = history.turns();
let mut signatures: Vec<u64> = Vec::new();
// Walk backwards and collect signatures from assistant turns with tool calls
let mut count = 0;
for turn in turns.iter().rev() {
if count >= window_size {
break;
}
let sigs = extract_signatures_from_assistant(turn);
if !sigs.is_empty() {
// Combine all tool call signatures for this turn into a single signature
let mut hasher = DefaultHasher::new();
for sig in &sigs {
sig.hash(&mut hasher);
}
signatures.push(hasher.finish());
count += 1;
}
}
// Signatures are in reverse order; reverse to chronological
signatures.reverse();
if signatures.len() < 2 {
return false;
}
// Check repeating patterns of length 1, 2, 3
for pattern_len in 1..=3 {
if signatures.len() < pattern_len * 2 {
continue;
}
if is_repeating_pattern(&signatures, pattern_len) {
return true;
}
}
false
}
fn is_repeating_pattern(signatures: &[u64], pattern_len: usize) -> bool {
if signatures.len() < pattern_len * 2 {
return false;
}
let pattern = &signatures[signatures.len() - pattern_len..];
// Check that the preceding chunk matches
let preceding_start = signatures.len() - pattern_len * 2;
let preceding = &signatures[preceding_start..preceding_start + pattern_len];
pattern == preceding
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::SystemTime;
use unified_llm::types::{ToolCall, Usage};
fn assistant_with_tool(name: &str, args: serde_json::Value) -> Turn {
Turn::Assistant {
content: String::new(),
tool_calls: vec![ToolCall::new("call_1", name, args)],
reasoning: None,
usage: Usage::default(),
response_id: "resp".into(),
timestamp: SystemTime::now(),
}
}
#[test]
fn too_few_turns_returns_false() {
let history = History::new();
assert!(!detect_loop(&history, 10));
}
#[test]
fn single_turn_returns_false() {
let mut history = History::new();
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
assert!(!detect_loop(&history, 10));
}
#[test]
fn pattern_1_repeating_detected() {
let mut history = History::new();
// Same tool call repeated 3 times
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
assert!(detect_loop(&history, 10));
}
#[test]
fn pattern_2_repeating_detected() {
let mut history = History::new();
// A-B-A-B pattern
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
history.push(assistant_with_tool("read_file", serde_json::json!({"path": "foo.rs"})));
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
history.push(assistant_with_tool("read_file", serde_json::json!({"path": "foo.rs"})));
assert!(detect_loop(&history, 10));
}
#[test]
fn pattern_3_repeating_detected() {
let mut history = History::new();
// A-B-C-A-B-C pattern
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
history.push(assistant_with_tool("read_file", serde_json::json!({"path": "a.rs"})));
history.push(assistant_with_tool("grep", serde_json::json!({"pattern": "fn"})));
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
history.push(assistant_with_tool("read_file", serde_json::json!({"path": "a.rs"})));
history.push(assistant_with_tool("grep", serde_json::json!({"pattern": "fn"})));
assert!(detect_loop(&history, 10));
}
#[test]
fn non_repeating_returns_false() {
let mut history = History::new();
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
history.push(assistant_with_tool("read_file", serde_json::json!({"path": "a.rs"})));
history.push(assistant_with_tool("grep", serde_json::json!({"pattern": "fn"})));
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "cat"})));
assert!(!detect_loop(&history, 10));
}
#[test]
fn same_name_different_args_are_different() {
let mut history = History::new();
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "pwd"})));
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "cat"})));
assert!(!detect_loop(&history, 10));
}
#[test]
fn tool_call_signature_same_input_same_output() {
let sig1 = tool_call_signature("shell", &serde_json::json!({"cmd": "ls"}));
let sig2 = tool_call_signature("shell", &serde_json::json!({"cmd": "ls"}));
assert_eq!(sig1, sig2);
}
#[test]
fn tool_call_signature_different_name_different_output() {
let sig1 = tool_call_signature("shell", &serde_json::json!({"cmd": "ls"}));
let sig2 = tool_call_signature("read_file", &serde_json::json!({"cmd": "ls"}));
assert_ne!(sig1, sig2);
}
#[test]
fn tool_call_signature_different_args_different_output() {
let sig1 = tool_call_signature("shell", &serde_json::json!({"cmd": "ls"}));
let sig2 = tool_call_signature("shell", &serde_json::json!({"cmd": "pwd"}));
assert_ne!(sig1, sig2);
}
#[test]
fn user_turns_are_ignored() {
let mut history = History::new();
history.push(Turn::User {
content: "hello".into(),
timestamp: SystemTime::now(),
});
history.push(Turn::User {
content: "hello".into(),
timestamp: SystemTime::now(),
});
history.push(Turn::User {
content: "hello".into(),
timestamp: SystemTime::now(),
});
assert!(!detect_loop(&history, 10));
}
#[test]
fn window_size_limits_lookback() {
let mut history = History::new();
// Add non-repeating turns first
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "unique1"})));
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "unique2"})));
// Then repeating turns
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
history.push(assistant_with_tool("shell", serde_json::json!({"cmd": "ls"})));
// With window=2, we only see the last 2 which are repeating
assert!(detect_loop(&history, 2));
}
}

View file

@ -0,0 +1,275 @@
use crate::execution_env::ExecutionEnvironment;
use crate::provider_profile::ProviderProfile;
use crate::tool_registry::ToolRegistry;
use unified_llm::types::ToolDefinition;
use super::{build_env_context_block, stub_tool};
pub struct AnthropicProfile {
model: String,
registry: ToolRegistry,
}
impl AnthropicProfile {
#[must_use]
pub fn new(model: impl Into<String>) -> Self {
let mut registry = ToolRegistry::new();
registry.register(stub_tool(
"read_file",
"Read the contents of a file at the given path",
serde_json::json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "Path to the file to read" }
},
"required": ["path"]
}),
));
registry.register(stub_tool(
"write_file",
"Write content to a file at the given path",
serde_json::json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "Path to the file to write" },
"content": { "type": "string", "description": "Content to write" }
},
"required": ["path", "content"]
}),
));
registry.register(stub_tool(
"edit_file",
"Edit a file by replacing old text with new text",
serde_json::json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "Path to the file to edit" },
"old_text": { "type": "string", "description": "Text to find and replace" },
"new_text": { "type": "string", "description": "Replacement text" }
},
"required": ["path", "old_text", "new_text"]
}),
));
registry.register(stub_tool(
"shell",
"Execute a shell command",
serde_json::json!({
"type": "object",
"properties": {
"command": { "type": "string", "description": "Shell command to execute" },
"timeout_ms": { "type": "integer", "description": "Timeout in milliseconds" }
},
"required": ["command"]
}),
));
registry.register(stub_tool(
"grep",
"Search for a pattern in files",
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" }
},
"required": ["pattern", "path"]
}),
));
registry.register(stub_tool(
"glob",
"Find files matching a glob pattern",
serde_json::json!({
"type": "object",
"properties": {
"pattern": { "type": "string", "description": "Glob pattern to match files" }
},
"required": ["pattern"]
}),
));
Self {
model: model.into(),
registry,
}
}
}
impl ProviderProfile for AnthropicProfile {
fn id(&self) -> String {
"anthropic".into()
}
fn model(&self) -> String {
self.model.clone()
}
fn tool_registry(&self) -> &ToolRegistry {
&self.registry
}
fn build_system_prompt(
&self,
env: &dyn ExecutionEnvironment,
project_docs: &[String],
) -> String {
let env_block = build_env_context_block(env);
let docs_section = if project_docs.is_empty() {
String::new()
} else {
format!("\n\n{}", project_docs.join("\n\n"))
};
format!(
"You are Claude, an AI assistant by Anthropic. You help users with software engineering tasks.\n\n\
{env_block}\n\n\
# Tools\n\
Use the provided tools to interact with the codebase and environment.\
{docs_section}"
)
}
fn tools(&self) -> Vec<ToolDefinition> {
self.registry.definitions()
}
fn provider_options(&self) -> Option<serde_json::Value> {
None
}
fn supports_reasoning(&self) -> bool {
true
}
fn supports_streaming(&self) -> bool {
true
}
fn supports_parallel_tool_calls(&self) -> bool {
true
}
fn context_window_size(&self) -> usize {
200_000
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::execution_env::*;
use async_trait::async_trait;
struct TestEnv;
#[async_trait]
impl ExecutionEnvironment for TestEnv {
async fn read_file(&self, _: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(
&self,
_: &str,
_: &[String],
_: u64,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(
&self,
_: &str,
_: &str,
_: &GrepOptions,
) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/home/test"
}
fn platform(&self) -> &str {
"linux"
}
fn os_version(&self) -> String {
"Linux 6.1.0".into()
}
}
#[test]
fn anthropic_profile_identity() {
let profile = AnthropicProfile::new("claude-sonnet-4-20250514");
assert_eq!(profile.id(), "anthropic");
assert_eq!(profile.model(), "claude-sonnet-4-20250514");
}
#[test]
fn anthropic_profile_capabilities() {
let profile = AnthropicProfile::new("claude-sonnet-4-20250514");
assert!(profile.supports_reasoning());
assert!(profile.supports_streaming());
assert!(profile.supports_parallel_tool_calls());
assert_eq!(profile.context_window_size(), 200_000);
}
#[test]
fn anthropic_system_prompt_contains_env_context() {
let profile = AnthropicProfile::new("claude-sonnet-4-20250514");
let env = TestEnv;
let prompt = profile.build_system_prompt(&env, &[]);
assert!(prompt.contains("You are Claude, an AI assistant by Anthropic"));
assert!(prompt.contains("# Environment"));
assert!(prompt.contains("linux"));
assert!(prompt.contains("/home/test"));
assert!(prompt.contains("# Tools"));
}
#[test]
fn anthropic_system_prompt_includes_project_docs() {
let profile = AnthropicProfile::new("claude-sonnet-4-20250514");
let env = TestEnv;
let docs = vec!["# Project README".into(), "# CONTRIBUTING guide".into()];
let prompt = profile.build_system_prompt(&env, &docs);
assert!(prompt.contains("# Project README"));
assert!(prompt.contains("# CONTRIBUTING guide"));
}
#[test]
fn anthropic_tools_registered() {
let profile = AnthropicProfile::new("claude-sonnet-4-20250514");
let names = profile.tool_registry().names();
assert_eq!(names.len(), 6);
assert!(names.contains(&"read_file".to_string()));
assert!(names.contains(&"write_file".to_string()));
assert!(names.contains(&"edit_file".to_string()));
assert!(names.contains(&"shell".to_string()));
assert!(names.contains(&"grep".to_string()));
assert!(names.contains(&"glob".to_string()));
}
}

View file

@ -0,0 +1,263 @@
use crate::execution_env::ExecutionEnvironment;
use crate::provider_profile::ProviderProfile;
use crate::tool_registry::ToolRegistry;
use unified_llm::types::ToolDefinition;
use super::{build_env_context_block, stub_tool};
pub struct GeminiProfile {
model: String,
registry: ToolRegistry,
}
impl GeminiProfile {
#[must_use]
pub fn new(model: impl Into<String>) -> Self {
let mut registry = ToolRegistry::new();
registry.register(stub_tool(
"read_file",
"Read the contents of a file at the given path",
serde_json::json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "Path to the file to read" }
},
"required": ["path"]
}),
));
registry.register(stub_tool(
"write_file",
"Write content to a file at the given path",
serde_json::json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "Path to the file to write" },
"content": { "type": "string", "description": "Content to write" }
},
"required": ["path", "content"]
}),
));
registry.register(stub_tool(
"edit_file",
"Edit a file by replacing old text with new text",
serde_json::json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "Path to the file to edit" },
"old_text": { "type": "string", "description": "Text to find and replace" },
"new_text": { "type": "string", "description": "Replacement text" }
},
"required": ["path", "old_text", "new_text"]
}),
));
registry.register(stub_tool(
"shell",
"Execute a shell command",
serde_json::json!({
"type": "object",
"properties": {
"command": { "type": "string", "description": "Shell command to execute" },
"timeout_ms": { "type": "integer", "description": "Timeout in milliseconds" }
},
"required": ["command"]
}),
));
registry.register(stub_tool(
"grep",
"Search for a pattern in files",
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" }
},
"required": ["pattern", "path"]
}),
));
registry.register(stub_tool(
"glob",
"Find files matching a glob pattern",
serde_json::json!({
"type": "object",
"properties": {
"pattern": { "type": "string", "description": "Glob pattern to match files" }
},
"required": ["pattern"]
}),
));
Self {
model: model.into(),
registry,
}
}
}
impl ProviderProfile for GeminiProfile {
fn id(&self) -> String {
"gemini".into()
}
fn model(&self) -> String {
self.model.clone()
}
fn tool_registry(&self) -> &ToolRegistry {
&self.registry
}
fn build_system_prompt(
&self,
env: &dyn ExecutionEnvironment,
project_docs: &[String],
) -> String {
let env_block = build_env_context_block(env);
let docs_section = if project_docs.is_empty() {
String::new()
} else {
format!("\n\n{}", project_docs.join("\n\n"))
};
format!(
"You are a coding assistant powered by Gemini. You help users with software engineering tasks.\n\n\
{env_block}\n\n\
# Tools\n\
Use the provided tools to interact with the codebase and environment.\
{docs_section}"
)
}
fn tools(&self) -> Vec<ToolDefinition> {
self.registry.definitions()
}
fn provider_options(&self) -> Option<serde_json::Value> {
None
}
fn supports_reasoning(&self) -> bool {
true
}
fn supports_streaming(&self) -> bool {
true
}
fn supports_parallel_tool_calls(&self) -> bool {
true
}
fn context_window_size(&self) -> usize {
1_000_000
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::execution_env::*;
use async_trait::async_trait;
struct TestEnv;
#[async_trait]
impl ExecutionEnvironment for TestEnv {
async fn read_file(&self, _: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(
&self,
_: &str,
_: &[String],
_: u64,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(
&self,
_: &str,
_: &str,
_: &GrepOptions,
) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/home/test"
}
fn platform(&self) -> &str {
"linux"
}
fn os_version(&self) -> String {
"Linux 6.1.0".into()
}
}
#[test]
fn gemini_profile_identity() {
let profile = GeminiProfile::new("gemini-2.0-flash");
assert_eq!(profile.id(), "gemini");
assert_eq!(profile.model(), "gemini-2.0-flash");
}
#[test]
fn gemini_profile_capabilities() {
let profile = GeminiProfile::new("gemini-2.0-flash");
assert!(profile.supports_reasoning());
assert!(profile.supports_streaming());
assert!(profile.supports_parallel_tool_calls());
assert_eq!(profile.context_window_size(), 1_000_000);
}
#[test]
fn gemini_system_prompt_contains_env_context() {
let profile = GeminiProfile::new("gemini-2.0-flash");
let env = TestEnv;
let prompt = profile.build_system_prompt(&env, &[]);
assert!(prompt.contains("powered by Gemini"));
assert!(prompt.contains("# Environment"));
assert!(prompt.contains("linux"));
}
#[test]
fn gemini_tools_registered() {
let profile = GeminiProfile::new("gemini-2.0-flash");
let names = profile.tool_registry().names();
assert_eq!(names.len(), 6);
assert!(names.contains(&"read_file".to_string()));
assert!(names.contains(&"write_file".to_string()));
assert!(names.contains(&"edit_file".to_string()));
assert!(names.contains(&"shell".to_string()));
assert!(names.contains(&"grep".to_string()));
assert!(names.contains(&"glob".to_string()));
}
}

View file

@ -0,0 +1,111 @@
pub mod anthropic;
pub mod gemini;
pub mod openai;
pub use anthropic::AnthropicProfile;
pub use gemini::GeminiProfile;
pub use openai::OpenAiProfile;
use crate::execution_env::ExecutionEnvironment;
use crate::tool_registry::RegisteredTool;
use std::sync::Arc;
use unified_llm::types::ToolDefinition;
#[must_use]
pub fn build_env_context_block(env: &dyn ExecutionEnvironment) -> String {
format!(
"# Environment\n- Working directory: {}\n- Platform: {}\n- OS: {}",
env.working_directory(),
env.platform(),
env.os_version()
)
}
#[must_use]
pub fn stub_tool(name: &str, description: &str, parameters: serde_json::Value) -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: name.into(),
description: description.into(),
parameters,
},
executor: Arc::new(|_args, _env| {
Box::pin(async { Err("Tool not yet connected to execution environment".into()) })
}),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::execution_env::*;
use async_trait::async_trait;
struct TestEnv;
#[async_trait]
impl ExecutionEnvironment for TestEnv {
async fn read_file(&self, _: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(
&self,
_: &str,
_: &[String],
_: u64,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(
&self,
_: &str,
_: &str,
_: &GrepOptions,
) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/home/test"
}
fn platform(&self) -> &str {
"linux"
}
fn os_version(&self) -> String {
"Linux 6.1.0".into()
}
}
#[test]
fn env_context_block_contains_platform() {
let env = TestEnv;
let block = build_env_context_block(&env);
assert!(block.contains("# Environment"));
assert!(block.contains("linux"));
assert!(block.contains("/home/test"));
assert!(block.contains("Linux 6.1.0"));
}
}

View file

@ -0,0 +1,705 @@
use crate::execution_env::ExecutionEnvironment;
use crate::provider_profile::ProviderProfile;
use crate::tool_registry::{RegisteredTool, ToolRegistry};
use std::sync::Arc;
use unified_llm::types::ToolDefinition;
use super::{build_env_context_block, stub_tool};
pub struct OpenAiProfile {
model: String,
registry: ToolRegistry,
}
impl OpenAiProfile {
#[must_use]
pub fn new(model: impl Into<String>) -> Self {
let mut registry = ToolRegistry::new();
registry.register(stub_tool(
"read_file",
"Read the contents of a file at the given path",
serde_json::json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "Path to the file to read" }
},
"required": ["path"]
}),
));
registry.register(stub_tool(
"write_file",
"Write content to a file at the given path",
serde_json::json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "Path to the file to write" },
"content": { "type": "string", "description": "Content to write" }
},
"required": ["path", "content"]
}),
));
registry.register(stub_tool(
"shell",
"Execute a shell command",
serde_json::json!({
"type": "object",
"properties": {
"command": { "type": "string", "description": "Shell command to execute" },
"timeout_ms": { "type": "integer", "description": "Timeout in milliseconds" }
},
"required": ["command"]
}),
));
registry.register(stub_tool(
"grep",
"Search for a pattern in files",
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" }
},
"required": ["pattern", "path"]
}),
));
registry.register(stub_tool(
"glob",
"Find files matching a glob pattern",
serde_json::json!({
"type": "object",
"properties": {
"pattern": { "type": "string", "description": "Glob pattern to match files" }
},
"required": ["pattern"]
}),
));
registry.register(make_apply_patch_tool());
Self {
model: model.into(),
registry,
}
}
}
impl ProviderProfile for OpenAiProfile {
fn id(&self) -> String {
"openai".into()
}
fn model(&self) -> String {
self.model.clone()
}
fn tool_registry(&self) -> &ToolRegistry {
&self.registry
}
fn build_system_prompt(
&self,
env: &dyn ExecutionEnvironment,
project_docs: &[String],
) -> String {
let env_block = build_env_context_block(env);
let docs_section = if project_docs.is_empty() {
String::new()
} else {
format!("\n\n{}", project_docs.join("\n\n"))
};
format!(
"You are a coding assistant. You help users with software engineering tasks.\n\n\
{env_block}\n\n\
# Tools\n\
Use the provided tools to interact with the codebase and environment.\n\
Use apply_patch for file edits when possible.\
{docs_section}"
)
}
fn tools(&self) -> Vec<ToolDefinition> {
self.registry.definitions()
}
fn provider_options(&self) -> Option<serde_json::Value> {
None
}
fn supports_reasoning(&self) -> bool {
true
}
fn supports_streaming(&self) -> bool {
true
}
fn supports_parallel_tool_calls(&self) -> bool {
true
}
fn context_window_size(&self) -> usize {
128_000
}
}
// -- apply_patch v4a format --
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Change {
Remove(String),
Add(String),
Context(String),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Hunk {
pub context_line: String,
pub changes: Vec<Change>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PatchOperation {
Add { path: String, content: String },
Delete { path: String },
Update { path: String, hunks: Vec<Hunk> },
}
/// Parses a v4a format patch string into a list of patch operations.
///
/// # Errors
/// Returns an error if the patch format is invalid.
pub fn parse_v4a_patch(text: &str) -> Result<Vec<PatchOperation>, String> {
let lines: Vec<&str> = text.lines().collect();
let mut ops = Vec::new();
let mut i = 0;
// Find "*** Begin Patch"
while i < lines.len() {
if lines[i].trim() == "*** Begin Patch" {
i += 1;
break;
}
i += 1;
}
while i < lines.len() {
let line = lines[i].trim();
if line == "*** End Patch" {
break;
}
if let Some(path) = line.strip_prefix("*** Add File: ") {
let path = path.to_string();
i += 1;
let mut content = String::new();
while i < lines.len() {
let l = lines[i];
if l.starts_with("*** ") {
break;
}
if let Some(text_line) = l.strip_prefix('+') {
if !content.is_empty() {
content.push('\n');
}
content.push_str(text_line);
} else {
return Err(format!("Expected '+' prefix in Add File block, got: {l}"));
}
i += 1;
}
ops.push(PatchOperation::Add { path, content });
} else if let Some(path) = line.strip_prefix("*** Delete File: ") {
ops.push(PatchOperation::Delete {
path: path.to_string(),
});
i += 1;
} else if let Some(path) = line.strip_prefix("*** Update File: ") {
let path = path.to_string();
i += 1;
let mut hunks = Vec::new();
while i < lines.len() {
let l = lines[i];
if l.starts_with("*** ") && !l.starts_with("@@ ") {
break;
}
if l.starts_with("@@ ") && l.ends_with(" @@") {
let context_line = l[3..l.len() - 3].to_string();
i += 1;
let mut changes = Vec::new();
while i < lines.len() {
let cl = lines[i];
if cl.starts_with("*** ") || (cl.starts_with("@@ ") && cl.ends_with(" @@"))
{
break;
}
if let Some(removed) = cl.strip_prefix('-') {
changes.push(Change::Remove(removed.to_string()));
} else if let Some(added) = cl.strip_prefix('+') {
changes.push(Change::Add(added.to_string()));
} else if let Some(ctx) = cl.strip_prefix(' ') {
changes.push(Change::Context(ctx.to_string()));
} else if cl.is_empty() {
changes.push(Change::Context(String::new()));
} else {
return Err(format!(
"Unexpected line in hunk (expected +, -, or space prefix): {cl}"
));
}
i += 1;
}
hunks.push(Hunk {
context_line,
changes,
});
} else {
return Err(format!("Expected @@ context @@ line, got: {l}"));
}
}
ops.push(PatchOperation::Update { path, hunks });
} else {
return Err(format!("Unexpected line in patch: {line}"));
}
}
Ok(ops)
}
/// Applies a list of patch operations using the given execution environment.
///
/// # Errors
/// Returns an error if any file operation fails.
pub async fn apply_patch_operations(
ops: &[PatchOperation],
env: &dyn ExecutionEnvironment,
) -> Result<String, String> {
let mut results = Vec::new();
for op in ops {
match op {
PatchOperation::Add { path, content } => {
env.write_file(path, content).await?;
results.push(format!("Added file: {path}"));
}
PatchOperation::Delete { path } => {
env.write_file(path, "").await?;
results.push(format!("Deleted file: {path}"));
}
PatchOperation::Update { path, hunks } => {
let original = env.read_file(path).await?;
let updated = apply_hunks(&original, hunks)?;
env.write_file(path, &updated).await?;
results.push(format!("Updated file: {path}"));
}
}
}
Ok(results.join("\n"))
}
fn apply_hunks(content: &str, hunks: &[Hunk]) -> Result<String, String> {
let mut lines: Vec<String> = content.lines().map(String::from).collect();
// Apply hunks in reverse order to preserve line indices
for hunk in hunks.iter().rev() {
let context_pos = lines
.iter()
.position(|l| l.trim() == hunk.context_line.trim())
.ok_or_else(|| {
format!(
"Could not find context line in file: '{}'",
hunk.context_line
)
})?;
// Build what we expect to find and what to replace with
let mut new_lines: Vec<String> = Vec::new();
// The context line itself is part of the hunk context
// We start replacing at context_pos
new_lines.push(lines[context_pos].clone());
let mut file_idx = context_pos + 1;
for change in &hunk.changes {
match change {
Change::Remove(_) => {
file_idx += 1;
}
Change::Add(text) => {
new_lines.push(text.clone());
}
Change::Context(_) => {
if file_idx < lines.len() {
new_lines.push(lines[file_idx].clone());
}
file_idx += 1;
}
}
}
// Calculate total lines consumed from original (context_line + removes + context changes)
let total_original_lines = 1 + hunk
.changes
.iter()
.filter(|c| matches!(c, Change::Remove(_) | Change::Context(_)))
.count();
// Replace the range
let end = (context_pos + total_original_lines).min(lines.len());
lines.splice(context_pos..end, new_lines);
}
Ok(lines.join("\n"))
}
fn make_apply_patch_tool() -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: "apply_patch".into(),
description: "Apply a v4a format patch to modify files".into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"patch": {
"type": "string",
"description": "The patch content in v4a format"
}
},
"required": ["patch"]
}),
},
executor: Arc::new(|args, env| {
Box::pin(async move {
let patch_text = args
.get("patch")
.and_then(|v| v.as_str())
.ok_or_else(|| "Missing required parameter: patch".to_string())?;
let ops = parse_v4a_patch(patch_text)?;
apply_patch_operations(&ops, env.as_ref()).await
})
}),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::execution_env::*;
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Mutex;
struct TestEnv;
#[async_trait]
impl ExecutionEnvironment for TestEnv {
async fn read_file(&self, _: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(
&self,
_: &str,
_: &[String],
_: u64,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(
&self,
_: &str,
_: &str,
_: &GrepOptions,
) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/home/test"
}
fn platform(&self) -> &str {
"linux"
}
fn os_version(&self) -> String {
"Linux 6.1.0".into()
}
}
struct MockFileEnv {
files: Mutex<HashMap<String, String>>,
}
impl MockFileEnv {
fn new(files: HashMap<String, String>) -> Self {
Self {
files: Mutex::new(files),
}
}
}
#[async_trait]
impl ExecutionEnvironment for MockFileEnv {
async fn read_file(&self, path: &str) -> Result<String, String> {
self.files
.lock()
.unwrap()
.get(path)
.cloned()
.ok_or_else(|| format!("File not found: {path}"))
}
async fn write_file(&self, path: &str, content: &str) -> Result<(), String> {
self.files
.lock()
.unwrap()
.insert(path.to_string(), content.to_string());
Ok(())
}
async fn file_exists(&self, path: &str) -> Result<bool, String> {
Ok(self.files.lock().unwrap().contains_key(path))
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(
&self,
_: &str,
_: &[String],
_: u64,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(
&self,
_: &str,
_: &str,
_: &GrepOptions,
) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp"
}
fn platform(&self) -> &str {
"linux"
}
fn os_version(&self) -> String {
"Linux 6.1.0".into()
}
}
#[test]
fn openai_profile_identity() {
let profile = OpenAiProfile::new("o3-mini");
assert_eq!(profile.id(), "openai");
assert_eq!(profile.model(), "o3-mini");
}
#[test]
fn openai_profile_capabilities() {
let profile = OpenAiProfile::new("o3-mini");
assert!(profile.supports_reasoning());
assert!(profile.supports_streaming());
assert!(profile.supports_parallel_tool_calls());
assert_eq!(profile.context_window_size(), 128_000);
}
#[test]
fn openai_system_prompt_contains_env_context() {
let profile = OpenAiProfile::new("o3-mini");
let env = TestEnv;
let prompt = profile.build_system_prompt(&env, &[]);
assert!(prompt.contains("You are a coding assistant"));
assert!(prompt.contains("# Environment"));
assert!(prompt.contains("linux"));
assert!(prompt.contains("apply_patch"));
}
#[test]
fn openai_tools_registered() {
let profile = OpenAiProfile::new("o3-mini");
let names = profile.tool_registry().names();
assert_eq!(names.len(), 6);
assert!(names.contains(&"read_file".to_string()));
assert!(names.contains(&"write_file".to_string()));
assert!(names.contains(&"shell".to_string()));
assert!(names.contains(&"grep".to_string()));
assert!(names.contains(&"glob".to_string()));
assert!(names.contains(&"apply_patch".to_string()));
}
#[test]
fn parse_v4a_add_file() {
let patch = "\
*** Begin Patch
*** Add File: src/new_file.rs
+fn main() {
+ println!(\"hello\");
+}
*** End Patch";
let ops = parse_v4a_patch(patch).unwrap();
assert_eq!(ops.len(), 1);
assert_eq!(
ops[0],
PatchOperation::Add {
path: "src/new_file.rs".into(),
content: "fn main() {\n println!(\"hello\");\n}".into(),
}
);
}
#[test]
fn parse_v4a_delete_file() {
let patch = "\
*** Begin Patch
*** Delete File: src/old_file.rs
*** End Patch";
let ops = parse_v4a_patch(patch).unwrap();
assert_eq!(ops.len(), 1);
assert_eq!(
ops[0],
PatchOperation::Delete {
path: "src/old_file.rs".into(),
}
);
}
#[test]
fn parse_v4a_update_file() {
let patch = "\
*** Begin Patch
*** Update File: src/lib.rs
@@ fn hello() @@
- println!(\"old\");
+ println!(\"new\");
*** End Patch";
let ops = parse_v4a_patch(patch).unwrap();
assert_eq!(ops.len(), 1);
match &ops[0] {
PatchOperation::Update { path, hunks } => {
assert_eq!(path, "src/lib.rs");
assert_eq!(hunks.len(), 1);
assert_eq!(hunks[0].context_line, "fn hello()");
assert_eq!(hunks[0].changes.len(), 2);
assert_eq!(
hunks[0].changes[0],
Change::Remove(" println!(\"old\");".into())
);
assert_eq!(
hunks[0].changes[1],
Change::Add(" println!(\"new\");".into())
);
}
_ => panic!("Expected Update operation"),
}
}
#[test]
fn parse_v4a_multi_operation() {
let patch = "\
*** Begin Patch
*** Add File: src/a.rs
+// file a
*** Delete File: src/b.rs
*** Update File: src/c.rs
@@ fn main() @@
- old_call();
+ new_call();
*** End Patch";
let ops = parse_v4a_patch(patch).unwrap();
assert_eq!(ops.len(), 3);
assert!(matches!(&ops[0], PatchOperation::Add { .. }));
assert!(matches!(&ops[1], PatchOperation::Delete { .. }));
assert!(matches!(&ops[2], PatchOperation::Update { .. }));
}
#[tokio::test]
async fn apply_patch_add_file() {
let env = MockFileEnv::new(HashMap::new());
let ops = vec![PatchOperation::Add {
path: "src/new.rs".into(),
content: "fn new() {}".into(),
}];
let result = apply_patch_operations(&ops, &env).await.unwrap();
assert!(result.contains("Added file: src/new.rs"));
let content = env.read_file("src/new.rs").await.unwrap();
assert_eq!(content, "fn new() {}");
}
#[tokio::test]
async fn apply_patch_update_file() {
let mut files = HashMap::new();
files.insert(
"src/lib.rs".to_string(),
"fn hello() {\n println!(\"old\");\n}".to_string(),
);
let env = MockFileEnv::new(files);
let ops = vec![PatchOperation::Update {
path: "src/lib.rs".into(),
hunks: vec![Hunk {
context_line: "fn hello() {".into(),
changes: vec![
Change::Remove(" println!(\"old\");".into()),
Change::Add(" println!(\"new\");".into()),
],
}],
}];
let result = apply_patch_operations(&ops, &env).await.unwrap();
assert!(result.contains("Updated file: src/lib.rs"));
let content = env.read_file("src/lib.rs").await.unwrap();
assert!(content.contains("println!(\"new\")"));
assert!(!content.contains("println!(\"old\")"));
}
}

View file

@ -0,0 +1,225 @@
use crate::execution_env::ExecutionEnvironment;
const BUDGET_BYTES: usize = 32768;
pub async fn discover_project_docs(
env: &dyn ExecutionEnvironment,
git_root: &str,
working_dir: &str,
provider_id: &str,
) -> Vec<String> {
let directories = build_directory_walk(git_root, working_dir);
let candidate_filenames: Vec<&str> = match provider_id {
"anthropic" => vec!["AGENTS.md", "CLAUDE.md"],
"openai" => vec!["AGENTS.md", ".github/copilot-instructions.md"],
"gemini" => vec!["AGENTS.md", "GEMINI.md"],
_ => vec!["AGENTS.md"],
};
let mut results = Vec::new();
let mut budget_remaining = BUDGET_BYTES;
for dir in &directories {
for filename in &candidate_filenames {
let path = format!("{dir}/{filename}");
if let Ok(content) = env.read_file(&path).await {
if content.is_empty() {
continue;
}
if content.len() <= budget_remaining {
budget_remaining -= content.len();
results.push(content);
} else if budget_remaining > 0 {
let truncated = truncate_to_budget(&content, budget_remaining);
budget_remaining = 0;
results.push(truncated);
}
}
}
}
results
}
fn build_directory_walk(git_root: &str, working_dir: &str) -> Vec<String> {
let mut dirs = vec![git_root.to_string()];
if working_dir == git_root {
return dirs;
}
// Strip git_root prefix to get relative path components
let relative = working_dir
.strip_prefix(git_root)
.and_then(|s| s.strip_prefix('/'))
.unwrap_or("");
if relative.is_empty() {
return dirs;
}
let mut current = git_root.to_string();
let parts: Vec<&str> = relative.split('/').collect();
for part in parts {
current = format!("{current}/{part}");
dirs.push(current.clone());
}
dirs
}
fn truncate_to_budget(content: &str, budget: usize) -> String {
const MARKER: &str = "... [truncated]";
if budget <= MARKER.len() {
return MARKER[..budget].to_string();
}
let usable = budget - MARKER.len();
// Find the last valid char boundary within usable bytes
let mut end = usable;
while end > 0 && !content.is_char_boundary(end) {
end -= 1;
}
format!("{}{MARKER}", &content[..end])
}
#[cfg(test)]
mod tests {
use super::*;
use crate::execution_env::*;
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
struct DocEnv {
files: HashMap<String, String>,
}
#[async_trait]
impl ExecutionEnvironment for DocEnv {
async fn read_file(&self, path: &str) -> Result<String, String> {
self.files
.get(path)
.cloned()
.ok_or_else(|| format!("not found: {path}"))
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, path: &str) -> Result<bool, String> {
Ok(self.files.contains_key(path))
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(&self, _: &str, _: &str, _: &GrepOptions) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp"
}
fn platform(&self) -> &str {
"darwin"
}
fn os_version(&self) -> String {
String::new()
}
}
#[tokio::test]
async fn discovers_agents_md() {
let mut files = HashMap::new();
files.insert("/repo/AGENTS.md".into(), "Agent instructions".into());
let env: Arc<dyn ExecutionEnvironment> = Arc::new(DocEnv { files });
let docs = discover_project_docs(env.as_ref(), "/repo", "/repo", "anthropic").await;
assert_eq!(docs.len(), 1);
assert_eq!(docs[0], "Agent instructions");
}
#[tokio::test]
async fn filters_by_provider() {
let mut files = HashMap::new();
files.insert("/repo/AGENTS.md".into(), "agents".into());
files.insert("/repo/CLAUDE.md".into(), "claude".into());
files.insert(
"/repo/.github/copilot-instructions.md".into(),
"copilot".into(),
);
files.insert("/repo/GEMINI.md".into(), "gemini".into());
let env: Arc<dyn ExecutionEnvironment> = Arc::new(DocEnv {
files: files.clone(),
});
let anthropic_docs =
discover_project_docs(env.as_ref(), "/repo", "/repo", "anthropic").await;
assert_eq!(anthropic_docs.len(), 2);
assert_eq!(anthropic_docs[0], "agents");
assert_eq!(anthropic_docs[1], "claude");
let env: Arc<dyn ExecutionEnvironment> = Arc::new(DocEnv {
files: files.clone(),
});
let openai_docs = discover_project_docs(env.as_ref(), "/repo", "/repo", "openai").await;
assert_eq!(openai_docs.len(), 2);
assert_eq!(openai_docs[0], "agents");
assert_eq!(openai_docs[1], "copilot");
let env: Arc<dyn ExecutionEnvironment> = Arc::new(DocEnv { files });
let gemini_docs = discover_project_docs(env.as_ref(), "/repo", "/repo", "gemini").await;
assert_eq!(gemini_docs.len(), 2);
assert_eq!(gemini_docs[0], "agents");
assert_eq!(gemini_docs[1], "gemini");
}
#[tokio::test]
async fn truncates_at_budget() {
let mut files = HashMap::new();
// Create content that exceeds 32KB budget
let large_content = "x".repeat(30000);
let second_content = "y".repeat(5000);
files.insert("/repo/AGENTS.md".into(), large_content.clone());
files.insert("/repo/CLAUDE.md".into(), second_content);
let env: Arc<dyn ExecutionEnvironment> = Arc::new(DocEnv { files });
let docs = discover_project_docs(env.as_ref(), "/repo", "/repo", "anthropic").await;
assert_eq!(docs.len(), 2);
assert_eq!(docs[0], large_content);
// Second doc should be truncated to fit remaining budget
assert!(docs[1].ends_with("... [truncated]"));
assert!(docs[0].len() + docs[1].len() <= BUDGET_BYTES);
}
#[tokio::test]
async fn walks_directory_hierarchy() {
let mut files = HashMap::new();
files.insert("/repo/AGENTS.md".into(), "root agents".into());
files.insert("/repo/src/AGENTS.md".into(), "src agents".into());
files.insert("/repo/src/app/AGENTS.md".into(), "app agents".into());
let env: Arc<dyn ExecutionEnvironment> = Arc::new(DocEnv { files });
let docs =
discover_project_docs(env.as_ref(), "/repo", "/repo/src/app", "anthropic").await;
assert_eq!(docs.len(), 3);
assert_eq!(docs[0], "root agents");
assert_eq!(docs[1], "src agents");
assert_eq!(docs[2], "app agents");
}
}

View file

@ -0,0 +1,176 @@
use crate::execution_env::ExecutionEnvironment;
use crate::tool_registry::ToolRegistry;
use unified_llm::types::ToolDefinition;
pub trait ProviderProfile: Send + Sync {
fn id(&self) -> String;
fn model(&self) -> String;
fn tool_registry(&self) -> &ToolRegistry;
fn build_system_prompt(
&self,
env: &dyn ExecutionEnvironment,
project_docs: &[String],
) -> String;
fn tools(&self) -> Vec<ToolDefinition>;
fn provider_options(&self) -> Option<serde_json::Value>;
fn supports_reasoning(&self) -> bool;
fn supports_streaming(&self) -> bool;
fn supports_parallel_tool_calls(&self) -> bool;
fn context_window_size(&self) -> usize;
}
#[cfg(test)]
mod tests {
use super::*;
use crate::execution_env::*;
use async_trait::async_trait;
struct TestEnv;
#[async_trait]
impl ExecutionEnvironment for TestEnv {
async fn read_file(&self, _: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(
&self,
_: &str,
_: &[String],
_: u64,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(
&self,
_: &str,
_: &str,
_: &GrepOptions,
) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/home/test"
}
fn platform(&self) -> &str {
"linux"
}
fn os_version(&self) -> String {
"Linux 6.1.0".into()
}
}
struct TestProfile {
registry: ToolRegistry,
}
impl TestProfile {
fn new() -> Self {
Self {
registry: ToolRegistry::new(),
}
}
}
impl ProviderProfile for TestProfile {
fn id(&self) -> String {
"test-provider".into()
}
fn model(&self) -> String {
"test-model".into()
}
fn tool_registry(&self) -> &ToolRegistry {
&self.registry
}
fn build_system_prompt(
&self,
env: &dyn ExecutionEnvironment,
project_docs: &[String],
) -> String {
format!(
"You are working on {}. Docs: {}",
env.platform(),
project_docs.len()
)
}
fn tools(&self) -> Vec<ToolDefinition> {
self.registry.definitions()
}
fn provider_options(&self) -> Option<serde_json::Value> {
None
}
fn supports_reasoning(&self) -> bool {
true
}
fn supports_streaming(&self) -> bool {
true
}
fn supports_parallel_tool_calls(&self) -> bool {
false
}
fn context_window_size(&self) -> usize {
200_000
}
}
#[test]
fn profile_id_and_model() {
let profile = TestProfile::new();
assert_eq!(profile.id(), "test-provider");
assert_eq!(profile.model(), "test-model");
}
#[test]
fn profile_capabilities() {
let profile = TestProfile::new();
assert!(profile.supports_reasoning());
assert!(profile.supports_streaming());
assert!(!profile.supports_parallel_tool_calls());
assert_eq!(profile.context_window_size(), 200_000);
}
#[test]
fn profile_build_system_prompt() {
let profile = TestProfile::new();
let env = TestEnv;
let docs = vec!["README.md contents".into()];
let prompt = profile.build_system_prompt(&env, &docs);
assert!(prompt.contains("linux"));
assert!(prompt.contains("1"));
}
#[test]
fn profile_provider_options_none() {
let profile = TestProfile::new();
assert!(profile.provider_options().is_none());
}
#[test]
fn profile_tools_empty_registry() {
let profile = TestProfile::new();
assert!(profile.tools().is_empty());
}
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,599 @@
use crate::error::AgentError;
use crate::session::Session;
use crate::tool_registry::RegisteredTool;
use crate::types::Turn;
use std::collections::{HashMap, VecDeque};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use unified_llm::types::ToolDefinition;
pub type SessionFactory = Arc<dyn Fn() -> Session + Send + Sync>;
pub struct SubAgent {
id: String,
depth: usize,
task: Option<tokio::task::JoinHandle<Result<String, AgentError>>>,
followup_queue: Arc<Mutex<VecDeque<String>>>,
abort_flag: Arc<AtomicBool>,
}
impl SubAgent {
pub fn id(&self) -> &str {
&self.id
}
pub fn depth(&self) -> usize {
self.depth
}
}
pub struct SubAgentManager {
agents: HashMap<String, SubAgent>,
max_depth: usize,
}
impl SubAgentManager {
pub fn new(max_depth: usize) -> Self {
Self {
agents: HashMap::new(),
max_depth,
}
}
pub fn spawn(
&mut self,
mut session: Session,
task_prompt: String,
depth: usize,
) -> Result<String, String> {
if depth >= self.max_depth {
return Err(format!(
"Maximum subagent depth ({}) reached",
self.max_depth
));
}
let agent_id = uuid::Uuid::new_v4().to_string();
let followup_queue = session.followup_queue_handle();
let abort_flag = session.abort_flag_handle();
let task = tokio::spawn(async move {
session.process_input(&task_prompt).await?;
let turns = session.history().turns();
let last_text = turns.iter().rev().find_map(|t| {
if let Turn::Assistant { content, .. } = t {
Some(content.clone())
} else {
None
}
});
Ok(last_text.unwrap_or_default())
});
self.agents.insert(
agent_id.clone(),
SubAgent {
id: agent_id.clone(),
depth,
task: Some(task),
followup_queue,
abort_flag,
},
);
Ok(agent_id)
}
pub fn send_input(&self, agent_id: &str, message: &str) -> Result<(), String> {
let agent = self
.agents
.get(agent_id)
.ok_or_else(|| format!("No agent found with id: {agent_id}"))?;
agent
.followup_queue
.lock()
.expect("followup queue lock poisoned")
.push_back(message.to_string());
Ok(())
}
pub async fn wait(&mut self, agent_id: &str) -> Result<String, String> {
let mut agent = self
.agents
.remove(agent_id)
.ok_or_else(|| format!("No agent found with id: {agent_id}"))?;
match agent.task.take() {
Some(join_handle) => match join_handle.await {
Ok(result) => result.map_err(|e| e.to_string()),
Err(e) => Err(format!("Agent task panicked: {e}")),
},
None => Err(format!("Agent {agent_id} has no running task")),
}
}
pub fn close(&mut self, agent_id: &str) -> Result<(), String> {
let agent = self
.agents
.remove(agent_id)
.ok_or_else(|| format!("No agent found with id: {agent_id}"))?;
agent.abort_flag.store(true, Ordering::SeqCst);
if let Some(join_handle) = agent.task {
join_handle.abort();
}
Ok(())
}
pub fn get(&self, agent_id: &str) -> Option<&SubAgent> {
self.agents.get(agent_id)
}
}
pub fn make_spawn_agent_tool(
manager: Arc<tokio::sync::Mutex<SubAgentManager>>,
session_factory: SessionFactory,
current_depth: usize,
) -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: "spawn_agent".into(),
description: "Spawn a subagent to work on a delegated task".into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"task": {
"type": "string",
"description": "The task description for the subagent"
}
},
"required": ["task"]
}),
},
executor: Arc::new(move |args, _env| {
let manager = manager.clone();
let session_factory = session_factory.clone();
Box::pin(async move {
let task = args
.get("task")
.and_then(|v| v.as_str())
.ok_or_else(|| "Missing required parameter: task".to_string())?;
let session = session_factory();
let mut mgr = manager.lock().await;
mgr.spawn(session, task.to_string(), current_depth)
})
}),
}
}
pub fn make_send_input_tool(
manager: Arc<tokio::sync::Mutex<SubAgentManager>>,
) -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: "send_input".into(),
description: "Send a follow-up message to a running subagent".into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"agent_id": {
"type": "string",
"description": "The ID of the agent to send input to"
},
"message": {
"type": "string",
"description": "The message to send to the agent"
}
},
"required": ["agent_id", "message"]
}),
},
executor: Arc::new(move |args, _env| {
let manager = manager.clone();
Box::pin(async move {
let agent_id = args
.get("agent_id")
.and_then(|v| v.as_str())
.ok_or_else(|| "Missing required parameter: agent_id".to_string())?;
let message = args
.get("message")
.and_then(|v| v.as_str())
.ok_or_else(|| "Missing required parameter: message".to_string())?;
let mgr = manager.lock().await;
mgr.send_input(agent_id, message)?;
Ok(format!("Message sent to agent {agent_id}"))
})
}),
}
}
pub fn make_wait_tool(
manager: Arc<tokio::sync::Mutex<SubAgentManager>>,
) -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: "wait".into(),
description: "Wait for a subagent to complete and return its result".into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"agent_id": {
"type": "string",
"description": "The ID of the agent to wait for"
}
},
"required": ["agent_id"]
}),
},
executor: Arc::new(move |args, _env| {
let manager = manager.clone();
Box::pin(async move {
let agent_id = args
.get("agent_id")
.and_then(|v| v.as_str())
.ok_or_else(|| "Missing required parameter: agent_id".to_string())?;
let mut mgr = manager.lock().await;
mgr.wait(agent_id).await
})
}),
}
}
pub fn make_close_agent_tool(
manager: Arc<tokio::sync::Mutex<SubAgentManager>>,
) -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: "close_agent".into(),
description: "Close a running subagent".into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"agent_id": {
"type": "string",
"description": "The ID of the agent to close"
}
},
"required": ["agent_id"]
}),
},
executor: Arc::new(move |args, _env| {
let manager = manager.clone();
Box::pin(async move {
let agent_id = args
.get("agent_id")
.and_then(|v| v.as_str())
.ok_or_else(|| "Missing required parameter: agent_id".to_string())?;
let mut mgr = manager.lock().await;
mgr.close(agent_id)?;
Ok(format!("Agent {agent_id} closed"))
})
}),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::SessionConfig;
use crate::execution_env::*;
use crate::provider_profile::ProviderProfile;
use crate::tool_registry::ToolRegistry;
use async_trait::async_trait;
use std::sync::atomic::AtomicUsize;
use unified_llm::client::Client;
use unified_llm::error::SdkError;
use unified_llm::provider::{ProviderAdapter, StreamEventStream};
use unified_llm::types::{FinishReason, Message, Response, Usage};
// --- Mock LLM Provider ---
struct MockLlmProvider {
responses: Vec<Response>,
call_index: AtomicUsize,
}
impl MockLlmProvider {
fn new(responses: Vec<Response>) -> Self {
Self {
responses,
call_index: AtomicUsize::new(0),
}
}
}
#[async_trait]
impl ProviderAdapter for MockLlmProvider {
fn name(&self) -> &str {
"mock"
}
async fn complete(
&self,
_request: &unified_llm::types::Request,
) -> Result<Response, SdkError> {
let idx = self.call_index.fetch_add(1, Ordering::SeqCst);
if idx < self.responses.len() {
Ok(self.responses[idx].clone())
} else {
Ok(self.responses[self.responses.len() - 1].clone())
}
}
async fn stream(
&self,
_request: &unified_llm::types::Request,
) -> Result<StreamEventStream, SdkError> {
Err(SdkError::Configuration {
message: "streaming not supported in mock".into(),
})
}
}
// --- Memory Execution Environment ---
struct MemoryExecutionEnvironment;
#[async_trait]
impl ExecutionEnvironment for MemoryExecutionEnvironment {
async fn read_file(&self, _path: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, _path: &str, _content: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _path: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _path: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(
&self,
_command: &str,
_args: &[String],
_timeout_ms: u64,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: "mock output".into(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 10,
})
}
async fn grep(
&self,
_pattern: &str,
_path: &str,
_options: &GrepOptions,
) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _pattern: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp/test"
}
fn platform(&self) -> &str {
"darwin"
}
fn os_version(&self) -> String {
"Darwin 24.0.0".into()
}
}
// --- Test Profile ---
struct TestProfile {
registry: ToolRegistry,
}
impl TestProfile {
fn new() -> Self {
Self {
registry: ToolRegistry::new(),
}
}
}
impl ProviderProfile for TestProfile {
fn id(&self) -> String {
"mock".into()
}
fn model(&self) -> String {
"mock-model".into()
}
fn tool_registry(&self) -> &ToolRegistry {
&self.registry
}
fn build_system_prompt(
&self,
_env: &dyn ExecutionEnvironment,
_project_docs: &[String],
) -> String {
"You are a test assistant.".into()
}
fn tools(&self) -> Vec<ToolDefinition> {
self.registry.definitions()
}
fn provider_options(&self) -> Option<serde_json::Value> {
None
}
fn supports_reasoning(&self) -> bool {
false
}
fn supports_streaming(&self) -> bool {
false
}
fn supports_parallel_tool_calls(&self) -> bool {
false
}
fn context_window_size(&self) -> usize {
200_000
}
}
// --- Helper functions ---
fn text_response(text: &str) -> Response {
Response {
id: format!("resp_{text}"),
model: "mock-model".into(),
provider: "mock".into(),
message: Message::assistant(text),
finish_reason: FinishReason::Stop,
usage: Usage {
input_tokens: 10,
output_tokens: 5,
total_tokens: 15,
..Default::default()
},
raw: None,
warnings: vec![],
rate_limit: None,
}
}
async fn make_client(provider: Arc<dyn ProviderAdapter>) -> Client {
let mut providers = HashMap::new();
providers.insert(provider.name().to_string(), provider);
Client::new(providers, Some("mock".into()), vec![])
}
async fn make_session(responses: Vec<Response>) -> Session {
let provider = Arc::new(MockLlmProvider::new(responses));
let client = make_client(provider).await;
let profile = Arc::new(TestProfile::new());
let env = Arc::new(MemoryExecutionEnvironment);
Session::new(client, profile, env, SessionConfig::default())
}
// --- Tests ---
#[test]
fn manager_creation() {
let manager = SubAgentManager::new(3);
assert_eq!(manager.max_depth, 3);
assert!(manager.agents.is_empty());
}
#[tokio::test]
async fn spawn_creates_agent_and_returns_id() {
let mut manager = SubAgentManager::new(3);
let session = make_session(vec![text_response("Hello")]).await;
let result = manager.spawn(session, "Do something".into(), 0);
assert!(result.is_ok());
let agent_id = result.unwrap();
assert!(!agent_id.is_empty());
assert!(manager.get(&agent_id).is_some());
assert_eq!(manager.get(&agent_id).unwrap().depth(), 0);
}
#[tokio::test]
async fn depth_limit_enforced() {
let mut manager = SubAgentManager::new(2);
let session = make_session(vec![text_response("Hello")]).await;
let result = manager.spawn(session, "Do something".into(), 2);
assert!(result.is_err());
assert!(result.unwrap_err().contains("Maximum subagent depth"));
}
#[tokio::test]
async fn close_removes_agent() {
let mut manager = SubAgentManager::new(3);
let session = make_session(vec![text_response("Hello")]).await;
let agent_id = manager.spawn(session, "Do something".into(), 0).unwrap();
assert!(manager.get(&agent_id).is_some());
let result = manager.close(&agent_id);
assert!(result.is_ok());
assert!(manager.get(&agent_id).is_none());
}
#[tokio::test]
async fn send_input_nonexistent_agent_errors() {
let manager = SubAgentManager::new(3);
let result = manager.send_input("nonexistent-id", "hello");
assert!(result.is_err());
assert!(result.unwrap_err().contains("No agent found"));
}
#[tokio::test]
async fn wait_nonexistent_agent_errors() {
let mut manager = SubAgentManager::new(3);
let result = manager.wait("nonexistent-id").await;
assert!(result.is_err());
assert!(result.unwrap_err().contains("No agent found"));
}
#[tokio::test]
async fn wait_returns_result() {
let mut manager = SubAgentManager::new(3);
let session =
make_session(vec![text_response("Task completed successfully")]).await;
let agent_id = manager.spawn(session, "Do something".into(), 0).unwrap();
let result = manager.wait(&agent_id).await;
assert!(result.is_ok());
assert_eq!(result.unwrap(), "Task completed successfully");
assert!(manager.get(&agent_id).is_none());
}
#[test]
fn tool_definitions_correct() {
let manager = Arc::new(tokio::sync::Mutex::new(SubAgentManager::new(3)));
let factory: SessionFactory = Arc::new(|| {
panic!("should not be called");
});
let spawn_tool = make_spawn_agent_tool(manager.clone(), factory, 0);
assert_eq!(spawn_tool.definition.name, "spawn_agent");
assert!(spawn_tool.definition.parameters["properties"]["task"].is_object());
let spawn_required = spawn_tool.definition.parameters["required"]
.as_array()
.unwrap();
assert!(spawn_required.contains(&serde_json::json!("task")));
let send_tool = make_send_input_tool(manager.clone());
assert_eq!(send_tool.definition.name, "send_input");
assert!(send_tool.definition.parameters["properties"]["agent_id"].is_object());
assert!(send_tool.definition.parameters["properties"]["message"].is_object());
let send_required = send_tool.definition.parameters["required"]
.as_array()
.unwrap();
assert!(send_required.contains(&serde_json::json!("agent_id")));
assert!(send_required.contains(&serde_json::json!("message")));
let wait_tool = make_wait_tool(manager.clone());
assert_eq!(wait_tool.definition.name, "wait");
assert!(wait_tool.definition.parameters["properties"]["agent_id"].is_object());
let wait_required = wait_tool.definition.parameters["required"]
.as_array()
.unwrap();
assert!(wait_required.contains(&serde_json::json!("agent_id")));
let close_tool = make_close_agent_tool(manager);
assert_eq!(close_tool.definition.name, "close_agent");
assert!(close_tool.definition.parameters["properties"]["agent_id"].is_object());
let close_required = close_tool.definition.parameters["required"]
.as_array()
.unwrap();
assert!(close_required.contains(&serde_json::json!("agent_id")));
}
}

View file

@ -0,0 +1,238 @@
use crate::execution_env::ExecutionEnvironment;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use unified_llm::types::ToolDefinition;
pub type ToolExecutor = Arc<
dyn Fn(
serde_json::Value,
Arc<dyn ExecutionEnvironment>,
) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>>
+ Send
+ Sync,
>;
pub struct RegisteredTool {
pub definition: ToolDefinition,
pub executor: ToolExecutor,
}
pub struct ToolRegistry {
tools: HashMap<String, RegisteredTool>,
}
impl ToolRegistry {
#[must_use]
pub fn new() -> Self {
Self {
tools: HashMap::new(),
}
}
pub fn register(&mut self, tool: RegisteredTool) {
self.tools.insert(tool.definition.name.clone(), tool);
}
pub fn unregister(&mut self, name: &str) -> Option<RegisteredTool> {
self.tools.remove(name)
}
#[must_use]
pub fn get(&self, name: &str) -> Option<&RegisteredTool> {
self.tools.get(name)
}
#[must_use]
pub fn definitions(&self) -> Vec<ToolDefinition> {
self.tools.values().map(|t| t.definition.clone()).collect()
}
#[must_use]
pub fn names(&self) -> Vec<String> {
self.tools.keys().cloned().collect()
}
}
impl Default for ToolRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_tool(name: &str) -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: name.into(),
description: format!("Tool {name}"),
parameters: serde_json::json!({"type": "object"}),
},
executor: Arc::new(|_args, _env| Box::pin(async { Ok("ok".into()) })),
}
}
#[test]
fn register_and_get() {
let mut registry = ToolRegistry::new();
registry.register(make_tool("read_file"));
let tool = registry.get("read_file");
assert!(tool.is_some());
assert_eq!(tool.unwrap().definition.name, "read_file");
}
#[test]
fn get_missing_returns_none() {
let registry = ToolRegistry::new();
assert!(registry.get("nonexistent").is_none());
}
#[test]
fn unregister_removes_tool() {
let mut registry = ToolRegistry::new();
registry.register(make_tool("read_file"));
let removed = registry.unregister("read_file");
assert!(removed.is_some());
assert!(registry.get("read_file").is_none());
}
#[test]
fn unregister_missing_returns_none() {
let mut registry = ToolRegistry::new();
assert!(registry.unregister("nonexistent").is_none());
}
#[test]
fn name_collision_overrides() {
let mut registry = ToolRegistry::new();
registry.register(RegisteredTool {
definition: ToolDefinition {
name: "tool_a".into(),
description: "version 1".into(),
parameters: serde_json::json!({}),
},
executor: Arc::new(|_args, _env| Box::pin(async { Ok("v1".into()) })),
});
registry.register(RegisteredTool {
definition: ToolDefinition {
name: "tool_a".into(),
description: "version 2".into(),
parameters: serde_json::json!({}),
},
executor: Arc::new(|_args, _env| Box::pin(async { Ok("v2".into()) })),
});
let tool = registry.get("tool_a").unwrap();
assert_eq!(tool.definition.description, "version 2");
}
#[test]
fn definitions_returns_all() {
let mut registry = ToolRegistry::new();
registry.register(make_tool("tool_a"));
registry.register(make_tool("tool_b"));
let defs = registry.definitions();
assert_eq!(defs.len(), 2);
let names: Vec<&str> = defs.iter().map(|d| d.name.as_str()).collect();
assert!(names.contains(&"tool_a"));
assert!(names.contains(&"tool_b"));
}
#[test]
fn names_returns_all() {
let mut registry = ToolRegistry::new();
registry.register(make_tool("tool_x"));
registry.register(make_tool("tool_y"));
let names = registry.names();
assert_eq!(names.len(), 2);
assert!(names.contains(&"tool_x".to_string()));
assert!(names.contains(&"tool_y".to_string()));
}
#[tokio::test]
async fn executor_can_be_called() {
let mut registry = ToolRegistry::new();
registry.register(make_tool("echo"));
let tool = registry.get("echo").unwrap();
use crate::execution_env::*;
use async_trait::async_trait;
struct DummyEnv;
#[async_trait]
impl ExecutionEnvironment for DummyEnv {
async fn read_file(&self, _: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(
&self,
_: &str,
_: &[String],
_: u64,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(
&self,
_: &str,
_: &str,
_: &GrepOptions,
) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp"
}
fn platform(&self) -> &str {
"darwin"
}
fn os_version(&self) -> String {
String::new()
}
}
let env: Arc<dyn ExecutionEnvironment> = Arc::new(DummyEnv);
let result = (tool.executor)(serde_json::json!({}), env).await;
assert_eq!(result.unwrap(), "ok");
}
#[test]
fn default_creates_empty_registry() {
let registry = ToolRegistry::default();
assert!(registry.names().is_empty());
assert!(registry.definitions().is_empty());
}
}

View file

@ -0,0 +1,856 @@
use crate::execution_env::GrepOptions;
use crate::tool_registry::RegisteredTool;
use std::fmt::Write;
use std::sync::Arc;
use unified_llm::types::ToolDefinition;
#[must_use]
pub fn make_read_file_tool() -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: "read_file".into(),
description: "Read the contents of a file".into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"file_path": {"type": "string", "description": "Absolute path to the file"},
"offset": {"type": "integer", "description": "1-based line number to start reading from"},
"limit": {"type": "integer", "description": "Number of lines to read"}
},
"required": ["file_path"]
}),
},
executor: Arc::new(|args, env| {
Box::pin(async move {
let file_path = args["file_path"]
.as_str()
.ok_or_else(|| "file_path is required".to_string())?;
let offset = args.get("offset").and_then(serde_json::Value::as_u64);
let limit = args.get("limit").and_then(serde_json::Value::as_u64);
let content = env.read_file(file_path).await?;
if offset.is_none() && limit.is_none() {
return Ok(content);
}
#[allow(clippy::cast_possible_truncation)]
let offset = offset.unwrap_or(1) as usize;
let lines: Vec<&str> = content.lines().collect();
let start = if offset > 0 { offset - 1 } else { 0 };
#[allow(clippy::cast_possible_truncation)]
let selected: Vec<&str> = match limit {
Some(lim) => lines.into_iter().skip(start).take(lim as usize).collect(),
None => lines.into_iter().skip(start).collect(),
};
Ok(selected.join("\n"))
})
}),
}
}
#[must_use]
pub fn make_write_file_tool() -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: "write_file".into(),
description: "Write content to a file".into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"file_path": {"type": "string", "description": "Absolute path to the file"},
"content": {"type": "string", "description": "Content to write to the file"}
},
"required": ["file_path", "content"]
}),
},
executor: Arc::new(|args, env| {
Box::pin(async move {
let file_path = args["file_path"]
.as_str()
.ok_or_else(|| "file_path is required".to_string())?;
let content = args["content"]
.as_str()
.ok_or_else(|| "content is required".to_string())?;
env.write_file(file_path, content).await?;
Ok(format!("Successfully wrote to {file_path}"))
})
}),
}
}
#[must_use]
pub fn make_edit_file_tool() -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: "edit_file".into(),
description: "Edit a file by replacing a string".into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"file_path": {"type": "string", "description": "Absolute path to the file"},
"old_string": {"type": "string", "description": "The string to find and replace"},
"new_string": {"type": "string", "description": "The replacement string"},
"replace_all": {"type": "boolean", "description": "Replace all occurrences (default false)"}
},
"required": ["file_path", "old_string", "new_string"]
}),
},
executor: Arc::new(|args, env| {
Box::pin(async move {
let file_path = args["file_path"]
.as_str()
.ok_or_else(|| "file_path is required".to_string())?;
let old_string = args["old_string"]
.as_str()
.ok_or_else(|| "old_string is required".to_string())?;
let new_string = args["new_string"]
.as_str()
.ok_or_else(|| "new_string is required".to_string())?;
let replace_all = args
.get("replace_all")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let numbered_content = env.read_file(file_path).await?;
// Strip line numbers: each line looks like " 1 | content" or " 10 | content"
let raw_lines: Vec<&str> = numbered_content
.lines()
.map(|line| {
line.find(" | ")
.map_or(line, |idx| &line[idx + 3..])
})
.collect();
let raw_content = raw_lines.join("\n");
let count = raw_content.matches(old_string).count();
if count == 0 {
return Err("old_string not found in file".to_string());
}
if count > 1 && !replace_all {
return Err(format!(
"old_string is not unique in file (found {count} occurrences). Use replace_all or provide more context"
));
}
let new_content = if replace_all {
raw_content.replace(old_string, new_string)
} else {
raw_content.replacen(old_string, new_string, 1)
};
env.write_file(file_path, &new_content).await?;
Ok(format!("Successfully edited {file_path}"))
})
}),
}
}
#[must_use]
pub fn make_shell_tool() -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: "shell".into(),
description: "Execute a shell command".into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"command": {"type": "string", "description": "The shell command to execute"},
"timeout_ms": {"type": "integer", "description": "Timeout in milliseconds (default 10000)"}
},
"required": ["command"]
}),
},
executor: Arc::new(|args, env| {
Box::pin(async move {
let command = args["command"]
.as_str()
.ok_or_else(|| "command is required".to_string())?;
let timeout_ms = args
.get("timeout_ms")
.and_then(serde_json::Value::as_u64)
.unwrap_or(10000);
let result = env
.exec_command(
"/bin/bash",
&["-c".into(), command.into()],
timeout_ms,
)
.await?;
let mut output = String::new();
if result.timed_out {
output.push_str("Command timed out.\n");
}
let _ = write!(
output,
"Exit code: {}\nstdout:\n{}\nstderr:\n{}",
result.exit_code, result.stdout, result.stderr
);
Ok(output)
})
}),
}
}
#[must_use]
pub fn make_grep_tool() -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: "grep".into(),
description: "Search file contents with a regex pattern".into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"pattern": {"type": "string", "description": "Regex pattern to search for"},
"path": {"type": "string", "description": "Path to search in (default \".\")"},
"glob_filter": {"type": "string", "description": "Glob pattern to filter files"},
"case_insensitive": {"type": "boolean", "description": "Case insensitive search"},
"max_results": {"type": "integer", "description": "Maximum number of results"}
},
"required": ["pattern"]
}),
},
executor: Arc::new(|args, env| {
Box::pin(async move {
let pattern = args["pattern"]
.as_str()
.ok_or_else(|| "pattern is required".to_string())?;
let path = args
.get("path")
.and_then(serde_json::Value::as_str)
.unwrap_or(".");
#[allow(clippy::cast_possible_truncation)]
let options = GrepOptions {
glob_filter: args
.get("glob_filter")
.and_then(serde_json::Value::as_str)
.map(String::from),
case_insensitive: args
.get("case_insensitive")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false),
max_results: args
.get("max_results")
.and_then(serde_json::Value::as_u64)
.map(|v| v as usize),
};
let results = env.grep(pattern, path, &options).await?;
Ok(results.join("\n"))
})
}),
}
}
#[must_use]
pub fn make_glob_tool() -> RegisteredTool {
RegisteredTool {
definition: ToolDefinition {
name: "glob".into(),
description: "Find files matching a glob pattern".into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"pattern": {"type": "string", "description": "Glob pattern to match files"}
},
"required": ["pattern"]
}),
},
executor: Arc::new(|args, env| {
Box::pin(async move {
let pattern = args["pattern"]
.as_str()
.ok_or_else(|| "pattern is required".to_string())?;
let results = env.glob(pattern).await?;
Ok(results.join("\n"))
})
}),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::execution_env::*;
use async_trait::async_trait;
use std::sync::Mutex;
struct ReadFileEnv {
content: String,
}
#[async_trait]
impl ExecutionEnvironment for ReadFileEnv {
async fn read_file(&self, _path: &str) -> Result<String, String> {
Ok(self.content.clone())
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(&self, _: &str, _: &str, _: &GrepOptions) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp"
}
fn platform(&self) -> &str {
"darwin"
}
fn os_version(&self) -> String {
String::new()
}
}
struct WriteFileEnv {
written: Mutex<Option<(String, String)>>,
}
#[async_trait]
impl ExecutionEnvironment for WriteFileEnv {
async fn read_file(&self, _path: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, path: &str, content: &str) -> Result<(), String> {
*self.written.lock().unwrap() = Some((path.into(), content.into()));
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(&self, _: &str, _: &str, _: &GrepOptions) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp"
}
fn platform(&self) -> &str {
"darwin"
}
fn os_version(&self) -> String {
String::new()
}
}
struct EditFileEnv {
content: String,
written: Mutex<Option<String>>,
}
#[async_trait]
impl ExecutionEnvironment for EditFileEnv {
async fn read_file(&self, _path: &str) -> Result<String, String> {
Ok(self.content.clone())
}
async fn write_file(&self, _path: &str, content: &str) -> Result<(), String> {
*self.written.lock().unwrap() = Some(content.into());
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(&self, _: &str, _: &str, _: &GrepOptions) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp"
}
fn platform(&self) -> &str {
"darwin"
}
fn os_version(&self) -> String {
String::new()
}
}
struct ShellEnv {
result: ExecResult,
}
#[async_trait]
impl ExecutionEnvironment for ShellEnv {
async fn read_file(&self, _: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
Ok(self.result.clone())
}
async fn grep(&self, _: &str, _: &str, _: &GrepOptions) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp"
}
fn platform(&self) -> &str {
"darwin"
}
fn os_version(&self) -> String {
String::new()
}
}
struct ShellCapturingEnv {
captured_timeout: Mutex<Option<u64>>,
}
#[async_trait]
impl ExecutionEnvironment for ShellCapturingEnv {
async fn read_file(&self, _: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(
&self,
_: &str,
_: &[String],
timeout_ms: u64,
) -> Result<ExecResult, String> {
*self.captured_timeout.lock().unwrap() = Some(timeout_ms);
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(&self, _: &str, _: &str, _: &GrepOptions) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp"
}
fn platform(&self) -> &str {
"darwin"
}
fn os_version(&self) -> String {
String::new()
}
}
struct GrepEnv {
results: Vec<String>,
}
#[async_trait]
impl ExecutionEnvironment for GrepEnv {
async fn read_file(&self, _: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(&self, _: &str, _: &str, _: &GrepOptions) -> Result<Vec<String>, String> {
Ok(self.results.clone())
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp"
}
fn platform(&self) -> &str {
"darwin"
}
fn os_version(&self) -> String {
String::new()
}
}
struct GlobEnv {
results: Vec<String>,
}
#[async_trait]
impl ExecutionEnvironment for GlobEnv {
async fn read_file(&self, _: &str) -> Result<String, String> {
Ok(String::new())
}
async fn write_file(&self, _: &str, _: &str) -> Result<(), String> {
Ok(())
}
async fn file_exists(&self, _: &str) -> Result<bool, String> {
Ok(false)
}
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 0,
})
}
async fn grep(&self, _: &str, _: &str, _: &GrepOptions) -> Result<Vec<String>, String> {
Ok(vec![])
}
async fn glob(&self, _: &str) -> Result<Vec<String>, String> {
Ok(self.results.clone())
}
async fn initialize(&self) -> Result<(), String> {
Ok(())
}
async fn cleanup(&self) -> Result<(), String> {
Ok(())
}
fn working_directory(&self) -> &str {
"/tmp"
}
fn platform(&self) -> &str {
"darwin"
}
fn os_version(&self) -> String {
String::new()
}
}
#[tokio::test]
async fn read_file_returns_content() {
let tool = make_read_file_tool();
let env: Arc<dyn ExecutionEnvironment> = Arc::new(ReadFileEnv {
content: " 1 | hello\n 2 | world".into(),
});
let result = (tool.executor)(serde_json::json!({"file_path": "/test.txt"}), env).await;
assert_eq!(result.unwrap(), " 1 | hello\n 2 | world");
}
#[tokio::test]
async fn read_file_with_offset_and_limit() {
let tool = make_read_file_tool();
let env: Arc<dyn ExecutionEnvironment> = Arc::new(ReadFileEnv {
content: " 1 | line1\n 2 | line2\n 3 | line3\n 4 | line4".into(),
});
let result = (tool.executor)(
serde_json::json!({"file_path": "/test.txt", "offset": 2, "limit": 2}),
env,
)
.await;
assert_eq!(result.unwrap(), " 2 | line2\n 3 | line3");
}
#[tokio::test]
async fn write_file_calls_env() {
let tool = make_write_file_tool();
let env = Arc::new(WriteFileEnv {
written: Mutex::new(None),
});
let env_clone: Arc<dyn ExecutionEnvironment> = env.clone();
let result = (tool.executor)(
serde_json::json!({"file_path": "/out.txt", "content": "hello"}),
env_clone,
)
.await;
assert_eq!(result.unwrap(), "Successfully wrote to /out.txt");
let written = env.written.lock().unwrap();
let (path, content) = written.as_ref().unwrap();
assert_eq!(path, "/out.txt");
assert_eq!(content, "hello");
}
#[tokio::test]
async fn edit_file_replaces_match() {
let tool = make_edit_file_tool();
let env = Arc::new(EditFileEnv {
content: " 1 | hello world".into(),
written: Mutex::new(None),
});
let env_clone: Arc<dyn ExecutionEnvironment> = env.clone();
let result = (tool.executor)(
serde_json::json!({
"file_path": "/f.txt",
"old_string": "hello",
"new_string": "goodbye"
}),
env_clone,
)
.await;
assert_eq!(result.unwrap(), "Successfully edited /f.txt");
let written = env.written.lock().unwrap();
assert_eq!(written.as_ref().unwrap(), "goodbye world");
}
#[tokio::test]
async fn edit_file_not_found_error() {
let tool = make_edit_file_tool();
let env: Arc<dyn ExecutionEnvironment> = Arc::new(EditFileEnv {
content: " 1 | hello world".into(),
written: Mutex::new(None),
});
let result = (tool.executor)(
serde_json::json!({
"file_path": "/f.txt",
"old_string": "missing",
"new_string": "replacement"
}),
env,
)
.await;
assert_eq!(result.unwrap_err(), "old_string not found in file");
}
#[tokio::test]
async fn edit_file_not_unique_error() {
let tool = make_edit_file_tool();
let env: Arc<dyn ExecutionEnvironment> = Arc::new(EditFileEnv {
content: " 1 | aa bb aa".into(),
written: Mutex::new(None),
});
let result = (tool.executor)(
serde_json::json!({
"file_path": "/f.txt",
"old_string": "aa",
"new_string": "cc"
}),
env,
)
.await;
let err = result.unwrap_err();
assert!(err.contains("not unique"));
assert!(err.contains("2 occurrences"));
}
#[tokio::test]
async fn edit_file_replace_all() {
let tool = make_edit_file_tool();
let env = Arc::new(EditFileEnv {
content: " 1 | aa bb aa".into(),
written: Mutex::new(None),
});
let env_clone: Arc<dyn ExecutionEnvironment> = env.clone();
let result = (tool.executor)(
serde_json::json!({
"file_path": "/f.txt",
"old_string": "aa",
"new_string": "cc",
"replace_all": true
}),
env_clone,
)
.await;
assert_eq!(result.unwrap(), "Successfully edited /f.txt");
let written = env.written.lock().unwrap();
assert_eq!(written.as_ref().unwrap(), "cc bb cc");
}
#[tokio::test]
async fn shell_basic_command() {
let tool = make_shell_tool();
let env: Arc<dyn ExecutionEnvironment> = Arc::new(ShellEnv {
result: ExecResult {
stdout: "hello".into(),
stderr: String::new(),
exit_code: 0,
timed_out: false,
duration_ms: 10,
},
});
let result = (tool.executor)(serde_json::json!({"command": "echo hello"}), env).await;
let output = result.unwrap();
assert!(output.contains("Exit code: 0"));
assert!(output.contains("hello"));
}
#[tokio::test]
async fn shell_with_timeout() {
let tool = make_shell_tool();
let env = Arc::new(ShellCapturingEnv {
captured_timeout: Mutex::new(None),
});
let env_clone: Arc<dyn ExecutionEnvironment> = env.clone();
let _result = (tool.executor)(
serde_json::json!({"command": "sleep 1", "timeout_ms": 5000}),
env_clone,
)
.await;
assert_eq!(*env.captured_timeout.lock().unwrap(), Some(5000));
}
#[tokio::test]
async fn shell_nonzero_exit_code() {
let tool = make_shell_tool();
let env: Arc<dyn ExecutionEnvironment> = Arc::new(ShellEnv {
result: ExecResult {
stdout: String::new(),
stderr: "error".into(),
exit_code: 1,
timed_out: false,
duration_ms: 10,
},
});
let result = (tool.executor)(serde_json::json!({"command": "false"}), env).await;
let output = result.unwrap();
assert!(output.contains("Exit code: 1"));
assert!(output.contains("error"));
}
#[tokio::test]
async fn shell_timeout_output() {
let tool = make_shell_tool();
let env: Arc<dyn ExecutionEnvironment> = Arc::new(ShellEnv {
result: ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: -1,
timed_out: true,
duration_ms: 10000,
},
});
let result = (tool.executor)(serde_json::json!({"command": "sleep 100"}), env).await;
let output = result.unwrap();
assert!(output.starts_with("Command timed out.\n"));
}
#[tokio::test]
async fn grep_basic() {
let tool = make_grep_tool();
let env: Arc<dyn ExecutionEnvironment> = Arc::new(GrepEnv {
results: vec!["src/main.rs:10:fn main()".into(), "src/lib.rs:5:pub fn".into()],
});
let result = (tool.executor)(serde_json::json!({"pattern": "fn"}), env).await;
let output = result.unwrap();
assert!(output.contains("src/main.rs:10:fn main()"));
assert!(output.contains("src/lib.rs:5:pub fn"));
}
#[tokio::test]
async fn glob_basic() {
let tool = make_glob_tool();
let env: Arc<dyn ExecutionEnvironment> = Arc::new(GlobEnv {
results: vec!["src/main.rs".into(), "src/lib.rs".into()],
});
let result = (tool.executor)(serde_json::json!({"pattern": "src/**/*.rs"}), env).await;
let output = result.unwrap();
assert!(output.contains("src/main.rs"));
assert!(output.contains("src/lib.rs"));
}
}

View file

@ -0,0 +1,221 @@
use crate::config::SessionConfig;
use std::collections::HashMap;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TruncationMode {
HeadTail,
Tail,
}
const TRUNCATION_WARNING_HEAD_TAIL: &str =
"\n\n[WARNING: Output truncated. Showing first and last portions.]";
const TRUNCATION_WARNING_TAIL: &str =
"\n\n[WARNING: Output truncated. Showing last portion only.]";
const LINE_TRUNCATION_WARNING: &str =
"\n\n[WARNING: Output truncated by line count. Showing first and last lines.]";
fn default_char_limits() -> HashMap<&'static str, usize> {
let mut m = HashMap::new();
m.insert("read_file", 50_000);
m.insert("shell", 30_000);
m.insert("grep", 20_000);
m.insert("glob", 20_000);
m.insert("edit_file", 10_000);
m.insert("write_file", 1_000);
m
}
fn default_line_limits() -> HashMap<&'static str, usize> {
let mut m = HashMap::new();
m.insert("shell", 256);
m.insert("grep", 200);
m.insert("glob", 500);
m
}
pub fn truncate_output(output: &str, max_chars: usize, mode: TruncationMode) -> String {
if output.len() <= max_chars {
return output.to_string();
}
match mode {
TruncationMode::HeadTail => {
let half = max_chars / 2;
let head = &output[..half];
let tail = &output[output.len() - half..];
format!("{head}{TRUNCATION_WARNING_HEAD_TAIL}\n\n{tail}")
}
TruncationMode::Tail => {
let tail = &output[output.len() - max_chars..];
format!("{TRUNCATION_WARNING_TAIL}\n\n{tail}")
}
}
}
pub fn truncate_lines(output: &str, max_lines: usize) -> String {
let lines: Vec<&str> = output.lines().collect();
if lines.len() <= max_lines {
return output.to_string();
}
let half = max_lines / 2;
let head: Vec<&str> = lines[..half].to_vec();
let tail: Vec<&str> = lines[lines.len() - half..].to_vec();
format!(
"{}{LINE_TRUNCATION_WARNING}\n\n{}",
head.join("\n"),
tail.join("\n")
)
}
pub fn truncate_tool_output(output: &str, tool_name: &str, config: &SessionConfig) -> String {
let builtin_char_limits = default_char_limits();
let builtin_line_limits = default_line_limits();
// Char truncation first
let char_limit = config
.tool_output_limits
.get(tool_name)
.copied()
.or_else(|| builtin_char_limits.get(tool_name).copied());
let after_chars = match char_limit {
Some(limit) => truncate_output(output, limit, TruncationMode::HeadTail),
None => output.to_string(),
};
// Then line truncation
let line_limit = config
.tool_line_limits
.get(tool_name)
.copied()
.or_else(|| builtin_line_limits.get(tool_name).copied());
match line_limit {
Some(limit) => truncate_lines(&after_chars, limit),
None => after_chars,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn under_limit_passthrough_chars() {
let output = "short output";
let result = truncate_output(output, 100, TruncationMode::HeadTail);
assert_eq!(result, output);
}
#[test]
fn under_limit_passthrough_lines() {
let output = "line1\nline2\nline3";
let result = truncate_lines(output, 10);
assert_eq!(result, output);
}
#[test]
fn head_tail_split() {
let output = "a".repeat(100);
let result = truncate_output(&output, 40, TruncationMode::HeadTail);
assert!(result.contains(&"a".repeat(20)));
assert!(result.contains(TRUNCATION_WARNING_HEAD_TAIL));
}
#[test]
fn tail_mode() {
let output = format!("{}BBB", "A".repeat(100));
let result = truncate_output(&output, 10, TruncationMode::Tail);
assert!(result.contains(TRUNCATION_WARNING_TAIL));
assert!(result.ends_with("AAAAAAABBB"));
}
#[test]
fn line_truncation_splits_head_tail() {
let lines: Vec<String> = (1..=20).map(|i| format!("line {i}")).collect();
let output = lines.join("\n");
let result = truncate_lines(&output, 6);
assert!(result.contains("line 1"));
assert!(result.contains("line 3"));
assert!(result.contains("line 18"));
assert!(result.contains("line 20"));
assert!(result.contains(LINE_TRUNCATION_WARNING));
}
#[test]
fn char_truncation_before_lines() {
// Create an output that is large in chars and many lines
let long_line = "x".repeat(50_000);
let output = format!("{long_line}\n{long_line}");
let config = SessionConfig::default();
let result = truncate_tool_output(&output, "shell", &config);
// Should have been char-truncated first (30k limit for shell)
assert!(result.len() < output.len());
}
#[test]
fn config_override_char_limit() {
let output = "x".repeat(200);
let mut config = SessionConfig::default();
config
.tool_output_limits
.insert("my_tool".into(), 50);
let result = truncate_tool_output(&output, "my_tool", &config);
assert!(result.len() < output.len());
assert!(result.contains(TRUNCATION_WARNING_HEAD_TAIL));
}
#[test]
fn config_override_line_limit() {
let lines: Vec<String> = (1..=100).map(|i| format!("line {i}")).collect();
let output = lines.join("\n");
let mut config = SessionConfig::default();
config.tool_line_limits.insert("my_tool".into(), 10);
let result = truncate_tool_output(&output, "my_tool", &config);
assert!(result.contains(LINE_TRUNCATION_WARNING));
}
#[test]
fn unknown_tool_no_truncation() {
let output = "x".repeat(200);
let config = SessionConfig::default();
let result = truncate_tool_output(&output, "unknown_tool", &config);
assert_eq!(result, output);
}
#[test]
fn default_char_limits_match_spec() {
let limits = default_char_limits();
assert_eq!(limits.get("read_file"), Some(&50_000));
assert_eq!(limits.get("shell"), Some(&30_000));
assert_eq!(limits.get("grep"), Some(&20_000));
assert_eq!(limits.get("glob"), Some(&20_000));
assert_eq!(limits.get("edit_file"), Some(&10_000));
assert_eq!(limits.get("write_file"), Some(&1_000));
}
#[test]
fn default_line_limits_match_spec() {
let limits = default_line_limits();
assert_eq!(limits.get("shell"), Some(&256));
assert_eq!(limits.get("grep"), Some(&200));
assert_eq!(limits.get("glob"), Some(&500));
}
#[test]
fn exact_limit_not_truncated() {
let output = "x".repeat(100);
let result = truncate_output(&output, 100, TruncationMode::HeadTail);
assert_eq!(result, output);
}
#[test]
fn exact_line_limit_not_truncated() {
let lines: Vec<String> = (1..=10).map(|i| format!("line {i}")).collect();
let output = lines.join("\n");
let result = truncate_lines(&output, 10);
assert_eq!(result, output);
}
}

View file

@ -0,0 +1,184 @@
use std::collections::HashMap;
use std::time::SystemTime;
use unified_llm::types::{ToolCall, ToolResult, Usage};
#[derive(Debug, Clone)]
pub enum Turn {
User {
content: String,
timestamp: SystemTime,
},
Assistant {
content: String,
tool_calls: Vec<ToolCall>,
reasoning: Option<String>,
usage: Usage,
response_id: String,
timestamp: SystemTime,
},
ToolResults {
results: Vec<ToolResult>,
timestamp: SystemTime,
},
System {
content: String,
timestamp: SystemTime,
},
Steering {
content: String,
timestamp: SystemTime,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionState {
Idle,
Processing,
AwaitingInput,
Closed,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EventKind {
SessionStart,
SessionEnd,
UserInput,
AssistantTextStart,
AssistantTextDelta,
AssistantTextEnd,
ToolCallStart,
ToolCallOutputDelta,
ToolCallEnd,
SteeringInjected,
TurnLimit,
LoopDetection,
ContextWindowWarning,
Error,
}
#[derive(Debug, Clone)]
pub struct SessionEvent {
pub kind: EventKind,
pub timestamp: SystemTime,
pub session_id: String,
pub data: HashMap<String, serde_json::Value>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn turn_user_construction() {
let turn = Turn::User {
content: "Hello".into(),
timestamp: SystemTime::now(),
};
match &turn {
Turn::User { content, .. } => assert_eq!(content, "Hello"),
_ => panic!("Expected User turn"),
}
}
#[test]
fn turn_assistant_construction() {
let turn = Turn::Assistant {
content: "Hi there".into(),
tool_calls: vec![],
reasoning: None,
usage: Usage::default(),
response_id: "resp_1".into(),
timestamp: SystemTime::now(),
};
match &turn {
Turn::Assistant {
content,
tool_calls,
reasoning,
response_id,
..
} => {
assert_eq!(content, "Hi there");
assert!(tool_calls.is_empty());
assert!(reasoning.is_none());
assert_eq!(response_id, "resp_1");
}
_ => panic!("Expected Assistant turn"),
}
}
#[test]
fn turn_tool_results_construction() {
let result = ToolResult {
tool_call_id: "call_1".into(),
content: serde_json::json!("result"),
is_error: false,
image_data: None,
image_media_type: None,
};
let turn = Turn::ToolResults {
results: vec![result],
timestamp: SystemTime::now(),
};
match &turn {
Turn::ToolResults { results, .. } => {
assert_eq!(results.len(), 1);
assert_eq!(results[0].tool_call_id, "call_1");
}
_ => panic!("Expected ToolResults turn"),
}
}
#[test]
fn turn_system_construction() {
let turn = Turn::System {
content: "System prompt".into(),
timestamp: SystemTime::now(),
};
match &turn {
Turn::System { content, .. } => assert_eq!(content, "System prompt"),
_ => panic!("Expected System turn"),
}
}
#[test]
fn turn_steering_construction() {
let turn = Turn::Steering {
content: "Focus on the task".into(),
timestamp: SystemTime::now(),
};
match &turn {
Turn::Steering { content, .. } => assert_eq!(content, "Focus on the task"),
_ => panic!("Expected Steering turn"),
}
}
#[test]
fn session_state_equality() {
assert_eq!(SessionState::Idle, SessionState::Idle);
assert_eq!(SessionState::Processing, SessionState::Processing);
assert_eq!(SessionState::AwaitingInput, SessionState::AwaitingInput);
assert_eq!(SessionState::Closed, SessionState::Closed);
assert_ne!(SessionState::Idle, SessionState::Closed);
}
#[test]
fn event_kind_equality() {
assert_eq!(EventKind::SessionStart, EventKind::SessionStart);
assert_ne!(EventKind::SessionStart, EventKind::SessionEnd);
assert_eq!(EventKind::LoopDetection, EventKind::LoopDetection);
}
#[test]
fn session_event_construction() {
let event = SessionEvent {
kind: EventKind::SessionStart,
timestamp: SystemTime::now(),
session_id: "sess_1".into(),
data: HashMap::new(),
};
assert_eq!(event.kind, EventKind::SessionStart);
assert_eq!(event.session_id, "sess_1");
assert!(event.data.is_empty());
}
}