Fix tests + move assistant request up to abstract base class

Author: nlarew

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<string, string>;
+    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<typeof this.argsShape>
-    ): Promise<CallToolResult> | 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),
+        });
     }
 }

src/tools/assistant/listKnowledgeSources.ts

@@ -28,10 +28,9 @@ export class ListKnowledgeSourcesTool extends AssistantToolBase {
     public operationType: OperationType = "read";
 
     protected async execute(): Promise<CallToolResult> {
-        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}`;

src/tools/assistant/searchKnowledge.ts

@@ -46,11 +46,10 @@ export class SearchKnowledgeTool extends AssistantToolBase {
     public operationType: OperationType = "read";
 
     protected async execute(args: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        const searchEndpoint = new URL("content/search", this.baseUrl);
-        const response = await fetch(searchEndpoint, {
+        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}`;

tests/integration/tools/assistant/assistantHelpers.ts

@@ -1,15 +1,20 @@
-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";
 
 export type IntegrationTestFunction = (integration: IntegrationTest) => void;
 
 export function describeWithAssistant(name: string, fn: IntegrationTestFunction): SuiteCollector<object> {
     const testDefinition = (): void => {
-        const integration = setupIntegrationTest(() => ({
-            ...defaultTestConfig,
-            assistantBaseUrl: "https://knowledge.test.mongodb.com/api/", // Use test URL
-        }));
+        const integration = setupIntegrationTest(
+            () => ({
+                ...defaultTestConfig,
+                assistantBaseUrl: "https://knowledge.test.mongodb.com/api/", // Use test URL
+            }),
+            () => ({
+                ...defaultDriverOptions,
+            })
+        );
 
         describe(name, () => {
             fn(integration);