From f405f80ce3eef1277402f1496c2eb3892fda5152 Mon Sep 17 00:00:00 2001 From: KhizarA77 Date: Wed, 10 Dec 2025 15:38:24 +0500 Subject: [PATCH 1/2] feat: add WebSocket server transport --- package-lock.json | 10 --- src/server/websocket.ts | 184 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 184 insertions(+), 10 deletions(-) create mode 100644 src/server/websocket.ts diff --git a/package-lock.json b/package-lock.json index d32963a73..30a051b54 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1319,7 +1319,6 @@ "integrity": "sha512-PC0PDZfJg8sP7cmKe6L3QIL8GZwU5aRvUFedqSIpw3B+QjRSUZeeITC2M5XKeMXEzL6wccN196iy3JLwKNvDVA==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@typescript-eslint/scope-manager": "8.48.1", "@typescript-eslint/types": "8.48.1", @@ -1751,7 +1750,6 @@ "integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==", "dev": true, "license": "MIT", - "peer": true, "bin": { "acorn": "bin/acorn" }, @@ -2294,7 +2292,6 @@ "integrity": "sha512-BhHmn2yNOFA9H9JmmIVKJmd288g9hrVRDkdoIgRCRuSySRUHH7r/DI6aAXW9T1WwUuY3DFgrcaqB+deURBLR5g==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@eslint-community/eslint-utils": "^4.8.0", "@eslint-community/regexpp": "^4.12.1", @@ -2613,7 +2610,6 @@ "resolved": "https://registry.npmjs.org/express/-/express-5.2.1.tgz", "integrity": "sha512-hIS4idWWai69NezIdRt2xFVofaF4j+6INOpJlVOLDO8zXGpUVEVzIYk12UUi2JzjEzWL3IOAxcTubgz9Po0yXw==", "license": "MIT", - "peer": true, "dependencies": { "accepts": "^2.0.0", "body-parser": "^2.2.1", @@ -4071,7 +4067,6 @@ "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "dev": true, "license": "MIT", - "peer": true, "engines": { "node": ">=12" }, @@ -4152,7 +4147,6 @@ "integrity": "sha512-4H8vUNGNjQ4V2EOoGw005+c+dGuPSnhpPBPHBtsZdGZBk/iJb4kguGlPWaZTZ3q5nMtFOEsY0nRDlh9PJyd6SQ==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "esbuild": "~0.25.0", "get-tsconfig": "^4.7.5" @@ -4198,7 +4192,6 @@ "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.6.3.tgz", "integrity": "sha512-hjcS1mhfuyi4WW8IWtjP7brDrG2cuDZukyrYrSauoXGNgx0S7zceP07adYkJycEr56BOUTNPzbInooiN3fn1qw==", "dev": true, - "peer": true, "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" @@ -4393,7 +4386,6 @@ "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "dev": true, "license": "MIT", - "peer": true, "engines": { "node": ">=12" }, @@ -4407,7 +4399,6 @@ "integrity": "sha512-BxAKBWmIbrDgrokdGZH1IgkIk/5mMHDreLDmCJ0qpyJaAteP8NvMhkwr/ZCQNqNH97bw/dANTE9PDzqwJghfMQ==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "esbuild": "^0.25.0", "fdir": "^6.5.0", @@ -4560,7 +4551,6 @@ "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz", "integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==", "license": "MIT", - "peer": true, "funding": { "url": "https://github.com/sponsors/colinhacks" } diff --git a/src/server/websocket.ts b/src/server/websocket.ts new file mode 100644 index 000000000..4e87ee502 --- /dev/null +++ b/src/server/websocket.ts @@ -0,0 +1,184 @@ +import type { Server as HttpServer } from 'node:http'; +import WebSocket, { WebSocketServer } from 'ws'; +import { Transport } from '../shared/transport.js'; +import { + JSONRPCMessage, + JSONRPCMessageSchema, + type MessageExtraInfo +} from '../types.js'; + +const SUBPROTOCOL = 'mcp'; + +export interface WebSocketServerTransportOptions { + /** + * Optional existing HTTP(S) server to attach the WebSocket server to. + * If provided, `port` and `host` are ignored. + */ + server?: HttpServer; + + /** + * Port to listen on if no HTTP server is provided. + * Defaults to 0 (OS picks a free port). + */ + port?: number; + + /** + * Host to bind to when creating a standalone WebSocket server. + */ + host?: string; + + /** + * Optional path for the WebSocket endpoint, e.g. "/mcp". + */ + path?: string; +} + +/** + * Server transport for WebSocket: this communicates with an MCP client + * over the WebSocket protocol. + * + * This is the WebSocket analogue of StdioServerTransport: it expects + * exactly one client per transport instance and delivers JSON-RPC + * messages via the Transport interface. + */ +export class WebSocketServerTransport implements Transport { + private _wss: WebSocketServer; + private _socket?: WebSocket; + private _started = false; + + // Transport interface fields / callbacks + sessionId?: string; + onclose?: () => void; + onerror?: (error: Error) => void; + onmessage?: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; + setProtocolVersion?: (version: string) => void; + + constructor(options: WebSocketServerTransportOptions = {}) { + const { server, port, host, path } = options; + + this._wss = new WebSocketServer({ + server, + port: server ? undefined : (port ?? 0), + host: server ? undefined : host, + path, + handleProtocols: (protocols /* , req */) => { + // Require the MCP subprotocol if offered + if (protocols.has(SUBPROTOCOL)) { + return SUBPROTOCOL; + } + // Reject if the client doesn't offer the MCP subprotocol + return false; + } + }); + } + + /** + * Starts listening for a single WebSocket client and sets up MCP message handling. + * + * Resolves once a client connects successfully. + */ + start(): Promise { + if (this._started) { + throw new Error( + 'WebSocketServerTransport already started! If using Server class, note that connect() calls start() automatically.' + ); + } + + this._started = true; + + return new Promise((resolve, reject) => { + const handleError = (err: Error) => { + this._wss.off('connection', handleConnection); + this.onerror?.(err); + reject(err); + }; + + const handleConnection = (socket: WebSocket) => { + // Only allow one client per transport instance + if (this._socket) { + socket.close(1013, 'Only one client is allowed per transport'); + return; + } + + // Enforce negotiated subprotocol + if (socket.protocol !== SUBPROTOCOL) { + socket.close(1002, 'MCP subprotocol (mcp) required'); + return; + } + + this._socket = socket; + + socket.on('message', data => { + try { + const parsed = JSON.parse(data.toString()); + const message = JSONRPCMessageSchema.parse(parsed); + this.onmessage?.(message); + } catch (error) { + this.onerror?.(error as Error); + } + }); + + socket.on('error', err => { + this.onerror?.(err as Error); + }); + + socket.on('close', () => { + this._socket = undefined; + this.onclose?.(); + }); + + this._wss.off('error', handleError); + this._wss.off('connection', handleConnection); + resolve(); + }; + + this._wss.on('connection', handleConnection); + this._wss.once('error', handleError); + }); + } + + /** + * Sends a JSON-RPC message to the connected WebSocket client. + */ + send(message: JSONRPCMessage): Promise { + return new Promise((resolve, reject) => { + if (!this._socket || this._socket.readyState !== WebSocket.OPEN) { + const error = new Error('Not connected'); + this.onerror?.(error); + reject(error); + return; + } + + const payload = JSON.stringify(message); + this._socket.send(payload, err => { + if (err) { + this.onerror?.(err); + reject(err); + } else { + resolve(); + } + }); + }); + } + + /** + * Closes the WebSocket connection and the underlying WebSocket server. + */ + async close(): Promise { + if (this._socket && this._socket.readyState === WebSocket.OPEN) { + this._socket.close(); + } + + await new Promise((resolve, reject) => { + this._wss.close(err => { + if (err) { + this.onerror?.(err); + reject(err); + } else { + this.onclose?.(); + resolve(); + } + }); + }); + } +} From 42a8d6db5f4dd0d73331e0411476137fcbffe9e6 Mon Sep 17 00:00:00 2001 From: KhizarA77 Date: Wed, 10 Dec 2025 15:59:32 +0500 Subject: [PATCH 2/2] chore: fix formatting for websocket server transport --- src/server/websocket.ts | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/server/websocket.ts b/src/server/websocket.ts index 4e87ee502..ef6839fdf 100644 --- a/src/server/websocket.ts +++ b/src/server/websocket.ts @@ -1,11 +1,7 @@ import type { Server as HttpServer } from 'node:http'; import WebSocket, { WebSocketServer } from 'ws'; import { Transport } from '../shared/transport.js'; -import { - JSONRPCMessage, - JSONRPCMessageSchema, - type MessageExtraInfo -} from '../types.js'; +import { JSONRPCMessage, JSONRPCMessageSchema, type MessageExtraInfo } from '../types.js'; const SUBPROTOCOL = 'mcp';