diff --git a/mcp/streamable.go b/mcp/streamable.go index 3087c46e..819bc08f 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -409,18 +409,7 @@ 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) - var supportedVersions []string - if server := h.getServer(req); 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 - } req = req.WithContext(context.WithValue(req.Context(), protocolVersionContextKey{}, protocolVersion)) - if h.opts.Stateless { h.serveStateless(w, req) } else { @@ -428,6 +417,18 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque } } +func (h *StreamableHTTPHandler) validateProtocolVersion(w http.ResponseWriter, protocolVersion string, server *Server) bool { + supportedVersions := supportedProtocolVersions + if server != nil { + supportedVersions = server.protocolVersions + } + 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 false + } + return true +} + // serveStateless handles requests for stateless servers. // Stateless servers only support POST. Each request creates a temporary // session that is closed when the request completes. @@ -462,6 +463,9 @@ func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http. } 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 @@ -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) @@ -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) @@ -720,6 +733,9 @@ 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 aa19aa3a..19f475e1 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -4973,3 +4973,65 @@ 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) + } + }) + } + + 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() + }) +}