summercms/modules/pier/session.go

package pier

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"slices"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/modelcontextprotocol/go-sdk/jsonrpc"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// Protocol constants the router reads.
const (
	sessionIDHeader       = "Mcp-Session-Id"
	protocolVersionHeader = "Mcp-Protocol-Version"
	protocolVersion202607 = "2026-07-28"
	protocolVersion202511 = "2025-11-25"
	protocolVersion202503 = "2025-03-26"
	methodInitialize      = "initialize"
	methodListTools       = "tools/list"
	methodCallTool        = "tools/call"
	methodListPrompts     = "prompts/list"
	methodGetPrompt       = "prompts/get"
	methodDiscover        = "server/discover"
)

// Keys of auth.TokenInfo.Extra set by the bearer middleware.
const (
	extraIdentity = "pier.identity"
	extraCall     = "pier.call"
)

// sessionRouter is the stateful streamable HTTP path (D-25): sessions live
// in the Store, bound to the identity that opened them, and are served by
// a per-session SDK server on whichever instance receives the request,
// rehydrated from the stored session state when it is not local.
type sessionRouter struct {
	p     *Pier
	mu    sync.Mutex
	local map[string]*localSession
}

// localSession is one session connected on this instance.
type localSession struct {
	id        string
	binding   string
	created   time.Time
	server    *mcp.Server
	transport *mcp.StreamableServerTransport
	ss        *mcp.ServerSession

	mu       sync.Mutex
	identity Identity
	tools    map[string]bool
	prompts  map[string]bool

	initialized atomic.Bool
	persisted   atomic.Bool
}

func newSessionRouter(p *Pier) *sessionRouter {
	return &sessionRouter{p: p, local: map[string]*localSession{}}
}

