summercms/modules/pier/stateless.go

package pier

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

	"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
	// active counts the requests being served on server (a
	// subscriptions/listen stream for its whole life); guarded by
	// statelessPath.mu. The LRU never evicts an active entry (RA-WR-04).
	active int

	// mu serializes re-filtering and guards fingerprint, widened, tools and
	// prompts (see localSession); identity has its own box.
	mu          sync.Mutex
	identity    identityBox
	fingerprint string // of the identity the tool set was last filtered for
	// widened lists the other fingerprints whose tools and prompts were
	// added since (RA-WR-02); a recompute clears it.
	widened map[string]bool
	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

// entryUseKey keys the entryUse of one stateless request on its context.
type entryUseKey struct{}

// entryUse lists the entries one stateless request is served by, so
// ServeHTTP releases them when the request ends; guarded by
// statelessPath.mu.
type entryUse struct {
	entries []*statelessEntry
}

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
}

// ServeHTTP serves one stateless request. A subscriptions/listen POST is
// long-lived, so like the stateful GET stream it re-verifies its bearer
// every recheck interval and ends when the bearer no longer resolves to
// the binding that opened it (A-WR-06).
func (sp *statelessPath) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	use := &entryUse{}
	r = r.WithContext(context.WithValue(r.Context(), entryUseKey{}, use))
	defer sp.release(use)
	if r.Method != http.MethodPost || r.Body == nil {
		sp.handler.ServeHTTP(w, r)
		return
	}
	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))
	id, ok := identityFromRequest(r)
	if !ok || !slices.Contains(peekMethods(body), methodListen) {
		sp.handler.ServeHTTP(w, r)
		return
	}
	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()
	raw := bearerOf(r)
	go func() {
		ticker := time.NewTicker(sp.p.recheck)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				cur, err := sp.p.identities.resolve(ctx, raw)
				if err != nil || cur.Binding != id.Binding {
					logSession(logger(), "stream_closed", "", id.Binding)
					cancel()
					return
				}
			}
		}
	}()
	sp.handler.ServeHTTP(w, r.WithContext(ctx))
}

// 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
	}
	if e := sp.hit(r, id.Binding); e != nil {
		sp.refresh(r.Context(), e, id)
		return e.server
	}
	// Built with no lock held: the guard checks may query a database
	// (A-WR-07). A concurrent request of the same binding may have
	// inserted an entry meanwhile; that one is adopted.
	srv, tools, prompts := sp.p.newServer(r.Context(), id, nil)
	e := &statelessEntry{binding: id.Binding, server: srv, identity: identityBox{id: id}, fingerprint: id.fingerprint(), tools: tools, prompts: prompts}
	sp.mu.Lock()
	if cur, ok := sp.entries[id.Binding]; ok {
		sp.order.MoveToFront(cur.elem)
		sp.track(r, cur)
		sp.mu.Unlock()
		sp.refresh(r.Context(), cur, id)
		return cur.server
	}
	e.elem = sp.order.PushFront(e)
	sp.entries[id.Binding] = e
	sp.track(r, e)
	// Past capacity the least recently used idle entries go. An entry that
	// is serving a request (a quiet subscriptions/listen stream, a tool call
	// in flight) is never evicted, so capacity never cuts a stream or a call
	// off and a revocation still finds every open stream (A-WR-06, RA-WR-04);
	// while every entry is busy the LRU stays above capacity.
	var evicted []*statelessEntry
	for el := sp.order.Back(); el != nil && sp.order.Len() > sp.capacity; {
		prev := el.Prev()
		if old := el.Value.(*statelessEntry); old.active == 0 {
			sp.order.Remove(el)
			delete(sp.entries, old.binding)
			evicted = append(evicted, old)
		}
		el = prev
	}
	sp.mu.Unlock()
	for _, old := range evicted {
		closeSessions(old.server)
	}
	return srv
}

