Label audit backend in Serve handlers to avoid re-aggregation

[tgrunnagle]

Summary

On the vMCP Serve path, audit backend-enrichment resolved a request's backend by calling core.Lookup*, which — by the core's deliberate stateless design — re-aggregates every backend's capabilities on every audited request (O(backends·capabilities), plus a remote discovery fan-out) purely to label the audit event. This was a documented deferral from PR #5491 (P2.4) and is a hard blocker before audit can ship on a live Serve path.

• Why: Eliminate the per-request second aggregation that audit enrichment incurred on the Serve path, where the backend-enrichment middleware ran for every MCP request when AuditConfig != nil — re-reading the request body and performing a full backend aggregation just to attach a backend label.
• What: Label the audit event at the call site instead of in a middleware. The serving backend is already known where each per-session tool/resource SDK handler is built (at registration), so the name is resolved once there and the handler writes it into the audit BackendInfo that the audit middleware created upstream and reads after the handler returns. The backend-enrichment middleware is dropped from the Serve path entirely; it stays wired only on the legacy server.New path (s.core == nil) until P3.2 Reduce server.New body to the wrapper #5445 retires it along with the rest of the discovery path.

This removes a layer rather than optimizing around it — no per-request body re-read, no core.Lookup*, and no cache. AC1 (no second aggregation on audited Serve requests) holds structurally: labelling is an O(1) write of a pre-resolved string.

Closes #5493

Type of change

• Refactoring (no behavior change)

Test plan

• Unit tests (task test)
• Linting (task lint-fix)
• Manual testing (describe below)

Verification run on the touched package:

• go test -race ./pkg/vmcp/server/... passes.
• golangci-lint run ./pkg/vmcp/server/... reports 0 issues.
• go build ./... is clean.
• A repo-wide task lint-fix surfaces only a pre-existing, unrelated G115 in cmd/thv/app/upgrade.go (not touched by this PR).

Acceptance criteria verified by tests:

1. Audited Serve-path requests no longer trigger a second backend aggregation — TestServeHandlersLabelAuditBackend asserts the handler writes the backend label into the audit BackendInfo and that core.ListTools stays at its single registration-time call (labelling never re-aggregates).
2. The documented hard-blocker for live-Serve audit is resolved.

Changes

File	Change
pkg/vmcp/server/serve_handlers.go	coreToolHandler / coreResourceHandler take the pre-resolved backend name and write it into the audit BackendInfo at handler entry; coreSessionTools / coreSessionResources resolve the name once (from the existing per-session aggregation) and capture it in the handler closure.
pkg/vmcp/server/backend_enrichment.go	Enrichment middleware is now legacy-only: new withBackendEnrichment helper wires it only when s.core == nil; backendEnrichmentMiddleware / resolveBackendName become free functions (they no longer read Server state).
pkg/vmcp/server/server.go	Handler delegates the enrichment wiring to withBackendEnrichment (gated to the legacy path).
pkg/vmcp/server/serve_session_test.go	Adds TestServeHandlersLabelAuditBackend (handlers label the audit event; no re-aggregation); handler call sites updated for the new backendName parameter.
pkg/vmcp/server/backend_enrichment_test.go	Legacy-path coverage retained; the Serve-path enrichment test is removed (the middleware no longer runs there).

Does this introduce a user-facing change?

No. This is an internal change. For a registered backend the audit backend_name is identical to the legacy path (both record backend.Name). The one deliberate divergence is the orphan-backend edge case — a backend advertised by the core but absent from the registry — where the Serve path records the raw BackendID while the legacy minimal-target fallback leaves it empty (see Special notes); recording an identifier rather than dropping it is the intended, arguably-better behavior.

Special notes for reviewers

• Commit history / evolution: the first three commits implemented issue P2.4b vMCP Serve: cache per-session advertised set to avoid audit re-aggregation #5493's prescribed per-session advertised-set cache; the final commit (Label audit backend in Serve handlers, drop cache) reworks it to call-site labelling per review feedback — the cache was compensating complexity for a misplaced middleware layer rather than removing it. The net diff is the handler-labelling design; the cache and its tests are gone.
• Why call-site labelling is correct (verified against the code before refactoring):
  • The audit middleware (auditor.Middleware) creates the *BackendInfo, calls next.ServeHTTP, and reads BackendName only afterward — and it already creates the struct "if not already present", so dropping the enrichment middleware is safe.
  • mcp-go (v0.54.1) threads the HTTP request context into the handler: handlePost builds ctx := s.server.WithContext(r.ctx(), session) → HandleMessage → handleToolCall → finalHandler(ctx, request). This is the same context path coreToolHandler already uses to read auth.IdentityFromContext, so the BackendInfo pointer is reachable and the post-handler read observes the handler's write.
