diff --git a/src/common/playground/playgroundClient.ts b/src/common/playground/playgroundClient.ts new file mode 100644 index 00000000..2803f118 --- /dev/null +++ b/src/common/playground/playgroundClient.ts @@ -0,0 +1,88 @@ +const PLAYGROUND_SEARCH_URL = "https://search-playground.mongodb.com/api/tools/code-playground/search"; + +/** + * Payload for the Playground endpoint. + */ +export interface PlaygroundRunRequest { + documents: string; + aggregationPipeline: string; + indexDefinition: string; + synonyms: string; +} + +/** + * Successful response from Playground server. + */ +export interface PlaygroundRunResponse { + documents: Array>; +} + +/** + * Error response from Playground server. + */ +interface PlaygroundRunErrorResponse { + code: string; + message: string; +} + +/** + * MCP specific Playground error public for tools. + */ +export class PlaygroundRunError extends Error implements PlaygroundRunErrorResponse { + constructor( + public message: string, + public code: string + ) { + super(message); + } +} + +export enum RunErrorCode { + NETWORK_ERROR = "NETWORK_ERROR", + UNKNOWN = "UNKNOWN", +} + +/** + * Handles Search Playground requests, abstracting low-level details from MCP tools. + * https://search-playground.mongodb.com + */ +export class PlaygroundClient { + async run(request: PlaygroundRunRequest): Promise { + const options: RequestInit = { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify(request), + }; + + let response: Response; + try { + response = await fetch(PLAYGROUND_SEARCH_URL, options); + } catch { + throw new PlaygroundRunError("Cannot run pipeline.", RunErrorCode.NETWORK_ERROR); + } + + if (!response.ok) { + const runErrorResponse = await this.getRunErrorResponse(response); + throw new PlaygroundRunError(runErrorResponse.message, runErrorResponse.code); + } + + try { + return (await response.json()) as PlaygroundRunResponse; + } catch { + throw new PlaygroundRunError("Response is not valid JSON.", RunErrorCode.UNKNOWN); + } + } + + private async getRunErrorResponse(response: Response): Promise { + try { + return (await response.json()) as PlaygroundRunErrorResponse; + } catch { + return { + message: `HTTP ${response.status} ${response.statusText}.`, + code: RunErrorCode.UNKNOWN, + }; + } + } +} diff --git a/src/tools/playground/runPipeline.ts b/src/tools/playground/runPipeline.ts index d00d9be8..85f9c43d 100644 --- a/src/tools/playground/runPipeline.ts +++ b/src/tools/playground/runPipeline.ts @@ -2,17 +2,11 @@ import { OperationType, TelemetryToolMetadata, ToolArgs, ToolBase, ToolCategory import { z } from "zod"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { EJSON } from "bson"; - -const PLAYGROUND_SEARCH_URL = "https://search-playground.mongodb.com/api/tools/code-playground/search"; - -const DEFAULT_DOCUMENTS = [ - { - name: "First document", - }, - { - name: "Second document", - }, -]; +import { + PlaygroundRunError, + PlaygroundRunRequest, + PlaygroundRunResponse, +} from "../../common/playground/playgroundClient.js"; const DEFAULT_SEARCH_INDEX_DEFINITION = { mappings: { @@ -20,32 +14,16 @@ const DEFAULT_SEARCH_INDEX_DEFINITION = { }, }; -const DEFAULT_PIPELINE = [ - { - $search: { - index: "default", - text: { - query: "first", - path: { - wildcard: "*", - }, - }, - }, - }, -]; - const DEFAULT_SYNONYMS: Array> = []; export const RunPipelineOperationArgs = { documents: z .array(z.record(z.string(), z.unknown())) .max(500) - .describe("Documents to run the pipeline against. 500 is maximum.") - .default(DEFAULT_DOCUMENTS), + .describe("Documents to run the pipeline against. 500 is maximum."), aggregationPipeline: z .array(z.record(z.string(), z.unknown())) - .describe("MongoDB aggregation pipeline to run on the provided documents.") - .default(DEFAULT_PIPELINE), + .describe("MongoDB aggregation pipeline to run on the provided documents."), searchIndexDefinition: z .record(z.string(), z.unknown()) .describe("MongoDB search index definition to create before running the pipeline.") @@ -58,22 +36,6 @@ export const RunPipelineOperationArgs = { .default(DEFAULT_SYNONYMS), }; -interface RunRequest { - documents: string; - aggregationPipeline: string; - indexDefinition: string; - synonyms: string; -} - -interface RunResponse { - documents: Array>; -} - -interface RunErrorResponse { - code: string; - message: string; -} - export class RunPipeline extends ToolBase { protected name = "run-pipeline"; protected description = @@ -93,47 +55,24 @@ export class RunPipeline extends ToolBase { return {}; } - private async runPipeline(runRequest: RunRequest): Promise { - const options: RequestInit = { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify(runRequest), - }; - - let response: Response; + private async runPipeline(runRequest: PlaygroundRunRequest): Promise { + // import PlaygroundClient dynamically so we can mock it properly in the tests + const { PlaygroundClient } = await import("../../common/playground/playgroundClient.js"); + const client = new PlaygroundClient(); try { - response = await fetch(PLAYGROUND_SEARCH_URL, options); - } catch { - throw new Error("Cannot run pipeline: network error."); - } + return await client.run(runRequest); + } catch (error: unknown) { + let message: string | undefined; - if (!response.ok) { - const errorMessage = await this.getPlaygroundResponseError(response); - throw new Error(`Pipeline run failed: ${errorMessage}`); - } + if (error instanceof PlaygroundRunError) { + message = `Error code: ${error.code}. Error message: ${error.message}.`; + } - try { - return (await response.json()) as RunResponse; - } catch { - throw new Error("Pipeline run failed: response is not valid JSON."); + throw new Error(message || "Cannot run pipeline."); } } - private async getPlaygroundResponseError(response: Response): Promise { - let errorMessage = `HTTP ${response.status} ${response.statusText}.`; - try { - const errorResponse = (await response.json()) as RunErrorResponse; - errorMessage += ` Error code: ${errorResponse.code}. Error message: ${errorResponse.message}`; - } catch { - // Ignore JSON parse errors - } - - return errorMessage; - } - - private convertToRunRequest(toolArgs: ToolArgs): RunRequest { + private convertToRunRequest(toolArgs: ToolArgs): PlaygroundRunRequest { try { return { documents: JSON.stringify(toolArgs.documents), @@ -146,7 +85,7 @@ export class RunPipeline extends ToolBase { } } - private convertToToolResult(runResponse: RunResponse): CallToolResult { + private convertToToolResult(runResponse: PlaygroundRunResponse): CallToolResult { const content: Array<{ text: string; type: "text" }> = [ { text: `Found ${runResponse.documents.length} documents":`, diff --git a/tests/integration/tools/playground/runPipeline.test.ts b/tests/integration/tools/playground/runPipeline.test.ts index 3036d42d..d76aba0e 100644 --- a/tests/integration/tools/playground/runPipeline.test.ts +++ b/tests/integration/tools/playground/runPipeline.test.ts @@ -1,9 +1,28 @@ +import { jest } from "@jest/globals"; import { describeWithMongoDB } from "../mongodb/mongodbHelpers.js"; import { getResponseElements } from "../../helpers.js"; +import { PlaygroundRunError } from "../../../../src/common/playground/playgroundClient.js"; + +const setupMockPlaygroundClient = (implementation: unknown) => { + // mock ESM modules https://jestjs.io/docs/ecmascript-modules#module-mocking-in-esm + jest.unstable_mockModule("../../../../src/common/playground/playgroundClient.js", () => ({ + PlaygroundClient: implementation, + })); +}; describeWithMongoDB("runPipeline tool", (integration) => { + beforeEach(() => { + jest.resetModules(); + }); + it("should return results", async () => { - await integration.connectMcpClient(); + class PlaygroundClientMock { + run = () => ({ + documents: [{ name: "First document" }], + }); + } + setupMockPlaygroundClient(PlaygroundClientMock); + const response = await integration.mcpClient().callTool({ name: "run-pipeline", arguments: { @@ -20,12 +39,6 @@ describeWithMongoDB("runPipeline tool", (integration) => { }, }, }, - { - $project: { - _id: 0, - name: 1, - }, - }, ], }, }); @@ -41,4 +54,27 @@ describeWithMongoDB("runPipeline tool", (integration) => { }, ]); }); + + it("should return error", async () => { + class PlaygroundClientMock { + run = () => { + throw new PlaygroundRunError("Test error message", "TEST_CODE"); + }; + } + setupMockPlaygroundClient(PlaygroundClientMock); + + const response = await integration.mcpClient().callTool({ + name: "run-pipeline", + arguments: { + documents: [], + aggregationPipeline: [], + }, + }); + expect(response.content).toEqual([ + { + type: "text", + text: "Error running run-pipeline: Error code: TEST_CODE. Error message: Test error message.", + }, + ]); + }); });