func (sr *sessionRouter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	id, ok := identityFromRequest(r)
	if !ok {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	sid := r.Header.Get(sessionIDHeader)
	switch r.Method {
	case http.MethodDelete:
		sr.serveDelete(w, r, id, sid)
	case http.MethodGet:
		if sid == "" {
			http.Error(w, "Bad Request: GET requires an Mcp-Session-Id header", http.StatusBadRequest)
			return
		}
		ls, ok := sr.lookup(w, r, id, sid)
		if !ok {
			return
		}
		sr.serveStream(w, r, ls)
	case http.MethodPost:
		if sid == "" {
			sr.serveNew(w, r, id)
			return
		}
		ls, ok := sr.lookup(w, r, id, sid)
		if !ok {
			return
		}
		ls.transport.ServeHTTP(w, r)
		sr.persist(r.Context(), ls)
	default:
		w.Header().Set("Allow", "GET, POST, DELETE")
		http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
	}
}

// serveNew handles a POST without a session id: an initialize opens a
// session (refused with 429 above mcp.sessions.max_per_binding); any other
// request is served by a throwaway connection and closed, as the SDK does,
// so a server/discover probe or a ping before initialize gets an answer.
func (sr *sessionRouter) serveNew(w http.ResponseWriter, r *http.Request, id Identity) {
	body, err := io.ReadAll(r.Body)
	if err != nil {
		var mbe *http.MaxBytesError
		if errors.As(err, &mbe) {
			http.Error(w, fmt.Sprintf("request body exceeds %d bytes", mbe.Limit), http.StatusRequestEntityTooLarge)
			return
		}
		http.Error(w, "failed to read body", http.StatusBadRequest)
		return
	}
	_ = r.Body.Close()
	r.Body = io.NopCloser(bytes.NewReader(body))
	if !slices.Contains(peekMethods(body), methodInitialize) {
		sr.serveEphemeral(w, r, id)
		return
	}
	ctx := r.Context()
	n, err := sr.p.store.Count(ctx, id.Binding)
	if err != nil {
		sr.p.storeUnavailable(w, err)
		return
	}
	if n >= sr.p.cfg.MaxPerBinding {
		writeJSON(w, http.StatusTooManyRequests, map[string]any{"error": "too_many_sessions", "max_per_binding": sr.p.cfg.MaxPerBinding})
		return
	}
	sid := rand.Text()
	sr.mu.Lock()
	ls, err := sr.connect(ctx, sid, id, nil)
	if err == nil {
		sr.local[sid] = ls
	}
	sr.mu.Unlock()
	if err != nil {
		logger().Error("pier: session connect failed", "error", err)
		http.Error(w, "failed connection", http.StatusInternalServerError)
		return
	}
	if err := sr.p.store.Put(ctx, sid, Record{Binding: id.Binding, Created: ls.created}, sr.p.cfg.IdleTTL); err != nil {
		sr.closeLocal(sid)
		sr.p.storeUnavailable(w, err)
		return
	}
	logSession(logger(), "open", sid, id.Binding)
	ls.transport.ServeHTTP(w, r)
	if ls.ss.InitializeParams() == nil {
		// initialize failed: nothing to keep (the SDK cleans up the same way).
		_ = sr.p.store.Delete(ctx, sid)
		sr.closeLocal(sid)
		return
	}
	sr.persist(ctx, ls)
}

// serveEphemeral answers one request outside any session with the state
// the SDK synthesises for a pre-initialize request.
func (sr *sessionRouter) serveEphemeral(w http.ResponseWriter, r *http.Request, id Identity) {
	pv := r.Header.Get(protocolVersionHeader)
	if pv == "" || pv >= protocolVersion202607 {
		pv = protocolVersion202503
	}
	srv, _, _ := sr.p.newServer(r.Context(), id, nil)
	tr := &mcp.StreamableServerTransport{}
	state := &mcp.ServerSessionState{
		InitializeParams:  &mcp.InitializeParams{ProtocolVersion: pv},
		InitializedParams: &mcp.InitializedParams{},
		LogLevel:          mcp.LoggingLevel("info"),
	}
	ss, err := srv.Connect(r.Context(), tr, &mcp.ServerSessionOptions{State: state})
	if err != nil {
		http.Error(w, "failed connection", http.StatusInternalServerError)
		return
	}
	defer ss.Close()
	tr.ServeHTTP(w, r)
}

// serveDelete ends a session everywhere: the record is removed, every
// instance is told and the local copy is closed.
func (sr *sessionRouter) serveDelete(w http.ResponseWriter, r *http.Request, id Identity, sid string) {
	if sid == "" {
		http.Error(w, "Bad Request: DELETE requires an Mcp-Session-Id header", http.StatusBadRequest)
		return
	}
	ctx := r.Context()
	rec, ok, err := sr.p.store.Get(ctx, sid)
	if err != nil {
		sr.p.storeUnavailable(w, err)
		return
	}
	if !ok {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}
	if rec.Binding != id.Binding {
		logSession(logger(), "binding_refused", sid, id.Binding)
		writeJSON(w, http.StatusForbidden, map[string]any{"error": "session_binding"})
		return
	}
	if err := sr.p.store.Delete(ctx, sid); err != nil {
		sr.p.storeUnavailable(w, err)
		return
	}
	_ = sr.p.store.Publish(ctx, Event{Kind: EventEnded, SessionID: sid, Instance: sr.p.instance})
	sr.closeLocal(sid)
	logSession(logger(), "end", sid, id.Binding)
	w.WriteHeader(http.StatusNoContent)
}

// lookup finds the session of a request: 404 when the record is gone (the
// client re-initialises), 403 when the binding differs, else the local
// copy, rehydrated from the stored state when this instance has none; the
// record's TTL slides on every request.
func (sr *sessionRouter) lookup(w http.ResponseWriter, r *http.Request, id Identity, sid string) (*localSession, bool) {
	ctx := r.Context()
	rec, ok, err := sr.p.store.Get(ctx, sid)
	if err != nil {
		sr.p.storeUnavailable(w, err)
		return nil, false
	}
	if !ok {
		sr.closeLocal(sid)
		http.Error(w, "session not found", http.StatusNotFound)
		return nil, false
	}
	if rec.Binding != id.Binding {
		logSession(logger(), "binding_refused", sid, id.Binding)
		writeJSON(w, http.StatusForbidden, map[string]any{"error": "session_binding"})
		return nil, false
	}
	sr.mu.Lock()
	ls := sr.local[sid]
	if ls == nil {
		if len(rec.State) == 0 {
			sr.mu.Unlock()
			http.Error(w, "session not found", http.StatusNotFound)
			return nil, false
		}
		var state mcp.ServerSessionState
		if err := json.Unmarshal(rec.State, &state); err != nil || state.InitializeParams == nil {
			sr.mu.Unlock()
			http.Error(w, "session not found", http.StatusNotFound)
			return nil, false
		}
		ls, err = sr.connect(ctx, sid, id, &state)
		if err != nil {
			sr.mu.Unlock()
			logger().Error("pier: session rehydrate failed", "error", err)
			http.Error(w, "failed connection", http.StatusInternalServerError)
			return nil, false
		}
		ls.created = rec.Created
		if state.InitializedParams != nil {
			ls.persisted.Store(true)
		}
		sr.local[sid] = ls
		logSession(logger(), "rehydrate", sid, id.Binding)
	}
	sr.mu.Unlock()
	ls.mu.Lock()
	ls.identity = id
	ls.mu.Unlock()
	if err := sr.p.store.Touch(ctx, sid, sr.p.cfg.IdleTTL); err != nil {
		sr.p.storeUnavailable(w, err)
		return nil, false
	}
	return ls, true
}

// connect builds the per-identity server and connects it over a transport
// with the session id, from the stored state when rehydrating. A goroutine
// drops the local entry when the session closes.
func (sr *sessionRouter) connect(ctx context.Context, sid string, id Identity, state *mcp.ServerSessionState) (*localSession, error) {
	ls := &localSession{id: sid, binding: id.Binding, created: time.Now().UTC(), identity: id}
	srv, tools, prompts := sr.p.newServer(ctx, id, func() { ls.initialized.Store(true) })
	ls.server, ls.tools, ls.prompts = srv, tools, prompts
	ls.transport = &mcp.StreamableServerTransport{SessionID: sid}
	var opts *mcp.ServerSessionOptions
	if state != nil {
		opts = &mcp.ServerSessionOptions{State: state}
		if state.InitializedParams != nil {
			ls.initialized.Store(true)
		}
	}
	ss, err := srv.Connect(context.WithoutCancel(ctx), ls.transport, opts)
	if err != nil {
		return nil, err
	}
	ls.ss = ss
	go func() {
		_ = ss.Wait()
		sr.mu.Lock()
		if sr.local[sid] == ls {
			delete(sr.local, sid)
		}
		sr.mu.Unlock()
	}()
	return ls, nil
}

// persist writes the session state once initialize answered and again once
// the client's initialized notification arrived, so another instance can
// rehydrate a usable session.
func (sr *sessionRouter) persist(ctx context.Context, ls *localSession) {
	if ls.persisted.Load() {
		return
	}
	params := ls.ss.InitializeParams()
	if params == nil {
		return
	}
	state := mcp.ServerSessionState{
		InitializeParams:          params,
		NegotiatedProtocolVersion: negotiatedVersion(params.ProtocolVersion),
		LogLevel:                  mcp.LoggingLevel("info"),
	}
	complete := ls.initialized.Load()
	if complete {
		state.InitializedParams = &mcp.InitializedParams{}
	}
	raw, err := json.Marshal(state)
	if err != nil {
		return
	}
	if err := sr.p.store.Put(ctx, ls.id, Record{Binding: ls.binding, State: raw, Created: ls.created}, sr.p.cfg.IdleTTL); err != nil {
		logger().Warn("pier: session state not stored", "session", sessionPrefix(ls.id), "error", err)
		return
	}
	if complete {
		ls.persisted.Store(true)
	}
}

// serveStream serves the standalone GET stream and re-verifies its bearer
// every recheck interval on a goroutine that holds the raw bearer only for
// the stream's lifetime; a bearer that no longer resolves to the session's
// binding ends the stream (Pitfall 10). The server's base context ends
// every stream at shutdown.
func (sr *sessionRouter) serveStream(w http.ResponseWriter, r *http.Request, ls *localSession) {
	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()
	raw := bearerOf(r)
	go func() {
		ticker := time.NewTicker(sr.p.recheck)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				id, err := sr.p.identities.resolve(ctx, raw)
				if err != nil || id.Binding != ls.binding {
					logSession(logger(), "stream_closed", ls.id, ls.binding)
					cancel()
					return
				}
				_ = sr.p.store.Touch(ctx, ls.id, sr.p.cfg.IdleTTL)
			}
		}
	}()
	ls.transport.ServeHTTP(w, r.WithContext(ctx))
}

