Repository: modelcontextprotocol/go-sdk

Issue:
add http.Header to RequestExtra
Make incoming HTTP headers visible to tool handlers and other user code. 

One use case: a tool can pass select headers on to backends.

See #292.

Discussion:
@jba: #292 must land before work on this can begin.

Solution patch:
diff --git a/mcp/client.go b/mcp/client.go
index d1d1750..2511c05 100644
--- a/mcp/client.go
+++ b/mcp/client.go
@@ -105,6 +105,15 @@ func (e unsupportedProtocolVersionError) Error() string {
 // ClientSessionOptions is reserved for future use.
 type ClientSessionOptions struct{}
 
+func (c *Client) capabilities() *ClientCapabilities {
+	caps := &ClientCapabilities{}
+	caps.Roots.ListChanged = true
+	if c.opts.CreateMessageHandler != nil {
+		caps.Sampling = &SamplingCapabilities{}
+	}
+	return caps
+}
+
 // Connect begins an MCP session by connecting to a server over the given
 // transport, and initializing the session.
 //
@@ -118,16 +127,10 @@ func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptio
 		return nil, err
 	}
 
-	caps := &ClientCapabilities{}
-	caps.Roots.ListChanged = true
-	if c.opts.CreateMessageHandler != nil {
-		caps.Sampling = &SamplingCapabilities{}
-	}
-
 	params := &InitializeParams{
 		ProtocolVersion: latestProtocolVersion,
 		ClientInfo:      c.impl,
-		Capabilities:    caps,
+		Capabilities:    c.capabilities(),
 	}
 	req := &ClientRequest[*InitializeParams]{Session: cs, Params: params}
 	res, err := handleSend[*InitializeResult](ctx, methodInitialize, req)
