@assistant-ui/mcp-docs-server 0.2.3 → 0.3.1

package/src/proxy.test.ts

@@ -0,0 +1,385 @@
+import { describe, expect, it, vi } from "vitest";
+import { once } from "node:events";
+import {
+  createServer,
+  type IncomingMessage,
+  type ServerResponse,
+} from "node:http";
+import { PassThrough } from "node:stream";
+import {
+  McpServer,
+  WebStandardStreamableHTTPServerTransport,
+  fromJsonSchema,
+} from "@modelcontextprotocol/server";
+import { runProxy } from "./proxy.js";
+
+type RpcResponse = {
+  jsonrpc: "2.0";
+  id: number;
+  result?: Record<string, unknown>;
+  error?: {
+    code: number;
+    message: string;
+  };
+};
+
+const toWebRequest = async (request: IncomingMessage) => {
+  const chunks: Buffer[] = [];
+  for await (const chunk of request) {
+    chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
+  }
+
+  const headers = new Headers();
+  for (const [name, value] of Object.entries(request.headers)) {
+    if (Array.isArray(value)) {
+      for (const item of value) headers.append(name, item);
+    } else if (value !== undefined) {
+      headers.set(name, value);
+    }
+  }
+
+  const init: RequestInit = {
+    method: request.method ?? "GET",
+    headers,
+  };
+  if (chunks.length > 0) {
+    init.body = Buffer.concat(chunks);
+  }
+
+  return new Request(new URL(request.url ?? "/", "http://127.0.0.1"), init);
+};
+
+const writeWebResponse = async (response: Response, output: ServerResponse) => {
+  output.statusCode = response.status;
+  response.headers.forEach((value, name) => {
+    output.setHeader(name, value);
+  });
+  output.end(Buffer.from(await response.arrayBuffer()));
+};
+
+const listen = (server: ReturnType<typeof createServer>) =>
+  new Promise<URL>((resolve, reject) => {
+    const onError = (error: Error) => {
+      reject(error);
+    };
+    server.once("error", onError);
+    server.listen(0, "127.0.0.1", () => {
+      server.off("error", onError);
+      const address = server.address();
+      if (address === null || typeof address === "string") {
+        reject(new Error("HTTP server did not bind to a TCP port"));
+        return;
+      }
+      resolve(new URL(`http://127.0.0.1:${address.port}/mcp`));
+    });
+  });
+
+const closeHttpServer = (server: ReturnType<typeof createServer>) =>
+  new Promise<void>((resolve, reject) => {
+    server.close((error) => {
+      if (error) reject(error);
+      else resolve();
+    });
+    server.closeAllConnections();
+  });
+
+const withTimeout = <T>(promise: Promise<T>, label: string) =>
+  new Promise<T>((resolve, reject) => {
+    const timeout = setTimeout(() => {
+      reject(new Error(`Timed out waiting for ${label}`));
+    }, 2_000);
+    void promise.then(
+      (value) => {
+        clearTimeout(timeout);
+        resolve(value);
+      },
+      (error: unknown) => {
+        clearTimeout(timeout);
+        reject(error);
+      },
+    );
+  });
+
+const createRpcClient = (stdin: PassThrough, stdout: PassThrough) => {
+  let nextId = 1;
+  let buffer = "";
+  const pending = new Map<
+    number,
+    {
+      resolve(response: RpcResponse): void;
+      reject(error: Error): void;
+      timeout: ReturnType<typeof setTimeout>;
+    }
+  >();
+
+  stdout.on("data", (chunk: Buffer) => {
+    buffer += chunk.toString();
+    const lines = buffer.split("\n");
+    buffer = lines.pop() ?? "";
+
+    for (const line of lines) {
+      if (line.length === 0) continue;
+      const response = JSON.parse(line) as RpcResponse;
+      const request = pending.get(response.id);
+      if (!request) continue;
+      clearTimeout(request.timeout);
+      pending.delete(response.id);
+      request.resolve(response);
+    }
+  });
+
+  return {
+    request(method: string, params: Record<string, unknown>) {
+      const id = nextId++;
+      return new Promise<RpcResponse>((resolve, reject) => {
+        const timeout = setTimeout(() => {
+          pending.delete(id);
+          reject(new Error(`Timed out waiting for ${method}`));
+        }, 2_000);
+        pending.set(id, { resolve, reject, timeout });
+        stdin.write(
+          `${JSON.stringify({ jsonrpc: "2.0", id, method, params })}\n`,
+        );
+      });
+    },
+    notify(method: string) {
+      stdin.write(`${JSON.stringify({ jsonrpc: "2.0", method })}\n`);
+    },
+  };
+};
+
+describe("runProxy", () => {
+  it.each([
+    {
+      scenario: "stdin reached EOF before startup",
+      createStdin: async () => {
+        const stdin = new PassThrough();
+        const ended = once(stdin, "end");
+        stdin.resume();
+        stdin.end();
+        await ended;
+        return stdin;
+      },
+      stopInput: (stdin: PassThrough) => stdin.end(),
+    },
+    {
+      scenario: "stdin reaches EOF after startup",
+      createStdin: async () => new PassThrough(),
+      stopInput: (stdin: PassThrough) => stdin.end(),
+    },
+    {
+      scenario: "stdin ends without emitting close",
+      createStdin: async () => new PassThrough({ autoDestroy: false }),
+      stopInput: (stdin: PassThrough) => stdin.end(),
+    },
+    {
+      scenario: "stdin is destroyed without EOF",
+      createStdin: async () => new PassThrough(),
+      stopInput: (stdin: PassThrough) => stdin.destroy(),
+    },
+  ])("closes when $scenario", async ({ createStdin, stopInput }) => {
+    const stdin = await createStdin();
+    const stdout = new PassThrough();
+    const stderr = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+    const proxy = runProxy({
+      url: new URL("https://example.invalid/mcp"),
+      stdin,
+      stdout,
+    });
+    let closed = false;
+
+    try {
+      await new Promise<void>((resolve) => setImmediate(resolve));
+      stopInput(stdin);
+      await expect(
+        withTimeout(proxy, "proxy shutdown after stdin closed"),
+      ).resolves.toBeUndefined();
+      closed = true;
+    } finally {
+      stdout.destroy(closed ? undefined : new Error("Proxy test cleanup"));
+      await proxy;
+      stdin.destroy();
+      stderr.mockRestore();
+    }
+  });
+
+  it("does not log errors caused by clean shutdown", async () => {
+    let resolvePendingRequest = () => {};
+    const pendingRequest = new Promise<void>((resolve) => {
+      resolvePendingRequest = resolve;
+    });
+    const httpServer = createServer(async (request, response) => {
+      if (request.method !== "POST") {
+        response.statusCode = 405;
+        response.end();
+        return;
+      }
+
+      let body = "";
+      for await (const chunk of request) body += chunk.toString();
+      const message = JSON.parse(body) as { method: string; id?: number };
+      if (message.method === "initialize") {
+        response.setHeader("content-type", "application/json");
+        response.end(
+          JSON.stringify({
+            jsonrpc: "2.0",
+            id: message.id,
+            result: {
+              protocolVersion: "2025-11-25",
+              capabilities: {},
+              serverInfo: { name: "proxy-test", version: "1.0.0" },
+            },
+          }),
+        );
+        return;
+      }
+
+      response.writeHead(200, {
+        "content-type": "text/event-stream",
+        "cache-control": "no-cache",
+      });
+      response.write(": keep-alive\n\n");
+      resolvePendingRequest();
+    });
+    const url = await listen(httpServer);
+    const proxyStdin = new PassThrough();
+    const proxyStdout = new PassThrough();
+    const stderr = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+    const proxy = runProxy({ url, stdin: proxyStdin, stdout: proxyStdout });
+    const client = createRpcClient(proxyStdin, proxyStdout);
+    let proxyClosed = false;
+
+    try {
+      const initialize = await client.request("initialize", {
+        protocolVersion: "2025-11-25",
+        capabilities: {},
+        clientInfo: { name: "proxy-test-client", version: "1.0.0" },
+      });
+      expect(initialize.error).toBeUndefined();
+
+      client.notify("notifications/initialized");
+      await withTimeout(pendingRequest, "pending request");
+      proxyStdin.end();
+      await expect(
+        withTimeout(proxy, "proxy shutdown after stdin closed"),
+      ).resolves.toBeUndefined();
+      proxyClosed = true;
+
+      expect(stderr).not.toHaveBeenCalled();
+    } finally {
+      if (!proxyClosed) {
+        proxyStdin.destroy();
+        await proxy.catch(() => undefined);
+      }
+      await closeHttpServer(httpServer);
+      proxyStdout.destroy();
+      proxyStdin.destroy();
+      stderr.mockRestore();
+    }
+  });
+
+  it("proxies MCP requests and returns transport failures", async () => {
+    const inputSchema = fromJsonSchema<{ text: string }>({
+      type: "object",
+      properties: {
+        text: { type: "string" },
+      },
+      required: ["text"],
+      additionalProperties: false,
+    });
+    const mcpServer = new McpServer({ name: "proxy-test", version: "1.0.0" });
+    mcpServer.registerTool(
+      "echo",
+      { description: "Echo text", inputSchema },
+      ({ text }) => ({ content: [{ type: "text", text }] }),
+    );
+    const transport = new WebStandardStreamableHTTPServerTransport({
+      sessionIdGenerator: undefined,
+      enableJsonResponse: true,
+    });
+    let resolveInitialized = () => {};
+    const initialized = new Promise<void>((resolve) => {
+      resolveInitialized = resolve;
+    });
+    mcpServer.server.oninitialized = resolveInitialized;
+    await mcpServer.connect(transport);
+
+    const httpServer = createServer((request, response) => {
+      void (async () => {
+        const webRequest = await toWebRequest(request);
+        const webResponse =
+          webRequest.method === "GET"
+            ? Response.json({
+                name: "assistant-ui",
+                transport: "streamable-http",
+              })
+            : await transport.handleRequest(webRequest);
+        await writeWebResponse(webResponse, response);
+      })().catch((error: unknown) => {
+        response.statusCode = 500;
+        response.end(error instanceof Error ? error.message : String(error));
+      });
+    });
+    const url = await listen(httpServer);
+    const proxyStdin = new PassThrough();
+    const proxyStdout = new PassThrough();
+    const stderr = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+    const proxy = runProxy({ url, stdin: proxyStdin, stdout: proxyStdout });
+    const client = createRpcClient(proxyStdin, proxyStdout);
+    let httpServerClosed = false;
+
+    try {
+      const initialize = await client.request("initialize", {
+        protocolVersion: "2025-11-25",
+        capabilities: {},
+        clientInfo: { name: "proxy-test-client", version: "1.0.0" },
+      });
+      expect(initialize.error).toBeUndefined();
+      expect(initialize.result).toMatchObject({
+        protocolVersion: "2025-11-25",
+        serverInfo: { name: "proxy-test", version: "1.0.0" },
+      });
+
+      client.notify("notifications/initialized");
+      await withTimeout(initialized, "initialized notification");
+
+      const tools = await client.request("tools/list", {});
+      expect(tools.error).toBeUndefined();
+      expect(tools.result?.tools).toContainEqual(
+        expect.objectContaining({ name: "echo" }),
+      );
+
+      const call = await client.request("tools/call", {
+        name: "echo",
+        arguments: { text: "hello through the proxy" },
+      });
+      expect(call.error).toBeUndefined();
+      expect(call.result?.content).toEqual([
+        { type: "text", text: "hello through the proxy" },
+      ]);
+
+      await closeHttpServer(httpServer);
+      httpServerClosed = true;
+
+      const failure = await client.request("tools/list", {});
+      expect(failure.error).toEqual({
+        code: -32603,
+        message: "Failed to proxy request to the hosted MCP server",
+      });
+      expect(stderr).toHaveBeenCalled();
+    } finally {
+      if (!httpServerClosed) await closeHttpServer(httpServer);
+      proxyStdout.destroy(new Error("Proxy test cleanup"));
+      await proxy;
+      proxyStdin.destroy();
+      await mcpServer.close();
+      stderr.mockRestore();
+    }
+  });
+});

