diff --git a/mcp/client_example_test.go b/mcp/client_example_test.go index 3c3c383..bba3da4 100644 --- a/mcp/client_example_test.go +++ b/mcp/client_example_test.go @@ -45,9 +45,12 @@ func Example_roots() { if _, err := s.Connect(ctx, t1, nil); err != nil { log.Fatal(err) } - if _, err := c.Connect(ctx, t2, nil); err != nil { + + clientSession, err := c.Connect(ctx, t2, nil) + if err != nil { log.Fatal(err) } + defer clientSession.Close() // ...and add a root. The server is notified about the change. c.AddRoots(&mcp.Root{URI: "file://b"}) diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index fa941bf..9cc105c 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -592,7 +592,10 @@ func errorCode(err error) int64 { // // The caller should cancel either the client connection or server connection // when the connections are no longer needed. -func basicConnection(t *testing.T, config func(*Server)) (*ClientSession, *ServerSession) { +// +// The returned func cleans up by closing the client and waiting for the server +// to shut down. +func basicConnection(t *testing.T, config func(*Server)) (*ClientSession, *ServerSession, func()) { return basicClientServerConnection(t, nil, nil, config) } @@ -604,7 +607,10 @@ func basicConnection(t *testing.T, config func(*Server)) (*ClientSession, *Serve // // The caller should cancel either the client connection or server connection // when the connections are no longer needed. -func basicClientServerConnection(t *testing.T, client *Client, server *Server, config func(*Server)) (*ClientSession, *ServerSession) { +// +// The returned func cleans up by closing the client and waiting for the server +// to shut down. +func basicClientServerConnection(t *testing.T, client *Client, server *Server, config func(*Server)) (*ClientSession, *ServerSession, func()) { t.Helper() ctx := context.Background() @@ -628,14 +634,17 @@ func basicClientServerConnection(t *testing.T, client *Client, server *Server, c if err != nil { t.Fatal(err) } - return cs, ss + return cs, ss, func() { + cs.Close() + ss.Wait() + } } func TestServerClosing(t *testing.T) { - cs, ss := basicConnection(t, func(s *Server) { + cs, ss, cleanup := basicConnection(t, func(s *Server) { AddTool(s, greetTool(), sayHi) }) - defer cs.Close() + defer cleanup() ctx := context.Background() var wg sync.WaitGroup @@ -715,10 +724,10 @@ func TestCancellation(t *testing.T) { } return nil, nil, nil } - cs, _ := basicConnection(t, func(s *Server) { + cs, _, cleanup := basicConnection(t, func(s *Server) { AddTool(s, &Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, slowTool) }) - defer cs.Close() + defer cleanup() ctx, cancel := context.WithCancel(context.Background()) go cs.CallTool(ctx, &CallToolParams{Name: "slow"}) @@ -741,13 +750,10 @@ func TestMiddleware(t *testing.T) { t.Fatal(err) } // Wait for the server to exit after the client closes its connection. - var clientWG sync.WaitGroup - clientWG.Add(1) - go func() { + defer func() { if err := ss.Wait(); err != nil { t.Errorf("server failed: %v", err) } - clientWG.Done() }() var sbuf, cbuf bytes.Buffer @@ -767,6 +773,8 @@ func TestMiddleware(t *testing.T) { if err != nil { t.Fatal(err) } + defer cs.Close() + if _, err := cs.ListTools(ctx, nil); err != nil { t.Fatal(err) } @@ -1511,7 +1519,7 @@ func TestKeepAliveFailure(t *testing.T) { func TestAddTool_DuplicateNoPanicAndNoDuplicate(t *testing.T) { // Adding the same tool pointer twice should not panic and should not // produce duplicates in the server's tool list. - cs, _ := basicConnection(t, func(s *Server) { + cs, _, cleanup := basicConnection(t, func(s *Server) { // Use two distinct Tool instances with the same name but different // descriptions to ensure the second replaces the first // This case was written specifically to reproduce a bug where duplicate tools where causing jsonschema errors @@ -1520,7 +1528,7 @@ func TestAddTool_DuplicateNoPanicAndNoDuplicate(t *testing.T) { s.AddTool(t1, nopHandler) s.AddTool(t2, nopHandler) }) - defer cs.Close() + defer cleanup() ctx := context.Background() res, err := cs.ListTools(ctx, nil) @@ -1568,7 +1576,7 @@ func TestSynchronousNotifications(t *testing.T) { }, } server := NewServer(testImpl, serverOpts) - cs, ss := basicClientServerConnection(t, client, server, func(s *Server) { + cs, ss, cleanup := basicClientServerConnection(t, client, server, func(s *Server) { AddTool(s, &Tool{Name: "tool"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { if !rootsChanged.Load() { return nil, nil, fmt.Errorf("didn't get root change notification") @@ -1576,6 +1584,7 @@ func TestSynchronousNotifications(t *testing.T) { return new(CallToolResult), nil, nil }) }) + defer cleanup() t.Run("from client", func(t *testing.T) { client.AddRoots(&Root{Name: "myroot", URI: "file://foo"}) @@ -1617,7 +1626,7 @@ func TestNoDistributedDeadlock(t *testing.T) { }, } client := NewClient(testImpl, clientOpts) - cs, _ := basicClientServerConnection(t, client, nil, func(s *Server) { + cs, _, cleanup := basicClientServerConnection(t, client, nil, func(s *Server) { AddTool(s, &Tool{Name: "tool1"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { req.Session.CreateMessage(ctx, new(CreateMessageParams)) return new(CallToolResult), nil, nil @@ -1627,7 +1636,7 @@ func TestNoDistributedDeadlock(t *testing.T) { return new(CallToolResult), nil, nil }) }) - defer cs.Close() + defer cleanup() ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() @@ -1651,7 +1660,7 @@ func TestPointerArgEquivalence(t *testing.T) { type output struct { Out string } - cs, _ := basicConnection(t, func(s *Server) { + cs, _, cleanup := basicConnection(t, func(s *Server) { // Add two equivalent tools, one of which operates in the 'pointer' realm, // the other of which does not. // @@ -1686,7 +1695,7 @@ func TestPointerArgEquivalence(t *testing.T) { } }) }) - defer cs.Close() + defer cleanup() ctx := context.Background() tools, err := cs.ListTools(ctx, nil) @@ -1758,7 +1767,9 @@ func TestComplete(t *testing.T) { }, } server := NewServer(testImpl, serverOpts) - cs, _ := basicClientServerConnection(t, nil, server, func(s *Server) {}) + cs, _, cleanup := basicClientServerConnection(t, nil, server, func(s *Server) {}) + defer cleanup() + result, err := cs.Complete(context.Background(), &CompleteParams{ Argument: CompleteParamsArgument{ Name: "language", diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index d6b4d7e..16f19e2 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -55,6 +55,7 @@ func Example_prompts() { if err != nil { log.Fatal(err) } + defer cs.Close() // List the prompts. for p, err := range cs.Prompts(ctx, nil) { diff --git a/mcp/transport_example_test.go b/mcp/transport_example_test.go index 064ab0f..ab54a42 100644 --- a/mcp/transport_example_test.go +++ b/mcp/transport_example_test.go @@ -24,16 +24,21 @@ func ExampleLoggingTransport() { ctx := context.Background() t1, t2 := mcp.NewInMemoryTransports() server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil) - if _, err := server.Connect(ctx, t1, nil); err != nil { + serverSession, err := server.Connect(ctx, t1, nil) + if err != nil { log.Fatal(err) } + defer serverSession.Wait() client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil) var b bytes.Buffer logTransport := &mcp.LoggingTransport{Transport: t2, Writer: &b} - if _, err := client.Connect(ctx, logTransport, nil); err != nil { + clientSession, err := client.Connect(ctx, logTransport, nil) + if err != nil { log.Fatal(err) } + defer clientSession.Close() + // Sort for stability: reads are concurrent to writes. for _, line := range slices.Sorted(strings.SplitSeq(b.String(), "\n")) { fmt.Println(line) diff --git a/mcp/transport_test.go b/mcp/transport_test.go index d40ce10..10804a8 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -25,6 +25,7 @@ func TestBatchFraming(t *testing.T) { r, w := io.Pipe() tport := newIOConn(rwc{r, w}) tport.outgoingBatch = make([]jsonrpc.Message, 0, 2) + defer tport.Close() // Read the two messages into a channel, for easy testing later. read := make(chan jsonrpc.Message)