@gtkx/mcp 1.4.0 → 1.6.0

package/src/transport.ts

@@ -1,30 +1,36 @@
+/* eslint-disable @typescript-eslint/no-empty-function -- the app link negotiates no MCP capabilities */
+import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
 import type { Socket } from "node:net";
-import type { Duplex } from "node:stream";
-import { ErrorCode, invalidRequestError, isErrorCode, ProtocolError, requestTimeoutError } from "./protocol/errors.js";
-import { type Message, type Request, RequestSchema, type Response, ResponseSchema } from "./protocol/schemas.js";
-
-type ProtocolConnectionEvents = {
-    request: Request;
-    invalid: { id: string; error: ProtocolError };
-};
-
-type PendingRequest = {
-    resolve: (result: unknown) => void;
-    reject: (error: Error) => void;
-    timeout: NodeJS.Timeout;
-};
+import { normalizeError } from "@gtkx/utils";
+import { Protocol } from "@modelcontextprotocol/sdk/shared/protocol.js";
+import { ReadBuffer, serializeMessage } from "@modelcontextprotocol/sdk/shared/stdio.js";
+import {
+    ErrorCode,
+    type JSONRPCMessage,
+    type JSONRPCRequest,
+    McpError,
+    type Notification,
+    type Request,
+    type Result,
+    ResultSchema,
+} from "@modelcontextprotocol/sdk/types.js";
 
 type ConnectionEvent = CustomEvent<ProtocolConnection>;
-type ConnectionRequestEvent = CustomEvent<{ connection: ProtocolConnection; request: Request }>;
 type ConnectionErrorEvent = CustomEvent<Error>;
+type ConnectionRequestHandler = (connection: ProtocolConnection, request: JSONRPCRequest) => Promise<Result>;
+type RequestParams = Request["params"];
+
+type ConnectionOptions = {
+    onClose: () => void;
+    onError: (error: Error) => void;
+};
 
 type AppConnections = {
-    send(connectionId: string, message: Message): void;
+    onRequest: ConnectionRequestHandler;
 } & EventTarget;
 
