summercms/modules/pier/events.go

package pier

import (
	"context"
	"sync"

	"git.golem15.com/golem15/summercms/modules/pact"
)

// listenEvents republishes every pact.MCPSubjectChanged fired on the
// application through the store, so every instance (this one included)
// applies it in handleEvent; when publishing fails the event is applied
// locally at least.
func (p *Pier) listenEvents() {
	if p.app == nil || p.app.Events == nil {
		return
	}
	p.app.Events.Listen(PluginID, func(ctx context.Context, ev *pact.MCPSubjectChanged) error {
		if ev == nil || ev.Kind == "" {
			return nil
		}
		event := Event{Kind: ev.Kind, Guard: ev.Guard, UserID: ev.UserID, Instance: p.instance, TokenID: ev.TokenID, ClientID: ev.ClientID}
		if err := p.store.Publish(ctx, event); err != nil {
			p.handleEvent(event)
		}
		return nil
	})
}

// handleEvent applies one event on this instance (D-26, D-27): a session
// ended elsewhere is closed here; a permission or link change recomputes
// the tools and prompts of every local session and stateless server of
// that subject, which makes the SDK send the list_changed notifications; a
// revocation or deletion closes those sessions, drops their records and
// evicts the stateless servers. A revocation naming a TokenID or a ClientID
// closes only the sessions that token or client opened (revokedMatch). A
// link event also drops the link cache.
func (p *Pier) handleEvent(ev Event) {
	ctx := context.Background()
	switch ev.Kind {
	case EventEnded:
		if ev.SessionID != "" && ev.Instance != p.instance {
			p.sessions.closeLocal(ev.SessionID)
		}
	case pact.MCPSubjectRevoked, pact.MCPSubjectDeleted:
		match := revokedMatch(ev)
		for _, ls := range p.sessions.matchingFunc(match) {
			_ = p.store.Delete(ctx, ls.id)
			p.sessions.closeLocal(ls.id)
			logSession(logger(), "revoked", ls.id, ls.binding)
		}
		if p.stateless != nil {
			p.stateless.evictFunc(match)
		}
	case pact.MCPSubjectLink, pact.MCPSubjectPermissions:
		if ev.Kind == pact.MCPSubjectLink {
			if ev.Guard == BackendGuard {
				p.identities.forgetAllLinks()
			} else {
				p.identities.forgetLink(ev.UserID)
			}
		}
		for _, ls := range p.sessions.matching(ev.Guard, ev.UserID) {
			p.recomputeSession(ctx, ls, ev.Kind == pact.MCPSubjectLink)
		}
		if p.stateless != nil {
			for _, e := range p.stateless.matching(ev.Guard, ev.UserID) {
				p.recomputeStateless(ctx, e, ev.Kind == pact.MCPSubjectLink)
			}
		}
	}
}

// revokedMatch selects the identities a revoked or deleted event ends
// (P-31). A revoked event with a TokenID matches the identities whose
// subject under ev.Guard is ev.UserID with that TokenID and no ClientID (the
// sessions of one personal token); one with a ClientID matches the subject
// with that ClientID (one OAuth client); otherwise, and for every deleted
// event, any identity that includes the subject.
func revokedMatch(ev Event) func(Identity) bool {
	if ev.Kind == pact.MCPSubjectRevoked && (ev.TokenID != 0 || ev.ClientID != "") {
		return func(id Identity) bool {
			s := id.subjectFor(ev.Guard)
			if s == nil || s.UserID != ev.UserID {
				return false
			}
			if ev.TokenID != 0 {
				return s.TokenID == ev.TokenID && s.ClientID == ""
			}
			return s.ClientID == ev.ClientID
		}
	}
	return func(id Identity) bool { return id.has(ev.Guard, ev.UserID) }
}

