use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use async_trait::async_trait;
use tokio_util::sync::CancellationToken;
use lingshu_state::SessionDb;
use lingshu_types::{Message, Platform, RunOutcome, ToolError, ToolSchema};
use serde_json::{Map, Value};
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 run_outcome: RunOutcome,
pub messages: Vec<Message>,
}
#[derive(Debug, Clone)]
pub enum DelegationEvent {
TaskStarted {
task_index: usize,
task_count: usize,
goal: String,
depth: u32,
agent_id: String,
parent_id: Option<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,
pub agent_id: String,
pub parent_id: Option<String>,
pub depth: u32,
}
#[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 aliases(&self) -> &'static [&'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 path_arguments(&self) -> &'static [&'static str] {
&[]
}
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 delegate_agent_id: Option<String>,
pub delegate_parent_id: Option<String>,
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<lingshu_types::OriginChat>,
pub session_key: Option<String>,
pub todo_store: Option<Arc<crate::tools::todo::TodoStore>>,
pub current_tool_call_id: Option<String>,
pub current_tool_name: Option<String>,
pub injected_messages: Option<Arc<tokio::sync::Mutex<Vec<Message>>>>,
pub tool_progress_tx: Option<tokio::sync::mpsc::UnboundedSender<ToolProgressUpdate>>,
pub watch_notification_tx:
Option<tokio::sync::mpsc::UnboundedSender<crate::process_table::WatchEvent>>,
pub mutation_turn: Option<Arc<crate::mutations::MutationTurnState>>,
pub lsp_gate: Option<Arc<dyn crate::lsp_gate::LspGate>>,
pub kanban_task_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolProgressUpdate {
pub tool_call_id: String,
pub tool_name: String,
pub message: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolInventoryEntry {
pub name: String,
pub toolset: String,
pub description: String,
pub emoji: String,
pub aliases: Vec<String>,
pub dynamic: bool,
pub policy_enabled: bool,
pub startup_available: bool,
pub check_allowed: bool,
}
impl ToolInventoryEntry {
pub fn exposed(&self) -> bool {
self.policy_enabled && self.startup_available && self.check_allowed
}
}
#[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 {
pub fn record_mutation(&self, record: crate::mutations::MutationRecord) {
if let Some(state) = &self.mutation_turn {
state.push_success(record);
}
}
#[cfg(test)]
pub fn test_context() -> Self {
crate::delegation_state::set_spawn_paused(false);
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,
delegate_agent_id: None,
delegate_parent_id: None,
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,
current_tool_call_id: None,
current_tool_name: None,
injected_messages: None,
tool_progress_tx: None,
watch_notification_tx: None,
mutation_turn: None,
lsp_gate: None,
kanban_task_id: None,
}
}
pub fn emit_progress(&self, message: impl AsRef<str>) {
crate::tool_progress_tail::emit_tool_progress(self, message.as_ref());
}
}
pub struct ToolRegistry {
tools: HashMap<&'static str, &'static dyn ToolHandler>,
tool_aliases: HashMap<&'static str, &'static str>,
toolset_index: HashMap<&'static str, Vec<&'static str>>,
dynamic_tools: HashMap<String, Box<dyn ToolHandler>>,
dynamic_tool_aliases: HashMap<String, String>,
}
const PATH_FIELD_ALIASES: &[&str] = &["file_path", "filepath", "filename"];
const CONTENT_FIELD_ALIASES: &[&str] = &["file_content", "text", "body"];
fn rename_missing_canonical_field(
map: &mut serde_json::Map<String, serde_json::Value>,
canonical: &str,
aliases: &[&str],
) {
if map.contains_key(canonical) {
return;
}
for alias in aliases {
if let Some(value) = map.remove(*alias) {
map.insert(canonical.to_string(), value);
return;
}
}
}
fn normalize_schema_field_aliases(args: &mut serde_json::Value, schema: &serde_json::Value) {
let Some(properties) = schema
.get("properties")
.and_then(serde_json::Value::as_object)
else {
return;
};
let serde_json::Value::Object(map) = args else {
return;
};
if properties.contains_key("path") {
rename_missing_canonical_field(map, "path", PATH_FIELD_ALIASES);
}
if properties.contains_key("content") {
rename_missing_canonical_field(map, "content", CONTENT_FIELD_ALIASES);
}
}
fn schema_allows_null(prop_schema: &serde_json::Value) -> bool {
match prop_schema.get("type") {
Some(serde_json::Value::String(t)) if t == "null" => true,
Some(serde_json::Value::Array(types)) if types.iter().any(|t| t.as_str() == Some("null")) => {
true
}
_ => prop_schema.get("nullable").and_then(|v| v.as_bool()) == Some(true),
}
}
fn coerce_string_value(value: &str, expected_type: &str, prop_schema: &serde_json::Value) -> Option<serde_json::Value> {
if schema_allows_null(prop_schema) && value.trim().eq_ignore_ascii_case("null") {
return Some(serde_json::Value::Null);
}
match expected_type {
"integer" => value
.parse::<i64>()
.ok()
.map(|n| serde_json::Value::Number(n.into())),
"number" => value
.parse::<f64>()
.ok()
.and_then(serde_json::Number::from_f64)
.map(serde_json::Value::Number),
"boolean" => match value {
"true" | "1" | "True" => Some(serde_json::Value::Bool(true)),
"false" | "0" | "False" => Some(serde_json::Value::Bool(false)),
_ => None,
},
"array" => serde_json::from_str::<Vec<serde_json::Value>>(value)
.ok()
.map(serde_json::Value::Array),
"object" => serde_json::from_str::<serde_json::Map<String, serde_json::Value>>(value)
.ok()
.map(serde_json::Value::Object),
_ => None,
}
}
fn coerce_value_for_schema(value: &mut serde_json::Value, prop_schema: &serde_json::Value) {
if let Some(serde_json::Value::Array(types)) = prop_schema.get("type")
&& let Some(s) = value.as_str()
{
for t in types {
if let Some(t) = t.as_str()
&& let Some(coerced) = coerce_string_value(s, t, prop_schema)
{
*value = coerced;
return;
}
}
return;
}
let Some(expected_type) = prop_schema.get("type").and_then(|t| t.as_str()) else {
return;
};
if let Some(s) = value.as_str()
&& let Some(coerced) = coerce_string_value(s, expected_type, prop_schema)
{
*value = coerced;
return;
}
match expected_type {
"string" if !value.is_string() => {
*value = serde_json::Value::String(value.to_string());
}
"array" if !value.is_array() && value.as_str().is_none() => {
let v = value.take();
*value = serde_json::Value::Array(vec![v]);
}
_ => {}
}
}
fn coerce_tool_args(args: &mut serde_json::Value, schema: &serde_json::Value) {
normalize_schema_field_aliases(args, schema);
let Some(properties) = schema
.get("properties")
.and_then(serde_json::Value::as_object)
else {
return;
};
let serde_json::Value::Object(map) = args else {
return;
};
for (key, prop_schema) in properties {
let Some(value) = map.get_mut(key) else {
continue;
};
coerce_value_for_schema(value, prop_schema);
}
}
const TOOL_CROSS_REFS: &[(&str, &[&str], &str)] = &[
(
"browser_navigate",
&["web_search", "web_extract"],
"For simple info retrieval, prefer web_search or web_extract (faster, cheaper). ",
),
];
pub fn strip_unavailable_cross_refs(definitions: &mut [ToolSchema]) {
use std::collections::HashSet;
let available: HashSet<String> = definitions.iter().map(|s| s.name.clone()).collect();
for schema in definitions.iter_mut() {
for &(tool_name, required_tools, text_to_strip) in TOOL_CROSS_REFS {
let none_available = required_tools.iter().all(|&t| !available.contains(t));
if schema.name == tool_name && none_available {
schema.description = schema.description.replace(text_to_strip, "");
}
}
}
}
impl ToolRegistry {
fn tool_allowed_in_ctx(tool_name: &str, toolset: &str, ctx: &ToolContext) -> bool {
ctx.config.is_tool_enabled(tool_name, toolset)
}
pub fn new() -> Self {
let mut tools = HashMap::new();
let mut tool_aliases = HashMap::new();
let mut toolset_index: HashMap<&str, Vec<&str>> = HashMap::new();
for handler in inventory::iter::<&dyn ToolHandler> {
for &alias in handler.aliases() {
if alias == handler.name() {
continue;
}
assert!(
!tools.contains_key(alias) && !tool_aliases.contains_key(alias),
"duplicate tool alias registered: {alias}"
);
tool_aliases.insert(alias, handler.name());
}
tools.insert(handler.name(), *handler);
toolset_index
.entry(handler.toolset())
.or_default()
.push(handler.name());
}
Self {
tools,
tool_aliases,
toolset_index,
dynamic_tools: HashMap::new(),
dynamic_tool_aliases: HashMap::new(),
}
}
pub fn get_definitions(
&self,
enabled: Option<&[String]>,
disabled: Option<&[String]>,
ctx: &ToolContext,
) -> Vec<ToolSchema> {
let trace_schema_resolution = std::env::var("EDGECRAB_TRACE_SCHEMA_RESOLUTION")
.ok()
.map(|value| {
matches!(
value.trim().to_ascii_lowercase().as_str(),
"1" | "true" | "yes" | "on"
)
})
.unwrap_or(false);
let is_eligible =
|tool_name: &str, toolset: &str, available: bool, passes_check: bool| -> bool {
let explicitly_enabled = ctx
.config
.enabled_tools
.iter()
.any(|candidate| candidate == tool_name);
let explicitly_disabled = ctx
.config
.disabled_tools
.iter()
.any(|candidate| candidate == tool_name);
if explicitly_disabled {
return false;
}
let toolset_allowed = explicitly_enabled
|| (Self::tool_allowed_in_ctx(tool_name, toolset, ctx)
&& enabled.is_none_or(|sets| sets.iter().any(|s| s == toolset))
&& disabled.is_none_or(|sets| !sets.iter().any(|s| s == toolset)));
available && passes_check && toolset_allowed
};
let mut schemas = Vec::new();
for handler in self.tools.values() {
let check_started = std::time::Instant::now();
if trace_schema_resolution {
tracing::info!(tool = handler.name(), "schema-resolution: checking tool");
}
let available = handler.is_available();
let passes_check = handler.check_fn(ctx);
let elapsed_ms = check_started.elapsed().as_millis() as u64;
if trace_schema_resolution || elapsed_ms >= 100 {
tracing::info!(
tool = handler.name(),
available,
passes_check,
elapsed_ms,
"schema-resolution: tool check completed"
);
}
if is_eligible(handler.name(), handler.toolset(), available, passes_check) {
schemas.push(handler.schema());
}
}
for handler in self.dynamic_tools.values() {
let check_started = std::time::Instant::now();
if trace_schema_resolution {
tracing::info!(
tool = handler.name(),
"schema-resolution: checking dynamic tool"
);
}
let available = handler.is_available();
let passes_check = handler.check_fn(ctx);
let elapsed_ms = check_started.elapsed().as_millis() as u64;
if trace_schema_resolution || elapsed_ms >= 100 {
tracing::info!(
tool = handler.name(),
available,
passes_check,
elapsed_ms,
"schema-resolution: dynamic tool check completed"
);
}
if is_eligible(handler.name(), handler.toolset(), available, passes_check) {
schemas.push(handler.schema());
}
}
strip_unavailable_cross_refs(&mut schemas);
schemas
}
pub fn prepare_tool_arguments(&self, name: &str, args: &mut serde_json::Value) {
let static_name = self.tool_aliases.get(name).copied().unwrap_or(name);
if let Some(handler) = self.tools.get(static_name) {
coerce_tool_args(args, &handler.schema().parameters);
return;
}
let dynamic_name = self
.dynamic_tool_aliases
.get(name)
.map(String::as_str)
.unwrap_or(name);
if let Some(handler) = self.dynamic_tools.get(dynamic_name) {
coerce_tool_args(args, &handler.schema().parameters);
}
}
pub async fn dispatch(
&self,
name: &str,
mut args: serde_json::Value,
ctx: &ToolContext,
) -> Result<String, ToolError> {
let resolved = self.resolve_tool_call_name(name);
let name = if resolved.canonical.is_empty() {
name
} else {
resolved.canonical.as_str()
};
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);
}
let static_name = self.tool_aliases.get(name).copied().unwrap_or(name);
if let Some(handler) = self.tools.get(static_name) {
if !Self::tool_allowed_in_ctx(handler.name(), handler.toolset(), ctx) {
return Err(ToolError::Unavailable {
tool: name.to_string(),
reason: format!(
"tool '{}' is disabled in this session policy",
handler.name()
),
});
}
if !handler.check_fn(ctx) {
return Err(ToolError::Unavailable {
tool: name.to_string(),
reason: "tool gating check failed".into(),
});
}
coerce_tool_args(&mut args, &handler.schema().parameters);
return handler.execute(args, ctx).await;
}
let dynamic_name = self
.dynamic_tool_aliases
.get(name)
.map(String::as_str)
.unwrap_or(name);
if let Some(handler) = self.dynamic_tools.get(dynamic_name) {
if !Self::tool_allowed_in_ctx(handler.name(), handler.toolset(), ctx) {
return Err(ToolError::Unavailable {
tool: name.to_string(),
reason: format!(
"tool '{}' is disabled in this session policy",
handler.name()
),
});
}
if !handler.check_fn(ctx) {
return Err(ToolError::Unavailable {
tool: name.to_string(),
reason: "tool gating check failed".into(),
});
}
coerce_tool_args(&mut args, &handler.schema().parameters);
return handler.execute(args, ctx).await;
}
if let Some(suggestion) = crate::tool_name_repair::fuzzy_match_tool_name(self, name) {
Err(ToolError::NotFound(format!(
"Unknown tool '{}'. Did you mean '{}'?",
name, suggestion
)))
} else {
Err(ToolError::NotFound(name.to_string()))
}
}
pub fn enrich_invalid_args_error(
&self,
tool_name: &str,
error: &ToolError,
args_json: Option<&str>,
) -> Option<lingshu_types::ToolErrorResponse> {
if !matches!(error.core_error(), ToolError::InvalidArgs { .. }) {
return None;
}
let handler: &dyn ToolHandler = if let Some(h) = self.tools.get(tool_name) {
*h
} else {
let h = self.dynamic_tools.get(tool_name)?;
h.as_ref()
};
let schema = handler.schema();
let required_fields =
required_fields_from_parameters(&schema.parameters, args_json);
let usage_hint = Self::build_usage_hint(&schema);
Some(error.to_llm_payload_enriched(required_fields, usage_hint))
}
fn build_usage_hint(schema: &lingshu_types::ToolSchema) -> Option<String> {
let props = schema.parameters.get("properties")?.as_object()?;
let required: Vec<&str> = schema
.parameters
.get("required")
.and_then(|v| v.as_array())
.map(|arr| arr.iter().filter_map(|v| v.as_str()).collect())
.unwrap_or_default();
let mut parts = Vec::new();
for name in &required {
if let Some(prop) = props.get(*name) {
let ty = prop
.get("type")
.map(|t| {
if let Some(s) = t.as_str() {
s.to_string()
} else {
t.to_string()
}
})
.unwrap_or_else(|| "any".to_string());
parts.push(format!("{name}: {ty}"));
}
}
if parts.is_empty() {
return None;
}
Some(format!("Required: {}", parts.join(", ")))
}
pub fn register_dynamic(&mut self, handler: Box<dyn ToolHandler>) {
let name = handler.name().to_string();
for &alias in handler.aliases() {
if alias != handler.name() {
self.dynamic_tool_aliases
.insert(alias.to_string(), name.clone());
}
}
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 lookup_tool_name(&self, name: &str) -> Option<String> {
if self.tools.contains_key(name) {
return Some(name.to_string());
}
if let Some(&canonical) = self.tool_aliases.get(name) {
return Some(canonical.to_string());
}
if self.dynamic_tools.contains_key(name) {
return Some(name.to_string());
}
self.dynamic_tool_aliases
.get(name)
.cloned()
}
pub fn resolve_tool_call_name(&self, raw: &str) -> crate::tool_name_repair::ResolvedToolName {
crate::tool_name_repair::resolve_tool_call_name(self, raw)
}
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(self.tool_aliases.get(name).copied().unwrap_or(name))
.map(|h| h.toolset().to_string())
.or_else(|| {
self.dynamic_tools
.get(
self.dynamic_tool_aliases
.get(name)
.map(String::as_str)
.unwrap_or(name),
)
.map(|h| h.toolset().to_string())
})
}
pub fn required_fields_for_tool(&self, name: &str) -> Option<Vec<String>> {
let static_name = self.tool_aliases.get(name).copied().unwrap_or(name);
let schema = if let Some(handler) = self.tools.get(static_name) {
handler.schema()
} else {
let dynamic_name = self
.dynamic_tool_aliases
.get(name)
.map(String::as_str)
.unwrap_or(name);
let handler = self.dynamic_tools.get(dynamic_name)?;
handler.schema()
};
Some(
required_fields_from_parameters(&schema.parameters, None).unwrap_or_default(),
)
}
pub fn toolset_summary(&self) -> Vec<(String, usize)> {
let mut counts: std::collections::BTreeMap<String, usize> = self
.toolset_index
.iter()
.map(|(name, tools)| (name.to_string(), tools.len()))
.collect();
for handler in self.dynamic_tools.values() {
*counts.entry(handler.toolset().to_string()).or_default() += 1;
}
let mut summary: Vec<(String, usize)> = counts.into_iter().collect();
summary.sort_by(|a, b| a.0.cmp(&b.0));
summary
}
pub fn tool_inventory(&self, ctx: &ToolContext) -> Vec<ToolInventoryEntry> {
let mut entries: Vec<ToolInventoryEntry> = self
.tools
.values()
.map(|handler| ToolInventoryEntry {
name: handler.name().to_string(),
toolset: handler.toolset().to_string(),
description: handler.schema().description,
emoji: handler.emoji().to_string(),
aliases: handler
.aliases()
.iter()
.map(|alias| (*alias).to_string())
.collect(),
dynamic: false,
policy_enabled: ctx
.config
.is_tool_enabled(handler.name(), handler.toolset()),
startup_available: handler.is_available(),
check_allowed: handler.check_fn(ctx),
})
.collect();
entries.extend(self.dynamic_tools.values().map(|handler| {
ToolInventoryEntry {
name: handler.name().to_string(),
toolset: handler.toolset().to_string(),
description: handler.schema().description,
emoji: handler.emoji().to_string(),
aliases: handler
.aliases()
.iter()
.map(|alias| (*alias).to_string())
.collect(),
dynamic: true,
policy_enabled: ctx
.config
.is_tool_enabled(handler.name(), handler.toolset()),
startup_available: handler.is_available(),
check_allowed: handler.check_fn(ctx),
}
}));
entries.sort_by(|left, right| left.name.cmp(&right.name));
entries
}
pub fn is_parallel_safe(&self, name: &str) -> bool {
self.tools
.get(self.tool_aliases.get(name).copied().unwrap_or(name))
.map(|h| h.parallel_safe())
.or_else(|| {
self.dynamic_tools
.get(
self.dynamic_tool_aliases
.get(name)
.map(String::as_str)
.unwrap_or(name),
)
.map(|h| h.parallel_safe())
})
.unwrap_or(false)
}
pub fn path_arguments(&self, name: &str) -> &'static [&'static str] {
self.tools
.get(self.tool_aliases.get(name).copied().unwrap_or(name))
.map(|h| h.path_arguments())
.or_else(|| {
self.dynamic_tools
.get(
self.dynamic_tool_aliases
.get(name)
.map(String::as_str)
.unwrap_or(name),
)
.map(|h| h.path_arguments())
})
.unwrap_or(&[])
}
pub fn can_parallelize_in_batch(
&self,
name: &str,
args_json: &str,
claimed_paths: &std::collections::HashSet<String>,
) -> bool {
if self.is_parallel_safe(name) {
return true;
}
let path_args = self.path_arguments(name);
if path_args.is_empty() {
return false; }
let Ok(args) = serde_json::from_str::<serde_json::Value>(args_json) else {
return false; };
for pa in path_args {
if let Some(path) = args.get(*pa).and_then(serde_json::Value::as_str) {
if paths_overlap(path, claimed_paths) {
return false; }
}
}
true
}
pub fn extract_paths(&self, name: &str, args_json: &str) -> Vec<String> {
let path_args = self.path_arguments(name);
if path_args.is_empty() {
return Vec::new();
}
let Ok(args) = serde_json::from_str::<serde_json::Value>(args_json) else {
return Vec::new();
};
path_args
.iter()
.filter_map(|pa| {
args.get(*pa)
.and_then(serde_json::Value::as_str)
.map(String::from)
})
.collect()
}
}
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,
openai_compatible_tool_parameters(&s.parameters),
)
})
.collect()
}
pub fn openai_compatible_tool_parameters(schema: &Value) -> Value {
if let Some(one_of) = schema.get("oneOf").and_then(Value::as_array) {
return flatten_oneof_object_schema(one_of);
}
normalize_json_schema(schema)
}
fn flatten_oneof_object_schema(branches: &[Value]) -> Value {
let mut properties = Map::new();
for branch in branches {
if let Some(props) = branch.get("properties").and_then(Value::as_object) {
for (key, value) in props {
properties.entry(key.clone()).or_insert_with(|| value.clone());
}
}
}
Value::Object(Map::from_iter([
("type".into(), Value::String("object".into())),
("additionalProperties".into(), Value::Bool(false)),
("properties".into(), Value::Object(properties)),
(
"required".into(),
Value::Array(vec![Value::String("mode".into())]),
),
]))
}
pub fn annotate_llm_definitions_for_local_turn(
defs: Vec<edgequake_llm::ToolDefinition>,
provider_name: &str,
local_turn_budget: Option<(usize, usize)>,
) -> Vec<edgequake_llm::ToolDefinition> {
let Some((max_arg_bytes, max_output_tokens)) = local_turn_budget else {
return defs;
};
if !matches!(provider_name, "lmstudio" | "ollama") {
return defs;
}
let suffix =
crate::mutation_turn_policy::local_tool_budget_schema_suffix(max_arg_bytes, max_output_tokens);
defs.into_iter()
.map(|mut def| {
if crate::mutation_turn_policy::is_large_payload_tool(&def.function.name) {
def.function.description.push_str(&suffix);
}
def
})
.collect()
}
fn patch_mode_required_fields(mode: &str) -> Vec<String> {
match mode {
"patch" => vec![
"mode".into(),
"patch".into(),
],
_ => vec![
"mode".into(),
"path".into(),
"old_string".into(),
"new_string".into(),
],
}
}
fn is_patch_flat_schema(parameters: &Value) -> bool {
parameters.get("type").and_then(Value::as_str) == Some("object")
&& parameters
.get("properties")
.and_then(|p| p.get("mode"))
.and_then(|m| m.get("enum"))
.is_some()
&& parameters
.get("properties")
.and_then(|p| p.get("patch"))
.is_some()
&& parameters
.get("properties")
.and_then(|p| p.get("old_string"))
.is_some()
}
fn required_fields_from_parameters(
parameters: &Value,
args_json: Option<&str>,
) -> Option<Vec<String>> {
if let Some(one_of) = parameters.get("oneOf").and_then(Value::as_array) {
let mode = args_json
.and_then(|raw| serde_json::from_str::<Value>(raw).ok())
.and_then(|v| v.get("mode").and_then(Value::as_str).map(String::from))
.unwrap_or_else(|| "replace".to_string());
for branch in one_of {
let branch_mode = branch
.get("properties")
.and_then(|p| p.get("mode"))
.and_then(|m| m.get("const"))
.and_then(Value::as_str);
if branch_mode == Some(mode.as_str()) {
return branch
.get("required")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect()
});
}
}
return one_of.first().and_then(|branch| {
branch
.get("required")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect()
})
});
}
if is_patch_flat_schema(parameters) {
let mode = args_json
.and_then(|raw| serde_json::from_str::<Value>(raw).ok())
.and_then(|v| v.get("mode").and_then(Value::as_str).map(String::from))
.unwrap_or_else(|| "replace".to_string());
return Some(patch_mode_required_fields(&mode));
}
parameters
.get("required")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect()
})
}
pub fn normalize_json_schema(schema: &Value) -> Value {
match schema {
Value::Object(object) => {
let mut normalized = Map::new();
for (key, value) in object {
normalized.insert(key.clone(), normalize_json_schema(value));
}
if normalized.get("type").and_then(Value::as_str) == Some("object") {
let needs_properties =
!matches!(normalized.get("properties"), Some(Value::Object(_)));
if needs_properties {
normalized.insert("properties".into(), Value::Object(Map::new()));
}
}
Value::Object(normalized)
}
Value::Array(items) => Value::Array(items.iter().map(normalize_json_schema).collect()),
_ => schema.clone(),
}
}
fn paths_overlap(path: &str, claimed_paths: &std::collections::HashSet<String>) -> bool {
let normalize = |p: &str| -> String {
let trimmed = p.trim_end_matches('/').trim_end_matches('\\');
if trimmed.is_empty() {
p.to_string()
} else {
trimmed.to_string()
}
};
let norm_path = normalize(path);
for claimed in claimed_paths {
let norm_claimed = normalize(claimed);
if norm_path == norm_claimed {
return true;
}
if norm_path.starts_with(&norm_claimed)
&& norm_path.as_bytes().get(norm_claimed.len()) == Some(&b'/')
{
return true;
}
if norm_claimed.starts_with(&norm_path)
&& norm_claimed.as_bytes().get(norm_path.len()) == Some(&b'/')
{
return true;
}
}
false
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn assert_provider_safe_top_level_schema(tool_name: &str, schema: &Value) {
assert_eq!(
schema.get("type").and_then(Value::as_str),
Some("object"),
"tool '{tool_name}' must export a top-level object schema"
);
assert!(
matches!(schema.get("properties"), Some(Value::Object(_))),
"tool '{tool_name}' must export an object-valued top-level properties map"
);
for forbidden in ["anyOf", "oneOf", "allOf", "not", "if", "then", "else"] {
assert!(
schema.get(forbidden).is_none(),
"tool '{tool_name}' must not use top-level {forbidden}"
);
}
}
struct TestTool;
#[async_trait]
impl ToolHandler for TestTool {
fn name(&self) -> &'static str {
"test_tool"
}
fn aliases(&self) -> &'static [&'static str] {
&["legacy_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(),
tool_aliases: HashMap::new(),
toolset_index: HashMap::new(),
dynamic_tools: HashMap::new(),
dynamic_tool_aliases: 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
.tool_aliases
.insert("legacy_test_tool", "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_accepts_legacy_alias() {
let registry = make_registry_with_tools();
let ctx = ToolContext::test_context();
let result = registry
.dispatch("legacy_test_tool", json!({"input": "hello"}), &ctx)
.await;
assert_eq!(result.expect("alias dispatch"), "echo: hello");
}
#[tokio::test]
async fn dispatch_unknown_tool_fuzzy() {
let registry = make_registry_with_tools();
let ctx = ToolContext::test_context();
let result = registry
.dispatch("test_tol", json!({"input": "hello"}), &ctx)
.await
.expect("fuzzy name repair should dispatch");
assert_eq!(result, "echo: hello");
}
#[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 get_definitions_includes_explicitly_enabled_tool_from_blocked_toolset() {
let registry = make_registry_with_tools();
let mut ctx = ToolContext::test_context();
ctx.config.parent_active_toolsets = vec!["other".into()];
ctx.config.enabled_tools = vec!["test_tool".into()];
let defs = registry.get_definitions(Some(&["other".to_string()]), None, &ctx);
assert!(defs.iter().any(|d| d.name == "test_tool"));
}
#[tokio::test]
async fn dispatch_rejects_explicitly_disabled_tool() {
let registry = make_registry_with_tools();
let mut ctx = ToolContext::test_context();
ctx.config.enabled_tools = vec!["test_tool".into()];
ctx.config.disabled_tools = vec!["test_tool".into()];
let err = registry
.dispatch("test_tool", json!({"input": "hello"}), &ctx)
.await
.expect_err("disabled tool should be blocked");
match err {
ToolError::Unavailable { reason, .. } => {
assert!(reason.contains("disabled"));
}
other => panic!("Expected Unavailable, got: {:?}", other),
}
}
#[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("legacy_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!(
crate::tool_name_repair::fuzzy_match_tool_name(®istry, "test_tol"),
Some("test_tool".into())
);
assert_eq!(
crate::tool_name_repair::fuzzy_match_tool_name(®istry, "tset_tool"),
Some("test_tool".into())
);
}
#[test]
fn fuzzy_match_too_far() {
let registry = make_registry_with_tools();
assert_eq!(
crate::tool_name_repair::fuzzy_match_tool_name(®istry, "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 normalize_json_schema_inserts_missing_object_properties_recursively() {
let normalized = normalize_json_schema(&json!({
"type": "object",
"additionalProperties": true,
"properties": {
"options": {
"type": "object",
"additionalProperties": true
},
"entries": {
"type": "array",
"items": {
"type": "object"
}
}
}
}));
assert_eq!(normalized["type"], "object");
assert!(normalized["properties"].is_object());
assert!(normalized["properties"]["options"]["properties"].is_object());
assert!(normalized["properties"]["entries"]["items"]["properties"].is_object());
}
#[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"));
}
#[test]
fn toolset_for_alias_resolves_to_canonical_toolset() {
let registry = make_registry_with_tools();
assert_eq!(
registry.toolset_for_tool("legacy_test_tool").as_deref(),
Some("test")
);
}
#[test]
fn registry_inventory_matches_get_definitions_for_all_tools() {
let registry = ToolRegistry::new();
let ctx = ToolContext::test_context();
let mut inventory_names: Vec<String> = registry
.tool_inventory(&ctx)
.into_iter()
.filter(|entry| entry.exposed())
.map(|entry| entry.name)
.collect();
inventory_names.sort();
let mut definition_names: Vec<String> = registry
.get_definitions(None, None, &ctx)
.into_iter()
.map(|schema| schema.name)
.collect();
definition_names.sort();
assert_eq!(
definition_names, inventory_names,
"get_definitions(None, None, ctx) should match the exposed tool inventory; `all` must not inject hidden tool contracts"
);
}
#[test]
fn all_exported_tool_definitions_have_provider_safe_top_level_schemas() {
let registry = ToolRegistry::new();
let ctx = ToolContext::test_context();
let schemas = registry.get_definitions(None, None, &ctx);
assert!(
!schemas.is_empty(),
"expected the registry to expose at least one tool definition"
);
for schema in &schemas {
assert!(
!schema.description.trim().is_empty(),
"tool '{}' must have a non-empty description",
schema.name
);
}
let llm_defs = to_llm_definitions(&schemas);
for definition in &llm_defs {
assert_provider_safe_top_level_schema(
&definition.function.name,
&definition.function.parameters,
);
}
}
#[test]
fn enrich_invalid_args_adds_required_fields_and_hint() {
let registry = ToolRegistry::new();
let err = ToolError::InvalidArgs {
tool: "read_file".into(),
message: "missing field `path`".into(),
};
let enriched = registry.enrich_invalid_args_error("read_file", &err, None);
assert!(
enriched.is_some(),
"should return enriched error for InvalidArgs"
);
let resp = enriched.unwrap();
assert!(
resp.required_fields.is_some(),
"should have required_fields"
);
let rf = resp.required_fields.unwrap();
assert!(
rf.contains(&"path".to_string()),
"read_file requires 'path'"
);
assert!(resp.usage_hint.is_some(), "should have a usage hint");
assert!(resp.usage_hint.unwrap().contains("path"));
}
#[test]
fn enrich_returns_none_for_non_invalid_args() {
let registry = ToolRegistry::new();
let err = ToolError::NotFound("read_file".into());
assert!(
registry
.enrich_invalid_args_error("read_file", &err, None)
.is_none()
);
}
#[test]
fn lh61_annotate_llm_definitions_appends_budget_suffix_for_local_mutation_tools() {
let defs = vec![edgequake_llm::ToolDefinition::function(
"write_file",
"Write a file.",
serde_json::json!({}),
)];
let out = annotate_llm_definitions_for_local_turn(defs, "lmstudio", Some((27_852, 8192)));
assert!(out[0].function.description.contains("27852"));
assert!(out[0].function.description.contains("8192"));
let unchanged =
annotate_llm_definitions_for_local_turn(out.clone(), "anthropic", Some((27_852, 8192)));
assert_eq!(unchanged[0].function.description, out[0].function.description);
let no_budget = annotate_llm_definitions_for_local_turn(
vec![edgequake_llm::ToolDefinition::function(
"write_file",
"Write a file.",
serde_json::json!({}),
)],
"lmstudio",
None,
);
assert!(!no_budget[0].function.description.contains("Local turn limit"));
}
#[test]
fn lh62_required_fields_from_patch_flat_schema_respects_mode() {
let params = crate::tools::file_patch::patch_tool_parameters_json();
let replace_args = r#"{"mode":"replace","path":"a.md"}"#;
let fields =
required_fields_from_parameters(¶ms, Some(replace_args)).expect("fields");
assert!(fields.contains(&"path".to_string()));
assert!(fields.contains(&"old_string".to_string()));
assert!(!fields.contains(&"patch".to_string()));
let patch_args = r#"{"mode":"patch"}"#;
let patch_fields =
required_fields_from_parameters(¶ms, Some(patch_args)).expect("patch fields");
assert_eq!(patch_fields, vec!["mode".to_string(), "patch".to_string()]);
}
#[test]
fn lh64_patch_schema_passes_openai_compatible_provider_safe_gate() {
let params = crate::tools::file_patch::patch_tool_parameters_json();
let wire = openai_compatible_tool_parameters(¶ms);
assert_provider_safe_top_level_schema("patch", &wire);
}
#[test]
fn openai_compatible_tool_parameters_flattens_top_level_oneof() {
let oneof = json!({
"oneOf": [
{
"type": "object",
"properties": {
"mode": { "const": "a" },
"foo": { "type": "string" }
}
},
{
"type": "object",
"properties": {
"mode": { "const": "b" },
"bar": { "type": "string" }
}
}
]
});
let flat = openai_compatible_tool_parameters(&oneof);
assert_eq!(flat["type"], "object");
assert!(flat.get("oneOf").is_none());
assert!(flat["properties"].get("foo").is_some());
assert!(flat["properties"].get("bar").is_some());
}
#[test]
fn coerce_string_to_integer() {
let schema = serde_json::json!({
"properties": { "line": { "type": "integer" } }
});
let mut args = serde_json::json!({"line": "42"});
coerce_tool_args(&mut args, &schema);
assert_eq!(args["line"], 42);
}
#[test]
fn coerce_union_integer_from_string() {
let schema = serde_json::json!({
"properties": { "line": { "type": ["integer", "string"] } }
});
let mut args = serde_json::json!({"line": "42"});
coerce_tool_args(&mut args, &schema);
assert_eq!(args["line"], 42);
}
#[test]
fn coerce_null_string_when_nullable() {
let schema = serde_json::json!({
"properties": { "limit": { "type": ["integer", "null"], "nullable": true } }
});
let mut args = serde_json::json!({"limit": "null"});
coerce_tool_args(&mut args, &schema);
assert!(args["limit"].is_null());
}
#[test]
fn prepare_tool_arguments_renames_file_path_before_dispatch() {
let registry = ToolRegistry::new();
let mut args = serde_json::json!({
"file_path": "notes.md",
"text": "hello"
});
registry.prepare_tool_arguments("write_file", &mut args);
assert_eq!(args["path"], "notes.md");
assert_eq!(args["content"], "hello");
}
#[test]
fn coerce_string_to_boolean() {
let schema = serde_json::json!({
"properties": { "flag": { "type": "boolean" } }
});
let mut args = serde_json::json!({"flag": "true"});
coerce_tool_args(&mut args, &schema);
assert_eq!(args["flag"], true);
let mut args2 = serde_json::json!({"flag": "0"});
coerce_tool_args(&mut args2, &schema);
assert_eq!(args2["flag"], false);
}
#[test]
fn coerce_string_to_number() {
let schema = serde_json::json!({
"properties": { "val": { "type": "number" } }
});
let mut args = serde_json::json!({"val": "3.14"});
coerce_tool_args(&mut args, &schema);
let expected = 314_f64 / 100.0;
assert!((args["val"].as_f64().unwrap() - expected).abs() < f64::EPSILON);
}
#[test]
fn coerce_number_to_string() {
let schema = serde_json::json!({
"properties": { "name": { "type": "string" } }
});
let mut args = serde_json::json!({"name": 42});
coerce_tool_args(&mut args, &schema);
assert_eq!(args["name"], "42");
}
#[test]
fn coerce_single_value_to_array() {
let schema = serde_json::json!({
"properties": { "tags": { "type": "array" } }
});
let mut args = serde_json::json!({"tags": "foo"});
coerce_tool_args(&mut args, &schema);
assert_eq!(args["tags"], serde_json::json!(["foo"]));
}
#[test]
fn coerce_already_correct_type_unchanged() {
let schema = serde_json::json!({
"properties": { "line": { "type": "integer" } }
});
let mut args = serde_json::json!({"line": 42});
coerce_tool_args(&mut args, &schema);
assert_eq!(args["line"], 42);
}
#[test]
fn coerce_unparseable_string_left_unchanged() {
let schema = serde_json::json!({
"properties": { "line": { "type": "integer" } }
});
let mut args = serde_json::json!({"line": "not_a_number"});
let original = args.clone();
coerce_tool_args(&mut args, &schema);
assert_eq!(args, original);
}
#[test]
fn extract_paths_returns_matching_args() {
let reg = ToolRegistry::new();
let paths = reg.extract_paths("write_file", r#"{"path":"foo.rs","content":"x"}"#);
assert!(paths.is_empty() || paths.contains(&"foo.rs".to_string()));
}
#[test]
fn can_parallelize_no_overlap() {
let reg = ToolRegistry::new();
let claimed: std::collections::HashSet<String> = std::collections::HashSet::new();
let result = reg.can_parallelize_in_batch("unknown_tool", "{}", &claimed);
assert!(!result, "unknown tool should not be parallelizable");
}
#[test]
fn can_parallelize_bad_json() {
let reg = ToolRegistry::new();
let claimed: std::collections::HashSet<String> = std::collections::HashSet::new();
let result = reg.can_parallelize_in_batch("write_file", "not json", &claimed);
assert!(!result, "bad json should not be parallelizable");
}
#[test]
fn required_fields_for_write_file_reflect_schema() {
let reg = ToolRegistry::new();
let required = reg
.required_fields_for_tool("write_file")
.expect("write_file should be registered");
assert!(required.contains(&"path".to_string()));
assert!(required.contains(&"content".to_string()));
assert!(!required.contains(&"create_dirs".to_string()));
}
#[test]
fn paths_overlap_exact_match() {
let claimed: std::collections::HashSet<String> =
["src/main.rs".to_string()].into_iter().collect();
assert!(super::paths_overlap("src/main.rs", &claimed));
}
#[test]
fn paths_overlap_parent_child() {
let claimed: std::collections::HashSet<String> = ["src".to_string()].into_iter().collect();
assert!(
super::paths_overlap("src/main.rs", &claimed),
"child of claimed parent should overlap"
);
}
#[test]
fn paths_overlap_child_parent() {
let claimed: std::collections::HashSet<String> =
["src/main.rs".to_string()].into_iter().collect();
assert!(
super::paths_overlap("src", &claimed),
"parent of claimed child should overlap"
);
}
#[test]
fn paths_overlap_trailing_slash() {
let claimed: std::collections::HashSet<String> = ["src/".to_string()].into_iter().collect();
assert!(
super::paths_overlap("src/main.rs", &claimed),
"trailing slash parent should overlap with child"
);
}
#[test]
fn paths_overlap_no_overlap_siblings() {
let claimed: std::collections::HashSet<String> =
["src/main.rs".to_string()].into_iter().collect();
assert!(
!super::paths_overlap("src/lib.rs", &claimed),
"sibling files should NOT overlap"
);
}
#[test]
fn paths_overlap_no_false_prefix() {
let claimed: std::collections::HashSet<String> = ["src".to_string()].into_iter().collect();
assert!(
!super::paths_overlap("src2/foo.rs", &claimed),
"src2 is NOT a child of src — no path separator"
);
}
#[test]
fn paths_overlap_empty_claimed() {
let claimed: std::collections::HashSet<String> = std::collections::HashSet::new();
assert!(
!super::paths_overlap("anything", &claimed),
"empty claimed set should never overlap"
);
}
fn make_browser_schema(description: &str) -> ToolSchema {
ToolSchema {
name: "browser_navigate".into(),
description: description.into(),
parameters: json!({ "type": "object", "properties": {} }),
strict: None,
}
}
fn make_tool_schema(name: &str) -> ToolSchema {
ToolSchema {
name: name.into(),
description: format!("Tool {name}"),
parameters: json!({ "type": "object", "properties": {} }),
strict: None,
}
}
#[test]
fn strip_cross_refs_removes_ref_when_referenced_tool_absent() {
let desc = "Navigate to a URL. \
For simple info retrieval, prefer web_search or web_extract (faster, cheaper). \
Use browser tools when interactive.";
let mut defs = vec![make_browser_schema(desc)];
super::strip_unavailable_cross_refs(&mut defs);
assert!(
!defs[0].description.contains("web_search"),
"cross-ref to missing web_search should be stripped"
);
assert!(
defs[0].description.contains("Use browser tools"),
"rest of description should be preserved"
);
}
#[test]
fn strip_cross_refs_preserves_ref_when_referenced_tool_present() {
let desc = "Navigate to a URL. \
For simple info retrieval, prefer web_search or web_extract (faster, cheaper). \
Use browser tools when interactive.";
let mut defs = vec![make_browser_schema(desc), make_tool_schema("web_search")];
super::strip_unavailable_cross_refs(&mut defs);
assert!(
defs[0].description.contains("web_search"),
"cross-ref should be kept when web_search is available (at least one present)"
);
}
#[test]
fn strip_cross_refs_handles_empty_definitions() {
let mut defs: Vec<ToolSchema> = vec![];
super::strip_unavailable_cross_refs(&mut defs);
}
#[test]
fn strip_cross_refs_handles_no_matching_cross_refs() {
let mut defs = vec![make_tool_schema("unrelated_tool")];
let original_desc = defs[0].description.clone();
super::strip_unavailable_cross_refs(&mut defs);
assert_eq!(
defs[0].description, original_desc,
"unrelated tool unchanged"
);
}
}