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>
This commit is contained in:
Hongxiang Jiang 2025-04-28 10:49:40 -04:00 • committed by Gopher Robot
commit fd3ce0d810
2 changed files with 63 additions and 6 deletions

View file

@ -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)

View file

@ -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) {