use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use async_trait::async_trait;
use tokio_util::sync::CancellationToken;
use edgecrab_state::SessionDb;
use edgecrab_types::{Message, Platform, ToolError, ToolSchema};
use crate::config_ref::AppConfigRef;
use crate::process_table::ProcessTable;
pub struct ClarifyRequest {
pub question: String,
pub choices: Option<Vec<String>>,
pub response_tx: tokio::sync::oneshot::Sender<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ApprovalResponse {
Once,
Session,
Always,
Deny,
}
pub struct ApprovalRequest {
pub command: String,
pub full_command: String,
pub reasons: Vec<String>,
pub response_tx: tokio::sync::oneshot::Sender<ApprovalResponse>,
}
pub const MAX_CLARIFY_CHOICES: usize = 4;
#[derive(Debug, Clone)]
pub struct SubAgentResult {
pub summary: String,
pub api_calls: u32,
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_read_tokens: u64,
pub cache_write_tokens: u64,
pub reasoning_tokens: u64,
pub model: Option<String>,
pub interrupted: bool,
pub budget_exhausted: bool,
pub messages: Vec<Message>,
}
#[derive(Debug, Clone)]
pub enum DelegationEvent {
TaskStarted {
task_index: usize,
task_count: usize,
goal: String,
},
Thinking {
task_index: usize,
task_count: usize,
text: String,
},
ToolCalled {
task_index: usize,
task_count: usize,
tool_name: String,
args_json: String,
},
TaskFinished {
task_index: usize,
task_count: usize,
status: String,
duration_ms: u64,
summary: String,
api_calls: u32,
model: Option<String>,
},
}
#[derive(Debug, Clone)]
pub struct SubAgentRunRequest {
pub goal: String,
pub system_prompt: String,
pub enabled_toolsets: Vec<String>,
pub max_iterations: u32,
pub model_override: Option<String>,
pub parent_cancel: CancellationToken,
pub progress_tx: Option<tokio::sync::mpsc::UnboundedSender<DelegationEvent>>,
pub task_index: usize,
pub task_count: usize,
}
#[async_trait]
pub trait SubAgentRunner: Send + Sync {
async fn run_task(&self, request: SubAgentRunRequest) -> Result<SubAgentResult, String>;
}
#[async_trait]
pub trait ToolHandler: Send + Sync + 'static {
fn name(&self) -> &'static str;
fn toolset(&self) -> &'static str;
fn schema(&self) -> ToolSchema;
async fn execute(
&self,
args: serde_json::Value,
ctx: &ToolContext,
) -> Result<String, ToolError>;
fn is_available(&self) -> bool {
true
}
fn check_fn(&self, _ctx: &ToolContext) -> bool {
true
}
fn parallel_safe(&self) -> bool {
false
}
fn emoji(&self) -> &'static str {
"⚡"
}
}
inventory::collect!(&'static dyn ToolHandler);
pub struct ToolContext {
pub task_id: String,
pub cwd: PathBuf,
pub session_id: String,
pub user_task: Option<String>,
pub cancel: CancellationToken,
pub config: AppConfigRef,
pub state_db: Option<Arc<SessionDb>>,
pub platform: Platform,
pub process_table: Option<Arc<ProcessTable>>,
pub provider: Option<Arc<dyn edgequake_llm::LLMProvider>>,
pub tool_registry: Option<Arc<ToolRegistry>>,
pub delegate_depth: u32,
pub sub_agent_runner: Option<Arc<dyn SubAgentRunner>>,
pub delegation_event_tx: Option<tokio::sync::mpsc::UnboundedSender<DelegationEvent>>,
pub clarify_tx: Option<tokio::sync::mpsc::UnboundedSender<ClarifyRequest>>,
pub approval_tx: Option<tokio::sync::mpsc::UnboundedSender<ApprovalRequest>>,
pub on_skills_changed: Option<Arc<dyn Fn() + Send + Sync>>,
pub gateway_sender: Option<Arc<dyn GatewaySender>>,
pub origin_chat: Option<(String, String)>,
pub session_key: Option<String>,
pub todo_store: Option<Arc<crate::tools::todo::TodoStore>>,
}
#[async_trait]
pub trait GatewaySender: Send + Sync + 'static {
async fn send_message(
&self,
platform: &str,
recipient: &str,
message: &str,
) -> Result<(), String>;
async fn list_targets(&self) -> Result<Vec<String>, String>;
}
impl ToolContext {
#[cfg(test)]
pub fn test_context() -> Self {
Self {
task_id: "test-task".into(),
cwd: std::env::temp_dir(),
session_id: "test-session".into(),
user_task: None,
cancel: CancellationToken::new(),
config: AppConfigRef::default(),
state_db: None,
platform: Platform::Cli,
process_table: None,
provider: None,
tool_registry: None,
delegate_depth: 0,
sub_agent_runner: None,
delegation_event_tx: None,
clarify_tx: None,
approval_tx: None,
on_skills_changed: None,
gateway_sender: None,
origin_chat: None,
session_key: None,
todo_store: None,
}
}
}
pub struct ToolRegistry {
tools: HashMap<&'static str, &'static dyn ToolHandler>,
toolset_index: HashMap<&'static str, Vec<&'static str>>,
dynamic_tools: HashMap<String, Box<dyn ToolHandler>>,
}
impl ToolRegistry {
fn toolset_allowed_in_ctx(toolset: &str, ctx: &ToolContext) -> bool {
ctx.config.is_toolset_enabled(toolset)
}
pub fn new() -> Self {
let mut tools = HashMap::new();
let mut toolset_index: HashMap<&str, Vec<&str>> = HashMap::new();
for handler in inventory::iter::<&dyn ToolHandler> {
tools.insert(handler.name(), *handler);
toolset_index
.entry(handler.toolset())
.or_default()
.push(handler.name());
}
Self {
tools,
toolset_index,
dynamic_tools: HashMap::new(),
}
}
pub fn get_definitions(
&self,
enabled: Option<&[String]>,
disabled: Option<&[String]>,
ctx: &ToolContext,
) -> Vec<ToolSchema> {
let is_eligible = |toolset: &str, available: bool, passes_check: bool| -> bool {
available
&& passes_check
&& Self::toolset_allowed_in_ctx(toolset, ctx)
&& enabled.is_none_or(|sets| sets.iter().any(|s| s == toolset))
&& disabled.is_none_or(|sets| !sets.iter().any(|s| s == toolset))
};
let static_schemas = self
.tools
.values()
.filter(|h| is_eligible(h.toolset(), h.is_available(), h.check_fn(ctx)))
.map(|h| h.schema());
let dynamic_schemas = self
.dynamic_tools
.values()
.filter(|h| is_eligible(h.toolset(), h.is_available(), h.check_fn(ctx)))
.map(|h| h.schema());
static_schemas.chain(dynamic_schemas).collect()
}
pub async fn dispatch(
&self,
name: &str,
args: serde_json::Value,
ctx: &ToolContext,
) -> Result<String, ToolError> {
const READ_LOOP_EXEMPT: &[&str] = &["read_file", "search_files"];
if !READ_LOOP_EXEMPT.contains(&name) {
crate::read_tracker::notify_other_tool_call(&ctx.session_id);
}
if let Some(handler) = self.tools.get(name) {
if !Self::toolset_allowed_in_ctx(handler.toolset(), ctx) {
return Err(ToolError::Unavailable {
tool: name.to_string(),
reason: format!(
"toolset '{}' is disabled in this session",
handler.toolset()
),
});
}
if !handler.check_fn(ctx) {
return Err(ToolError::Unavailable {
tool: name.to_string(),
reason: "tool gating check failed".into(),
});
}
return handler.execute(args, ctx).await;
}
if let Some(handler) = self.dynamic_tools.get(name) {
if !Self::toolset_allowed_in_ctx(handler.toolset(), ctx) {
return Err(ToolError::Unavailable {
tool: name.to_string(),
reason: format!(
"toolset '{}' is disabled in this session",
handler.toolset()
),
});
}
if !handler.check_fn(ctx) {
return Err(ToolError::Unavailable {
tool: name.to_string(),
reason: "tool gating check failed".into(),
});
}
return handler.execute(args, ctx).await;
}
if let Some(suggestion) = self.fuzzy_match(name) {
Err(ToolError::NotFound(format!(
"Unknown tool '{}'. Did you mean '{}'?",
name, suggestion
)))
} else {
Err(ToolError::NotFound(name.to_string()))
}
}
pub fn register_dynamic(&mut self, handler: Box<dyn ToolHandler>) {
let name = handler.name().to_string();
self.dynamic_tools.insert(name, handler);
}
pub fn tool_names(&self) -> Vec<&str> {
let mut names: Vec<&str> = self.tools.keys().copied().collect();
names.extend(self.dynamic_tools.keys().map(|s| s.as_str()));
names.sort();
names
}
pub fn toolset_names(&self) -> Vec<&str> {
let mut names: Vec<&str> = self.toolset_index.keys().copied().collect();
names.sort();
names
}
pub fn tools_in_toolset(&self, toolset: &str) -> Vec<&str> {
self.toolset_index.get(toolset).cloned().unwrap_or_default()
}
pub fn toolset_for_tool(&self, name: &str) -> Option<String> {
self.tools
.get(name)
.map(|h| h.toolset().to_string())
.or_else(|| {
self.dynamic_tools
.get(name)
.map(|h| h.toolset().to_string())
})
}
pub fn toolset_summary(&self) -> Vec<(String, usize)> {
let mut summary: Vec<(String, usize)> = self
.toolset_index
.iter()
.map(|(name, tools)| (name.to_string(), tools.len()))
.collect();
summary.sort_by(|a, b| a.0.cmp(&b.0));
summary
}
pub fn is_parallel_safe(&self, name: &str) -> bool {
self.tools
.get(name)
.map(|h| h.parallel_safe())
.or_else(|| self.dynamic_tools.get(name).map(|h| h.parallel_safe()))
.unwrap_or(false)
}
fn fuzzy_match(&self, name: &str) -> Option<&str> {
let threshold = 3;
let mut best: Option<(&str, usize)> = None;
for &tool_name in self.tools.keys() {
let dist = strsim::levenshtein(name, tool_name);
if dist <= threshold
&& (best.is_none() || dist < best.as_ref().map_or(usize::MAX, |b| b.1))
{
best = Some((tool_name, dist));
}
}
for tool_name in self.dynamic_tools.keys() {
let dist = strsim::levenshtein(name, tool_name);
if dist <= threshold
&& (best.is_none() || dist < best.as_ref().map_or(usize::MAX, |b| b.1))
{
best = Some((tool_name.as_str(), dist));
}
}
best.map(|(name, _)| name)
}
}
impl Default for ToolRegistry {
fn default() -> Self {
Self::new()
}
}
pub fn to_llm_definitions(schemas: &[ToolSchema]) -> Vec<edgequake_llm::ToolDefinition> {
schemas
.iter()
.map(|s| {
edgequake_llm::ToolDefinition::function(&s.name, &s.description, s.parameters.clone())
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
struct TestTool;
#[async_trait]
impl ToolHandler for TestTool {
fn name(&self) -> &'static str {
"test_tool"
}
fn toolset(&self) -> &'static str {
"test"
}
fn schema(&self) -> ToolSchema {
ToolSchema {
name: "test_tool".into(),
description: "A test tool".into(),
parameters: json!({
"type": "object",
"properties": {
"input": { "type": "string" }
},
"required": ["input"]
}),
strict: None,
}
}
async fn execute(
&self,
args: serde_json::Value,
_ctx: &ToolContext,
) -> Result<String, ToolError> {
let input = args["input"]
.as_str()
.ok_or_else(|| ToolError::InvalidArgs {
tool: "test_tool".into(),
message: "input required".into(),
})?;
Ok(format!("echo: {}", input))
}
fn parallel_safe(&self) -> bool {
true
}
fn emoji(&self) -> &'static str {
"🧪"
}
}
struct GatedTool;
#[async_trait]
impl ToolHandler for GatedTool {
fn name(&self) -> &'static str {
"gated_tool"
}
fn toolset(&self) -> &'static str {
"gated"
}
fn schema(&self) -> ToolSchema {
ToolSchema {
name: "gated_tool".into(),
description: "A gated tool".into(),
parameters: json!({"type": "object"}),
strict: None,
}
}
async fn execute(
&self,
_args: serde_json::Value,
_ctx: &ToolContext,
) -> Result<String, ToolError> {
Ok("gated result".into())
}
fn check_fn(&self, _ctx: &ToolContext) -> bool {
false }
}
fn make_registry_with_tools() -> ToolRegistry {
let mut registry = ToolRegistry {
tools: HashMap::new(),
toolset_index: HashMap::new(),
dynamic_tools: HashMap::new(),
};
let test_tool: &'static dyn ToolHandler = Box::leak(Box::new(TestTool));
let gated_tool: &'static dyn ToolHandler = Box::leak(Box::new(GatedTool));
registry.tools.insert(test_tool.name(), test_tool);
registry.tools.insert(gated_tool.name(), gated_tool);
registry
.toolset_index
.entry("test")
.or_default()
.push("test_tool");
registry
.toolset_index
.entry("gated")
.or_default()
.push("gated_tool");
registry
}
#[test]
fn registry_default_builds() {
let registry = ToolRegistry::new();
assert!(registry.tools.is_empty() || !registry.tools.is_empty());
}
#[tokio::test]
async fn dispatch_exact_match() {
let registry = make_registry_with_tools();
let ctx = ToolContext::test_context();
let result = registry
.dispatch("test_tool", json!({"input": "hello"}), &ctx)
.await;
assert_eq!(result.expect("dispatch"), "echo: hello");
}
#[tokio::test]
async fn dispatch_rejects_disabled_toolset_even_when_tool_exists() {
let registry = make_registry_with_tools();
let mut ctx = ToolContext::test_context();
ctx.config.disabled_toolsets = vec!["test".into()];
let err = registry
.dispatch("test_tool", json!({"input": "hello"}), &ctx)
.await
.expect_err("disabled toolset should block dispatch");
match err {
ToolError::Unavailable { reason, .. } => {
assert!(reason.contains("disabled"));
}
other => panic!("Expected Unavailable, got: {:?}", other),
}
}
#[tokio::test]
async fn dispatch_unknown_tool_fuzzy() {
let registry = make_registry_with_tools();
let ctx = ToolContext::test_context();
let err = registry
.dispatch("test_tol", json!({}), &ctx) .await
.expect_err("should suggest similar tool name");
match err {
ToolError::NotFound(msg) => {
assert!(msg.contains("Did you mean"), "Got: {}", msg);
assert!(msg.contains("test_tool"), "Got: {}", msg);
}
other => panic!("Expected NotFound, got: {:?}", other),
}
}
#[tokio::test]
async fn dispatch_completely_unknown() {
let registry = make_registry_with_tools();
let ctx = ToolContext::test_context();
let err = registry
.dispatch("zzzzzzzzzzzzz", json!({}), &ctx)
.await
.expect_err("completely unknown tool should error");
match err {
ToolError::NotFound(msg) => {
assert!(!msg.contains("Did you mean"), "Got: {}", msg);
}
other => panic!("Expected NotFound, got: {:?}", other),
}
}
#[tokio::test]
async fn dispatch_gated_tool_blocked() {
let registry = make_registry_with_tools();
let ctx = ToolContext::test_context();
let err = registry
.dispatch("gated_tool", json!({}), &ctx)
.await
.expect_err("gated tool should be unavailable");
match err {
ToolError::Unavailable { tool, .. } => assert_eq!(tool, "gated_tool"),
other => panic!("Expected Unavailable, got: {:?}", other),
}
}
#[test]
fn get_definitions_filters_by_toolset() {
let registry = make_registry_with_tools();
let ctx = ToolContext::test_context();
let defs = registry.get_definitions(Some(&["test".to_string()]), None, &ctx);
assert!(defs.iter().any(|d| d.name == "test_tool"));
assert!(!defs.iter().any(|d| d.name == "gated_tool"));
}
#[test]
fn get_definitions_excludes_disabled() {
let registry = make_registry_with_tools();
let ctx = ToolContext::test_context();
let defs = registry.get_definitions(None, Some(&["gated".to_string()]), &ctx);
assert!(!defs.iter().any(|d| d.name == "gated_tool"));
}
#[test]
fn get_definitions_excludes_check_fn_failed() {
let registry = make_registry_with_tools();
let ctx = ToolContext::test_context();
let defs = registry.get_definitions(None, None, &ctx);
assert!(!defs.iter().any(|d| d.name == "gated_tool"));
assert!(defs.iter().any(|d| d.name == "test_tool"));
}
#[test]
fn get_definitions_excludes_ctx_disabled_toolsets() {
let registry = make_registry_with_tools();
let mut ctx = ToolContext::test_context();
ctx.config.disabled_toolsets = vec!["test".into()];
let defs = registry.get_definitions(None, None, &ctx);
assert!(!defs.iter().any(|d| d.name == "test_tool"));
}
#[test]
fn tool_names_sorted() {
let registry = make_registry_with_tools();
let names = registry.tool_names();
let mut sorted = names.clone();
sorted.sort();
assert_eq!(names, sorted);
}
#[test]
fn parallel_safe_check() {
let registry = make_registry_with_tools();
assert!(registry.is_parallel_safe("test_tool"));
assert!(!registry.is_parallel_safe("gated_tool"));
assert!(!registry.is_parallel_safe("nonexistent"));
}
#[test]
fn register_dynamic_tool() {
let mut registry = make_registry_with_tools();
let ctx = ToolContext::test_context();
registry.register_dynamic(Box::new(TestTool));
assert!(registry.tool_names().contains(&"test_tool"));
let defs = registry.get_definitions(None, None, &ctx);
assert!(defs.iter().any(|d| d.name == "test_tool"));
}
#[test]
fn fuzzy_match_close_typo() {
let registry = make_registry_with_tools();
assert_eq!(registry.fuzzy_match("test_tol"), Some("test_tool"));
assert_eq!(registry.fuzzy_match("tset_tool"), Some("test_tool"));
}
#[test]
fn fuzzy_match_too_far() {
let registry = make_registry_with_tools();
assert_eq!(registry.fuzzy_match("completely_different"), None);
}
#[test]
fn to_llm_definitions_conversion() {
let schemas = vec![ToolSchema {
name: "read_file".into(),
description: "Read a file".into(),
parameters: json!({
"type": "object",
"properties": {
"path": { "type": "string" }
},
"required": ["path"]
}),
strict: None,
}];
let defs = to_llm_definitions(&schemas);
assert_eq!(defs.len(), 1);
assert_eq!(defs[0].function.name, "read_file");
}
#[test]
fn toolset_names_and_tools() {
let registry = make_registry_with_tools();
let toolsets = registry.toolset_names();
assert!(toolsets.contains(&"test"));
assert!(toolsets.contains(&"gated"));
let test_tools = registry.tools_in_toolset("test");
assert!(test_tools.contains(&"test_tool"));
}
}