-function connectionRequestEvent(connection: ProtocolConnection, request: Request): ConnectionRequestEvent {
-    return new CustomEvent("request", { detail: { connection, request } });
-}
+const MAX_MESSAGE_BYTES = 64 * 1024 * 1024;
+const MAX_PENDING_WRITE_BYTES = 8 * 1024 * 1024;
 
 function connectionDisconnectionEvent(connection: ProtocolConnection): ConnectionEvent {
     return new CustomEvent("disconnection", { detail: connection });
@@ -34,208 +40,137 @@ function connectionErrorEvent(error: Error): ConnectionErrorEvent {
     return new CustomEvent("error", { detail: error });
 }
 
-class ConnectionClosedError extends Error {
-    constructor() {
-        super("Connection stream is not writable");
-        this.name = "ConnectionClosedError";
-    }
-}
-
-class ProtocolConnection extends EventTarget {
-    static fromSocket(
-        socket: Socket,
-        options: {
-            onClose?: () => void;
-            onError?: (error: Error) => void;
-        } = {},
-    ): ProtocolConnection {
-        const connection = new ProtocolConnection(socket);
-
-        socket.on("data", (data: Buffer) => {
-            connection.feed(data);
-        });
-
-        socket.on("close", () => {
-            connection.rejectPending(new Error("Connection closed"));
-            options.onClose?.();
-        });
-
-        if (options.onError) {
-            socket.on("error", options.onError);
-        }
-
-        return connection;
-    }
-
-    private buffer = "";
-    private pending: Map<string, PendingRequest> = new Map();
-    private writer: Duplex;
-
-    id: string = crypto.randomUUID();
-
-    constructor(writer: Duplex) {
-        super();
-        this.writer = writer;
-    }
-
-    private notify<K extends keyof ProtocolConnectionEvents>(type: K, detail: ProtocolConnectionEvents[K]): void {
-        this.dispatchEvent(new CustomEvent(type, { detail }));
-    }
+class SocketTransport implements Transport {
+    private readonly readBuffer: ReadBuffer = new ReadBuffer({ maxBufferSize: MAX_MESSAGE_BYTES });
+    private readonly socket: Socket;
 
-    private rejectWhenClosed(id: string, timeoutHandle: NodeJS.Timeout, reject: (error: Error) => void): void {
-        if (this.writer.writable) {
-            return;
-        }
+    onclose?: () => void;
+    onerror?: (error: Error) => void;
+    onmessage?: (message: JSONRPCMessage) => void;
 
-        clearTimeout(timeoutHandle);
-        this.pending.delete(id);
-        reject(new ConnectionClosedError());
+    constructor(socket: Socket) {
+        this.socket = socket;
     }
 
-    private dispatchParsed(parsed: unknown): boolean {
-        const message = parsed as Record<string, unknown>;
-
-        if (typeof message.method === "string") {
-            const requestResult = RequestSchema.safeParse(parsed);
+    private read(): void {
+        for (;;) {
+            let message: JSONRPCMessage | null;
 
-            if (!requestResult.success) {
-                return false;
+            try {
+                message = this.readBuffer.readMessage();
+            } catch (error) {
+                this.onerror?.(normalizeError(error));
+                continue;
             }
 
-            this.notify("request", requestResult.data);
-
-            return true;
-        }
-
-        const responseResult = ResponseSchema.safeParse(parsed);
+            if (message === null) {
+                return;
+            }
 
-        if (!responseResult.success) {
-            return false;
+            this.onmessage?.(message);
         }
+    }
 
-        this.handleResponse(responseResult.data);
+    private drop(reason: string): Promise<void> {
+        const error = new McpError(ErrorCode.ConnectionClosed, reason);
+        this.onerror?.(error);
+        this.socket.destroy();
 
-        return true;
+        return Promise.reject(error);
     }
 
-    private processLine(line: string): void {
-        let parsed: unknown;
-
+    private receive(chunk: Buffer): void {
         try {
-            parsed = JSON.parse(line);
-        } catch {
-            this.notify("invalid", { id: "unknown", error: invalidRequestError("Invalid JSON") });
-
-            return;
-        }
+            this.readBuffer.append(chunk);
+        } catch (error) {
+            this.onerror?.(normalizeError(error));
+            this.socket.destroy();
 
-        if (this.dispatchParsed(parsed)) {
             return;
         }
 
-        const message = parsed as Record<string, unknown>;
-        const id = typeof message.id === "string" ? message.id : "unknown";
-        this.notify("invalid", { id, error: invalidRequestError("Invalid message format") });
+        this.read();
     }
 
-    private handleResponse(response: Response): void {
-        const entry = this.pending.get(response.id);
+    start(): Promise<void> {
+        this.socket.on("data", (chunk: Buffer) => {
+            this.receive(chunk);
+        });
 
-        if (!entry) {
-            return;
-        }
+        this.socket.on("close", () => {
+            this.onclose?.();
+        });
 
-        clearTimeout(entry.timeout);
-        this.pending.delete(response.id);
+        this.socket.on("error", (error) => {
+            this.onerror?.(error);
+        });
 
-        if (response.error) {
-            const err = response.error;
+        return Promise.resolve();
+    }
 
-            entry.reject(
-                new ProtocolError(isErrorCode(err.code) ? err.code : ErrorCode.INTERNAL_ERROR, err.message, err.data),
-            );
-        } else {
-            entry.resolve(response.result);
+    send(message: JSONRPCMessage): Promise<void> {
+        if (!this.socket.writable) {
+            return Promise.reject(new McpError(ErrorCode.ConnectionClosed, "Connection stream is not writable"));
         }
-    }
 
-    on<K extends keyof ProtocolConnectionEvents>(
-        type: K,
-        listener: (detail: ProtocolConnectionEvents[K]) => void,
-    ): void {
-        this.addEventListener(type, (event) => {
-            listener((event as CustomEvent<ProtocolConnectionEvents[K]>).detail);
-        });
-    }
+        if (this.socket.writableLength > MAX_PENDING_WRITE_BYTES) {
+            return this.drop("Connection stream is not draining");
+        }
 
-    feed(data: Buffer | string): void {
-        this.buffer += typeof data === "string" ? data : data.toString();
-        let newlineIndex = this.buffer.indexOf("\n");
+        this.socket.write(serializeMessage(message));
 
-        while (newlineIndex !== -1) {
-            const line = this.buffer.slice(0, newlineIndex);
-            this.buffer = this.buffer.slice(newlineIndex + 1);
+        return Promise.resolve();
+    }
 
-            if (line.trim()) {
-                this.processLine(line);
-            }
+    close(): Promise<void> {
+        this.socket.destroy();
 
-            newlineIndex = this.buffer.indexOf("\n");
-        }
+        return Promise.resolve();
     }
+}
 
-    write(message: Message): void {
-        if (!this.writer.writable) {
-            return;
-        }
+class ProtocolConnection extends Protocol<Request, Notification, Result> {
+    static fromSocket(socket: Socket, options: ConnectionOptions): ProtocolConnection {
+        const connection = Object.assign(new ProtocolConnection(), {
+            onclose: options.onClose,
+            onerror: options.onError,
+        });
+
+        void connection.connect(new SocketTransport(socket));
 
-        this.writer.write(`${JSON.stringify(message)}\n`);
+        return connection;
     }
 
-    send<T = unknown>(method: string, params: unknown, timeout: number): Promise<T> {
-        return new Promise<T>((resolve, reject) => {
-            if (!this.writer.writable) {
-                reject(new ConnectionClosedError());
+    id: string = crypto.randomUUID();
 
-                return;
-            }
+    protected assertCapabilityForMethod(): void {}
 
-            const id = crypto.randomUUID();
+    protected assertNotificationCapability(): void {}
 
-            const timeoutHandle = setTimeout(() => {
-                this.pending.delete(id);
-                reject(requestTimeoutError(timeout));
-            }, timeout);
+    protected assertRequestHandlerCapability(): void {}
 
-            this.pending.set(id, {
-                resolve: resolve as (result: unknown) => void,
-                reject,
-                timeout: timeoutHandle,
-            });
+    protected assertTaskCapability(): void {}
 
-            this.write({ id, method, params });
-            this.rejectWhenClosed(id, timeoutHandle, reject);
-        });
-    }
+    protected assertTaskHandlerCapability(): void {}
 
-    rejectPending(error: Error): void {
-        for (const entry of this.pending.values()) {
-            clearTimeout(entry.timeout);
-            entry.reject(error);
-        }
+    async send<T>(method: string, params?: RequestParams, timeout?: number): Promise<T> {
+        const result = await this.request(
+            { method, ...(params !== undefined && { params }) },
+            ResultSchema,
+            { ...(timeout !== undefined && { timeout }) },
+        );
 
-        this.pending.clear();
+        return result as T;
     }
 }
 
 export {
-    ConnectionClosedError,
     connectionDisconnectionEvent,
     connectionErrorEvent,
-    connectionRequestEvent,
     ProtocolConnection,
     type AppConnections,
     type ConnectionEvent,
     type ConnectionErrorEvent,
-    type ConnectionRequestEvent,
+    type ConnectionRequestHandler,
+    type RequestParams,
 };