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.
This commit is contained in:
parent
31e97ad69e
commit
9d6e3074f9
1 changed files with 51 additions and 0 deletions
|
|
@ -192,6 +192,57 @@ func TestStreamableTransports(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestStreamableConcurrentHandling(t *testing.T) {
|
||||||
|
// This test checks that the streamable server and client transports can
|
||||||
|
// communicate.
|
||||||
|
type count struct {
|
||||||
|
Count int
|
||||||
|
}
|
||||||
|
|
||||||
|
var mu sync.Mutex
|
||||||
|
counts := make(map[string]int)
|
||||||
|
|
||||||
|
server := NewServer(testImpl, nil)
|
||||||
|
AddTool(server, &Tool{Name: "inc"}, func(ctx context.Context, req *CallToolRequest, _ any) (*CallToolResult, count, error) {
|
||||||
|
id := req.Session.ID()
|
||||||
|
mu.Lock()
|
||||||
|
defer mu.Unlock()
|
||||||
|
c := counts[id]
|
||||||
|
counts[id] = c + 1
|
||||||
|
return nil, count{c}, nil
|
||||||
|
})
|
||||||
|
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
|
||||||
|
httpServer := httptest.NewServer(mustNotPanic(t, handler))
|
||||||
|
defer httpServer.Close()
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
client := NewClient(testImpl, nil)
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
for range 100 {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("Connect failed: %v", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer clientSession.Close()
|
||||||
|
for i := range 10 {
|
||||||
|
res, err := clientSession.CallTool(ctx, &CallToolParams{Name: "inc"})
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("CallTool failed: %v", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if got := int(res.StructuredContent.(map[string]any)["Count"].(float64)); got != i {
|
||||||
|
t.Errorf("got count %d, want %d", got, i)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
}
|
||||||
|
|
||||||
func TestStreamableServerShutdown(t *testing.T) {
|
func TestStreamableServerShutdown(t *testing.T) {
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue