From c613c636e1072d8a1d3e36f5646456e97dee392a Mon Sep 17 00:00:00 2001 From: Danil Nikolaev Date: Tue, 19 May 2026 08:33:18 +0300 Subject: [PATCH] Add local tool filtering --- .env.example | 9 +++++ README.md | 7 ++++ bun.lock | 8 ++-- src/ai/tools/registry.ts | 83 ++++++++++++++++++++++++++++----------- src/common/environment.ts | 6 +++ 5 files changed, 87 insertions(+), 26 deletions(-) diff --git a/.env.example b/.env.example index 5599b82..d92b8aa 100644 --- a/.env.example +++ b/.env.example @@ -46,6 +46,15 @@ USE_NAMES_IN_PROMPT=true # Disable all built-in local tools and keep only MCP tools DISABLE_LOCAL_TOOLS=false +# Filter built-in local tools by name. +# LOCAL_TOOL_ALLOWLIST lets through only the listed tools. +# LOCAL_TOOL_DENYLIST removes the listed tools. +# Examples: +# LOCAL_TOOL_ALLOWLIST=get_datetime,web_search +# LOCAL_TOOL_DENYLIST=shell_execute,python_interpreter +LOCAL_TOOL_ALLOWLIST= +LOCAL_TOOL_DENYLIST= + # Custom system prompt for AI (or put it into data/SYSTEM_PROMPT.md) SYSTEM_PROMPT= diff --git a/README.md b/README.md index 4ee6fec..5675abd 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,13 @@ If you want to disable all built-in local tools and use only MCP tools, set: DISABLE_LOCAL_TOOLS=true ``` +If you want a partial filter instead, use tool names: + +```bash +LOCAL_TOOL_ALLOWLIST=get_datetime,web_search +LOCAL_TOOL_DENYLIST=shell_execute,python_interpreter +``` + For local Ollama document RAG, install an embedding model locally and set it in `.env`: ```bash diff --git a/bun.lock b/bun.lock index a9d356f..8a11f2c 100644 --- a/bun.lock +++ b/bun.lock @@ -14,8 +14,8 @@ "emoji-regex": "^10.6.0", "fluent-ffmpeg": "^2.1.3", "ollama": "^0.6.3", - "openai": "^6.37.0", - "pg": "^8.20.0", + "openai": "^6.38.0", + "pg": "^8.21.0", "qrcode": "^1.5.4", "sharp": "^0.34.5", "systeminformation": "^5.31.6", @@ -27,12 +27,12 @@ "@eslint/js": "^9.39.4", "@types/bun": "^1.3.14", "@types/fluent-ffmpeg": "^2.1.28", - "@types/node": "^25.8.0", + "@types/node": "^25.9.0", "@types/pg": "^8.20.0", "@types/qrcode": "^1.5.6", "eslint": "^9.39.4", "typescript": "^6.0.3", - "typescript-eslint": "^8.59.3", + "typescript-eslint": "^8.59.4", }, }, }, diff --git a/src/ai/tools/registry.ts b/src/ai/tools/registry.ts index dce51c5..4d89019 100644 --- a/src/ai/tools/registry.ts +++ b/src/ai/tools/registry.ts @@ -73,35 +73,62 @@ export const fileTools = [ deletePathTool, ] satisfies AiTool[]; +function parseToolNameSet(raw: string | undefined): Set | undefined { + if (!raw?.trim()) return undefined; + + const names = raw + .split(",") + .map(item => item.trim().toLowerCase()) + .filter(Boolean); + + return names.length ? new Set(names) : undefined; +} + +function isLocalToolEnabled(toolName: string): boolean { + if (Environment.DISABLE_LOCAL_TOOLS) return false; + + const allowlist = parseToolNameSet(Environment.LOCAL_TOOL_ALLOWLIST); + if (allowlist && !allowlist.has(toolName.toLowerCase())) return false; + + const denylist = parseToolNameSet(Environment.LOCAL_TOOL_DENYLIST); + if (denylist && denylist.has(toolName.toLowerCase())) return false; + + return true; +} + +function filterEnabledTools(tools: AiTool[]): AiTool[] { + return tools.filter(tool => isLocalToolEnabled(tool.function.name)); +} + export const getTools = (forCreator?: boolean) => { - const tools: AiTool[] = Environment.DISABLE_LOCAL_TOOLS ? [] : [ - ...defaultTools, - ]; + const tools: AiTool[] = []; if (Environment.DISABLE_LOCAL_TOOLS) { tools.push(...getMcpTools()); return tools; } + tools.push(...filterEnabledTools(defaultTools)); + if (Environment.BRAVE_SEARCH_API_KEY) { - tools.push(webSearchTool); + tools.push(...filterEnabledTools([webSearchTool])); } if (Environment.OPEN_WEATHER_MAP_API_KEY) { - tools.push(getWeatherTool); + tools.push(...filterEnabledTools([getWeatherTool])); } if (Environment.FILE_TOOLS_ROOT_DIR && Environment.ENABLE_FS_TOOLS) { - tools.push(...fileTools); + tools.push(...filterEnabledTools(fileTools)); } if (forCreator) { if (Environment.ENABLE_PYTHON_INTERPRETER) { - tools.push(pythonInterpreterTool); + tools.push(...filterEnabledTools([pythonInterpreterTool])); } if (Environment.ENABLE_UNSAFE_EVAL) { - tools.push(shellExecuteTool); + tools.push(...filterEnabledTools([shellExecuteTool])); } } @@ -132,7 +159,7 @@ export const fileToolHandlers = { }; export const getToolHandlers = () => { - let handlers: Record = { + const handlers: Record = { ...getMcpToolHandlers(), }; @@ -140,21 +167,29 @@ export const getToolHandlers = () => { return handlers; } - handlers = { - ...handlers, - get_datetime: getCurrentDateTime, - get_financial_market_data: getMarketRates, + if (isLocalToolEnabled("get_datetime")) handlers.get_datetime = getCurrentDateTime; + if (isLocalToolEnabled("get_financial_market_data")) handlers.get_financial_market_data = getMarketRates; - ...fileToolHandlers, + if (isLocalToolEnabled("read_file")) handlers.read_file = readFile; + if (isLocalToolEnabled("list_directory")) handlers.list_directory = listDirectory; + if (isLocalToolEnabled("search_files")) handlers.search_files = searchFiles; + if (isLocalToolEnabled("create_file")) handlers.create_file = createFile; + if (isLocalToolEnabled("begin_file_write")) handlers.begin_file_write = beginFileWrite; + if (isLocalToolEnabled("write_file_chunk")) handlers.write_file_chunk = writeFileChunk; + if (isLocalToolEnabled("finish_file_write")) handlers.finish_file_write = finishFileWrite; + if (isLocalToolEnabled("cancel_file_write")) handlers.cancel_file_write = cancelFileWrite; + if (isLocalToolEnabled("send_file_as_attachment")) handlers.send_file_as_attachment = sendFileAsAttachment; + if (isLocalToolEnabled("create_directory")) handlers.create_directory = createDirectory; + if (isLocalToolEnabled("copy_path")) handlers.copy_path = copyPath; + if (isLocalToolEnabled("update_file")) handlers.update_file = updateFile; + if (isLocalToolEnabled("edit_file_patch")) handlers.edit_file_patch = editFilePatch; + if (isLocalToolEnabled("rename_path")) handlers.rename_path = renamePath; + if (isLocalToolEnabled("delete_path")) handlers.delete_path = deletePath; - python_interpreter: runPythonInterpreter, - - shell_execute: shellExecute, - - web_search: webSearch, - - get_weather: getWeather, - }; + if (isLocalToolEnabled("python_interpreter")) handlers.python_interpreter = runPythonInterpreter; + if (isLocalToolEnabled("shell_execute")) handlers.shell_execute = shellExecute; + if (isLocalToolEnabled("web_search")) handlers.web_search = webSearch; + if (isLocalToolEnabled("get_weather")) handlers.get_weather = getWeather; return handlers; }; @@ -167,6 +202,10 @@ export function getToolPrompts(toolNames: string[]): string[] { const prompts: string[] = []; for (const toolName of toolNames) { + if (!isLocalToolEnabled(toolName)) { + continue; + } + if (!prompts.includes(fileToolsToolPrompt) && fileTools.map(t => t.function.name).includes(toolName)) { prompts.push(fileToolsToolPrompt); diff --git a/src/common/environment.ts b/src/common/environment.ts index 03d82f6..a998b47 100644 --- a/src/common/environment.ts +++ b/src/common/environment.ts @@ -215,6 +215,8 @@ const RuntimeEnvSchema = z.object({ ENABLE_PYTHON_INTERPRETER: optionalBooleanSchema, DISABLE_LOCAL_TOOLS: optionalBooleanSchema, + LOCAL_TOOL_ALLOWLIST: optionalStringSchema, + LOCAL_TOOL_DENYLIST: optionalStringSchema, MCP_SERVERS: optionalStringSchema, OLLAMA_API_KEY: optionalStringSchema, @@ -311,6 +313,8 @@ export class Environment { static ENABLE_PYTHON_INTERPRETER: boolean = false; static DISABLE_LOCAL_TOOLS: boolean = false; + static LOCAL_TOOL_ALLOWLIST?: string; + static LOCAL_TOOL_DENYLIST?: string; static MCP_SERVERS?: string; static OLLAMA_API_KEY?: string; @@ -1847,6 +1851,8 @@ export class Environment { Environment.ENABLE_PYTHON_INTERPRETER = env.ENABLE_PYTHON_INTERPRETER ?? false; Environment.DISABLE_LOCAL_TOOLS = env.DISABLE_LOCAL_TOOLS ?? false; + Environment.LOCAL_TOOL_ALLOWLIST = env.LOCAL_TOOL_ALLOWLIST; + Environment.LOCAL_TOOL_DENYLIST = env.LOCAL_TOOL_DENYLIST; Environment.MCP_SERVERS = env.MCP_SERVERS; Environment.OLLAMA_API_KEY = env.OLLAMA_API_KEY;