Http headers (#480)

* Add HTTP headers to Request

- Pass headers from request via ctx
- Add headers from ctx to CallToolRequest to further pass it to handlers.

* Add http-headers field to all Request Structs

* Use request_handler.tmpl to add http headers

* Add tests and documentation for Header passthrough

* Chore: fix comment in CallToolRequest struct

* Fix Race condition while setting context

Author: adityamj

mcp/prompts.go

@@ -1,11 +1,14 @@
 package mcp
 
+import "net/http"
+
 /* Prompts */
 
 // ListPromptsRequest is sent from the client to request a list of prompts and
 // prompt templates the server has.
 type ListPromptsRequest struct {
 	PaginatedRequest
+	Header http.Header `json:"-"`
 }
 
 // ListPromptsResult is the server's response to a prompts/list request from
@@ -20,6 +23,7 @@ type ListPromptsResult struct {
 type GetPromptRequest struct {
 	Request
 	Params GetPromptParams `json:"params"`
+	Header http.Header     `json:"-"`
 }
 
 type GetPromptParams struct {

mcp/tools.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"net/http"
 	"reflect"
 	"strconv"
 )
@@ -14,6 +15,7 @@ var errToolSchemaConflict = errors.New("provide either InputSchema or RawInputSc
 // server has.
 type ListToolsRequest struct {
 	PaginatedRequest
+	Header http.Header `json:"-"`
 }
 
 // ListToolsResult is the server's response to a tools/list request from the
@@ -45,6 +47,7 @@ type CallToolResult struct {
 // CallToolRequest is used by the client to invoke a tool provided by the server.
 type CallToolRequest struct {
 	Request
+	Header http.Header    `json:"-"` // HTTP headers from the original request
 	Params CallToolParams `json:"params"`
 }
 

mcp/types.go

@@ -9,6 +9,7 @@ import (
 	"strconv"
 
 	"github.com/yosida95/uritemplate/v3"
+	"net/http"
 )
 
 type MCPMethod string
@@ -399,6 +400,7 @@ type CancelledNotificationParams struct {
 type InitializeRequest struct {
 	Request
 	Params InitializeParams `json:"params"`
+	Header http.Header      `json:"-"`
 }
 
 type InitializeParams struct {
@@ -489,6 +491,7 @@ type Implementation struct {
 // or else may be disconnected.
 type PingRequest struct {
 	Request
+	Header http.Header `json:"-"`
 }
 
 /* Progress notifications */
@@ -541,6 +544,7 @@ type PaginatedResult struct {
 // the server has.
 type ListResourcesRequest struct {
 	PaginatedRequest
+	Header http.Header `json:"-"`
 }
 
 // ListResourcesResult is the server's response to a resources/list request
@@ -554,6 +558,7 @@ type ListResourcesResult struct {
 // resource templates the server has.
 type ListResourceTemplatesRequest struct {
 	PaginatedRequest
+	Header http.Header `json:"-"`
 }
 
 // ListResourceTemplatesResult is the server's response to a
@@ -567,6 +572,7 @@ type ListResourceTemplatesResult struct {
 // specific resource URI.
 type ReadResourceRequest struct {
 	Request
+	Header http.Header        `json:"-"`
 	Params ReadResourceParams `json:"params"`
 }
 
@@ -598,6 +604,7 @@ type ResourceListChangedNotification struct {
 type SubscribeRequest struct {
 	Request
 	Params SubscribeParams `json:"params"`
+	Header http.Header     `json:"-"`
 }
 
 type SubscribeParams struct {
@@ -612,6 +619,7 @@ type SubscribeParams struct {
 type UnsubscribeRequest struct {
 	Request
 	Params UnsubscribeParams `json:"params"`
+	Header http.Header       `json:"-"`
 }
 
 type UnsubscribeParams struct {
@@ -717,6 +725,7 @@ func (BlobResourceContents) isResourceContents() {}
 type SetLevelRequest struct {
 	Request
 	Params SetLevelParams `json:"params"`
+	Header http.Header    `json:"-"`
 }
 
 type SetLevelParams struct {
@@ -980,6 +989,7 @@ type ModelHint struct {
 type CompleteRequest struct {
 	Request
 	Params CompleteParams `json:"params"`
+	Header http.Header    `json:"-"`
 }
 
 type CompleteParams struct {
@@ -1032,6 +1042,7 @@ type PromptReference struct {
 // structure or access specific locations that the client has permission to read from.
 type ListRootsRequest struct {
 	Request
+	Header http.Header `json:"-"`
 }
 
 // ListRootsResult is the client's response to a roots/list request from the server.

server/ctx.go

@@ -0,0 +1,8 @@
+package server
+
+type contextKey int
+
+const (
+	// This const is used as key for context value lookup
+	requestHeader contextKey = iota
+)

server/internal/gen/request_handler.go.tmpl

@@ -72,6 +72,14 @@ func (s *MCPServer) HandleMessage(
     	)
     }
 
+    // Get request header from ctx
+    h := ctx.Value(requestHeader)
+	headers, ok := h.(http.Header)
+
+	if headers == nil || !ok {
+		headers = make(http.Header)
+	}
+
 	switch baseMessage.Method {
 	{{- range .}}
 	case mcp.{{.MethodName}}:
@@ -90,6 +98,7 @@ func (s *MCPServer) HandleMessage(
 				err:  &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method},
 			}
 		} else {
+            request.Header = headers
 			s.hooks.before{{.HookName}}(ctx, baseMessage.ID, &request)
 			result, err = s.{{.HandlerFunc}}(ctx, baseMessage.ID, request)
 		}

server/request_handler.go

@@ -504,7 +504,8 @@ func (s *SSEServer) handleMessage(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusAccepted)
 
 	// Create a new context for handling the message that will be canceled when the message handling is done
-	messageCtx, cancel := context.WithCancel(detachedCtx)
+	messageCtx := context.WithValue(detachedCtx, requestHeader, r.Header)
+	messageCtx, cancel := context.WithCancel(messageCtx)
 
 	go func(ctx context.Context) {
 		defer cancel()