@assistant-ui/mcp-docs-server 0.1.17 → 0.1.19

package/src/tools/examples.ts

@@ -0,0 +1,107 @@
+import { z } from "zod/v3";
+import { readFile, readdir, lstat } from "node:fs/promises";
+import { join, extname } from "node:path";
+import { CODE_EXAMPLES_PATH, MAX_FILE_SIZE } from "../constants.js";
+import { logger } from "../utils/logger.js";
+import { formatMCPResponse } from "../utils/mcp-format.js";
+import { sanitizePath } from "../utils/security.js";
+
+const examplesInputSchema = z.object({
+  example: z
+    .string()
+    .optional()
+    .describe(
+      'Example name (e.g., "with-ai-sdk-v6"). Leave empty to list all examples.',
+    ),
+});
+
+async function listCodeExamples(): Promise<string[]> {
+  try {
+    const files = await readdir(CODE_EXAMPLES_PATH);
+    return files
+      .filter((file) => extname(file) === ".md")
+      .map((file) => file.replace(".md", ""))
+      .sort();
+  } catch (error) {
+    logger.error("Failed to list code examples", error);
+    return [];
+  }
+}
+
+async function readCodeExample(exampleName: string): Promise<string | null> {
+  try {
+    const sanitized = sanitizePath(exampleName);
+    const filePath = join(CODE_EXAMPLES_PATH, `${sanitized}.md`);
+
+    const stats = await lstat(filePath);
+    if (stats.isSymbolicLink()) {
+      logger.warn(`Attempted to read symlink: ${filePath}`);
+      return null;
+    }
+
+    if (stats.size > MAX_FILE_SIZE) {
+      logger.warn(`File size exceeds limit: ${filePath} (${stats.size} bytes)`);
+      return null;
+    }
+
+    const content = await readFile(filePath, "utf-8");
+    return content;
+  } catch (error) {
+    logger.error(`Failed to read example: ${exampleName}`, error);
+    return null;
+  }
+}
+
+export const examplesTools = {
+  name: "assistantUIExamples",
+  description:
+    "List available examples or retrieve complete code for a specific example",
+  parameters: examplesInputSchema.shape,
+  execute: async ({ example }: z.infer<typeof examplesInputSchema>) => {
+    try {
+      if (!example) {
+        logger.info("Listing all available examples");
+        const examples = await listCodeExamples();
+
+        if (examples.length === 0) {
+          return formatMCPResponse({
+            error:
+              "No examples found. Please run documentation preparation first.",
+            hint: "Run: pnpm prepare-docs",
+          });
+        }
+
+        return formatMCPResponse({
+          type: "list",
+          examples,
+          total: examples.length,
+          hint: "Use example parameter to get complete code for any example",
+        });
+      }
+
+      logger.info(`Retrieving example: ${example}`);
+      const content = await readCodeExample(example);
+
+      if (!content) {
+        const availableExamples = await listCodeExamples();
+        return formatMCPResponse({
+          error: `Example not found: ${example}`,
+          availableExamples,
+          hint: "Use without example parameter to list all available examples",
+        });
+      }
+
+      return formatMCPResponse({
+        type: "example",
+        name: example,
+        content,
+      });
+    } catch (error) {
+      logger.error("Failed to retrieve examples", error);
+      return formatMCPResponse({
+        error: "Failed to retrieve examples",
+        message: error instanceof Error ? error.message : String(error),
+      });
+    }
+  },
+};

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

