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
2026-04-25 17:15:37 +02:00 · 2025-12-05 11:01:20 -08:00
parent b54ac2de58
commit 92752a2a15
15 changed files with 370 additions and 95 deletions
--- 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
+      # - 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
--- 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"
+# 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
--- 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
--- 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"
+# 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
--- 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"
  }
-}
+}
--- 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<StdioClientTransport | StreamableHTTPClientTransport | SSEClientTransport>} - 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;
  }

--- 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;
--- 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.";
--- 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.";
              }
--- 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!`;
              }
--- 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: {
--- 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: {
--- 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());
  }
--- 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.
--- 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"