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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user