Close spec compliance gaps in coding-agent-loop

Replace all stub tools in profiles with real make_*_tool() factories,
wire up project docs discovery in session, add missing events
(SessionStart, SteeringInjected, Error), enrich environment context
block, fix per-tool truncation modes, improve loop detection to check
all groups, add SIGTERM-before-SIGKILL on timeout, and update subagent
with SubAgentResult struct.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Bryan Helmkamp 2026-02-20 13:54:50 -04:00
parent 890a6e9ccb
commit a8c576deba
19 changed files with 392 additions and 309 deletions

1
Cargo.lock generated
View file

@ -202,6 +202,7 @@ version = "0.1.0"
dependencies = [
"async-trait",
"futures",
"libc",
"serde",
"serde_json",
"thiserror",

View file

@ -19,6 +19,9 @@ uuid.workspace = true
futures.workspace = true
async-trait.workspace = true
[target.'cfg(unix)'.dependencies]
libc = "0.2"
[dev-dependencies]
tokio = { workspace = true, features = ["test-util", "macros"] }

View file

@ -12,6 +12,7 @@ pub struct SessionConfig {
pub enable_loop_detection: bool,
pub loop_detection_window: usize,
pub max_subagent_depth: usize,
pub git_root: Option<String>,
}
impl Default for SessionConfig {
@ -27,6 +28,7 @@ impl Default for SessionConfig {
enable_loop_detection: true,
loop_detection_window: 10,
max_subagent_depth: 1,
git_root: None,
}
}
}

View file

