internal/mcp: rename ServerConnection->ServerSession

Bring the Server implementation up to date with the spec by renaming
ServerConnection to ServerSession. Also update variable names, which
were in various states of disarray from the multiple renamings of this
type.

Change-Id: I4aec883db6d0500fb91e6daa51b702614877cf8a
Reviewed-on: https://go-review.googlesource.com/c/tools/+/672655
Auto-Submit: Robert Findley <rfindley@google.com>
TryBot-Bypass: Robert Findley <rfindley@google.com>
Commit-Queue: Robert Findley <rfindley@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>

README.md

@@ -59,7 +59,7 @@ type HiParams struct {
 	Name string `json:"name"`
 }
 
-func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) ([]mcp.Content, error) {
+func SayHi(ctx context.Context, cc *mcp.ServerSession, params *HiParams) ([]mcp.Content, error) {
 	return []mcp.Content{
 		mcp.TextContent{Text: "Hi " + params.Name},
 	}, nil

examples/hello/main.go

@@ -20,13 +20,13 @@ type HiParams struct {
 	Name string `json:"name"`
 }
 
-func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) ([]*mcp.Content, error) {
+func SayHi(ctx context.Context, cc *mcp.ServerSession, params *HiParams) ([]*mcp.Content, error) {
 	return []*mcp.Content{
 		mcp.NewTextContent("Hi " + params.Name),
 	}, nil
 }
 
-func PromptHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) (*mcp.GetPromptResult, error) {
+func PromptHi(ctx context.Context, cc *mcp.ServerSession, params *HiParams) (*mcp.GetPromptResult, error) {
 	return &mcp.GetPromptResult{
 		Description: "Code review prompt",
 		Messages: []*mcp.PromptMessage{

examples/sse/main.go

@@ -19,7 +19,7 @@ type SayHiParams struct {
 	Name string `json:"name" mcp:"the name to say hi to"`
 }
 
-func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *SayHiParams) ([]*mcp.Content, error) {
+func SayHi(ctx context.Context, cc *mcp.ServerSession, params *SayHiParams) ([]*mcp.Content, error) {
 	return []*mcp.Content{
 		mcp.NewTextContent("Hi " + params.Name),
 	}, nil

mcp/mcp.go

@@ -10,9 +10,9 @@
 // To get started, create either a [Client] or [Server], and connect it to a
 // peer using a [Transport]. The diagram below illustrates how this works:
 //
-//	Client                                    Server
+//	Client                                   Server
-//	  ⇅          (jsonrpc2)                     ⇅
+//	  ⇅          (jsonrpc2)                    ⇅
-//	Client Transport ⇄ Server Transport ⇄ ServerConnection
+//	Client Transport ⇄ Server Transport ⇄ ServerSession
 //
 // A [Client] is an MCP client, which can be configured with various client
 // capabilities. Clients may be connected to a [Server] instance
@@ -21,7 +21,7 @@
 // Similarly, a [Server] is an MCP server, which can be configured with various
 // server capabilities. Servers may be connected to one or more [Client]
 // instances using the [Server.Connect] method, which creates a
-// [ServerConnection].
+// [ServerSession].
 //
 // A [Transport] connects a bidirectional [Stream] of jsonrpc2 messages. In
 // practice, transports in the MCP spec are are either client transports or

mcp/mcp_test.go

@@ -24,7 +24,7 @@ type hiParams struct {
 	Name string
 }
 
-func sayHi(ctx context.Context, cc *ServerConnection, v hiParams) ([]*Content, error) {
+func sayHi(ctx context.Context, cc *ServerSession, v hiParams) ([]*Content, error) {
 	if err := cc.Ping(ctx, nil); err != nil {
 		return nil, fmt.Errorf("ping failed: %v", err)
 	}
@@ -43,13 +43,13 @@ func TestEndToEnd(t *testing.T) {
 	// The 'fail' tool returns this error.
 	failure := errors.New("mcp failure")
 	s.AddTools(
-		NewTool("fail", "just fail", func(context.Context, *ServerConnection, struct{}) ([]*Content, error) {
+		NewTool("fail", "just fail", func(context.Context, *ServerSession, struct{}) ([]*Content, error) {
 			return nil, failure
 		}),
 	)
 
 	s.AddPrompts(
-		NewPrompt("code_review", "do a code review", func(_ context.Context, _ *ServerConnection, params struct{ Code string }) (*GetPromptResult, error) {
+		NewPrompt("code_review", "do a code review", func(_ context.Context, _ *ServerSession, params struct{ Code string }) (*GetPromptResult, error) {
 			return &GetPromptResult{
 				Description: "Code review prompt",
 				Messages: []*PromptMessage{
@@ -57,17 +57,17 @@ func TestEndToEnd(t *testing.T) {
 				},
 			}, nil
 		}),
-		NewPrompt("fail", "", func(_ context.Context, _ *ServerConnection, params struct{}) (*GetPromptResult, error) {
+		NewPrompt("fail", "", func(_ context.Context, _ *ServerSession, params struct{}) (*GetPromptResult, error) {
 			return nil, failure
 		}),
 	)
 
 	// Connect the server.
-	sc, err := s.Connect(ctx, st, nil)
+	ss, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
-	if got := slices.Collect(s.Clients()); len(got) != 1 {
+	if got := slices.Collect(s.Sessions()); len(got) != 1 {
 		t.Errorf("after connection, Clients() has length %d, want 1", len(got))
 	}
 
@@ -75,7 +75,7 @@ func TestEndToEnd(t *testing.T) {
 	var clientWG sync.WaitGroup
 	clientWG.Add(1)
 	go func() {
-		if err := sc.Wait(); err != nil {
+		if err := ss.Wait(); err != nil {
 			t.Errorf("server failed: %v", err)
 		}
 		clientWG.Done()
@@ -198,7 +198,7 @@ func TestEndToEnd(t *testing.T) {
 			URI:      "file:///nonexistent.txt",
 		}
 
-		readHandler := func(_ context.Context, _ *ServerConnection, p *ReadResourceParams) (*ReadResourceResult, error) {
+		readHandler := func(_ context.Context, _ *ServerSession, p *ReadResourceParams) (*ReadResourceResult, error) {
 			if p.URI == "file:///file1.txt" {
 				return &ReadResourceResult{
 					Contents: &ResourceContents{
@@ -249,9 +249,9 @@ func TestEndToEnd(t *testing.T) {
 		}
 	})
 	t.Run("roots", func(t *testing.T) {
-		// Take the server's first ServerConnection.
+		// Take the server's first ServerSession.
-		var sc *ServerConnection
+		var sc *ServerSession
-		for sc = range s.Clients() {
+		for sc = range s.Sessions() {
 			break
 		}
 
@@ -272,7 +272,7 @@ func TestEndToEnd(t *testing.T) {
 
 	// After disconnecting, neither client nor server should have any
 	// connections.
-	for range s.Clients() {
+	for range s.Sessions() {
 		t.Errorf("unexpected client after disconnection")
 	}
 }
@@ -282,7 +282,7 @@ func TestEndToEnd(t *testing.T) {
 //
 // The caller should cancel either the client connection or server connection
 // when the connections are no longer needed.
-func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerConnection, *Client) {
+func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerSession, *Client) {
 	t.Helper()
 
 	ctx := context.Background()
@@ -292,7 +292,7 @@ func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerConnection, *Cl
 
 	// The 'greet' tool says hi.
 	s.AddTools(tools...)
-	cc, err := s.Connect(ctx, st, nil)
+	ss, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -301,7 +301,7 @@ func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerConnection, *Cl
 	if err := c.Start(ctx); err != nil {
 		t.Fatal(err)
 	}
-	return cc, c
+	return ss, c
 }
 
 func TestServerClosing(t *testing.T) {
@@ -370,7 +370,7 @@ func TestCancellation(t *testing.T) {
 		cancelled = make(chan struct{}, 1) // don't block the request
 	)
 
-	slowRequest := func(ctx context.Context, cc *ServerConnection, v struct{}) ([]*Content, error) {
+	slowRequest := func(ctx context.Context, cc *ServerSession, v struct{}) ([]*Content, error) {
 		start <- struct{}{}
 		select {
 		case <-ctx.Done():

mcp/prompt.go

@@ -16,7 +16,7 @@ import (
 )
 
 // A PromptHandler handles a call to prompts/get.
-type PromptHandler func(context.Context, *ServerConnection, map[string]string) (*GetPromptResult, error)
+type PromptHandler func(context.Context, *ServerSession, map[string]string) (*GetPromptResult, error)
 
 // A Prompt is a prompt definition bound to a prompt handler.
 type ServerPrompt struct {
@@ -32,7 +32,7 @@ type ServerPrompt struct {
 // of type string or *string. The argument names for the resulting prompt
 // definition correspond to the JSON names of the request fields, and any
 // fields that are not marked "omitempty" are considered required.
-func NewPrompt[TReq any](name, description string, handler func(context.Context, *ServerConnection, TReq) (*GetPromptResult, error), opts ...PromptOption) *ServerPrompt {
+func NewPrompt[TReq any](name, description string, handler func(context.Context, *ServerSession, TReq) (*GetPromptResult, error), opts ...PromptOption) *ServerPrompt {
 	schema, err := jsonschema.For[TReq]()
 	if err != nil {
 		panic(err)
@@ -60,7 +60,7 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context,
 			Required:    required[name],
 		})
 	}
-	prompt.Handler = func(ctx context.Context, cc *ServerConnection, args map[string]string) (*GetPromptResult, error) {
+	prompt.Handler = func(ctx context.Context, cc *ServerSession, args map[string]string) (*GetPromptResult, error) {
 		// For simplicity, just marshal and unmarshal the arguments.
 		// This could be avoided in the future.
 		rawArgs, err := json.Marshal(args)

mcp/prompt_test.go

@@ -13,7 +13,7 @@ import (
 )
 
 // testPromptHandler is used for type inference in TestNewPrompt.
-func testPromptHandler[T any](context.Context, *mcp.ServerConnection, T) (*mcp.GetPromptResult, error) {
+func testPromptHandler[T any](context.Context, *mcp.ServerSession, T) (*mcp.GetPromptResult, error) {
 	panic("not implemented")
 }
 

mcp/server.go

@@ -30,7 +30,7 @@ type Server struct {
 	prompts   *featureSet[*ServerPrompt]
 	tools     *featureSet[*ServerTool]
 	resources *featureSet[*ServerResource]
-	conns     []*ServerConnection
+	conns     []*ServerSession
 }
 
 // ServerOptions is used to configure behavior of the server.
@@ -121,7 +121,7 @@ const codeResourceNotFound = -31002
 
 // A ResourceHandler is a function that reads a resource.
 // If it cannot find the resource, it should return the result of calling [ResourceNotFoundError].
-type ResourceHandler func(context.Context, *ServerConnection, *ReadResourceParams) (*ReadResourceResult, error)
+type ResourceHandler func(context.Context, *ServerSession, *ReadResourceParams) (*ReadResourceResult, error)
 
 // A ServerResource associates a Resource with its handler.
 type ServerResource struct {
@@ -157,16 +157,15 @@ func (s *Server) RemoveResources(uris ...string) {
 	s.resources.remove(uris...)
 }
 
-// Clients returns an iterator that yields the current set of client
+// Sessions returns an iterator that yields the current set of server sessions.
-// connections.
+func (s *Server) Sessions() iter.Seq[*ServerSession] {
-func (s *Server) Clients() iter.Seq[*ServerConnection] {
 	s.mu.Lock()
 	clients := slices.Clone(s.conns)
 	s.mu.Unlock()
 	return slices.Values(clients)
 }
 
-func (s *Server) listPrompts(_ context.Context, _ *ServerConnection, params *ListPromptsParams) (*ListPromptsResult, error) {
+func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPromptsParams) (*ListPromptsResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	res := new(ListPromptsResult)
@@ -176,7 +175,7 @@ func (s *Server) listPrompts(_ context.Context, _ *ServerConnection, params *Lis
 	return res, nil
 }
 
-func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *GetPromptParams) (*GetPromptResult, error) {
+func (s *Server) getPrompt(ctx context.Context, cc *ServerSession, params *GetPromptParams) (*GetPromptResult, error) {
 	s.mu.Lock()
 	prompt, ok := s.prompts.get(params.Name)
 	s.mu.Unlock()
@@ -187,7 +186,7 @@ func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *Ge
 	return prompt.Handler(ctx, cc, params.Arguments)
 }
 
-func (s *Server) listTools(_ context.Context, _ *ServerConnection, params *ListToolsParams) (*ListToolsResult, error) {
+func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListToolsParams) (*ListToolsResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	res := new(ListToolsResult)
@@ -197,7 +196,7 @@ func (s *Server) listTools(_ context.Context, _ *ServerConnection, params *ListT
 	return res, nil
 }
 
-func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *CallToolParams) (*CallToolResult, error) {
+func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParams) (*CallToolResult, error) {
 	s.mu.Lock()
 	tool, ok := s.tools.get(params.Name)
 	s.mu.Unlock()
@@ -207,7 +206,7 @@ func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *Cal
 	return tool.Handler(ctx, cc, params)
 }
 
-func (s *Server) listResources(_ context.Context, _ *ServerConnection, params *ListResourcesParams) (*ListResourcesResult, error) {
+func (s *Server) listResources(_ context.Context, _ *ServerSession, params *ListResourcesParams) (*ListResourcesResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	res := new(ListResourcesResult)
@@ -217,7 +216,7 @@ func (s *Server) listResources(_ context.Context, _ *ServerConnection, params *L
 	return res, nil
 }
 
-func (s *Server) readResource(ctx context.Context, ss *ServerConnection, params *ReadResourceParams) (*ReadResourceResult, error) {
+func (s *Server) readResource(ctx context.Context, ss *ServerSession, params *ReadResourceParams) (*ReadResourceResult, error) {
 	uri := params.URI
 	// Look up the resource URI in the list we have.
 	// This is a security check as well as an information lookup.
@@ -250,29 +249,29 @@ func (s *Server) readResource(ctx context.Context, ss *ServerConnection, params
 //
 // Run blocks until the client terminates the connection.
 func (s *Server) Run(ctx context.Context, t Transport, opts *ConnectionOptions) error {
-	cc, err := s.Connect(ctx, t, opts)
+	ss, err := s.Connect(ctx, t, opts)
 	if err != nil {
 		return err
 	}
-	return cc.Wait()
+	return ss.Wait()
 }
 
-// bind implements the binder[*ServerConnection] interface, so that Servers can
+// bind implements the binder[*ServerSession] interface, so that Servers can
 // be connected using [connect].
-func (s *Server) bind(conn *jsonrpc2.Connection) *ServerConnection {
+func (s *Server) bind(conn *jsonrpc2.Connection) *ServerSession {
-	cc := &ServerConnection{conn: conn, server: s}
+	cc := &ServerSession{conn: conn, server: s}
 	s.mu.Lock()
 	s.conns = append(s.conns, cc)
 	s.mu.Unlock()
 	return cc
 }
 
-// disconnect implements the binder[*ServerConnection] interface, so that
+// disconnect implements the binder[*ServerSession] interface, so that
 // Servers can be connected using [connect].
-func (s *Server) disconnect(cc *ServerConnection) {
+func (s *Server) disconnect(cc *ServerSession) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	s.conns = slices.DeleteFunc(s.conns, func(cc2 *ServerConnection) bool {
+	s.conns = slices.DeleteFunc(s.conns, func(cc2 *ServerSession) bool {
 		return cc2 == cc
 	})
 }
@@ -283,17 +282,17 @@ func (s *Server) disconnect(cc *ServerConnection) {
 // It returns a connection object that may be used to terminate the connection
 // (with [Connection.Close]), or await client termination (with
 // [Connection.Wait]).
-func (s *Server) Connect(ctx context.Context, t Transport, opts *ConnectionOptions) (*ServerConnection, error) {
+func (s *Server) Connect(ctx context.Context, t Transport, opts *ConnectionOptions) (*ServerSession, error) {
 	return connect(ctx, t, opts, s)
 }
 
-// A ServerConnection is a connection from a single MCP client. Its methods can
+// A ServerSession is a logical connection from a single MCP client. Its
-// be used to send requests or notifications to the client. Create a connection
+// methods can be used to send requests or notifications to the client. Create
-// by calling [Server.Connect].
+// a session by calling [Server.Connect].
 //
-// Call [ServerConnection.Close] to close the connection, or await client
+// Call [ServerSession.Close] to close the connection, or await client
-// termination with [ServerConnection.Wait].
+// termination with [ServerSession.Wait].
-type ServerConnection struct {
+type ServerSession struct {
 	server *Server
 	conn   *jsonrpc2.Connection
 
@@ -303,15 +302,15 @@ type ServerConnection struct {
 }
 
 // Ping makes an MCP "ping" request to the client.
-func (cc *ServerConnection) Ping(ctx context.Context, _ *PingParams) error {
+func (cc *ServerSession) Ping(ctx context.Context, _ *PingParams) error {
 	return call(ctx, cc.conn, "ping", nil, nil)
 }
 
-func (cc *ServerConnection) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) {
+func (cc *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) {
 	return standardCall[ListRootsResult](ctx, cc.conn, "roots/list", params)
 }
 
-func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
+func (cc *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
 	cc.mu.Lock()
 	initialized := cc.initialized
 	cc.mu.Unlock()
@@ -362,7 +361,7 @@ func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (
 	return nil, jsonrpc2.ErrNotHandled
 }
 
-func (cc *ServerConnection) initialize(ctx context.Context, _ *ServerConnection, params *initializeParams) (*initializeResult, error) {
+func (cc *ServerSession) initialize(ctx context.Context, _ *ServerSession, params *initializeParams) (*initializeResult, error) {
 	cc.mu.Lock()
 	cc.initializeParams = params
 	cc.mu.Unlock()
@@ -400,12 +399,12 @@ func (cc *ServerConnection) initialize(ctx context.Context, _ *ServerConnection,
 // Close performs a graceful shutdown of the connection, preventing new
 // requests from being handled, and waiting for ongoing requests to return.
 // Close then terminates the connection.
-func (cc *ServerConnection) Close() error {
+func (cc *ServerSession) Close() error {
 	return cc.conn.Close()
 }
 
 // Wait waits for the connection to be closed by the client.
-func (cc *ServerConnection) Wait() error {
+func (cc *ServerSession) Wait() error {
 	return cc.conn.Wait()
 }
 
@@ -413,7 +412,7 @@ func (cc *ServerConnection) Wait() error {
 //
 // Importantly, it returns nil if the handler returned an error, which is a
 // requirement of the jsonrpc2 package.
-func dispatch[TParams, TResult any](ctx context.Context, conn *ServerConnection, req *jsonrpc2.Request, f func(context.Context, *ServerConnection, TParams) (TResult, error)) (any, error) {
+func dispatch[TParams, TResult any](ctx context.Context, conn *ServerSession, req *jsonrpc2.Request, f func(context.Context, *ServerSession, TParams) (TResult, error)) (any, error) {
 	var params TParams
 	if err := json.Unmarshal(req.Params, &params); err != nil {
 		return nil, err

mcp/server_example_test.go

@@ -16,7 +16,7 @@ type SayHiParams struct {
 	Name string `json:"name" mcp:"the name to say hi to"`
 }
 
-func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *SayHiParams) ([]*mcp.Content, error) {
+func SayHi(ctx context.Context, cc *mcp.ServerSession, params *SayHiParams) ([]*mcp.Content, error) {
 	return []*mcp.Content{
 		mcp.NewTextContent("Hi " + params.Name),
 	}, nil
@@ -29,7 +29,7 @@ func ExampleServer() {
 	server := mcp.NewServer("greeter", "v0.0.1", nil)
 	server.AddTools(mcp.NewTool("greet", "say hi", SayHi))
 
-	clientConnection, err := server.Connect(ctx, serverTransport, nil)
+	serverSession, err := server.Connect(ctx, serverTransport, nil)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -46,7 +46,7 @@ func ExampleServer() {
 	fmt.Println(res.Content[0].Text)
 
 	client.Close()
-	clientConnection.Wait()
+	serverSession.Wait()
 
 	// Output: Hi user
 }

mcp/sse.go

@@ -64,7 +64,7 @@ func writeEvent(w io.Writer, evt event) (int, error) {
 // https://modelcontextprotocol.io/specification/2024-11-05/basic/transports
 type SSEHandler struct {
 	getServer    func(request *http.Request) *Server
-	onConnection func(*ServerConnection) // for testing; must not block
+	onConnection func(*ServerSession) // for testing; must not block
 
 	mu       sync.Mutex
 	sessions map[string]*SSEServerTransport
@@ -229,15 +229,15 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 
 	// TODO(hxjiang): getServer returns nil will panic.
 	server := h.getServer(req)
-	cc, err := server.Connect(req.Context(), transport, nil)
+	ss, err := server.Connect(req.Context(), transport, nil)
 	if err != nil {
 		http.Error(w, "connection failed", http.StatusInternalServerError)
 		return
 	}
 	if h.onConnection != nil {
-		h.onConnection(cc)
+		h.onConnection(ss)
 	}
-	defer cc.Close() // close the transport when the GET exits
+	defer ss.Close() // close the transport when the GET exits
 
 	select {
 	case <-req.Context().Done():

mcp/sse_example_test.go

@@ -18,7 +18,7 @@ type AddParams struct {
 	X, Y int
 }
 
-func Add(ctx context.Context, cc *mcp.ServerConnection, params *AddParams) ([]*mcp.Content, error) {
+func Add(ctx context.Context, cc *mcp.ServerSession, params *AddParams) ([]*mcp.Content, error) {
 	return []*mcp.Content{
 		mcp.NewTextContent(fmt.Sprintf("%d", params.X+params.Y)),
 	}, nil

mcp/sse_test.go

@@ -23,8 +23,8 @@ func TestSSEServer(t *testing.T) {
 
 			sseHandler := NewSSEHandler(func(*http.Request) *Server { return server })
 
-			conns := make(chan *ServerConnection, 1)
+			conns := make(chan *ServerSession, 1)
-			sseHandler.onConnection = func(cc *ServerConnection) {
+			sseHandler.onConnection = func(cc *ServerSession) {
 				select {
 				case conns <- cc:
 				default:

mcp/tool.go

@@ -13,7 +13,7 @@ import (
 )
 
 // A ToolHandler handles a call to tools/call.
-type ToolHandler func(context.Context, *ServerConnection, *CallToolParams) (*CallToolResult, error)
+type ToolHandler func(context.Context, *ServerSession, *CallToolParams) (*CallToolResult, error)
 
 // A Tool is a tool definition that is bound to a tool handler.
 type ServerTool struct {
@@ -34,12 +34,12 @@ type ServerTool struct {
 //
 // TODO: just have the handler return a CallToolResult: returning []Content is
 // going to be inconsistent with other server features.
-func NewTool[TReq any](name, description string, handler func(context.Context, *ServerConnection, TReq) ([]*Content, error), opts ...ToolOption) *ServerTool {
+func NewTool[TReq any](name, description string, handler func(context.Context, *ServerSession, TReq) ([]*Content, error), opts ...ToolOption) *ServerTool {
 	schema, err := jsonschema.For[TReq]()
 	if err != nil {
 		panic(err)
 	}
-	wrapped := func(ctx context.Context, cc *ServerConnection, params *CallToolParams) (*CallToolResult, error) {
+	wrapped := func(ctx context.Context, cc *ServerSession, params *CallToolParams) (*CallToolResult, error) {
 		var v TReq
 		if err := unmarshalSchema(params.Arguments, schema, &v); err != nil {
 			return nil, err

mcp/tool_test.go

@@ -15,7 +15,7 @@ import (
 )
 
 // testToolHandler is used for type inference in TestNewTool.
-func testToolHandler[T any](context.Context, *mcp.ServerConnection, T) ([]*mcp.Content, error) {
+func testToolHandler[T any](context.Context, *mcp.ServerSession, T) ([]*mcp.Content, error) {
 	panic("not implemented")
 }