Skip to content

Commit 9d6e307

Browse files
authored
mcp: add a explicit test for concurrent streamable sessions (#584)
We seem not to have a test that explicitly exercises concurrent streamable sessions. Add one.
1 parent 31e97ad commit 9d6e307

File tree

1 file changed

+51
-0
lines changed

1 file changed

+51
-0
lines changed

mcp/streamable_test.go

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,57 @@ func TestStreamableTransports(t *testing.T) {
192192
}
193193
}
194194

195+
func TestStreamableConcurrentHandling(t *testing.T) {
196+
// This test checks that the streamable server and client transports can
197+
// communicate.
198+
type count struct {
199+
Count int
200+
}
201+
202+
var mu sync.Mutex
203+
counts := make(map[string]int)
204+
205+
server := NewServer(testImpl, nil)
206+
AddTool(server, &Tool{Name: "inc"}, func(ctx context.Context, req *CallToolRequest, _ any) (*CallToolResult, count, error) {
207+
id := req.Session.ID()
208+
mu.Lock()
209+
defer mu.Unlock()
210+
c := counts[id]
211+
counts[id] = c + 1
212+
return nil, count{c}, nil
213+
})
214+
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
215+
httpServer := httptest.NewServer(mustNotPanic(t, handler))
216+
defer httpServer.Close()
217+
218+
ctx := context.Background()
219+
client := NewClient(testImpl, nil)
220+
var wg sync.WaitGroup
221+
for range 100 {
222+
wg.Add(1)
223+
go func() {
224+
defer wg.Done()
225+
clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil)
226+
if err != nil {
227+
t.Errorf("Connect failed: %v", err)
228+
return
229+
}
230+
defer clientSession.Close()
231+
for i := range 10 {
232+
res, err := clientSession.CallTool(ctx, &CallToolParams{Name: "inc"})
233+
if err != nil {
234+
t.Errorf("CallTool failed: %v", err)
235+
return
236+
}
237+
if got := int(res.StructuredContent.(map[string]any)["Count"].(float64)); got != i {
238+
t.Errorf("got count %d, want %d", got, i)
239+
}
240+
}
241+
}()
242+
}
243+
wg.Wait()
244+
}
245+
195246
func TestStreamableServerShutdown(t *testing.T) {
196247
ctx := context.Background()
197248

0 commit comments

Comments
 (0)