// recomputeSession re-filters a local session's identity (re-resolving
// the backend link after a link event) and applies the difference with
// AddTool, RemoveTools, AddPrompt and RemovePrompts.
func (p *Pier) recomputeSession(ctx context.Context, ls *localSession, relink bool) {
	ctx, cancel := p.guardContext(ctx)
	defer cancel()
	ls.mu.Lock()
	id := ls.identity.get()
	if relink {
		p.identities.refreshBackend(ctx, &id)
		ls.identity.set(id)
	}
	tools, prompts, err := p.catalog.allowedReport(ctx, id)
	applyDiff(ls.server, p, id, tools, prompts, ls.tools, ls.prompts)
	ls.fingerprint = recomputedFingerprint(id, err)
	ls.mu.Unlock()
}

// recomputeStateless does the same for one per-identity stateless server.
func (p *Pier) recomputeStateless(ctx context.Context, e *statelessEntry, relink bool) {
	ctx, cancel := p.guardContext(ctx)
	defer cancel()
	e.mu.Lock()
	id := e.identity.get()
	if relink {
		p.identities.refreshBackend(ctx, &id)
		e.identity.set(id)
	}
	tools, prompts, err := p.catalog.allowedReport(ctx, id)
	applyDiff(e.server, p, id, tools, prompts, e.tools, e.prompts)
	e.fingerprint = recomputedFingerprint(id, err)
	e.widened = nil
	e.mu.Unlock()
}

// recomputedFingerprint is the fingerprint a recompute records: the
// identity's, or none when a guard check failed (the failing items stay
// hidden, which is safe for an event), so the next request of the binding
// re-filters instead of keeping them hidden until another event (RA-WR-03).
func recomputedFingerprint(id Identity, err error) string {
	if err != nil {
		logger().Warn("pier: recompute incomplete, re-filtered on the next request", "binding", id.Binding, "error", err)
		return ""
	}
	return id.fingerprint()
}

// widen re-filters a stateful session's server for the identity of the
// request that reached it when that identity's fingerprint differs from
// the one the server was last filtered for (A-WR-04): a binding that
// returns with more scope, or with a backend subject added, gains the newly
// allowed tools and prompts at once, and one that returns with less loses
// them. A session serves one client; the stateless path, whose server every
// token of a binding shares, only adds (statelessPath.refresh). The caller
// holds the session's own mutex, never a router-wide one.
func (p *Pier) widen(ctx context.Context, srv serverFeatures, id Identity, fp *string, haveTools, havePrompts map[string]bool) {
	next := id.fingerprint()
	if *fp == next {
		return
	}
	// A client that hangs up mid-check must not leave the server filtered
	// by cancelled guard checks (an error hides the item), but the checks
	// still end within mcp.call_timeout: the caller holds a lock.
	ctx, cancel := p.guardContext(ctx)
	defer cancel()
	// A failed check is not a refusal: the server and *fp stay as they are,
	// so the next request retries instead of the tools staying hidden until
	// the fingerprint changes again (RA-WR-03).
	tools, prompts, err := p.catalog.allowedReport(ctx, id)
	if err != nil {
		logger().Warn("pier: widening deferred", "binding", id.Binding, "error", err)
		return
	}
	applyDiff(srv, p, id, tools, prompts, haveTools, havePrompts)
	*fp = next
}

// guardContext is the context of guard checks run while a session's or
// entry's mutex is held (widening and fan-out recomputes): not cancelled
// with the request that triggered them, but bounded by mcp.call_timeout,
// so a slow provider never holds that mutex indefinitely (RA-WR-01).
func (p *Pier) guardContext(ctx context.Context) (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.WithoutCancel(ctx), p.cfg.CallTimeout)
}

// addMissing adds the allowed tools and prompts a server lacks, updating
// the tracking maps in place; it removes nothing.
func addMissing(srv serverFeatures, p *Pier, id Identity, tools []*tool, prompts []*prompt, haveTools, havePrompts map[string]bool) {
	for _, t := range tools {
		if !haveTools[t.name] {
			srv.AddTool(t.desc, p.toolHandler(t, id))
			haveTools[t.name] = true
		}
	}
	for _, pr := range prompts {
		if !havePrompts[pr.name] {
			srv.AddPrompt(pr.desc, p.promptHandler(pr))
			havePrompts[pr.name] = true
		}
	}
}

