From 2338504db3a3fba80bb45955dc448fccddd9a31b Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Mon, 18 Aug 2025 12:20:25 -0500 Subject: [PATCH 01/12] (EAI-1266) Knowledge Search on MCP --- src/common/config.ts | 2 + src/server.ts | 3 +- src/tools/assistant/assistantTool.ts | 45 ++++++++++++ src/tools/assistant/list_knowledge_sources.ts | 51 ++++++++++++++ src/tools/assistant/search_knowledge.ts | 68 +++++++++++++++++++ src/tools/assistant/tools.ts | 4 ++ src/tools/tool.ts | 2 +- 7 files changed, 173 insertions(+), 2 deletions(-) create mode 100644 src/tools/assistant/assistantTool.ts create mode 100644 src/tools/assistant/list_knowledge_sources.ts create mode 100644 src/tools/assistant/search_knowledge.ts create mode 100644 src/tools/assistant/tools.ts diff --git a/src/common/config.ts b/src/common/config.ts index f5c6a079d..4f6381214 100644 --- a/src/common/config.ts +++ b/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 diff --git a/src/server.ts b/src/server.ts index bf41b26d1..810024c9e 100644 --- a/src/server.ts +++ b/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); diff --git a/src/tools/assistant/assistantTool.ts b/src/tools/assistant/assistantTool.ts new file mode 100644 index 000000000..3d61da72b --- /dev/null +++ b/src/tools/assistant/assistantTool.ts @@ -0,0 +1,45 @@ +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"; + +export abstract class AssistantToolBase extends ToolBase { + protected server?: Server; + public category: ToolCategory = "assistant"; + protected baseUrl: URL; + protected requiredHeaders: Record; + + constructor( + protected readonly session: Session, + protected readonly config: UserConfig, + protected readonly telemetry: Telemetry + ) { + super(session, config, telemetry); + this.baseUrl = new URL(config.assistantBaseUrl); + this.requiredHeaders = { + "x-request-origin": "mongodb-mcp-server", + "user-agent": "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 + ): TelemetryToolMetadata { + return {}; + } + + protected handleError( + error: unknown, + args: ToolArgs + ): Promise | CallToolResult { + return super.handleError(error, args); + } +} diff --git a/src/tools/assistant/list_knowledge_sources.ts b/src/tools/assistant/list_knowledge_sources.ts new file mode 100644 index 000000000..5becc53c1 --- /dev/null +++ b/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 { + 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, + }, + })), + }; + } +} diff --git a/src/tools/assistant/search_knowledge.ts b/src/tools/assistant/search_knowledge.ts new file mode 100644 index 000000000..1b181615f --- /dev/null +++ b/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): Promise { + 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, + }, + })), + }; + } +} diff --git a/src/tools/assistant/tools.ts b/src/tools/assistant/tools.ts new file mode 100644 index 000000000..907d05eff --- /dev/null +++ b/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]; diff --git a/src/tools/tool.ts b/src/tools/tool.ts index 2e7d6cf27..0b7c267f2 100644 --- a/src/tools/tool.ts +++ b/src/tools/tool.ts @@ -11,7 +11,7 @@ import { Server } from "../server.js"; export type ToolArgs = z.objectOutputType; 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; From 881c4eeefc2823d9127a9ef91a07ed46156e2761 Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Mon, 18 Aug 2025 20:16:46 -0500 Subject: [PATCH 02/12] Add tests for assistant tools --- tests/integration/helpers.ts | 16 +- .../tools/assistant/assistantHelpers.ts | 78 +++++++ .../assistant/listKnowledgeSources.test.ts | 125 ++++++++++ .../tools/assistant/searchKnowledge.test.ts | 217 ++++++++++++++++++ 4 files changed, 430 insertions(+), 6 deletions(-) create mode 100644 tests/integration/tools/assistant/assistantHelpers.ts create mode 100644 tests/integration/tools/assistant/listKnowledgeSources.test.ts create mode 100644 tests/integration/tools/assistant/searchKnowledge.test.ts diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index 738cbdfdc..12995b59b 100644 --- a/tests/integration/helpers.ts +++ b/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"); diff --git a/tests/integration/tools/assistant/assistantHelpers.ts b/tests/integration/tools/assistant/assistantHelpers.ts new file mode 100644 index 000000000..816d6d2c5 --- /dev/null +++ b/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 { + 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; +} + +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, + }; +} diff --git a/tests/integration/tools/assistant/listKnowledgeSources.test.ts b/tests/integration/tools/assistant/listKnowledgeSources.test.ts new file mode 100644 index 000000000..69c3715c1 --- /dev/null +++ b/tests/integration/tools/assistant/listKnowledgeSources.test.ts @@ -0,0 +1,125 @@ +import { describe, expect, it } from "vitest"; +import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; +import { expectDefined, validateToolMetadata, getResponseElements } from "../../helpers.js"; +import { describeWithAssistant, makeMockAssistantAPI } from "./assistantHelpers.js"; + +describeWithAssistant("list_knowledge_sources", (integration) => { + const { mockListSources, mockAPIError, mockNetworkError } = makeMockAssistantAPI(); + + validateToolMetadata( + integration, + "list_knowledge_sources", + "List available data sources in the MongoDB Assistant knowledge base", + [] + ); + + describe("happy path", () => { + it("returns list of data sources with metadata", async () => { + const mockSources = [ + { + id: "mongodb-manual", + type: "documentation", + versions: [ + { label: "7.0", isCurrent: true }, + { label: "6.0", isCurrent: false }, + ], + }, + { + id: "node-driver", + type: "driver", + versions: [ + { label: "6.0", isCurrent: true }, + { label: "5.0", isCurrent: false }, + ], + }, + ]; + + mockListSources(mockSources); + + const response = (await integration + .mcpClient() + .callTool({ name: "list_knowledge_sources", arguments: {} })) as CallToolResult; + + expect(response.isError).toBeFalsy(); + expect(response.content).toBeInstanceOf(Array); + expect(response.content).toHaveLength(2); + + const elements = getResponseElements(response.content); + + // Check first data source + expect(elements[0]?.text).toBe("mongodb-manual"); + expect(elements[0]?._meta).toEqual({ + type: "documentation", + versions: [ + { label: "7.0", isCurrent: true }, + { label: "6.0", isCurrent: false }, + ], + }); + + // Check second data source + expect(elements[1]?.text).toBe("node-driver"); + expect(elements[1]?._meta).toEqual({ + type: "driver", + versions: [ + { label: "6.0", isCurrent: true }, + { label: "5.0", isCurrent: false }, + ], + }); + }); + + it("handles empty data sources list", async () => { + mockListSources([]); + + const response = (await integration + .mcpClient() + .callTool({ name: "list_knowledge_sources", arguments: {} })) as CallToolResult; + + expect(response.isError).toBeFalsy(); + expect(response.content).toBeInstanceOf(Array); + expect(response.content).toHaveLength(0); + }); + }); + + describe("error handling", () => { + it("handles API error responses", async () => { + mockAPIError(500, "Internal Server Error"); + + const response = (await integration + .mcpClient() + .callTool({ name: "list_knowledge_sources", arguments: {} })) as CallToolResult; + + expect(response.isError).toBe(true); + expectDefined(response.content); + expect(response.content[0]).toHaveProperty("text"); + expect(response.content[0]?.text).toContain("Failed to list knowledge sources: Internal Server Error"); + }); + + it("handles network errors", async () => { + mockNetworkError(new Error("Network connection failed")); + + const response = (await integration + .mcpClient() + .callTool({ name: "list_knowledge_sources", arguments: {} })) as CallToolResult; + + expect(response.isError).toBe(true); + expectDefined(response.content); + expect(response.content[0]).toHaveProperty("text"); + expect(response.content[0]?.text).toContain("Network connection failed"); + }); + + it("handles malformed API response", async () => { + // Mock a response that doesn't match the expected schema + mockListSources(["invalid-response"]); + + const response = (await integration + .mcpClient() + .callTool({ name: "list_knowledge_sources", arguments: {} })) as CallToolResult; + + expect(response.isError).toBe(true); + expectDefined(response.content); + expect(response.content[0]).toHaveProperty("text"); + // Should contain some indication of a parsing/validation error + expect(response.content[0]?.text).toMatch(/error/i); + }); + }); +}); diff --git a/tests/integration/tools/assistant/searchKnowledge.test.ts b/tests/integration/tools/assistant/searchKnowledge.test.ts new file mode 100644 index 000000000..9a2d355e8 --- /dev/null +++ b/tests/integration/tools/assistant/searchKnowledge.test.ts @@ -0,0 +1,217 @@ +import { describe, expect, it } from "vitest"; +import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; +import { + expectDefined, + validateToolMetadata, + validateThrowsForInvalidArguments, + getResponseElements, +} from "../../helpers.js"; +import { describeWithAssistant, makeMockAssistantAPI } from "./assistantHelpers.js"; + +describeWithAssistant("search_knowledge", (integration) => { + const { mockSearchResults, mockAPIError, mockNetworkError } = makeMockAssistantAPI(); + + validateToolMetadata( + integration, + "search_knowledge", + "Search for information in the MongoDB Assistant knowledge base", + [ + { + name: "dataSources", + description: + "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.", + type: "array", + required: false, + }, + { + name: "limit", + description: "The maximum number of results to return", + type: "number", + required: false, + }, + { + name: "query", + description: "A natural language query to search for in the knowledge base", + type: "string", + required: true, + }, + ] + ); + + validateThrowsForInvalidArguments(integration, "search_knowledge", [ + {}, // missing required query + { query: 123 }, // invalid query type + { query: "test", limit: -1 }, // invalid limit + { query: "test", limit: 101 }, // limit too high + { query: "test", dataSources: "invalid" }, // invalid dataSources type + { query: "test", dataSources: [{ name: 123 }] }, // invalid dataSource name type + { query: "test", dataSources: [{}] }, // missing required name field + ]); + + describe("Success Cases", () => { + it("searches with query only", async () => { + const mockResults = [ + { + url: "https://docs.mongodb.com/manual/aggregation/", + title: "Aggregation Pipeline", + text: "The aggregation pipeline is a framework for data aggregation modeled on the concept of data processing pipelines.", + metadata: { + tags: ["aggregation", "pipeline"], + source: "mongodb-manual", + }, + }, + { + url: "https://docs.mongodb.com/manual/reference/operator/aggregation/", + title: "Aggregation Pipeline Operators", + text: "Aggregation pipeline operations have an array of operators available.", + metadata: { + tags: ["aggregation", "operators"], + source: "mongodb-manual", + }, + }, + ]; + + mockSearchResults(mockResults); + + const response = (await integration.mcpClient().callTool({ + name: "search_knowledge", + arguments: { query: "aggregation pipeline" }, + })) as CallToolResult; + + expect(response.isError).toBeFalsy(); + expect(response.content).toBeInstanceOf(Array); + expect(response.content).toHaveLength(2); + + const elements = getResponseElements(response.content); + + // Check first result + expect(elements[0]?.text).toBe( + "The aggregation pipeline is a framework for data aggregation modeled on the concept of data processing pipelines." + ); + expect(elements[0]?._meta).toEqual({ + tags: ["aggregation", "pipeline"], + source: "mongodb-manual", + }); + + // Check second result + expect(elements[1]?.text).toBe("Aggregation pipeline operations have an array of operators available."); + expect(elements[1]?._meta).toEqual({ + tags: ["aggregation", "operators"], + source: "mongodb-manual", + }); + }); + + it("searches with query, limit, and dataSources", async () => { + const mockResults = [ + { + url: "https://mongodb.github.io/node-mongodb-native/", + title: "Node.js Driver", + text: "The official MongoDB driver for Node.js provides a high-level API on top of mongodb-core.", + metadata: { + tags: ["driver", "nodejs"], + source: "node-driver", + }, + }, + ]; + + mockSearchResults(mockResults); + + const response = (await integration.mcpClient().callTool({ + name: "search_knowledge", + arguments: { + query: "node.js driver", + limit: 1, + dataSources: [{ name: "node-driver", versionLabel: "6.0" }], + }, + })) as CallToolResult; + + expect(response.isError).toBeFalsy(); + expect(response.content).toBeInstanceOf(Array); + expect(response.content).toHaveLength(1); + + const elements = getResponseElements(response.content); + expect(elements[0]?.text).toBe( + "The official MongoDB driver for Node.js provides a high-level API on top of mongodb-core." + ); + expect(elements[0]?._meta).toEqual({ + tags: ["driver", "nodejs"], + source: "node-driver", + }); + }); + + it("handles empty search results", async () => { + mockSearchResults([]); + + const response = (await integration + .mcpClient() + .callTool({ name: "search_knowledge", arguments: { query: "nonexistent topic" } })) as CallToolResult; + + expect(response.isError).toBeFalsy(); + expect(response.content).toBeInstanceOf(Array); + expect(response.content).toHaveLength(0); + }); + + it("uses default limit when not specified", async () => { + const mockResults = Array(5) + .fill(null) + .map((_, i) => ({ + url: `https://docs.mongodb.com/result${i}`, + title: `Result ${i}`, + text: `Search result number ${i}`, + metadata: { tags: [`tag${i}`] }, + })); + + mockSearchResults(mockResults); + + const response = (await integration + .mcpClient() + .callTool({ name: "search_knowledge", arguments: { query: "test query" } })) as CallToolResult; + + expect(response.isError).toBeFalsy(); + expect(response.content).toHaveLength(5); + }); + }); + + describe("error handling", () => { + it("handles API error responses", async () => { + mockAPIError(404, "Not Found"); + + const response = (await integration + .mcpClient() + .callTool({ name: "search_knowledge", arguments: { query: "test query" } })) as CallToolResult; + + expect(response.isError).toBe(true); + expectDefined(response.content); + expect(response.content[0]).toHaveProperty("text"); + expect(response.content[0]?.text).toContain("Failed to search knowledge base: Not Found"); + }); + + it("handles network errors", async () => { + mockNetworkError(new Error("Connection timeout")); + + const response = (await integration + .mcpClient() + .callTool({ name: "search_knowledge", arguments: { query: "test query" } })) as CallToolResult; + + expect(response.isError).toBe(true); + expectDefined(response.content); + expect(response.content[0]).toHaveProperty("text"); + expect(response.content[0]?.text).toContain("Connection timeout"); + }); + + it("handles malformed API response", async () => { + // Mock a response that doesn't match the expected schema + mockSearchResults(["invalid-response"]); + + const response = (await integration + .mcpClient() + .callTool({ name: "search_knowledge", arguments: { query: "test query" } })) as CallToolResult; + + expect(response.isError).toBe(true); + expectDefined(response.content); + expect(response.content[0]).toHaveProperty("text"); + // Should contain some indication of a parsing/validation error + expect(response.content[0]?.text).toMatch(/error/i); + }); + }); +}); From 3bdd3dfa80a82e6749bf65b1bb3c4a2e6978f77e Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Wed, 20 Aug 2025 00:55:26 -0500 Subject: [PATCH 03/12] Include server version in assistant user-agent header --- src/helpers/getServerVersion.ts | 24 ++++++++++++++++++++++++ src/tools/assistant/assistantTool.ts | 4 +++- 2 files changed, 27 insertions(+), 1 deletion(-) create mode 100644 src/helpers/getServerVersion.ts diff --git a/src/helpers/getServerVersion.ts b/src/helpers/getServerVersion.ts new file mode 100644 index 000000000..0c6b5dd92 --- /dev/null +++ b/src/helpers/getServerVersion.ts @@ -0,0 +1,24 @@ +import { fileURLToPath } from "url"; +import { readFileSync } from "fs"; +import { dirname, resolve } from "path"; +import { z } from "zod"; + +const packageJsonSchema = z + .object({ + version: z.string(), + }) + .passthrough(); + +export function getServerVersion(): string | null { + try { + const __filename = fileURLToPath(import.meta.url); + const __dirname = dirname(__filename); + const pkgPath = resolve(__dirname, "..", "..", "package.json"); + + const packageJson = packageJsonSchema.parse(JSON.parse(readFileSync(pkgPath, "utf-8"))); + return packageJson.version; + } catch (err) { + console.warn("Could not read package.json:", err); + return null; + } +} diff --git a/src/tools/assistant/assistantTool.ts b/src/tools/assistant/assistantTool.ts index 3d61da72b..0ab8087f6 100644 --- a/src/tools/assistant/assistantTool.ts +++ b/src/tools/assistant/assistantTool.ts @@ -4,6 +4,7 @@ import { Server } from "../../server.js"; import { Session } from "../../common/session.js"; import { UserConfig } from "../../common/config.js"; import { Telemetry } from "../../telemetry/telemetry.js"; +import { getServerVersion } from "../../helpers/getServerVersion.js"; export abstract class AssistantToolBase extends ToolBase { protected server?: Server; @@ -18,9 +19,10 @@ export abstract class AssistantToolBase extends ToolBase { ) { super(session, config, telemetry); this.baseUrl = new URL(config.assistantBaseUrl); + const serverVersion = getServerVersion(); this.requiredHeaders = { "x-request-origin": "mongodb-mcp-server", - "user-agent": "mongodb-mcp-server", + "user-agent": serverVersion ? `mongodb-mcp-server/v${serverVersion}` : "mongodb-mcp-server", }; } From 59b3c636ffb6662340455cde4f3ec6f4600757e3 Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Fri, 22 Aug 2025 10:18:39 -0500 Subject: [PATCH 04/12] get server version from packageInfo --- src/helpers/getServerVersion.ts | 24 ------------------------ src/tools/assistant/assistantTool.ts | 4 ++-- 2 files changed, 2 insertions(+), 26 deletions(-) delete mode 100644 src/helpers/getServerVersion.ts diff --git a/src/helpers/getServerVersion.ts b/src/helpers/getServerVersion.ts deleted file mode 100644 index 0c6b5dd92..000000000 --- a/src/helpers/getServerVersion.ts +++ /dev/null @@ -1,24 +0,0 @@ -import { fileURLToPath } from "url"; -import { readFileSync } from "fs"; -import { dirname, resolve } from "path"; -import { z } from "zod"; - -const packageJsonSchema = z - .object({ - version: z.string(), - }) - .passthrough(); - -export function getServerVersion(): string | null { - try { - const __filename = fileURLToPath(import.meta.url); - const __dirname = dirname(__filename); - const pkgPath = resolve(__dirname, "..", "..", "package.json"); - - const packageJson = packageJsonSchema.parse(JSON.parse(readFileSync(pkgPath, "utf-8"))); - return packageJson.version; - } catch (err) { - console.warn("Could not read package.json:", err); - return null; - } -} diff --git a/src/tools/assistant/assistantTool.ts b/src/tools/assistant/assistantTool.ts index 0ab8087f6..c07e44fd2 100644 --- a/src/tools/assistant/assistantTool.ts +++ b/src/tools/assistant/assistantTool.ts @@ -4,7 +4,7 @@ import { Server } from "../../server.js"; import { Session } from "../../common/session.js"; import { UserConfig } from "../../common/config.js"; import { Telemetry } from "../../telemetry/telemetry.js"; -import { getServerVersion } from "../../helpers/getServerVersion.js"; +import { packageInfo } from "../../common/packageInfo.js"; export abstract class AssistantToolBase extends ToolBase { protected server?: Server; @@ -19,7 +19,7 @@ export abstract class AssistantToolBase extends ToolBase { ) { super(session, config, telemetry); this.baseUrl = new URL(config.assistantBaseUrl); - const serverVersion = getServerVersion(); + const serverVersion = packageInfo.version; this.requiredHeaders = { "x-request-origin": "mongodb-mcp-server", "user-agent": serverVersion ? `mongodb-mcp-server/v${serverVersion}` : "mongodb-mcp-server", From bb5e585799bfaebdefd0a94d51f1e6fffa0584b1 Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Fri, 26 Sep 2025 13:42:55 -0500 Subject: [PATCH 05/12] Review --- README.md | 6 ++++++ src/common/logger.ts | 3 +++ src/tools/assistant/assistantTool.ts | 7 +++---- ...dge_sources.ts => listKnowledgeSources.ts} | 19 ++++++++++++++++-- ...search_knowledge.ts => searchKnowledge.ts} | 19 ++++++++++++++++-- src/tools/assistant/tools.ts | 4 ++-- .../assistant/listKnowledgeSources.test.ts | 14 ++++++------- .../tools/assistant/searchKnowledge.test.ts | 20 +++++++++---------- 8 files changed, 65 insertions(+), 27 deletions(-) rename src/tools/assistant/{list_knowledge_sources.ts => listKnowledgeSources.ts} (74%) rename src/tools/assistant/{search_knowledge.ts => searchKnowledge.ts} (81%) diff --git a/README.md b/README.md index d508b755d..2e2633f70 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,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) @@ -308,6 +309,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 +- `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`. diff --git a/src/common/logger.ts b/src/common/logger.ts index b172ec54c..d3e93c674 100644 --- a/src/common/logger.ts +++ b/src/common/logger.ts @@ -62,6 +62,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; interface LogPayload { diff --git a/src/tools/assistant/assistantTool.ts b/src/tools/assistant/assistantTool.ts index c07e44fd2..8cfca4723 100644 --- a/src/tools/assistant/assistantTool.ts +++ b/src/tools/assistant/assistantTool.ts @@ -31,10 +31,9 @@ export abstract class AssistantToolBase extends ToolBase { return super.register(server); } - protected resolveTelemetryMetadata( - // eslint-disable-next-line @typescript-eslint/no-unused-vars - args: ToolArgs - ): TelemetryToolMetadata { + protected resolveTelemetryMetadata(_args: ToolArgs): 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 {}; } diff --git a/src/tools/assistant/list_knowledge_sources.ts b/src/tools/assistant/listKnowledgeSources.ts similarity index 74% rename from src/tools/assistant/list_knowledge_sources.ts rename to src/tools/assistant/listKnowledgeSources.ts index 5becc53c1..13c02b4b6 100644 --- a/src/tools/assistant/list_knowledge_sources.ts +++ b/src/tools/assistant/listKnowledgeSources.ts @@ -2,6 +2,7 @@ import { z } from "zod"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { OperationType } from "../tool.js"; import { AssistantToolBase } from "./assistantTool.js"; +import { LogId } from "../../common/logger.js"; export const dataSourceMetadataSchema = z.object({ id: z.string().describe("The name of the data source"), @@ -21,7 +22,7 @@ export const listDataSourcesResponseSchema = z.object({ }); export class ListKnowledgeSourcesTool extends AssistantToolBase { - public name = "list_knowledge_sources"; + public name = "list-knowledge-sources"; protected description = "List available data sources in the MongoDB Assistant knowledge base"; protected argsShape = {}; public operationType: OperationType = "read"; @@ -33,7 +34,21 @@ export class ListKnowledgeSourcesTool extends AssistantToolBase { headers: this.requiredHeaders, }); if (!response.ok) { - throw new Error(`Failed to list knowledge sources: ${response.statusText}`); + 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 } = listDataSourcesResponseSchema.parse(await response.json()); diff --git a/src/tools/assistant/search_knowledge.ts b/src/tools/assistant/searchKnowledge.ts similarity index 81% rename from src/tools/assistant/search_knowledge.ts rename to src/tools/assistant/searchKnowledge.ts index 1b181615f..93abcb68a 100644 --- a/src/tools/assistant/search_knowledge.ts +++ b/src/tools/assistant/searchKnowledge.ts @@ -2,6 +2,7 @@ import { z } from "zod"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { ToolArgs, OperationType } from "../tool.js"; import { AssistantToolBase } from "./assistantTool.js"; +import { LogId } from "../../common/logger.js"; export const SearchKnowledgeToolArgs = { query: z.string().describe("A natural language query to search for in the knowledge base"), @@ -37,7 +38,7 @@ export const searchResponseSchema = z.object({ }); export class SearchKnowledgeTool extends AssistantToolBase { - public name = "search_knowledge"; + public name = "search-knowledge"; protected description = "Search for information in the MongoDB Assistant knowledge base"; protected argsShape = { ...SearchKnowledgeToolArgs, @@ -52,7 +53,21 @@ export class SearchKnowledgeTool extends AssistantToolBase { body: JSON.stringify(args), }); if (!response.ok) { - throw new Error(`Failed to search knowledge base: ${response.statusText}`); + 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 } = searchResponseSchema.parse(await response.json()); return { diff --git a/src/tools/assistant/tools.ts b/src/tools/assistant/tools.ts index 907d05eff..12f0b9e89 100644 --- a/src/tools/assistant/tools.ts +++ b/src/tools/assistant/tools.ts @@ -1,4 +1,4 @@ -import { ListKnowledgeSourcesTool } from "./list_knowledge_sources.js"; -import { SearchKnowledgeTool } from "./search_knowledge.js"; +import { ListKnowledgeSourcesTool } from "./listKnowledgeSources.js"; +import { SearchKnowledgeTool } from "./searchKnowledge.js"; export const AssistantTools = [ListKnowledgeSourcesTool, SearchKnowledgeTool]; diff --git a/tests/integration/tools/assistant/listKnowledgeSources.test.ts b/tests/integration/tools/assistant/listKnowledgeSources.test.ts index 69c3715c1..f1761c3a3 100644 --- a/tests/integration/tools/assistant/listKnowledgeSources.test.ts +++ b/tests/integration/tools/assistant/listKnowledgeSources.test.ts @@ -3,12 +3,12 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { expectDefined, validateToolMetadata, getResponseElements } from "../../helpers.js"; import { describeWithAssistant, makeMockAssistantAPI } from "./assistantHelpers.js"; -describeWithAssistant("list_knowledge_sources", (integration) => { +describeWithAssistant("list-knowledge-sources", (integration) => { const { mockListSources, mockAPIError, mockNetworkError } = makeMockAssistantAPI(); validateToolMetadata( integration, - "list_knowledge_sources", + "list-knowledge-sources", "List available data sources in the MongoDB Assistant knowledge base", [] ); @@ -38,7 +38,7 @@ describeWithAssistant("list_knowledge_sources", (integration) => { const response = (await integration .mcpClient() - .callTool({ name: "list_knowledge_sources", arguments: {} })) as CallToolResult; + .callTool({ name: "list-knowledge-sources", arguments: {} })) as CallToolResult; expect(response.isError).toBeFalsy(); expect(response.content).toBeInstanceOf(Array); @@ -72,7 +72,7 @@ describeWithAssistant("list_knowledge_sources", (integration) => { const response = (await integration .mcpClient() - .callTool({ name: "list_knowledge_sources", arguments: {} })) as CallToolResult; + .callTool({ name: "list-knowledge-sources", arguments: {} })) as CallToolResult; expect(response.isError).toBeFalsy(); expect(response.content).toBeInstanceOf(Array); @@ -86,7 +86,7 @@ describeWithAssistant("list_knowledge_sources", (integration) => { const response = (await integration .mcpClient() - .callTool({ name: "list_knowledge_sources", arguments: {} })) as CallToolResult; + .callTool({ name: "list-knowledge-sources", arguments: {} })) as CallToolResult; expect(response.isError).toBe(true); expectDefined(response.content); @@ -99,7 +99,7 @@ describeWithAssistant("list_knowledge_sources", (integration) => { const response = (await integration .mcpClient() - .callTool({ name: "list_knowledge_sources", arguments: {} })) as CallToolResult; + .callTool({ name: "list-knowledge-sources", arguments: {} })) as CallToolResult; expect(response.isError).toBe(true); expectDefined(response.content); @@ -113,7 +113,7 @@ describeWithAssistant("list_knowledge_sources", (integration) => { const response = (await integration .mcpClient() - .callTool({ name: "list_knowledge_sources", arguments: {} })) as CallToolResult; + .callTool({ name: "list-knowledge-sources", arguments: {} })) as CallToolResult; expect(response.isError).toBe(true); expectDefined(response.content); diff --git a/tests/integration/tools/assistant/searchKnowledge.test.ts b/tests/integration/tools/assistant/searchKnowledge.test.ts index 9a2d355e8..43552e682 100644 --- a/tests/integration/tools/assistant/searchKnowledge.test.ts +++ b/tests/integration/tools/assistant/searchKnowledge.test.ts @@ -8,12 +8,12 @@ import { } from "../../helpers.js"; import { describeWithAssistant, makeMockAssistantAPI } from "./assistantHelpers.js"; -describeWithAssistant("search_knowledge", (integration) => { +describeWithAssistant("search-knowledge", (integration) => { const { mockSearchResults, mockAPIError, mockNetworkError } = makeMockAssistantAPI(); validateToolMetadata( integration, - "search_knowledge", + "search-knowledge", "Search for information in the MongoDB Assistant knowledge base", [ { @@ -38,7 +38,7 @@ describeWithAssistant("search_knowledge", (integration) => { ] ); - validateThrowsForInvalidArguments(integration, "search_knowledge", [ + validateThrowsForInvalidArguments(integration, "search-knowledge", [ {}, // missing required query { query: 123 }, // invalid query type { query: "test", limit: -1 }, // invalid limit @@ -74,7 +74,7 @@ describeWithAssistant("search_knowledge", (integration) => { mockSearchResults(mockResults); const response = (await integration.mcpClient().callTool({ - name: "search_knowledge", + name: "search-knowledge", arguments: { query: "aggregation pipeline" }, })) as CallToolResult; @@ -117,7 +117,7 @@ describeWithAssistant("search_knowledge", (integration) => { mockSearchResults(mockResults); const response = (await integration.mcpClient().callTool({ - name: "search_knowledge", + name: "search-knowledge", arguments: { query: "node.js driver", limit: 1, @@ -144,7 +144,7 @@ describeWithAssistant("search_knowledge", (integration) => { const response = (await integration .mcpClient() - .callTool({ name: "search_knowledge", arguments: { query: "nonexistent topic" } })) as CallToolResult; + .callTool({ name: "search-knowledge", arguments: { query: "nonexistent topic" } })) as CallToolResult; expect(response.isError).toBeFalsy(); expect(response.content).toBeInstanceOf(Array); @@ -165,7 +165,7 @@ describeWithAssistant("search_knowledge", (integration) => { const response = (await integration .mcpClient() - .callTool({ name: "search_knowledge", arguments: { query: "test query" } })) as CallToolResult; + .callTool({ name: "search-knowledge", arguments: { query: "test query" } })) as CallToolResult; expect(response.isError).toBeFalsy(); expect(response.content).toHaveLength(5); @@ -178,7 +178,7 @@ describeWithAssistant("search_knowledge", (integration) => { const response = (await integration .mcpClient() - .callTool({ name: "search_knowledge", arguments: { query: "test query" } })) as CallToolResult; + .callTool({ name: "search-knowledge", arguments: { query: "test query" } })) as CallToolResult; expect(response.isError).toBe(true); expectDefined(response.content); @@ -191,7 +191,7 @@ describeWithAssistant("search_knowledge", (integration) => { const response = (await integration .mcpClient() - .callTool({ name: "search_knowledge", arguments: { query: "test query" } })) as CallToolResult; + .callTool({ name: "search-knowledge", arguments: { query: "test query" } })) as CallToolResult; expect(response.isError).toBe(true); expectDefined(response.content); @@ -205,7 +205,7 @@ describeWithAssistant("search_knowledge", (integration) => { const response = (await integration .mcpClient() - .callTool({ name: "search_knowledge", arguments: { query: "test query" } })) as CallToolResult; + .callTool({ name: "search-knowledge", arguments: { query: "test query" } })) as CallToolResult; expect(response.isError).toBe(true); expectDefined(response.content); From 2736bcad9b22e63cb2b30669d032fde714cd7a5a Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Mon, 29 Sep 2025 16:51:29 -0500 Subject: [PATCH 06/12] Fix tests + move assistant request up to abstract base class --- src/tools/assistant/assistantTool.ts | 25 +++++++++++-------- src/tools/assistant/listKnowledgeSources.ts | 5 ++-- src/tools/assistant/searchKnowledge.ts | 7 +++--- .../tools/assistant/assistantHelpers.ts | 15 +++++++---- 4 files changed, 29 insertions(+), 23 deletions(-) diff --git a/src/tools/assistant/assistantTool.ts b/src/tools/assistant/assistantTool.ts index 359a60626..8e9078d2e 100644 --- a/src/tools/assistant/assistantTool.ts +++ b/src/tools/assistant/assistantTool.ts @@ -7,25 +7,22 @@ import { } 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; + protected requiredHeaders: Headers; constructor({ session, config, telemetry, elicitation }: ToolConstructorParams) { super({ session, config, telemetry, elicitation }); this.baseUrl = new URL(config.assistantBaseUrl); const serverVersion = packageInfo.version; - this.requiredHeaders = { + this.requiredHeaders = new Headers({ "x-request-origin": "mongodb-mcp-server", "user-agent": serverVersion ? `mongodb-mcp-server/v${serverVersion}` : "mongodb-mcp-server", - }; + }); } public register(server: Server): boolean { @@ -39,10 +36,16 @@ export abstract class AssistantToolBase extends ToolBase { return {}; } - protected handleError( - error: unknown, - args: ToolArgs - ): Promise | CallToolResult { - return super.handleError(error, args); + 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"); + } + return await fetch(endpoint, { + method: args.method, + headers, + body: JSON.stringify(args.body), + }); } } diff --git a/src/tools/assistant/listKnowledgeSources.ts b/src/tools/assistant/listKnowledgeSources.ts index 13c02b4b6..624c22542 100644 --- a/src/tools/assistant/listKnowledgeSources.ts +++ b/src/tools/assistant/listKnowledgeSources.ts @@ -28,10 +28,9 @@ export class ListKnowledgeSourcesTool extends AssistantToolBase { public operationType: OperationType = "read"; protected async execute(): Promise { - const searchEndpoint = new URL("content/sources", this.baseUrl); - const response = await fetch(searchEndpoint, { + const response = await this.callAssistantApi({ method: "GET", - headers: this.requiredHeaders, + endpoint: "content/sources", }); if (!response.ok) { const message = `Failed to list knowledge sources: ${response.statusText}`; diff --git a/src/tools/assistant/searchKnowledge.ts b/src/tools/assistant/searchKnowledge.ts index 93abcb68a..532cb3071 100644 --- a/src/tools/assistant/searchKnowledge.ts +++ b/src/tools/assistant/searchKnowledge.ts @@ -46,11 +46,10 @@ export class SearchKnowledgeTool extends AssistantToolBase { public operationType: OperationType = "read"; protected async execute(args: ToolArgs): Promise { - const searchEndpoint = new URL("content/search", this.baseUrl); - const response = await fetch(searchEndpoint, { + const response = await this.callAssistantApi({ method: "POST", - headers: new Headers({ ...this.requiredHeaders, "Content-Type": "application/json" }), - body: JSON.stringify(args), + endpoint: "content/search", + body: args, }); if (!response.ok) { const message = `Failed to search knowledge base: ${response.statusText}`; diff --git a/tests/integration/tools/assistant/assistantHelpers.ts b/tests/integration/tools/assistant/assistantHelpers.ts index 816d6d2c5..2f629de7d 100644 --- a/tests/integration/tools/assistant/assistantHelpers.ts +++ b/tests/integration/tools/assistant/assistantHelpers.ts @@ -1,4 +1,4 @@ -import { setupIntegrationTest, IntegrationTest, defaultTestConfig } from "../../helpers.js"; +import { setupIntegrationTest, IntegrationTest, defaultTestConfig, defaultDriverOptions } from "../../helpers.js"; import { describe, SuiteCollector } from "vitest"; import { vi, beforeAll, afterAll } from "vitest"; @@ -6,10 +6,15 @@ export type IntegrationTestFunction = (integration: IntegrationTest) => void; export function describeWithAssistant(name: string, fn: IntegrationTestFunction): SuiteCollector { const testDefinition = (): void => { - const integration = setupIntegrationTest(() => ({ - ...defaultTestConfig, - assistantBaseUrl: "https://knowledge.test.mongodb.com/api/", // Use test URL - })); + const integration = setupIntegrationTest( + () => ({ + ...defaultTestConfig, + assistantBaseUrl: "https://knowledge.test.mongodb.com/api/", // Use test URL + }), + () => ({ + ...defaultDriverOptions, + }) + ); describe(name, () => { fn(integration); From f8365acae7d161d698563080dba20521323bce0f Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Tue, 30 Sep 2025 16:45:18 -0500 Subject: [PATCH 07/12] Don't use zod to parse Assistant API responses --- src/tools/assistant/listKnowledgeSources.ts | 33 +++++++++-------- src/tools/assistant/searchKnowledge.ts | 36 ++++++++++--------- .../assistant/listKnowledgeSources.test.ts | 15 -------- .../tools/assistant/searchKnowledge.test.ts | 15 -------- 4 files changed, 35 insertions(+), 64 deletions(-) diff --git a/src/tools/assistant/listKnowledgeSources.ts b/src/tools/assistant/listKnowledgeSources.ts index 624c22542..47ae79162 100644 --- a/src/tools/assistant/listKnowledgeSources.ts +++ b/src/tools/assistant/listKnowledgeSources.ts @@ -4,22 +4,21 @@ import { OperationType } from "../tool.js"; import { AssistantToolBase } from "./assistantTool.js"; import { LogId } from "../../common/logger.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 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 class ListKnowledgeSourcesTool extends AssistantToolBase { public name = "list-knowledge-sources"; @@ -49,7 +48,7 @@ export class ListKnowledgeSourcesTool extends AssistantToolBase { isError: true, }; } - const { dataSources } = listDataSourcesResponseSchema.parse(await response.json()); + const { dataSources } = (await response.json()) as ListKnowledgeSourcesResponse; return { content: dataSources.map(({ id, type, versions }) => ({ diff --git a/src/tools/assistant/searchKnowledge.ts b/src/tools/assistant/searchKnowledge.ts index 532cb3071..2b17c1654 100644 --- a/src/tools/assistant/searchKnowledge.ts +++ b/src/tools/assistant/searchKnowledge.ts @@ -20,22 +20,24 @@ export const SearchKnowledgeToolArgs = { ), }; -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 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"; @@ -68,7 +70,7 @@ export class SearchKnowledgeTool extends AssistantToolBase { isError: true, }; } - const { results } = searchResponseSchema.parse(await response.json()); + const { results } = (await response.json()) as SearchKnowledgeResponse; return { content: results.map(({ text, metadata }) => ({ type: "text", diff --git a/tests/integration/tools/assistant/listKnowledgeSources.test.ts b/tests/integration/tools/assistant/listKnowledgeSources.test.ts index f1761c3a3..015050645 100644 --- a/tests/integration/tools/assistant/listKnowledgeSources.test.ts +++ b/tests/integration/tools/assistant/listKnowledgeSources.test.ts @@ -106,20 +106,5 @@ describeWithAssistant("list-knowledge-sources", (integration) => { expect(response.content[0]).toHaveProperty("text"); expect(response.content[0]?.text).toContain("Network connection failed"); }); - - it("handles malformed API response", async () => { - // Mock a response that doesn't match the expected schema - mockListSources(["invalid-response"]); - - const response = (await integration - .mcpClient() - .callTool({ name: "list-knowledge-sources", arguments: {} })) as CallToolResult; - - expect(response.isError).toBe(true); - expectDefined(response.content); - expect(response.content[0]).toHaveProperty("text"); - // Should contain some indication of a parsing/validation error - expect(response.content[0]?.text).toMatch(/error/i); - }); }); }); diff --git a/tests/integration/tools/assistant/searchKnowledge.test.ts b/tests/integration/tools/assistant/searchKnowledge.test.ts index 43552e682..e65781920 100644 --- a/tests/integration/tools/assistant/searchKnowledge.test.ts +++ b/tests/integration/tools/assistant/searchKnowledge.test.ts @@ -198,20 +198,5 @@ describeWithAssistant("search-knowledge", (integration) => { expect(response.content[0]).toHaveProperty("text"); expect(response.content[0]?.text).toContain("Connection timeout"); }); - - it("handles malformed API response", async () => { - // Mock a response that doesn't match the expected schema - mockSearchResults(["invalid-response"]); - - const response = (await integration - .mcpClient() - .callTool({ name: "search-knowledge", arguments: { query: "test query" } })) as CallToolResult; - - expect(response.isError).toBe(true); - expectDefined(response.content); - expect(response.content[0]).toHaveProperty("text"); - // Should contain some indication of a parsing/validation error - expect(response.content[0]?.text).toMatch(/error/i); - }); }); }); From 46713471cc6996e92306368335634967029c7789 Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Thu, 2 Oct 2025 01:02:34 -0500 Subject: [PATCH 08/12] Give example KB sources in README --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 24da7dc25..9de6572e5 100644 --- a/README.md +++ b/README.md @@ -323,8 +323,8 @@ NOTE: atlas tools are only available when you set credentials on [configuration] #### MongoDB Assistant Tools -- `list-knowledge-sources` - List available data sources in the MongoDB Assistant knowledge base -- `search-knowledge` - Search for information in the MongoDB Assistant knowledge base +- `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 From 4d887c94aead7b353b00e8e8d807405e46e7991a Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Fri, 3 Oct 2025 13:29:19 -0500 Subject: [PATCH 09/12] Use custom fetch + treat results as untrusted --- package-lock.json | 4 +- package.json | 3 +- src/tools/assistant/assistantTool.ts | 15 ++++-- src/tools/assistant/listKnowledgeSources.ts | 54 ++++++++++++--------- src/tools/assistant/searchKnowledge.ts | 17 ++++--- 5 files changed, 54 insertions(+), 39 deletions(-) diff --git a/package-lock.json b/package-lock.json index 66405c533..1d9d9f492 100644 --- a/package-lock.json +++ b/package-lock.json @@ -26,6 +26,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" }, @@ -15363,10 +15364,7 @@ "version": "2.8.1", "resolved": "https://registry.npmjs.org/yaml/-/yaml-2.8.1.tgz", "integrity": "sha512-lcYcMxX2PO9XMGvAJkJ3OsNMw+/7FKes7/hgerGUYWIoWu5j/+YQqcZr5JnPZWzOsEBgMbSbiSTn/dv/69Mkpw==", - "dev": true, "license": "ISC", - "optional": true, - "peer": true, "bin": { "yaml": "bin.mjs" }, diff --git a/package.json b/package.json index e689ea9cb..f2a906db0 100644 --- a/package.json +++ b/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" }, diff --git a/src/tools/assistant/assistantTool.ts b/src/tools/assistant/assistantTool.ts index 8e9078d2e..a368a77ad 100644 --- a/src/tools/assistant/assistantTool.ts +++ b/src/tools/assistant/assistantTool.ts @@ -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), diff --git a/src/tools/assistant/listKnowledgeSources.ts b/src/tools/assistant/listKnowledgeSources.ts index 47ae79162..98c536597 100644 --- a/src/tools/assistant/listKnowledgeSources.ts +++ b/src/tools/assistant/listKnowledgeSources.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 + ), }; } } diff --git a/src/tools/assistant/searchKnowledge.ts b/src/tools/assistant/searchKnowledge.ts index 2b17c1654..21627f9c8 100644 --- a/src/tools/assistant/searchKnowledge.ts +++ b/src/tools/assistant/searchKnowledge.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 + ), }; } } From caa77921059abffc01603105d65d1b87943b38df Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Fri, 3 Oct 2025 13:40:58 -0500 Subject: [PATCH 10/12] Update tests for custom fetch + untrusted content formatting --- .../tools/assistant/assistantHelpers.ts | 17 ++-- .../assistant/listKnowledgeSources.test.ts | 38 +++++++-- .../tools/assistant/searchKnowledge.test.ts | 78 ++++++++++++++----- 3 files changed, 100 insertions(+), 33 deletions(-) diff --git a/tests/integration/tools/assistant/assistantHelpers.ts b/tests/integration/tools/assistant/assistantHelpers.ts index 2f629de7d..15efb804a 100644 --- a/tests/integration/tools/assistant/assistantHelpers.ts +++ b/tests/integration/tools/assistant/assistantHelpers.ts @@ -1,6 +1,6 @@ import { setupIntegrationTest, IntegrationTest, defaultTestConfig, defaultDriverOptions } from "../../helpers.js"; import { describe, SuiteCollector } from "vitest"; -import { vi, beforeAll, afterAll } from "vitest"; +import { vi, beforeAll, afterAll, beforeEach } from "vitest"; export type IntegrationTestFunction = (integration: IntegrationTest) => void; @@ -9,7 +9,7 @@ export function describeWithAssistant(name: string, fn: IntegrationTestFunction) const integration = setupIntegrationTest( () => ({ ...defaultTestConfig, - assistantBaseUrl: "https://knowledge.test.mongodb.com/api/", // Use test URL + assistantBaseUrl: "https://knowledge.test.mongodb.com/api/", // Not a real URL }), () => ({ ...defaultDriverOptions, @@ -39,8 +39,13 @@ interface MockedAssistantAPI { export function makeMockAssistantAPI(): MockedAssistantAPI { const mockFetch = vi.fn(); - beforeAll(() => { - global.fetch = mockFetch; + beforeAll(async () => { + const { createFetch } = await import("@mongodb-js/devtools-proxy-support"); + vi.mocked(createFetch).mockReturnValue(mockFetch as never); + }); + + beforeEach(() => { + mockFetch.mockClear(); }); afterAll(() => { @@ -50,14 +55,14 @@ export function makeMockAssistantAPI(): MockedAssistantAPI { const mockListSources: MockedAssistantAPI["mockListSources"] = (sources) => { mockFetch.mockResolvedValueOnce({ ok: true, - json: () => Promise.resolve({ dataSources: sources }), + json: async () => ({ dataSources: sources }), }); }; const mockSearchResults: MockedAssistantAPI["mockSearchResults"] = (results) => { mockFetch.mockResolvedValueOnce({ ok: true, - json: () => Promise.resolve({ results }), + json: async () => ({ results }), }); }; diff --git a/tests/integration/tools/assistant/listKnowledgeSources.test.ts b/tests/integration/tools/assistant/listKnowledgeSources.test.ts index 015050645..630a0f2ed 100644 --- a/tests/integration/tools/assistant/listKnowledgeSources.test.ts +++ b/tests/integration/tools/assistant/listKnowledgeSources.test.ts @@ -1,7 +1,18 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; -import { expectDefined, validateToolMetadata, getResponseElements } from "../../helpers.js"; +import { + expectDefined, + validateToolMetadata, + getResponseElements, + getDataFromUntrustedContent, +} from "../../helpers.js"; import { describeWithAssistant, makeMockAssistantAPI } from "./assistantHelpers.js"; +import { parse as yamlParse } from "yaml"; + +// Mock the devtools-proxy-support module +vi.mock("@mongodb-js/devtools-proxy-support", () => ({ + createFetch: vi.fn(), +})); describeWithAssistant("list-knowledge-sources", (integration) => { const { mockListSources, mockAPIError, mockNetworkError } = makeMockAssistantAPI(); @@ -46,10 +57,19 @@ describeWithAssistant("list-knowledge-sources", (integration) => { const elements = getResponseElements(response.content); + // First element is the description + expect(elements[0]?.text).toBe("Found 2 data sources in the MongoDB Assistant knowledge base."); + + // Second element contains the YAML data + expect(elements[1]?.text).toContain(" { }); // Check second data source - expect(elements[1]?.text).toBe("node-driver"); - expect(elements[1]?._meta).toEqual({ + expect(dataSources[1]).toMatchObject({ + id: "node-driver", type: "driver", + currentVersion: "6.0", versions: [ { label: "6.0", isCurrent: true }, { label: "5.0", isCurrent: false }, @@ -76,7 +97,10 @@ describeWithAssistant("list-knowledge-sources", (integration) => { expect(response.isError).toBeFalsy(); expect(response.content).toBeInstanceOf(Array); - expect(response.content).toHaveLength(0); + expect(response.content).toHaveLength(2); + + const elements = getResponseElements(response.content); + expect(elements[0]?.text).toBe("Found 0 data sources in the MongoDB Assistant knowledge base."); }); }); diff --git a/tests/integration/tools/assistant/searchKnowledge.test.ts b/tests/integration/tools/assistant/searchKnowledge.test.ts index e65781920..75931f3e0 100644 --- a/tests/integration/tools/assistant/searchKnowledge.test.ts +++ b/tests/integration/tools/assistant/searchKnowledge.test.ts @@ -1,12 +1,19 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { expectDefined, validateToolMetadata, validateThrowsForInvalidArguments, getResponseElements, + getDataFromUntrustedContent, } from "../../helpers.js"; import { describeWithAssistant, makeMockAssistantAPI } from "./assistantHelpers.js"; +import { parse as yamlParse } from "yaml"; + +// Mock the devtools-proxy-support module +vi.mock("@mongodb-js/devtools-proxy-support", () => ({ + createFetch: vi.fn(), +})); describeWithAssistant("search-knowledge", (integration) => { const { mockSearchResults, mockAPIError, mockNetworkError } = makeMockAssistantAPI(); @@ -84,20 +91,34 @@ describeWithAssistant("search-knowledge", (integration) => { const elements = getResponseElements(response.content); + // First element is the description + expect(elements[0]?.text).toBe("Found 2 results in the MongoDB Assistant knowledge base."); + + // Second element contains the YAML data + expect(elements[1]?.text).toContain(" { expect(response.isError).toBeFalsy(); expect(response.content).toBeInstanceOf(Array); - expect(response.content).toHaveLength(1); + expect(response.content).toHaveLength(2); const elements = getResponseElements(response.content); - expect(elements[0]?.text).toBe( - "The official MongoDB driver for Node.js provides a high-level API on top of mongodb-core." - ); - expect(elements[0]?._meta).toEqual({ - tags: ["driver", "nodejs"], - source: "node-driver", + expect(elements[0]?.text).toBe("Found 1 results in the MongoDB Assistant knowledge base."); + + const yamlData = getDataFromUntrustedContent(elements[1]?.text ?? ""); + const results = yamlParse(yamlData); + + expect(results[0]).toMatchObject({ + url: "https://mongodb.github.io/node-mongodb-native/", + title: "Node.js Driver", + text: "The official MongoDB driver for Node.js provides a high-level API on top of mongodb-core.", + metadata: { + tags: ["driver", "nodejs"], + source: "node-driver", + }, }); }); @@ -148,7 +176,10 @@ describeWithAssistant("search-knowledge", (integration) => { expect(response.isError).toBeFalsy(); expect(response.content).toBeInstanceOf(Array); - expect(response.content).toHaveLength(0); + expect(response.content).toHaveLength(2); + + const elements = getResponseElements(response.content); + expect(elements[0]?.text).toBe("Found 0 results in the MongoDB Assistant knowledge base."); }); it("uses default limit when not specified", async () => { @@ -168,7 +199,14 @@ describeWithAssistant("search-knowledge", (integration) => { .callTool({ name: "search-knowledge", arguments: { query: "test query" } })) as CallToolResult; expect(response.isError).toBeFalsy(); - expect(response.content).toHaveLength(5); + expect(response.content).toHaveLength(2); + + const elements = getResponseElements(response.content); + expect(elements[0]?.text).toBe("Found 5 results in the MongoDB Assistant knowledge base."); + + const yamlData = getDataFromUntrustedContent(elements[1]?.text ?? ""); + const results = yamlParse(yamlData); + expect(results).toHaveLength(5); }); }); From 141520599ffc0c30bbbf3c185628a4b03019b1d3 Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Fri, 3 Oct 2025 14:07:36 -0500 Subject: [PATCH 11/12] Integration -> Unit --- .../tools => unit}/assistant/assistantHelpers.ts | 15 ++++++++++----- .../assistant/listKnowledgeSources.test.ts | 2 +- .../assistant/searchKnowledge.test.ts | 2 +- 3 files changed, 12 insertions(+), 7 deletions(-) rename tests/{integration/tools => unit}/assistant/assistantHelpers.ts (81%) rename tests/{integration/tools => unit}/assistant/listKnowledgeSources.test.ts (99%) rename tests/{integration/tools => unit}/assistant/searchKnowledge.test.ts (99%) diff --git a/tests/integration/tools/assistant/assistantHelpers.ts b/tests/unit/assistant/assistantHelpers.ts similarity index 81% rename from tests/integration/tools/assistant/assistantHelpers.ts rename to tests/unit/assistant/assistantHelpers.ts index 15efb804a..68b361243 100644 --- a/tests/integration/tools/assistant/assistantHelpers.ts +++ b/tests/unit/assistant/assistantHelpers.ts @@ -1,15 +1,20 @@ -import { setupIntegrationTest, IntegrationTest, defaultTestConfig, defaultDriverOptions } from "../../helpers.js"; +import { + setupIntegrationTest, + IntegrationTest, + defaultTestConfig, + defaultDriverOptions, +} from "../../integration/helpers.js"; import { describe, SuiteCollector } from "vitest"; import { vi, beforeAll, afterAll, beforeEach } from "vitest"; -export type IntegrationTestFunction = (integration: IntegrationTest) => void; +export type MockIntegrationTestFunction = (integration: IntegrationTest) => void; -export function describeWithAssistant(name: string, fn: IntegrationTestFunction): SuiteCollector { +export function describeWithAssistant(name: string, fn: MockIntegrationTestFunction): SuiteCollector { const testDefinition = (): void => { const integration = setupIntegrationTest( () => ({ ...defaultTestConfig, - assistantBaseUrl: "https://knowledge.test.mongodb.com/api/", // Not a real URL + assistantBaseUrl: "https://knowledge-mock.mongodb.com/api/v1", // Not a real URL }), () => ({ ...defaultDriverOptions, @@ -22,7 +27,7 @@ export function describeWithAssistant(name: string, fn: IntegrationTestFunction) }; // eslint-disable-next-line vitest/valid-describe-callback - return describe("assistant", testDefinition); + return describe("assistant (mocked)", testDefinition); } /** diff --git a/tests/integration/tools/assistant/listKnowledgeSources.test.ts b/tests/unit/assistant/listKnowledgeSources.test.ts similarity index 99% rename from tests/integration/tools/assistant/listKnowledgeSources.test.ts rename to tests/unit/assistant/listKnowledgeSources.test.ts index 630a0f2ed..0b656e6de 100644 --- a/tests/integration/tools/assistant/listKnowledgeSources.test.ts +++ b/tests/unit/assistant/listKnowledgeSources.test.ts @@ -5,7 +5,7 @@ import { validateToolMetadata, getResponseElements, getDataFromUntrustedContent, -} from "../../helpers.js"; +} from "../../integration/helpers.js"; import { describeWithAssistant, makeMockAssistantAPI } from "./assistantHelpers.js"; import { parse as yamlParse } from "yaml"; diff --git a/tests/integration/tools/assistant/searchKnowledge.test.ts b/tests/unit/assistant/searchKnowledge.test.ts similarity index 99% rename from tests/integration/tools/assistant/searchKnowledge.test.ts rename to tests/unit/assistant/searchKnowledge.test.ts index 75931f3e0..40ed3798d 100644 --- a/tests/integration/tools/assistant/searchKnowledge.test.ts +++ b/tests/unit/assistant/searchKnowledge.test.ts @@ -6,7 +6,7 @@ import { validateThrowsForInvalidArguments, getResponseElements, getDataFromUntrustedContent, -} from "../../helpers.js"; +} from "../../integration/helpers.js"; import { describeWithAssistant, makeMockAssistantAPI } from "./assistantHelpers.js"; import { parse as yamlParse } from "yaml"; From f4bec2c2704a2a6e53691e45049047d69a7ef34e Mon Sep 17 00:00:00 2001 From: Nick Larew Date: Fri, 3 Oct 2025 14:11:54 -0500 Subject: [PATCH 12/12] Tweak tool descriptions --- src/tools/assistant/listKnowledgeSources.ts | 4 +++- src/tools/assistant/searchKnowledge.ts | 9 +++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/src/tools/assistant/listKnowledgeSources.ts b/src/tools/assistant/listKnowledgeSources.ts index 98c536597..9343c1c57 100644 --- a/src/tools/assistant/listKnowledgeSources.ts +++ b/src/tools/assistant/listKnowledgeSources.ts @@ -22,8 +22,10 @@ export type ListKnowledgeSourcesResponse = { dataSources: KnowledgeSource[]; }; +export const ListKnowledgeSourcesToolName = "list-knowledge-sources"; + export class ListKnowledgeSourcesTool extends AssistantToolBase { - public name = "list-knowledge-sources"; + public name = ListKnowledgeSourcesToolName; protected description = "List available data sources in the MongoDB Assistant knowledge base"; protected argsShape = {}; public operationType: OperationType = "read"; diff --git a/src/tools/assistant/searchKnowledge.ts b/src/tools/assistant/searchKnowledge.ts index 21627f9c8..3dd813c5e 100644 --- a/src/tools/assistant/searchKnowledge.ts +++ b/src/tools/assistant/searchKnowledge.ts @@ -4,9 +4,14 @@ import { type ToolArgs, type OperationType, formatUntrustedData } from "../tool. 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 knowledge base"), + 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( @@ -17,7 +22,7 @@ export const SearchKnowledgeToolArgs = { ) .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." + `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.` ), };