@@ -226,7 +229,6 @@ func (c *Client) AddRoots(roots ...*Root) {
 // RemoveRoots removes the roots with the given URIs,
 // and notifies any connected servers if the list has changed.
 // It is not an error to remove a nonexistent root.
-// TODO: notification
 func (c *Client) RemoveRoots(uris ...string) {
 	changeAndNotify(c, notificationRootsListChanged, &RootsListChangedParams{},
 		func() bool { return c.roots.remove(uris...) })
diff --git a/mcp/features.go b/mcp/features.go
index 43c5885..438370f 100644
--- a/mcp/features.go
+++ b/mcp/features.go
@@ -17,7 +17,9 @@ import (
 // A featureSet is a collection of features of type T.
 // Every feature has a unique ID, and the spec never mentions
 // an ordering for the List calls, so what it calls a "list" is actually a set.
-// TODO: switch to an ordered map
+//
+// An alternative implementation would use an ordered map, but that's probably
+// not necessary as adds and removes are rare, and usually batched.
 type featureSet[T any] struct {
 	uniqueID   func(T) string
 	features   map[string]T
diff --git a/mcp/server.go b/mcp/server.go
index 8802133..d98ff8a 100644
--- a/mcp/server.go
+++ b/mcp/server.go
@@ -89,17 +89,18 @@ type ServerOptions struct {
 // The first argument must not be nil.
 //
 // If non-nil, the provided options are used to configure the server.
-func NewServer(impl *Implementation, opts *ServerOptions) *Server {
+func NewServer(impl *Implementation, options *ServerOptions) *Server {
 	if impl == nil {
 		panic("nil Implementation")
 	}
-	if opts == nil {
-		opts = new(ServerOptions)
+	var opts ServerOptions
+	if options != nil {
+		opts = *options
 	}
+	options = nil // prevent reuse
 	if opts.PageSize < 0 {
 		panic(fmt.Errorf("invalid page size %d", opts.PageSize))
 	}
-	// TODO(jba): don't modify opts, modify Server.opts.
 	if opts.PageSize == 0 {
 		opts.PageSize = DefaultPageSize
 	}
@@ -111,7 +112,7 @@ func NewServer(impl *Implementation, opts *ServerOptions) *Server {
 	}
 	return &Server{
 		impl:                    impl,
-		opts:                    *opts,
+		opts:                    opts,
 		prompts:                 newFeatureSet(func(p *serverPrompt) string { return p.prompt.Name }),
 		tools:                   newFeatureSet(func(t *serverTool) string { return t.tool.Name }),
 		resources:               newFeatureSet(func(r *serverResource) string { return r.resource.URI }),
@@ -328,8 +329,11 @@ func (s *Server) getPrompt(ctx context.Context, req *ServerRequest[*GetPromptPar
 	prompt, ok := s.prompts.get(req.Params.Name)
 	s.mu.Unlock()
 	if !ok {
-		// TODO: surface the error code over the wire, instead of flattening it into the string.
-		return nil, fmt.Errorf("%s: unknown prompt %q", jsonrpc2.ErrInvalidParams, req.Params.Name)
+		// Return a proper JSON-RPC error with the correct error code
+		return nil, &jsonrpc2.WireError{
+			Code:    CodeInvalidParams,
+			Message: fmt.Sprintf("unknown prompt %q", req.Params.Name),
+		}
 	}
 	return prompt.handler(ctx, req.Session, req.Params)
 }
@@ -353,7 +357,10 @@ func (s *Server) callTool(ctx context.Context, req *ServerRequest[*CallToolParam
 	st, ok := s.tools.get(req.Params.Name)
 	s.mu.Unlock()
 	if !ok {
-		return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, req.Params.Name)
+		return nil, &jsonrpc2.WireError{
+			Code:    CodeInvalidParams,
+			Message: fmt.Sprintf("unknown tool %q", req.Params.Name),
+		}
 	}
 	// TODO: if handler returns nil content, it will serialize as null.
 	// Add a test and fix.
@@ -457,7 +464,7 @@ func fileResourceHandler(dir string) ResourceHandler {
 	return func(ctx context.Context, req *ServerRequest[*ReadResourceParams]) (_ *ReadResourceResult, err error) {
 		defer util.Wrapf(&err, "reading resource %s", req.Params.URI)
 
-		// TODO: use a memoizing API here.
+		// TODO(#25): use a memoizing API here.
 		rootRes, err := req.Session.ListRoots(ctx, nil)
 		if err != nil {
 			return nil, fmt.Errorf("listing roots: %w", err)
@@ -842,6 +849,8 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any,
 	return handleReceive(ctx, ss, req)
 }
 
+func (ss *ServerSession) InitializeParams() *InitializeParams { return ss.state.InitializeParams }
+
 func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParams) (*InitializeResult, error) {
 	if params == nil {
 		return nil, fmt.Errorf("%w: \"params\" must be be provided", jsonrpc2.ErrInvalidParams)
diff --git a/mcp/shared.go b/mcp/shared.go
index bda631f..b7b8bda 100644
--- a/mcp/shared.go
+++ b/mcp/shared.go
@@ -4,8 +4,10 @@
 
 // This file contains code shared between client and server, including
 // method handler and middleware definitions.
-// TODO: much of this is here so that we can factor out commonalities using
-// generics. Perhaps it can be simplified with reflection.
+//
+// Much of this is here so that we can factor out commonalities using
+// generics. If this becomes unwieldy, it can perhaps be simplified with
+// reflection.
 
 package mcp
 
@@ -14,6 +16,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log"
+	"net/http"
 	"reflect"
 	"slices"
 	"strings"
@@ -311,6 +314,8 @@ const (
 	CodeResourceNotFound = -32002
 	// The error code if the method exists and was called properly, but the peer does not support it.
 	CodeUnsupportedMethod = -31001
+	// The error code for invalid parameters
+	CodeInvalidParams = -32602
 )
 
 // notifySessions calls Notify on all the sessions.
@@ -396,6 +401,7 @@ type ServerRequest[P Params] struct {
 // the transport layer.
 type RequestExtra struct {
 	TokenInfo *auth.TokenInfo // bearer token info (e.g. from OAuth) if any
+	Header    http.Header     // header from HTTP request, if any
 }
 
 func (*ClientRequest[P]) isRequest() {}
diff --git a/mcp/streamable.go b/mcp/streamable.go
index c51f3cc..eb42ca5 100644
--- a/mcp/streamable.go
+++ b/mcp/streamable.go
@@ -580,10 +580,6 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 	// This also requires access to the negotiated version, which would either be
 	// set by the MCP-Protocol-Version header, or would require peeking into the
 	// session.
-	if err != nil {
-		http.Error(w, fmt.Sprintf("malformed payload: %v", err), http.StatusBadRequest)
-		return
-	}
 	incoming, _, err := readBatch(body)
 	if err != nil {
 		http.Error(w, fmt.Sprintf("malformed payload: %v", err), http.StatusBadRequest)
@@ -592,17 +588,20 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 	requests := make(map[jsonrpc.ID]struct{})
 	tokenInfo := auth.TokenInfoFromContext(req.Context())
 	for _, msg := range incoming {
-		if req, ok := msg.(*jsonrpc.Request); ok {
+		if jreq, ok := msg.(*jsonrpc.Request); ok {
 			// Preemptively check that this is a valid request, so that we can fail
 			// the HTTP request. If we didn't do this, a request with a bad method or
 			// missing ID could be silently swallowed.
-			if _, err := checkRequest(req, serverMethodInfos); err != nil {
+			if _, err := checkRequest(jreq, serverMethodInfos); err != nil {
 				http.Error(w, err.Error(), http.StatusBadRequest)
 				return
 			}
-			req.Extra = &RequestExtra{TokenInfo: tokenInfo}
-			if req.IsCall() {
-				requests[req.ID] = struct{}{}
+			jreq.Extra = &RequestExtra{
+				TokenInfo: tokenInfo,
+				Header:    req.Header,
+			}
+			if jreq.IsCall() {
+				requests[jreq.ID] = struct{}{}
 			}
 		}
 	}
@@ -1152,7 +1151,6 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e
 	}
 
 	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-		// TODO: do a best effort read of the body here, and format it in the error.
 		resp.Body.Close()
 		return fmt.Errorf("broken session: %v", resp.Status)
 	}
diff --git a/mcp/tool.go b/mcp/tool.go
index 7173b8a..893b48f 100644
--- a/mcp/tool.go
+++ b/mcp/tool.go
@@ -12,6 +12,7 @@ import (
 	"reflect"
 
 	"github.com/google/jsonschema-go/jsonschema"
+	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
 )
 
 // A ToolHandler handles a call to tools/call.
@@ -70,9 +71,16 @@ func newServerTool[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*serverTool
 			Params:  params,
 			Extra:   req.Extra,
 		})
-		// TODO(rfindley): investigate why server errors are embedded in this strange way,
-		// rather than returned as jsonrpc2 server errors.
+		// Handle server errors appropriately:
+		// - If the handler returns a structured error (like jsonrpc2.WireError), return it directly
+		// - If the handler returns a regular error, wrap it in a CallToolResult with IsError=true
+		// - This allows tools to distinguish between protocol errors and tool execution errors
 		if err != nil {
+			// Check if this is already a structured JSON-RPC error
+			if wireErr, ok := err.(*jsonrpc2.WireError); ok {
+				return nil, wireErr
+			}
+			// For regular errors, embed them in the tool result as per MCP spec
 			return &CallToolResult{
 				Content: []Content{&TextContent{Text: err.Error()}},
 				IsError: true,



Test patch:
diff --git a/mcp/client_test.go b/mcp/client_test.go
index 7920c55..469fa3f 100644
--- a/mcp/client_test.go
+++ b/mcp/client_test.go
@@ -190,3 +190,48 @@ func TestClientPaginateVariousPageSizes(t *testing.T) {
 		})
 	}
 }
+
+func TestClientCapabilities(t *testing.T) {
+	testCases := []struct {
+		name             string
+		configureClient  func(s *Client)
+		clientOpts       ClientOptions
+		wantCapabilities *ClientCapabilities
+	}{
+		{
+			name:            "With initial capabilities",
+			configureClient: func(s *Client) {},
+			wantCapabilities: &ClientCapabilities{
+				Roots: struct {
+					ListChanged bool "json:\"listChanged,omitempty\""
+				}{ListChanged: true},
+			},
+		},
+		{
+			name:            "With sampling",
+			configureClient: func(s *Client) {},
+			clientOpts: ClientOptions{
+				CreateMessageHandler: func(context.Context, *ClientRequest[*CreateMessageParams]) (*CreateMessageResult, error) {
+					return nil, nil
+				},
+			},
+			wantCapabilities: &ClientCapabilities{
+				Roots: struct {
+					ListChanged bool "json:\"listChanged,omitempty\""
+				}{ListChanged: true},
+				Sampling: &SamplingCapabilities{},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			client := NewClient(testImpl, &tc.clientOpts)
+			tc.configureClient(client)
+			gotCapabilities := client.capabilities()
+			if diff := cmp.Diff(tc.wantCapabilities, gotCapabilities); diff != "" {
+				t.Errorf("capabilities() mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
diff --git a/mcp/tool_test.go b/mcp/tool_test.go
index 609536c..4c73ec6 100644
--- a/mcp/tool_test.go
+++ b/mcp/tool_test.go
@@ -7,12 +7,16 @@ package mcp
 import (
 	"context"
 	"encoding/json"
+	"errors"
+	"fmt"
 	"reflect"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/google/jsonschema-go/jsonschema"
+	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
 )
 
 // testToolHandler is used for type inference in TestNewServerTool.
@@ -132,3 +136,91 @@ func TestUnmarshalSchema(t *testing.T) {
 
 	}
 }
+
+func TestToolErrorHandling(t *testing.T) {
+	// Construct server and add both tools at the top level
+	server := NewServer(testImpl, nil)
+
+	// Create a tool that returns a structured error
+	structuredErrorHandler := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResultFor[any], error) {
+		return nil, &jsonrpc2.WireError{
+			Code:    CodeInvalidParams,
+			Message: "internal server error",
+		}
+	}
+
+	// Create a tool that returns a regular error
+	regularErrorHandler := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResultFor[any], error) {
+		return nil, fmt.Errorf("tool execution failed")
+	}
+
+	AddTool(server, &Tool{Name: "error_tool", Description: "returns structured error"}, structuredErrorHandler)
+	AddTool(server, &Tool{Name: "regular_error_tool", Description: "returns regular error"}, regularErrorHandler)
+
+	// Connect server and client once
+	ct, st := NewInMemoryTransports()
+	_, err := server.Connect(context.Background(), st, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	client := NewClient(testImpl, nil)
+	cs, err := client.Connect(context.Background(), ct, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cs.Close()
+
+	// Test that structured JSON-RPC errors are returned directly
+	t.Run("structured_error", func(t *testing.T) {
+		// Call the tool
+		_, err = cs.CallTool(context.Background(), &CallToolParams{
+			Name:      "error_tool",
+			Arguments: map[string]any{},
+		})
+
+		// Should get the structured error directly
+		if err == nil {
+			t.Fatal("expected error, got nil")
+		}
+
+		var wireErr *jsonrpc2.WireError
+		if !errors.As(err, &wireErr) {
+			t.Fatalf("expected WireError, got %[1]T: %[1]v", err)
+		}
+
+		if wireErr.Code != CodeInvalidParams {
+			t.Errorf("expected error code %d, got %d", CodeInvalidParams, wireErr.Code)
+		}
+	})
+
+	// Test that regular errors are embedded in tool results
+	t.Run("regular_error", func(t *testing.T) {
+		// Call the tool
+		result, err := cs.CallTool(context.Background(), &CallToolParams{
+			Name:      "regular_error_tool",
+			Arguments: map[string]any{},
+		})
+
+		// Should not get an error at the protocol level
+		if err != nil {
+			t.Fatalf("unexpected protocol error: %v", err)
+		}
+
+		// Should get a result with IsError=true
+		if !result.IsError {
+			t.Error("expected IsError=true, got false")
+		}
+
+		// Should have error message in content
+		if len(result.Content) == 0 {
+			t.Error("expected error content, got empty")
+		}
+
+		if textContent, ok := result.Content[0].(*TextContent); !ok {
+			t.Error("expected TextContent")
+		} else if !strings.Contains(textContent.Text, "tool execution failed") {
+			t.Errorf("expected error message in content, got: %s", textContent.Text)
+		}
+	})
+}