@ -34,6 +34,8 @@ pub trait ExecutionEnvironment: Send + Sync {
command: &str,
args: &[String],
timeout_ms: u64,
working_dir: Option<&str>,
env_vars: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String>;
async fn grep(
&self,
@ -79,6 +81,8 @@ mod tests {
_command: &str,
_args: &[String],
_timeout_ms: u64,
_working_dir: Option<&str>,
_env_vars: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: "output".into(),
@ -126,7 +130,7 @@ mod tests {
#[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();
let result = env.exec_command("echo", &[], 5000, None, None).await.unwrap();
assert_eq!(result.exit_code, 0);
assert!(!result.timed_out);
}

View file

@ -73,7 +73,7 @@ impl History {
}
Turn::System { content, .. } => Message::system(content),
Turn::Steering { content, .. } => Message {
role: Role::Developer,
role: Role::User,
content: vec![ContentPart::text(content)],
name: None,
tool_call_id: None,
@ -202,7 +202,7 @@ mod tests {
}
#[test]
fn steering_turn_maps_to_developer_message() {
fn steering_turn_maps_to_user_message() {
let mut history = History::new();
history.push(Turn::Steering {
content: "Focus on the main task".into(),
@ -210,7 +210,7 @@ mod tests {
});
let messages = history.convert_to_messages();
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].role, Role::Developer);
assert_eq!(messages[0].role, Role::User);
assert_eq!(messages[0].text(), "Focus on the main task");
}

View file

@ -23,14 +23,14 @@ 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 profiles::{AnthropicProfile, EnvContext, GeminiProfile, OpenAiProfile};
pub use provider_profile::ProviderProfile;
pub use session::Session;
pub use subagent::{SubAgent, SubAgentManager};
pub use subagent::{SubAgent, SubAgentManager, SubAgentResult};
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,
make_shell_tool_with_config, make_write_file_tool,
};
pub use truncation::{truncate_lines, truncate_output, truncate_tool_output, TruncationMode};
pub use types::{EventKind, SessionEvent, SessionState, Turn};

View file

@ -109,16 +109,29 @@ impl ExecutionEnvironment for LocalExecutionEnvironment {
command: &str,
args: &[String],
timeout_ms: u64,
working_dir: Option<&str>,
env_vars: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
let start = Instant::now();
let filtered_env: Vec<(String, String)> = std::env::vars()
let mut filtered_env: Vec<(String, String)> = std::env::vars()
.filter(|(key, _)| !Self::should_filter_env_var(key))
.collect();
if let Some(extra) = env_vars {
for (k, v) in extra {
filtered_env.push((k.clone(), v.clone()));
}
}
let effective_dir = working_dir.map_or_else(
|| self.working_directory.clone(),
std::path::PathBuf::from,
);
let mut cmd = Command::new(command);
cmd.args(args)
.current_dir(&self.working_directory)
.current_dir(&effective_dir)
.env_clear()
.envs(filtered_env)
.stdout(std::process::Stdio::piped())
@ -136,8 +149,33 @@ impl ExecutionEnvironment for LocalExecutionEnvironment {
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;
// SIGTERM first, then SIGKILL after 2 seconds
#[cfg(unix)]
if let Some(pid) = child.id() {
#[allow(clippy::cast_possible_wrap)]
unsafe {
libc::kill(pid as i32, libc::SIGTERM);
}
// Wait 2 seconds for graceful shutdown
if tokio::time::timeout(
std::time::Duration::from_secs(2),
child.wait(),
)
.await
.is_err()
{
let _ = child.kill().await;
let _ = child.wait().await;
}
} else {
let _ = child.kill().await;
let _ = child.wait().await;
}
#[cfg(not(unix))]
{
let _ = child.kill().await;
let _ = child.wait().await;
}
(true, -1)
};
@ -209,7 +247,19 @@ impl ExecutionEnvironment for LocalExecutionEnvironment {
.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();
let mut results: Vec<String> = stdout.lines().map(String::from).filter(|l| !l.is_empty()).collect();
// Sort by mtime (newest first)
results.sort_by(|a, b| {
let mtime_a = std::fs::metadata(a)
.and_then(|m| m.modified())
.unwrap_or(std::time::SystemTime::UNIX_EPOCH);
let mtime_b = std::fs::metadata(b)
.and_then(|m| m.modified())
.unwrap_or(std::time::SystemTime::UNIX_EPOCH);
mtime_b.cmp(&mtime_a)
});
Ok(results)
}
@ -361,7 +411,7 @@ mod tests {
let dir = temp_dir();
let env = LocalExecutionEnvironment::new(dir.clone());
let result = env
.exec_command("echo", &["hello".into()], 5000)
.exec_command("echo", &["hello".into()], 5000, None, None)
.await
.unwrap();
@ -377,7 +427,7 @@ mod tests {
let dir = temp_dir();
let env = LocalExecutionEnvironment::new(dir.clone());
let result = env
.exec_command("sh", &["-c".into(), "exit 42".into()], 5000)
.exec_command("sh", &["-c".into(), "exit 42".into()], 5000, None, None)
.await
.unwrap();
@ -391,7 +441,7 @@ mod tests {
let dir = temp_dir();
let env = LocalExecutionEnvironment::new(dir.clone());
let result = env
.exec_command("sleep", &["10".into()], 200)
.exec_command("sleep", &["10".into()], 200, None, None)
.await
.unwrap();
@ -405,7 +455,7 @@ mod tests {
let dir = temp_dir();
let env = LocalExecutionEnvironment::new(dir.clone());
let result = env
.exec_command("sh", &["-c".into(), "echo err >&2".into()], 5000)
.exec_command("sh", &["-c".into(), "echo err >&2".into()], 5000, None, None)
.await
.unwrap();

View file

@ -71,11 +71,23 @@ fn is_repeating_pattern(signatures: &[u64], pattern_len: usize) -> bool {
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];
// Check ALL preceding groups in window match, not just the last 2
let num_groups = signatures.len() / pattern_len;
if num_groups < 2 {
return false;
}
pattern == preceding
// Walk backwards through all complete groups
let groups_start = signatures.len() - (num_groups * pattern_len);
for group_idx in 0..num_groups - 1 {
let start = groups_start + group_idx * pattern_len;
let group = &signatures[start..start + pattern_len];
if group != pattern {
return false;
}
}
true
}
#[cfg(test)]

View file

@ -1,9 +1,14 @@
use crate::config::SessionConfig;
use crate::execution_env::ExecutionEnvironment;
use crate::provider_profile::ProviderProfile;
use crate::tool_registry::ToolRegistry;
use crate::tools::{
make_edit_file_tool, make_glob_tool, make_grep_tool, make_read_file_tool,
make_shell_tool_with_config, make_write_file_tool,
};
use unified_llm::types::ToolDefinition;
use super::{build_env_context_block, stub_tool};
use super::build_env_context_block;
pub struct AnthropicProfile {
model: String,
@ -13,84 +18,18 @@ pub struct AnthropicProfile {
impl AnthropicProfile {
#[must_use]
pub fn new(model: impl Into<String>) -> Self {
let config = SessionConfig {
default_command_timeout_ms: 120_000,
..SessionConfig::default()
};
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"]
}),
));
registry.register(make_read_file_tool());
registry.register(make_write_file_tool());
registry.register(make_edit_file_tool());
registry.register(make_shell_tool_with_config(&config));
registry.register(make_grep_tool());
registry.register(make_glob_tool());
Self {
model: model.into(),
@ -112,6 +51,10 @@ impl ProviderProfile for AnthropicProfile {
&self.registry
}
fn tool_registry_mut(&mut self) -> &mut ToolRegistry {
&mut self.registry
}
fn build_system_prompt(
&self,
env: &dyn ExecutionEnvironment,
@ -185,6 +128,8 @@ mod tests {
_: &str,
_: &[String],
_: u64,
_: Option<&str>,
_: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),

View file

@ -1,9 +1,13 @@
use crate::execution_env::ExecutionEnvironment;
use crate::provider_profile::ProviderProfile;
use crate::tool_registry::ToolRegistry;
use crate::tools::{
make_edit_file_tool, make_glob_tool, make_grep_tool, make_read_file_tool, make_shell_tool,
make_write_file_tool,
};
use unified_llm::types::ToolDefinition;
use super::{build_env_context_block, stub_tool};
use super::build_env_context_block;
pub struct GeminiProfile {
model: String,
@ -15,82 +19,12 @@ impl GeminiProfile {
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"]
}),
));
registry.register(make_read_file_tool());
registry.register(make_write_file_tool());
registry.register(make_edit_file_tool());
registry.register(make_shell_tool());
registry.register(make_grep_tool());
registry.register(make_glob_tool());
Self {
model: model.into(),
@ -112,6 +46,10 @@ impl ProviderProfile for GeminiProfile {
&self.registry
}
fn tool_registry_mut(&mut self) -> &mut ToolRegistry {
&mut self.registry
}
fn build_system_prompt(
&self,
env: &dyn ExecutionEnvironment,
@ -185,6 +123,8 @@ mod tests {
_: &str,
_: &[String],
_: u64,
_: Option<&str>,
_: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),

View file

@ -7,32 +7,44 @@ 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()
)
/// Additional context for building environment blocks
#[derive(Default)]
pub struct EnvContext {
pub git_branch: Option<String>,
pub is_git_repo: bool,
pub date: String,
pub model_name: String,
}
#[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()) })
}),
pub fn build_env_context_block(env: &dyn ExecutionEnvironment) -> String {
build_env_context_block_with(env, &EnvContext::default())
}
#[must_use]
pub fn build_env_context_block_with(env: &dyn ExecutionEnvironment, ctx: &EnvContext) -> String {
let mut lines = vec![
"# Environment".to_string(),
format!("- Working directory: {}", env.working_directory()),
format!("- Platform: {}", env.platform()),
format!("- OS: {}", env.os_version()),
];
if ctx.is_git_repo {
lines.push(format!("- Is a git repository: {}", ctx.is_git_repo));
}
if let Some(ref branch) = ctx.git_branch {
lines.push(format!("- Git branch: {branch}"));
}
if !ctx.date.is_empty() {
lines.push(format!("- Date: {}", ctx.date));
}
if !ctx.model_name.is_empty() {
lines.push(format!("- Model: {}", ctx.model_name));
}
lines.join("\n")
}
#[cfg(test)]
@ -62,6 +74,8 @@ mod tests {
_: &str,
_: &[String],
_: u64,
_: Option<&str>,
_: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
@ -108,4 +122,20 @@ mod tests {
assert!(block.contains("/home/test"));
assert!(block.contains("Linux 6.1.0"));
}
#[test]
fn env_context_block_with_extra_context() {
let env = TestEnv;
let ctx = EnvContext {
git_branch: Some("main".into()),
is_git_repo: true,
date: "2026-02-20".into(),
model_name: "claude-opus-4-6".into(),
};
let block = build_env_context_block_with(&env, &ctx);
assert!(block.contains("Git branch: main"));
assert!(block.contains("Is a git repository: true"));
assert!(block.contains("Date: 2026-02-20"));
assert!(block.contains("Model: claude-opus-4-6"));
}
}

View file

@ -1,10 +1,13 @@
use crate::execution_env::ExecutionEnvironment;
use crate::provider_profile::ProviderProfile;
use crate::tool_registry::{RegisteredTool, ToolRegistry};
use crate::tools::{
make_glob_tool, make_grep_tool, make_read_file_tool, make_shell_tool, make_write_file_tool,
};
use std::sync::Arc;
use unified_llm::types::ToolDefinition;
use super::{build_env_context_block, stub_tool};
use super::build_env_context_block;
pub struct OpenAiProfile {
model: String,
@ -16,69 +19,11 @@ impl OpenAiProfile {
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_read_file_tool());
registry.register(make_write_file_tool());
registry.register(make_shell_tool());
registry.register(make_grep_tool());
registry.register(make_glob_tool());
registry.register(make_apply_patch_tool());
Self {
@ -101,6 +46,10 @@ impl ProviderProfile for OpenAiProfile {
&self.registry
}
fn tool_registry_mut(&mut self) -> &mut ToolRegistry {
&mut self.registry
}
fn build_system_prompt(
&self,
env: &dyn ExecutionEnvironment,
@ -417,6 +366,8 @@ mod tests {
_: &str,
_: &[String],
_: u64,
_: Option<&str>,
_: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
@ -494,6 +445,8 @@ mod tests {
_: &str,
_: &[String],
_: u64,
_: Option<&str>,
_: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),

View file

@ -12,7 +12,7 @@ pub async fn discover_project_docs(
let candidate_filenames: Vec<&str> = match provider_id {
"anthropic" => vec!["AGENTS.md", "CLAUDE.md"],
"openai" => vec!["AGENTS.md", ".github/copilot-instructions.md"],
"openai" => vec!["AGENTS.md", ".codex/instructions.md"],
"gemini" => vec!["AGENTS.md", "GEMINI.md"],
_ => vec!["AGENTS.md"],
};
@ -70,7 +70,7 @@ fn build_directory_walk(git_root: &str, working_dir: &str) -> Vec<String> {
}
fn truncate_to_budget(content: &str, budget: usize) -> String {
const MARKER: &str = "... [truncated]";
const MARKER: &str = "[Project instructions truncated at 32KB]";
if budget <= MARKER.len() {
return MARKER[..budget].to_string();
}
@ -112,7 +112,7 @@ mod tests {
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
async fn exec_command(&self, _: &str, _: &[String], _: u64, _: Option<&str>, _: Option<&std::collections::HashMap<String, String>>) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
@ -160,7 +160,7 @@ mod tests {
files.insert("/repo/AGENTS.md".into(), "agents".into());
files.insert("/repo/CLAUDE.md".into(), "claude".into());
files.insert(
"/repo/.github/copilot-instructions.md".into(),
"/repo/.codex/instructions.md".into(),
"copilot".into(),
);
files.insert("/repo/GEMINI.md".into(), "gemini".into());
@ -203,7 +203,7 @@ mod tests {
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[1].ends_with("[Project instructions truncated at 32KB]"));
assert!(docs[0].len() + docs[1].len() <= BUDGET_BYTES);
}

View file

@ -6,6 +6,7 @@ pub trait ProviderProfile: Send + Sync {
fn id(&self) -> String;
fn model(&self) -> String;
fn tool_registry(&self) -> &ToolRegistry;
fn tool_registry_mut(&mut self) -> &mut ToolRegistry;
fn build_system_prompt(
&self,
env: &dyn ExecutionEnvironment,
@ -46,6 +47,8 @@ mod tests {
_: &str,
_: &[String],
_: u64,
_: Option<&str>,
_: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
@ -105,6 +108,9 @@ mod tests {
fn tool_registry(&self) -> &ToolRegistry {
&self.registry
}
fn tool_registry_mut(&mut self) -> &mut ToolRegistry {
&mut self.registry
}
fn build_system_prompt(
&self,
env: &dyn ExecutionEnvironment,

View file

@ -3,6 +3,7 @@ use crate::error::AgentError;
use crate::event::EventEmitter;
use crate::history::History;
use crate::loop_detection::detect_loop;
use crate::project_docs::discover_project_docs;
use crate::provider_profile::ProviderProfile;
use crate::truncation::truncate_tool_output;
use crate::types::{EventKind, SessionEvent, SessionState, Turn};
@ -28,6 +29,7 @@ pub struct Session {
steering_queue: Arc<Mutex<VecDeque<String>>>,
followup_queue: Arc<Mutex<VecDeque<String>>>,
abort_flag: Arc<AtomicBool>,
project_docs: Vec<String>,
}
impl Session {
@ -50,6 +52,20 @@ impl Session {
steering_queue: Arc::new(Mutex::new(VecDeque::new())),
followup_queue: Arc::new(Mutex::new(VecDeque::new())),
abort_flag: Arc::new(AtomicBool::new(false)),
project_docs: Vec::new(),
}
}
/// Initialize session by discovering project docs. Call before `process_input`.
pub async fn initialize(&mut self) {
if let Some(ref git_root) = self.config.git_root {
self.project_docs = discover_project_docs(
self.execution_env.as_ref(),
git_root,
self.execution_env.working_directory(),
&self.provider_profile.id(),
)
.await;
}
}
@ -96,6 +112,12 @@ impl Session {
}
pub async fn process_input(&mut self, input: &str) -> Result<(), AgentError> {
self.event_emitter.emit(
EventKind::SessionStart,
self.id.clone(),
HashMap::new(),
);
// Use a queue to avoid recursive async calls for followups
let mut current_input = input.to_string();
@ -182,6 +204,16 @@ impl Session {
let response = match self.llm_client.complete(&request).await {
Ok(resp) => resp,
Err(err) => {
let mut error_data = HashMap::new();
error_data.insert(
"error".to_string(),
serde_json::json!(err.to_string()),
);
self.event_emitter.emit(
EventKind::Error,
self.id.clone(),
error_data,
);
if is_auth_error(&err) {
self.state = SessionState::Closed;
}
@ -264,13 +296,18 @@ impl Session {
content: msg,
timestamp: SystemTime::now(),
});
self.event_emitter.emit(
EventKind::SteeringInjected,
self.id.clone(),
HashMap::new(),
);
}
}
fn build_request(&self) -> Request {
let system_prompt = self
.provider_profile
.build_system_prompt(self.execution_env.as_ref(), &[]);
.build_system_prompt(self.execution_env.as_ref(), &self.project_docs);
let mut messages = vec![Message::system(system_prompt)];
messages.extend(self.history.convert_to_messages());
@ -518,7 +555,7 @@ impl Session {
fn estimate_token_count(&self) -> usize {
let system_prompt = self
.provider_profile
.build_system_prompt(self.execution_env.as_ref(), &[]);
.build_system_prompt(self.execution_env.as_ref(), &self.project_docs);
let mut total_chars = system_prompt.len();
for turn in self.history.turns() {
@ -710,6 +747,8 @@ mod tests {
_command: &str,
_args: &[String],
_timeout_ms: u64,
_working_dir: Option<&str>,
_env_vars: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: "mock output".into(),
@ -785,6 +824,10 @@ mod tests {
&self.registry
}
fn tool_registry_mut(&mut self) -> &mut ToolRegistry {
&mut self.registry
}
fn build_system_prompt(
&self,
_env: &dyn ExecutionEnvironment,
@ -1340,6 +1383,10 @@ mod tests {
&self.registry
}
fn tool_registry_mut(&mut self) -> &mut ToolRegistry {
&mut self.registry
}
fn build_system_prompt(
&self,
_env: &dyn ExecutionEnvironment,

View file

@ -9,10 +9,17 @@ use unified_llm::types::ToolDefinition;
pub type SessionFactory = Arc<dyn Fn() -> Session + Send + Sync>;
#[derive(Debug, Clone)]
pub struct SubAgentResult {
pub output: String,
pub success: bool,
pub turns_used: usize,
}
pub struct SubAgent {
id: String,
depth: usize,
task: Option<tokio::task::JoinHandle<Result<String, AgentError>>>,
task: Option<tokio::task::JoinHandle<Result<SubAgentResult, AgentError>>>,
followup_queue: Arc<Mutex<VecDeque<String>>>,
abort_flag: Arc<AtomicBool>,
}
@ -58,8 +65,9 @@ impl SubAgentManager {
let abort_flag = session.abort_flag_handle();
let task = tokio::spawn(async move {
session.process_input(&task_prompt).await?;
let result = session.process_input(&task_prompt).await;
let turns = session.history().turns();
let turns_used = turns.len();
let last_text = turns.iter().rev().find_map(|t| {
if let Turn::Assistant { content, .. } = t {
Some(content.clone())
@ -67,7 +75,15 @@ impl SubAgentManager {
None
}
});
Ok(last_text.unwrap_or_default())
let success = result.is_ok();
if let Err(e) = result {
return Err(e);
}
Ok(SubAgentResult {
output: last_text.unwrap_or_default(),
success,
turns_used,
})
});
self.agents.insert(
@ -99,7 +115,7 @@ impl SubAgentManager {
Ok(())
}
pub async fn wait(&mut self, agent_id: &str) -> Result<String, String> {
pub async fn wait(&mut self, agent_id: &str) -> Result<SubAgentResult, String> {
let mut agent = self
.agents
.remove(agent_id)
@ -149,6 +165,18 @@ pub fn make_spawn_agent_tool(
"task": {
"type": "string",
"description": "The task description for the subagent"
},
"working_dir": {
"type": "string",
"description": "Working directory for the subagent"
},
"model": {
"type": "string",
"description": "Model to use for the subagent"
},
"max_turns": {
"type": "integer",
"description": "Maximum number of turns for the subagent"
}
},
"required": ["task"]
@ -240,7 +268,11 @@ pub fn make_wait_tool(
.ok_or_else(|| "Missing required parameter: agent_id".to_string())?;
let mut mgr = manager.lock().await;
mgr.wait(agent_id).await
let result = mgr.wait(agent_id).await?;
Ok(format!(
"Agent completed (success: {}, turns: {})\n\n{}",
result.success, result.turns_used, result.output
))
})
}),
}
@ -361,6 +393,8 @@ mod tests {
_command: &str,
_args: &[String],
_timeout_ms: u64,
_working_dir: Option<&str>,
_env_vars: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: "mock output".into(),
@ -422,6 +456,9 @@ mod tests {
fn tool_registry(&self) -> &ToolRegistry {
&self.registry
}
fn tool_registry_mut(&mut self) -> &mut ToolRegistry {
&mut self.registry
}
fn build_system_prompt(
&self,
_env: &dyn ExecutionEnvironment,
@ -551,7 +588,10 @@ mod tests {
let result = manager.wait(&agent_id).await;
assert!(result.is_ok());
assert_eq!(result.unwrap(), "Task completed successfully");
let agent_result = result.unwrap();
assert_eq!(agent_result.output, "Task completed successfully");
assert!(agent_result.success);
assert!(agent_result.turns_used > 0);
assert!(manager.get(&agent_id).is_none());
}

View file

@ -187,6 +187,8 @@ mod tests {
_: &str,
_: &[String],
_: u64,
_: Option<&str>,
_: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),

View file

@ -1,3 +1,4 @@
use crate::config::SessionConfig;
use crate::execution_env::GrepOptions;
use crate::tool_registry::RegisteredTool;
use std::fmt::Write;
@ -15,7 +16,7 @@ pub fn make_read_file_tool() -> RegisteredTool {
"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"}
"limit": {"type": "integer", "description": "Number of lines to read (default 2000)"}
},
"required": ["file_path"]
}),
@ -30,19 +31,19 @@ pub fn make_read_file_tool() -> RegisteredTool {
let content = env.read_file(file_path).await?;
if offset.is_none() && limit.is_none() {
return Ok(content);
}
// Default limit of 2000 lines when no limit param provided
let effective_limit = limit.unwrap_or(2000);
#[allow(clippy::cast_possible_truncation)]
let offset = offset.unwrap_or(1) as usize;
let offset_val = offset.unwrap_or(1) as usize;
let lines: Vec<&str> = content.lines().collect();
let start = if offset > 0 { offset - 1 } else { 0 };
let start = if offset_val > 0 { offset_val - 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(),
};
let selected: Vec<&str> = lines
.into_iter()
.skip(start)
.take(effective_limit as usize)
.collect();
Ok(selected.join("\n"))
})
}),
@ -150,6 +151,13 @@ pub fn make_edit_file_tool() -> RegisteredTool {
#[must_use]
pub fn make_shell_tool() -> RegisteredTool {
make_shell_tool_with_config(&SessionConfig::default())
}
#[must_use]
pub fn make_shell_tool_with_config(config: &SessionConfig) -> RegisteredTool {
let default_timeout = config.default_command_timeout_ms;
let max_timeout = config.max_command_timeout_ms;
RegisteredTool {
definition: ToolDefinition {
name: "shell".into(),
@ -158,12 +166,13 @@ pub fn make_shell_tool() -> RegisteredTool {
"type": "object",
"properties": {
"command": {"type": "string", "description": "The shell command to execute"},
"timeout_ms": {"type": "integer", "description": "Timeout in milliseconds (default 10000)"}
"timeout_ms": {"type": "integer", "description": "Timeout in milliseconds"},
"description": {"type": "string", "description": "Description of what this command does"}
},
"required": ["command"]
}),
},
executor: Arc::new(|args, env| {
executor: Arc::new(move |args, env| {
Box::pin(async move {
let command = args["command"]
.as_str()
@ -171,13 +180,16 @@ pub fn make_shell_tool() -> RegisteredTool {
let timeout_ms = args
.get("timeout_ms")
.and_then(serde_json::Value::as_u64)
.unwrap_or(10000);
.unwrap_or(default_timeout)
.min(max_timeout);
let result = env
.exec_command(
"/bin/bash",
&["-c".into(), command.into()],
timeout_ms,
None,
None,
)
.await?;
@ -256,7 +268,8 @@ pub fn make_glob_tool() -> RegisteredTool {
parameters: serde_json::json!({
"type": "object",
"properties": {
"pattern": {"type": "string", "description": "Glob pattern to match files"}
"pattern": {"type": "string", "description": "Glob pattern to match files"},
"path": {"type": "string", "description": "Directory to search in (default: working directory)"}
},
"required": ["pattern"]
}),
@ -266,8 +279,16 @@ pub fn make_glob_tool() -> RegisteredTool {
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);
let results = env.glob(pattern).await?;
let full_pattern = match path {
Some(dir) => format!("{dir}/{pattern}"),
None => pattern.to_string(),
};
let results = env.glob(&full_pattern).await?;
Ok(results.join("\n"))
})
}),
@ -299,7 +320,7 @@ mod tests {
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
async fn exec_command(&self, _: &str, _: &[String], _: u64, _: Option<&str>, _: Option<&std::collections::HashMap<String, String>>) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
@ -350,7 +371,7 @@ mod tests {
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
async fn exec_command(&self, _: &str, _: &[String], _: u64, _: Option<&str>, _: Option<&std::collections::HashMap<String, String>>) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
@ -402,7 +423,7 @@ mod tests {
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
async fn exec_command(&self, _: &str, _: &[String], _: u64, _: Option<&str>, _: Option<&std::collections::HashMap<String, String>>) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
@ -452,7 +473,7 @@ mod tests {
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
async fn exec_command(&self, _: &str, _: &[String], _: u64, _: Option<&str>, _: Option<&std::collections::HashMap<String, String>>) -> Result<ExecResult, String> {
Ok(self.result.clone())
}
async fn grep(&self, _: &str, _: &str, _: &GrepOptions) -> Result<Vec<String>, String> {
@ -501,6 +522,8 @@ mod tests {
_: &str,
_: &[String],
timeout_ms: u64,
_: Option<&str>,
_: Option<&std::collections::HashMap<String, String>>,
) -> Result<ExecResult, String> {
*self.captured_timeout.lock().unwrap() = Some(timeout_ms);
Ok(ExecResult {
@ -552,7 +575,7 @@ mod tests {
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
async fn exec_command(&self, _: &str, _: &[String], _: u64, _: Option<&str>, _: Option<&std::collections::HashMap<String, String>>) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
@ -602,7 +625,7 @@ mod tests {
async fn list_directory(&self, _: &str) -> Result<Vec<DirEntry>, String> {
Ok(vec![])
}
async fn exec_command(&self, _: &str, _: &[String], _: u64) -> Result<ExecResult, String> {
async fn exec_command(&self, _: &str, _: &[String], _: u64, _: Option<&str>, _: Option<&std::collections::HashMap<String, String>>) -> Result<ExecResult, String> {
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),

View file

@ -7,13 +7,6 @@ pub enum TruncationMode {
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);
@ -22,6 +15,8 @@ fn default_char_limits() -> HashMap<&'static str, usize> {
m.insert("glob", 20_000);
m.insert("edit_file", 10_000);
m.insert("write_file", 1_000);
m.insert("apply_patch", 10_000);
m.insert("spawn_agent", 30_000);
m
}
@ -33,21 +28,40 @@ fn default_line_limits() -> HashMap<&'static str, usize> {
m
}
fn default_truncation_modes() -> HashMap<&'static str, TruncationMode> {
let mut m = HashMap::new();
m.insert("grep", TruncationMode::Tail);
m.insert("glob", TruncationMode::Tail);
m.insert("edit_file", TruncationMode::Tail);
m.insert("apply_patch", TruncationMode::Tail);
m.insert("write_file", TruncationMode::Tail);
m
}
pub fn truncate_output(output: &str, max_chars: usize, mode: TruncationMode) -> String {
if output.len() <= max_chars {
return output.to_string();
}
let removed = output.len() - max_chars;
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}")
format!(
"{head}\n\n[WARNING: Output truncated. {removed} characters removed. \
Full output available in event stream. Retry with smaller scope if needed.]\n\n{tail}"
)
}
TruncationMode::Tail => {
let tail = &output[output.len() - max_chars..];
format!("{TRUNCATION_WARNING_TAIL}\n\n{tail}")
format!(
"\n\n[WARNING: Output truncated. {removed} characters removed. \
Showing last portion only. Full output available in event stream. \
Retry with smaller scope if needed.]\n\n{tail}"
)
}
}
}
@ -61,9 +75,11 @@ pub fn truncate_lines(output: &str, max_lines: usize) -> String {
let half = max_lines / 2;
let head: Vec<&str> = lines[..half].to_vec();
let tail: Vec<&str> = lines[lines.len() - half..].to_vec();
let omitted = lines.len() - max_lines;
format!(
"{}{LINE_TRUNCATION_WARNING}\n\n{}",
"{}\n\n[WARNING: Output truncated by line count. {omitted} lines omitted. \
Showing first and last lines.]\n\n{}",
head.join("\n"),
tail.join("\n")
)
@ -72,6 +88,13 @@ pub fn truncate_lines(output: &str, max_lines: usize) -> String {
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();
let builtin_modes = default_truncation_modes();
// Determine truncation mode for this tool (default HeadTail)
let mode = builtin_modes
.get(tool_name)
.copied()
.unwrap_or(TruncationMode::HeadTail);
// Char truncation first
let char_limit = config
@ -81,7 +104,7 @@ pub fn truncate_tool_output(output: &str, tool_name: &str, config: &SessionConfi
.or_else(|| builtin_char_limits.get(tool_name).copied());
let after_chars = match char_limit {
Some(limit) => truncate_output(output, limit, TruncationMode::HeadTail),
Some(limit) => truncate_output(output, limit, mode),
None => output.to_string(),
};
@ -121,14 +144,16 @@ mod tests {
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));
assert!(result.contains("Output truncated"));
assert!(result.contains("60 characters removed"));
}
#[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.contains("Output truncated"));
assert!(result.contains("Showing last portion only"));
assert!(result.ends_with("AAAAAAABBB"));
}
@ -141,7 +166,7 @@ mod tests {
assert!(result.contains("line 3"));
assert!(result.contains("line 18"));
assert!(result.contains("line 20"));
assert!(result.contains(LINE_TRUNCATION_WARNING));
assert!(result.contains("14 lines omitted"));
}
#[test]
@ -164,7 +189,7 @@ mod tests {
.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));
assert!(result.contains("Output truncated"));
}
#[test]
@ -174,7 +199,7 @@ mod tests {
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));
assert!(result.contains("lines omitted"));
}
#[test]