From 92752a2a15b70fa4cc961a7b828e624378e6b791 Mon Sep 17 00:00:00 2001 From: Timothy Carambat Date: Fri, 5 Dec 2025 11:01:20 -0800 Subject: [PATCH] Improve MCP functionality (#4709) * Improve MCP functionality * apply ANSI strip path * attempt ARM build * improve dockerfile IO build time and migrate to ARM build * fix comment * add ability to disable MCP cooldown feature * update devbuild name * move chromium arm build patch to CDN --- .github/workflows/dev-build.yaml | 72 +++++------ docker/.env.example | 7 +- docker/Dockerfile | 6 +- server/.env.example | 7 +- server/package.json | 4 +- server/utils/MCP/hypervisor/index.js | 116 +++++++++++++++--- server/utils/MCP/index.js | 53 ++++++-- server/utils/agents/aibitat/plugins/memory.js | 6 +- .../utils/agents/aibitat/plugins/rechart.js | 4 +- .../aibitat/plugins/save-file-browser.js | 10 +- .../aibitat/providers/helpers/untooled.js | 44 ++++++- .../utils/agents/aibitat/providers/ollama.js | 14 ++- server/utils/agents/aibitat/utils/dedupe.js | 79 ++++++++++-- server/utils/helpers/updateENV.js | 3 + server/yarn.lock | 40 ++++++ 15 files changed, 370 insertions(+), 95 deletions(-) diff --git a/.github/workflows/dev-build.yaml b/.github/workflows/dev-build.yaml index fbcdcb7c..23223834 100644 --- a/.github/workflows/dev-build.yaml +++ b/.github/workflows/dev-build.yaml @@ -1,4 +1,4 @@ -name: AnythingLLM Development Docker image (amd64) +name: AnythingLLM Development Docker image (arm64) concurrency: group: build-${{ github.ref }} @@ -6,7 +6,7 @@ concurrency: on: push: - branches: ['4686-feat-migrate-react-router-to-use-createbrowserrouter'] # put your current branch to create a build. Core team only. + branches: ['mcp-improvements'] # put your current branch to create a build. Core team only. paths-ignore: - '**.md' - 'cloud-deployments/*' @@ -21,9 +21,9 @@ on: - 'extras/**/*' # Extra is just for news and other local content. jobs: - push_multi_platform_to_registries: - name: Push Docker multi-platform image to multiple registries - runs-on: ubuntu-latest + push_dev_build_to_dockerhub: + name: Push development build image to Docker Hub + runs-on: ubuntu-22.04-arm permissions: packages: write contents: read @@ -78,8 +78,8 @@ jobs: push: true sbom: true provenance: mode=max - platforms: linux/amd64 - # platforms: linux/amd64,linux/arm64 + # platforms: linux/amd64 + platforms: linux/arm64 tags: ${{ steps.meta.outputs.tags }} labels: ${{ steps.meta.outputs.labels }} cache-from: type=gha @@ -88,37 +88,37 @@ jobs: # For Docker scout there are some intermediary reported CVEs which exists outside # of execution content or are unreachable by an attacker but exist in image. # We create VEX files for these so they don't show in scout summary. - - name: Collect known and verified CVE exceptions - id: cve-list - run: | - # Collect CVEs from filenames in vex folder - CVE_NAMES="" - for file in ./docker/vex/*.vex.json; do - [ -e "$file" ] || continue - filename=$(basename "$file") - stripped_filename=${filename%.vex.json} - CVE_NAMES+=" $stripped_filename" - done - echo "CVE_EXCEPTIONS=$CVE_NAMES" >> $GITHUB_OUTPUT - shell: bash + # - name: Collect known and verified CVE exceptions + # id: cve-list + # run: | + # # Collect CVEs from filenames in vex folder + # CVE_NAMES="" + # for file in ./docker/vex/*.vex.json; do + # [ -e "$file" ] || continue + # filename=$(basename "$file") + # stripped_filename=${filename%.vex.json} + # CVE_NAMES+=" $stripped_filename" + # done + # echo "CVE_EXCEPTIONS=$CVE_NAMES" >> $GITHUB_OUTPUT + # shell: bash # About VEX attestations https://docs.docker.com/scout/explore/exceptions/ # Justifications https://github.com/openvex/spec/blob/main/OPENVEX-SPEC.md#status-justifications # Fixed to use v1.15.1 of scout-cli as v1.16.0 install script is broken # https://github.com/docker/scout-cli - - name: Add VEX attestations - env: - CVE_EXCEPTIONS: ${{ steps.cve-list.outputs.CVE_EXCEPTIONS }} - run: | - echo $CVE_EXCEPTIONS - curl -sSfL https://raw.githubusercontent.com/docker/scout-cli/main/install.sh | sh -s -- - for cve in $CVE_EXCEPTIONS; do - for tag in "${{ join(fromJSON(steps.meta.outputs.json).tags, ' ') }}"; do - echo "Attaching VEX exception $cve to $tag" - docker scout attestation add \ - --file "./docker/vex/$cve.vex.json" \ - --predicate-type https://openvex.dev/ns/v0.2.0 \ - $tag - done - done - shell: bash \ No newline at end of file + # - name: Add VEX attestations + # env: + # CVE_EXCEPTIONS: ${{ steps.cve-list.outputs.CVE_EXCEPTIONS }} + # run: | + # echo $CVE_EXCEPTIONS + # curl -sSfL https://raw.githubusercontent.com/docker/scout-cli/main/install.sh | sh -s -- + # for cve in $CVE_EXCEPTIONS; do + # for tag in "${{ join(fromJSON(steps.meta.outputs.json).tags, ' ') }}"; do + # echo "Attaching VEX exception $cve to $tag" + # docker scout attestation add \ + # --file "./docker/vex/$cve.vex.json" \ + # --predicate-type https://openvex.dev/ns/v0.2.0 \ + # $tag + # done + # done + # shell: bash \ No newline at end of file diff --git a/docker/.env.example b/docker/.env.example index 76131cff..c3fa5544 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -400,4 +400,9 @@ GID='1000' # Disable Swagger API documentation endpoint. # Set to "true" to disable the /api/docs endpoint (recommended for production deployments). -# DISABLE_SWAGGER_DOCS="true" \ No newline at end of file +# DISABLE_SWAGGER_DOCS="true" + +# Disable MCP cooldown timer for agent calls +# this can lead to infinite recursive calls of the same function +# for some model/provider combinations +# MCP_NO_COOLDOWN="true \ No newline at end of file diff --git a/docker/Dockerfile b/docker/Dockerfile index 07a5d63e..ed26c288 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -58,7 +58,7 @@ WORKDIR /app # so web-scraping would be broken in arm docker containers unless we patch it # by manually installing a compatible chromedriver. RUN echo "Need to patch Puppeteer x Chromium support for ARM86 - installing dep!" && \ - curl https://playwright.azureedge.net/builds/chromium/1088/chromium-linux-arm64.zip -o chrome-linux.zip && \ + curl -fSL https://webassets.anythingllm.com/chromium-1088-linux-arm64.zip -o chrome-linux.zip && \ unzip chrome-linux.zip && \ rm -rf chrome-linux.zip @@ -161,10 +161,6 @@ USER anythingllm FROM backend-build AS production-build WORKDIR /app COPY --chown=anythingllm:anythingllm --from=frontend-build /app/frontend/dist /app/server/public -USER root -RUN chown -R anythingllm:anythingllm /app/server && \ - chown -R anythingllm:anythingllm /app/collector -USER anythingllm # Setup the environment ENV NODE_ENV=production diff --git a/server/.env.example b/server/.env.example index 47aef59d..2f7b96db 100644 --- a/server/.env.example +++ b/server/.env.example @@ -403,4 +403,9 @@ TTS_PROVIDER="native" # Disable Swagger API documentation endpoint. # Set to "true" to disable the /api/docs endpoint (recommended for production deployments). -# DISABLE_SWAGGER_DOCS="true" \ No newline at end of file +# DISABLE_SWAGGER_DOCS="true" + +# Disable MCP cooldown timer for agent calls +# this can lead to infinite recursive calls of the same function +# for some model/provider combinations +# MCP_NO_COOLDOWN="true \ No newline at end of file diff --git a/server/package.json b/server/package.json index 95513ac4..0a723c6f 100644 --- a/server/package.json +++ b/server/package.json @@ -53,6 +53,7 @@ "express": "^4.18.2", "extract-json-from-string": "^1.0.1", "fast-levenshtein": "^3.0.0", + "fix-path": "^4.0.0", "graphql": "^16.7.1", "ip": "^2.0.1", "joi": "^17.11.0", @@ -74,6 +75,7 @@ "posthog-node": "^3.1.1", "prisma": "5.3.1", "slugify": "^1.6.6", + "strip-ansi": "^7.1.2", "swagger-autogen": "^2.23.5", "swagger-ui-express": "^5.0.0", "truncate": "^3.0.0", @@ -101,4 +103,4 @@ "nodemon": "^2.0.22", "prettier": "^3.0.3" } -} \ No newline at end of file +} diff --git a/server/utils/MCP/hypervisor/index.js b/server/utils/MCP/hypervisor/index.js index 861274f2..7720bb97 100644 --- a/server/utils/MCP/hypervisor/index.js +++ b/server/utils/MCP/hypervisor/index.js @@ -193,8 +193,9 @@ class MCPHypervisor { this.log(`Pruning MCP server: ${name}`); const mcp = this.mcps[name]; + if (!mcp.transport) return true; const childProcess = mcp.transport._process; - if (childProcess) childProcess.kill(1); + if (childProcess) childProcess.kill("SIGTERM"); mcp.transport.close(); delete this.mcps[name]; @@ -215,10 +216,11 @@ class MCPHypervisor { for (const name of Object.keys(this.mcps)) { if (!this.mcps[name]) continue; const mcp = this.mcps[name]; + if (!mcp.transport) continue; const childProcess = mcp.transport._process; if (childProcess) this.log(`Killing MCP ${name} (PID: ${childProcess.pid})`, { - killed: childProcess.kill(1), + killed: childProcess.kill("SIGTERM"), }); mcp.transport.close(); @@ -228,18 +230,51 @@ class MCPHypervisor { this.mcpLoadingResults = {}; } + /** + * Load shell environment for desktop applications. + * MacOS and Linux don't inherit login shell environment. So this function + * fixes the PATH and accessible commands when running AnythingLLM outside of Docker during development on Mac/Linux and in-container (Linux). + * @returns {Promise<{[key: string]: string}>} - Environment variables from shell + */ + async #loadShellEnvironment() { + try { + if (process.platform === "win32") return process.env; + const { default: fixPath } = await import("fix-path"); + const { default: stripAnsi } = await import("strip-ansi"); + fixPath(); + + // Due to node v20 requirement to have a minimum version of fix-path v5, we need to strip ANSI codes manually + // which was the only patch between v4 and v5. Here we just apply manually. + // https://github.com/sindresorhus/fix-path/issues/6 + if (process.env.PATH) process.env.PATH = stripAnsi(process.env.PATH); + return process.env; + } catch (error) { + console.warn( + "Failed to load shell environment, using process.env:", + error.message + ); + return process.env; + } + } + /** * Build the MCP server environment variables - ensures proper PATH and NODE_PATH * inheritance across all platforms and deployment scenarios. * @param {Object} server - The server definition - * @returns {{env: { [key: string]: string } | {}}} - The environment variables + * @returns {Promise<{env: { [key: string]: string } | {}}}> - The environment variables */ - #buildMCPServerENV(server) { - // Start with essential environment variables, inheriting from current process - // This ensures GUI applications on macOS/Linux get proper PATH inheritance + async #buildMCPServerENV(server) { + const shellEnv = await this.#loadShellEnvironment(); let baseEnv = { - PATH: process.env.PATH || "/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin", - NODE_PATH: process.env.NODE_PATH || "/usr/local/lib/node_modules", + PATH: + shellEnv.PATH || + process.env.PATH || + "/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin", + NODE_PATH: + shellEnv.NODE_PATH || + process.env.NODE_PATH || + "/usr/local/lib/node_modules", + ...shellEnv, // Include all shell environment variables }; // Docker-specific environment setup @@ -273,21 +308,50 @@ class MCPHypervisor { * @returns {MCPServerTypes | null} - The server type */ #parseServerType(server) { - if (server.hasOwnProperty("command")) return "stdio"; - if (server.hasOwnProperty("url")) return "http"; + if ( + server.type === "sse" || + server.type === "streamable" || + server.type === "http" + ) + return "http"; + if (Object.prototype.hasOwnProperty.call(server, "command")) return "stdio"; + if (Object.prototype.hasOwnProperty.call(server, "url")) return "http"; return "sse"; } /** * Validate the server definition by type * - Will throw an error if the server definition is invalid + * @param {string} name - The name of the MCP server * @param {Object} server - The server definition * @param {MCPServerTypes} type - The server type * @returns {void} */ - #validateServerDefinitionByType(server, type) { + #validateServerDefinitionByType(name, server, type) { + if ( + server.type === "sse" || + server.type === "streamable" || + server.type === "http" + ) { + if (!server.url) { + throw new Error( + `MCP server "${name}": missing required "url" for ${server.type} transport` + ); + } + + try { + new URL(server.url); + } catch (error) { + throw new Error(`MCP server "${name}": invalid URL "${server.url}"`); + } + return; + } + if (type === "stdio") { - if (server.hasOwnProperty("args") && !Array.isArray(server.args)) + if ( + Object.prototype.hasOwnProperty.call(server, "args") && + !Array.isArray(server.args) + ) throw new Error("MCP server args must be an array"); } @@ -295,7 +359,6 @@ class MCPHypervisor { if (!["sse", "streamable"].includes(server?.type)) throw new Error("MCP server type must have sse or streamable value."); } - if (type === "sse") return; return; } @@ -304,16 +367,16 @@ class MCPHypervisor { * Setup the server transport by type and server definition * @param {Object} server - The server definition * @param {MCPServerTypes} type - The server type - * @returns {StdioClientTransport | StreamableHTTPClientTransport | SSEClientTransport} - The server transport + * @returns {Promise} - The server transport */ - #setupServerTransport(server, type) { + async #setupServerTransport(server, type) { // if not stdio then it is http or sse if (type !== "stdio") return this.createHttpTransport(server); return new StdioClientTransport({ command: server.command, args: server?.args ?? [], - ...this.#buildMCPServerENV(server), + ...(await this.#buildMCPServerENV(server)), }); } @@ -328,6 +391,7 @@ class MCPHypervisor { // If the server block has a type property then use that to determine the transport type switch (server.type) { case "streamable": + case "http": return new StreamableHTTPClientTransport(url, { requestInit: { headers: server.headers, @@ -354,10 +418,10 @@ class MCPHypervisor { const serverType = this.#parseServerType(server); if (!serverType) throw new Error("MCP server command or url is required"); - this.#validateServerDefinitionByType(server, serverType); + this.#validateServerDefinitionByType(name, server, serverType); this.log(`Attempting to start MCP server: ${name}`); const mcp = new Client({ name: name, version: "1.0.0" }); - const transport = this.#setupServerTransport(server, serverType); + const transport = await this.#setupServerTransport(server, serverType); // Add connection event listeners transport.onclose = () => this.log(`${name} - Transport closed`); @@ -369,10 +433,22 @@ class MCPHypervisor { // Connect and await the connection with a timeout this.mcps[name] = mcp; const connectionPromise = mcp.connect(transport); + + let timeoutId; const timeoutPromise = new Promise((_, reject) => { - setTimeout(() => reject(new Error("Connection timeout")), 30_000); // 30 second timeout + timeoutId = setTimeout( + () => reject(new Error("Connection timeout")), + 30_000 + ); // 30 second timeout }); - await Promise.race([connectionPromise, timeoutPromise]); + + try { + await Promise.race([connectionPromise, timeoutPromise]); + if (timeoutId) clearTimeout(timeoutId); + } catch (error) { + if (timeoutId) clearTimeout(timeoutId); + throw error; + } return true; } diff --git a/server/utils/MCP/index.js b/server/utils/MCP/index.js index 03e72a86..8ac7d7d8 100644 --- a/server/utils/MCP/index.js +++ b/server/utils/MCP/index.js @@ -29,8 +29,15 @@ class MCPCompatibilityLayer extends MCPHypervisor { const mcp = this.mcps[name]; if (!mcp) return null; - const tools = (await mcp.listTools()).tools; - if (!tools.length) return null; + let tools; + try { + const response = await mcp.listTools(); + tools = response.tools; + } catch (error) { + this.log(`Failed to list tools for MCP server ${name}:`, error); + return null; + } + if (!tools || !tools.length) return null; const plugins = []; for (const tool of tools) { @@ -46,6 +53,7 @@ class MCPCompatibilityLayer extends MCPHypervisor { name: `${name}-${tool.name}`, controller: new AbortController(), description: tool.description, + isMCPTool: true, examples: [], parameters: { $schema: "http://json-schema.org/draft-07/schema#", @@ -53,6 +61,13 @@ class MCPCompatibilityLayer extends MCPHypervisor { }, handler: async function (args = {}) { try { + const mcpLayer = new MCPCompatibilityLayer(); + const currentMcp = mcpLayer.mcps[name]; + if (!currentMcp) + throw new Error( + `MCP server ${name} is not currently running` + ); + aibitat.handlerProps.log( `Executing MCP server: ${name}:${tool.name} with args:`, args @@ -60,7 +75,7 @@ class MCPCompatibilityLayer extends MCPHypervisor { aibitat.introspect( `Executing MCP server: ${name} with ${JSON.stringify(args, null, 2)}` ); - const result = await mcp.callTool({ + const result = await currentMcp.callTool({ name: tool.name, arguments: args, }); @@ -71,9 +86,7 @@ class MCPCompatibilityLayer extends MCPHypervisor { aibitat.introspect( `MCP server: ${name}:${tool.name} completed successfully` ); - return typeof result === "object" - ? JSON.stringify(result) - : String(result); + return MCPCompatibilityLayer.returnMCPResult(result); } catch (error) { aibitat.handlerProps.log( `MCP server: ${name}:${tool.name} failed with error:`, @@ -134,7 +147,9 @@ class MCPCompatibilityLayer extends MCPHypervisor { } const online = !!(await mcp.ping()); - const tools = online ? (await mcp.listTools()).tools : []; + const tools = (online ? (await mcp.listTools()).tools : []).filter( + (tool) => !tool.name.startsWith("handle_mcp_connection_mcp_") + ); servers.push({ name, config: config?.server || null, @@ -199,5 +214,29 @@ class MCPCompatibilityLayer extends MCPHypervisor { this.log(`MCP server was killed and removed from config file: ${name}`); return { success: true, error: null }; } + + /** + * Return the result of an MCP server call as a string + * This will handle circular references and bigints since an MCP server can return any type of data. + * @param {Object} result - The result to return + * @returns {string} The result as a string + */ + static returnMCPResult(result) { + if (typeof result !== "object" || result === null) return String(result); + + const seen = new WeakSet(); + try { + return JSON.stringify(result, (key, value) => { + if (typeof value === "bigint") return value.toString(); + if (typeof value === "object" && value !== null) { + if (seen.has(value)) return "[Circular]"; + seen.add(value); + } + return value; + }); + } catch (e) { + return `[Unserializable: ${e.message}]`; + } + } } module.exports = MCPCompatibilityLayer; diff --git a/server/utils/agents/aibitat/plugins/memory.js b/server/utils/agents/aibitat/plugins/memory.js index b229af3c..1856de39 100644 --- a/server/utils/agents/aibitat/plugins/memory.js +++ b/server/utils/agents/aibitat/plugins/memory.js @@ -67,7 +67,11 @@ const memory = { }, handler: async function ({ action = "", content = "" }) { try { - if (this.tracker.isDuplicate(this.name, { action, content })) + const { isDuplicate } = this.tracker.isDuplicate(this.name, { + action, + content, + }); + if (isDuplicate) return `This was a duplicated call and it's output will be ignored.`; let response = "There was nothing to do."; diff --git a/server/utils/agents/aibitat/plugins/rechart.js b/server/utils/agents/aibitat/plugins/rechart.js index a41ddd65..b9b4ff73 100644 --- a/server/utils/agents/aibitat/plugins/rechart.js +++ b/server/utils/agents/aibitat/plugins/rechart.js @@ -55,9 +55,9 @@ Make sure the format use double quotes and property names are string literals. P required: ["type", "title", "dataset"], handler: async function ({ type, dataset, title }) { try { - if (!this.tracker.isUnique(this.name)) { + if (this.tracker.isMarkedUnique(this.name)) { this.super.handlerProps.log( - `${this.name} has been run for this chat response already. It can only be called once per chat.` + `${this.name} has been called for this chat response already. It can only be called once per chat.` ); return "The chart was generated and returned to the user. This function completed successfully. Do not call this function again."; } diff --git a/server/utils/agents/aibitat/plugins/save-file-browser.js b/server/utils/agents/aibitat/plugins/save-file-browser.js index b0ef039f..c319b42b 100644 --- a/server/utils/agents/aibitat/plugins/save-file-browser.js +++ b/server/utils/agents/aibitat/plugins/save-file-browser.js @@ -60,11 +60,13 @@ const saveFileInBrowser = { }, handler: async function ({ file_content = "", filename }) { try { - if ( - this.tracker.isDuplicate(this.name, { file_content, filename }) - ) { + const { isDuplicate, reason } = this.tracker.isDuplicate( + this.name, + { file_content, filename } + ); + if (isDuplicate) { this.super.handlerProps.log( - `${this.name} was called, but exited early since it was not a unique call.` + `${this.name} was called, but exited early because ${reason}.` ); return `${filename} file has been saved successfully!`; } diff --git a/server/utils/agents/aibitat/providers/helpers/untooled.js b/server/utils/agents/aibitat/providers/helpers/untooled.js index 2da3e4f3..ee4dfd88 100644 --- a/server/utils/agents/aibitat/providers/helpers/untooled.js +++ b/server/utils/agents/aibitat/providers/helpers/untooled.js @@ -45,6 +45,30 @@ ${JSON.stringify(def.parameters.properties, null, 4)}\n`; return output; } + /** + * Check if a function call is an MCP tool. + * We do this because some MCP tools dont return values and will cause infinite loops in calling for Untooled to call the same function over and over again. + * Any MCP tool is automatically marked with a cooldown to prevent infinite loops of the same function over and over again. + * + * This can lead to unexpected behavior if you want a model using Untooled to call a repeat action multiple times. + * eg: Create 3 Jira tickets about x, y, and z. -> will skip y and z if you don't disable the cooldown. + * + * You can disable this check by setting the `MCP_NO_COOLDOWN` flag to any value in the ENV. + * + * @param {{name: string, arguments: Object}} functionCall - The function call to check. + * @param {Object[]} functions - The list of functions definitions to check against. + * @return {boolean} - True if the function call is an MCP tool, false otherwise. + */ + isMCPTool(functionCall = {}, functions = []) { + if (process.env.MCP_NO_COOLDOWN) return false; + + const foundFunc = functions.find( + (def) => def?.name?.toLowerCase() === functionCall.name?.toLowerCase() + ); + if (!foundFunc) return false; + return foundFunc?.isMCPTool || false; + } + /** * Validate a function call against a list of functions. * @param {{name: string, arguments: Object}} functionCall - The function call to validate. @@ -135,9 +159,11 @@ ${JSON.stringify(def.parameters.properties, null, 4)}\n`; return { toolCall: null, text: null }; } - if (this.deduplicator.isDuplicate(call.name, call.arguments)) { + const { isDuplicate, reason: duplicateReason } = + this.deduplicator.isDuplicate(call.name, call.arguments); + if (isDuplicate) { this.providerLog( - `Function tool with exact arguments has already been called this stack.` + `Cannot call ${call.name} again because ${duplicateReason}.` ); return { toolCall: null, text: null }; } @@ -197,9 +223,11 @@ ${JSON.stringify(def.parameters.properties, null, 4)}\n`; return { toolCall: null, text: null, uuid: msgUUID }; } - if (this.deduplicator.isDuplicate(call.name, call.arguments)) { + const { isDuplicate, reason: duplicateReason } = + this.deduplicator.isDuplicate(call.name, call.arguments); + if (isDuplicate) { this.providerLog( - `Function tool with exact arguments has already been called this stack.` + `Cannot call ${call.name} again because ${duplicateReason}.` ); eventHandler?.("reportStreamEvent", { type: "removeStatusResponse", @@ -251,7 +279,9 @@ ${JSON.stringify(def.parameters.properties, null, 4)}\n`; if (toolCall !== null) { this.providerLog(`Valid tool call found - running ${toolCall.name}.`); - this.deduplicator.trackRun(toolCall.name, toolCall.arguments); + this.deduplicator.trackRun(toolCall.name, toolCall.arguments, { + cooldown: this.isMCPTool(toolCall, functions), + }); return { result: null, functionCall: { @@ -349,7 +379,9 @@ ${JSON.stringify(def.parameters.properties, null, 4)}\n`; if (toolCall !== null) { this.providerLog(`Valid tool call found - running ${toolCall.name}.`); - this.deduplicator.trackRun(toolCall.name, toolCall.arguments); + this.deduplicator.trackRun(toolCall.name, toolCall.arguments, { + cooldown: this.isMCPTool(toolCall, functions), + }); return { result: null, functionCall: { diff --git a/server/utils/agents/aibitat/providers/ollama.js b/server/utils/agents/aibitat/providers/ollama.js index 65a914fd..e39a6704 100644 --- a/server/utils/agents/aibitat/providers/ollama.js +++ b/server/utils/agents/aibitat/providers/ollama.js @@ -146,9 +146,11 @@ class OllamaProvider extends InheritMultiple([Provider, UnTooled]) { return { toolCall: null, text: null, uuid: msgUUID }; } - if (this.deduplicator.isDuplicate(call.name, call.arguments)) { + const { isDuplicate, reason: duplicateReason } = + this.deduplicator.isDuplicate(call.name, call.arguments); + if (isDuplicate) { this.providerLog( - `Function tool with exact arguments has already been called this stack.` + `Cannot call ${call.name} again because ${duplicateReason}.` ); eventHandler?.("reportStreamEvent", { type: "removeStatusResponse", @@ -197,7 +199,9 @@ class OllamaProvider extends InheritMultiple([Provider, UnTooled]) { if (toolCall !== null) { this.providerLog(`Valid tool call found - running ${toolCall.name}.`); - this.deduplicator.trackRun(toolCall.name, toolCall.arguments); + this.deduplicator.trackRun(toolCall.name, toolCall.arguments, { + cooldown: this.isMCPTool(toolCall, functions), + }); return { result: null, functionCall: { @@ -314,7 +318,9 @@ class OllamaProvider extends InheritMultiple([Provider, UnTooled]) { if (toolCall !== null) { this.providerLog(`Valid tool call found - running ${toolCall.name}.`); - this.deduplicator.trackRun(toolCall.name, toolCall.arguments); + this.deduplicator.trackRun(toolCall.name, toolCall.arguments, { + cooldown: this.isMCPTool(toolCall, functions), + }); return { result: null, functionCall: { diff --git a/server/utils/agents/aibitat/utils/dedupe.js b/server/utils/agents/aibitat/utils/dedupe.js index 321f6929..366559d8 100644 --- a/server/utils/agents/aibitat/utils/dedupe.js +++ b/server/utils/agents/aibitat/utils/dedupe.js @@ -12,9 +12,9 @@ // Track Run/isDuplicate prevents _exact_ data re-runs based on the SHA of their inputs // StartCooldown/isOnCooldown does prevention of _near-duplicate_ runs based on only the function name that is running. -// isUnique/markUnique/removeUniqueConstraint prevents one-time functions from re-running. EG: charting. +// isMarkedUnique/markUnique/removeUniqueConstraint prevents one-time functions from re-running. EG: charting. const crypto = require("crypto"); -const DEFAULT_COOLDOWN_MS = 5 * 1000; +const DEFAULT_COOLDOWN_MS = 30 * 1000; class Deduplicator { #hashes = {}; @@ -22,20 +22,59 @@ class Deduplicator { #uniques = {}; constructor() {} - trackRun(key, params = {}) { + log(message, ...args) { + console.log(`\x1b[36m[Deduplicator]\x1b[0m ${message}`, ...args); + } + + trackRun( + key, + params = {}, + options = { + cooldown: false, + cooldownInMs: DEFAULT_COOLDOWN_MS, + markUnique: false, + } + ) { const hash = crypto .createHash("sha256") .update(JSON.stringify({ key, params })) .digest("hex"); this.#hashes[hash] = Number(new Date()); + if (options.cooldown) + this.startCooldown(key, { cooldownInMs: options.cooldownInMs }); + if (options.markUnique) this.markUnique(key); } + /** + * Checks if a key and params are: + * - exactly the same as a previous run. + * - on cooldown. + * - marked as unique. + * @param {string} key - The key to check. + * @param {Object} params - The parameters to check. + * @returns {{isDuplicate: boolean, reason: string}} - The result of the check. + */ isDuplicate(key, params = {}) { const newSig = crypto .createHash("sha256") .update(JSON.stringify({ key, params })) .digest("hex"); - return this.#hashes.hasOwnProperty(newSig); + if (this.#hashes.hasOwnProperty(newSig)) + return { + isDuplicate: true, + reason: `an exact duplicate of previous run of ${key}`, + }; + if (this.isOnCooldown(key)) + return { + isDuplicate: true, + reason: `the function is on cooldown for ${key}.`, + }; + if (this.isMarkedUnique(key)) + return { + isDuplicate: true, + reason: `the function is marked as unique for ${key}. Can only be called once per agent session.`, + }; + return { isDuplicate: false, reason: "" }; } /** @@ -57,28 +96,54 @@ class Deduplicator { return; } + /** + * Starts a cooldown for a key. + * @param {string} key - The key to start the cooldown for (string key of the function name). + * @param {Object} parameters - The parameters for the cooldown. + * @param {number} parameters.cooldownInMs - The cooldown in milliseconds. + */ startCooldown( key, parameters = { cooldownInMs: DEFAULT_COOLDOWN_MS, } ) { - this.#cooldowns[key] = Number(new Date()) + Number(parameters.cooldownInMs); + const cooldownDelay = parameters.cooldownInMs || DEFAULT_COOLDOWN_MS; + this.log(`Starting cooldown for ${key} for ${cooldownDelay}ms`); + this.#cooldowns[key] = Number(new Date()) + Number(cooldownDelay); } + /** + * Checks if a key is on cooldown. + * @param {string} key - The key to check. + * @returns {boolean} - True if the key is on cooldown, false otherwise. + */ isOnCooldown(key) { if (!this.#cooldowns.hasOwnProperty(key)) return false; return Number(new Date()) <= this.#cooldowns[key]; } - isUnique(key) { - return !this.#uniques.hasOwnProperty(key); + /** + * Checks if a key is marked as unique and currently tracked by the deduplicator. + * @param {string} key - The key to check. + * @returns {boolean} - True if the key is marked as unique, false otherwise. + */ + isMarkedUnique(key) { + return this.#uniques.hasOwnProperty(key); } + /** + * Removes the unique constraint for a key. + * @param {string} key - The key to remove the unique constraint for. + */ removeUniqueConstraint(key) { delete this.#uniques[key]; } + /** + * Marks a key as unique and currently tracked by the deduplicator. + * @param {string} key - The key to mark as unique. + */ markUnique(key) { this.#uniques[key] = Number(new Date()); } diff --git a/server/utils/helpers/updateENV.js b/server/utils/helpers/updateENV.js index 43b48794..a170a9f3 100644 --- a/server/utils/helpers/updateENV.js +++ b/server/utils/helpers/updateENV.js @@ -1248,6 +1248,9 @@ function dumpENV() { // Allow setting a custom response timeout for Ollama "OLLAMA_RESPONSE_TIMEOUT", + + // Allow disabling of MCP tool cooldown + "MCP_NO_COOLDOWN", ]; // Simple sanitization of each value to prevent ENV injection via newline or quote escaping. diff --git a/server/yarn.lock b/server/yarn.lock index 68a798d7..529331d9 100644 --- a/server/yarn.lock +++ b/server/yarn.lock @@ -3693,6 +3693,11 @@ ansi-regex@^5.0.1: resolved "https://registry.npmjs.org/ansi-regex/-/ansi-regex-5.0.1.tgz" integrity sha512-quJQXlTSUGL2LH9SUXo8VwsY4soanhgo6LNSm84E1LBcE8s3O0wpdiRzyR9z/ZZJMlMWv37qOOb9pdJlMUEKFQ== +ansi-regex@^6.0.1: + version "6.2.2" + resolved "https://registry.yarnpkg.com/ansi-regex/-/ansi-regex-6.2.2.tgz#60216eea464d864597ce2832000738a0589650c1" + integrity sha512-Bq3SmSpyFHaWjPk8If9yc6svM8c56dB5BAtW4Qbw5jHTwwXXcTLoRMkpDJp6VL0XzlWaCHTXrkFURMYmD0sLqg== + ansi-styles@^4.0.0, ansi-styles@^4.1.0: version "4.3.0" resolved "https://registry.npmjs.org/ansi-styles/-/ansi-styles-4.3.0.tgz" @@ -4656,6 +4661,11 @@ deepmerge@^4.2.2: resolved "https://registry.npmjs.org/deepmerge/-/deepmerge-4.3.1.tgz" integrity sha512-3sUqbMEc77XqpdNO7FRyRog+eW3ph+GYCbj+rK+uYyRMuwsVy0rMiVtPn+QJlKFvWP/1PYpapqYn0Me2knFn+A== +default-shell@^2.0.0: + version "2.2.0" + resolved "https://registry.yarnpkg.com/default-shell/-/default-shell-2.2.0.tgz#31481c19747bfe59319b486591643eaf115a1864" + integrity sha512-sPpMZcVhRQ0nEMDtuMJ+RtCxt7iHPAMBU+I4tAlo5dU1sjRpNax0crj6nR3qKpvVnckaQ9U38enXcwW9nZJeCw== + define-data-property@^1.0.1, define-data-property@^1.1.4: version "1.1.4" resolved "https://registry.npmjs.org/define-data-property/-/define-data-property-1.1.4.tgz" @@ -5454,6 +5464,13 @@ find-up@^5.0.0: locate-path "^6.0.0" path-exists "^4.0.0" +fix-path@^4.0.0: + version "4.0.0" + resolved "https://registry.yarnpkg.com/fix-path/-/fix-path-4.0.0.tgz#bc1d14f038edb734ac46944a45454106952ca429" + integrity sha512-g31GX207Tt+psI53ZSaB1egprYbEN0ZYl90aKcO22A2LmCNnFsSq3b5YpoKp3E/QEiWByTXGJOkFQG4S07Bc1A== + dependencies: + shell-path "^3.0.0" + flat-cache@^3.0.4: version "3.2.0" resolved "https://registry.npmjs.org/flat-cache/-/flat-cache-3.2.0.tgz" @@ -8129,6 +8146,22 @@ shebang-regex@^3.0.0: resolved "https://registry.npmjs.org/shebang-regex/-/shebang-regex-3.0.0.tgz" integrity sha512-7++dFhtcx3353uBaq8DDR4NuxBetBzC7ZQOhmTQInHEd6bSrXdiEyzCvG07Z44UYdLShWUyXt5M/yhz8ekcb1A== +shell-env@^4.0.1: + version "4.0.1" + resolved "https://registry.yarnpkg.com/shell-env/-/shell-env-4.0.1.tgz#883302d9426095d398a39b102a851adb306b8cb8" + integrity sha512-w3oeZ9qg/P6Lu6qqwavvMnB/bwfsz67gPB3WXmLd/n6zuh7TWQZtGa3iMEdmua0kj8rivkwl+vUjgLWlqZOMPw== + dependencies: + default-shell "^2.0.0" + execa "^5.1.1" + strip-ansi "^7.0.1" + +shell-path@^3.0.0: + version "3.1.0" + resolved "https://registry.yarnpkg.com/shell-path/-/shell-path-3.1.0.tgz#950671fe15de70fb4d984b886d55e8a2f10bfe33" + integrity sha512-s/9q9PEtcRmDTz69+cJ3yYBAe9yGrL7e46gm2bU4pQ9N48ecPK9QrGFnLwYgb4smOHskx4PL7wCNMktW2AoD+g== + dependencies: + shell-env "^4.0.1" + side-channel-list@^1.0.0: version "1.0.0" resolved "https://registry.yarnpkg.com/side-channel-list/-/side-channel-list-1.0.0.tgz#10cb5984263115d3b7a0e336591e290a830af8ad" @@ -8353,6 +8386,13 @@ strip-ansi@^6.0.0, strip-ansi@^6.0.1: dependencies: ansi-regex "^5.0.1" +strip-ansi@^7.0.1, strip-ansi@^7.1.2: + version "7.1.2" + resolved "https://registry.yarnpkg.com/strip-ansi/-/strip-ansi-7.1.2.tgz#132875abde678c7ea8d691533f2e7e22bb744dba" + integrity sha512-gmBGslpoQJtgnMAvOVqGZpEz9dyoKTCzy2nfz/n8aIFhN/jCE/rCmcxabB6jOOHV+0WNnylOxaxBQPSvcWklhA== + dependencies: + ansi-regex "^6.0.1" + strip-final-newline@^2.0.0: version "2.0.0" resolved "https://registry.npmjs.org/strip-final-newline/-/strip-final-newline-2.0.0.tgz"