// hit answers the entry of a binding, most recently used now and counted as
// serving r, or nil.
func (sp *statelessPath) hit(r *http.Request, binding string) *statelessEntry {
	sp.mu.Lock()
	defer sp.mu.Unlock()
	e, ok := sp.entries[binding]
	if !ok {
		return nil
	}
	sp.order.MoveToFront(e.elem)
	sp.track(r, e)
	return e
}

// track counts e as serving r until ServeHTTP releases r's entries; the
// caller holds sp.mu. A request reaching getServer without ServeHTTP's
// entryUse is not counted.
func (sp *statelessPath) track(r *http.Request, e *statelessEntry) {
	use, _ := r.Context().Value(entryUseKey{}).(*entryUse)
	if use == nil || slices.Contains(use.entries, e) {
		return
	}
	e.active++
	use.entries = append(use.entries, e)
}

// release ends the request's use of its entries.
func (sp *statelessPath) release(use *entryUse) {
	sp.mu.Lock()
	for _, e := range use.entries {
		e.active--
	}
	use.entries = nil
	sp.mu.Unlock()
}

// refresh stores the request's identity on an entry and, when the
// identity's fingerprint is new to the entry, adds the tools and prompts
// it allows (A-WR-04). It never removes any (RA-WR-02): one server serves
// every live token of the binding, and two of them may differ in reach (two
// devices, an older narrower token next to a re-consented one), so removing
// on every alternation would pull a tool from under another token's call in
// flight. The receiving middleware hides, per request, whatever the
// request's identity may not use; a fan-out recompute still narrows.
func (sp *statelessPath) refresh(ctx context.Context, e *statelessEntry, id Identity) {
	e.mu.Lock()
	defer e.mu.Unlock()
	e.identity.set(id)
	fp := id.fingerprint()
	if fp == e.fingerprint || e.widened[fp] {
		return
	}
	ctx, cancel := sp.p.guardContext(ctx)
	defer cancel()
	tools, prompts, err := sp.p.catalog.allowedReport(ctx, id)
	if err != nil {
		// Not recorded as widened: the next request retries (RA-WR-03).
		logger().Warn("pier: widening deferred", "binding", id.Binding, "error", err)
		return
	}
	addMissing(e.server, sp.p, id, tools, prompts, e.tools, e.prompts)
	if e.widened == nil {
		e.widened = map[string]bool{}
	}
	e.widened[fp] = true
}

// closeSessions closes every session of a server.
func closeSessions(srv *mcp.Server) {
	for ss := range srv.Sessions() {
		_ = ss.Close()
	}
}

// matching lists the entries whose identity includes the subject.
func (sp *statelessPath) matching(guard string, userID uint) []*statelessEntry {
	return sp.matchingFunc(func(id Identity) bool { return id.has(guard, userID) })
}

// matchingFunc lists the entries whose identity satisfies match, matched
// after the path lock is released (see sessionRouter.matchingFunc).
func (sp *statelessPath) matchingFunc(match func(Identity) bool) []*statelessEntry {
	sp.mu.Lock()
	all := make([]*statelessEntry, 0, len(sp.entries))
	for _, e := range sp.entries {
		all = append(all, e)
	}
	sp.mu.Unlock()
	var out []*statelessEntry
	for _, e := range all {
		if match(e.identity.get()) {
			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) {
	sp.evictFunc(func(id Identity) bool { return id.has(guard, userID) })
}

// evictFunc drops the entries whose identity satisfies match and closes
// their sessions.
func (sp *statelessPath) evictFunc(match func(Identity) bool) {
	for _, e := range sp.matchingFunc(match) {
		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()
		closeSessions(e.server)
	}
}

// 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) {
	// The first request on any path opens a Redis store that the mount
	// warm-up has not opened yet, without waiting for it, so an instance
	// serving only stateless clients still subscribes to fan-out (A-WR-05).
	if ls, ok := p.store.(*lazyStore); ok {
		ls.kick()
	}
	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)
}