Add PlaygroundClient client. Add tests for RunPipeline tool.

vm-mishchenko · vm-mishchenko · commit e74e7d20f370 · 2025-05-15T11:59:13.000-07:00
diff --git a/src/common/playground/playgroundClient.ts 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<Record<string, unknown>>;
+}
+
+/**
+ * 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<PlaygroundRunResponse> {
+        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<PlaygroundRunErrorResponse> {
+        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
@@ -2,50 +2,28 @@ 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: {
         dynamic: true,
     },
 };
 
-const DEFAULT_PIPELINE = [
-    {
-        $search: {
-            index: "default",
-            text: {
-                query: "first",
-                path: {
-                    wildcard: "*",
-                },
-            },
-        },
-    },
-];
-
 const DEFAULT_SYNONYMS: Array<Record<string, unknown>> = [];
 
 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<Record<string, unknown>>;
-}
-
-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<RunResponse> {
-        const options: RequestInit = {
-            method: "POST",
-            headers: {
-                "Content-Type": "application/json",
-            },
-            body: JSON.stringify(runRequest),
-        };
-
-        let response: Response;
+    private async runPipeline(runRequest: PlaygroundRunRequest): Promise<PlaygroundRunResponse> {
+        // 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<string> {
-        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<typeof this.argsShape>): RunRequest {
+    private convertToRunRequest(toolArgs: ToolArgs<typeof this.argsShape>): 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
@@ -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.",
+            },
+        ]);
+    });
 });