@@ -0,0 +1,124 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+import { testContext } from "./test-setup.js";
+import * as fs from "node:fs/promises";
+
+vi.mock("fs/promises", async () => {
+  const { readdir, readFile, stat, lstat } = await import("node:fs/promises");
+  return {
+    readdir,
+    readFile,
+    stat,
+    lstat: vi.fn(lstat),
+  };
+});
+
+describe("assistantUIDocs", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+  it("should list root directory contents", async () => {
+    const result = await testContext.callTool("assistantUIDocs", {
+      paths: ["/"],
+    });
+
+    expect(result.path).toBe("/");
+    expect(result.found).toBe(true);
+    expect(result.type).toBe("directory");
+    expect(result.directories).toContain("(docs)");
+    expect(result.directories).toContain("(reference)");
+    expect(result.directories).toContain("ui");
+  });
+
+  it("should retrieve specific documentation file", async () => {
+    const result = await testContext.callTool("assistantUIDocs", {
+      paths: ["(docs)/index"],
+    });
+
+    expect(result.path).toBe("(docs)/index");
+    expect(result.found).toBe(true);
+    expect(result.type).toBe("file");
+    expect(result.content).toBeDefined();
+    expect(result.content).toContain("assistant-ui");
+  });
+
+  it("should handle non-existent paths", async () => {
+    const result = await testContext.callTool("assistantUIDocs", {
+      paths: ["non-existent-path"],
+    });
+
+    expect(result.error).toBeDefined();
+    expect(result.error).toContain("Documentation not found");
+    expect(result.suggestions).toBeDefined();
+  });
+
+  it("should support multiple path requests", async () => {
+    const result = await testContext.callTool("assistantUIDocs", {
+      paths: ["(docs)/index", "(reference)/api-reference/primitives/thread"],
+    });
+
+    expect(result.results).toBeDefined();
+    expect(result.results).toHaveLength(2);
+    expect(result.results[0].path).toBe("(docs)/index");
+    expect(result.results[1].path).toBe(
+      "(reference)/api-reference/primitives/thread",
+    );
+  });
+
+  it("should list directory contents with files", async () => {
+    const result = await testContext.callTool("assistantUIDocs", {
+      paths: ["(reference)/api-reference/primitives"],
+    });
+
+    expect(result.path).toBe("(reference)/api-reference/primitives");
+    expect(result.found).toBe(true);
+    expect(result.type).toBe("directory");
+    expect(result.files).toContain("thread");
+    expect(result.files).toContain("message");
+    expect(result.files).toContain("composer");
+  });
+
+  it("should parse MDX files with frontmatter", async () => {
+    const result = await testContext.callTool("assistantUIDocs", {
+      paths: ["(docs)/index"],
+    });
+
+    expect(result.content).toBeDefined();
+    expect(result.content).toContain("title:");
+    expect(result.content).toContain("assistant-ui");
+  });
+
+  it("should skip symlinks and large files", async () => {
+    const mockedLstat = vi.mocked(fs.lstat);
+
+    mockedLstat.mockResolvedValueOnce({
+      isSymbolicLink: () => true,
+      isFile: () => false,
+      isDirectory: () => false,
+    } as any);
+
+    const symlinkResult = await testContext.callTool("assistantUIDocs", {
+      paths: ["symlink-test"],
+    });
+    expect(symlinkResult.error).toBe(
+      "Symlinks are not allowed for security reasons",
+    );
+
+    mockedLstat.mockRejectedValueOnce(new Error("ENOENT"));
+    mockedLstat.mockResolvedValueOnce({
+      isSymbolicLink: () => false,
+      isFile: () => true,
+      size: 11 * 1024 * 1024,
+    } as any);
+
+    const largeFileResult = await testContext.callTool("assistantUIDocs", {
+      paths: ["large-file"],
+    });
+    expect(largeFileResult.error).toContain(
+      "File size exceeds maximum allowed size",
+    );
+  });
+});

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

