Jakub Zych
afc7510d1b
The stateful GET stream and the stateless subscriptions/listen stream capture the client address when they open and resolve every periodic re-check with towel.WithClientIP, as the bearer middleware does, so a failing re-check (a revoked admin token, for example) is logged with the remote address. TestStreamRecheckClientIP records the address the verifier sees on every verification of both streams.
772 lines
24 KiB
Go
772 lines
24 KiB
Go
package pier
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"slices"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"git.golem15.com/golem15/summercms/modules/towel"
|
|
"github.com/modelcontextprotocol/go-sdk/auth"
|
|
"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"
|
|
methodListen = "subscriptions/listen"
|
|
)
|
|
|
|
// ephemeralSession is the session value of the audit line of a tool call
|
|
// served on the stateful path without a session (serveEphemeral).
|
|
const ephemeralSession = "ephemeral"
|
|
|
|
// 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 serializes re-filtering (widening, fan-out recomputes) and guards
|
|
// fingerprint, tools and prompts; it may be held across guard checks,
|
|
// bounded by mcp.call_timeout. identity has its own box, so fan-out
|
|
// matching never waits on mu.
|
|
mu sync.Mutex
|
|
identity identityBox
|
|
fingerprint string // of the identity the tool set was last filtered for
|
|
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()
|
|
sid := rand.Text()
|
|
rec := Record{Binding: id.Binding, Created: time.Now().UTC()}
|
|
// The slot is reserved atomically before the server is built
|
|
// (A-WR-03); a store without Reserve falls back to Count then Put.
|
|
res, reserves := sr.p.store.(reserver)
|
|
if reserves {
|
|
ok, err := res.Reserve(ctx, id.Binding, sid, rec, sr.p.cfg.MaxPerBinding, sr.p.cfg.IdleTTL)
|
|
if err != nil {
|
|
sr.p.storeUnavailable(w, err)
|
|
return
|
|
}
|
|
if !ok {
|
|
sr.tooMany(w)
|
|
return
|
|
}
|
|
} else {
|
|
n, err := sr.p.store.Count(ctx, id.Binding)
|
|
if err != nil {
|
|
sr.p.storeUnavailable(w, err)
|
|
return
|
|
}
|
|
if n >= sr.p.cfg.MaxPerBinding {
|
|
sr.tooMany(w)
|
|
return
|
|
}
|
|
}
|
|
release := func() {
|
|
if reserves {
|
|
_ = res.Release(context.WithoutCancel(ctx), id.Binding, sid)
|
|
} else {
|
|
_ = sr.p.store.Delete(context.WithoutCancel(ctx), sid)
|
|
}
|
|
}
|
|
// The server is built with no router lock held: its guard checks may
|
|
// query a database (A-WR-07). The id is new, so nothing else holds it.
|
|
ls, err := sr.connect(ctx, sid, id, nil)
|
|
if err != nil {
|
|
release()
|
|
logger().Error("pier: session connect failed", "error", err)
|
|
http.Error(w, "failed connection", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
ls.created = rec.Created
|
|
sr.mu.Lock()
|
|
sr.local[sid] = ls
|
|
sr.mu.Unlock()
|
|
if !reserves {
|
|
if err := sr.p.store.Put(ctx, sid, rec, 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), and the reserved slot frees.
|
|
release()
|
|
sr.closeLocal(sid)
|
|
return
|
|
}
|
|
sr.persist(ctx, ls)
|
|
}
|
|
|
|
// tooMany answers 429 for an initialize past mcp.sessions.max_per_binding.
|
|
func (sr *sessionRouter) tooMany(w http.ResponseWriter) {
|
|
writeJSON(w, http.StatusTooManyRequests, map[string]any{"error": "too_many_sessions", "max_per_binding": sr.p.cfg.MaxPerBinding})
|
|
}
|
|
|
|
// 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) {
|
|
// The audit line of a tool call served here names the session
|
|
// "ephemeral", not "stateless": it ran on the stateful path without a
|
|
// session. The token info is built per request, so it is ours to mark.
|
|
if ti := auth.TokenInfoFromContext(r.Context()); ti != nil && ti.Extra != nil {
|
|
info, _ := ti.Extra[extraCall].(callInfo)
|
|
info.session = ephemeralSession
|
|
ti.Extra[extraCall] = info
|
|
}
|
|
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]
|
|
sr.mu.Unlock()
|
|
if ls == nil {
|
|
if len(rec.State) == 0 {
|
|
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 {
|
|
http.Error(w, "session not found", http.StatusNotFound)
|
|
return nil, false
|
|
}
|
|
// Rehydrate with no router lock held (A-WR-07), then insert unless
|
|
// a concurrent request of the same session got there first.
|
|
built, err := sr.connect(ctx, sid, id, &state)
|
|
if err != nil {
|
|
logger().Error("pier: session rehydrate failed", "error", err)
|
|
http.Error(w, "failed connection", http.StatusInternalServerError)
|
|
return nil, false
|
|
}
|
|
built.created = rec.Created
|
|
if state.InitializedParams != nil {
|
|
built.persisted.Store(true)
|
|
}
|
|
sr.mu.Lock()
|
|
if cur := sr.local[sid]; cur != nil {
|
|
sr.mu.Unlock()
|
|
_ = built.ss.Close()
|
|
ls = cur
|
|
} else {
|
|
sr.local[sid] = built
|
|
sr.mu.Unlock()
|
|
ls = built
|
|
logSession(logger(), "rehydrate", sid, id.Binding)
|
|
}
|
|
}
|
|
ls.mu.Lock()
|
|
ls.identity.set(id)
|
|
sr.p.widen(ctx, ls.server, id, &ls.fingerprint, ls.tools, ls.prompts)
|
|
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: identityBox{id: id}, fingerprint: id.fingerprint()}
|
|
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
|
|
}
|
|
rec := Record{Binding: ls.binding, State: raw, Created: ls.created}
|
|
stored := true
|
|
// The record is only rewritten, never created here: a session that ended
|
|
// while this request ran (DELETE, revocation, deletion, expiry) stays
|
|
// ended instead of being written back by its last request.
|
|
if u, ok := sr.p.store.(updater); ok {
|
|
stored, err = u.Update(ctx, ls.id, rec, sr.p.cfg.IdleTTL)
|
|
} else {
|
|
err = sr.p.store.Put(ctx, ls.id, rec, sr.p.cfg.IdleTTL)
|
|
}
|
|
if err != nil {
|
|
logger().Warn("pier: session state not stored", "session", sessionPrefix(ls.id), "error", err)
|
|
return
|
|
}
|
|
if !stored {
|
|
sr.closeLocal(ls.id)
|
|
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 re-check resolves with the
|
|
// stream's client address on the context, as the bearer middleware does, so
|
|
// a failed re-check is logged with it. 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)
|
|
ip := sr.p.clientIP(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(towel.WithClientIP(ctx, ip), raw)
|
|
if err != nil || id.Binding != ls.binding {
|
|
logSession(logger(), "stream_closed", ls.id, ls.binding)
|
|
cancel()
|
|
return
|
|
}
|
|
// A record that is gone (an ended event lost between instances,
|
|
// an expiry) ends the stream and the local copy.
|
|
if _, ok, err := sr.p.store.Get(ctx, ls.id); err == nil && !ok {
|
|
logSession(logger(), "stream_closed", ls.id, ls.binding)
|
|
cancel()
|
|
sr.closeLocal(ls.id)
|
|
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()
|
|
}
|
|
}
|
|
|
|
// reap closes, every interval until ctx ends, the local sessions whose
|
|
// store record has expired (P-24, A-CR-01). The store's TTL is the single
|
|
// source of truth across instances: a session another instance keeps alive
|
|
// still has its record, so this instance keeps its copy.
|
|
func (sr *sessionRouter) reap(ctx context.Context, every time.Duration) {
|
|
ticker := time.NewTicker(every)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-ticker.C:
|
|
sr.reapOnce(ctx)
|
|
}
|
|
}
|
|
}
|
|
|
|
// reapOnce runs one reaper pass: each local session whose record is gone is
|
|
// closed (its SDK server session ends, it leaves the local map and so the
|
|
// fan-out matching set) and logged as expired. A store error keeps the
|
|
// local copy (never close on doubt) and is logged once per pass. It returns
|
|
// how many sessions it closed.
|
|
func (sr *sessionRouter) reapOnce(ctx context.Context) int {
|
|
type entry struct{ id, binding string }
|
|
sr.mu.Lock()
|
|
entries := make([]entry, 0, len(sr.local))
|
|
for id, ls := range sr.local {
|
|
entries = append(entries, entry{id, ls.binding})
|
|
}
|
|
sr.mu.Unlock()
|
|
closed := 0
|
|
warned := false
|
|
for _, e := range entries {
|
|
if ctx.Err() != nil {
|
|
break
|
|
}
|
|
_, ok, err := sr.p.store.Get(ctx, e.id)
|
|
if err != nil {
|
|
if !warned {
|
|
logger().Warn("pier: session reaper could not read the store", "error", err)
|
|
warned = true
|
|
}
|
|
continue
|
|
}
|
|
if ok {
|
|
continue
|
|
}
|
|
sr.closeLocal(e.id)
|
|
logSession(logger(), "expired", e.id, e.binding)
|
|
closed++
|
|
}
|
|
return closed
|
|
}
|
|
|
|
// reapInterval is how often the reaper runs: half the idle TTL, at most a
|
|
// minute and at least a second.
|
|
func reapInterval(idleTTL time.Duration) time.Duration {
|
|
return max(min(idleTTL/2, time.Minute), time.Second)
|
|
}
|
|
|
|
// 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 && req.Session.ID() != "" {
|
|
info.session = req.Session.ID()
|
|
}
|
|
return p.dispatch.call(ctx, id, info, t, req.Params.Arguments), nil
|
|
}
|
|
}
|
|
|
|
// promptRenderFailed is the JSON-RPC error text of a prompt whose Render
|
|
// failed; the cause is logged, never sent (RA-WR-05).
|
|
const promptRenderFailed = "pier: the prompt could not be rendered"
|
|
|
|
// promptHandler checks the required arguments and renders the prompt. A
|
|
// Render error is logged server-side and answered with a fixed text, so a
|
|
// database or provider error never reaches the agent (as for dispatch).
|
|
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 {
|
|
logger().Warn("pier: prompt render failed", "prompt", pr.name, "error", err)
|
|
return nil, &jsonrpc.Error{Code: jsonrpc.CodeInternalError, Message: promptRenderFailed}
|
|
}
|
|
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)
|
|
}
|