// bearerOf reads the bearer of a request's Authorization header.
func bearerOf(r *http.Request) string {
	fields := strings.Fields(r.Header.Get("Authorization"))
	if len(fields) != 2 || !strings.EqualFold(fields[0], "bearer") {
		return ""
	}
	return fields[1]
}

// closeLocal closes and forgets this instance's copy of a session.
func (sr *sessionRouter) closeLocal(sid string) {
	sr.mu.Lock()
	ls := sr.local[sid]
	delete(sr.local, sid)
	sr.mu.Unlock()
	if ls != nil && ls.ss != nil {
		_ = ls.ss.Close()
	}
}

// negotiatedVersion mirrors the SDK's handshake rule: the client's version
// when the SDK supports it and it predates 2026-07-28, else the newest
// handshake-era version.
func negotiatedVersion(client string) string {
	supported := mcp.SupportedProtocolVersions()
	if slices.Contains(supported, client) && client < protocolVersion202607 {
		return client
	}
	for _, v := range supported {
		if v < protocolVersion202607 {
			return v
		}
	}
	return protocolVersion202511
}

// peekMethods lists the JSON-RPC methods of a request body (one message or
// a batch) without consuming it; a malformed body lists none.
func peekMethods(body []byte) []string {
	type msg struct {
		Method string `json:"method"`
	}
	trimmed := bytes.TrimSpace(body)
	if len(trimmed) == 0 {
		return nil
	}
	if trimmed[0] == '[' {
		var batch []msg
		if err := json.Unmarshal(trimmed, &batch); err != nil {
			return nil
		}
		out := make([]string, 0, len(batch))
		for _, m := range batch {
			out = append(out, m.Method)
		}
		return out
	}
	var m msg
	if err := json.Unmarshal(trimmed, &m); err != nil {
		return nil
	}
	return []string{m.Method}
}