@@ -0,0 +1,94 @@
+import { describe, it, expect, vi } from "vitest";
+import { testContext } from "./test-setup.js";
+import * as fs from "node:fs/promises";
+
+vi.mock("fs/promises", async () => {
+  const { readdir, readFile, lstat } = await import("node:fs/promises");
+  return {
+    readdir,
+    readFile,
+    lstat: vi.fn(lstat),
+  };
+});
+
+describe("assistantUIExamples", () => {
+  it("should list all available examples", async () => {
+    const result = await testContext.callTool("assistantUIExamples", {});
+
+    expect(result.type).toBe("list");
+    expect(result.examples).toBeDefined();
+    expect(Array.isArray(result.examples)).toBe(true);
+    expect(result.examples.length).toBeGreaterThan(0);
+    expect(result.examples).toContain("with-ai-sdk-v6");
+    expect(result.examples).toContain("with-langgraph");
+    expect(result.total).toBe(result.examples.length);
+  });
+
+  it("should retrieve specific example code", async () => {
+    const result = await testContext.callTool("assistantUIExamples", {
+      example: "with-ai-sdk-v6",
+    });
+
+    expect(result.type).toBe("example");
+    expect(result.name).toBe("with-ai-sdk-v6");
+    expect(result.content).toBeDefined();
+    expect(result.content).toContain("# Example: with-ai-sdk-v6");
+    expect(result.content).toContain("app/api/chat/route.ts");
+    expect(result.content).toContain("streamText");
+  });
+
+  it("should handle non-existent examples", async () => {
+    const result = await testContext.callTool("assistantUIExamples", {
+      example: "non-existent-example",
+    });
+
+    expect(result.error).toBeDefined();
+    expect(result.error).toContain("Example not found");
+    expect(result.availableExamples).toBeDefined();
+    expect(Array.isArray(result.availableExamples)).toBe(true);
+  });
+
+  it("should include all files in example", async () => {
+    const result = await testContext.callTool("assistantUIExamples", {
+      example: "with-ai-sdk-v6",
+    });
+
+    expect(result.content).toContain("package.json");
+    expect(result.content).toContain("components/assistant-ui/thread.tsx");
+  });
+
+  it("should handle empty example parameter", async () => {
+    const result = await testContext.callTool("assistantUIExamples", {
+      example: undefined,
+    });
+
+    expect(result.type).toBe("list");
+    expect(result.hint).toBeDefined();
+  });
+
+  it("should skip symlinks and large files", async () => {
+    const mockedLstat = vi.mocked(fs.lstat);
+
+    mockedLstat.mockResolvedValueOnce({
+      isSymbolicLink: () => true,
+      isFile: () => false,
+      size: 0,
+    } as any);
+
+    const symlinkResult = await testContext.callTool("assistantUIExamples", {
+      example: "symlink-example",
+    });
+    expect(symlinkResult.error).toContain("Example not found");
+
+    mockedLstat.mockResolvedValueOnce({
+      isSymbolicLink: () => false,
+      isFile: () => true,
+      size: 11 * 1024 * 1024, // 11MB - exceeds MAX_FILE_SIZE
+    } as any);
+
+    const largeFileResult = await testContext.callTool("assistantUIExamples", {
+      example: "large-example",
+    });
+    expect(largeFileResult.error).toContain("Example not found");
+  });
+});

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

@@ -0,0 +1,46 @@
+import { describe, it, expect } from "vitest";
+import { server } from "../../index.js";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { docsTools } from "../docs.js";
+import { examplesTools } from "../examples.js";
+
+describe("MCP Server Integration", () => {
+  it("should be an instance of McpServer", () => {
+    expect(server).toBeInstanceOf(McpServer);
+  });
+
+  it("should have tools with correct properties", () => {
+    expect(docsTools.name).toBe("assistantUIDocs");
+    expect(docsTools.description).toContain(
+      "Retrieve assistant-ui documentation",
+    );
+    expect(docsTools.parameters).toBeDefined();
+    expect(docsTools.execute).toBeInstanceOf(Function);
+
+    expect(examplesTools.name).toBe("assistantUIExamples");
+    expect(examplesTools.description).toContain("List available examples");
+    expect(examplesTools.parameters).toBeDefined();
+    expect(examplesTools.execute).toBeInstanceOf(Function);
+  });
+
+  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(examplesTools.parameters).toBeDefined();
+    expect(Object.keys(examplesTools.parameters)).toContain("example");
+  });
+
+  it("should execute tools successfully", async () => {
+    const docsResult = await docsTools.execute({ paths: ["/"] });
+    expect(docsResult).toBeDefined();
+    expect(docsResult.content).toBeDefined();
+    expect(docsResult.content[0].type).toBe("text");
+
+    const examplesResult = await examplesTools.execute({});
+    expect(examplesResult).toBeDefined();
+    expect(examplesResult.content).toBeDefined();
+    expect(examplesResult.content[0].type).toBe("text");
+  });
+});

