From c297dd092fd834b56de5ca0d04affc0c2d24dd78 Mon Sep 17 00:00:00 2001 From: MrAvalonApple <74775400+ibrahimkazimov@users.noreply.github.com> Date: Tue, 7 Apr 2026 20:42:08 +0300 Subject: [PATCH] fix: enhance tool registration and filtering for runtime-added tools --- src/agent/agent.ts | 2 +- src/agent/runner.ts | 21 ++++++++++----------- src/tool/framework.ts | 20 ++++++++++++++++++-- tests/tool-filtering.test.ts | 30 +++++++++++++++++++++++++++--- tests/trace.test.ts | 1 + 5 files changed, 57 insertions(+), 17 deletions(-) diff --git a/src/agent/agent.ts b/src/agent/agent.ts index 78d5cca..8c1007c 100644 --- a/src/agent/agent.ts +++ b/src/agent/agent.ts @@ -263,7 +263,7 @@ export class Agent { * The tool becomes available to the next LLM call — no restart required. */ addTool(tool: FrameworkToolDefinition): void { - this._toolRegistry.register(tool) + this._toolRegistry.register(tool, { runtimeAdded: true }) } /** diff --git a/src/agent/runner.ts b/src/agent/runner.ts index f1b6030..81155e8 100644 --- a/src/agent/runner.ts +++ b/src/agent/runner.ts @@ -47,8 +47,6 @@ export const TOOL_PRESETS = { full: ['file_read', 'file_write', 'file_edit', 'grep', 'glob', 'bash'], } as const satisfies Record -const BUILT_IN_TOOL_NAMES = new Set(TOOL_PRESETS.full as readonly string[]) - /** Framework-level disallowed tools for safety rails. */ export const AGENT_FRAMEWORK_DISALLOWED: readonly string[] = [ // Empty for now, infrastructure for future built-in tools @@ -228,39 +226,40 @@ export class AgentRunner { ) if (overlap.length > 0) { console.warn( - `AgentRunner: tool "${overlap[0]}" appears in both allowedTools and disallowedTools. ` + + `AgentRunner: tools [${overlap.map(name => `"${name}"`).join(', ')}] appear in both allowedTools and disallowedTools. ` + 'This is contradictory and may lead to unexpected behavior.' ) } } const allTools = this.toolRegistry.toToolDefs() - const customTools = allTools.filter(t => !BUILT_IN_TOOL_NAMES.has(t.name)) - let builtInTools = allTools.filter(t => BUILT_IN_TOOL_NAMES.has(t.name)) + const runtimeCustomTools = this.toolRegistry.toRuntimeToolDefs() + const runtimeCustomToolNames = new Set(runtimeCustomTools.map(t => t.name)) + let filteredTools = allTools.filter(t => !runtimeCustomToolNames.has(t.name)) // 1. Apply preset filter if set if (this.options.toolPreset) { const presetTools = new Set(TOOL_PRESETS[this.options.toolPreset] as readonly string[]) - builtInTools = builtInTools.filter(t => presetTools.has(t.name)) + filteredTools = filteredTools.filter(t => presetTools.has(t.name)) } // 2. Apply allowlist filter if set if (this.options.allowedTools) { - builtInTools = builtInTools.filter(t => this.options.allowedTools!.includes(t.name)) + filteredTools = filteredTools.filter(t => this.options.allowedTools!.includes(t.name)) } // 3. Apply denylist filter if set if (this.options.disallowedTools) { const denied = new Set(this.options.disallowedTools) - builtInTools = builtInTools.filter(t => !denied.has(t.name)) + filteredTools = filteredTools.filter(t => !denied.has(t.name)) } // 4. Apply framework-level safety rails const frameworkDenied = new Set(AGENT_FRAMEWORK_DISALLOWED) - builtInTools = builtInTools.filter(t => !frameworkDenied.has(t.name)) + filteredTools = filteredTools.filter(t => !frameworkDenied.has(t.name)) - // Custom tools stay available regardless of built-in filtering rules. - return [...builtInTools, ...customTools] + // Runtime-added custom tools stay available regardless of filtering rules. + return [...filteredTools, ...runtimeCustomTools] } // ------------------------------------------------------------------------- diff --git a/src/tool/framework.ts b/src/tool/framework.ts index 6b6a574..3b25c97 100644 --- a/src/tool/framework.ts +++ b/src/tool/framework.ts @@ -93,13 +93,17 @@ export function defineTool(config: { export class ToolRegistry { // eslint-disable-next-line @typescript-eslint/no-explicit-any private readonly tools = new Map>() + private readonly runtimeToolNames = new Set() /** * Add a tool to the registry. Throws if a tool with the same name has * already been registered — prevents silent overwrites. */ // eslint-disable-next-line @typescript-eslint/no-explicit-any - register(tool: ToolDefinition): void { + register( + tool: ToolDefinition, + options?: { runtimeAdded?: boolean }, + ): void { if (this.tools.has(tool.name)) { throw new Error( `ToolRegistry: a tool named "${tool.name}" is already registered. ` + @@ -107,6 +111,9 @@ export class ToolRegistry { ) } this.tools.set(tool.name, tool) + if (options?.runtimeAdded === true) { + this.runtimeToolNames.add(tool.name) + } } /** Return a tool by name, or `undefined` if not found. */ @@ -147,11 +154,12 @@ export class ToolRegistry { */ unregister(name: string): void { this.tools.delete(name) + this.runtimeToolNames.delete(name) } /** Alias for {@link unregister} — available for symmetry with `register`. */ deregister(name: string): void { - this.tools.delete(name) + this.unregister(name) } /** @@ -170,6 +178,14 @@ export class ToolRegistry { }) } + /** + * Return only tools that were added dynamically at runtime (e.g. via + * `agent.addTool()`), in LLM definition format. + */ + toRuntimeToolDefs(): LLMToolDef[] { + return this.toToolDefs().filter(tool => this.runtimeToolNames.has(tool.name)) + } + /** * Convert all registered tools to the Anthropic-style `input_schema` * format. Prefer {@link toToolDefs} for normal use; this method is exposed diff --git a/tests/tool-filtering.test.ts b/tests/tool-filtering.test.ts index 5cf09c0..4f3b6f0 100644 --- a/tests/tool-filtering.test.ts +++ b/tests/tool-filtering.test.ts @@ -74,7 +74,7 @@ function createTestTools() { description: 'Custom tool', inputSchema: z.object({ input: z.string() }), execute: async () => ({ data: 'custom', isError: false }), - })) + }), { runtimeAdded: true }) return registry } @@ -248,6 +248,29 @@ describe('Tool filtering', () => { expect(toolNames).toEqual(['custom_tool']) }) + + it('runtime-added tools bypass filtering regardless of tool name', () => { + const runtimeBuiltinNamedRegistry = new ToolRegistry() + runtimeBuiltinNamedRegistry.register(defineTool({ + name: 'file_read', + description: 'Runtime override', + inputSchema: z.object({ path: z.string() }), + execute: async () => ({ data: 'runtime', isError: false }), + }), { runtimeAdded: true }) + + const runtimeBuiltinNamedRunner = new AgentRunner( + mockAdapter, + runtimeBuiltinNamedRegistry, + new ToolExecutor(runtimeBuiltinNamedRegistry), + { + model: 'test-model', + disallowedTools: ['file_read'], + }, + ) + + const tools = (runtimeBuiltinNamedRunner as any).resolveTools() as LLMToolDef[] + expect(tools.map(t => t.name)).toEqual(['file_read']) + }) }) describe('resolveTools - validation warnings', () => { @@ -271,7 +294,7 @@ describe('Tool filtering', () => { ;(runner as any).resolveTools() expect(consoleWarnSpy).toHaveBeenCalledWith( - expect.stringContaining('tool "bash" appears in both allowedTools and disallowedTools') + expect.stringContaining('tools ["bash"] appear in both allowedTools and disallowedTools') ) }) @@ -301,4 +324,5 @@ describe('Tool filtering', () => { expect(consoleWarnSpy).not.toHaveBeenCalled() }) }) -}) \ No newline at end of file +}) + diff --git a/tests/trace.test.ts b/tests/trace.test.ts index 00e8330..a7a5257 100644 --- a/tests/trace.test.ts +++ b/tests/trace.test.ts @@ -451,3 +451,4 @@ describe('Agent trace events', () => { expect(llmTraces[1]!.turn).toBe(2) }) }) +