1
+ import { describe , expect , test , beforeEach , jest } from "@jest/globals" ;
2
+ import { Protocol } from "./protocol.js" ;
3
+ import { Transport } from "./transport.js" ;
4
+ import { Request , Notification , Result , JSONRPCMessage } from "../types.js" ;
5
+ import { z } from "zod" ;
6
+
7
+ // Mock Transport class
8
+ class MockTransport implements Transport {
9
+ id : string ;
10
+ onclose ?: ( ) => void ;
11
+ onerror ?: ( error : Error ) => void ;
12
+ onmessage ?: ( message : unknown ) => void ;
13
+ sentMessages : JSONRPCMessage [ ] = [ ] ;
14
+
15
+ constructor ( id : string ) {
16
+ this . id = id ;
17
+ }
18
+
19
+ async start ( ) : Promise < void > { }
20
+
21
+ async close ( ) : Promise < void > {
22
+ this . onclose ?.( ) ;
23
+ }
24
+
25
+ async send ( message : JSONRPCMessage ) : Promise < void > {
26
+ this . sentMessages . push ( message ) ;
27
+ }
28
+ }
29
+
30
+ describe ( "Protocol transport handling bug" , ( ) => {
31
+ let protocol : Protocol < Request , Notification , Result > ;
32
+ let transportA : MockTransport ;
33
+ let transportB : MockTransport ;
34
+
35
+ beforeEach ( ( ) => {
36
+ protocol = new ( class extends Protocol < Request , Notification , Result > {
37
+ protected assertCapabilityForMethod ( ) : void { }
38
+ protected assertNotificationCapability ( ) : void { }
39
+ protected assertRequestHandlerCapability ( ) : void { }
40
+ } ) ( ) ;
41
+
42
+ transportA = new MockTransport ( "A" ) ;
43
+ transportB = new MockTransport ( "B" ) ;
44
+ } ) ;
45
+
46
+ test ( "should send response to the correct transport when multiple clients are connected" , async ( ) => {
47
+ // Set up a request handler that simulates processing time
48
+ let resolveHandler : any ;
49
+ const handlerPromise = new Promise < Result > ( ( resolve ) => {
50
+ resolveHandler = resolve ;
51
+ } ) ;
52
+
53
+ const TestRequestSchema = z . object ( {
54
+ method : z . literal ( "test/method" ) ,
55
+ params : z . object ( {
56
+ from : z . string ( )
57
+ } ) . optional ( )
58
+ } ) ;
59
+
60
+ protocol . setRequestHandler (
61
+ TestRequestSchema ,
62
+ async ( request ) => {
63
+ console . log ( `Processing request from ${ request . params ?. from } ` ) ;
64
+ return handlerPromise ;
65
+ }
66
+ ) ;
67
+
68
+ // Client A connects and sends a request
69
+ await protocol . connect ( transportA ) ;
70
+
71
+ const requestFromA = {
72
+ jsonrpc : "2.0" as const ,
73
+ method : "test/method" ,
74
+ params : { from : "clientA" } ,
75
+ id : 1
76
+ } ;
77
+
78
+ // Simulate client A sending a request
79
+ transportA . onmessage ?.( requestFromA ) ;
80
+
81
+ // While A's request is being processed, client B connects
82
+ // This overwrites the transport reference in the protocol
83
+ await protocol . connect ( transportB ) ;
84
+
85
+ const requestFromB = {
86
+ jsonrpc : "2.0" as const ,
87
+ method : "test/method" ,
88
+ params : { from : "clientB" } ,
89
+ id : 2
90
+ } ;
91
+
92
+ // Client B sends its own request
93
+ transportB . onmessage ?.( requestFromB ) ;
94
+
95
+ // Now complete A's request
96
+ resolveHandler ! ( { data : "responseForA" } as Result ) ;
97
+
98
+ // Wait for async operations to complete
99
+ await new Promise ( resolve => setTimeout ( resolve , 10 ) ) ;
100
+
101
+ // Check where the responses went
102
+ console . log ( "Transport A received:" , transportA . sentMessages ) ;
103
+ console . log ( "Transport B received:" , transportB . sentMessages ) ;
104
+
105
+ // BUG: The response for client A's request will be sent to transport B
106
+ // because the protocol's transport reference was overwritten
107
+
108
+ // What happens (bug):
109
+ // - Transport A should have received the response for request ID 1, but it's empty
110
+ expect ( transportA . sentMessages . length ) . toBe ( 0 ) ;
111
+
112
+ // - Transport B incorrectly receives BOTH responses
113
+ expect ( transportB . sentMessages . length ) . toBe ( 2 ) ;
114
+ expect ( transportB . sentMessages [ 0 ] ) . toMatchObject ( {
115
+ jsonrpc : "2.0" ,
116
+ id : 1 , // This is A's request ID!
117
+ result : { data : "responseForA" }
118
+ } ) ;
119
+
120
+ // What SHOULD happen (after fix):
121
+ // - Transport A should receive response for request ID 1
122
+ // - Transport B should receive response for request ID 2
123
+ } ) ;
124
+
125
+ test ( "demonstrates the timing issue with multiple rapid connections" , async ( ) => {
126
+ const delays : number [ ] = [ ] ;
127
+ const results : { transport : string ; response : any } [ ] = [ ] ;
128
+
129
+ const DelayedRequestSchema = z . object ( {
130
+ method : z . literal ( "test/delayed" ) ,
131
+ params : z . object ( {
132
+ delay : z . number ( ) ,
133
+ client : z . string ( )
134
+ } ) . optional ( )
135
+ } ) ;
136
+
137
+ // Set up handler with variable delay
138
+ protocol . setRequestHandler (
139
+ DelayedRequestSchema ,
140
+ async ( request , extra ) => {
141
+ const delay = request . params ?. delay || 0 ;
142
+ delays . push ( delay ) ;
143
+
144
+ await new Promise ( resolve => setTimeout ( resolve , delay ) ) ;
145
+
146
+ return {
147
+ processedBy : `handler-${ extra . requestId } ` ,
148
+ delay : delay
149
+ } as Result ;
150
+ }
151
+ ) ;
152
+
153
+ // Rapid succession of connections and requests
154
+ await protocol . connect ( transportA ) ;
155
+ transportA . onmessage ?.( {
156
+ jsonrpc : "2.0" as const ,
157
+ method : "test/delayed" ,
158
+ params : { delay : 50 , client : "A" } ,
159
+ id : 1
160
+ } ) ;
161
+
162
+ // Connect B while A is processing
163
+ setTimeout ( async ( ) => {
164
+ await protocol . connect ( transportB ) ;
165
+ transportB . onmessage ?.( {
166
+ jsonrpc : "2.0" as const ,
167
+ method : "test/delayed" ,
168
+ params : { delay : 10 , client : "B" } ,
169
+ id : 2
170
+ } ) ;
171
+ } , 10 ) ;
172
+
173
+ // Wait for all processing
174
+ await new Promise ( resolve => setTimeout ( resolve , 100 ) ) ;
175
+
176
+ // Collect results
177
+ if ( transportA . sentMessages . length > 0 ) {
178
+ results . push ( { transport : "A" , response : transportA . sentMessages } ) ;
179
+ }
180
+ if ( transportB . sentMessages . length > 0 ) {
181
+ results . push ( { transport : "B" , response : transportB . sentMessages } ) ;
182
+ }
183
+
184
+ console . log ( "Timing test results:" , results ) ;
185
+
186
+ // BUG: All responses go to transport B
187
+ expect ( transportA . sentMessages . length ) . toBe ( 0 ) ;
188
+ expect ( transportB . sentMessages . length ) . toBe ( 2 ) ; // Gets both responses
189
+ } ) ;
190
+ } ) ;
0 commit comments