// newServer builds one SDK server for an identity: the tools and prompts
// allowed(...) grants it, the receiving middleware that re-checks every
// list and call for the current request's identity, and the handlers that
// dispatch calls and render prompts. It returns the server and the names
// it holds, for later diffs.
func (p *Pier) newServer(ctx context.Context, id Identity, onInitialized func()) (*mcp.Server, map[string]bool, map[string]bool) {
	opts := &mcp.ServerOptions{
		Capabilities: &mcp.ServerCapabilities{
			Tools:   &mcp.ToolCapabilities{ListChanged: true},
			Prompts: &mcp.PromptCapabilities{ListChanged: true},
		},
	}
	if onInitialized != nil {
		opts.InitializedHandler = func(context.Context, *mcp.InitializedRequest) { onInitialized() }
	}
	srv := mcp.NewServer(&mcp.Implementation{Name: p.appName, Version: p.version}, opts)
	srv.AddReceivingMiddleware(p.receiving(id))
	tools, prompts := p.catalog.allowed(ctx, id)
	toolSet := map[string]bool{}
	for _, t := range tools {
		srv.AddTool(t.desc, p.toolHandler(t, id))
		toolSet[t.name] = true
	}
	promptSet := map[string]bool{}
	for _, pr := range prompts {
		srv.AddPrompt(pr.desc, p.promptHandler(pr))
		promptSet[pr.name] = true
	}
	return srv, toolSet, promptSet
}

