1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
|
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import WebSocket from 'ws';
import {
Transport,
TransportSendOptions,
} from '@modelcontextprotocol/sdk/shared/transport.js';
import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
import { AuthInfo } from '@modelcontextprotocol/sdk/server/auth/types.js';
export class WebSocketClientTransport implements Transport {
private socket: WebSocket | null = null;
onclose?: () => void;
onerror?: (error: Error) => void;
onmessage?: (
message: JSONRPCMessage,
extra?: { authInfo?: AuthInfo },
) => void;
constructor(private readonly url: URL) {}
async start(): Promise<void> {
return new Promise((resolve, reject) => {
const handshakeTimeoutDuration = 10000;
let connectionTimeout: NodeJS.Timeout | null = null;
try {
this.socket = new WebSocket(this.url.toString(), {
handshakeTimeout: handshakeTimeoutDuration,
});
connectionTimeout = setTimeout(() => {
this.socket?.close();
reject(
new Error(
`WebSocket connection timed out after ${handshakeTimeoutDuration}ms`,
),
);
}, handshakeTimeoutDuration);
this.socket.on('open', () => {
clearTimeout(connectionTimeout!);
resolve();
});
this.socket.on('message', (data) => {
try {
const parsedMessage: JSONRPCMessage = JSON.parse(data.toString());
this.onmessage?.(parsedMessage, { authInfo: undefined }); // Auth unsupported currently
} catch (error: unknown) {
this.onerror?.(
error instanceof Error ? error : new Error(String(error)),
);
}
});
this.socket.on('error', (error) => {
clearTimeout(connectionTimeout!);
this.onerror?.(error);
reject(error);
});
this.socket.on('close', () => {
clearTimeout(connectionTimeout!);
this.onclose?.();
this.socket = null;
});
} catch (error: unknown) {
clearTimeout(connectionTimeout!);
reject(error instanceof Error ? error : new Error(String(error)));
}
});
}
async close(): Promise<void> {
if (this.socket) {
this.socket.close();
this.socket = null;
}
}
async send(
message: JSONRPCMessage,
_options?: TransportSendOptions,
): Promise<void> {
if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
throw new Error(
'WebSocket is not connected or not open. Cannot send message.',
);
}
this.socket.send(JSON.stringify(message));
}
}
|
