From 45ae3c9c09532b053766cdcda6710161d4d79b48 Mon Sep 17 00:00:00 2001 From: Adam Jamiu Date: Sun, 18 Oct 2026 20:38:54 +0100 Subject: [PATCH 1/2] mcp: reuse server selection for new streamable requests Reuse the server selected for protocol-version validation when creating a new stateful or stateless session. This prevents duplicate getServer callback invocations for a single POST. Fixes #1413 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mcp/streamable.go | 19 +++++++++---------- mcp/streamable_test.go | 22 ++++++++++++++++++++++ 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/mcp/streamable.go b/mcp/streamable.go index 3087c46e5..544969651 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -409,8 +409,10 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque // // [§2.7]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#protocol-version-header protocolVersion := req.Header.Get(protocolVersionHeader) + req = req.WithContext(context.WithValue(req.Context(), protocolVersionContextKey{}, protocolVersion)) var supportedVersions []string - if server := h.getServer(req); server != nil { + server := h.getServer(req) + if server != nil { supportedVersions = server.protocolVersions } else { supportedVersions = supportedProtocolVersions @@ -419,12 +421,11 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque http.Error(w, fmt.Sprintf("Bad Request: Unsupported protocol version (supported versions: %s)", strings.Join(supportedVersions, ",")), http.StatusBadRequest) return } - req = req.WithContext(context.WithValue(req.Context(), protocolVersionContextKey{}, protocolVersion)) if h.opts.Stateless { - h.serveStateless(w, req) + h.serveStateless(w, req, server) } else { - h.serveStateful(w, req) + h.serveStateful(w, req, server) } } @@ -434,7 +435,7 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque // // When the allowsessionsinstateless compatibility flag is set, DELETE is also // accepted (as a no-op) and session IDs are read from the request header. -func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http.Request) { +func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http.Request, server *Server) { legacySessions := allowsessionsinstateless == "1" if req.Method == http.MethodDelete && legacySessions { @@ -461,7 +462,6 @@ func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http. return } - server := h.getServer(req) if server == nil { http.Error(w, "no server available", http.StatusBadRequest) return @@ -611,12 +611,12 @@ func (h *StreamableHTTPHandler) newStreamableServerTransport(sessionID string, s // serveStateful handles requests for stateful servers. // Stateful servers support GET, POST, and DELETE, and maintain persistent // sessions keyed by session ID. -func (h *StreamableHTTPHandler) serveStateful(w http.ResponseWriter, req *http.Request) { +func (h *StreamableHTTPHandler) serveStateful(w http.ResponseWriter, req *http.Request, server *Server) { switch req.Method { case http.MethodGet: h.serveStatefulGET(w, req) case http.MethodPost: - h.serveStatefulPOST(w, req) + h.serveStatefulPOST(w, req, server) case http.MethodDelete: h.serveStatefulDELETE(w, req) default: @@ -692,7 +692,7 @@ func (h *StreamableHTTPHandler) serveStatefulDELETE(w http.ResponseWriter, req * // POST may arrive with or without a Mcp-Session-Id header. Without a session // ID, a new session is created (this is the normal path for the first // initialize request). -func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *http.Request) { +func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *http.Request, server *Server) { if baseMediaType(req.Header.Get("Content-Type")) != "application/json" { http.Error(w, "Content-Type must be 'application/json'", http.StatusUnsupportedMediaType) return @@ -719,7 +719,6 @@ func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *ht } // No session ID: create a new session. - server := h.getServer(req) if server == nil { http.Error(w, "no server available", http.StatusBadRequest) return diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index aa19aa3ac..8904be433 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -4973,3 +4973,25 @@ func TestStreamableSupportedProtocolVersions_Header(t *testing.T) { } }) } + +func TestStreamableGetServerCalledOnceForNewPOST(t *testing.T) { + for _, stateless := range []bool{false, true} { + t.Run(fmt.Sprintf("stateless=%v", stateless), func(t *testing.T) { + server := NewServer(testImpl, nil) + var calls atomic.Int32 + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { + calls.Add(1) + return server + }, &StreamableHTTPOptions{Stateless: stateless}) + + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(`{}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + handler.ServeHTTP(httptest.NewRecorder(), req) + + if got := calls.Load(); got != 1 { + t.Errorf("getServer called %d times, want 1", got) + } + }) + } +} From 5f175f587031f62458e35345e9e40d8bae5ac13e Mon Sep 17 00:00:00 2001 From: Adam Jamiu Date: Mon, 19 Oct 2026 18:30:52 +0100 Subject: [PATCH 2/2] mcp: resolve streamable server per request path Validate protocol versions using the server selected for new POSTs, or the server associated with an existing session. Avoid invoking getServer for requests routed to established sessions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mcp/streamable.go | 47 ++++++++++++++++++++++++++++-------------- mcp/streamable_test.go | 40 +++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 15 deletions(-) diff --git a/mcp/streamable.go b/mcp/streamable.go index 544969651..819bc08f4 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -410,23 +410,23 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque // [§2.7]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#protocol-version-header protocolVersion := req.Header.Get(protocolVersionHeader) req = req.WithContext(context.WithValue(req.Context(), protocolVersionContextKey{}, protocolVersion)) - var supportedVersions []string - server := h.getServer(req) + if h.opts.Stateless { + h.serveStateless(w, req) + } else { + h.serveStateful(w, req) + } +} + +func (h *StreamableHTTPHandler) validateProtocolVersion(w http.ResponseWriter, protocolVersion string, server *Server) bool { + supportedVersions := supportedProtocolVersions if server != nil { supportedVersions = server.protocolVersions - } else { - supportedVersions = supportedProtocolVersions } if protocolVersion != "" && !slices.Contains(supportedVersions, protocolVersion) && protocolVersion < protocolVersion20260728 { http.Error(w, fmt.Sprintf("Bad Request: Unsupported protocol version (supported versions: %s)", strings.Join(supportedVersions, ",")), http.StatusBadRequest) - return - } - - if h.opts.Stateless { - h.serveStateless(w, req, server) - } else { - h.serveStateful(w, req, server) + return false } + return true } // serveStateless handles requests for stateless servers. @@ -435,7 +435,7 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque // // When the allowsessionsinstateless compatibility flag is set, DELETE is also // accepted (as a no-op) and session IDs are read from the request header. -func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http.Request, server *Server) { +func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http.Request) { legacySessions := allowsessionsinstateless == "1" if req.Method == http.MethodDelete && legacySessions { @@ -462,6 +462,10 @@ func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http. return } + server := h.getServer(req) + if !h.validateProtocolVersion(w, protocolVersionFromContext(req.Context()), server) { + return + } if server == nil { http.Error(w, "no server available", http.StatusBadRequest) return @@ -611,12 +615,12 @@ func (h *StreamableHTTPHandler) newStreamableServerTransport(sessionID string, s // serveStateful handles requests for stateful servers. // Stateful servers support GET, POST, and DELETE, and maintain persistent // sessions keyed by session ID. -func (h *StreamableHTTPHandler) serveStateful(w http.ResponseWriter, req *http.Request, server *Server) { +func (h *StreamableHTTPHandler) serveStateful(w http.ResponseWriter, req *http.Request) { switch req.Method { case http.MethodGet: h.serveStatefulGET(w, req) case http.MethodPost: - h.serveStatefulPOST(w, req, server) + h.serveStatefulPOST(w, req) case http.MethodDelete: h.serveStatefulDELETE(w, req) default: @@ -666,6 +670,9 @@ func (h *StreamableHTTPHandler) serveStatefulGET(w http.ResponseWriter, req *htt if !ok { return } + if !h.validateProtocolVersion(w, protocolVersionFromContext(req.Context()), sessInfo.session.server) { + return + } sessInfo.transport.ServeHTTP(w, req) } @@ -683,6 +690,9 @@ func (h *StreamableHTTPHandler) serveStatefulDELETE(w http.ResponseWriter, req * if !ok { return } + if !h.validateProtocolVersion(w, protocolVersionFromContext(req.Context()), sessInfo.session.server) { + return + } sessInfo.session.Close() w.WriteHeader(http.StatusNoContent) @@ -692,7 +702,7 @@ func (h *StreamableHTTPHandler) serveStatefulDELETE(w http.ResponseWriter, req * // POST may arrive with or without a Mcp-Session-Id header. Without a session // ID, a new session is created (this is the normal path for the first // initialize request). -func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *http.Request, server *Server) { +func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *http.Request) { if baseMediaType(req.Header.Get("Content-Type")) != "application/json" { http.Error(w, "Content-Type must be 'application/json'", http.StatusUnsupportedMediaType) return @@ -712,6 +722,9 @@ func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *ht if !ok { return } + if !h.validateProtocolVersion(w, protocolVersionFromContext(req.Context()), sessInfo.session.server) { + return + } sessInfo.startPOST() defer sessInfo.endPOST() sessInfo.transport.ServeHTTP(w, req) @@ -719,6 +732,10 @@ func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *ht } // No session ID: create a new session. + server := h.getServer(req) + if !h.validateProtocolVersion(w, protocolVersionFromContext(req.Context()), server) { + return + } if server == nil { http.Error(w, "no server available", http.StatusBadRequest) return diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 8904be433..19f475e15 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -4994,4 +4994,44 @@ func TestStreamableGetServerCalledOnceForNewPOST(t *testing.T) { } }) } + + t.Run("existing stateful session", func(t *testing.T) { + server := NewServer(testImpl, nil) + var calls atomic.Int32 + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { + calls.Add(1) + return server + }, &StreamableHTTPOptions{JSONResponse: true}) + + post := func(body string, sessionID string) *httptest.ResponseRecorder { + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + if sessionID != "" { + req.Header.Set(sessionIDHeader, sessionID) + } + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + return rec + } + + initBody := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}` + initRec := post(initBody, "") + if initRec.Code != http.StatusOK { + t.Fatalf("initialize status = %d, want %d: %s", initRec.Code, http.StatusOK, initRec.Body) + } + sessionID := initRec.Header().Get(sessionIDHeader) + if sessionID == "" { + t.Fatal("initialize response missing session ID") + } + + pingRec := post(`{"jsonrpc":"2.0","id":2,"method":"ping"}`, sessionID) + if pingRec.Code != http.StatusOK { + t.Fatalf("ping status = %d, want %d: %s", pingRec.Code, http.StatusOK, pingRec.Body) + } + if got := calls.Load(); got != 1 { + t.Errorf("getServer called %d times, want 1 for new session only", got) + } + handler.closeAll() + }) }