// applyDiff brings a server's tool and prompt sets to the allowed ones,
// updating the tracking maps in place.
func applyDiff(srv serverFeatures, p *Pier, id Identity, tools []*tool, prompts []*prompt, haveTools, havePrompts map[string]bool) {
	wantTools := map[string]bool{}
	for _, t := range tools {
		wantTools[t.name] = true
		if !haveTools[t.name] {
			srv.AddTool(t.desc, p.toolHandler(t, id))
			haveTools[t.name] = true
		}
	}
	var dropTools []string
	for name := range haveTools {
		if !wantTools[name] {
			dropTools = append(dropTools, name)
			delete(haveTools, name)
		}
	}
	if len(dropTools) > 0 {
		srv.RemoveTools(dropTools...)
	}
	wantPrompts := map[string]bool{}
	for _, pr := range prompts {
		wantPrompts[pr.name] = true
		if !havePrompts[pr.name] {
			srv.AddPrompt(pr.desc, p.promptHandler(pr))
			havePrompts[pr.name] = true
		}
	}
	var dropPrompts []string
	for name := range havePrompts {
		if !wantPrompts[name] {
			dropPrompts = append(dropPrompts, name)
			delete(havePrompts, name)
		}
	}
	if len(dropPrompts) > 0 {
		srv.RemovePrompts(dropPrompts...)
	}
}

// matching lists the local sessions whose identity includes the subject.
func (sr *sessionRouter) matching(guard string, userID uint) []*localSession {
	return sr.matchingFunc(func(id Identity) bool { return id.has(guard, userID) })
}

// matchingFunc lists the local sessions whose identity satisfies match.
// The sessions are snapshotted under the router lock and matched after it
// is released, through each session's identity box, which is never held
// across I/O: an event never waits on a session that is re-filtering, and
// never while it holds the router lock (RA-WR-01).
func (sr *sessionRouter) matchingFunc(match func(Identity) bool) []*localSession {
	sr.mu.Lock()
	all := make([]*localSession, 0, len(sr.local))
	for _, ls := range sr.local {
		all = append(all, ls)
	}
	sr.mu.Unlock()
	var out []*localSession
	for _, ls := range all {
		if match(ls.identity.get()) {
			out = append(out, ls)
		}
	}
	return out
}

// identityBox holds the identity of a session or stateless entry behind a
// mutex of its own that is never held across I/O, so fan-out matching
// reads it without waiting for a guard check (RA-WR-01).
type identityBox struct {
	mu sync.Mutex
	id Identity
}

func (b *identityBox) get() Identity {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.id
}

func (b *identityBox) set(id Identity) {
	b.mu.Lock()
	b.id = id
	b.mu.Unlock()
}

// refreshBackend re-resolves an OAuth identity's backend subject through
// the link resolver (after a link event), adding or dropping it.
func (r *identityResolver) refreshBackend(ctx context.Context, id *Identity) {
	if id.TokenKind != TokenKindOAuth || id.Frontend == nil {
		return
	}
	id.Backend = nil
	for _, s := range id.Scopes {
		if s != "backend" {
			continue
		}
		if link, ok := r.linkedBackend(ctx, id.Frontend.UserID); ok {
			id.Backend = &pact.MCPSubject{
				Guard:    BackendGuard,
				UserID:   link.ID,
				Scopes:   id.Scopes,
				ClientID: id.Frontend.ClientID,
				TokenID:  id.Frontend.TokenID,
			}
		}
		return
	}
}

// forgetAllLinks drops the whole link cache.
func (r *identityResolver) forgetAllLinks() {
	r.mu.Lock()
	r.links = map[uint]linkEntry{}
	r.mu.Unlock()
}