// receiving is the per-request re-check (D-08, D-25): tools/list and
// prompts/list are filtered for the identity of the request that carried
// them, and tools/call or prompts/get for an item outside that list is a
// JSON-RPC invalid-params error with a fixed text.
func (p *Pier) receiving(session Identity) mcp.Middleware {
	return func(next mcp.MethodHandler) mcp.MethodHandler {
		return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
			switch method {
			case methodListTools, methodListPrompts, methodCallTool, methodGetPrompt:
			default:
				return next(ctx, method, req)
			}
			id, _ := identityFromExtra(req.GetExtra(), session)
			tools, prompts := p.catalog.allowed(ctx, id)
			switch method {
			case methodCallTool:
				params, _ := req.GetParams().(*mcp.CallToolParamsRaw)
				if params == nil || !containsTool(tools, params.Name) {
					return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: "pier: tool is not available to this caller"}
				}
				return next(ctx, method, req)
			case methodGetPrompt:
				params, _ := req.GetParams().(*mcp.GetPromptParams)
				if params == nil || !containsPrompt(prompts, params.Name) {
					return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: "pier: prompt is not available to this caller"}
				}
				return next(ctx, method, req)
			case methodListTools:
				res, err := next(ctx, method, req)
				if err != nil {
					return nil, err
				}
				if lr, ok := res.(*mcp.ListToolsResult); ok {
					kept := make([]*mcp.Tool, 0, len(lr.Tools))
					for _, t := range lr.Tools {
						if containsTool(tools, t.Name) {
							kept = append(kept, t)
						}
					}
					lr.Tools = kept
				}
				return res, nil
			default:
				res, err := next(ctx, method, req)
				if err != nil {
					return nil, err
				}
				if lr, ok := res.(*mcp.ListPromptsResult); ok {
					kept := make([]*mcp.Prompt, 0, len(lr.Prompts))
					for _, pr := range lr.Prompts {
						if containsPrompt(prompts, pr.Name) {
							kept = append(kept, pr)
						}
					}
					lr.Prompts = kept
				}
				return res, nil
			}
		}
	}
}

func containsTool(tools []*tool, name string) bool {
	for _, t := range tools {
		if t.name == name {
			return true
		}
	}
	return false
}

func containsPrompt(prompts []*prompt, name string) bool {
	for _, p := range prompts {
		if p.name == name {
			return true
		}
	}
	return false
}

// toolHandler dispatches one call as the current request's identity.
func (p *Pier) toolHandler(t *tool, session Identity) mcp.ToolHandler {
	return func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		id, info := identityFromExtra(req.Extra, session)
		if req.Session != nil {
			info.session = req.Session.ID()
		}
		return p.dispatch.call(ctx, id, info, t, req.Params.Arguments), nil
	}
}

// promptHandler checks the required arguments and renders the prompt.
func (p *Pier) promptHandler(pr *prompt) mcp.PromptHandler {
	return func(ctx context.Context, req *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
		args := map[string]string{}
		if req.Params != nil {
			for k, v := range req.Params.Arguments {
				args[k] = v
			}
		}
		for _, a := range pr.decl.Arguments {
			if a.Required && args[a.Name] == "" {
				return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("pier: prompt argument %q is required", a.Name)}
			}
		}
		messages, err := pr.decl.Render(ctx, args)
		if err != nil {
			return nil, err
		}
		res := &mcp.GetPromptResult{Description: pr.decl.Description, Messages: []*mcp.PromptMessage{}}
		for _, m := range messages {
			role := mcp.Role(m.Role)
			if role != "assistant" {
				role = "user"
			}
			res.Messages = append(res.Messages, &mcp.PromptMessage{Role: role, Content: &mcp.TextContent{Text: m.Text}})
		}
		return res, nil
	}
}

// identityFromExtra reads the request's identity and call info off the
// token info the bearer middleware stored; the session's identity is the
// fallback for a request that carries none.
func identityFromExtra(extra *mcp.RequestExtra, fallback Identity) (Identity, callInfo) {
	if extra == nil || extra.TokenInfo == nil || extra.TokenInfo.Extra == nil {
		return fallback, callInfo{}
	}
	id, _ := extra.TokenInfo.Extra[extraIdentity].(Identity)
	info, _ := extra.TokenInfo.Extra[extraCall].(callInfo)
	if id.Binding == "" {
		return fallback, info
	}
	return id, info
}

// writeJSON writes a small JSON body with a status.
func writeJSON(w http.ResponseWriter, status int, body any) {
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(body)
}