diff --git a/admin/server/mcp.go b/admin/server/mcp.go index 190a3610f334..ca788080c93b 100644 --- a/admin/server/mcp.go +++ b/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) 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/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/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/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 new file mode 100644 index 000000000000..0b0eeb3a448a --- /dev/null +++ b/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) + }) +} diff --git a/runtime/pkg/middleware/mcp_session_test.go b/runtime/pkg/middleware/mcp_session_test.go new file mode 100644 index 000000000000..a9fe8758baf2 --- /dev/null +++ b/runtime/pkg/middleware/mcp_session_test.go @@ -0,0 +1,47 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "strings" + "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")) + + // 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 ef9653edc7b3..3bc93d811227 100644 --- a/runtime/server/mcp.go +++ b/runtime/server/mcp.go @@ -19,7 +19,13 @@ import ( func (s *Server) mcpHandler() http.Handler { runner := ai.NewRunner(s.runtime, s.activity) - return mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { + 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 == "" { @@ -29,8 +35,8 @@ func (s *Server) mcpHandler() http.Handler { instanceID, _ = s.runtime.DefaultInstanceID() } - // Get session ID (will be empty if it's the first request) - sessionID := r.Header.Get("Mcp-Session-Id") + // The session ID is issued by MCPSessionIDMiddleware if the client did not send one. + 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. @@ -68,7 +74,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..f1059ae800b8 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,83 @@ 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) + + // 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) + }) + } +}