@gtkx/mcp 1.3.0 → 1.5.0

package/src/connection-registry.ts

@@ -1,10 +1,10 @@
 import type { Socket } from "node:net";
-import type { Message } from "./protocol/schemas.js";
+import { methodNotFoundError } from "./protocol/errors.js";
 import {
     type AppConnections,
     connectionDisconnectionEvent,
     connectionErrorEvent,
-    connectionRequestEvent,
+    type ConnectionRequestHandler,
     ProtocolConnection,
 } from "./transport.js";
 
@@ -12,6 +12,8 @@ class ConnectionRegistry extends EventTarget implements AppConnections {
     private connections: Map<string, ProtocolConnection> = new Map();
     private sockets: Map<string, Socket> = new Map();
 
+    onRequest: ConnectionRequestHandler = (_connection, request) => Promise.reject(methodNotFoundError(request.method));
+
     register(socket: Socket): ProtocolConnection {
         const connection = ProtocolConnection.fromSocket(socket, {
             onClose: () => {
@@ -24,30 +26,12 @@ class ConnectionRegistry extends EventTarget implements AppConnections {
 
         this.connections.set(connection.id, connection);
         this.sockets.set(connection.id, socket);
-        connection.on("request", (request) => this.dispatchEvent(connectionRequestEvent(connection, request)));
-
-        connection.on("invalid", ({ id: badId, error }) => {
-            connection.write({ id: badId, error: error.toErrorObject() });
-        });
+        connection.fallbackRequestHandler = (request) => this.onRequest(connection, request);
 
         return connection;
     }
 
-    send(connectionId: string, message: Message): void {
-        const connection = this.connections.get(connectionId);
-
-        if (!connection) {
-            return;
-        }
-
-        connection.write(message);
-    }
-
-    dispose(reason: string): void {
-        for (const connection of this.connections.values()) {
-            connection.rejectPending(new Error(reason));
-        }
-
+    dispose(): void {
         for (const socket of this.sockets.values()) {
             socket.destroy();
         }

package/src/internal.ts

@@ -1,8 +1,7 @@
 export {
-    ErrorCode,
     invalidRequestError,
+    isConnectionClosedError,
     methodNotFoundError,
-    ProtocolError,
     propertyNotFoundError,
     widgetNotFoundError,
 } from "./protocol/errors.js";
@@ -11,7 +10,6 @@ export {
     DEFAULT_SUBTREE_DEPTH,
     MAX_SUBTREE_WIDGETS,
     type ParamsSchema,
-    type Request,
     type SerializedProperty,
     type SerializedWidget,
     type ServerInitiatedMethod,
@@ -19,3 +17,4 @@ export {
     ServerRequestParamsSchemas,
 } from "./protocol/schemas.js";
 export { ProtocolConnection } from "./transport.js";
+export type { JSONRPCRequest, Result } from "@modelcontextprotocol/sdk/types.js";

package/src/protocol/errors.ts

@@ -1,3 +1,5 @@
+import { McpError, ErrorCode as SdkErrorCode } from "@modelcontextprotocol/sdk/types.js";
+
 type ErrorCode = (typeof ErrorCode)[keyof typeof ErrorCode];
 
 const ErrorCode = {
@@ -12,6 +14,9 @@ const ErrorCode = {
     PROPERTY_NOT_FOUND: 1008,
 } as const;
 
+const CONNECTION_CLOSED_CODE: number = SdkErrorCode.ConnectionClosed;
+const REQUEST_TIMEOUT_CODE: number = SdkErrorCode.RequestTimeout;
+
 function isErrorCode(code: number): code is ErrorCode {
     return (Object.values(ErrorCode) as number[]).includes(code);
 }
@@ -64,6 +69,17 @@ function methodNotFoundError(method: string): ProtocolError {
     return new ProtocolError(ErrorCode.METHOD_NOT_FOUND, `Method '${method}' not found`, { method });
 }
 
+function isConnectionClosedError(value: unknown): boolean {
+    return value instanceof McpError && value.code === CONNECTION_CLOSED_CODE;
+}
+
+function protocolErrorFrom(error: McpError): ProtocolError {
+    const prefix = `MCP error ${String(error.code)}: `;
+    const message = error.message.startsWith(prefix) ? error.message.slice(prefix.length) : error.message;
+
+    return new ProtocolError(isErrorCode(error.code) ? error.code : ErrorCode.INTERNAL_ERROR, message, error.data);
+}
+
 class ProtocolError extends Error {
     code: ErrorCode;
     data?: unknown;
@@ -74,19 +90,12 @@ class ProtocolError extends Error {
         this.data = data;
         this.name = "ProtocolError";
     }
-
-    toErrorObject(): { code: number; message: string; data?: unknown } {
-        return {
-            code: this.code,
-            message: this.message,
-            ...(this.data !== undefined && { data: this.data }),
-        };
-    }
 }
 
 export {
-    ErrorCode,
-    isErrorCode,
+    CONNECTION_CLOSED_CODE,
+    isConnectionClosedError,
+    REQUEST_TIMEOUT_CODE,
     noAppConnectedError,
     appNotFoundError,
     connectionWriteFailedError,
@@ -96,4 +105,5 @@ export {
     invalidRequestError,
     methodNotFoundError,
     ProtocolError,
+    protocolErrorFrom,
 };

package/src/protocol/schemas.ts

@@ -2,9 +2,6 @@ import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { z } from "zod";
 
-type Request = z.infer<typeof RequestSchema>;
-type Response = z.infer<typeof ResponseSchema>;
-
 type SerializedWidget = {
     id: string;
     type: string;
@@ -37,39 +34,6 @@ type ServerRequestParams<Method extends keyof typeof ServerRequestParamsSchemas>
 
 type ParamsSchema<Output> = z.ZodType<Output>;
 type ServerInitiatedMethod = keyof typeof ServerRequestParamsSchemas;
-type Message = Request | Response;
-
-const RequestSchema: z.ZodObject<
-    {
-        id: z.ZodString;
-        method: z.ZodString;
-        params: z.ZodOptional<z.ZodUnknown>;
-    }
-> = z.object({
-    id: z.string(),
-    method: z.string(),
-    params: z.unknown().optional(),
-});
-
-const ErrorSchema: z.ZodObject<
-    { code: z.ZodNumber; message: z.ZodString; data: z.ZodOptional<z.ZodUnknown> }
-> = z.object({
-    code: z.number(),
-    message: z.string(),
-    data: z.unknown().optional(),
-});
-
-const ResponseSchema: z.ZodObject<
-    {
-        id: z.ZodString;
-        result: z.ZodOptional<z.ZodUnknown>;
-        error: z.ZodOptional<typeof ErrorSchema>;
-    }
-> = z.object({
-    id: z.string(),
-    result: z.unknown().optional(),
-    error: ErrorSchema.optional(),
-});
 
 const RegisterParamsSchema: z.ZodObject<
     {
@@ -175,8 +139,6 @@ function getRuntimeDir(): string {
 export {
     DEFAULT_SUBTREE_DEPTH,
     MAX_SUBTREE_WIDGETS,
-    RequestSchema,
-    ResponseSchema,
     RegisterParamsSchema,
     widgetIdParams,
     widgetPropsParams,
@@ -187,13 +149,10 @@ export {
     screenshotParams,
     ServerRequestParamsSchemas,
     DEFAULT_SOCKET_PATH,
-    type Request,
-    type Response,
     type SerializedWidget,
     type SerializedProperty,
     type AppInfo,
     type ServerRequestParams,
     type ParamsSchema,
     type ServerInitiatedMethod,
-    type Message,
 };

package/src/socket-server.ts

@@ -353,7 +353,7 @@ class SocketServer {
         }
 
         this.server = null;
-        this.registry.dispose("Server stopping");
+        this.registry.dispose();
         await closeServer(server);
         await this.release();
     }

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,
 };