Skip to content

Add PlaygroundClient client. Add tests for RunPipeline tool. #255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: search-skunkworks-2025
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 88 additions & 0 deletions src/common/playground/playgroundClient.ts
Original file line number Diff line number Diff line change
@@ -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,
};
}
}
}
101 changes: 20 additions & 81 deletions src/tools/playground/runPipeline.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.")
Expand All @@ -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 =
Expand All @@ -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),
Expand All @@ -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":`,
Expand Down
50 changes: 43 additions & 7 deletions tests/integration/tools/playground/runPipeline.test.ts
Original file line number Diff line number Diff line change
@@ -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: {
Expand All @@ -20,12 +39,6 @@ describeWithMongoDB("runPipeline tool", (integration) => {
},
},
},
{
$project: {
_id: 0,
name: 1,
},
},
],
},
});
Expand All @@ -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.",
},
]);
});
});
Loading