package pier import ( "bytes" "context" "crypto/rand" "encoding/json" "errors" "fmt" "io" "net/http" "slices" "strings" "sync" "sync/atomic" "time" "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 sync.Mutex identity Identity 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 = 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: 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 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) 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(ctx, 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 } } // promptHandler checks the required arguments and renders the prompt. 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 { return nil, err } 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) }