package/src/tools/tests/json-parsing.test.ts

@@ -0,0 +1,23 @@
+import { describe, it, expect, vi, afterEach } from "vitest";
+import { testContext } from "./test-setup.js";
+import { docsTools } from "../docs.js";
+
+describe("JSON parsing error handling", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("should provide helpful error message for invalid JSON", async () => {
+    vi.spyOn(docsTools, "execute").mockResolvedValue({
+      content: [{ text: "invalid json {not valid}" }],
+    });
+
+    await expect(
+      testContext.callTool("assistantUIDocs", { paths: ["/"] }),
+    ).rejects.toThrow(/Tool assistantUIDocs returned invalid JSON/);
+
+    await expect(
+      testContext.callTool("assistantUIDocs", { paths: ["/"] }),
+    ).rejects.toThrow(/invalid json \{not valid\}/);
+  });
+});

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

@@ -0,0 +1,133 @@
+import { describe, it, expect } from "vitest";
+import { server } from "../../index.js";
+import {
+  InitializeRequestSchema,
+  ListToolsRequestSchema,
+  CallToolRequestSchema,
+  type InitializeRequest,
+  type ListToolsRequest,
+  type CallToolRequest,
+} from "@modelcontextprotocol/sdk/types.js";
+
+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 = {
+      method: "initialize",
+      params: {
+        protocolVersion: "2024-11-05",
+        capabilities: {
+          tools: {},
+        },
+        clientInfo: {
+          name: "test-client",
+          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();
+  });
+
+  it("should handle ListTools request", async () => {
+    const request: ListToolsRequest = {
+      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, {});
+
+    expect(result).toBeDefined();
+    expect(result.tools).toBeInstanceOf(Array);
+    expect(result.tools).toHaveLength(2);
+
+    // 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();
+  });
+
+  it("should handle CallTool request for assistantUIDocs", async () => {
+    const request: CallToolRequest = {
+      method: "tools/call",
+      params: {
+        name: "assistantUIDocs",
+        arguments: {
+          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();
+
+    // 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");
+  });
+
+  it("should handle CallTool request for assistantUIExamples with no arguments", async () => {
+    const request: CallToolRequest = {
+      method: "tools/call",
+      params: {
+        name: "assistantUIExamples",
+        arguments: {},
+      },
+    };
+
+    // 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();
+
+    // 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");
+  });
+});

package/src/tools/tests/path-traversal.test.ts

@@ -0,0 +1,81 @@
+import { describe, it, expect } from "vitest";
+import { testContext } from "./test-setup.js";
+
+describe("Path Traversal Security", () => {
+  describe("assistantUIDocs tool", () => {
+    const maliciousPaths = [
+      "../../../../etc/passwd",
+      "../../../package.json",
+      "..\\..\\..\\windows\\system32",
+      "/etc/passwd",
+      "docs/../../../sensitive-file",
+      "./../../private/keys",
+    ];
+
+    maliciousPaths.forEach((path) => {
+      it(`should block path traversal attempt: ${path}`, async () => {
+        const result = await testContext.callTool("assistantUIDocs", {
+          paths: [path],
+        });
+
+        expect(result.error).toBeDefined();
+        expect(result.error).toContain("Invalid path");
+        expect(result.content).toBeUndefined();
+      });
+    });
+
+    it("should handle multiple paths with one malicious", async () => {
+      const result = await testContext.callTool("assistantUIDocs", {
+        paths: ["(docs)/index", "../../../../etc/passwd", "(docs)/guides"],
+      });
+
+      expect(result.results).toBeDefined();
+      expect(result.results).toHaveLength(3);
+
+      expect(result.results[0].found).toBe(true);
+      expect(result.results[1].error).toContain("Invalid path");
+      expect(result.results[2].found).toBe(true);
+    });
+  });
+
+  describe("assistantUIExamples tool", () => {
+    const maliciousExamples = [
+      "../../../../etc/passwd",
+      "../../../src/index",
+      "..\\..\\..\\config",
+      "/root/.ssh/id_rsa",
+      "examples/../../../private",
+    ];
+
+    maliciousExamples.forEach((example) => {
+      it(`should block path traversal attempt: ${example}`, async () => {
+        const result = await testContext.callTool("assistantUIExamples", {
+          example,
+        });
+
+        expect(result.error).toBeDefined();
+        expect(result.error).toContain("Example not found");
+        expect(result.content).toBeUndefined();
+      });
+    });
+  });
+
+  describe("Valid paths should still work", () => {
+    it("should allow valid documentation paths", async () => {
+      const result = await testContext.callTool("assistantUIDocs", {
+        paths: ["(docs)/index", "(reference)/api-reference/primitives/thread"],
+      });
+
+      expect(result.results).toBeDefined();
+      expect(result.results).toHaveLength(2);
+      expect(result.results.every((r: any) => r.found)).toBe(true);
+    });
+
+    it("should allow valid example names", async () => {
+      const result = await testContext.callTool("assistantUIExamples", {});
+
+      expect(result.examples).toBeDefined();
+      expect(Array.isArray(result.examples)).toBe(true);
+    });
+  });
+});

package/src/tools/tests/test-setup.ts

@@ -0,0 +1,40 @@
+import { beforeAll } from "vitest";
+import { existsSync } from "node:fs";
+import { join } from "node:path";
+import { PACKAGE_DIR } from "../../constants.js";
+import { docsTools } from "../docs.js";
+import { examplesTools } from "../examples.js";
+
+const tools = {
+  assistantUIDocs: docsTools,
+  assistantUIExamples: examplesTools,
+};
+
+export const testContext = {
+  callTool: async (name: string, args: any) => {
+    const tool = tools[name as keyof typeof tools];
+    if (!tool) {
+      throw new Error(`Tool ${name} not found`);
+    }
+    const result = await tool.execute(args);
+
+    const text = result.content?.[0]?.text;
+    if (text === undefined) {
+      throw new Error(`Tool ${name} returned no content`);
+    }
+    try {
+      return JSON.parse(text);
+    } catch (error) {
+      throw new Error(
+        `Tool ${name} returned invalid JSON. Output: ${text}\nParse error: ${error instanceof Error ? error.message : String(error)}`,
+      );
+    }
+  },
+};
+
+beforeAll(() => {
+  const docsPath = join(PACKAGE_DIR, ".docs");
+  if (!existsSync(docsPath)) {
+    throw new Error("Documentation not prepared. Run: pnpm build");
+  }
+});

package/src/utils/logger.ts

@@ -0,0 +1,20 @@
+import { IS_PREPARE_MODE } from "../constants.js";
+
+export const logger = {
+  debug: (message: string, ...args: any[]) => {
+    if (process.env["DEBUG"]) {
+      console.debug(`[DEBUG] ${message}`, ...args);
+    }
+  },
+  info: (message: string, ...args: any[]) => {
+    if (IS_PREPARE_MODE) {
+      console.log(`[INFO] ${message}`, ...args);
+    }
+  },
+  error: (message: string, ...args: any[]) => {
+    console.error(`[ERROR] ${message}`, ...args);
+  },
+  warn: (message: string, ...args: any[]) => {
+    console.warn(`[WARN] ${message}`, ...args);
+  },
+};

package/src/utils/mcp-format.ts

@@ -0,0 +1,12 @@
+export function formatMCPResponse(data: any): {
+  content: Array<{ type: "text"; text: string }>;
+} {
+  return {
+    content: [
+      {
+        type: "text",
+        text: typeof data === "string" ? data : JSON.stringify(data, null, 2),
+      },
+    ],
+  };
+}