Add MongoDB Assistant Tools

src/common/config.ts

@@ -104,6 +104,7 @@ export interface UserConfig extends CliOptions {
     apiBaseUrl: string;
     apiClientId?: string;
     apiClientSecret?: string;
+    assistantBaseUrl: string;
     telemetry: "enabled" | "disabled";
     logPath: string;
     exportsPath: string;
@@ -123,6 +124,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: 300000, // 5 minutes

src/server.ts

@@ -17,6 +17,7 @@ import {
 } from "@modelcontextprotocol/sdk/types.js";
 import assert from "assert";
 import { ToolBase } from "./tools/tool.js";
+import { AssistantTools } from "./tools/assistant/tools.js";
 
 export interface ServerOptions {
     session: Session;
@@ -172,7 +173,7 @@ export class Server {
     }
 
     private registerTools(): void {
-        for (const toolConstructor of [...AtlasTools, ...MongoDbTools]) {
+        for (const toolConstructor of [...AtlasTools, ...MongoDbTools, ...AssistantTools]) {
             const tool = new toolConstructor(this.session, this.userConfig, this.telemetry);
             if (tool.register(this)) {
                 this.tools.push(tool);

src/tools/assistant/assistantTool.ts

@@ -0,0 +1,47 @@
+import { TelemetryToolMetadata, ToolArgs, ToolBase, ToolCategory } from "../tool.js";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { Server } from "../../server.js";
+import { Session } from "../../common/session.js";
+import { UserConfig } from "../../common/config.js";
+import { Telemetry } from "../../telemetry/telemetry.js";
+import { packageInfo } from "../../common/packageInfo.js";
+
+export abstract class AssistantToolBase extends ToolBase {
+    protected server?: Server;
+    public category: ToolCategory = "assistant";
+    protected baseUrl: URL;
+    protected requiredHeaders: Record<string, string>;
+
+    constructor(
+        protected readonly session: Session,
+        protected readonly config: UserConfig,
+        protected readonly telemetry: Telemetry
+    ) {
+        super(session, config, telemetry);
+        this.baseUrl = new URL(config.assistantBaseUrl);
+        const serverVersion = packageInfo.version;
+        this.requiredHeaders = {
+            "x-request-origin": "mongodb-mcp-server",
+            "user-agent": serverVersion ? `mongodb-mcp-server/v${serverVersion}` : "mongodb-mcp-server",
+        };
+    }
+
+    public register(server: Server): boolean {
+        this.server = server;
+        return super.register(server);
+    }
+
+    protected resolveTelemetryMetadata(
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        args: ToolArgs<typeof this.argsShape>
+    ): TelemetryToolMetadata {
+        return {};
+    }
+
+    protected handleError(
+        error: unknown,
+        args: ToolArgs<typeof this.argsShape>
+    ): Promise<CallToolResult> | CallToolResult {
+        return super.handleError(error, args);
+    }
+}

src/tools/assistant/list_knowledge_sources.ts

@@ -0,0 +1,51 @@
+import { z } from "zod";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { OperationType } from "../tool.js";
+import { AssistantToolBase } from "./assistantTool.js";
+
+export const dataSourceMetadataSchema = z.object({
+    id: z.string().describe("The name of the data source"),
+    type: z.string().optional().describe("The type of the data source"),
+    versions: z
+        .array(
+            z.object({
+                label: z.string().describe("The version label of the data source"),
+                isCurrent: z.boolean().describe("Whether this version is current active version"),
+            })
+        )
+        .describe("A list of available versions for this data source"),
+});
+
+export const listDataSourcesResponseSchema = z.object({
+    dataSources: z.array(dataSourceMetadataSchema).describe("A list of data sources"),
+});
+
+export class ListKnowledgeSourcesTool extends AssistantToolBase {
+    public name = "list_knowledge_sources";
+    protected description = "List available data sources in the MongoDB Assistant knowledge base";
+    protected argsShape = {};
+    public operationType: OperationType = "read";
+
+    protected async execute(): Promise<CallToolResult> {
+        const searchEndpoint = new URL("content/sources", this.baseUrl);
+        const response = await fetch(searchEndpoint, {
+            method: "GET",
+            headers: this.requiredHeaders,
+        });
+        if (!response.ok) {
+            throw new Error(`Failed to list knowledge sources: ${response.statusText}`);
+        }
+        const { dataSources } = listDataSourcesResponseSchema.parse(await response.json());
+
+        return {
+            content: dataSources.map(({ id, type, versions }) => ({
+                type: "text",
+                text: id,
+                _meta: {
+                    type,
+                    versions,
+                },
+            })),
+        };
+    }
+}

src/tools/assistant/search_knowledge.ts

@@ -0,0 +1,68 @@
+import { z } from "zod";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { ToolArgs, OperationType } from "../tool.js";
+import { AssistantToolBase } from "./assistantTool.js";
+
+export const SearchKnowledgeToolArgs = {
+    query: z.string().describe("A natural language query to search for in the knowledge base"),
+    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 search in. 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."
+        ),
+};
+
+export const knowledgeChunkSchema = z
+    .object({
+        url: z.string().describe("The URL of the search result"),
+        title: z.string().describe("Title of the search result"),
+        text: z.string().describe("Chunk text"),
+        metadata: z
+            .object({
+                tags: z.array(z.string()).describe("The tags of the source"),
+            })
+            .passthrough(),
+    })
+    .passthrough();
+
+export const searchResponseSchema = z.object({
+    results: z.array(knowledgeChunkSchema).describe("A list of search results"),
+});
+
+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 searchEndpoint = new URL("content/search", this.baseUrl);
+        const response = await fetch(searchEndpoint, {
+            method: "POST",
+            headers: new Headers({ ...this.requiredHeaders, "Content-Type": "application/json" }),
+            body: JSON.stringify(args),
+        });
+        if (!response.ok) {
+            throw new Error(`Failed to search knowledge base: ${response.statusText}`);
+        }
+        const { results } = searchResponseSchema.parse(await response.json());
+        return {
+            content: results.map(({ text, metadata }) => ({
+                type: "text",
+                text,
+                _meta: {
+                    ...metadata,
+                },
+            })),
+        };
+    }
+}

src/tools/assistant/tools.ts

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

src/tools/tool.ts

@@ -11,7 +11,7 @@ import { Server } from "../server.js";
 export type ToolArgs<Args extends ZodRawShape> = z.objectOutputType<Args, ZodNever>;
 
 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

@@ -129,22 +129,26 @@ export function setupIntegrationTest(getUserConfig: () => UserConfig): Integrati
     };
 }
 
-// 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");

tests/integration/tools/assistant/assistantHelpers.ts

@@ -0,0 +1,78 @@
+import { setupIntegrationTest, IntegrationTest, defaultTestConfig } from "../../helpers.js";
+import { describe, SuiteCollector } from "vitest";
+import { vi, beforeAll, afterAll } from "vitest";
+
+export type IntegrationTestFunction = (integration: IntegrationTest) => void;
+
+export function describeWithAssistant(name: string, fn: IntegrationTestFunction): SuiteCollector<object> {
+    const testDefinition = (): void => {
+        const integration = setupIntegrationTest(() => ({
+            ...defaultTestConfig,
+            assistantBaseUrl: "https://knowledge.test.mongodb.com/api/", // Use test URL
+        }));
+
+        describe(name, () => {
+            fn(integration);
+        });
+    };
+
+    // eslint-disable-next-line vitest/valid-describe-callback
+    return describe("assistant", testDefinition);
+}
+
+/**
+ * Mocks fetch for assistant API calls
+ */
+interface MockedAssistantAPI {
+    mockListSources: (sources: unknown[]) => void;
+    mockSearchResults: (results: unknown[]) => void;
+    mockAPIError: (status: number, statusText: string) => void;
+    mockNetworkError: (error: Error) => void;
+    mockFetch: ReturnType<typeof vi.fn>;
+}
+
+export function makeMockAssistantAPI(): MockedAssistantAPI {
+    const mockFetch = vi.fn();
+
+    beforeAll(() => {
+        global.fetch = mockFetch;
+    });
+
+    afterAll(() => {
+        vi.restoreAllMocks();
+    });
+
+    const mockListSources: MockedAssistantAPI["mockListSources"] = (sources) => {
+        mockFetch.mockResolvedValueOnce({
+            ok: true,
+            json: () => Promise.resolve({ dataSources: sources }),
+        });
+    };
+
+    const mockSearchResults: MockedAssistantAPI["mockSearchResults"] = (results) => {
+        mockFetch.mockResolvedValueOnce({
+            ok: true,
+            json: () => Promise.resolve({ results }),
+        });
+    };
+
+    const mockAPIError: MockedAssistantAPI["mockAPIError"] = (status, statusText) => {
+        mockFetch.mockResolvedValueOnce({
+            ok: false,
+            status,
+            statusText,
+        });
+    };
+
+    const mockNetworkError: MockedAssistantAPI["mockNetworkError"] = (error) => {
+        mockFetch.mockRejectedValueOnce(error);
+    };
+
+    return {
+        mockListSources,
+        mockSearchResults,
+        mockAPIError,
+        mockNetworkError,
+        mockFetch,
+    };
+}