Add MongoDB Assistant Tools

README.md

@@ -14,6 +14,7 @@ A Model Context Protocol server for interacting with MongoDB Databases and Mongo
 - [🛠️ Supported Tools](#supported-tools)
   - [MongoDB Atlas Tools](#mongodb-atlas-tools)
   - [MongoDB Database Tools](#mongodb-database-tools)
+  - [MongoDB Assistant Tools](#mongodb-assistant-tools)
 - [📄 Supported Resources](#supported-resources)
 - [⚙️ Configuration](#configuration)
   - [Configuration Options](#configuration-options)
@@ -320,6 +321,11 @@ NOTE: atlas tools are only available when you set credentials on [configuration]
 - `db-stats` - Return statistics about a MongoDB database
 - `export` - Export query or aggregation results to EJSON format. Creates a uniquely named export accessible via the `exported-data` resource.
 
+#### MongoDB Assistant Tools
+
+- `list-knowledge-sources` - List available data sources in the MongoDB Assistant knowledge base. Example sources include various MongoDB documentation sites, MongoDB University courses, and other useful learning resources.
+- `search-knowledge` - Search for information in the MongoDB Assistant knowledge base.
+
 ## 📄 Supported Resources
 
 - `config` - Server configuration, supplied by the user either as environment variables or as startup arguments with sensitive parameters redacted. The resource can be accessed under URI `config://config`.

package.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"
   },

src/common/config.ts

@@ -160,6 +160,7 @@ export interface UserConfig extends CliOptions {
     apiBaseUrl: string;
     apiClientId?: string;
     apiClientSecret?: string;
+    assistantBaseUrl: string;
     telemetry: "enabled" | "disabled";
     logPath: string;
     exportsPath: string;
@@ -185,6 +186,7 @@ export interface UserConfig extends CliOptions {
 
 export const defaultUserConfig: UserConfig = {
     apiBaseUrl: "https://cloud.mongodb.com/",
+    assistantBaseUrl: "https://knowledge.mongodb.com/api/v1/",
     logPath: getLogPath(),
     exportsPath: getExportsPath(),
     exportTimeoutMs: 5 * 60 * 1000, // 5 minutes

src/common/logger.ts

@@ -70,6 +70,9 @@ export const LogId = {
     exportLockError: mongoLogId(1_007_008),
 
     oidcFlow: mongoLogId(1_008_001),
+
+    assistantListKnowledgeSourcesError: mongoLogId(1_009_001),
+    assistantSearchKnowledgeError: mongoLogId(1_009_002),
 } as const;
 
 export interface LogPayload {

src/server.ts

@@ -19,6 +19,7 @@ import {
 } from "@modelcontextprotocol/sdk/types.js";
 import assert from "assert";
 import type { ToolBase } from "./tools/tool.js";
+import { AssistantTools } from "./tools/assistant/tools.js";
 import { validateConnectionString } from "./helpers/connectionOptions.js";
 import { packageInfo } from "./common/packageInfo.js";
 import { type ConnectionErrorHandler } from "./common/connectionErrorHandler.js";
@@ -206,7 +207,7 @@ export class Server {
     }
 
     private registerTools(): void {
-        for (const toolConstructor of [...AtlasTools, ...MongoDbTools]) {
+        for (const toolConstructor of [...AtlasTools, ...MongoDbTools, ...AssistantTools]) {
             const tool = new toolConstructor({
                 session: this.session,
                 config: this.userConfig,
@@ -302,6 +303,7 @@ export class Server {
                     context: "server",
                     message: `Failed to connect to MongoDB instance using the connection string from the config: ${error instanceof Error ? error.message : String(error)}`,
                 });
+                j;
             }
         }
     }

src/tools/assistant/assistantTool.ts

@@ -0,0 +1,58 @@
+import {
+    ToolBase,
+    type TelemetryToolMetadata,
+    type ToolArgs,
+    type ToolCategory,
+    type ToolConstructorParams,
+} from "../tool.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;
+    public category: ToolCategory = "assistant";
+    protected baseUrl: URL;
+    protected requiredHeaders: Headers;
+
+    constructor({ session, config, telemetry, elicitation }: ToolConstructorParams) {
+        super({ session, config, telemetry, elicitation });
+        this.baseUrl = new URL(config.assistantBaseUrl);
+        this.requiredHeaders = new Headers({
+            "x-request-origin": "mongodb-mcp-server",
+            "user-agent": packageInfo.version ? `mongodb-mcp-server/v${packageInfo.version}` : "mongodb-mcp-server",
+        });
+    }
+
+    public register(server: Server): boolean {
+        this.server = server;
+        return super.register(server);
+    }
+
+    protected resolveTelemetryMetadata(_args: ToolArgs<typeof this.argsShape>): TelemetryToolMetadata {
+        // Assistant tool calls are not associated with a specific project or organization
+        // Therefore, we don't have any values to add to the telemetry metadata
+        return {};
+    }
+
+    protected async callAssistantApi(args: { method: "GET" | "POST"; endpoint: string; body?: unknown }) {
+        const endpoint = new URL(args.endpoint, this.baseUrl);
+        const headers = new Headers(this.requiredHeaders);
+        if (args.method === "POST") {
+            headers.set("Content-Type", "application/json");
+        }
+
+        // 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/listKnowledgeSources.ts

@@ -0,0 +1,74 @@
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.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 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 const ListKnowledgeSourcesToolName = "list-knowledge-sources";
+
+export class ListKnowledgeSourcesTool extends AssistantToolBase {
+    public name = ListKnowledgeSourcesToolName;
+    protected description = "List available data sources in the MongoDB Assistant knowledge base";
+    protected argsShape = {};
+    public operationType: OperationType = "read";
+
+    protected async execute(): Promise<CallToolResult> {
+        const response = await this.callAssistantApi({
+            method: "GET",
+            endpoint: "content/sources",
+        });
+        if (!response.ok) {
+            const message = `Failed to list knowledge sources: ${response.statusText}`;
+            this.session.logger.debug({
+                id: LogId.assistantListKnowledgeSourcesError,
+                context: "assistant-list-knowledge-sources",
+                message,
+            });
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: message,
+                    },
+                ],
+                isError: true,
+            };
+        }
+        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: formatUntrustedData(
+                `Found ${dataSources.length} data sources in the MongoDB Assistant knowledge base.`,
+                text
+            ),
+        };
+    }
+}

src/tools/assistant/searchKnowledge.ts

@@ -0,0 +1,90 @@
+import { z } from "zod";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.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";
+import { ListKnowledgeSourcesToolName } from "./listKnowledgeSources.js";
+
+export const SearchKnowledgeToolArgs = {
+    query: z
+        .string()
+        .describe(
+            "A natural language query to search for in the MongoDB Assistant knowledge base. This should be a single question or a topic that is relevant to the user's MongoDB use case."
+        ),
+    limit: z.number().min(1).max(100).optional().default(5).describe("The maximum number of results to return"),
+    dataSources: z
+        .array(
+            z.object({
+                name: z.string().describe("The name of the data source"),
+                versionLabel: z.string().optional().describe("The version label of the data source"),
+            })
+        )
+        .optional()
+        .describe(
+            `A list of one or more data sources to limit the search to. You can specify a specific version of a data source by providing the version label. If not provided, the latest version of all data sources will be searched. Available data sources and their versions can be listed by calling the ${ListKnowledgeSourcesToolName} tool.`
+        ),
+};
+
+export type SearchKnowledgeResponse = {
+    /** A list of search results */
+    results: {
+        /** The URL of the search result */
+        url: string;
+        /** The page title of the search result */
+        title: string;
+        /** The text of the page chunk returned from the search */
+        text: string;
+        /** Metadata for the search result */
+        metadata: {
+            /** A list of tags that describe the page */
+            tags: string[];
+            /** Additional metadata */
+            [key: string]: unknown;
+        };
+    }[];
+};
+
+export class SearchKnowledgeTool extends AssistantToolBase {
+    public name = "search-knowledge";
+    protected description = "Search for information in the MongoDB Assistant knowledge base";
+    protected argsShape = {
+        ...SearchKnowledgeToolArgs,
+    };
+    public operationType: OperationType = "read";
+
+    protected async execute(args: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
+        const response = await this.callAssistantApi({
+            method: "POST",
+            endpoint: "content/search",
+            body: args,
+        });
+        if (!response.ok) {
+            const message = `Failed to search knowledge base: ${response.statusText}`;
+            this.session.logger.debug({
+                id: LogId.assistantSearchKnowledgeError,
+                context: "assistant-search-knowledge",
+                message,
+            });
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: message,
+                    },
+                ],
+                isError: true,
+            };
+        }
+        const { results } = (await response.json()) as SearchKnowledgeResponse;
+
+        const text = yamlStringify(results);
+
+        return {
+            content: formatUntrustedData(
+                `Found ${results.length} results in the MongoDB Assistant knowledge base.`,
+                text
+            ),
+        };
+    }
+}

src/tools/assistant/tools.ts

@@ -0,0 +1,4 @@
+import { ListKnowledgeSourcesTool } from "./listKnowledgeSources.js";
+import { SearchKnowledgeTool } from "./searchKnowledge.js";
+
+export const AssistantTools = [ListKnowledgeSourcesTool, SearchKnowledgeTool];

src/tools/tool.ts

@@ -16,7 +16,7 @@ export type ToolCallbackArgs<Args extends ZodRawShape> = Parameters<ToolCallback
 export type ToolExecutionContext<Args extends ZodRawShape = ZodRawShape> = Parameters<ToolCallback<Args>>[1];
 
 export type OperationType = "metadata" | "read" | "create" | "delete" | "update" | "connect";
-export type ToolCategory = "mongodb" | "atlas";
+export type ToolCategory = "mongodb" | "atlas" | "assistant";
 export type TelemetryToolMetadata = {
     projectId?: string;
     orgId?: string;

tests/integration/helpers.ts

@@ -179,22 +179,26 @@ export function setupIntegrationTest(
     };
 }
 
-// eslint-disable-next-line @typescript-eslint/no-redundant-type-constituents
-export function getResponseContent(content: unknown | { content: unknown }): string {
+export function getResponseContent(content: unknown): string {
     return getResponseElements(content)
         .map((item) => item.text)
         .join("\n");
 }
 
-// eslint-disable-next-line @typescript-eslint/no-redundant-type-constituents
-export function getResponseElements(content: unknown | { content: unknown }): { type: string; text: string }[] {
+export interface ResponseElement {
+    type: string;
+    text: string;
+    _meta?: unknown;
+}
+
+export function getResponseElements(content: unknown): ResponseElement[] {
     if (typeof content === "object" && content !== null && "content" in content) {
-        content = (content as { content: unknown }).content;
+        content = content.content;
     }
 
     expect(content).toBeInstanceOf(Array);
 
-    const response = content as { type: string; text: string }[];
+    const response = content as ResponseElement[];
     for (const item of response) {
         expect(item).toHaveProperty("type");
         expect(item).toHaveProperty("text");