summercms/modules/pier/stateless.go

package pier

import (
	"bytes"
	"container/list"
	"errors"
	"io"
	"net/http"
	"slices"
	"sync"

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

// serverFeatures is what applyDiff needs from an SDK server.
type serverFeatures interface {
	AddTool(t *mcp.Tool, h mcp.ToolHandler)
	RemoveTools(names ...string)
	AddPrompt(p *mcp.Prompt, h mcp.PromptHandler)
	RemovePrompts(names ...string)
}

// statelessEntry is one per-identity server of the stateless path, keyed
// by binding, with the names it holds for fan-out diffs.
type statelessEntry struct {
	binding string
	server  *mcp.Server
	elem    *list.Element

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

// statelessPath serves clients of protocol 2026-07-28 and later (D-32,
// P-08): the SDK's stateless handler over an in-process LRU of
// per-identity servers (capacity mcp.sessions.max_per_binding * 16),
// evicted on revocation, mutated on fan-out events so subscriptions/listen
// clients receive the list changes. Every request is authorized afresh by
// the bearer middleware in front of it.
type statelessPath struct {
	p        *Pier
	handler  http.Handler
	capacity int

	mu      sync.Mutex
	entries map[string]*statelessEntry
	order   *list.List // front = most recently used
}

// statelessCapacityFactor multiplies mcp.sessions.max_per_binding into the
// LRU capacity.
const statelessCapacityFactor = 16

func newStatelessPath(p *Pier) *statelessPath {
	sp := &statelessPath{
		p:        p,
		capacity: p.cfg.MaxPerBinding * statelessCapacityFactor,
		entries:  map[string]*statelessEntry{},
		order:    list.New(),
	}
	sp.handler = mcp.NewStreamableHTTPHandler(sp.getServer, &mcp.StreamableHTTPOptions{
		Stateless: true,
		// pier runs behind a reverse proxy and checks Origin itself; the
		// SDK's loopback rule would refuse every proxied request of a
		// binary listening on 127.0.0.1.
		DisableLocalhostProtection: true,
		MaxRequestBodyBytes:        maxRequestBody,
	})
	return sp
}

func (sp *statelessPath) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	sp.handler.ServeHTTP(w, r)
}

// getServer returns the per-identity server of the request's binding,
// building it on a miss and refreshing the identity on a hit.
func (sp *statelessPath) getServer(r *http.Request) *mcp.Server {
	id, ok := identityFromRequest(r)
	if !ok {
		return nil
	}
	sp.mu.Lock()
	defer sp.mu.Unlock()
	if e, ok := sp.entries[id.Binding]; ok {
		sp.order.MoveToFront(e.elem)
		e.mu.Lock()
		e.identity = id
		e.mu.Unlock()
		return e.server
	}
	srv, tools, prompts := sp.p.newServer(r.Context(), id, nil)
	e := &statelessEntry{binding: id.Binding, server: srv, identity: id, tools: tools, prompts: prompts}
	e.elem = sp.order.PushFront(e)
	sp.entries[id.Binding] = e
	for sp.order.Len() > sp.capacity {
		last := sp.order.Back()
		if last == nil {
			break
		}
		old := last.Value.(*statelessEntry)
		sp.order.Remove(last)
		delete(sp.entries, old.binding)
	}
	return srv
}

// matching lists the entries whose identity includes the subject.
func (sp *statelessPath) matching(guard string, userID uint) []*statelessEntry {
	sp.mu.Lock()
	defer sp.mu.Unlock()
	var out []*statelessEntry
	for _, e := range sp.entries {
		e.mu.Lock()
		hit := e.identity.has(guard, userID)
		e.mu.Unlock()
		if hit {
			out = append(out, e)
		}
	}
	return out
}

// evict drops the entries of a subject (a revocation or deletion) and
// closes their sessions, so a subscriptions/listen stream ends.
func (sp *statelessPath) evict(guard string, userID uint) {
	for _, e := range sp.matching(guard, userID) {
		sp.mu.Lock()
		if cur, ok := sp.entries[e.binding]; ok && cur == e {
			sp.order.Remove(e.elem)
			delete(sp.entries, e.binding)
		}
		sp.mu.Unlock()
		for ss := range e.server.Sessions() {
			_ = ss.Close()
		}
	}
}

// route sends a request to the stateless path when its protocol version
// header is 2026-07-28 or later, or when the first JSON-RPC method of a
// POST is server/discover; everything else goes to the stateful router
// (P-08).
func (p *Pier) route(w http.ResponseWriter, r *http.Request) {
	if p.stateless == nil {
		p.sessions.ServeHTTP(w, r)
		return
	}
	if pv := r.Header.Get(protocolVersionHeader); pv >= protocolVersion202607 {
		p.stateless.ServeHTTP(w, r)
		return
	}
	if r.Method == http.MethodPost && r.Body != nil {
		body, err := io.ReadAll(r.Body)
		if err != nil {
			var mbe *http.MaxBytesError
			if errors.As(err, &mbe) {
				http.Error(w, "request body too large", 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), methodDiscover) {
			p.stateless.ServeHTTP(w, r)
			return
		}
	}
	p.sessions.ServeHTTP(w, r)
}