package/src/proxy.ts

@@ -0,0 +1,126 @@
+import {
+  StreamableHTTPClientTransport,
+  isInitializeRequest,
+  isJSONRPCRequest,
+  isJSONRPCResultResponse,
+} from "@modelcontextprotocol/client";
+import { StdioServerTransport } from "@modelcontextprotocol/server/stdio";
+import { type Readable, type Writable, finished } from "node:stream";
+
+const DEFAULT_URL = "https://www.assistant-ui.com/mcp";
+
+const logError = (message: string, error: unknown) => {
+  console.error(`assistant-ui MCP proxy: ${message}`, error);
+};
+
+export async function runProxy({
+  url = new URL(process.env.ASSISTANT_UI_MCP_URL ?? DEFAULT_URL),
+  stdin = process.stdin,
+  stdout = process.stdout,
+}: {
+  url?: URL;
+  stdin?: NodeJS.ReadableStream;
+  stdout?: NodeJS.WritableStream;
+} = {}) {
+  const input = stdin as Readable;
+  const stdio = new StdioServerTransport(input, stdout as Writable);
+  const http = new StreamableHTTPClientTransport(url);
+  let initializeRequestId: string | number | undefined;
+  let closing = false;
+  let resolveClosed = () => {};
+  const closed = new Promise<void>((resolve) => {
+    resolveClosed = resolve;
+  });
+
+  const closeCounterpart = (counterpart: { close(): Promise<void> }) => {
+    if (closing) return;
+    closing = true;
+    void counterpart
+      .close()
+      .catch((error: unknown) => {
+        logError("failed to close transport", error);
+      })
+      .then(resolveClosed);
+  };
+
+  const logUnlessClosing = (message: string, error: unknown) => {
+    if (closing) return;
+    logError(message, error);
+  };
+
+  stdio.onmessage = (message) => {
+    if (isJSONRPCRequest(message) && isInitializeRequest(message)) {
+      initializeRequestId = message.id;
+    }
+
+    void http.send(message).catch(async (error: unknown) => {
+      if (closing) return;
+      logError("failed to forward message to hosted endpoint", error);
+      if (!isJSONRPCRequest(message)) return;
+
+      await stdio
+        .send({
+          jsonrpc: "2.0",
+          id: message.id,
+          error: {
+            code: -32603,
+            message: "Failed to proxy request to the hosted MCP server",
+          },
+        })
+        .catch((sendError: unknown) => {
+          logUnlessClosing("failed to return proxy error response", sendError);
+        });
+    });
+  };
+
+  http.onmessage = (message) => {
+    if (
+      isJSONRPCResultResponse(message) &&
+      message.id === initializeRequestId
+    ) {
+      const protocolVersion = (
+        message.result as { protocolVersion?: unknown } | null
+      )?.protocolVersion;
+      if (typeof protocolVersion === "string") {
+        http.setProtocolVersion(protocolVersion);
+      }
+    }
+
+    void stdio.send(message).catch((error: unknown) => {
+      logUnlessClosing("failed to forward message to stdio", error);
+    });
+  };
+
+  stdio.onerror = (error) => {
+    logUnlessClosing("stdio transport error", error);
+  };
+  http.onerror = (error) => {
+    logUnlessClosing("HTTP transport error", error);
+  };
+  stdio.onclose = () => {
+    closeCounterpart(http);
+  };
+  http.onclose = () => {
+    closeCounterpart(stdio);
+  };
+
+  let releaseInput: (() => void) | undefined;
+
+  try {
+    await http.start();
+    await stdio.start();
+    releaseInput = finished(input, { writable: false }, () => {
+      void stdio.close().catch((error: unknown) => {
+        logError("failed to close transport", error);
+      });
+    });
+    await closed;
+  } catch (error) {
+    closing = true;
+    await Promise.allSettled([http.close(), stdio.close()]);
+    resolveClosed();
+    throw error;
+  } finally {
+    releaseInput?.();
+  }
+}