Use custom fetch + treat results as untrusted

Author: nlarew

package-lock.json

@@ -73,6 +73,7 @@
     "@types/yargs-parser": "^21.0.3",
     "@typescript-eslint/parser": "^8.44.0",
     "@vitest/coverage-v8": "^3.2.4",
+    "@vitest/eslint-plugin": "^1.3.4",
     "ai": "^4.3.17",
     "duplexpair": "^1.0.2",
     "eslint": "^9.34.0",
@@ -92,7 +93,6 @@
     "tsx": "^4.20.5",
     "typescript": "^5.9.2",
     "typescript-eslint": "^8.41.0",
-    "@vitest/eslint-plugin": "^1.3.4",
     "uuid": "^13.0.0",
     "vitest": "^3.2.4"
   },
@@ -114,6 +114,7 @@
     "oauth4webapi": "^3.8.0",
     "openapi-fetch": "^0.14.0",
     "ts-levenshtein": "^1.0.7",
+    "yaml": "^2.8.1",
     "yargs-parser": "^22.0.0",
     "zod": "^3.25.76"
   },

package.json

@@ -5,9 +5,10 @@ import {
     type ToolCategory,
     type ToolConstructorParams,
 } from "../tool.js";
-import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { createFetch } from "@mongodb-js/devtools-proxy-support";
 import { Server } from "../../server.js";
 import { packageInfo } from "../../common/packageInfo.js";
+import { formatUntrustedData } from "../tool.js";
 
 export abstract class AssistantToolBase extends ToolBase {
     protected server?: Server;
@@ -18,10 +19,9 @@ export abstract class AssistantToolBase extends ToolBase {
     constructor({ session, config, telemetry, elicitation }: ToolConstructorParams) {
         super({ session, config, telemetry, elicitation });
         this.baseUrl = new URL(config.assistantBaseUrl);
-        const serverVersion = packageInfo.version;
         this.requiredHeaders = new Headers({
             "x-request-origin": "mongodb-mcp-server",
-            "user-agent": serverVersion ? `mongodb-mcp-server/v${serverVersion}` : "mongodb-mcp-server",
+            "user-agent": packageInfo.version ? `mongodb-mcp-server/v${packageInfo.version}` : "mongodb-mcp-server",
         });
     }
 
@@ -42,7 +42,14 @@ export abstract class AssistantToolBase extends ToolBase {
         if (args.method === "POST") {
             headers.set("Content-Type", "application/json");
         }
-        return await fetch(endpoint, {
+
+        // Use the same custom fetch implementation as the Atlas API client.
+        // We need this to support enterprise proxies.
+        const customFetch = createFetch({
+            useEnvironmentVariableProxies: true,
+        }) as unknown as typeof fetch;
+
+        return await customFetch(endpoint, {
             method: args.method,
             headers,
             body: JSON.stringify(args.body),

src/tools/assistant/assistantTool.ts

@@ -1,25 +1,27 @@
-import { z } from "zod";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
-import { OperationType } from "../tool.js";
+import { formatUntrustedData, type OperationType } from "../tool.js";
 import { AssistantToolBase } from "./assistantTool.js";
 import { LogId } from "../../common/logger.js";
+import { stringify as yamlStringify } from "yaml";
 
-export type ListKnowledgeSourcesResponse = {
-    dataSources: {
-        /** The name of the data source */
-        id: string;
-        /** The type of the data source */
-        type: string;
-        /** A list of available versions for this data source */
-        versions: {
-            /** The version label of the data source */
-            label: string;
-            /** Whether this version is the current/default version */
-            isCurrent: boolean;
-        }[];
+export type KnowledgeSource = {
+    /** The name of the data source */
+    id: string;
+    /** The type of the data source */
+    type: string;
+    /** A list of available versions for this data source */
+    versions: {
+        /** The version label of the data source */
+        label: string;
+        /** Whether this version is the current/default version */
+        isCurrent: boolean;
     }[];
 };
 
+export type ListKnowledgeSourcesResponse = {
+    dataSources: KnowledgeSource[];
+};
+
 export class ListKnowledgeSourcesTool extends AssistantToolBase {
     public name = "list-knowledge-sources";
     protected description = "List available data sources in the MongoDB Assistant knowledge base";
@@ -50,15 +52,21 @@ export class ListKnowledgeSourcesTool extends AssistantToolBase {
         }
         const { dataSources } = (await response.json()) as ListKnowledgeSourcesResponse;
 
+        const text = yamlStringify(
+            dataSources.map((ds) => {
+                const currentVersion = ds.versions.find(({ isCurrent }) => isCurrent)?.label;
+                if (currentVersion) {
+                    (ds as KnowledgeSource & { currentVersion: string }).currentVersion = currentVersion;
+                }
+                return ds;
+            })
+        );
+
         return {
-            content: dataSources.map(({ id, type, versions }) => ({
-                type: "text",
-                text: id,
-                _meta: {
-                    type,
-                    versions,
-                },
-            })),
+            content: formatUntrustedData(
+                `Found ${dataSources.length} data sources in the MongoDB Assistant knowledge base.`,
+                text
+            ),
         };
     }
 }

src/tools/assistant/listKnowledgeSources.ts

@@ -1,8 +1,9 @@
 import { z } from "zod";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
-import { ToolArgs, OperationType } from "../tool.js";
+import { type ToolArgs, type OperationType, formatUntrustedData } from "../tool.js";
 import { AssistantToolBase } from "./assistantTool.js";
 import { LogId } from "../../common/logger.js";
+import { stringify as yamlStringify } from "yaml";
 
 export const SearchKnowledgeToolArgs = {
     query: z.string().describe("A natural language query to search for in the knowledge base"),
@@ -71,14 +72,14 @@ export class SearchKnowledgeTool extends AssistantToolBase {
             };
         }
         const { results } = (await response.json()) as SearchKnowledgeResponse;
+
+        const text = yamlStringify(results);
+
         return {
-            content: results.map(({ text, metadata }) => ({
-                type: "text",
-                text,
-                _meta: {
-                    ...metadata,
-                },
-            })),
+            content: formatUntrustedData(
+                `Found ${results.length} results in the MongoDB Assistant knowledge base.`,
+                text
+            ),
         };
     }
 }