mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-09-12 23:02:41 +00:00
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:
parent
255a93c2e8
commit
6a53634b05
21 changed files with 6930 additions and 0 deletions
26
crates/coding-agent-loop/Cargo.toml
Normal file
26
crates/coding-agent-loop/Cargo.toml
Normal 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
|
||||
64
crates/coding-agent-loop/src/config.rs
Normal file
64
crates/coding-agent-loop/src/config.rs
Normal 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);
|
||||
}
|
||||
}
|
||||
67
crates/coding-agent-loop/src/error.rs
Normal file
67
crates/coding-agent-loop/src/error.rs
Normal 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");
|
||||
}
|
||||
}
|
||||
109
crates/coding-agent-loop/src/event.rs
Normal file
109
crates/coding-agent-loop/src/event.rs
Normal 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();
|
||||
}
|
||||
}
|
||||
184
crates/coding-agent-loop/src/execution_env.rs
Normal file
184
crates/coding-agent-loop/src/execution_env.rs
Normal 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");
|
||||
}
|
||||
}
|
||||
278
crates/coding-agent-loop/src/history.rs
Normal file
278
crates/coding-agent-loop/src/history.rs
Normal 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);
|
||||
}
|
||||
}
|
||||
36
crates/coding-agent-loop/src/lib.rs
Normal file
36
crates/coding-agent-loop/src/lib.rs
Normal 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};
|
||||
572
crates/coding-agent-loop/src/local_env.rs
Normal file
572
crates/coding-agent-loop/src/local_env.rs
Normal 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");
|
||||
}
|
||||
}
|
||||
215
crates/coding-agent-loop/src/loop_detection.rs
Normal file
215
crates/coding-agent-loop/src/loop_detection.rs
Normal 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));
|
||||
}
|
||||
}
|
||||
275
crates/coding-agent-loop/src/profiles/anthropic.rs
Normal file
275
crates/coding-agent-loop/src/profiles/anthropic.rs
Normal 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()));
|
||||
}
|
||||
}
|
||||
263
crates/coding-agent-loop/src/profiles/gemini.rs
Normal file
263
crates/coding-agent-loop/src/profiles/gemini.rs
Normal 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()));
|
||||
}
|
||||
}
|
||||
111
crates/coding-agent-loop/src/profiles/mod.rs
Normal file
111
crates/coding-agent-loop/src/profiles/mod.rs
Normal 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"));
|
||||
}
|
||||
}
|
||||
705
crates/coding-agent-loop/src/profiles/openai.rs
Normal file
705
crates/coding-agent-loop/src/profiles/openai.rs
Normal 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\")"));
|
||||
}
|
||||
}
|
||||
225
crates/coding-agent-loop/src/project_docs.rs
Normal file
225
crates/coding-agent-loop/src/project_docs.rs
Normal 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");
|
||||
}
|
||||
}
|
||||
176
crates/coding-agent-loop/src/provider_profile.rs
Normal file
176
crates/coding-agent-loop/src/provider_profile.rs
Normal 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());
|
||||
}
|
||||
}
|
||||
1526
crates/coding-agent-loop/src/session.rs
Normal file
1526
crates/coding-agent-loop/src/session.rs
Normal file
File diff suppressed because it is too large
Load diff
599
crates/coding-agent-loop/src/subagent.rs
Normal file
599
crates/coding-agent-loop/src/subagent.rs
Normal 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")));
|
||||
}
|
||||
}
|
||||
238
crates/coding-agent-loop/src/tool_registry.rs
Normal file
238
crates/coding-agent-loop/src/tool_registry.rs
Normal 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());
|
||||
}
|
||||
}
|
||||
856
crates/coding-agent-loop/src/tools.rs
Normal file
856
crates/coding-agent-loop/src/tools.rs
Normal 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"));
|
||||
}
|
||||
}
|
||||
221
crates/coding-agent-loop/src/truncation.rs
Normal file
221
crates/coding-agent-loop/src/truncation.rs
Normal 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);
|
||||
}
|
||||
}
|
||||
184
crates/coding-agent-loop/src/types.rs
Normal file
184
crates/coding-agent-loop/src/types.rs
Normal 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());
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue