@assistant-ui/mcp-docs-server 0.1.39 → 0.2.1

package/src/tools/tests/completions.test.ts

@@ -1,46 +1,60 @@
-import { describe, it, expect } from "vitest";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
 import { server } from "../../index.js";
+import { createMcpTestClient, type McpTestClient } from "./mcp-test-client.js";
 
-function handler(method: string) {
-  const handlers = (server as any).server._requestHandlers;
-  const h = handlers.get(method);
-  if (!h) throw new Error(`No handler for ${method}`);
-  return h;
-}
+type CompletionResult = {
+  completion: {
+    values: string[];
+  };
+};
 
-async function complete(uri: string, name: string, value: string) {
-  return handler("completion/complete")(
-    {
+describe("resource argument completions", () => {
+  let client: McpTestClient;
+
+  beforeEach(async () => {
+    client = await createMcpTestClient(server);
+    await client.initialize();
+  });
+
+  afterEach(async () => {
+    await client?.close();
+  });
+
+  async function complete(uri: string, name: string, value: string) {
+    return client.request<CompletionResult>({
       method: "completion/complete",
-      params: { ref: { type: "ref/resource", uri }, argument: { name, value } },
-    },
-    {},
-  );
-}
+      params: {
+        ref: { type: "ref/resource", uri },
+        argument: { name, value },
+      },
+    });
+  }
 
-describe("resource argument completions", () => {
   it("suggests doc paths matching the partial value", async () => {
-    const res = await complete("aui-docs:///{+path}", "path", "api");
-    expect(res.completion.values.length).toBeGreaterThan(0);
+    const result = await complete("aui-docs:///{+path}", "path", "api");
+
+    expect(result.completion.values.length).toBeGreaterThan(0);
     expect(
-      res.completion.values.every((v: string) =>
-        v.toLowerCase().includes("api"),
+      result.completion.values.every((value) =>
+        value.toLowerCase().includes("api"),
       ),
     ).toBe(true);
   });
 
   it("suggests example names matching the partial value", async () => {
-    const res = await complete("aui-example:///{+name}", "name", "with");
-    expect(res.completion.values.length).toBeGreaterThan(0);
+    const result = await complete("aui-example:///{+name}", "name", "with");
+
+    expect(result.completion.values.length).toBeGreaterThan(0);
     expect(
-      res.completion.values.every((v: string) =>
-        v.toLowerCase().includes("with"),
+      result.completion.values.every((value) =>
+        value.toLowerCase().includes("with"),
       ),
     ).toBe(true);
   });
 
   it("returns no values for a nonsense partial", async () => {
-    const res = await complete("aui-docs:///{+path}", "path", "zzzqqxnope");
-    expect(res.completion.values).toHaveLength(0);
+    const result = await complete("aui-docs:///{+path}", "path", "zzzqqxnope");
+
+    expect(result.completion.values).toHaveLength(0);
   });
 });

package/src/tools/tests/docs.test.ts

@@ -42,7 +42,7 @@ describe("assistantUIDocs", () => {
     expect(result.found).toBe(true);
     expect(result.type).toBe("file");
     expect(result.content).toBeDefined();
-    expect(result.content).toContain("assistant-ui");
+    expect(result.content).toContain("title: Documentation");
   });
 
   it("should handle non-existent paths", async () => {
@@ -88,7 +88,7 @@ describe("assistantUIDocs", () => {
 
     expect(result.content).toBeDefined();
     expect(result.content).toContain("title:");
-    expect(result.content).toContain("assistant-ui");
+    expect(result.content).toContain("description:");
   });
 
   it("includes title and excerpt on a file response", async () => {

package/src/tools/tests/integration.test.ts

@@ -1,6 +1,6 @@
 import { describe, it, expect } from "vitest";
 import { server } from "../../index.js";
-import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { McpServer } from "@modelcontextprotocol/server";
 import { docsTools } from "../docs.js";
 import { examplesTools } from "../examples.js";
 
@@ -24,12 +24,11 @@ describe("MCP Server Integration", () => {
   });
 
   it("should have valid input schemas", () => {
-    // docsTools.parameters is the Zod schema shape
     expect(docsTools.parameters).toBeDefined();
-    expect(Object.keys(docsTools.parameters)).toContain("paths");
+    expect(docsTools.parameters.shape.paths).toBeDefined();
 
     expect(examplesTools.parameters).toBeDefined();
-    expect(Object.keys(examplesTools.parameters)).toContain("example");
+    expect(examplesTools.parameters.shape.example).toBeDefined();
   });
 
   it("should execute tools successfully", async () => {

package/src/tools/tests/mcp-protocol.test.ts

@@ -1,130 +1,160 @@
-import { describe, it, expect } from "vitest";
-import { server } from "../../index.js";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
 import {
-  InitializeRequestSchema,
-  ListToolsRequestSchema,
   CallToolRequestSchema,
-  ListPromptsRequestSchema,
   GetPromptRequestSchema,
-  type InitializeRequest,
-  type ListToolsRequest,
-  type CallToolRequest,
-  type ListPromptsRequest,
-  type GetPromptRequest,
-} from "@modelcontextprotocol/sdk/types.js";
+  InitializeRequestSchema,
+  ListPromptsRequestSchema,
+  ListToolsRequestSchema,
+} from "@modelcontextprotocol/core";
+import { server } from "../../index.js";
+import { createMcpTestClient, type McpTestClient } from "./mcp-test-client.js";
+
+type Tool = {
+  name: string;
+  title?: string;
+  annotations?: {
+    readOnlyHint?: boolean;
+    openWorldHint?: boolean;
+  };
+  inputSchema: {
+    type?: string;
+    properties?: Record<string, { description?: string }>;
+    required?: string[];
+  };
+};
+
+type ToolListResult = {
+  tools: Tool[];
+};
+
+type ToolCallResult = {
+  content: Array<{
+    type: string;
+  }>;
+};
+
+type PromptListResult = {
+  prompts: Array<{
+    name: string;
+    title?: string;
+  }>;
+};
+
+type PromptGetResult = {
+  messages: Array<{
+    content: {
+      type: string;
+      text?: string;
+    };
+  }>;
+};
+
+type InitializeResult = {
+  protocolVersion: string;
+  serverInfo: {
+    name: string;
+  };
+};
+
+function requireTool(tools: Tool[], name: string): Tool {
+  const tool = tools.find((candidate) => candidate.name === name);
+  expect(tool).toBeDefined();
+  if (!tool) throw new Error(`Missing tool: ${name}`);
+  return tool;
+}
 
 describe("MCP Protocol Integration", () => {
-  // These tests verify the MCP protocol layer handles requests correctly
-  // and that parameter schemas are properly converted to JSON schemas
-  it("should handle Initialize request", async () => {
-    const request: InitializeRequest = {
+  let client: McpTestClient;
+  let initializeResult: Promise<InitializeResult>;
+
+  beforeEach(async () => {
+    client = await createMcpTestClient(server);
+    initializeResult = client.initialize() as Promise<InitializeResult>;
+    await initializeResult;
+  });
+
+  afterEach(async () => {
+    await client?.close();
+  });
+
+  it("handles Initialize requests", async () => {
+    const request = {
       method: "initialize",
       params: {
-        protocolVersion: "2024-11-05",
+        protocolVersion: "2025-11-25",
         capabilities: {},
         clientInfo: {
-          name: "test-client",
+          name: "mcp-docs-server-tests",
           version: "1.0.0",
         },
       },
     };
 
-    // Parse and validate the request
-    const parsed = InitializeRequestSchema.parse(request);
-    expect(parsed).toBeDefined();
-
-    // The server should have an initialize handler set up
-    const handlers = (server as any).server._requestHandlers;
-    expect(handlers).toBeDefined();
-    expect(handlers.get("initialize")).toBeDefined();
+    expect(InitializeRequestSchema.parse(request)).toBeDefined();
+    await expect(initializeResult).resolves.toMatchObject({
+      protocolVersion: "2025-11-25",
+      serverInfo: { name: "assistant-ui-docs" },
+    });
   });
 
-  it("should handle ListTools request", async () => {
-    const request: ListToolsRequest = {
+  it("lists tool metadata and JSON schemas", async () => {
+    const request = ListToolsRequestSchema.parse({
       method: "tools/list",
       params: {},
-    };
-
-    // Parse and validate the request
-    const parsed = ListToolsRequestSchema.parse(request);
-    expect(parsed).toBeDefined();
-
-    // The server should have a tools/list handler
-    const handlers = (server as any).server._requestHandlers;
-    expect(handlers.get("tools/list")).toBeDefined();
-
-    // Call the handler
-    const handler = handlers.get("tools/list");
-    const result = await handler(parsed, {});
+    });
+    const result = await client.request<ToolListResult>(request);
 
-    expect(result).toBeDefined();
-    expect(result.tools).toBeInstanceOf(Array);
     expect(result.tools).toHaveLength(6);
 
-    // Check the tools have proper JSON schemas
-    const docsTool = result.tools.find(
-      (t: any) => t.name === "assistantUIDocs",
-    );
-    expect(docsTool).toBeDefined();
-    expect(docsTool.inputSchema).toBeDefined();
-    expect(docsTool.inputSchema.type).toBe("object");
-    expect(docsTool.inputSchema.properties).toBeDefined();
-
-    const examplesTool = result.tools.find(
-      (t: any) => t.name === "assistantUIExamples",
-    );
-    expect(examplesTool).toBeDefined();
-    expect(examplesTool.inputSchema).toBeDefined();
-    expect(examplesTool.inputSchema.type).toBe("object");
-    expect(examplesTool.inputSchema.properties).toBeDefined();
-
-    const searchTool = result.tools.find(
-      (t: any) => t.name === "assistantUISearch",
-    );
-    expect(searchTool).toBeDefined();
-    expect(searchTool.inputSchema).toBeDefined();
-    expect(searchTool.inputSchema.type).toBe("object");
-    expect(searchTool.inputSchema.properties).toBeDefined();
-
-    // registerTool metadata is surfaced on tools/list: `title` is a
-    // top-level field, `annotations` carries only the hint flags.
-    expect(docsTool.title).toBe("assistant-ui Documentation");
-    expect(docsTool.annotations?.readOnlyHint).toBe(true);
-    expect(docsTool.annotations?.openWorldHint).toBe(false);
-    expect(examplesTool.title).toBe("assistant-ui Examples");
-    expect(examplesTool.annotations?.readOnlyHint).toBe(true);
-    expect(examplesTool.annotations?.openWorldHint).toBe(false);
-    expect(searchTool.title).toBe("Search assistant-ui Documentation");
-    expect(searchTool.annotations?.readOnlyHint).toBe(true);
-    expect(searchTool.annotations?.openWorldHint).toBe(false);
-
-    const xuluxListTool = result.tools.find(
-      (t: any) => t.name === "assistantUITemplates",
-    );
-    expect(xuluxListTool).toBeDefined();
-    expect(xuluxListTool.title).toBe("assistant-ui Templates");
-    expect(xuluxListTool.annotations?.readOnlyHint).toBe(true);
-    expect(xuluxListTool.annotations?.openWorldHint).toBe(true);
-
-    const xuluxDetailsTool = result.tools.find(
-      (t: any) => t.name === "assistantUITemplateDetails",
-    );
-    expect(xuluxDetailsTool).toBeDefined();
-    expect(xuluxDetailsTool.title).toBe("assistant-ui Template Details");
-    expect(xuluxDetailsTool.annotations?.readOnlyHint).toBe(true);
-    expect(xuluxDetailsTool.annotations?.openWorldHint).toBe(true);
-
-    const xuluxPreviewTool = result.tools.find(
-      (t: any) => t.name === "assistantUITemplatePreview",
-    );
-    expect(xuluxPreviewTool).toBeDefined();
-    expect(xuluxPreviewTool.title).toBe("assistant-ui Template Preview URLs");
-    expect(xuluxPreviewTool.annotations?.readOnlyHint).toBe(false);
-    expect(xuluxPreviewTool.annotations?.openWorldHint).toBe(true);
+    for (const expected of [
+      {
+        name: "assistantUIDocs",
+        title: "assistant-ui Documentation",
+        annotations: { readOnlyHint: true, openWorldHint: false },
+      },
+      {
+        name: "assistantUIExamples",
+        title: "assistant-ui Examples",
+        annotations: { readOnlyHint: true, openWorldHint: false },
+      },
+      {
+        name: "assistantUISearch",
+        title: "Search assistant-ui Documentation",
+        annotations: { readOnlyHint: true, openWorldHint: false },
+      },
+      {
+        name: "assistantUITemplates",
+        title: "assistant-ui Templates",
+        annotations: { readOnlyHint: true, openWorldHint: true },
+      },
+      {
+        name: "assistantUITemplateDetails",
+        title: "assistant-ui Template Details",
+        annotations: { readOnlyHint: true, openWorldHint: true },
+      },
+      {
+        name: "assistantUITemplatePreview",
+        title: "assistant-ui Template Preview URLs",
+        annotations: { readOnlyHint: false, openWorldHint: true },
+      },
+    ]) {
+      const tool = requireTool(result.tools, expected.name);
+
+      expect(tool).toMatchObject({
+        title: expected.title,
+        annotations: expected.annotations,
+      });
+      expect(tool.inputSchema.type).toBe("object");
+      expect(tool.inputSchema.properties).toBeDefined();
+    }
+
+    const search = requireTool(result.tools, "assistantUISearch");
+    expect(search.inputSchema.required ?? []).not.toContain("limit");
+    expect(search.inputSchema.properties?.["query"]?.description).toBeDefined();
   });
 
-  it("should handle CallTool request for assistantUIDocs", async () => {
-    const request: CallToolRequest = {
+  it("calls assistantUIDocs", async () => {
+    const request = CallToolRequestSchema.parse({
       method: "tools/call",
       params: {
         name: "assistantUIDocs",
@@ -132,100 +162,55 @@ describe("MCP Protocol Integration", () => {
           paths: ["/"],
         },
       },
-    };
-
-    // Parse and validate the request
-    const parsed = CallToolRequestSchema.parse(request);
-    expect(parsed).toBeDefined();
-
-    // The server should have a tools/call handler
-    const handlers = (server as any).server._requestHandlers;
-    expect(handlers.get("tools/call")).toBeDefined();
+    });
+    const result = await client.request<ToolCallResult>(request);
 
-    // Call the handler through the MCP protocol layer
-    const handler = handlers.get("tools/call");
-    const result = await handler(parsed, {});
-
-    expect(result).toBeDefined();
-    expect(result.content).toBeDefined();
-    expect(result.content[0].type).toBe("text");
+    expect(result.content[0]?.type).toBe("text");
   });
 
-  it("should handle CallTool request for assistantUIExamples with no arguments", async () => {
-    const request: CallToolRequest = {
+  it("calls assistantUIExamples without arguments", async () => {
+    const request = CallToolRequestSchema.parse({
       method: "tools/call",
       params: {
         name: "assistantUIExamples",
         arguments: {},
       },
-    };
-
-    // Parse and validate the request
-    const parsed = CallToolRequestSchema.parse(request);
-    expect(parsed).toBeDefined();
+    });
+    const result = await client.request<ToolCallResult>(request);
 
-    // The server should have a tools/call handler
-    const handlers = (server as any).server._requestHandlers;
-    expect(handlers.get("tools/call")).toBeDefined();
-
-    // Call the handler
-    const handler = handlers.get("tools/call");
-    const result = await handler(parsed, {});
-
-    expect(result).toBeDefined();
-    expect(result.content).toBeDefined();
-    expect(result.content[0].type).toBe("text");
+    expect(result.content[0]?.type).toBe("text");
   });
 
-  it("should handle ListPrompts request", async () => {
-    const request: ListPromptsRequest = {
+  it("lists prompts", async () => {
+    const request = ListPromptsRequestSchema.parse({
       method: "prompts/list",
       params: {},
-    };
-
-    const parsed = ListPromptsRequestSchema.parse(request);
-    expect(parsed).toBeDefined();
-
-    const handlers = (server as any).server._requestHandlers;
-    expect(handlers.get("prompts/list")).toBeDefined();
-
-    const handler = handlers.get("prompts/list");
-    const result = await handler(parsed, {});
-
-    expect(result.prompts).toBeInstanceOf(Array);
+    });
+    const result = await client.request<PromptListResult>(request);
     const prompt = result.prompts.find(
-      (p: any) => p.name === "assistant-ui-template-workflow",
+      (candidate) => candidate.name === "assistant-ui-template-workflow",
     );
-    expect(prompt).toBeDefined();
-    expect(prompt.title).toBe("assistant-ui Template Workflow");
+
+    expect(prompt).toMatchObject({
+      title: "assistant-ui Template Workflow",
+    });
   });
 
-  it("should handle GetPrompt request for assistant-ui-template-workflow", async () => {
-    const request: GetPromptRequest = {
+  it("gets the assistant-ui-template-workflow prompt", async () => {
+    const request = GetPromptRequestSchema.parse({
       method: "prompts/get",
       params: {
         name: "assistant-ui-template-workflow",
         arguments: {},
       },
-    };
-
-    const parsed = GetPromptRequestSchema.parse(request);
-    expect(parsed).toBeDefined();
-
-    const handlers = (server as any).server._requestHandlers;
-    expect(handlers.get("prompts/get")).toBeDefined();
-
-    const handler = handlers.get("prompts/get");
-    const result = await handler(parsed, {});
+    });
+    const result = await client.request<PromptGetResult>(request);
+    const text = result.messages[0]?.content.text;
 
     expect(result.messages).toHaveLength(1);
-    expect(result.messages[0].content.type).toBe("text");
-    expect(result.messages[0].content.text).toContain("assistantUITemplates");
-    expect(result.messages[0].content.text).toContain(
-      "assistantUITemplateDetails",
-    );
-    expect(result.messages[0].content.text).toContain(
-      "assistantUITemplatePreview",
-    );
+    expect(result.messages[0]?.content.type).toBe("text");
+    expect(text).toContain("assistantUITemplates");
+    expect(text).toContain("assistantUITemplateDetails");
+    expect(text).toContain("assistantUITemplatePreview");
   });
 });

package/src/tools/tests/mcp-test-client.ts

@@ -0,0 +1,111 @@
+import {
+  InMemoryTransport,
+  type McpServer,
+} from "@modelcontextprotocol/server";
+
+type RawMcpRequest = {
+  method: string;
+  params?: unknown;
+};
+
+type JsonRpcResponse = {
+  id: number;
+  result?: unknown;
+  error?: {
+    code: number;
+    message: string;
+  };
+};
+
+type PendingRequest = {
+  resolve: (result: unknown) => void;
+  reject: (error: Error) => void;
+};
+
+export type McpTestClient = {
+  initialize: () => Promise<unknown>;
+  request: <Result>(request: RawMcpRequest) => Promise<Result>;
+  close: () => Promise<void>;
+};
+
+function isJsonRpcResponse(message: unknown): message is JsonRpcResponse {
+  return (
+    typeof message === "object" &&
+    message !== null &&
+    "id" in message &&
+    typeof message.id === "number" &&
+    ("result" in message || "error" in message)
+  );
+}
+
+export async function createMcpTestClient(
+  server: McpServer,
+): Promise<McpTestClient> {
+  const [clientTransport, serverTransport] =
+    InMemoryTransport.createLinkedPair();
+  const pendingRequests = new Map<number, PendingRequest>();
+  let nextRequestId = 1;
+
+  clientTransport.onmessage = (message) => {
+    if (!isJsonRpcResponse(message)) return;
+
+    const pendingRequest = pendingRequests.get(message.id);
+    if (!pendingRequest) return;
+
+    pendingRequests.delete(message.id);
+    if (message.error) {
+      pendingRequest.reject(
+        new Error(
+          `MCP request failed (${message.error.code}): ${message.error.message}`,
+        ),
+      );
+      return;
+    }
+
+    pendingRequest.resolve(message.result);
+  };
+
+  await clientTransport.start();
+  await server.connect(serverTransport);
+
+  const request = <Result>(rawRequest: RawMcpRequest): Promise<Result> => {
+    const id = nextRequestId++;
+
+    return new Promise((resolve, reject) => {
+      pendingRequests.set(id, { resolve, reject });
+      void clientTransport
+        .send({ jsonrpc: "2.0", id, ...rawRequest })
+        .catch((error: unknown) => {
+          pendingRequests.delete(id);
+          reject(error instanceof Error ? error : new Error(String(error)));
+        });
+    }) as Promise<Result>;
+  };
+
+  return {
+    initialize: async () => {
+      const result = await request({
+        method: "initialize",
+        params: {
+          protocolVersion: "2025-11-25",
+          capabilities: {},
+          clientInfo: {
+            name: "mcp-docs-server-tests",
+            version: "1.0.0",
+          },
+        },
+      });
+
+      await clientTransport.send({
+        jsonrpc: "2.0",
+        method: "notifications/initialized",
+      });
+
+      return result;
+    },
+    request,
+    close: async () => {
+      await clientTransport.close();
+    },
+  };
+}