From fd3ce0d810e52a8fc533c31c534f81dc41318ea6 Mon Sep 17 00:00:00 2001
From: Hongxiang Jiang <hxjiang@golang.org>
Date: Mon, 28 Apr 2025 10:49:40 -0400
Subject: [PATCH] internal/mcp: dynamically determine the mcp server based on
 URL

Based on my understanding, the getServer function is used to determine
which MCP server should be serving the session.

User might want to dynamically determine which MCP server to serve when
a new request comes in. In mcp/examples, two MCP servers are created
serving at the same port but under different URL "/greeter1",
"/greeter2". The MCP client will reach out to different URLs for
different MCP servers.

As a TODO, if the getServer() returns nil, the program will crash when
handing MCP request.

Change-Id: I4431266666e9ea28118e827a94eb1359b1fb3a81
Reviewed-on: https://go-review.googlesource.com/c/tools/+/668535
Reviewed-by: Robert Findley <rfindley@google.com>
Auto-Submit: Hongxiang Jiang <hxjiang@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
---
 examples/hello/main.go |  2 +-
 examples/sse/main.go   | 55 ++++++++++++++++++++++++++++++++++++++++++
 mcp/sse.go             |  9 ++++---
 mcp/sse_test.go        |  3 ++-
 4 files changed, 63 insertions(+), 6 deletions(-)
 create mode 100644 examples/sse/main.go

diff --git a/examples/hello/main.go b/examples/hello/main.go
index d661bc9..7ea8353 100644
--- a/examples/hello/main.go
+++ b/examples/hello/main.go
@@ -33,7 +33,7 @@ func main() {
 	server.AddTools(mcp.MakeTool("greet", "say hi", SayHi))
 
 	if *httpAddr != "" {
-		handler := mcp.NewSSEHandler(func() *mcp.Server {
+		handler := mcp.NewSSEHandler(func(*http.Request) *mcp.Server {
 			return server
 		})
 		http.ListenAndServe(*httpAddr, handler)
diff --git a/examples/sse/main.go b/examples/sse/main.go
new file mode 100644
index 0000000..3bd75c3
--- /dev/null
+++ b/examples/sse/main.go
@@ -0,0 +1,55 @@
+// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+package main
+
+import (
+	"context"
+	"flag"
+	"log"
+	"net/http"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+)
+
+var httpAddr = flag.String("http", "", "use SSE HTTP at this address")
+
+type SayHiParams struct {
+	Name string `json:"name" mcp:"the name to say hi to"`
+}
+
+func SayHi(ctx context.Context, cc *mcp.ClientConnection, params *SayHiParams) ([]mcp.Content, error) {
+	return []mcp.Content{
+		mcp.TextContent{Text: "Hi " + params.Name},
+	}, nil
+}
+
+func main() {
+	flag.Parse()
+
+	if httpAddr == nil || *httpAddr == "" {
+		log.Fatal("http address not set")
+	}
+
+	server1 := mcp.NewServer("greeter1", "v0.0.1", nil)
+	server1.AddTools(mcp.MakeTool("greet1", "say hi", SayHi))
+
+	server2 := mcp.NewServer("greeter2", "v0.0.1", nil)
+	server2.AddTools(mcp.MakeTool("greet2", "say hello", SayHi))
+
+	log.Printf("MCP servers serving at %s\n", *httpAddr)
+	handler := mcp.NewSSEHandler(func(request *http.Request) *mcp.Server {
+		url := request.URL.Path
+		log.Printf("Handling request for URL %s\n", url)
+		switch url {
+		case "/greeter1":
+			return server1
+		case "/greeter2":
+			return server2
+		default:
+			return nil
+		}
+	})
+	http.ListenAndServe(*httpAddr, handler)
+}
diff --git a/mcp/sse.go b/mcp/sse.go
index 36cfffd..c900012 100644
--- a/mcp/sse.go
+++ b/mcp/sse.go
@@ -72,7 +72,7 @@ func writeEvent(w io.Writer, evt event) (int, error) {
 //
 // https://modelcontextprotocol.io/specification/2024-11-05/basic/transports
 type SSEHandler struct {
-	getServer func() *Server
+	getServer func(request *http.Request) *Server
 	onClient  func(*ClientConnection) // for testing; must not block
 
 	mu       sync.Mutex
@@ -81,9 +81,9 @@ type SSEHandler struct {
 
 // NewSSEHandler returns a new [SSEHandler] that is ready to serve HTTP.
 //
-// The getServer function is used to bind create servers for new sessions. It
+// The getServer function is used to bind created servers for new sessions. It
 // is OK for getServer to return the same server multiple times.
-func NewSSEHandler(getServer func() *Server) *SSEHandler {
+func NewSSEHandler(getServer func(request *http.Request) *Server) *SSEHandler {
 	return &SSEHandler{
 		getServer: getServer,
 		sessions:  make(map[string]*sseSession),
@@ -179,7 +179,8 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	}
 	defer session.Close()
 
-	server := h.getServer()
+	// TODO(hxjiang): getServer returns nil will panic.
+	server := h.getServer(req)
 	cc, err := server.Connect(req.Context(), session, nil)
 	if err != nil {
 		http.Error(w, "connection failed", http.StatusInternalServerError)
diff --git a/mcp/sse_test.go b/mcp/sse_test.go
index 13196ab..a814afb 100644
--- a/mcp/sse_test.go
+++ b/mcp/sse_test.go
@@ -7,6 +7,7 @@ package mcp
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"net/http/httptest"
 	"testing"
 
@@ -20,7 +21,7 @@ func TestSSEServer(t *testing.T) {
 			server := NewServer("testServer", "v1.0.0", nil)
 			server.AddTools(MakeTool("greet", "say hi", sayHi))
 
-			sseHandler := NewSSEHandler(func() *Server { return server })
+			sseHandler := NewSSEHandler(func(*http.Request) *Server { return server })
 
 			clients := make(chan *ClientConnection, 1)
 			sseHandler.onClient = func(cc *ClientConnection) {