Repository: mark3labs/mcp-go

Issue:
Allow customising server/message context
It would be useful to be able to manipulate the context used when handling messages, to e.g. extract config in a transport-specific way and inject it into context so that tools can access it. This can be used to e.g. update settings from environment variables or headers depending on the transport (stdio vs sse).

I have a [commit](https://github.com/grafana/mcp-go/commit/5e1fa6f6f8afc981fb95f0cc5a42d376456501f8) which does something like this and works fine, but I'm unsure on the naming. Happy to create a PR if you'd accept it though.

Discussion:
@ezynda3: Would definitely welcome a PR for this

Solution patch:
diff --git a/examples/custom_context/main.go b/examples/custom_context/main.go
new file mode 100644
index 0000000..1eef274
--- /dev/null
+++ b/examples/custom_context/main.go
@@ -0,0 +1,164 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"os"
+
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/server"
+)
+
+// authKey is a custom context key for storing the auth token.
+type authKey struct{}
+
+// withAuthKey adds an auth key to the context.
+func withAuthKey(ctx context.Context, auth string) context.Context {
+	return context.WithValue(ctx, authKey{}, auth)
+}
+
+// authFromRequest extracts the auth token from the request headers.
+func authFromRequest(ctx context.Context, r *http.Request) context.Context {
+	return withAuthKey(ctx, r.Header.Get("Authorization"))
+}
+
+// authFromEnv extracts the auth token from the environment
+func authFromEnv(ctx context.Context) context.Context {
+	return withAuthKey(ctx, os.Getenv("API_KEY"))
+}
+
+// tokenFromContext extracts the auth token from the context.
+// This can be used by tools to extract the token regardless of the
+// transport being used by the server.
+func tokenFromContext(ctx context.Context) (string, error) {
+	auth, ok := ctx.Value(authKey{}).(string)
+	if !ok {
+		return "", fmt.Errorf("missing auth")
+	}
+	return auth, nil
+}
+
+type response struct {
+	Args    map[string]interface{} `json:"args"`
+	Headers map[string]string      `json:"headers"`
+}
+
+// makeRequest makes a request to httpbin.org including the auth token in the request
+// headers and the message in the query string.
+func makeRequest(ctx context.Context, message, token string) (*response, error) {
+	req, err := http.NewRequestWithContext(ctx, "GET", "https://httpbin.org/anything", nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Authorization", token)
+	req.URL.Query().Add("message", message)
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+	var r *response
+	if err := json.Unmarshal(body, r); err != nil {
+		return nil, err
+	}
+	return r, nil
+}
+
+// handleMakeAuthenticatedRequestTool is a tool that makes an authenticated request
+// using the token from the context.
+func handleMakeAuthenticatedRequestTool(
+	ctx context.Context,
+	request mcp.CallToolRequest,
+) (*mcp.CallToolResult, error) {
+	message, ok := request.Params.Arguments["message"].(string)
+	if !ok {
+		return nil, fmt.Errorf("missing message")
+	}
+	token, err := tokenFromContext(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("missing token: %v", err)
+	}
+	// Now our tool can make a request with the token, irrespective of where it came from.
+	resp, err := makeRequest(ctx, message, token)
+	if err != nil {
+		return nil, err
+	}
+	return mcp.NewToolResultText(fmt.Sprintf("%+v", resp)), nil
+}
+
+type MCPServer struct {
+	server *server.MCPServer
+}
+
+func NewMCPServer() *MCPServer {
+	mcpServer := server.NewMCPServer(
+		"example-server",
+		"1.0.0",
+		server.WithResourceCapabilities(true, true),
+		server.WithPromptCapabilities(true),
+		server.WithToolCapabilities(true),
+	)
+	mcpServer.AddTool(mcp.NewTool("make_authenticated_request",
+		mcp.WithDescription("Makes an authenticated request"),
+		mcp.WithString("message",
+			mcp.Description("Message to echo"),
+			mcp.Required(),
+		),
+	), handleMakeAuthenticatedRequestTool)
+
+	return &MCPServer{
+		server: mcpServer,
+	}
+}
+
+func (s *MCPServer) ServeSSE(addr string) *server.SSEServer {
+	return server.NewSSEServer(s.server,
+		server.WithBaseURL(fmt.Sprintf("http://%s", addr)),
+		server.WithSSEContextFunc(authFromRequest),
+	)
+}
+
+func (s *MCPServer) ServeStdio() error {
+	return server.ServeStdio(s.server, server.WithStdioContextFunc(authFromEnv))
+}
+
+func main() {
+	var transport string
+	flag.StringVar(&transport, "t", "stdio", "Transport type (stdio or sse)")
+	flag.StringVar(
+		&transport,
+		"transport",
+		"stdio",
+		"Transport type (stdio or sse)",
+	)
+	flag.Parse()
+
+	s := NewMCPServer()
+
+	switch transport {
+	case "stdio":
+		if err := s.ServeStdio(); err != nil {
+			log.Fatalf("Server error: %v", err)
+		}
+	case "sse":
+		sseServer := s.ServeSSE("localhost:8080")
+		log.Printf("SSE server listening on :8080")
+		if err := sseServer.Start(":8080"); err != nil {
+			log.Fatalf("Server error: %v", err)
+		}
+	default:
+		log.Fatalf(
+			"Invalid transport type: %s. Must be 'stdio' or 'sse'",
+			transport,
+		)
+	}
+}
diff --git a/server/sse.go b/server/sse.go
index 5a28aaa..0e2f8b8 100644
--- a/server/sse.go
+++ b/server/sse.go
@@ -21,6 +21,11 @@ type sseSession struct {
 	eventQueue chan string // Channel for queuing events
 }
 
+// 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 {
@@ -31,20 +36,21 @@ type SSEServer struct {
 	sseEndpoint     string
 	sessions        sync.Map
 	srv             *http.Server
+	contextFunc     SSEContextFunc
 }
 
-// Option defines a function type for configuring SSEServer
-type Option func(*SSEServer)
+// SSEOption defines a function type for configuring SSEServer
+type SSEOption func(*SSEServer)
 
 // WithBaseURL sets the base URL for the SSE server
-func WithBaseURL(baseURL string) Option {
+func WithBaseURL(baseURL string) SSEOption {
 	return func(s *SSEServer) {
 		s.baseURL = baseURL
 	}
 }
 
 // Add a new option for setting base path
-func WithBasePath(basePath string) Option {
+func WithBasePath(basePath string) SSEOption {
 	return func(s *SSEServer) {
 		// Ensure the path starts with / and doesn't end with /
 		if !strings.HasPrefix(basePath, "/") {
@@ -56,28 +62,36 @@ func WithBasePath(basePath string) Option {
 }
 
 // WithMessageEndpoint sets the message endpoint path
-func WithMessageEndpoint(endpoint string) Option {
+func WithMessageEndpoint(endpoint string) SSEOption {
 	return func(s *SSEServer) {
 		s.messageEndpoint = endpoint
 	}
 }
 
 // WithSSEEndpoint sets the SSE endpoint path
-func WithSSEEndpoint(endpoint string) Option {
+func WithSSEEndpoint(endpoint string) SSEOption {
 	return func(s *SSEServer) {
 		s.sseEndpoint = endpoint
 	}
 }
 
 // WithHTTPServer sets the HTTP server instance
-func WithHTTPServer(srv *http.Server) Option {
+func WithHTTPServer(srv *http.Server) SSEOption {
 	return func(s *SSEServer) {
 		s.srv = srv
 	}
 }
 
+// WithContextFunc sets a function that will be called to customise the context
+// to the server using the incoming request.
+func WithSSEContextFunc(fn SSEContextFunc) SSEOption {
+	return func(s *SSEServer) {
+		s.contextFunc = fn
+	}
+}
+
 // NewSSEServer creates a new SSE server instance with the given MCP server and options.
-func NewSSEServer(server *MCPServer, opts ...Option) *SSEServer {
+func NewSSEServer(server *MCPServer, opts ...SSEOption) *SSEServer {
 	s := &SSEServer{
 		server:          server,
 		sseEndpoint:     "/sse",
@@ -94,8 +108,11 @@ func NewSSEServer(server *MCPServer, opts ...Option) *SSEServer {
 }
 
 // NewTestServer creates a test server for testing purposes
-func NewTestServer(server *MCPServer) *httptest.Server {
+func NewTestServer(server *MCPServer, opts ...SSEOption) *httptest.Server {
 	sseServer := NewSSEServer(server)
+	for _, opt := range opts {
+		opt(sseServer)
+	}
 
 	testServer := httptest.NewServer(sseServer)
 	sseServer.baseURL = testServer.URL
@@ -230,6 +247,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")
diff --git a/server/stdio.go b/server/stdio.go
index 5de08c2..5f9221c 100644
--- a/server/stdio.go
+++ b/server/stdio.go
@@ -14,12 +14,38 @@ 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
+}
+
+// StdioOption defines a function type for configuring StdioServer
+type StdioOption func(*StdioServer)
+
+// WithErrorLogger sets the error logger for the server
+func WithErrorLogger(logger *log.Logger) StdioOption {
+	return func(s *StdioServer) {
+		s.errLogger = logger
+	}
+}
+
+// WithContextFunc 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 WithStdioContextFunc(fn StdioContextFunc) StdioOption {
+	return func(s *StdioServer) {
+		s.contextFunc = fn
+	}
 }
 
 // NewStdioServer creates a new stdio server wrapper around an MCPServer.
@@ -41,6 +67,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 +88,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
@@ -171,10 +209,14 @@ func (s *StdioServer) writeResponse(
 // ServeStdio is a convenience function that creates and starts a StdioServer with os.Stdin and os.Stdout.
 // It sets up signal handling for graceful shutdown on SIGTERM and SIGINT.
 // Returns an error if the server encounters any issues during operation.
-func ServeStdio(server *MCPServer) error {
+func ServeStdio(server *MCPServer, opts ...StdioOption) error {
 	s := NewStdioServer(server)
 	s.SetErrorLogger(log.New(os.Stderr, "", log.LstdFlags))
 
+	for _, opt := range opts {
+		opt(s)
+	}
+
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 



Test patch:
diff --git a/server/sse_test.go b/server/sse_test.go
index bd88cac..7ce113e 100644
--- a/server/sse_test.go
+++ b/server/sse_test.go
@@ -11,6 +11,8 @@ import (
 	"sync"
 	"testing"
 	"time"
+
+	"github.com/mark3labs/mcp-go/mcp"
 )
 
 func TestSSEServer(t *testing.T) {
@@ -468,4 +470,148 @@ func TestSSEServer(t *testing.T) {
 		cancel()
 		<-done
 	})
+
+	t.Run("Can use a custom context function", func(t *testing.T) {
+		// Use a custom context key to store a test value.
+		type testContextKey struct{}
+		testValFromContext := func(ctx context.Context) string {
+			val := ctx.Value(testContextKey{})
+			if val == nil {
+				return ""
+			}
+			return val.(string)
+		}
+		// Create a context function that sets a test value from the request.
+		// In real life this could be used to send configuration using headers
+		// or query parameters.
+		const testHeader = "X-Test-Header"
+		setTestValFromRequest := func(ctx context.Context, r *http.Request) context.Context {
+			return context.WithValue(ctx, testContextKey{}, r.Header.Get(testHeader))
+		}
+
+		mcpServer := NewMCPServer("test", "1.0.0",
+			WithResourceCapabilities(true, true),
+		)
+		// Add a tool which uses the context function.
+		mcpServer.AddTool(mcp.NewTool("test_tool"), func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			// Note this is agnostic to the transport type i.e. doesn't know about request headers.
+			testVal := testValFromContext(ctx)
+			return mcp.NewToolResultText(testVal), nil
+		})
+
+		testServer := NewTestServer(mcpServer, WithSSEContextFunc(setTestValFromRequest))
+		defer testServer.Close()
+
+		// Connect to SSE endpoint
+		sseResp, err := http.Get(fmt.Sprintf("%s/sse", testServer.URL))
+		if err != nil {
+			t.Fatalf("Failed to connect to SSE endpoint: %v", err)
+		}
+		defer sseResp.Body.Close()
+
+		// Read the endpoint event
+		buf := make([]byte, 1024)
+		n, err := sseResp.Body.Read(buf)
+		if err != nil {
+			t.Fatalf("Failed to read SSE response: %v", err)
+		}
+
+		endpointEvent := string(buf[:n])
+		messageURL := strings.TrimSpace(
+			strings.Split(strings.Split(endpointEvent, "data: ")[1], "\n")[0],
+		)
+
+		// Send initialize request
+		initRequest := map[string]interface{}{
+			"jsonrpc": "2.0",
+			"id":      1,
+			"method":  "initialize",
+			"params": map[string]interface{}{
+				"protocolVersion": "2024-11-05",
+				"clientInfo": map[string]interface{}{
+					"name":    "test-client",
+					"version": "1.0.0",
+				},
+			},
+		}
+
+		requestBody, err := json.Marshal(initRequest)
+		if err != nil {
+			t.Fatalf("Failed to marshal request: %v", err)
+		}
+
+		resp, err := http.Post(
+			messageURL,
+			"application/json",
+			bytes.NewBuffer(requestBody),
+		)
+
+		if err != nil {
+			t.Fatalf("Failed to send message: %v", err)
+		}
+		defer resp.Body.Close()
+
+		if resp.StatusCode != http.StatusAccepted {
+			t.Errorf("Expected status 202, got %d", resp.StatusCode)
+		}
+
+		// Verify response
+		var response map[string]interface{}
+		if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
+			t.Fatalf("Failed to decode response: %v", err)
+		}
+
+		if response["jsonrpc"] != "2.0" {
+			t.Errorf("Expected jsonrpc 2.0, got %v", response["jsonrpc"])
+		}
+		if response["id"].(float64) != 1 {
+			t.Errorf("Expected id 1, got %v", response["id"])
+		}
+
+		// Call the tool.
+		toolRequest := map[string]interface{}{
+			"jsonrpc": "2.0",
+			"id":      2,
+			"method":  "tools/call",
+			"params": map[string]interface{}{
+				"name": "test_tool",
+			},
+		}
+		requestBody, err = json.Marshal(toolRequest)
+		if err != nil {
+			t.Fatalf("Failed to marshal tool request: %v", err)
+		}
+
+		req, err := http.NewRequest(http.MethodPost, messageURL, bytes.NewBuffer(requestBody))
+		if err != nil {
+			t.Fatalf("Failed to create tool request: %v", err)
+		}
+		// Set the test header to a custom value.
+		req.Header.Set(testHeader, "test_value")
+
+		resp, err = http.DefaultClient.Do(req)
+		if err != nil {
+			t.Fatalf("Failed to call tool: %v", err)
+		}
+		defer resp.Body.Close()
+
+		response = make(map[string]interface{})
+		if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
+			t.Fatalf("Failed to decode response: %v", err)
+		}
+
+		if response["jsonrpc"] != "2.0" {
+			t.Errorf("Expected jsonrpc 2.0, got %v", response["jsonrpc"])
+		}
+		if response["id"].(float64) != 2 {
+			t.Errorf("Expected id 2, got %v", response["id"])
+		}
+		if response["result"].(map[string]interface{})["content"].([]interface{})[0].(map[string]interface{})["text"] != "test_value" {
+			t.Errorf("Expected result 'test_value', got %v", response["result"])
+		}
+		if response["error"] != nil {
+			t.Errorf("Expected no error, got %v", response["error"])
+		}
+	})
+
 }
diff --git a/server/stdio_test.go b/server/stdio_test.go
index 8de840a..6113174 100644
--- a/server/stdio_test.go
+++ b/server/stdio_test.go
@@ -6,7 +6,10 @@ import (
 	"encoding/json"
 	"io"
 	"log"
+	"os"
 	"testing"
+
+	"github.com/mark3labs/mcp-go/mcp"
 )
 
 func TestStdioServer(t *testing.T) {
@@ -110,4 +113,158 @@ func TestStdioServer(t *testing.T) {
 			t.Errorf("unexpected server error: %v", err)
 		}
 	})
+
+	t.Run("Can use a custom context function", func(t *testing.T) {
+		// Use a custom context key to store a test value.
+		type testContextKey struct{}
+		testValFromContext := func(ctx context.Context) string {
+			val := ctx.Value(testContextKey{})
+			if val == nil {
+				return ""
+			}
+			return val.(string)
+		}
+		// Create a context function that sets a test value from the environment.
+		// In real life this could be used to send configuration in a similar way,
+		// or from a config file.
+		const testEnvVar = "TEST_ENV_VAR"
+		setTestValFromEnv := func(ctx context.Context) context.Context {
+			return context.WithValue(ctx, testContextKey{}, os.Getenv(testEnvVar))
+		}
+		t.Setenv(testEnvVar, "test_value")
+
+		// Create pipes for stdin and stdout
+		stdinReader, stdinWriter := io.Pipe()
+		stdoutReader, stdoutWriter := io.Pipe()
+
+		// Create server
+		mcpServer := NewMCPServer("test", "1.0.0")
+		// Add a tool which uses the context function.
+		mcpServer.AddTool(mcp.NewTool("test_tool"), func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			// Note this is agnostic to the transport type i.e. doesn't know about request headers.
+			testVal := testValFromContext(ctx)
+			return mcp.NewToolResultText(testVal), nil
+		})
+		stdioServer := NewStdioServer(mcpServer)
+		stdioServer.SetErrorLogger(log.New(io.Discard, "", 0))
+		stdioServer.SetContextFunc(setTestValFromEnv)
+
+		// Create context with cancel
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		// Create error channel to catch server errors
+		serverErrCh := make(chan error, 1)
+
+		// Start server in goroutine
+		go func() {
+			err := stdioServer.Listen(ctx, stdinReader, stdoutWriter)
+			if err != nil && err != io.EOF && err != context.Canceled {
+				serverErrCh <- err
+			}
+			close(serverErrCh)
+		}()
+
+		// Create test message
+		initRequest := map[string]interface{}{
+			"jsonrpc": "2.0",
+			"id":      1,
+			"method":  "initialize",
+			"params": map[string]interface{}{
+				"protocolVersion": "2024-11-05",
+				"clientInfo": map[string]interface{}{
+					"name":    "test-client",
+					"version": "1.0.0",
+				},
+			},
+		}
+
+		// Send request
+		requestBytes, err := json.Marshal(initRequest)
+		if err != nil {
+			t.Fatal(err)
+		}
+		_, err = stdinWriter.Write(append(requestBytes, '\n'))
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Read response
+		scanner := bufio.NewScanner(stdoutReader)
+		if !scanner.Scan() {
+			t.Fatal("failed to read response")
+		}
+		responseBytes := scanner.Bytes()
+
+		var response map[string]interface{}
+		if err := json.Unmarshal(responseBytes, &response); err != nil {
+			t.Fatalf("failed to unmarshal response: %v", err)
+		}
+
+		// Verify response structure
+		if response["jsonrpc"] != "2.0" {
+			t.Errorf("expected jsonrpc version 2.0, got %v", response["jsonrpc"])
+		}
+		if response["id"].(float64) != 1 {
+			t.Errorf("expected id 1, got %v", response["id"])
+		}
+		if response["error"] != nil {
+			t.Errorf("unexpected error in response: %v", response["error"])
+		}
+		if response["result"] == nil {
+			t.Error("expected result in response")
+		}
+
+		// Call the tool.
+		toolRequest := map[string]interface{}{
+			"jsonrpc": "2.0",
+			"id":      2,
+			"method":  "tools/call",
+			"params": map[string]interface{}{
+				"name": "test_tool",
+			},
+		}
+		requestBytes, err = json.Marshal(toolRequest)
+		if err != nil {
+			t.Fatalf("Failed to marshal tool request: %v", err)
+		}
+
+		_, err = stdinWriter.Write(append(requestBytes, '\n'))
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if !scanner.Scan() {
+			t.Fatal("failed to read response")
+		}
+		responseBytes = scanner.Bytes()
+
+		response = map[string]interface{}{}
+		if err := json.Unmarshal(responseBytes, &response); err != nil {
+			t.Fatalf("failed to unmarshal response: %v", err)
+		}
+
+		if response["jsonrpc"] != "2.0" {
+			t.Errorf("Expected jsonrpc 2.0, got %v", response["jsonrpc"])
+		}
+		if response["id"].(float64) != 2 {
+			t.Errorf("Expected id 2, got %v", response["id"])
+		}
+		if response["result"].(map[string]interface{})["content"].([]interface{})[0].(map[string]interface{})["text"] != "test_value" {
+			t.Errorf("Expected result 'test_value', got %v", response["result"])
+		}
+		if response["error"] != nil {
+			t.Errorf("Expected no error, got %v", response["error"])
+		}
+
+		// Clean up
+		cancel()
+		stdinWriter.Close()
+		stdoutWriter.Close()
+
+		// Check for server errors
+		if err := <-serverErrCh; err != nil {
+			t.Errorf("unexpected server error: %v", err)
+		}
+	})
 }