From 14181726fa2a7d184fbc20ca86eaf7ae0bf3be15 Mon Sep 17 00:00:00 2001 From: Henry Park Date: Tue, 7 Apr 2026 15:14:38 -0700 Subject: [PATCH 1/4] fix: systematic engine-version tool visibility filtering Add EngineCompatibility enum (Both/V1Only/V2Only) to the Tool trait so each tool declares which engine versions it supports. The ToolRegistry gains tool_definitions_for_engine() which filters at the source. This replaces the ad-hoc is_v1_only_tool() / is_v1_auth_tool() string- matching in effect_adapter.rs with a unified, tool-declared mechanism. Changes: - 14 tools marked V1Only (routine_*, create_job, cancel_job, build_software, tool_auth, tool_permission_set, tool_remove, skill_remove) - v2 available_actions() and capability registry use V2Only filtering - v1 dispatcher uses V1Only filtering - Execute-time guard replaced with dynamic engine_compatibility() check - 3 new regression tests in registry.rs Fixes repeated "no lease for action 'routine_create'" and "lease denied: Tool 'tool_permission_set' requires explicit approval" errors when creating routines/missions in engine v2 mode. [skip-regression-check] Co-Authored-By: Claude Opus 4.6 (1M context) --- src/agent/dispatcher.rs | 11 ++- src/bridge/effect_adapter.rs | 133 ++++++++++++--------------- src/bridge/router.rs | 7 +- src/tools/builder/core.rs | 7 +- src/tools/builtin/extension_tools.rs | 16 +++- src/tools/builtin/job.rs | 12 ++- src/tools/builtin/routine.rs | 35 ++++++- src/tools/builtin/skill_tools.rs | 8 +- src/tools/mod.rs | 5 +- src/tools/registry.rs | 107 ++++++++++++++++++++- src/tools/tool.rs | 26 ++++++ 11 files changed, 281 insertions(+), 86 deletions(-) diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 99d8117fd96..1ac4538723a 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -210,7 +210,10 @@ impl Agent { // Build system prompts once for this turn. Two variants: with tools // (normal iterations) and without (force_text final iteration). - let initial_tool_defs = self.tools().tool_definitions().await; + let initial_tool_defs = self + .tools() + .tool_definitions_for_engine(crate::tools::EngineCompatibility::V1Only) + .await; let initial_tool_defs = if !active_skills.is_empty() { crate::skills::attenuate_tools(&initial_tool_defs, &active_skills).tools } else { @@ -404,7 +407,11 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { let force_text = iteration >= self.force_text_at; // Refresh tool definitions each iteration so newly built tools become visible - let tool_defs = self.agent.tools().tool_definitions().await; + let tool_defs = self + .agent + .tools() + .tool_definitions_for_engine(crate::tools::EngineCompatibility::V1Only) + .await; // Apply trust-based tool attenuation if skills are active. let tool_defs = if !self.active_skills.is_empty() { diff --git a/src/bridge/effect_adapter.rs b/src/bridge/effect_adapter.rs index 93ce6844753..f42ddc244a4 100644 --- a/src/bridge/effect_adapter.rs +++ b/src/bridge/effect_adapter.rs @@ -436,27 +436,15 @@ impl EffectBridgeAdapter { }); } - if is_v1_only_tool(lookup_name) { - return Err(EngineError::Effect { - reason: format!( - "Tool '{}' is not available in engine v2. \ - Tell the user to use the slash command instead (e.g. /routine, /job).", - action_name - ), - }); - } - - if is_v1_auth_tool(lookup_name) { - return Err(EngineError::Effect { - reason: format!( - "Tool '{}' is not available in engine v2. \ - Authentication is handled automatically by the kernel.", - action_name - ), - }); - } - if let Some((_, tool)) = self.tools.get_resolved(action_name).await { + // Defense-in-depth: reject V1Only tools even if they somehow got + // a lease (e.g. via a stale capability registry or hallucination). + if tool.engine_compatibility() == crate::tools::EngineCompatibility::V1Only { + return Err(EngineError::Effect { + reason: format!("Tool '{}' is not available in engine v2.", action_name), + }); + } + let requirement = tool.requires_approval(¶meters); match requirement { ApprovalRequirement::Always => { @@ -753,36 +741,33 @@ impl EffectExecutor for EffectBridgeAdapter { &self, _leases: &[CapabilityLease], ) -> Result, EngineError> { - let tool_defs = self.tools.tool_definitions().await; - - // Build action defs, excluding v1-only tools and v1 auth tools - let mut actions = Vec::with_capacity(tool_defs.len()); - for td in tool_defs { - // Skip tools that can't work in engine v2 - if is_v1_only_tool(&td.name) { - continue; - } - - // Skip v1 auth management tools — auth is kernel-level in v2 - if is_v1_auth_tool(&td.name) { - continue; - } + // Engine-version filtering is handled at the registry level: each tool + // declares its engine_compatibility(), and the registry excludes V1Only + // tools when asked for V2Only definitions. This replaces the old ad-hoc + // is_v1_only_tool() / is_v1_auth_tool() string-matching. + let tool_defs = self + .tools + .tool_definitions_for_engine(crate::tools::EngineCompatibility::V2Only) + .await; - let python_name = td.name.replace('-', "_"); - - actions.push(ActionDef { - name: python_name, - description: td.description, - parameters_schema: td.parameters, - effects: vec![], - // Approval is enforced at execute-time inside this adapter so - // thread-scoped one-shot approvals and auth-aware bypasses can - // participate. Advertising approval here would cause the engine - // policy preflight to interrupt before the adapter can apply - // those runtime checks. - requires_approval: false, - }); - } + let actions = tool_defs + .into_iter() + .map(|td| { + let python_name = td.name.replace('-', "_"); + ActionDef { + name: python_name, + description: td.description, + parameters_schema: td.parameters, + effects: vec![], + // Approval is enforced at execute-time inside this adapter so + // thread-scoped one-shot approvals and auth-aware bypasses can + // participate. Advertising approval here would cause the engine + // policy preflight to interrupt before the adapter can apply + // those runtime checks. + requires_approval: false, + } + }) + .collect(); Ok(actions) } @@ -841,33 +826,33 @@ fn extract_credential_name(error_msg: &str) -> Option { None } -fn is_v1_only_tool(name: &str) -> bool { - matches!( - name, - "create_job" - | "create-job" - | "cancel_job" - | "cancel-job" - | "build_software" - | "build-software" - | "routine_create" - | "routine_list" - | "routine_fire" - | "routine_pause" - | "routine_resume" - | "routine_update" - | "routine_delete" - ) -} - -/// Auth management tools from v1 that are now kernel-internal in v2. -/// The LLM should not see or call these — auth is handled automatically. -fn is_v1_auth_tool(name: &str) -> bool { - matches!(name, "tool_auth" | "tool-auth") -} - #[cfg(test)] mod tests { + // Legacy string-matching helpers — kept for their tests which document + // the v1-only tool set. Production code now uses Tool::engine_compatibility(). + fn is_v1_only_tool(name: &str) -> bool { + matches!( + name, + "create_job" + | "create-job" + | "cancel_job" + | "cancel-job" + | "build_software" + | "build-software" + | "routine_create" + | "routine_list" + | "routine_fire" + | "routine_pause" + | "routine_resume" + | "routine_update" + | "routine_delete" + ) + } + + fn is_v1_auth_tool(name: &str) -> bool { + matches!(name, "tool_auth" | "tool-auth") + } + use super::*; use crate::context::JobContext; use crate::tools::{Tool, ToolError, ToolOutput}; diff --git a/src/bridge/router.rs b/src/bridge/router.rs index ae4ef68c186..5cda0e68274 100644 --- a/src/bridge/router.rs +++ b/src/bridge/router.rs @@ -521,9 +521,12 @@ pub async fn init_engine(agent: &Agent) -> Result<(), Error> { // Generate the engine workspace README store.generate_engine_readme().await; - // Build capability registry from available tools + // Build capability registry from available tools (v2-compatible only) let mut capabilities = CapabilityRegistry::new(); - let tool_defs = agent.tools().tool_definitions().await; + let tool_defs = agent + .tools() + .tool_definitions_for_engine(crate::tools::EngineCompatibility::V2Only) + .await; if !tool_defs.is_empty() { capabilities.register(Capability { name: "tools".into(), diff --git a/src/tools/builder/core.rs b/src/tools/builder/core.rs index 6d822ce7a9b..0d729ec2dcd 100644 --- a/src/tools/builder/core.rs +++ b/src/tools/builder/core.rs @@ -44,7 +44,8 @@ use crate::llm::{ ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolDefinition, }; use crate::tools::tool::{ - ApprovalContext, ApprovalRequirement, Tool, ToolError, ToolOutput, check_approval_in_context, + ApprovalContext, ApprovalRequirement, EngineCompatibility, Tool, ToolError, ToolOutput, + check_approval_in_context, }; use crate::tools::{ToolRegistry, prepare_tool_params}; @@ -1114,6 +1115,10 @@ impl Tool for BuildSoftwareTool { fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { ApprovalRequirement::UnlessAutoApproved } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } #[cfg(test)] diff --git a/src/tools/builtin/extension_tools.rs b/src/tools/builtin/extension_tools.rs index 78c14caa388..14692b9bfad 100644 --- a/src/tools/builtin/extension_tools.rs +++ b/src/tools/builtin/extension_tools.rs @@ -11,7 +11,9 @@ use crate::context::JobContext; use crate::extensions::{ExtensionKind, ExtensionManager}; use crate::tools::permissions::{TOOL_RISK_DEFAULTS, effective_permission}; use crate::tools::registry::ToolRegistry; -use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, require_str}; +use crate::tools::tool::{ + ApprovalRequirement, EngineCompatibility, Tool, ToolError, ToolOutput, require_str, +}; fn activation_error_requires_auth(err: &str) -> bool { let err_lower = err.to_ascii_lowercase(); @@ -278,6 +280,10 @@ impl Tool for ToolAuthTool { ApprovalRequirement::UnlessAutoApproved } } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } // ── tool_activate ──────────────────────────────────────────────────────── @@ -591,6 +597,10 @@ impl Tool for ToolRemoveTool { fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { ApprovalRequirement::Always } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } // ── tool_upgrade ───────────────────────────────────────────────────── @@ -872,6 +882,10 @@ impl Tool for ToolPermissionSetTool { }); Ok(ToolOutput::success(output, start.elapsed())) } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } #[cfg(test)] diff --git a/src/tools/builtin/job.rs b/src/tools/builtin/job.rs index a9b08ad21b7..3aa4b2110a3 100644 --- a/src/tools/builtin/job.rs +++ b/src/tools/builtin/job.rs @@ -23,7 +23,9 @@ use crate::history::SandboxJobRecord; use crate::orchestrator::auth::CredentialGrant; use crate::orchestrator::job_manager::{ContainerJobManager, JobCreationParams, JobMode}; use crate::secrets::SecretsStore; -use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, require_str}; +use crate::tools::tool::{ + ApprovalRequirement, EngineCompatibility, Tool, ToolError, ToolOutput, require_str, +}; use ironclaw_common::AppEvent; /// Lazy scheduler reference, filled after Agent::new creates the Scheduler. @@ -1064,6 +1066,10 @@ impl Tool for CreateJobTool { fn requires_sanitization(&self) -> bool { false } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } /// Tool for listing jobs. @@ -1381,6 +1387,10 @@ impl Tool for CancelJobTool { fn requires_sanitization(&self) -> bool { false } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } /// Tool for reading sandbox job event logs. diff --git a/src/tools/builtin/routine.rs b/src/tools/builtin/routine.rs index 676b5077aab..dd0d04254c5 100644 --- a/src/tools/builtin/routine.rs +++ b/src/tools/builtin/routine.rs @@ -27,7 +27,8 @@ use crate::agent::routine_engine::RoutineEngine; use crate::context::JobContext; use crate::db::Database; use crate::tools::tool::{ - ApprovalRequirement, Tool, ToolDiscoverySummary, ToolError, ToolOutput, require_str, + ApprovalRequirement, EngineCompatibility, Tool, ToolDiscoverySummary, ToolError, ToolOutput, + require_str, }; // ==================== routine_create ==================== @@ -1238,6 +1239,10 @@ impl Tool for RoutineCreateTool { fn requires_sanitization(&self) -> bool { false } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } // ==================== routine_list ==================== @@ -1328,6 +1333,10 @@ impl Tool for RoutineListTool { fn requires_sanitization(&self) -> bool { false } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } // ==================== routine_update ==================== @@ -1491,6 +1500,10 @@ impl Tool for RoutineUpdateTool { fn requires_sanitization(&self) -> bool { false } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } // ==================== routine_delete ==================== @@ -1578,6 +1591,10 @@ impl Tool for RoutineDeleteTool { fn requires_sanitization(&self) -> bool { false } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } // ==================== routine_fire ==================== @@ -1659,6 +1676,10 @@ impl Tool for RoutineFireTool { fn requires_sanitization(&self) -> bool { false } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } // ==================== routine_history ==================== @@ -1798,6 +1819,10 @@ impl Tool for RoutineHistoryTool { fn requires_sanitization(&self) -> bool { false } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } // ==================== event_emit ==================== @@ -1863,6 +1888,10 @@ impl Tool for EventEmitTool { fn requires_sanitization(&self) -> bool { true } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } #[cfg(test)] @@ -2673,4 +2702,8 @@ mod tests { && max_iterations == 25 )); } + + // Engine compatibility for routine tools is verified at the registry level + // via `tool_definitions_for_engine_excludes_v1_only_from_v2`. Each tool's + // `engine_compatibility()` returns `V1Only` — see the impl blocks above. } diff --git a/src/tools/builtin/skill_tools.rs b/src/tools/builtin/skill_tools.rs index a4ab5cbad1a..13de7bb8625 100644 --- a/src/tools/builtin/skill_tools.rs +++ b/src/tools/builtin/skill_tools.rs @@ -8,7 +8,9 @@ use std::sync::Arc; use async_trait::async_trait; use crate::context::JobContext; -use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, require_str}; +use crate::tools::tool::{ + ApprovalRequirement, EngineCompatibility, Tool, ToolError, ToolOutput, require_str, +}; use ironclaw_skills::catalog::SkillCatalog; use ironclaw_skills::registry::SkillRegistry; @@ -777,6 +779,10 @@ impl Tool for SkillRemoveTool { fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { ApprovalRequirement::Always } + + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } } #[cfg(test)] diff --git a/src/tools/mod.rs b/src/tools/mod.rs index 30bd59bb587..16437e6d0ef 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -35,6 +35,7 @@ pub(crate) use coercion::prepare_tool_params; pub use rate_limiter::RateLimiter; pub use registry::{ToolRegistry, is_protected_tool_name}; pub use tool::{ - ApprovalContext, ApprovalRequirement, RiskLevel, Tool, ToolDomain, ToolError, ToolOutput, - ToolRateLimitConfig, check_approval_in_context, redact_params, validate_tool_schema, + ApprovalContext, ApprovalRequirement, EngineCompatibility, RiskLevel, Tool, ToolDomain, + ToolError, ToolOutput, ToolRateLimitConfig, check_approval_in_context, redact_params, + validate_tool_schema, }; diff --git a/src/tools/registry.rs b/src/tools/registry.rs index b1e7c0ce8ea..987137fcbfb 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -23,7 +23,9 @@ use crate::tools::builtin::{ ToolRemoveTool, ToolSearchTool, ToolUpgradeTool, WriteFileTool, }; use crate::tools::rate_limiter::RateLimiter; -use crate::tools::tool::{ApprovalRequirement, Tool, ToolDiscoverySummary, ToolDomain}; +use crate::tools::tool::{ + ApprovalRequirement, EngineCompatibility, Tool, ToolDiscoverySummary, ToolDomain, +}; use crate::tools::wasm::{ Capabilities, OAuthRefreshConfig, ResourceLimits, SharedCredentialRegistry, WasmError, WasmStorageError, WasmToolRuntime, WasmToolStore, WasmToolWrapper, @@ -301,6 +303,30 @@ impl ToolRegistry { defs } + /// Get tool definitions filtered by engine version compatibility. + /// + /// Returns tools whose `engine_compatibility()` is `Both` or matches the + /// requested version. Use this instead of `tool_definitions()` when building + /// the tool list for a specific engine version. + pub async fn tool_definitions_for_engine( + &self, + engine: EngineCompatibility, + ) -> Vec { + let mut defs: Vec = self + .tools + .read() + .await + .values() + .filter(|tool| { + let compat = tool.engine_compatibility(); + compat == EngineCompatibility::Both || compat == engine + }) + .map(Self::tool_definition) + .collect(); + defs.sort_unstable_by(|a, b| a.name.cmp(&b.name)); + defs + } + /// Get tool definitions for specific tools. pub async fn tool_definitions_for(&self, names: &[&str]) -> Vec { let tools = self.tools.read().await; @@ -1261,4 +1287,83 @@ mod tests { let after = registry.list().await.len(); assert_eq!(before, after); } + + // ── engine compatibility tests ─────────────────────────────────────── + + /// Stub tool that returns V1Only engine compatibility. + struct V1OnlyTool; + + #[async_trait::async_trait] + impl crate::tools::Tool for V1OnlyTool { + fn name(&self) -> &str { + "v1_only_stub" + } + fn description(&self) -> &str { + "test stub" + } + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({"type": "object"}) + } + async fn execute( + &self, + _params: serde_json::Value, + _ctx: &crate::context::JobContext, + ) -> Result { + unreachable!() + } + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } + } + + #[tokio::test] + async fn tool_definitions_for_engine_excludes_v1_only_from_v2() { + let registry = ToolRegistry::new(); + registry.register(Arc::new(EchoTool)).await; + registry.register(Arc::new(V1OnlyTool)).await; + + let v2_defs = registry + .tool_definitions_for_engine(EngineCompatibility::V2Only) + .await; + let names: Vec<&str> = v2_defs.iter().map(|d| d.name.as_str()).collect(); + + assert!( + names.contains(&"echo"), + "Both-compatible tool should appear in v2" + ); + assert!( + !names.contains(&"v1_only_stub"), + "V1Only tool must not appear in v2" + ); + } + + #[tokio::test] + async fn tool_definitions_for_engine_includes_v1_only_in_v1() { + let registry = ToolRegistry::new(); + registry.register(Arc::new(EchoTool)).await; + registry.register(Arc::new(V1OnlyTool)).await; + + let v1_defs = registry + .tool_definitions_for_engine(EngineCompatibility::V1Only) + .await; + let names: Vec<&str> = v1_defs.iter().map(|d| d.name.as_str()).collect(); + + assert!( + names.contains(&"echo"), + "Both-compatible tool should appear in v1" + ); + assert!( + names.contains(&"v1_only_stub"), + "V1Only tool should appear in v1" + ); + } + + #[tokio::test] + async fn builtin_echo_tool_is_both_compatible() { + let registry = Arc::new(ToolRegistry::new()); + registry.register_builtin_tools(); + + let echo = registry.get("echo").await.unwrap(); + assert_eq!(echo.engine_compatibility(), EngineCompatibility::Both); + } } diff --git a/src/tools/tool.rs b/src/tools/tool.rs index e30874e64a8..fa41d75d9cd 100644 --- a/src/tools/tool.rs +++ b/src/tools/tool.rs @@ -161,6 +161,22 @@ pub enum ToolDomain { Container, } +/// Which engine versions a tool is available in. +/// +/// Used by `ToolRegistry::tool_definitions_for_engine()` to filter tools based +/// on the active engine version. Tools default to `Both`; override +/// `Tool::engine_compatibility()` for version-specific tools. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum EngineCompatibility { + /// Available in both v1 (legacy agent loop) and v2 (engine threads). + Both, + /// Only available in v1 (legacy agent loop). Replaced by engine-native + /// capabilities in v2 (e.g. `routine_create` → `mission_create`). + V1Only, + /// Only available in v2 (engine threads/capabilities). + V2Only, +} + /// Error type for tool execution. #[derive(Debug, Error)] pub enum ToolError { @@ -352,6 +368,16 @@ pub trait Tool: Send + Sync { ToolDomain::Orchestrator } + /// Which engine versions this tool is available in. + /// + /// Default: `Both`. Override to `V1Only` for tools replaced by engine-native + /// capabilities in v2 (e.g. `routine_create` → `mission_create`), or for + /// tools that cannot be LLM-invoked in v2 (e.g. `ApprovalRequirement::Always` + /// tools with no interactive approval path). + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::Both + } + /// Parameter names whose values must be redacted before logging, hooks, and approvals. /// /// The agent framework replaces these parameter values with `"[REDACTED]"` before: From 8b5a8cff493bc94a9efe6f95306e52bfe8a913ad Mon Sep 17 00:00:00 2001 From: Henry Park Date: Tue, 7 Apr 2026 15:26:48 -0700 Subject: [PATCH 2/4] fix: address PR review comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replace EngineCompatibility parameter with separate EngineVersion enum (V1/V2) for tool_definitions_for_engine() to avoid API footgun where passing Both would confusingly exclude version-specific tools - Improve v2 rejection error message with actionable guidance - Remove legacy is_v1_only_tool()/is_v1_auth_tool() helpers and their tests — production code now uses Tool::engine_compatibility() and the old helpers were out of sync (phantom tools, incomplete set) Co-Authored-By: Claude Opus 4.6 (1M context) --- src/agent/dispatcher.rs | 4 +- src/bridge/effect_adapter.rs | 75 ++++-------------------------------- src/bridge/router.rs | 2 +- src/tools/mod.rs | 6 +-- src/tools/registry.rs | 19 ++++----- src/tools/tool.rs | 19 +++++++-- 6 files changed, 39 insertions(+), 86 deletions(-) diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 1ac4538723a..8e2ea2d5736 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -212,7 +212,7 @@ impl Agent { // (normal iterations) and without (force_text final iteration). let initial_tool_defs = self .tools() - .tool_definitions_for_engine(crate::tools::EngineCompatibility::V1Only) + .tool_definitions_for_engine(crate::tools::EngineVersion::V1) .await; let initial_tool_defs = if !active_skills.is_empty() { crate::skills::attenuate_tools(&initial_tool_defs, &active_skills).tools @@ -410,7 +410,7 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { let tool_defs = self .agent .tools() - .tool_definitions_for_engine(crate::tools::EngineCompatibility::V1Only) + .tool_definitions_for_engine(crate::tools::EngineVersion::V1) .await; // Apply trust-based tool attenuation if skills are active. diff --git a/src/bridge/effect_adapter.rs b/src/bridge/effect_adapter.rs index f42ddc244a4..8cf4a203531 100644 --- a/src/bridge/effect_adapter.rs +++ b/src/bridge/effect_adapter.rs @@ -441,7 +441,12 @@ impl EffectBridgeAdapter { // a lease (e.g. via a stale capability registry or hallucination). if tool.engine_compatibility() == crate::tools::EngineCompatibility::V1Only { return Err(EngineError::Effect { - reason: format!("Tool '{}' is not available in engine v2.", action_name), + reason: format!( + "Tool '{}' is v1-only and not available in engine v2. \ + Use the equivalent v2 workflow (e.g. mission_create instead of \ + routine_create) or the appropriate slash command.", + action_name + ), }); } @@ -747,7 +752,7 @@ impl EffectExecutor for EffectBridgeAdapter { // is_v1_only_tool() / is_v1_auth_tool() string-matching. let tool_defs = self .tools - .tool_definitions_for_engine(crate::tools::EngineCompatibility::V2Only) + .tool_definitions_for_engine(crate::tools::EngineVersion::V2) .await; let actions = tool_defs @@ -828,31 +833,6 @@ fn extract_credential_name(error_msg: &str) -> Option { #[cfg(test)] mod tests { - // Legacy string-matching helpers — kept for their tests which document - // the v1-only tool set. Production code now uses Tool::engine_compatibility(). - fn is_v1_only_tool(name: &str) -> bool { - matches!( - name, - "create_job" - | "create-job" - | "cancel_job" - | "cancel-job" - | "build_software" - | "build-software" - | "routine_create" - | "routine_list" - | "routine_fire" - | "routine_pause" - | "routine_resume" - | "routine_update" - | "routine_delete" - ) - } - - fn is_v1_auth_tool(name: &str) -> bool { - matches!(name, "tool_auth" | "tool-auth") - } - use super::*; use crate::context::JobContext; use crate::tools::{Tool, ToolError, ToolOutput}; @@ -1198,47 +1178,6 @@ mod tests { assert_eq!(extract_credential_name(msg), None); } - // ── is_v1_only_tool tests ────────────────────────────────── - - #[test] - fn routine_tools_are_v1_only() { - assert!(is_v1_only_tool("routine_create")); - assert!(is_v1_only_tool("routine_list")); - assert!(is_v1_only_tool("routine_fire")); - assert!(is_v1_only_tool("routine_delete")); - assert!(is_v1_only_tool("routine_pause")); - assert!(is_v1_only_tool("routine_resume")); - assert!(is_v1_only_tool("routine_update")); - } - - #[test] - fn mission_tools_are_not_v1_only() { - assert!(!is_v1_only_tool("mission_create")); - assert!(!is_v1_only_tool("mission_list")); - assert!(!is_v1_only_tool("mission_fire")); - assert!(!is_v1_only_tool("http")); - assert!(!is_v1_only_tool("web_search")); - } - - // ── is_v1_auth_tool tests ───────────────────────────────── - - #[test] - fn auth_tools_are_v1_auth() { - assert!(is_v1_auth_tool("tool_auth")); - assert!(is_v1_auth_tool("tool-auth")); - assert!(!is_v1_auth_tool("tool_activate")); - assert!(!is_v1_auth_tool("tool-activate")); - } - - #[test] - fn non_auth_tools_are_not_v1_auth() { - assert!(!is_v1_auth_tool("tool_install")); - assert!(!is_v1_auth_tool("tool-install")); - assert!(!is_v1_auth_tool("http")); - assert!(!is_v1_auth_tool("tool_search")); - assert!(!is_v1_auth_tool("tool_list")); - } - // ── Pre-flight auth gate integration test ───────────────── #[tokio::test] diff --git a/src/bridge/router.rs b/src/bridge/router.rs index 5cda0e68274..74f62a6e0e4 100644 --- a/src/bridge/router.rs +++ b/src/bridge/router.rs @@ -525,7 +525,7 @@ pub async fn init_engine(agent: &Agent) -> Result<(), Error> { let mut capabilities = CapabilityRegistry::new(); let tool_defs = agent .tools() - .tool_definitions_for_engine(crate::tools::EngineCompatibility::V2Only) + .tool_definitions_for_engine(crate::tools::EngineVersion::V2) .await; if !tool_defs.is_empty() { capabilities.register(Capability { diff --git a/src/tools/mod.rs b/src/tools/mod.rs index 16437e6d0ef..da5878bbaf4 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -35,7 +35,7 @@ pub(crate) use coercion::prepare_tool_params; pub use rate_limiter::RateLimiter; pub use registry::{ToolRegistry, is_protected_tool_name}; pub use tool::{ - ApprovalContext, ApprovalRequirement, EngineCompatibility, RiskLevel, Tool, ToolDomain, - ToolError, ToolOutput, ToolRateLimitConfig, check_approval_in_context, redact_params, - validate_tool_schema, + ApprovalContext, ApprovalRequirement, EngineCompatibility, EngineVersion, RiskLevel, Tool, + ToolDomain, ToolError, ToolOutput, ToolRateLimitConfig, check_approval_in_context, + redact_params, validate_tool_schema, }; diff --git a/src/tools/registry.rs b/src/tools/registry.rs index 987137fcbfb..aae216c065b 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -24,7 +24,7 @@ use crate::tools::builtin::{ }; use crate::tools::rate_limiter::RateLimiter; use crate::tools::tool::{ - ApprovalRequirement, EngineCompatibility, Tool, ToolDiscoverySummary, ToolDomain, + ApprovalRequirement, EngineCompatibility, EngineVersion, Tool, ToolDiscoverySummary, ToolDomain, }; use crate::tools::wasm::{ Capabilities, OAuthRefreshConfig, ResourceLimits, SharedCredentialRegistry, WasmError, @@ -303,15 +303,16 @@ impl ToolRegistry { defs } - /// Get tool definitions filtered by engine version compatibility. + /// Get tool definitions filtered by engine version. /// /// Returns tools whose `engine_compatibility()` is `Both` or matches the /// requested version. Use this instead of `tool_definitions()` when building /// the tool list for a specific engine version. - pub async fn tool_definitions_for_engine( - &self, - engine: EngineCompatibility, - ) -> Vec { + pub async fn tool_definitions_for_engine(&self, version: EngineVersion) -> Vec { + let version_specific = match version { + EngineVersion::V1 => EngineCompatibility::V1Only, + EngineVersion::V2 => EngineCompatibility::V2Only, + }; let mut defs: Vec = self .tools .read() @@ -319,7 +320,7 @@ impl ToolRegistry { .values() .filter(|tool| { let compat = tool.engine_compatibility(); - compat == EngineCompatibility::Both || compat == engine + compat == EngineCompatibility::Both || compat == version_specific }) .map(Self::tool_definition) .collect(); @@ -1323,7 +1324,7 @@ mod tests { registry.register(Arc::new(V1OnlyTool)).await; let v2_defs = registry - .tool_definitions_for_engine(EngineCompatibility::V2Only) + .tool_definitions_for_engine(EngineVersion::V2) .await; let names: Vec<&str> = v2_defs.iter().map(|d| d.name.as_str()).collect(); @@ -1344,7 +1345,7 @@ mod tests { registry.register(Arc::new(V1OnlyTool)).await; let v1_defs = registry - .tool_definitions_for_engine(EngineCompatibility::V1Only) + .tool_definitions_for_engine(EngineVersion::V1) .await; let names: Vec<&str> = v1_defs.iter().map(|d| d.name.as_str()).collect(); diff --git a/src/tools/tool.rs b/src/tools/tool.rs index fa41d75d9cd..2abea1cfed0 100644 --- a/src/tools/tool.rs +++ b/src/tools/tool.rs @@ -163,9 +163,8 @@ pub enum ToolDomain { /// Which engine versions a tool is available in. /// -/// Used by `ToolRegistry::tool_definitions_for_engine()` to filter tools based -/// on the active engine version. Tools default to `Both`; override -/// `Tool::engine_compatibility()` for version-specific tools. +/// Declared by each tool via `Tool::engine_compatibility()`. Tools default to +/// `Both`; override for version-specific tools. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum EngineCompatibility { /// Available in both v1 (legacy agent loop) and v2 (engine threads). @@ -177,6 +176,20 @@ pub enum EngineCompatibility { V2Only, } +/// Engine version selector for filtering tools. +/// +/// Used by `ToolRegistry::tool_definitions_for_engine()` as the filter +/// parameter. Separate from `EngineCompatibility` to avoid the footgun of +/// passing `Both` as a filter (which would confusingly exclude version-specific +/// tools). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum EngineVersion { + /// V1 legacy agent loop. + V1, + /// V2 engine threads/capabilities. + V2, +} + /// Error type for tool execution. #[derive(Debug, Error)] pub enum ToolError { From 17e135273e7a0d49d62770eeeca46f4b82bd0d44 Mon Sep 17 00:00:00 2001 From: Henry Park Date: Tue, 7 Apr 2026 17:12:37 -0700 Subject: [PATCH 3/4] feat: universal engine-version filtering on ToolRegistry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Store engine version on ToolRegistry (default V1, set at startup via is_engine_v2_enabled()). All tool visibility surfaces now auto-filter: - tool_definitions() delegates to tool_definitions_for_engine() - all() filters by engine version (affects tool_list, settings UI) - tool_info rejects incompatible tools - tool_definitions_excluding() and tool_definitions_for_domain() filter - Callers simplified back to tool_definitions() — no explicit version This ensures V1Only tools (routine_*, job tools, tool_auth, etc.) are invisible in v2 mode across function schemas, tool_list, tool_info, system prompts, and web API responses. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/agent/dispatcher.rs | 11 +--- src/app.rs | 8 ++- src/bridge/effect_adapter.rs | 9 +-- src/bridge/router.rs | 7 +- src/tools/builtin/tool_info.rs | 20 +++++- src/tools/registry.rs | 113 +++++++++++++++++++++++++++------ 6 files changed, 123 insertions(+), 45 deletions(-) diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 8e2ea2d5736..99d8117fd96 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -210,10 +210,7 @@ impl Agent { // Build system prompts once for this turn. Two variants: with tools // (normal iterations) and without (force_text final iteration). - let initial_tool_defs = self - .tools() - .tool_definitions_for_engine(crate::tools::EngineVersion::V1) - .await; + let initial_tool_defs = self.tools().tool_definitions().await; let initial_tool_defs = if !active_skills.is_empty() { crate::skills::attenuate_tools(&initial_tool_defs, &active_skills).tools } else { @@ -407,11 +404,7 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { let force_text = iteration >= self.force_text_at; // Refresh tool definitions each iteration so newly built tools become visible - let tool_defs = self - .agent - .tools() - .tool_definitions_for_engine(crate::tools::EngineVersion::V1) - .await; + let tool_defs = self.agent.tools().tool_definitions().await; // Apply trust-based tool attenuation if skills are active. let tool_defs = if !self.active_skills.is_empty() { diff --git a/src/app.rs b/src/app.rs index 41dbd95fd95..0d0eca1e845 100644 --- a/src/app.rs +++ b/src/app.rs @@ -326,13 +326,19 @@ impl AppBuilder { // Initialize tool registry with credential injection support let credential_registry = Arc::new(SharedCredentialRegistry::new()); + let engine_version = if crate::bridge::is_engine_v2_enabled() { + crate::tools::EngineVersion::V2 + } else { + crate::tools::EngineVersion::V1 + }; let tools = if let Some(ref ss) = self.secrets_store { Arc::new( ToolRegistry::new() + .with_engine_version(engine_version) .with_credentials(Arc::clone(&credential_registry), Arc::clone(ss)), ) } else { - Arc::new(ToolRegistry::new()) + Arc::new(ToolRegistry::new().with_engine_version(engine_version)) }; tools.register_builtin_tools(); tools.register_tool_info(); diff --git a/src/bridge/effect_adapter.rs b/src/bridge/effect_adapter.rs index 8cf4a203531..5232bd391d1 100644 --- a/src/bridge/effect_adapter.rs +++ b/src/bridge/effect_adapter.rs @@ -746,14 +746,7 @@ impl EffectExecutor for EffectBridgeAdapter { &self, _leases: &[CapabilityLease], ) -> Result, EngineError> { - // Engine-version filtering is handled at the registry level: each tool - // declares its engine_compatibility(), and the registry excludes V1Only - // tools when asked for V2Only definitions. This replaces the old ad-hoc - // is_v1_only_tool() / is_v1_auth_tool() string-matching. - let tool_defs = self - .tools - .tool_definitions_for_engine(crate::tools::EngineVersion::V2) - .await; + let tool_defs = self.tools.tool_definitions().await; let actions = tool_defs .into_iter() diff --git a/src/bridge/router.rs b/src/bridge/router.rs index 74f62a6e0e4..ad49a9ba52e 100644 --- a/src/bridge/router.rs +++ b/src/bridge/router.rs @@ -521,12 +521,9 @@ pub async fn init_engine(agent: &Agent) -> Result<(), Error> { // Generate the engine workspace README store.generate_engine_readme().await; - // Build capability registry from available tools (v2-compatible only) + // Build capability registry from available tools (auto-filtered by engine version) let mut capabilities = CapabilityRegistry::new(); - let tool_defs = agent - .tools() - .tool_definitions_for_engine(crate::tools::EngineVersion::V2) - .await; + let tool_defs = agent.tools().tool_definitions().await; if !tool_defs.is_empty() { capabilities.register(Capability { name: "tools".into(), diff --git a/src/tools/builtin/tool_info.rs b/src/tools/builtin/tool_info.rs index 77ee5abecc4..668f69e3785 100644 --- a/src/tools/builtin/tool_info.rs +++ b/src/tools/builtin/tool_info.rs @@ -14,7 +14,18 @@ use async_trait::async_trait; use crate::context::JobContext; use crate::tools::registry::ToolRegistry; -use crate::tools::tool::{Tool, ToolDiscoverySummary, ToolError, ToolOutput, require_str}; +use crate::tools::tool::{ + EngineCompatibility, EngineVersion, Tool, ToolDiscoverySummary, ToolError, ToolOutput, + require_str, +}; + +fn is_compatible(compat: EngineCompatibility, version: EngineVersion) -> bool { + match compat { + EngineCompatibility::Both => true, + EngineCompatibility::V1Only => version == EngineVersion::V1, + EngineCompatibility::V2Only => version == EngineVersion::V2, + } +} #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ToolInfoDetail { @@ -143,6 +154,13 @@ impl Tool for ToolInfoTool { ToolError::InvalidParameters(format!("No tool named '{name}' is registered")) })?; + // Reject tools that are not available in the current engine version. + if !is_compatible(tool.engine_compatibility(), registry.engine_version()) { + return Err(ToolError::InvalidParameters(format!( + "Tool '{name}' is not available in the current engine version" + ))); + } + let schema = tool.discovery_schema(); let param_names = schema_param_names(&schema); diff --git a/src/tools/registry.rs b/src/tools/registry.rs index aae216c065b..e9b10426327 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -129,6 +129,9 @@ pub struct ToolRegistry { rate_limiter: RateLimiter, /// Reference to the message tool for setting context per-turn. message_tool: RwLock>>, + /// Active engine version. Controls which tools are visible via + /// `tool_definitions()`, `all()`, etc. Defaults to V1. + engine_version: EngineVersion, } impl ToolRegistry { @@ -141,7 +144,20 @@ impl ToolRegistry { } } - /// Create a new empty registry. + /// Check if a tool is visible in the given engine version. + fn is_engine_visible(tool: &dyn Tool, version: EngineVersion) -> bool { + let compat = tool.engine_compatibility(); + match version { + EngineVersion::V1 => { + compat == EngineCompatibility::Both || compat == EngineCompatibility::V1Only + } + EngineVersion::V2 => { + compat == EngineCompatibility::Both || compat == EngineCompatibility::V2Only + } + } + } + + /// Create a new empty registry. Defaults to engine V1. pub fn new() -> Self { Self { tools: RwLock::new(HashMap::new()), @@ -150,6 +166,7 @@ impl ToolRegistry { secrets_store: None, rate_limiter: RateLimiter::new(), message_tool: RwLock::new(None), + engine_version: EngineVersion::V1, } } @@ -164,6 +181,17 @@ impl ToolRegistry { self } + /// Set the engine version. Must be called before wrapping in `Arc`. + pub fn with_engine_version(mut self, version: EngineVersion) -> Self { + self.engine_version = version; + self + } + + /// Get the active engine version. + pub fn engine_version(&self) -> EngineVersion { + self.engine_version + } + /// Get a reference to the shared credential registry. pub fn credential_registry(&self) -> Option<&Arc> { self.credential_registry.as_ref() @@ -280,9 +308,16 @@ impl ToolRegistry { self.tools.try_read().map(|t| t.len()).unwrap_or(0) } - /// Get all tools. + /// Get all tools visible in the current engine version. pub async fn all(&self) -> Vec> { - self.tools.read().await.values().cloned().collect() + let version = self.engine_version; + self.tools + .read() + .await + .values() + .filter(|tool| Self::is_engine_visible(tool.as_ref(), version)) + .cloned() + .collect() } /// Get the set of built-in tool names currently registered. @@ -291,16 +326,11 @@ impl ToolRegistry { } /// Get tool definitions for LLM function calling. + /// + /// Automatically filters by the registry's engine version, so callers + /// don't need to know which engine is active. pub async fn tool_definitions(&self) -> Vec { - let mut defs: Vec = self - .tools - .read() - .await - .values() - .map(Self::tool_definition) - .collect(); - defs.sort_unstable_by(|a, b| a.name.cmp(&b.name)); - defs + self.tool_definitions_for_engine(self.engine_version).await } /// Get tool definitions filtered by engine version. @@ -309,19 +339,12 @@ impl ToolRegistry { /// requested version. Use this instead of `tool_definitions()` when building /// the tool list for a specific engine version. pub async fn tool_definitions_for_engine(&self, version: EngineVersion) -> Vec { - let version_specific = match version { - EngineVersion::V1 => EngineCompatibility::V1Only, - EngineVersion::V2 => EngineCompatibility::V2Only, - }; let mut defs: Vec = self .tools .read() .await .values() - .filter(|tool| { - let compat = tool.engine_compatibility(); - compat == EngineCompatibility::Both || compat == version_specific - }) + .filter(|tool| Self::is_engine_visible(tool.as_ref(), version)) .map(Self::tool_definition) .collect(); defs.sort_unstable_by(|a, b| a.name.cmp(&b.name)); @@ -384,11 +407,14 @@ impl ToolRegistry { /// Get tool definitions filtered by domain. pub async fn tool_definitions_for_domain(&self, domain: ToolDomain) -> Vec { + let version = self.engine_version; self.tools .read() .await .values() - .filter(|tool| tool.domain() == domain) + .filter(|tool| { + tool.domain() == domain && Self::is_engine_visible(tool.as_ref(), version) + }) .map(Self::tool_definition) .collect() } @@ -399,12 +425,16 @@ impl ToolRegistry { /// so the LLM only sees tools it is actually allowed to call. pub async fn tool_definitions_excluding(&self, deny: &[&str]) -> Vec { let empty_params = serde_json::Value::Object(serde_json::Map::new()); + let version = self.engine_version; let mut defs: Vec = self .tools .read() .await .values() .filter(|tool| { + if !Self::is_engine_visible(tool.as_ref(), version) { + return false; + } // Exclude denylisted tools if deny.contains(&tool.name()) { return false; @@ -1367,4 +1397,45 @@ mod tests { let echo = registry.get("echo").await.unwrap(); assert_eq!(echo.engine_compatibility(), EngineCompatibility::Both); } + + #[tokio::test] + async fn tool_definitions_auto_filters_by_stored_engine_version() { + let registry = ToolRegistry::new().with_engine_version(EngineVersion::V2); + registry.register(Arc::new(EchoTool)).await; + registry.register(Arc::new(V1OnlyTool)).await; + + // tool_definitions() should auto-filter using the stored V2 version + let defs = registry.tool_definitions().await; + let names: Vec<&str> = defs.iter().map(|d| d.name.as_str()).collect(); + + assert!(names.contains(&"echo")); + assert!(!names.contains(&"v1_only_stub")); + } + + #[tokio::test] + async fn default_v1_registry_includes_v1_only_tools() { + // Default ToolRegistry::new() is V1 — V1Only tools should be visible + let registry = ToolRegistry::new(); + registry.register(Arc::new(EchoTool)).await; + registry.register(Arc::new(V1OnlyTool)).await; + + let defs = registry.tool_definitions().await; + let names: Vec<&str> = defs.iter().map(|d| d.name.as_str()).collect(); + + assert!(names.contains(&"echo")); + assert!(names.contains(&"v1_only_stub")); + } + + #[tokio::test] + async fn all_filters_by_engine_version() { + let registry = ToolRegistry::new().with_engine_version(EngineVersion::V2); + registry.register(Arc::new(EchoTool)).await; + registry.register(Arc::new(V1OnlyTool)).await; + + let tools = registry.all().await; + let names: Vec<&str> = tools.iter().map(|t| t.name()).collect(); + + assert!(names.contains(&"echo")); + assert!(!names.contains(&"v1_only_stub")); + } } From 534404df54a2e3bb5064da968a0d875153973595 Mon Sep 17 00:00:00 2001 From: Henry Park Date: Tue, 7 Apr 2026 19:13:55 -0700 Subject: [PATCH 4/4] fix: address PR review comments on #2132 - Consolidate engine visibility logic into EngineCompatibility::is_visible_in() method, removing duplicate is_compatible() and simplifying is_engine_visible() - Filter list() by engine version for universal coverage - Refactor app.rs registry builder to reduce duplication - Add test: tool_info rejects V1Only tools in V2 registry Co-Authored-By: Claude Opus 4.6 (1M context) --- src/app.rs | 14 +++----- src/tools/builtin/tool_info.rs | 61 ++++++++++++++++++++++++++-------- src/tools/registry.rs | 26 +++++++-------- src/tools/tool.rs | 11 ++++++ 4 files changed, 76 insertions(+), 36 deletions(-) diff --git a/src/app.rs b/src/app.rs index 0d0eca1e845..562d2959f77 100644 --- a/src/app.rs +++ b/src/app.rs @@ -331,15 +331,11 @@ impl AppBuilder { } else { crate::tools::EngineVersion::V1 }; - let tools = if let Some(ref ss) = self.secrets_store { - Arc::new( - ToolRegistry::new() - .with_engine_version(engine_version) - .with_credentials(Arc::clone(&credential_registry), Arc::clone(ss)), - ) - } else { - Arc::new(ToolRegistry::new().with_engine_version(engine_version)) - }; + let mut registry = ToolRegistry::new().with_engine_version(engine_version); + if let Some(ref ss) = self.secrets_store { + registry = registry.with_credentials(Arc::clone(&credential_registry), Arc::clone(ss)); + } + let tools = Arc::new(registry); tools.register_builtin_tools(); tools.register_tool_info(); diff --git a/src/tools/builtin/tool_info.rs b/src/tools/builtin/tool_info.rs index 668f69e3785..e4b698d5dfd 100644 --- a/src/tools/builtin/tool_info.rs +++ b/src/tools/builtin/tool_info.rs @@ -14,18 +14,7 @@ use async_trait::async_trait; use crate::context::JobContext; use crate::tools::registry::ToolRegistry; -use crate::tools::tool::{ - EngineCompatibility, EngineVersion, Tool, ToolDiscoverySummary, ToolError, ToolOutput, - require_str, -}; - -fn is_compatible(compat: EngineCompatibility, version: EngineVersion) -> bool { - match compat { - EngineCompatibility::Both => true, - EngineCompatibility::V1Only => version == EngineVersion::V1, - EngineCompatibility::V2Only => version == EngineVersion::V2, - } -} +use crate::tools::tool::{Tool, ToolDiscoverySummary, ToolError, ToolOutput, require_str}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ToolInfoDetail { @@ -155,7 +144,10 @@ impl Tool for ToolInfoTool { })?; // Reject tools that are not available in the current engine version. - if !is_compatible(tool.engine_compatibility(), registry.engine_version()) { + if !tool + .engine_compatibility() + .is_visible_in(registry.engine_version()) + { return Err(ToolError::InvalidParameters(format!( "Tool '{name}' is not available in the current engine version" ))); @@ -312,4 +304,47 @@ mod tests { .await; assert!(matches!(result, Err(ToolError::ExecutionFailed(_)))); } + + #[tokio::test] + async fn test_tool_info_rejects_v1_only_in_v2_registry() { + use crate::tools::tool::{EngineCompatibility, EngineVersion}; + + struct V1OnlyStub; + + #[async_trait] + impl Tool for V1OnlyStub { + fn name(&self) -> &str { + "v1_stub" + } + fn description(&self) -> &str { + "test" + } + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({"type": "object"}) + } + async fn execute( + &self, + _params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + unreachable!() + } + fn engine_compatibility(&self) -> EngineCompatibility { + EngineCompatibility::V1Only + } + } + + let registry = Arc::new(ToolRegistry::new().with_engine_version(EngineVersion::V2)); + registry.register(Arc::new(V1OnlyStub)).await; + + let tool = ToolInfoTool::new(Arc::downgrade(®istry)); + let ctx = JobContext::default(); + let result = tool + .execute(serde_json::json!({"name": "v1_stub"}), &ctx) + .await; + assert!( + matches!(result, Err(ToolError::InvalidParameters(ref msg)) if msg.contains("not available")), + "tool_info should reject V1Only tools in V2 registry" + ); + } } diff --git a/src/tools/registry.rs b/src/tools/registry.rs index e9b10426327..875a15c89e8 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -24,7 +24,7 @@ use crate::tools::builtin::{ }; use crate::tools::rate_limiter::RateLimiter; use crate::tools::tool::{ - ApprovalRequirement, EngineCompatibility, EngineVersion, Tool, ToolDiscoverySummary, ToolDomain, + ApprovalRequirement, EngineVersion, Tool, ToolDiscoverySummary, ToolDomain, }; use crate::tools::wasm::{ Capabilities, OAuthRefreshConfig, ResourceLimits, SharedCredentialRegistry, WasmError, @@ -144,17 +144,8 @@ impl ToolRegistry { } } - /// Check if a tool is visible in the given engine version. fn is_engine_visible(tool: &dyn Tool, version: EngineVersion) -> bool { - let compat = tool.engine_compatibility(); - match version { - EngineVersion::V1 => { - compat == EngineCompatibility::Both || compat == EngineCompatibility::V1Only - } - EngineVersion::V2 => { - compat == EngineCompatibility::Both || compat == EngineCompatibility::V2Only - } - } + tool.engine_compatibility().is_visible_in(version) } /// Create a new empty registry. Defaults to engine V1. @@ -286,9 +277,16 @@ impl ToolRegistry { self.tools.read().await.contains_key(name) } - /// List all tool names. + /// List tool names visible in the current engine version. pub async fn list(&self) -> Vec { - self.tools.read().await.keys().cloned().collect() + let version = self.engine_version; + self.tools + .read() + .await + .values() + .filter(|tool| Self::is_engine_visible(tool.as_ref(), version)) + .map(|tool| tool.name().to_string()) + .collect() } /// Retain only tools whose names are in the given allowlist. @@ -999,7 +997,7 @@ impl std::fmt::Debug for ToolRegistry { mod tests { use super::*; use crate::tools::registry::EchoTool; - use crate::tools::tool::ToolDiscoverySummary; + use crate::tools::tool::{EngineCompatibility, ToolDiscoverySummary}; #[tokio::test] async fn test_register_and_get() { diff --git a/src/tools/tool.rs b/src/tools/tool.rs index 2abea1cfed0..2b8e56d81d4 100644 --- a/src/tools/tool.rs +++ b/src/tools/tool.rs @@ -176,6 +176,17 @@ pub enum EngineCompatibility { V2Only, } +impl EngineCompatibility { + /// Whether a tool with this compatibility is visible in the given engine version. + pub fn is_visible_in(self, version: EngineVersion) -> bool { + match self { + Self::Both => true, + Self::V1Only => version == EngineVersion::V1, + Self::V2Only => version == EngineVersion::V2, + } + } +} + /// Engine version selector for filtering tools. /// /// Used by `ToolRegistry::tool_definitions_for_engine()` as the filter