fix(pier): reap stateful sessions whose store record expired

- a reaper started at Mount checks every local session against the store
  every min(idle_ttl/2, 1m) and closes each copy whose record is gone,
  logging the expired step (A-CR-01, A-IN-06)
- the GET stream re-check ends the stream when the record is gone (A-IN-01)
- tool calls served without a session on the stateful path log the
  session as ephemeral (A-IN-06)
- TestPierReaperSmoke and the ReapNow, LocalSessions, StopReaper seams
This commit is contained in:
Jakub Zych
2026-10-23 15:50:04 +02:00
parent d18d8f0dc2
commit 2bf9dca603
6 changed files with 190 additions and 8 deletions

View File

@@ -15,6 +15,7 @@ import (
"sync/atomic"
"time"
"github.com/modelcontextprotocol/go-sdk/auth"
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
"github.com/modelcontextprotocol/go-sdk/mcp"
)
@@ -34,6 +35,10 @@ const (
methodDiscover = "server/discover"
)
// 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"
@@ -171,6 +176,14 @@ func (sr *sessionRouter) serveNew(w http.ResponseWriter, r *http.Request, id Ide
// 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
@@ -384,6 +397,14 @@ func (sr *sessionRouter) serveStream(w http.ResponseWriter, r *http.Request, ls
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)
}
}
@@ -411,6 +432,66 @@ func (sr *sessionRouter) closeLocal(sid string) {
}
}
// 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.
@@ -570,7 +651,7 @@ func containsPrompt(prompts []*prompt, name string) bool {
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 {
if req.Session != nil && req.Session.ID() != "" {
info.session = req.Session.ID()
}
return p.dispatch.call(ctx, id, info, t, req.Params.Arguments), nil