forgecode / crates /forge_app /src /tool_registry.rs
SaylorTwift's picture
SaylorTwift HF Staff
Add files using upload-large-folder tool
e5034c3 verified
Raw History Blame Contribute Delete
39.2 kB
use std::sync::Arc;
use std::time::Duration;
use anyhow::Context;
use console::style;
use forge_domain::{
Agent, AgentId, AgentInput, ChatResponse, ChatResponseContent, Environment, InputModality,
Model, SystemContext, TemplateConfig, ToolCallContext, ToolCallFull, ToolCatalog,
ToolDefinition, ToolKind, ToolName, ToolOutput, ToolResult,
};
use forge_template::Element;
use futures::future::join_all;
use serde_json::{Map, Value, json};
use strum::IntoEnumIterator;
use tokio::time::timeout;
use crate::agent_executor::AgentExecutor;
use crate::dto::ToolsOverview;
use crate::error::Error;
use crate::fmt::content::FormatContent;
use crate::mcp_executor::McpExecutor;
use crate::tool_executor::ToolExecutor;
use crate::{
AgentRegistry, EnvironmentInfra, McpService, PolicyService, ProviderService, Services,
ToolResolver, WorkspaceService,
};
pub struct ToolRegistry<S> {
tool_executor: ToolExecutor<S>,
agent_executor: AgentExecutor<S>,
mcp_executor: McpExecutor<S>,
services: Arc<S>,
}
impl<S: Services + EnvironmentInfra<Config = forge_config::ForgeConfig>> ToolRegistry<S> {
pub fn new(services: Arc<S>) -> Self {
Self {
services: services.clone(),
tool_executor: ToolExecutor::new(services.clone()),
agent_executor: AgentExecutor::new(services.clone()),
mcp_executor: McpExecutor::new(services.clone()),
}
}
async fn call_with_timeout<F, Fut>(
&self,
tool_name: &ToolName,
future: F,
) -> anyhow::Result<ToolOutput>
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = anyhow::Result<ToolOutput>>,
{
let tool_timeout = Duration::from_secs(self.services.get_config()?.tool_timeout_secs);
timeout(tool_timeout, future())
.await
.context(Error::CallTimeout {
timeout: tool_timeout.as_secs() / 60,
tool_name: tool_name.clone(),
})?
}
/// Check if a tool operation is allowed based on the workflow policies
async fn check_tool_permission(
&self,
tool_input: &ToolCatalog,
context: &ToolCallContext,
) -> anyhow::Result<bool> {
let cwd = self.services.get_environment().cwd;
let operation = tool_input.to_policy_operation(cwd.clone());
if let Some(operation) = operation {
let decision = self.services.check_operation_permission(&operation).await?;
// Send custom policy message to the user when a policy file was created
if let Some(policy_path) = decision.path {
use forge_domain::TitleFormat;
use crate::utils::format_display_path;
context
.send_tool_input(
TitleFormat::debug("Permissions Update")
.sub_title(format_display_path(policy_path.as_path(), &cwd)),
)
.await?;
}
if !decision.allowed {
return Ok(true);
}
}
Ok(false)
}
async fn call_inner(
&self,
agent: &Agent,
input: ToolCallFull,
context: &ToolCallContext,
) -> anyhow::Result<ToolOutput> {
Self::validate_tool_call(agent, &input.name)?;
tracing::info!(tool_name = %input.name, arguments = %input.arguments.clone().into_string(), "Executing tool call");
let tool_name = input.name.clone();
// First, try to call a Forge tool
if ToolCatalog::contains(&input.name) {
let tool_input: ToolCatalog = ToolCatalog::try_from(input)?;
// Special handling for Task tool - delegate to AgentExecutor
if let ToolCatalog::Task(task_input) = tool_input {
let executor = self.agent_executor.clone();
let session_id = task_input.session_id.clone();
let agent_id = task_input.agent_id.clone();
// Parse session_id into ConversationId if present
let conversation_id = session_id
.map(|id| forge_domain::ConversationId::parse(&id))
.transpose()
.ok()
.flatten();
// NOTE: Agents should not timeout
let outputs = join_all(task_input.tasks.into_iter().map(|task| {
let agent_id = agent_id.clone();
let executor = executor.clone();
async move {
executor
.execute(AgentId::new(&agent_id), task, context, conversation_id)
.await
}
}))
.await
.into_iter()
.collect::<anyhow::Result<Vec<_>>>()?;
return Ok(ToolOutput::from(outputs.into_iter()));
}
let env = self.services.get_environment();
if let Some(content) = tool_input.to_content(&env) {
context.send(content).await?;
}
// Check permissions before executing the tool (only in restricted mode)
// This is done BEFORE the timeout to ensure permissions are never timed out
let is_restricted = self.services.get_config()?.restricted;
if is_restricted && self.check_tool_permission(&tool_input, context).await? {
// Send formatted output message for policy denial
context
.send(forge_domain::TitleFormat::error("Permission Denied"))
.await?;
return Ok(ToolOutput::text(
Element::new("permission_denied")
.cdata("User has denied the permission to execute this tool"),
));
}
// Validate tool modality support before execution
// Only resolve the current model when modality validation is needed.
if matches!(&tool_input, ToolCatalog::Read(input) if Self::has_image_extension(&input.file_path))
{
let model = self.get_current_model().await;
Self::validate_tool_modality(&tool_input, model.as_ref())?;
}
self.call_with_timeout(&tool_name, || {
self.tool_executor.execute(tool_input, context)
})
.await
} else if self.agent_executor.contains_tool(&input.name).await? {
// Handle agent delegation tool calls
let agent_input = AgentInput::try_from(&input)?;
let executor = self.agent_executor.clone();
let agent_name = input.name.as_str().to_string();
// NOTE: Agents should not timeout
let outputs = join_all(agent_input.tasks.into_iter().map(|task| {
let agent_name = agent_name.clone();
let executor = executor.clone();
async move {
executor
.execute(AgentId::new(&agent_name), task, context, None)
.await
}
}))
.await
.into_iter()
.collect::<anyhow::Result<Vec<_>>>()?;
Ok(ToolOutput::from(outputs.into_iter()))
} else if self.mcp_executor.contains_tool(&input.name).await? {
let output = self
.call_with_timeout(&tool_name, || self.mcp_executor.execute(input, context))
.await?;
let text = output
.values
.iter()
.filter_map(|output| output.as_str())
.fold(String::new(), |mut a, b| {
a.push('\n');
a.push_str(b);
a
});
if !text.trim().is_empty() {
let text = style(text).cyan().dim().to_string();
context
.send(ChatResponse::TaskMessage {
content: ChatResponseContent::ToolOutput(text),
})
.await?;
}
Ok(output)
} else {
Err(Error::NotFound(input.name).into())
}
}
pub async fn call(
&self,
agent: &Agent,
context: &ToolCallContext,
call: ToolCallFull,
) -> ToolResult {
let call_id = call.call_id.clone();
let tool_name = call.name.clone();
let output = self.call_inner(agent, call, context).await;
ToolResult::new(tool_name).call_id(call_id).output(output)
}
pub async fn list(&self) -> anyhow::Result<Vec<ToolDefinition>> {
Ok(self.tools_overview().await?.into())
}
/// Gets the model for the currently active agent by looking up the agent
/// and fetching its model from the provider's model list.
///
/// Returns None if no active agent, agent not found, or model not in
/// provider list.
async fn get_current_model(&self) -> Option<Model> {
let agent_id = self.services.get_active_agent_id().await.ok()??;
let agent = self.services.get_agent(&agent_id).await.ok()??;
let provider = self.services.get_provider(agent.provider).await.ok()?;
let models = self.services.models(provider).await.ok()?;
models.iter().find(|m| m.id == agent.model).cloned()
}
pub async fn tools_overview(&self) -> anyhow::Result<ToolsOverview> {
let mcp_tools = self.services.get_mcp_servers().await?;
let agent_tools = self.agent_executor.agent_definitions().await?;
// Get agents for template rendering in Task tool description
let mut agents = self.services.get_agents().await?;
// Check if current working directory is indexed
let environment = self.services.get_environment();
let cwd = environment.cwd.clone();
let is_indexed = self.services.is_indexed(&cwd).await.unwrap_or(false);
let is_authenticated = self.services.is_authenticated().await.unwrap_or(false);
// Get current model for dynamic tool descriptions
let model = self.get_current_model().await;
// Build TemplateConfig from ForgeConfig for tool description templates
let config = self.services.get_config()?;
// Filter out research subagents from task tool description when disabled
if !config.research_subagent {
agents.retain(|agent| {
let id = agent.id.as_str();
id != "sage" && id != "agent"
});
}
let template_config = TemplateConfig {
max_read_size: config.max_read_lines as usize,
max_line_length: config.max_line_chars,
max_image_size: config.max_image_size_bytes as usize,
stdout_max_prefix_length: config.max_stdout_prefix_lines,
stdout_max_suffix_length: config.max_stdout_suffix_lines,
stdout_max_line_length: config.max_stdout_line_chars,
};
Ok(ToolsOverview::new()
.system(Self::get_system_tools(
is_indexed && is_authenticated,
&environment,
model,
agents,
&template_config,
))
.agents(agent_tools)
.mcp(mcp_tools))
}
}
impl<S> ToolRegistry<S> {
fn get_system_tools(
sem_search_supported: bool,
env: &Environment,
model: Option<Model>,
agents: Vec<forge_domain::Agent>,
template_config: &TemplateConfig,
) -> Vec<ToolDefinition> {
use crate::TemplateEngine;
let handlebars = TemplateEngine::handlebar_instance();
let mut agents = agents;
agents.sort_by(|left, right| left.id.as_str().cmp(right.id.as_str()));
// Build tool_names map from all available tools
let tool_names: Map<String, Value> = ToolCatalog::iter()
.filter(|tool| {
// Only include tools that are supported (filter sem_search if not supported)
if matches!(tool, ToolCatalog::SemSearch(_)) {
sem_search_supported
} else {
true
}
})
.map(|tool| {
let def = tool.definition();
(def.name.to_string(), json!(def.name.to_string()))
})
.collect();
// Create template data with environment nested under "env"
let ctx = SystemContext {
env: Some(env.clone()),
model,
tool_names,
agents,
config: Some(template_config.clone()),
..Default::default()
};
ToolCatalog::iter()
.filter(|tool| {
// Filter out sem_search if cwd is not indexed
if matches!(tool, ToolCatalog::SemSearch(_)) {
sem_search_supported
} else {
true
}
})
.map(|tool| {
let mut def = tool.definition();
// Render template variables in description
if let Ok(rendered) = handlebars.render_template(&def.description, &ctx) {
def.description = rendered;
}
def
})
.collect::<Vec<_>>()
}
/// Validates if a tool is supported by both the agent and the system.
///
/// # Validation Process
/// Verifies the tool is supported by the agent specified in the context
fn validate_tool_call(agent: &Agent, tool_name: &ToolName) -> Result<(), Error> {
// Check if tool matches any pattern (supports globs like "mcp_*")
let matches = ToolResolver::is_allowed(agent, tool_name);
if !matches {
tracing::error!(tool_name = %tool_name, "No tool with name");
let supported_tools = agent
.tools
.iter()
.flatten()
.map(|t| t.as_str())
.collect::<Vec<_>>()
.join(", ");
return Err(Error::NotAllowed { name: tool_name.clone(), supported_tools });
}
Ok(())
}
/// Checks if a file path has an image extension.
/// This is a lightweight check that doesn't require reading the file.
fn has_image_extension(path: &str) -> bool {
const IMAGE_EXTENSIONS: &[&str] = &[
".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".svg", ".pdf",
];
let path_lower = path.to_lowercase();
IMAGE_EXTENSIONS.iter().any(|ext| path_lower.ends_with(ext))
}
/// Validates if a tool's modality requirements are supported by the current
/// model.
///
/// # Validation Process
/// Checks if the tool requires image input support and if the model
/// supports it. Currently, only the `read` tool can potentially require
/// image modality.
fn validate_tool_modality(
tool_input: &ToolCatalog,
model: Option<&Model>,
) -> Result<(), Error> {
// Check if this tool might require image support
// Currently, only the read tool can return image content
if let ToolCatalog::Read(input) = tool_input {
// Check if the file extension suggests it's an image
if Self::has_image_extension(&input.file_path) {
// Check if the model supports image input
let supports_image = model
.and_then(|m| {
m.input_modalities
.iter()
.find(|im| matches!(im, InputModality::Image))
})
.is_some();
if !supports_image {
let tool_name = ToolKind::Read.name();
let required_modality = "image".to_string();
let supported_modalities = model
.map(|m| {
m.input_modalities
.iter()
.map(|im| match im {
InputModality::Text => "text".to_string(),
InputModality::Image => "image".to_string(),
})
.collect::<Vec<_>>()
.join(", ")
})
.unwrap_or_else(|| "unknown".to_string());
return Err(Error::UnsupportedModality {
tool_name,
required_modality,
supported_modalities,
});
}
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use forge_domain::{
Agent, AgentId, Environment, ModelId, ProviderId, TemplateConfig, ToolCatalog, ToolName,
};
use pretty_assertions::assert_eq;
use crate::error::Error;
use crate::tool_registry::{ToolRegistry, create_test_agents};
fn agent() -> Agent {
// only allow read and search tools for this agent
Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.tools(vec![ToolName::new("read"), ToolName::new("fs_search")])
}
#[tokio::test]
async fn test_restricted_tool_call() {
let result = ToolRegistry::<()>::validate_tool_call(
&agent(),
&ToolName::new(ToolCatalog::Read(Default::default())),
);
assert!(result.is_ok(), "Tool call should be valid");
}
#[tokio::test]
async fn test_restricted_tool_call_err() {
let error = ToolRegistry::<()>::validate_tool_call(&agent(), &ToolName::new("write"))
.unwrap_err()
.to_string();
assert_eq!(
error,
"Tool 'write' is not available. Please try again with one of these tools: [read, fs_search]"
);
}
#[test]
fn test_validate_tool_call_with_glob_pattern_wildcard() {
let fixture = Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.tools(vec![ToolName::new("mcp_*"), ToolName::new("read")]);
let actual = ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("mcp_foo"));
assert!(actual.is_ok());
}
#[test]
fn test_validate_tool_call_with_glob_pattern_multiple_tools() {
let fixture = Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.tools(vec![ToolName::new("mcp_*"), ToolName::new("read")]);
let actual_mcp_read =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("mcp_read"));
let actual_mcp_write =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("mcp_write"));
let actual_read = ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("read"));
assert!(actual_mcp_read.is_ok());
assert!(actual_mcp_write.is_ok());
assert!(actual_read.is_ok());
}
#[test]
fn test_validate_tool_call_with_glob_pattern_no_match() {
let fixture = Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.tools(vec![ToolName::new("mcp_*"), ToolName::new("read")]);
let actual = ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("write"));
let expected = Error::NotAllowed {
name: ToolName::new("write"),
supported_tools: "mcp_*, read".to_string(),
}
.to_string();
assert_eq!(actual.unwrap_err().to_string(), expected);
}
#[test]
fn test_validate_tool_call_with_glob_pattern_question_mark() {
let fixture = Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.tools(vec![ToolName::new("read?"), ToolName::new("write")]);
let actual_read1 =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("read1"));
let actual_readx =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("readx"));
let actual_read = ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("read"));
assert!(actual_read1.is_ok());
assert!(actual_readx.is_ok());
assert!(actual_read.is_err());
}
#[test]
fn test_validate_tool_call_with_glob_pattern_character_class() {
let fixture = Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.tools(vec![ToolName::new("tool_[abc]"), ToolName::new("write")]);
let actual_tool_a =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("tool_a"));
let actual_tool_b =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("tool_b"));
let actual_tool_c =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("tool_c"));
let actual_tool_d =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("tool_d"));
assert!(actual_tool_a.is_ok());
assert!(actual_tool_b.is_ok());
assert!(actual_tool_c.is_ok());
assert!(actual_tool_d.is_err());
}
#[test]
fn test_validate_tool_call_with_glob_pattern_double_wildcard() {
let fixture = Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.tools(vec![ToolName::new("**"), ToolName::new("read")]);
let actual_any_tool =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("any_tool_name"));
let actual_nested =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("nested/tool"));
assert!(actual_any_tool.is_ok());
assert!(actual_nested.is_ok());
}
#[test]
fn test_validate_tool_call_exact_match_with_special_chars() {
let fixture = Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.tools(vec![ToolName::new("tool_[special]"), ToolName::new("read")]);
let actual =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("tool_[special]"));
// The glob pattern "tool_[special]" will match "tool_s", "tool_p", etc., not
// the literal string So this test verifies that exact matching doesn't
// work when the pattern is a valid glob
assert!(actual.is_err());
}
#[test]
fn test_validate_tool_call_backward_compatibility_exact_match() {
let fixture = Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.tools(vec![
ToolName::new("read"),
ToolName::new("write"),
ToolName::new("fs_search"),
]);
let actual_read = ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("read"));
let actual_write =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("write"));
let actual_invalid =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("delete"));
assert!(actual_read.is_ok());
assert!(actual_write.is_ok());
assert!(actual_invalid.is_err());
}
#[test]
fn test_validate_tool_call_empty_tools_list() {
let fixture = Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
);
let actual = ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("read"));
assert!(actual.is_err());
}
#[test]
fn test_validate_tool_call_glob_with_prefix_suffix() {
let fixture = Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.tools(vec![ToolName::new("mcp_*_tool")]);
let actual_match =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("mcp_read_tool"));
let actual_no_match =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("mcp_read"));
assert!(actual_match.is_ok());
assert!(actual_no_match.is_err());
}
#[test]
fn test_validate_tool_call_capitalized_read_write() {
// Test that capitalized "Read" and "Write" are accepted when agent has
// lowercase versions
let fixture = Agent::new(
AgentId::new("test_agent"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.tools(vec![ToolName::new("read"), ToolName::new("write")]);
let actual_read = ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("Read"));
let actual_write =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("Write"));
let actual_lowercase_read =
ToolRegistry::<()>::validate_tool_call(&fixture, &ToolName::new("read"));
assert!(actual_read.is_ok(), "Capitalized 'Read' should be accepted");
assert!(
actual_write.is_ok(),
"Capitalized 'Write' should be accepted"
);
assert!(
actual_lowercase_read.is_ok(),
"Lowercase 'read' should still be accepted"
);
}
#[test]
fn test_sem_search_included_when_supported() {
use fake::{Fake, Faker};
let env: Environment = Faker.fake();
let template_config = TemplateConfig::default();
let actual = ToolRegistry::<()>::get_system_tools(
true,
&env,
None,
create_test_agents(),
&template_config,
);
assert!(actual.iter().any(|t| t.name.as_str() == "sem_search"));
}
#[test]
fn test_sem_search_filtered_when_not_supported() {
use fake::{Fake, Faker};
let env: Environment = Faker.fake();
let template_config = TemplateConfig::default();
let actual = ToolRegistry::<()>::get_system_tools(
false,
&env,
None,
create_test_agents(),
&template_config,
);
assert!(actual.iter().all(|t| t.name.as_str() != "sem_search"));
}
#[test]
fn test_task_tool_description_is_stable_across_agent_order() {
use fake::{Fake, Faker};
let env: Environment = Faker.fake();
let template_config = TemplateConfig::default();
let agents = create_test_agents();
let mut reversed_agents = agents.clone();
reversed_agents.reverse();
let fixture =
ToolRegistry::<()>::get_system_tools(true, &env, None, agents, &template_config);
let actual = ToolRegistry::<()>::get_system_tools(
true,
&env,
None,
reversed_agents,
&template_config,
);
let expected = fixture
.iter()
.find(|tool| tool.name.as_str() == "task")
.expect("Task tool should exist")
.description
.clone();
let actual = actual
.iter()
.find(|tool| tool.name.as_str() == "task")
.expect("Task tool should exist")
.description
.clone();
assert_eq!(actual, expected);
}
}
#[cfg(test)]
fn create_test_agents() -> Vec<forge_domain::Agent> {
use forge_domain::{Agent, AgentId, ModelId, ProviderId, ToolName};
vec![
Agent::new(
AgentId::new("sage"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.id(AgentId::new("sage"))
.title("Research Agent")
.description("Specialized in researching codebases")
.tools(vec![
ToolName::new("read"),
ToolName::new("fs_search"),
ToolName::new("sem_search"),
ToolName::new("fetch"),
]),
Agent::new(
AgentId::new("debug"),
ProviderId::ANTHROPIC,
ModelId::new("claude-3-5-sonnet-20241022"),
)
.id(AgentId::new("debug"))
.title("Debug Agent")
.description("Specialized in debugging issues")
.tools(vec![
ToolName::new("read"),
ToolName::new("shell"),
ToolName::new("fs_search"),
ToolName::new("sem_search"),
ToolName::new("fetch"),
]),
]
}
#[cfg(test)]
fn create_test_model(
id: &str,
modalities: Vec<forge_domain::InputModality>,
) -> forge_domain::Model {
use forge_domain::{Model, ModelId};
Model {
id: ModelId::new(id),
name: Some(format!("Test {}", id)),
description: None,
context_length: Some(128000),
tools_supported: Some(true),
supports_parallel_tool_calls: Some(true),
supports_reasoning: Some(false),
input_modalities: modalities,
}
}
#[test]
fn test_template_rendering_in_tool_descriptions() {
use fake::{Fake, Faker};
let env: Environment = Faker.fake();
let template_config = TemplateConfig { max_line_length: 2000, ..Default::default() };
let actual = ToolRegistry::<()>::get_system_tools(
true,
&env,
None,
create_test_agents(),
&template_config,
);
let fs_search_tool = actual
.iter()
.find(|t| t.name.as_str() == "fs_search")
.unwrap();
// The description should not contain unrendered template variables
assert!(
!fs_search_tool.description.contains("{{"),
"Description should not contain unrendered template variable: {}",
fs_search_tool.description
);
// The description should contain the expected usage info
assert!(
fs_search_tool.description.contains("ripgrep"),
"Description should mention ripgrep: {}",
fs_search_tool.description
);
}
#[test]
fn test_dynamic_tool_description_with_vision_model() {
use fake::{Fake, Faker};
use forge_domain::InputModality;
let env: Environment = Faker.fake();
let template_config = TemplateConfig {
max_read_size: 2000,
max_line_length: 2000,
max_image_size: 5000,
..Default::default()
};
let vision_model = create_test_model("gpt-4o", vec![InputModality::Text, InputModality::Image]);
let tools_with_vision = ToolRegistry::<()>::get_system_tools(
true,
&env,
Some(vision_model),
create_test_agents(),
&template_config,
);
let read_tool = tools_with_vision
.iter()
.find(|t| t.name.as_str() == "read")
.unwrap();
insta::assert_snapshot!(read_tool.description);
}
#[test]
fn test_dynamic_tool_description_with_text_only_model() {
use fake::{Fake, Faker};
use forge_domain::InputModality;
let env: Environment = Faker.fake();
let template_config = TemplateConfig {
max_read_size: 2000,
max_line_length: 2000,
max_image_size: 5000,
..Default::default()
};
let text_only_model = create_test_model("gpt-3.5-turbo", vec![InputModality::Text]);
let tools_text_only = ToolRegistry::<()>::get_system_tools(
true,
&env,
Some(text_only_model),
create_test_agents(),
&template_config,
);
let read_tool = tools_text_only
.iter()
.find(|t| t.name.as_str() == "read")
.unwrap();
// Text-only model should NOT see image and PDF support
insta::assert_snapshot!(read_tool.description);
}
#[test]
fn test_validate_tool_modality_with_image_file_and_vision_model() {
use forge_domain::{InputModality, ToolCatalog};
let vision_model = create_test_model("gpt-4o", vec![InputModality::Text, InputModality::Image]);
let tool_input = ToolCatalog::Read(forge_domain::FSRead {
file_path: "/home/user/test.png".to_string(),
..Default::default()
});
let result = ToolRegistry::<()>::validate_tool_modality(&tool_input, Some(&vision_model));
assert!(result.is_ok(), "Vision model should support image files");
}
#[test]
fn test_validate_tool_modality_with_image_file_and_text_only_model() {
use forge_domain::{InputModality, ToolCatalog};
let text_only_model = create_test_model("gpt-3.5-turbo", vec![InputModality::Text]);
let tool_input = ToolCatalog::Read(forge_domain::FSRead {
file_path: "/home/user/test.png".to_string(),
..Default::default()
});
let result = ToolRegistry::<()>::validate_tool_modality(&tool_input, Some(&text_only_model));
assert!(
result.is_err(),
"Text-only model should not support image files"
);
let error = result.unwrap_err();
assert!(error.to_string().contains("requires image modality"));
assert!(error.to_string().contains("read"));
}
#[test]
fn test_validate_tool_modality_with_text_file_and_text_only_model() {
use forge_domain::{InputModality, ToolCatalog};
let text_only_model = create_test_model("gpt-3.5-turbo", vec![InputModality::Text]);
let tool_input = ToolCatalog::Read(forge_domain::FSRead {
file_path: "/home/user/test.txt".to_string(),
..Default::default()
});
let result = ToolRegistry::<()>::validate_tool_modality(&tool_input, Some(&text_only_model));
assert!(result.is_ok(), "Text-only model should support text files");
}
#[test]
fn test_validate_tool_modality_with_no_model() {
use forge_domain::ToolCatalog;
let tool_input = ToolCatalog::Read(forge_domain::FSRead {
file_path: "/home/user/test.png".to_string(),
..Default::default()
});
let result = ToolRegistry::<()>::validate_tool_modality(&tool_input, None);
assert!(result.is_err(), "Should error when no model is available");
let error = result.unwrap_err();
assert!(error.to_string().contains("requires image modality"));
assert!(error.to_string().contains("unknown"));
}
#[test]
fn test_validate_tool_modality_with_non_read_tool() {
use forge_domain::{InputModality, ToolCatalog};
let text_only_model = create_test_model("gpt-3.5-turbo", vec![InputModality::Text]);
let tool_input = ToolCatalog::Write(forge_domain::FSWrite {
file_path: "/home/user/test.png".to_string(),
content: "test".to_string(),
..Default::default()
});
let result = ToolRegistry::<()>::validate_tool_modality(&tool_input, Some(&text_only_model));
assert!(
result.is_ok(),
"Non-read tools should pass modality validation"
);
}
#[test]
fn test_has_image_extension() {
// Test various image extensions (case-insensitive)
assert!(ToolRegistry::<()>::has_image_extension("/path/to/file.png"));
assert!(ToolRegistry::<()>::has_image_extension("/path/to/file.PNG"));
assert!(ToolRegistry::<()>::has_image_extension("/path/to/file.jpg"));
assert!(ToolRegistry::<()>::has_image_extension(
"/path/to/file.jpeg"
));
assert!(ToolRegistry::<()>::has_image_extension(
"/path/to/file.JPEG"
));
assert!(ToolRegistry::<()>::has_image_extension("/path/to/file.gif"));
assert!(ToolRegistry::<()>::has_image_extension("/path/to/file.bmp"));
assert!(ToolRegistry::<()>::has_image_extension(
"/path/to/file.webp"
));
assert!(ToolRegistry::<()>::has_image_extension("/path/to/file.svg"));
// Test relative paths
assert!(ToolRegistry::<()>::has_image_extension("image.png"));
assert!(ToolRegistry::<()>::has_image_extension(
"../images/photo.jpg"
));
assert!(ToolRegistry::<()>::has_image_extension("/path/to/file.pdf"));
// Test non-image files
assert!(!ToolRegistry::<()>::has_image_extension(
"/path/to/file.txt"
));
assert!(!ToolRegistry::<()>::has_image_extension("/path/to/file.rs"));
assert!(!ToolRegistry::<()>::has_image_extension("/path/to/file"));
assert!(!ToolRegistry::<()>::has_image_extension("README.md"));
// Test edge cases
assert!(!ToolRegistry::<()>::has_image_extension(""));
assert!(ToolRegistry::<()>::has_image_extension(
"file.with.dots.png"
));
assert!(ToolRegistry::<()>::has_image_extension(".png")); // Hidden file with .png extension
}
#[test]
fn test_dynamic_tool_description_without_model() {
use fake::{Fake, Faker};
let env: Environment = Faker.fake();
let template_config = TemplateConfig {
max_read_size: 2000,
max_image_size: 5000,
max_line_length: 2000,
..Default::default()
};
// When no model is provided, should default to showing minimal capabilities
let tools_no_model = ToolRegistry::<()>::get_system_tools(
true,
&env,
None,
create_test_agents(),
&template_config,
);
let read_tool = tools_no_model
.iter()
.find(|t| t.name.as_str() == "read")
.unwrap();
// Without model info, should show basic text file support
insta::assert_snapshot!(read_tool.description);
}
#[test]
fn test_all_rendered_tool_descriptions() {
use fake::{Fake, Faker};
let mut env: Environment = Faker.fake();
env.cwd = "/home/user/project".into();
let template_config = TemplateConfig {
max_read_size: 2000,
max_line_length: 2000,
max_image_size: 5000,
stdout_max_prefix_length: 200,
stdout_max_suffix_length: 200,
stdout_max_line_length: 2000,
};
let tools = ToolRegistry::<()>::get_system_tools(
true,
&env,
None,
create_test_agents(),
&template_config,
);
// Verify all tools have rendered descriptions (no template syntax left)
for tool in &tools {
assert!(
!tool.description.contains("{{"),
"Tool '{}' has unrendered template variables:\n{}",
tool.name,
tool.description
);
}
// Snapshot all rendered tool descriptions for visual verification
// This will fail if a tool is renamed and descriptions reference the old name
let all_descriptions: Vec<_> = tools
.iter()
.map(|t| format!("### {}\n\n{}\n", t.name, t.description))
.collect();
insta::assert_snapshot!(
"all_rendered_tool_descriptions",
all_descriptions.join("\n---\n\n")
);
}