feat: allow customising context via callback on transport servers

This commit adds an callback function to the transport servers that allows
developers to inject context values into the server context.

This can be used to inject context values extracted from environment variables
(in stdio mode) or from headers (in sse mode), and access them in tools
using the provided context.

Author: sd2k

server/sse.go

@@ -12,13 +12,19 @@ import (
 	"github.com/mark3labs/mcp-go/mcp"
 )
 
+// SSEContextFunc is a function that takes an existing context and the current
+// request and returns a potentially modified context based on the request
+// content. This can be used to inject context values from headers, for example.
+type SSEContextFunc func(ctx context.Context, r *http.Request) context.Context
+
 // SSEServer implements a Server-Sent Events (SSE) based MCP server.
 // It provides real-time communication capabilities over HTTP using the SSE protocol.
 type SSEServer struct {
-	server   *MCPServer
-	baseURL  string
-	sessions sync.Map
-	srv      *http.Server
+	server      *MCPServer
+	baseURL     string
+	sessions    sync.Map
+	srv         *http.Server
+	contextFunc SSEContextFunc
 }
 
 // sseSession represents an active SSE connection.
@@ -36,6 +42,12 @@ func NewSSEServer(server *MCPServer, baseURL string) *SSEServer {
 	}
 }
 
+// SetContextFunc sets a function that will be called to customise the context
+// to the server using the incoming request.
+func (s *SSEServer) SetContextFunc(fn SSEContextFunc) {
+	s.contextFunc = fn
+}
+
 // NewTestServer creates a test server for testing purposes
 func NewTestServer(server *MCPServer) *httptest.Server {
 	sseServer := &SSEServer{
@@ -172,6 +184,10 @@ func (s *SSEServer) handleMessage(w http.ResponseWriter, r *http.Request) {
 		SessionID: sessionID,
 	})
 
+	if s.contextFunc != nil {
+		ctx = s.contextFunc(ctx, r)
+	}
+
 	sessionI, ok := s.sessions.Load(sessionID)
 	if !ok {
 		s.writeJSONRPCError(w, nil, mcp.INVALID_PARAMS, "Invalid session ID")

server/stdio.go

@@ -14,12 +14,19 @@ import (
 	"github.com/mark3labs/mcp-go/mcp"
 )
 
+// StdioContextFunc is a function that takes an existing context and returns
+// a potentially modified context.
+// This can be used to inject context values from environment variables,
+// for example.
+type StdioContextFunc func(ctx context.Context) context.Context
+
 // StdioServer wraps a MCPServer and handles stdio communication.
 // It provides a simple way to create command-line MCP servers that
 // communicate via standard input/output streams using JSON-RPC messages.
 type StdioServer struct {
-	server    *MCPServer
-	errLogger *log.Logger
+	server      *MCPServer
+	errLogger   *log.Logger
+	contextFunc StdioContextFunc
 }
 
 // NewStdioServer creates a new stdio server wrapper around an MCPServer.
@@ -41,6 +48,13 @@ func (s *StdioServer) SetErrorLogger(logger *log.Logger) {
 	s.errLogger = logger
 }
 
+// SetContextFunc sets a function that will be called to customise the context
+// to the server. Note that the stdio server uses the same context for all requests,
+// so this function will only be called once per server instance.
+func (s *StdioServer) SetContextFunc(fn StdioContextFunc) {
+	s.contextFunc = fn
+}
+
 // Listen starts listening for JSON-RPC messages on the provided input and writes responses to the provided output.
 // It runs until the context is cancelled or an error occurs.
 // Returns an error if there are issues with reading input or writing output.
@@ -55,6 +69,11 @@ func (s *StdioServer) Listen(
 		SessionID: "stdio",
 	})
 
+	// Add in any custom context.
+	if s.contextFunc != nil {
+		ctx = s.contextFunc(ctx)
+	}
+
 	reader := bufio.NewReader(stdin)
 
 	// Start notification handler