• Deviation from P2.4b vMCP Serve: cache per-session advertised set to avoid audit re-aggregation #5493: the issue prescribes a per-session cache. This PR meets the issue's goal (no second aggregation on audited Serve requests) more directly, without a cache. Flagging for the issue author since the prescribed solution changed.
• Orphan-backend audit divergence: when a backend is advertised by the core but absent from the registry, backendDisplayName records the raw BackendID; the legacy path's minimal-target fallback leaves WorkloadName empty. So Serve and legacy audit values differ in this edge case (pre-existing helper behavior, carried over). It's covered and locked in by TestBackendDisplayName. Audit-log consumers correlating across the two paths should expect an identifier (not empty) for an orphan backend on the Serve path.
• Follow-ups (out of scope, both LOW):
  • core.Lookup* now has no production callers — a candidate for retirement.
  • Cross-pod resources/read audit labels: only the tools mapping is re-injected on lazy re-hydration, so a resources/read audit on a non-owning replica records no backend label. Revisit when the live Serve composition root lands.
• Part of the vMCP core-extraction epic (Phase 2: Serve transport helper, re-home transport, replace discovery #5431).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 69.82%. Comparing base (0dd6069) to head (004ae22).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5512      +/-   ##
==========================================
+ Coverage   69.80%   69.82%   +0.01%     
==========================================
  Files         646      646              
  Lines       65771    65760      -11     
==========================================
+ Hits        45912    45917       +5     
+ Misses      16512    16499      -13     
+ Partials     3347     3344       -3     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[tgrunnagle]

Multi-agent review — Cache vMCP advertised set (#5512)

Reviewed by 6 specialist agents (concurrency/lifecycle, vMCP architecture & anti-patterns, Go style/comments, test coverage, security/audit, general quality). No HIGH-severity findings; no authorization, routing, or admission bypass. The design is sound and a near-textbook satisfaction of the vMCP anti-pattern rules it operates under: the optimization is evidence-backed and was explicitly deferred to this exact landing point (#9 satisfied), the abstraction is package-internal and touches no shared interface (#8), sessionID is threaded as an explicit parameter rather than via ctx.Value (#1 improved), and advertisedSet is immutable-after-construction with reconstruct-and-replace updates (#10).

The findings below are non-blocking for the current test-only Serve path, but several should be resolved before audit ships on a live Serve composition root.

Consensus findings (≥7/10)

#	Severity	Finding
F1	MEDIUM	lazyInjectSessionTools tools-only store can clobber a fuller cached entry on the owning replica for a zero-tool/has-resource session; the "never clobbers a fuller entry" comment is false in that case
F2	LOW	Lifecycle comment claims DELETE/TTL paths "do not notify the transport" — the mcp-go SDK in fact fires OnUnregisterSession on both; vMCP simply doesn't register the hook
F3	MEDIUM	Backend label is silently dropped on a cache miss (cross-pod / post-eviction); the path it replaces always resolved one — an audit-completeness regression gated by #5493's own "hard blocker before live audit"
F4	LOW	core.LookupTool/LookupResource/LookupPrompt now have zero production callers — dead interface surface (acknowledged out-of-scope; track + annotate)
F5	LOW	Test gaps: lazy-inject cache re-population and registration-failure eviction are both untested
F6	LOW	Orphan-backend (advertised but absent from registry) records the raw BackendID on Serve but empty WorkloadName on legacy — the "audit parity" comment overstates equivalence (pre-existing behavior, carried over)
F7	LOW	fakeCore.lookupResCalls/lookupPromptCalls counters are incremented but never asserted; the resources/read re-aggregation guard is missing

What this PR gets right (verified)

• AC1 is genuinely guarded: the "does not re-aggregate" test asserts lookupToolCalls==0/listToolsCalls==0 against a fresh fakeCore whose counters are actually wired — not vacuous.
• Cache is bounded (LRU, validated capacity; 0→default, negative rejected) and nil-receiver-safe (the legacy server.New path leaves it nil and calls evict unconditionally).
• Stores logical backend display names keyed by a server-generated UUID session ID — no internal addressing in shared state; a miss never affects enforcement (calls still route through enforceSessionBinding + core admission).

Dropped (< 7, for audit only)

const-coupling drift (F8), 22-field Server struct (F9, pre-existing), capacity divergence when audit is disabled (F10), forged-header audit-label confusion (F11, no auth impact), minor test nits (F12), lazy re-inject identity mismatch (F13, mitigated by binding check).

Overall: COMMENT — approve-with-suggestions. None of the findings block merge for a test-only path; F1/F3/F5 are worth resolving (or explicitly tracking against the live-audit blocker) before the production Serve composition root lands.

🤖 Generated with Claude Code

[jerm-dro]

blocker: This is the overuse of middleware coming home to roost, and the cache is complexity added to compensate for it rather than removing it. The backend-enrichment middleware checks the request method and no-ops for most of them, re-reads and re-Unmarshals the request body, and exists solely to stuff one value into BackendInfo context state — and now it needs a per-session cache to do that without re-aggregating on the Serve path. We're spending a cache, four eviction sites, an LRU backstop, and a capacity knob to keep a misplaced layer performant.

The backend name is produced at the call path, and the audit middleware emits after the handler returns (next.ServeHTTP → logAuditEvent, reading BackendName at emit time). So the handler can write the label directly into the BackendInfo the audit middleware already created — and the enrichment middleware comes off the Serve chain entirely.

// coreSessionTools — pass the name (already computed here) into the closure:
Handler: s.coreToolHandler(sessionID, domainTool.Name,
    s.backendDisplayName(ctx, domainTool.BackendID)),

// coreToolHandler — label where the call is handled:
func (s *Server) coreToolHandler(sessionID, toolName, backendName string) server.ToolHandlerFunc {
    return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
        if bi, ok := audit.BackendInfoFromContext(ctx); ok && bi != nil {
            bi.BackendName = backendName
        }
        // ... existing binding check + core.CallTool, unchanged ...
    }
}

Same one-liner in coreResourceHandler. Perf should be unchanged (the name is computed once at registration; per-request is an O(1) write, no core.Lookup*) and the log identical (same field, value, emitter).

Net effect: removes serve_advertised_set.go + tests, the advertisedSets field, all four evict calls + the LRU backstop, the resolveBackendName/cachedBackendName split, and drops the enrichment middleware from the Serve path. A layer comes out and no cache goes in. The middleware stays wired only for the legacy path (s.core == nil) until #5445 retires it.

This rests on two things I sketched but didn't implement: that the audit event is emitted after the handler runs, and that the SDK propagates the request context into the handler so BackendInfo is reachable there (coreToolHandler already reads identity from ctx, which suggests it does). Is this approach feasible? WDYT?

[tgrunnagle]

Implemented in bb48cfe. Feasible — confirmed both assumptions against the code before refactoring:

• Audit emits after the handler, with shared-pointer semantics. auditor.Middleware (auditor.go:191-254) creates the *BackendInfo, calls next.ServeHTTP, and reads BackendName only afterward (auditor.go:571). It already creates the struct "if not already present", so dropping the enrichment middleware is fine — the audit middleware itself supplies it (gated on AuditConfig != nil, server.go).
• The SDK threads r.Context() into the handler. handlePost builds ctx := s.server.WithContext(r.ctx(), session) → HandleMessage(ctx, …) → handleToolCall → finalHandler(ctx, request) (mcp-go v0.54.1). Same context path coreToolHandler already uses to read auth.IdentityFromContext, so the BackendInfo pointer is reachable and the post-handler read sees the write.

What landed:

• coreToolHandler/coreResourceHandler write the pre-resolved backend name into the audit BackendInfo at handler entry; the name is resolved once at registration (coreSessionTools/coreSessionResources) and captured in the closure.
• The enrichment middleware is gated to the legacy path (s.core == nil) via a small withBackendEnrichment helper and stays there until P3.2 Reduce server.New body to the wrapper #5445; backendEnrichmentMiddleware/resolveBackendName are now free functions (no Server state).
• Removed serve_advertised_set.go + tests, the advertisedSets field, the OnUnregisterSession hook, and all evict sites. Net −251 lines. AC1 (no second aggregation on audited Serve requests) now holds structurally — labelling is an O(1) write, no core.Lookup*. Verified by TestServeHandlersLabelAuditBackend (asserts the label is written and core.ListTools stays at its single registration-time call).

go test -race ./pkg/vmcp/server/... passes; golangci-lint run ./pkg/vmcp/server/... → 0 issues.

Note: this deviates from issue #5493's prescribed per-session cache. It achieves the issue's goal (no re-aggregation on audited Serve requests) more directly — flagging for the issue author since the prescribed solution changed.

[jerm-dro]

Approving — please add a direct test for backendDisplayName; its only coverage went away with the deleted Serve-path enrichment test (see inline).

[tgrunnagle]

Multi-agent review — Label audit backend in Serve handlers (#5512)

Reviewed by 5 specialist agents (vMCP architecture/anti-patterns, Go correctness/concurrency, audit completeness/security, test coverage, general quality). No HIGH-severity runtime-correctness or security findings — the call-site-labelling design is sound and verified end-to-end:

• The audit lifecycle guarantees the handler's write is observed: auditor.Middleware creates &BackendInfo{} and attaches it before next.ServeHTTP, then reads BackendName in addMetadata after the handler returns.
• The write is nil-safe and per-request (no data race); backendName is captured by value at registration via an O(1) registry lookup — no re-aggregation, so AC1 of #5493 holds structurally.
• No internal addressing leaks into the audit record; a label is never derived from a capability the caller can't reach; prompts are correctly out of scope (the Serve path never serves them).
• Removing the middleware (vs. caching around it) is the correct resolution of anti-patterns #2/#4/#8/#9 — the divergence from #5493's prescribed cache is well-justified.

The actionable surface is test coverage + comment hygiene. The headline item (F1) is this PR's approving reviewer's own outstanding request.

Consensus findings (≥7/10)

#	Severity	Finding
F1	HIGH	backendDisplayName has no direct test — the approving reviewer's explicit, unaddressed ask; its 3 branches are uncovered (the new test passes the name as a literal)
F2	MEDIUM	pkg/audit/auditor.go:570 comment is now false on the Serve path (detailed below — line is outside the diff, so noted here rather than inline)
F3	MEDIUM	The resource subtest lacks the listResourcesCalls re-aggregation guard the tool subtest has
F4	MEDIUM	No end-to-end test that a real BackendID flows through backendDisplayName to the audit label (literals bypass it)
F5	LOW	The documented orphan non-parity contradicts the PR's "no behavior change / identical to legacy" framing

F2 (line outside the diff — noted here, can't anchor inline): stale comment in pkg/audit/auditor.go:570

addMetadata reads backend_name under the comment "Backend info is populated by the backend enrichment middleware." After this PR, on the Serve path it's populated by the per-session handlers (serve_handlers.go) — the middleware isn't wired there. Suggest a source-agnostic wording, e.g. "populated upstream — by the enrichment middleware on the legacy path, or by the Serve-path session handlers." (go-style "Keep Comments Synchronized With Code", which this PR invokes elsewhere.)

Dropped (<7, for audit only)

coreSessionTools comment conflates ID source (aggregation) vs name source (registry); mutating a context-carried *BackendInfo is pre-existing anti-pattern #7 (accept as-is); core.Lookup* now has zero production callers (acknowledged follow-up); the local PR_SUMMARY.md describes the abandoned cache design (untracked — not part of this PR).

Overall: COMMENT — approve-with-one-required-test. The design is correct and already human-approved; F1 (the reviewer's ask), ideally with F3 + F4, is closed by a small test addition.

🤖 Generated with Claude Code

[tgrunnagle]

F2 (stale auditor.go comment) addressed in 68ca785 — reworded to be source-agnostic: "BackendInfo is populated upstream: by the backend-enrichment middleware on the vMCP legacy path, or by the vMCP Serve-path session handlers (which label the request directly)."