From 21f141f9b525c099e5c9ce90b551383eb0716de3 Mon Sep 17 00:00:00 2001 From: Nishant Bangarwa Date: Tue, 20 Oct 2026 11:46:45 +0530 Subject: [PATCH 1/2] fix: MCP requests no longer share an AI session after the go-sdk bump to v1.8.0 Since go-sdk v1.8.0, a stateless streamable HTTP handler neither reads nor sets the Mcp-Session-Id header (SEP-2567), and ServerOptions.GetSessionID is not consulted. Rill runs both its MCP servers stateless and keys AI sessions on that header, so after #10019 every request to the runtime's /mcp endpoint created its own ai_sessions row, and the admin service's unified MCP server stopped issuing session IDs, so it never forwarded one to the runtime. Issue and preserve the header in Rill's own HTTP middleware instead of relying on the SDK: a request without a session ID gets a new one, and every response carries the session ID it was served with, which MCP clients echo on subsequent requests. This does not depend on the SDK's MCPGODEBUG=allowsessionsinstateless compatibility flag, which is removed in v1.9. Tests assert that all the requests of one client connection share one AI session, on both the 2025 protocol versions (initialize handshake) and the 2026-07-28 version. --- admin/server/mcp.go | 10 ++-- admin/server/mcp_test.go | 40 +++++++++++++ runtime/ai/mcp.go | 3 - runtime/pkg/middleware/mcp_session.go | 29 +++++++++ runtime/pkg/middleware/mcp_session_test.go | 34 +++++++++++ runtime/server/mcp.go | 7 ++- runtime/server/mcp_test.go | 70 ++++++++++++++++++++++ 7 files changed, 182 insertions(+), 11 deletions(-) create mode 100644 runtime/pkg/middleware/mcp_session.go create mode 100644 runtime/pkg/middleware/mcp_session_test.go diff --git a/admin/server/mcp.go b/admin/server/mcp.go index 190a3610f334..d316449690d6 100644 --- a/admin/server/mcp.go +++ b/admin/server/mcp.go @@ -79,7 +79,9 @@ func (s *Server) mcpHandler() (http.Handler, error) { } } - return mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { + // The handler is stateless, so the session ID is issued and read by us, not the SDK. + // Clients pass it back on subsequent requests, enabling us to keep one AI session per project. + return middleware.MCPSessionIDMiddleware(mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { srv := mcp.NewServer( &mcp.Implementation{ Name: "rill", @@ -89,8 +91,6 @@ func (s *Server) mcpHandler() (http.Handler, error) { &mcp.ServerOptions{ Instructions: mcpInstructions, HasTools: true, - // Issue a session ID so clients pass it back on subsequent requests, enabling us to keep one AI session per project. - GetSessionID: uuid.NewString, }, ) @@ -102,7 +102,7 @@ func (s *Server) mcpHandler() (http.Handler, error) { })) client := &mcpClient{ - sessionID: r.Header.Get("Mcp-Session-Id"), + sessionID: r.Header.Get("Mcp-Session-Id"), // Issued by MCPSessionIDMiddleware if the client did not send one userAgent: r.UserAgent(), protocolVersion: r.Header.Get("Mcp-Protocol-Version"), } @@ -117,7 +117,7 @@ func (s *Server) mcpHandler() (http.Handler, error) { return srv }, &mcp.StreamableHTTPOptions{ Stateless: true, // Required since the server serves multiple users. - }), nil + })), nil } // mcpRequireAuth rejects unauthenticated requests to the MCP server. diff --git a/admin/server/mcp_test.go b/admin/server/mcp_test.go index 00ee5ae7239f..c6d455de5298 100644 --- a/admin/server/mcp_test.go +++ b/admin/server/mcp_test.go @@ -5,11 +5,14 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "strings" "testing" "github.com/google/jsonschema-go/jsonschema" + "github.com/google/uuid" "github.com/modelcontextprotocol/go-sdk/jsonrpc" "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/rilldata/rill/admin" "github.com/rilldata/rill/admin/database" "github.com/rilldata/rill/runtime/ai" "github.com/stretchr/testify/require" @@ -59,7 +62,9 @@ func TestCallRuntimeToolErrors(t *testing.T) { } // The server serves a single tool, so that calls to any other tool exercise the runtime's unknown tool error. + var runtimeSessionID string handler := mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { + runtimeSessionID = r.Header.Get("Mcp-Session-Id") srv := mcp.NewServer(&mcp.Implementation{Name: "runtime"}, &mcp.ServerOptions{HasTools: true}) mcp.AddTool(srv, &mcp.Tool{Name: "echo"}, func(ctx context.Context, req *mcp.CallToolRequest, args echoArgs) (*mcp.CallToolResult, echoResult, error) { return nil, echoResult{Name: args.Name}, nil @@ -78,6 +83,8 @@ func TestCallRuntimeToolErrors(t *testing.T) { require.NoError(t, err) require.False(t, res.IsError) require.Equal(t, map[string]any{"name": "hello"}, res.StructuredContent) + // The call carries a session ID derived from the client's, so the runtime keeps one AI session per client and project. + require.Equal(t, uuid.NewSHA1(mcpSessionNamespace, []byte("default/session")).String(), runtimeSessionID) // Calling a tool the runtime does not serve returns a tool error the client can act on, not a protocol error. res, err = s.callRuntimeTool(t.Context(), depl, "jwt", client, ai.ListSkillsName, json.RawMessage(`{}`)) @@ -105,3 +112,36 @@ func TestCallRuntimeToolErrors(t *testing.T) { _, err = s.callRuntimeTool(t.Context(), &database.Deployment{RuntimeHost: internal.URL, RuntimeInstanceID: "default"}, "jwt", client, "echo", json.RawMessage(`{}`)) require.ErrorContains(t, err, "boom") } + +// TestMCPHandlerSessionID asserts that the MCP server issues a session ID on a client's first request and keeps the one a client sends. +// The server is stateless, so the SDK does not do this itself; without it, every tool call would become its own AI session. +func TestMCPHandlerSessionID(t *testing.T) { + s := &Server{admin: &admin.Service{}} + handler, err := s.mcpHandler() + require.NoError(t, err) + srv := httptest.NewServer(handler) + t.Cleanup(srv.Close) + + post := func(t *testing.T, sessionID, body string) *http.Response { + req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, srv.URL, strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + if sessionID != "" { + req.Header.Set("Mcp-Session-Id", sessionID) + } + res, err := http.DefaultClient.Do(req) + require.NoError(t, err) + t.Cleanup(func() { res.Body.Close() }) + require.Equal(t, http.StatusOK, res.StatusCode) + return res + } + + res := post(t, "", `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`) + sessionID := res.Header.Get("Mcp-Session-Id") + _, err = uuid.Parse(sessionID) + require.NoError(t, err) + + res = post(t, sessionID, `{"jsonrpc":"2.0","id":2,"method":"tools/list"}`) + require.Equal(t, sessionID, res.Header.Get("Mcp-Session-Id")) +} diff --git a/runtime/ai/mcp.go b/runtime/ai/mcp.go index ed1e4c6cf28c..dd67971e03b7 100644 --- a/runtime/ai/mcp.go +++ b/runtime/ai/mcp.go @@ -88,9 +88,6 @@ func (s *Session) MCPServer(ctx context.Context) *mcp.Server { } }, HasTools: true, - GetSessionID: func() string { - return s.id - }, }, ) diff --git a/runtime/pkg/middleware/mcp_session.go b/runtime/pkg/middleware/mcp_session.go new file mode 100644 index 000000000000..4fb3a49ce226 --- /dev/null +++ b/runtime/pkg/middleware/mcp_session.go @@ -0,0 +1,29 @@ +package middleware + +import ( + "net/http" + + "github.com/google/uuid" +) + +// mcpSessionIDHeader is the header MCP clients use to carry a session ID. +// The MCP spec requires a client that receives it to send it on all its subsequent requests. +const mcpSessionIDHeader = "Mcp-Session-Id" + +// MCPSessionIDMiddleware issues and preserves the Mcp-Session-Id header for a stateless MCP server. +// Since go-sdk v1.8, a stateless streamable HTTP handler neither reads nor sets the header, +// since the MCP spec is moving to sessionless servers (see SEP-2567). +// Rill nonetheless keys its AI sessions on the header, so that all the requests of one client connection share a session. +// A request without the header gets a new session ID, which is set on the request for the wrapped handler to read, +// and every response carries the session ID the request was served with. +func MCPSessionIDMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + sessionID := r.Header.Get(mcpSessionIDHeader) + if sessionID == "" { + sessionID = uuid.NewString() + r.Header.Set(mcpSessionIDHeader, sessionID) + } + w.Header().Set(mcpSessionIDHeader, sessionID) + next.ServeHTTP(w, r) + }) +} diff --git a/runtime/pkg/middleware/mcp_session_test.go b/runtime/pkg/middleware/mcp_session_test.go new file mode 100644 index 000000000000..b6140aae02b1 --- /dev/null +++ b/runtime/pkg/middleware/mcp_session_test.go @@ -0,0 +1,34 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestMCPSessionIDMiddleware(t *testing.T) { + var seen string + handler := MCPSessionIDMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seen = r.Header.Get("Mcp-Session-Id") + w.WriteHeader(http.StatusOK) + })) + + // A request without a session ID gets a new one, visible to the handler and returned to the client. + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/mcp", nil)) + issued := rec.Header().Get("Mcp-Session-Id") + _, err := uuid.Parse(issued) + require.NoError(t, err) + require.Equal(t, issued, seen) + + // A request with a session ID keeps it. + req := httptest.NewRequest(http.MethodPost, "/mcp", nil) + req.Header.Set("Mcp-Session-Id", "SESSION-123") + rec = httptest.NewRecorder() + handler.ServeHTTP(rec, req) + require.Equal(t, "SESSION-123", seen) + require.Equal(t, "SESSION-123", rec.Header().Get("Mcp-Session-Id")) +} diff --git a/runtime/server/mcp.go b/runtime/server/mcp.go index ef9653edc7b3..56e82123b7f0 100644 --- a/runtime/server/mcp.go +++ b/runtime/server/mcp.go @@ -19,7 +19,8 @@ import ( func (s *Server) mcpHandler() http.Handler { runner := ai.NewRunner(s.runtime, s.activity) - return mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { + // The handler is stateless, so the session ID that keys the AI session must be issued and read by us, not the SDK. + return middleware.MCPSessionIDMiddleware(mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { // Extract instance ID from the request path instanceID := r.PathValue("instance_id") if instanceID == "" { @@ -29,7 +30,7 @@ func (s *Server) mcpHandler() http.Handler { instanceID, _ = s.runtime.DefaultInstanceID() } - // Get session ID (will be empty if it's the first request) + // The session ID is issued by MCPSessionIDMiddleware if the client did not send one. sessionID := r.Header.Get("Mcp-Session-Id") // It's just preliminary: for direct connections, the MCP server updates it with the actual user agent after the initialization handshake. @@ -68,7 +69,7 @@ func (s *Server) mcpHandler() http.Handler { // It only affects the response: as the spec requires, requests must still accept both content types. // It lets the admin service's unified MCP server forward tool calls with a plain HTTP request instead of parsing an SSE stream. JSONResponse: true, - }) + })) } // mcpRequestSourceMiddleware tags MCP requests with the "mcp" source so billable queries they trigger diff --git a/runtime/server/mcp_test.go b/runtime/server/mcp_test.go index 1c56c5b5534b..3ba625941a68 100644 --- a/runtime/server/mcp_test.go +++ b/runtime/server/mcp_test.go @@ -9,6 +9,7 @@ import ( "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/rilldata/rill/runtime" "github.com/rilldata/rill/runtime/ai" + "github.com/rilldata/rill/runtime/drivers" "github.com/rilldata/rill/runtime/pkg/activity" "github.com/rilldata/rill/runtime/pkg/ratelimit" "github.com/rilldata/rill/runtime/server/auth" @@ -226,3 +227,72 @@ ARPU excludes trial users.`, type roundTripperFunc func(*http.Request) (*http.Response, error) func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } + +// TestMCPSessionContinuity asserts that all the requests of one MCP client connection share one AI session, +// and that separate connections get separate sessions. +// The server runs the go-sdk's streamable HTTP handler in stateless mode, which since go-sdk v1.8 no longer issues the Mcp-Session-Id header itself. +func TestMCPSessionContinuity(t *testing.T) { + // Clients on the 2025 protocol versions (Claude, ChatGPT, Cursor, ...) do an initialize handshake. + // Clients on the 2026-07-28 protocol version (the go-sdk client's default) do not: the server only ever sees their tool calls. + for _, protocolVersion := range []string{"2025-06-18", "2026-07-28"} { + t.Run(protocolVersion, func(t *testing.T) { + rt, instanceID := testruntime.NewInstanceWithOptions(t, testruntime.InstanceOptions{ + Files: map[string]string{"rill.yaml": ""}, + }) + testruntime.RequireReconcileState(t, rt, instanceID, 1, 0, 0) + + srv, err := NewServer(context.Background(), &Options{}, rt, zap.NewNop(), ratelimit.NewNoop(), activity.NewNoopClient()) + require.NoError(t, err) + + httpSrv := httptest.NewServer(auth.HTTPMiddleware(srv.aud, srv.mcpHandler())) + defer httpSrv.Close() + + catalog, release, err := rt.Catalog(t.Context(), instanceID) + require.NoError(t, err) + defer release() + sessions := func(t *testing.T) []*drivers.AISession { + res, err := catalog.FindAISessions(t.Context(), "", "") + require.NoError(t, err) + return res + } + require.Empty(t, sessions(t)) + + // The first request is answered with a session ID, which the client echoes on every subsequent request. + mcpClient := mcp.NewClient(&mcp.Implementation{Name: "mcp/test", Version: "1.0.0"}, nil) + connOpts := &mcp.ClientSessionOptions{ProtocolVersion: protocolVersion} + conn, err := mcpClient.Connect(t.Context(), &mcp.StreamableClientTransport{Endpoint: httpSrv.URL}, connOpts) + require.NoError(t, err) + defer conn.Close() + require.Equal(t, protocolVersion, conn.InitializeResult().ProtocolVersion) + require.NotEmpty(t, conn.ID()) + + _, err = conn.ListTools(t.Context(), &mcp.ListToolsParams{}) + require.NoError(t, err) + _, err = conn.CallTool(t.Context(), &mcp.CallToolParams{Name: ai.ListMetricsViewsName}) + require.NoError(t, err) + _, err = conn.CallTool(t.Context(), &mcp.CallToolParams{Name: ai.ListMetricsViewsName}) + require.NoError(t, err) + + // All the requests above land in a single session. + res := sessions(t) + require.Len(t, res, 1) + require.Equal(t, conn.ID(), res[0].ID) + // Note the user agent is the HTTP User-Agent header, not the client info from the initialize handshake: + // a stateless server handles the initialized notification in a session synthesized without the initialize params, so InitializedHandler never sees the client info. + require.NotEmpty(t, res[0].UserAgent) + msgs, err := catalog.FindAIMessages(t.Context(), conn.ID()) + require.NoError(t, err) + require.Len(t, msgs, 4) // two tool calls, each with a call and a result message + + // A second connection gets its own session. + conn2, err := mcpClient.Connect(t.Context(), &mcp.StreamableClientTransport{Endpoint: httpSrv.URL}, connOpts) + require.NoError(t, err) + defer conn2.Close() + require.NotEmpty(t, conn2.ID()) + require.NotEqual(t, conn.ID(), conn2.ID()) + _, err = conn2.CallTool(t.Context(), &mcp.CallToolParams{Name: ai.ListMetricsViewsName}) + require.NoError(t, err) + require.Len(t, sessions(t), 2) + }) + } +} From 40d1be3d30e9891a18f073c872bf64af8435088a Mon Sep 17 00:00:00 2001 From: Nishant Bangarwa Date: Tue, 20 Oct 2026 15:22:34 +0530 Subject: [PATCH 2/2] fix: address review on MCP session IDs - Retry the session lookup when the insert races with a concurrent request that addresses the same new session ID - Return no server for non-POST MCP requests, so probes do not create AI sessions - Reject malformed client-supplied session IDs in the middleware - Export the header name and drop duplicate comments and a redundant guard --- admin/server/mcp.go | 14 +-- runtime/ai/ai.go | 117 ++++++++++++--------- runtime/ai/session_test.go | 58 ++++++++++ runtime/pkg/middleware/mcp_session.go | 20 +++- runtime/pkg/middleware/mcp_session_test.go | 13 +++ runtime/server/mcp.go | 9 +- runtime/server/mcp_test.go | 11 ++ 7 files changed, 176 insertions(+), 66 deletions(-) create mode 100644 runtime/ai/session_test.go diff --git a/admin/server/mcp.go b/admin/server/mcp.go index d316449690d6..ca788080c93b 100644 --- a/admin/server/mcp.go +++ b/admin/server/mcp.go @@ -79,8 +79,6 @@ func (s *Server) mcpHandler() (http.Handler, error) { } } - // The handler is stateless, so the session ID is issued and read by us, not the SDK. - // Clients pass it back on subsequent requests, enabling us to keep one AI session per project. return middleware.MCPSessionIDMiddleware(mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { srv := mcp.NewServer( &mcp.Implementation{ @@ -102,7 +100,7 @@ func (s *Server) mcpHandler() (http.Handler, error) { })) client := &mcpClient{ - sessionID: r.Header.Get("Mcp-Session-Id"), // Issued by MCPSessionIDMiddleware if the client did not send one + sessionID: r.Header.Get(middleware.MCPSessionIDHeader), // Issued by MCPSessionIDMiddleware if the client did not send one userAgent: r.UserAgent(), protocolVersion: r.Header.Get("Mcp-Protocol-Version"), } @@ -263,12 +261,10 @@ func (s *Server) callRuntimeTool(ctx context.Context, depl *database.Deployment, if client.userAgent != "" { req.Header.Set("User-Agent", client.userAgent) } - if client.sessionID != "" { - // Derive a session ID per project instead of passing the client's through. - // A runtime may serve several projects from one catalog, where AI session IDs must be unique. - sessionID := uuid.NewSHA1(mcpSessionNamespace, []byte(depl.RuntimeInstanceID+"/"+client.sessionID)) - req.Header.Set("Mcp-Session-Id", sessionID.String()) - } + // Derive a session ID per project instead of passing the client's through. + // A runtime may serve several projects from one catalog, where AI session IDs must be unique. + sessionID := uuid.NewSHA1(mcpSessionNamespace, []byte(depl.RuntimeInstanceID+"/"+client.sessionID)) + req.Header.Set(middleware.MCPSessionIDHeader, sessionID.String()) // Uses http.DefaultClient to get caching/pooling of TCP connections. res, err := http.DefaultClient.Do(req) diff --git a/runtime/ai/ai.go b/runtime/ai/ai.go index 7a55a39dfe36..c289c4f5a5a5 100644 --- a/runtime/ai/ai.go +++ b/runtime/ai/ai.go @@ -110,55 +110,10 @@ func (r *Runner) Session(ctx context.Context, opts *SessionOptions) (res *Sessio } defer release() - // Create or load the session in the catalog - var session *drivers.AISession - var messages []*Message - if opts.SessionID != "" { - session, err = catalog.FindAISession(ctx, opts.SessionID) - if err != nil { - // If CreateIfNotExists is set, an unknown session ID is created below instead of erroring. - // This lets a caller that doesn't know Rill's session IDs address a stable session, - // which the unified MCP server in the admin service relies on to keep one session per project. - if !errors.Is(err, drivers.ErrNotFound) || !opts.CreateIfNotExists { - return nil, fmt.Errorf("failed to find session %q: %w", opts.SessionID, err) - } - } - } - if session != nil { - // Check access: you can access anonymous sessions, your own sessions, and shared sessions. - // For shared sessions, if you are not the owner, you can only see messages up to the SharedUntilMessageID (inclusive). - // For sessions without an owner (unauthenticated users using a public project), we don't check access and rely on security by obscurity (generally a decent trade-off, but specifically introduced to get citation links over MCP working for unauthenticated demos). - // It's important to respect SkipChecks to ensure access in Rill Developer (where auth is disabled, but SkipChecks is true). - var retrieveUntilMessageID string - if session.OwnerID != "" && session.OwnerID != opts.Claims.UserID && !opts.Claims.SkipChecks { - if session.SharedUntilMessageID == "" { - return nil, fmt.Errorf("%w: access denied to session %q", runtime.ErrForbidden, session.ID) - } - retrieveUntilMessageID = session.SharedUntilMessageID - } - - ms, err := catalog.FindAIMessages(ctx, opts.SessionID) - if err != nil { - return nil, fmt.Errorf("failed to find messages for session %q: %w", opts.SessionID, err) - } - for _, m := range ms { - messages = append(messages, &Message{ - ID: m.ID, - ParentID: m.ParentID, - SessionID: m.SessionID, - Time: m.CreatedOn, - Index: m.Index, - Role: Role(m.Role), - Type: MessageType(m.Type), - Tool: m.Tool, - ContentType: MessageContentType(m.ContentType), - Content: m.Content, - }) - // only load messages up to and including that retrieveUntilMessageID; messages are ordered by "Index" ascending. - if m.ID == retrieveUntilMessageID { - break - } - } + // Load the session from the catalog, or create it if it does not exist. + session, messages, err := findSession(ctx, catalog, opts) + if err != nil { + return nil, err } if session == nil { id := opts.SessionID @@ -176,7 +131,14 @@ func (r *Runner) Session(ctx context.Context, opts *SessionOptions) (res *Sessio } err = catalog.InsertAISession(ctx, session) if err != nil { - return nil, fmt.Errorf("failed to create session: %w", err) + // Concurrent requests addressing the same new session ID race to create it, + // e.g. parallel tool calls forwarded by the unified MCP server for a project it has not called before. + // If another request won the race, continue with the session it created. + var retryErr error + session, messages, retryErr = findSession(ctx, catalog, opts) + if retryErr != nil || session == nil { + return nil, fmt.Errorf("failed to create session: %w", errors.Join(err, retryErr)) + } } } @@ -239,6 +201,61 @@ func (r *Runner) Session(ctx context.Context, opts *SessionOptions) (res *Sessio }, nil } +// findSession loads the session identified by opts.SessionID from the catalog, along with the messages the caller may see. +// It returns a nil session if opts.SessionID is empty, or if the session does not exist and opts.CreateIfNotExists is set. +func findSession(ctx context.Context, catalog drivers.CatalogStore, opts *SessionOptions) (*drivers.AISession, []*Message, error) { + if opts.SessionID == "" { + return nil, nil, nil + } + session, err := catalog.FindAISession(ctx, opts.SessionID) + if err != nil { + // If CreateIfNotExists is set, an unknown session ID is created by the caller instead of erroring. + // This lets a caller that doesn't know Rill's session IDs address a stable session, + // which the unified MCP server in the admin service relies on to keep one session per project. + if errors.Is(err, drivers.ErrNotFound) && opts.CreateIfNotExists { + return nil, nil, nil + } + return nil, nil, fmt.Errorf("failed to find session %q: %w", opts.SessionID, err) + } + + // Check access: you can access anonymous sessions, your own sessions, and shared sessions. + // For shared sessions, if you are not the owner, you can only see messages up to the SharedUntilMessageID (inclusive). + // For sessions without an owner (unauthenticated users using a public project), we don't check access and rely on security by obscurity (generally a decent trade-off, but specifically introduced to get citation links over MCP working for unauthenticated demos). + // It's important to respect SkipChecks to ensure access in Rill Developer (where auth is disabled, but SkipChecks is true). + var retrieveUntilMessageID string + if session.OwnerID != "" && session.OwnerID != opts.Claims.UserID && !opts.Claims.SkipChecks { + if session.SharedUntilMessageID == "" { + return nil, nil, fmt.Errorf("%w: access denied to session %q", runtime.ErrForbidden, session.ID) + } + retrieveUntilMessageID = session.SharedUntilMessageID + } + + ms, err := catalog.FindAIMessages(ctx, opts.SessionID) + if err != nil { + return nil, nil, fmt.Errorf("failed to find messages for session %q: %w", opts.SessionID, err) + } + var messages []*Message + for _, m := range ms { + messages = append(messages, &Message{ + ID: m.ID, + ParentID: m.ParentID, + SessionID: m.SessionID, + Time: m.CreatedOn, + Index: m.Index, + Role: Role(m.Role), + Type: MessageType(m.Type), + Tool: m.Tool, + ContentType: MessageContentType(m.ContentType), + Content: m.Content, + }) + // only load messages up to and including that retrieveUntilMessageID; messages are ordered by "Index" ascending. + if m.ID == retrieveUntilMessageID { + break + } + } + return session, messages, nil +} + // ForkSession creates a new session cloned from an existing session. func (r *Runner) ForkSession(ctx context.Context, opts *SessionOptions) (string, error) { if opts.SessionID == "" { diff --git a/runtime/ai/session_test.go b/runtime/ai/session_test.go new file mode 100644 index 000000000000..a555037a8191 --- /dev/null +++ b/runtime/ai/session_test.go @@ -0,0 +1,58 @@ +package ai_test + +import ( + "fmt" + "testing" + + "github.com/google/uuid" + "github.com/rilldata/rill/runtime" + "github.com/rilldata/rill/runtime/ai" + "github.com/rilldata/rill/runtime/pkg/activity" + "github.com/rilldata/rill/runtime/testruntime" + "github.com/stretchr/testify/require" + "golang.org/x/sync/errgroup" +) + +// TestSessionConcurrentCreate ensures concurrent requests addressing the same new session ID all get that session. +// It happens when an MCP client makes parallel tool calls on a session the runtime has not seen before, +// e.g. the first calls the unified MCP server in the admin service forwards for a project. +func TestSessionConcurrentCreate(t *testing.T) { + rt, instanceID := testruntime.NewInstanceWithOptions(t, testruntime.InstanceOptions{ + Files: map[string]string{"rill.yaml": ""}, + }) + testruntime.RequireReconcileState(t, rt, instanceID, 1, 0, 0) + r := ai.NewRunner(rt, activity.NewNoopClient()) + sessionID := uuid.NewString() + + start := make(chan struct{}) + var g errgroup.Group + for range 10 { + g.Go(func() error { + <-start + s, err := r.Session(t.Context(), &ai.SessionOptions{ + InstanceID: instanceID, + SessionID: sessionID, + CreateIfNotExists: true, + Claims: &runtime.SecurityClaims{SkipChecks: true}, + UserAgent: "test", + }) + if err != nil { + return err + } + if s.ID() != sessionID { + return fmt.Errorf("got session %q, want %q", s.ID(), sessionID) + } + return nil + }) + } + close(start) + require.NoError(t, g.Wait()) + + catalog, release, err := rt.Catalog(t.Context(), instanceID) + require.NoError(t, err) + defer release() + sessions, err := catalog.FindAISessions(t.Context(), "", "") + require.NoError(t, err) + require.Len(t, sessions, 1) + require.Equal(t, sessionID, sessions[0].ID) +} diff --git a/runtime/pkg/middleware/mcp_session.go b/runtime/pkg/middleware/mcp_session.go index 4fb3a49ce226..0b0eeb3a448a 100644 --- a/runtime/pkg/middleware/mcp_session.go +++ b/runtime/pkg/middleware/mcp_session.go @@ -2,13 +2,18 @@ package middleware import ( "net/http" + "strings" "github.com/google/uuid" ) -// mcpSessionIDHeader is the header MCP clients use to carry a session ID. +// MCPSessionIDHeader is the header MCP clients use to carry a session ID. // The MCP spec requires a client that receives it to send it on all its subsequent requests. -const mcpSessionIDHeader = "Mcp-Session-Id" +const MCPSessionIDHeader = "Mcp-Session-Id" + +// maxMCPSessionIDLength bounds client-supplied session IDs, which the runtime stores as AI session IDs. +// Rill's own IDs are UUIDs, so the bound only matters for IDs that did not originate here. +const maxMCPSessionIDLength = 256 // MCPSessionIDMiddleware issues and preserves the Mcp-Session-Id header for a stateless MCP server. // Since go-sdk v1.8, a stateless streamable HTTP handler neither reads nor sets the header, @@ -16,14 +21,19 @@ const mcpSessionIDHeader = "Mcp-Session-Id" // Rill nonetheless keys its AI sessions on the header, so that all the requests of one client connection share a session. // A request without the header gets a new session ID, which is set on the request for the wrapped handler to read, // and every response carries the session ID the request was served with. +// A request with a malformed session ID is rejected: the spec restricts session IDs to visible ASCII, +// and the length bound protects the catalog, which stores the ID as the AI session's ID. func MCPSessionIDMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - sessionID := r.Header.Get(mcpSessionIDHeader) + sessionID := r.Header.Get(MCPSessionIDHeader) if sessionID == "" { sessionID = uuid.NewString() - r.Header.Set(mcpSessionIDHeader, sessionID) + r.Header.Set(MCPSessionIDHeader, sessionID) + } else if len(sessionID) > maxMCPSessionIDLength || strings.ContainsFunc(sessionID, func(c rune) bool { return c < 0x21 || c > 0x7E }) { + http.Error(w, "Bad Request: invalid Mcp-Session-Id header", http.StatusBadRequest) + return } - w.Header().Set(mcpSessionIDHeader, sessionID) + w.Header().Set(MCPSessionIDHeader, sessionID) next.ServeHTTP(w, r) }) } diff --git a/runtime/pkg/middleware/mcp_session_test.go b/runtime/pkg/middleware/mcp_session_test.go index b6140aae02b1..a9fe8758baf2 100644 --- a/runtime/pkg/middleware/mcp_session_test.go +++ b/runtime/pkg/middleware/mcp_session_test.go @@ -3,6 +3,7 @@ package middleware import ( "net/http" "net/http/httptest" + "strings" "testing" "github.com/google/uuid" @@ -31,4 +32,16 @@ func TestMCPSessionIDMiddleware(t *testing.T) { handler.ServeHTTP(rec, req) require.Equal(t, "SESSION-123", seen) require.Equal(t, "SESSION-123", rec.Header().Get("Mcp-Session-Id")) + + // A malformed session ID is rejected before reaching the handler. + for _, invalid := range []string{"session with spaces", "s\u00e9ssion", "session\x7f", strings.Repeat("s", maxMCPSessionIDLength+1)} { + seen = "" + req = httptest.NewRequest(http.MethodPost, "/mcp", nil) + req.Header.Set("Mcp-Session-Id", invalid) + rec = httptest.NewRecorder() + handler.ServeHTTP(rec, req) + require.Equal(t, http.StatusBadRequest, rec.Code, invalid) + require.Empty(t, seen, invalid) + require.Empty(t, rec.Header().Get("Mcp-Session-Id"), invalid) + } } diff --git a/runtime/server/mcp.go b/runtime/server/mcp.go index 56e82123b7f0..3bc93d811227 100644 --- a/runtime/server/mcp.go +++ b/runtime/server/mcp.go @@ -19,8 +19,13 @@ import ( func (s *Server) mcpHandler() http.Handler { runner := ai.NewRunner(s.runtime, s.activity) - // The handler is stateless, so the session ID that keys the AI session must be issued and read by us, not the SDK. return middleware.MCPSessionIDMiddleware(mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { + // A stateless server only accepts POST, but the SDK calls this function for every request before rejecting other methods with 405. + // Returning early keeps probes and legacy GET/DELETE requests from creating AI sessions. + if r.Method != http.MethodPost { + return nil + } + // Extract instance ID from the request path instanceID := r.PathValue("instance_id") if instanceID == "" { @@ -31,7 +36,7 @@ func (s *Server) mcpHandler() http.Handler { } // The session ID is issued by MCPSessionIDMiddleware if the client did not send one. - sessionID := r.Header.Get("Mcp-Session-Id") + sessionID := r.Header.Get(middleware.MCPSessionIDHeader) // It's just preliminary: for direct connections, the MCP server updates it with the actual user agent after the initialization handshake. // Requests proxied by the admin service's unified MCP server don't carry that handshake, so the forwarded user agent is the only signal of the real client. diff --git a/runtime/server/mcp_test.go b/runtime/server/mcp_test.go index 3ba625941a68..f1059ae800b8 100644 --- a/runtime/server/mcp_test.go +++ b/runtime/server/mcp_test.go @@ -293,6 +293,17 @@ func TestMCPSessionContinuity(t *testing.T) { _, err = conn2.CallTool(t.Context(), &mcp.CallToolParams{Name: ai.ListMetricsViewsName}) require.NoError(t, err) require.Len(t, sessions(t), 2) + + // Requests other than POST are rejected without creating a session. + for _, method := range []string{http.MethodGet, http.MethodDelete} { + req, err := http.NewRequestWithContext(t.Context(), method, httpSrv.URL, nil) + require.NoError(t, err) + res, err := http.DefaultClient.Do(req) + require.NoError(t, err) + res.Body.Close() + require.Equal(t, http.StatusMethodNotAllowed, res.StatusCode, method) + } + require.Len(t, sessions(t), 2) }) } }