mcp: pass TokenInfo to server handler

[jba]

If there is a TokenInfo in the request context of a StreamableServerTransport, then propagate it through to the ServerRequest that is passed to server methods like callTool.

Fixes #317.

[jba]

Sorry about the force push! I rebased onto main.
Diffs:

diff --git a/mcp/shared.go b/mcp/shared.go
index 43b3026..e685c2f 100644
--- a/mcp/shared.go
+++ b/mcp/shared.go
@@ -133,8 +133,8 @@ func handleReceive[S Session](ctx context.Context, session S, jreq *jsonrpc.Requ
 	}
 
 	mh := session.receivingMethodHandler().(MethodHandler)
-	ti, _ := jreq.Extra.(*RequestExtra)
-	req := info.newRequest(session, params, ti)
+	re, _ := jreq.Extra.(*RequestExtra)
+	req := info.newRequest(session, params, re)
 	// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
 	res, err := mh(ctx, jreq.Method, req)
 	if err != nil {
@@ -181,7 +181,7 @@ type methodInfo struct {
 	// Unmarshal params from the wire into a Params struct.
 	// Used on the receive side.
 	unmarshalParams func(json.RawMessage) (Params, error)
-	newRequest      func(Session, Params, *auth.TokenInfo) Request
+	newRequest      func(Session, Params, *RequestExtra) Request
 	// Run the code when a call to the method is received.
 	// Used on the receive side.
 	handleMethod methodHandler
@@ -216,7 +216,7 @@ const (
 
 func newClientMethodInfo[P paramsPtr[T], R Result, T any](d typedClientMethodHandler[P, R], flags methodFlags) methodInfo {
 	mi := newMethodInfo[P, R](flags)
-	mi.newRequest = func(s Session, p Params, _ *auth.TokenInfo) Request {
+	mi.newRequest = func(s Session, p Params, _ *RequestExtra) Request {
 		r := &ClientRequest[P]{Session: s.(*ClientSession)}
 		if p != nil {
 			r.Params = p.(P)
@@ -231,8 +231,8 @@ func newClientMethodInfo[P paramsPtr[T], R Result, T any](d typedClientMethodHan
 
 func newServerMethodInfo[P paramsPtr[T], R Result, T any](d typedServerMethodHandler[P, R], flags methodFlags) methodInfo {
 	mi := newMethodInfo[P, R](flags)
-	mi.newRequest = func(s Session, p Params, ti *auth.TokenInfo) Request {
-		r := &ServerRequest[P]{Session: s.(*ServerSession), Extra: RequestExtra{TokenInfo: ti}}
+	mi.newRequest = func(s Session, p Params, re *RequestExtra) Request {
+		r := &ServerRequest[P]{Session: s.(*ServerSession), Extra: re}
 		if p != nil {
 			r.Params = p.(P)
 		}
@@ -395,7 +395,7 @@ type ClientRequest[P Params] struct {
 type ServerRequest[P Params] struct {
 	Session *ServerSession
 	Params  P
-	Extra   RequestExtra
+	Extra   *RequestExtra
 }
 
 // RequestExtra is extra information included in requests, typically from
diff --git a/mcp/streamable.go b/mcp/streamable.go
index 5757be5..365d1b8 100644
--- a/mcp/streamable.go
+++ b/mcp/streamable.go
@@ -511,7 +511,7 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 				http.Error(w, err.Error(), http.StatusBadRequest)
 				return
 			}
-			req.Extra = tokenInfo
+			req.Extra = &RequestExtra{TokenInfo: tokenInfo}
 			if req.ID.IsValid() {
 				requests[req.ID] = struct{}{}
 			}
diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go
index d0f3210..dff7cc4 100644
--- a/mcp/streamable_test.go
+++ b/mcp/streamable_test.go
@@ -1047,7 +1047,7 @@ func TestTokenInfo(t *testing.T) {
 
 	// Create a server with a tool that returns TokenInfo.
 	tokenInfo := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[struct{}]]) (*CallToolResultFor[any], error) {
-		return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: fmt.Sprintf("%v", req.TokenInfo)}}}, nil
+		return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: fmt.Sprintf("%v", req.Extra.TokenInfo)}}}, nil
 	}
 	server := NewServer(testImpl, nil)
 	AddTool(server, &Tool{Name: "tokenInfo", Description: "return token info"}, tokenInfo)
diff --git a/mcp/tool.go b/mcp/tool.go
index 052e987..7173b8a 100644
--- a/mcp/tool.go
+++ b/mcp/tool.go
@@ -66,9 +66,9 @@ func newServerTool[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*serverTool
 		}
 		// TODO(jba): improve copy
 		res, err := h(ctx, &ServerRequest[*CallToolParamsFor[In]]{
-			Session:   req.Session,
-			Params:    params,
-			TokenInfo: req.TokenInfo,
+			Session: req.Session,
+			Params:  params,
+			Extra:   req.Extra,
 		})
 		// TODO(rfindley): investigate why server errors are embedded in this strange way,
 		// rather than returned as jsonrpc2 server errors.
diff --git a/mcp/transport.go b/mcp/transport.go
index 6d25de3..1d2da5d 100644
--- a/mcp/transport.go
+++ b/mcp/transport.go
@@ -86,8 +86,7 @@ type serverConnection interface {
 
 // A StdioTransport is a [Transport] that communicates over stdin/stdout using
 // newline-delimited JSON.
-type StdioTransport struct {
-}
+type StdioTransport struct{}
 
 // Connect implements the [Transport] interface.
 func (*StdioTransport) Connect(context.Context) (Connection, error) {

[pagarwal-tibco]

@jba just wanted to check/confirm that you are not limiting to specific http headers and also support any custom headers to be propagated to tool implementation

[jba]

We can do that, but can you give an example of how you'd use it?

[findleyr]

LGTM, but I know nothing about oauth. Add dneil for that, either in this CL or as an audit pass?

[jba]

The only OAuth-relevant part of this PR, AFAICT, is RequiredBearerTokenOptions.Scopes. That comes right from type typescript:
https://github.com/modelcontextprotocol/typescript-sdk/blob/main/src/server/auth/middleware/bearerAuth.ts#L6-L15

I guess we could bikeshed the name ("RequiredScopes"?) but otherwise there's nothing to see here.

[pagarwal-tibco]

We can do that, but can you give an example of how you'd use it?

We have usecase where we need to propagate the incoming header to backend APIs invoked by tool implementation