fix: enhance tool registration and filtering for runtime-added tools

This commit is contained in:
MrAvalonApple 2026-04-07 20:42:08 +03:00
parent ff008aad31
commit c297dd092f
5 changed files with 57 additions and 17 deletions

View File

@ -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 })
}
/**

View File

@ -47,8 +47,6 @@ export const TOOL_PRESETS = {
full: ['file_read', 'file_write', 'file_edit', 'grep', 'glob', 'bash'],
} as const satisfies Record<string, readonly string[]>
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]
}
// -------------------------------------------------------------------------

View File

@ -93,13 +93,17 @@ export function defineTool<TInput>(config: {
export class ToolRegistry {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
private readonly tools = new Map<string, ToolDefinition<any>>()
private readonly runtimeToolNames = new Set<string>()
/**
* 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<any>): void {
register(
tool: ToolDefinition<any>,
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

View File

@ -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')
)
})
@ -302,3 +325,4 @@ describe('Tool filtering', () => {
})
})
})

View File

@ -451,3 +451,4 @@ describe('Agent trace events', () => {
expect(llmTraces[1]!.turn).toBe(2)
})
})