fix: MCP requests no longer share an AI session after the go-sdk bump to v1.8.0

admin/server/mcp.go

@@ -79,7 +79,7 @@ func (s *Server) mcpHandler() (http.Handler, error) {
 		}
 	}
 
-	return mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server {
+	return middleware.MCPSessionIDMiddleware(mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server {
 		srv := mcp.NewServer(
 			&mcp.Implementation{
 				Name:    "rill",
@@ -89,8 +89,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 +100,7 @@ func (s *Server) mcpHandler() (http.Handler, error) {
 		}))
 
 		client := &mcpClient{
-			sessionID:       r.Header.Get("Mcp-Session-Id"),
+			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"),
 		}
@@ -117,7 +115,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.
@@ -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)

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"))
+}

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 == "" {

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
-			},
 		},
 	)
 

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)
+}

runtime/pkg/middleware/mcp_session.go

@@ -0,0 +1,39 @@
+package middleware
+
+import (
+	"net/http"
+	"strings"
+
+	"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"
+
+// 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,
+// 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.
+// 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)
+		if sessionID == "" {
+			sessionID = uuid.NewString()
+			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)
+		next.ServeHTTP(w, r)
+	})
+}