feat(mcptest): Change Server.Start to accept a context.Context. (#339)

* feat(mcptest): Change `Server.Start` to accept a `context.Context`.

Previously, `mcptest` used `context.TODO()` as the context with this server
code was executed (with a comment to upgrade to `testing.T.Context()`
eventually.

This effectively prevents tests from preparing a context to be used with server
code. This is imporant, because `WithHTTPContextFunc` and friends appear to be
the intended way to extract things like authentication information from the raw
HTTP request and pass it to the server code. Without this change, such code is
effectively untestable, at least with this package.

The `NewServer` convenience method has been left unchanged (i.e. it does not
accept a context) to avoid breaking users taking the happy path. Chaning the
signature of `Start()` is however a breaking change.

* test(mcptest): Update tests to use the new `Start()` signature.

---------

Co-authored-by: Navendu Pottekkat <[email protected]>

Author: octo

mcptest/mcptest.go

@@ -24,7 +24,6 @@ type Server struct {
 	prompts   []server.ServerPrompt
 	resources []server.ServerResource
 
-	ctx    context.Context
 	cancel func()
 
 	serverReader *io.PipeReader
@@ -45,7 +44,8 @@ func NewServer(t *testing.T, tools ...server.ServerTool) (*Server, error) {
 	server := NewUnstartedServer(t)
 	server.AddTools(tools...)
 
-	if err := server.Start(); err != nil {
+	// TODO: use t.Context() once go.mod is upgraded to go 1.24+
+	if err := server.Start(context.TODO()); err != nil {
 		return nil, err
 	}
 
@@ -59,12 +59,6 @@ func NewUnstartedServer(t *testing.T) *Server {
 		name: t.Name(),
 	}
 
-	// Use t.Context() once we switch to go >= 1.24
-	ctx := context.TODO()
-
-	// Set up context with cancellation, used to stop the server
-	server.ctx, server.cancel = context.WithCancel(ctx)
-
 	// Set up pipes for client-server communication
 	server.serverReader, server.clientWriter = io.Pipe()
 	server.clientReader, server.serverWriter = io.Pipe()
@@ -114,9 +108,11 @@ func (s *Server) AddResources(resources ...server.ServerResource) {
 
 // Start starts the server in a goroutine. Make sure to defer Close() after Start().
 // When using NewServer(), the returned server is already started.
-func (s *Server) Start() error {
+func (s *Server) Start(ctx context.Context) error {
 	s.wg.Add(1)
 
+	ctx, s.cancel = context.WithCancel(ctx)
+
 	// Start the MCP server in a goroutine
 	go func() {
 		defer s.wg.Done()
@@ -132,21 +128,21 @@ func (s *Server) Start() error {
 		stdioServer := server.NewStdioServer(mcpServer)
 		stdioServer.SetErrorLogger(logger)
 
-		if err := stdioServer.Listen(s.ctx, s.serverReader, s.serverWriter); err != nil {
+		if err := stdioServer.Listen(ctx, s.serverReader, s.serverWriter); err != nil {
 			logger.Println("StdioServer.Listen failed:", err)
 		}
 	}()
 
 	s.transport = transport.NewIO(s.clientReader, s.clientWriter, io.NopCloser(&s.logBuffer))
-	if err := s.transport.Start(s.ctx); err != nil {
+	if err := s.transport.Start(ctx); err != nil {
 		return fmt.Errorf("transport.Start(): %w", err)
 	}
 
 	s.client = client.NewClient(s.transport)
 
 	var initReq mcp.InitializeRequest
 	initReq.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
-	if _, err := s.client.Initialize(s.ctx, initReq); err != nil {
+	if _, err := s.client.Initialize(ctx, initReq); err != nil {
 		return fmt.Errorf("client.Initialize(): %w", err)
 	}
 

mcptest/mcptest_test.go

@@ -109,7 +109,7 @@ func TestServerWithPrompt(t *testing.T) {
 
 	srv.AddPrompt(prompt, handler)
 
-	err := srv.Start()
+	err := srv.Start(ctx)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -164,7 +164,7 @@ func TestServerWithResource(t *testing.T) {
 
 	srv.AddResource(resource, handler)
 
-	err := srv.Start()
+	err := srv.Start(ctx)
 	if err != nil {
 		t.Fatal(err)
 	}