P2.2 Move SDK hooks + two-phase session creation under Serve

[tgrunnagle]

Description

Relocate the three mcp-go SDK hooks and the two-phase session-creation wiring out of
server.New and into the new Serve transport helper. The OnRegisterSession hook
(capability registration) and the OnBeforeListTools/OnBeforeCallTool hooks (cross-pod
Redis re-hydration) — together with the transport session manager, the session data
storage backend, and the vMCP session manager — are SDK-lifecycle concerns that belong in
the transport layer of the New/Serve split, not in the core VMCP domain object.

Context

This is step P2.2 of Phase 2. #5439 introduced Serve(ctx, VMCP, *ServerConfig) -> *Server plus the mcp-go server, route mux, and HTTP lifecycle skeleton, but the SDK hooks
and session creation still live in server.New. This task moves them under Serve.

mcp-go's two-phase session creation is the canonical SDK-coupling anti-pattern (#5): the
session ID is sent to the client during initialize, and the OnRegisterSession hook then
fires after registration to attach the per-session capability set. Because the SDK only
fires OnRegisterSession during initialize, a session re-hydrated from Redis on a
different pod has an empty per-session tool store; the OnBeforeListTools/
OnBeforeCallTool hooks lazily re-inject the tools before the SDK handler reads them. All
of this — hooks, two-phase dance, heartbeat, session-ID manager, and the Redis path — must
stay isolated inside Serve/the adapter and never cross the VMCP boundary (architecture.md
Core Principle #1; research.md lines 221-226).

The session layer's "fixed at initialize" capability set, identity binding
(MetadataKeyIdentityBinding), and the cross-pod Redis re-hydration path stay in
Serve/session. server.New remains unchanged in observable behavior throughout Phase 2.

Parent Story: #5431
Dependencies: #5439 (needs the Serve skeleton + ServerConfig to register hooks into)
Blocks: #5442

Acceptance Criteria

• The three SDK hooks (AddOnRegisterSession -> handleSessionRegistration,
  AddBeforeListTools / AddBeforeCallTool -> lazyInjectSessionTools) are registered
  from within Serve (on the server.Hooks{} wired into the mcp-go server).
• Two-phase session creation wiring lives under Serve: the transport session manager
  (transportsession.NewManager), the session data storage backend
  (buildSessionDataStorage, memory/Redis incl. THV_SESSION_REDIS_PASSWORD), and the vMCP
  session manager (sessionmanager.New).
• Session-scoped state — the "fixed at initialize" capability set, identity binding, and
  the cross-pod Redis re-hydration path — remains in Serve/session; nothing about hooks,
  two-phase creation, or session lifecycle crosses the VMCP boundary (no mcp-go types on
  the interface).
• Session-management integration tests stay green.
• PR is ≤ 400 LOC and ≤ 10 files changed (excluding tests/docs/generated)
• server.New signature and observable behavior unchanged
• All tests pass (task test); lint clean (task lint-fix)
• Code reviewed and approved

Technical Approach

Recommended Implementation

Move the hook registration block and the session-creation wiring from server.New into
Serve (in pkg/vmcp/server/serve.go), keeping the *Server receiver methods
(handleSessionRegistration / handleSessionRegistrationImpl, lazyInjectSessionTools)
and the package-level helper setSessionToolsDirect (server.go:1048 — a package-level
func, not a *Server method) unchanged — they already operate against
s.vmcpSessionMgr and the SDK server.ClientSession, so relocation is mechanical.

Sequence within Serve (matching today's order in server.New):

1. Construct the server.Hooks{} (today at server.go:330) and pass it to NewMCPServer
   via server.WithHooks(hooks) (server.go:333-340) — these are co-located because the
   hook callbacks close over srv, so register them after the *Server is assembled.
2. Build transportsession.NewManager(cfg.SessionTTL, transportsession.NewStreamableSession)
   (server.go:413).
3. Build buildSessionDataStorage(ctx, cfg) (server.go:415 -> :265), preserving the
   memory/Redis branch and the THV_SESSION_REDIS_PASSWORD read, and keep the
   closeStorageOnErr defer-guard so the cleanup goroutine does not leak on an error return
   (go-style: pair acquisition with release).
4. Build sessionmanager.New(sessionDataStorage, sessMgrCfg, backendRegistry) (server.go:462)
   and append optimizerCleanup to shutdownFuncs.
5. Register the three hooks (server.go:491-511) against the *Server so they delegate to
   the receiver methods.

Keep this purely a relocation — do not change the hook callbacks, the two-phase semantics,
or the Redis re-hydration logic. The session-creation collaborators may be sourced from the
ServerConfig slots #5439 established (SessionFactory, SessionStorage, SessionTTL,
optimizer factory/config). Because server.New is not yet routed through Serve (that is
Phase 3), server.New keeps its own copy until P3.2; this task wires Serve so the hooks +
session creation execute identically when Serve is exercised directly.

Patterns & Frameworks

• mcp-go SDK boundary stays in Serve/the adapter — hooks, two-phase creation, and the
  THV_SESSION_REDIS_PASSWORD Redis path are transport concerns (RFC transport-extension
  case Implement secret store #4, "intentionally closed"; architecture.md Core Principle fix(typo): corrects readme #1).
• Anti-pattern Implement secret injection #5: the two-phase dance must remain invisible to the core; do not let
  SDK-specific session concepts leak onto the VMCP interface.
• Resource-leak guard: preserve the closeStorageOnErr defer so session data storage is
  closed on any error path out of Serve (.claude/rules/go-style.md "Resource Leaks").
• stdlib testing.T + testify in pkg/vmcp/server; no Ginkgo (architecture.md Core
  Principle Create proxy subcommand #7 / R8). Mocks via gomock + task gen.
• Conventions: .claude/rules/go-style.md, .claude/rules/vmcp-anti-patterns.md,
  .claude/rules/security.md.

Code Pointers

• pkg/vmcp/server/server.go:489-511 — the three hook registrations
  (AddOnRegisterSession, AddBeforeListTools, AddBeforeCallTool). Relocate into Serve.
• pkg/vmcp/server/server.go:330,333-340 — hooks := &server.Hooks{} and NewMCPServer(... WithHooks(hooks)). The hooks object and its mcp-go server are co-located with the hook
  registrations (callbacks close over srv).
• pkg/vmcp/server/server.go:1130 — handleSessionRegistrationImpl (and
  handleSessionRegistration at :1103): phase-2 of two-phase creation; calls
  vmcpSessionMgr.CreateSession + GetAdaptedTools/GetAdaptedResources. Stays a *Server
  method; no change.
• pkg/vmcp/server/server.go:1078 — lazyInjectSessionTools: cross-pod Redis re-hydration
  target of the OnBeforeListTools/OnBeforeCallTool hooks. Stays a *Server method; no
  change.
• pkg/vmcp/server/server.go:413 — transportsession.NewManager(cfg.SessionTTL, transportsession.NewStreamableSession). Relocate.
• pkg/vmcp/server/server.go:415 + :265 — buildSessionDataStorage (memory/Redis;
  THV_SESSION_REDIS_PASSWORD via vmcpconfig.RedisPasswordEnvVar). Relocate with the
  closeStorageOnErr guard.
• pkg/vmcp/server/server.go:462 — sessionmanager.New(sessionDataStorage, sessMgrCfg, backendRegistry) + optimizerCleanup appended to shutdownFuncs. Relocate.
• pkg/vmcp/server/serve.go — Serve + ServerConfig from P2.1 Serve skeleton + ServerConfig #5439; the relocation
  target.
• pkg/vmcp/server/server.go:301 — server.New (7-param signature; stays stable, behavior
  unchanged).
• research.md lines 221-226 — "Session-scoped concerns that MUST stay in Serve/session
  layer" (capabilities fixed at initialize, identity binding, cross-pod Redis storage,
  two-phase creation + SDK hooks + heartbeat + session-ID manager).
• architecture.md lines 285-290 ("Serve (transport) wiring relocated"), 300-303 ("Session
  layer … stays in Serve"), 318-321 ("fixed at initialize cache stays in Serve").
• pkg/vmcp/server/session_management_integration_test.go (907 lines),
  session_management_realbackend_integration_test.go (261) — the integration suite that
  must stay green; the parity gate for this task.

Component Interfaces

No interface changes. This is a relocation of existing wiring; the hook callbacks and the
session-creation collaborators keep their current shapes. The hooks attach to the SDK via
the mcp-go server.Hooks value already built in Serve (#5439) and delegate to the
*Server receiver methods:

// Inside Serve (pkg/vmcp/server/serve.go), after srv is constructed.
// hooks is the *server.Hooks passed to NewMCPServer(server.WithHooks(hooks)).

// Phase-2 of two-phase creation: fires AFTER the SDK registers the session
// (the Mcp-Session-Id has already been sent to the client during initialize).
hooks.AddOnRegisterSession(func(ctx context.Context, session server.ClientSession) {
    srv.handleSessionRegistration(ctx, session)
})

// Cross-pod Redis re-hydration: lazily re-inject tools when a session re-hydrated
// from Redis on another pod has an empty per-session SDK tool store.
hooks.AddBeforeListTools(func(ctx context.Context, _ any, _ *mcp.ListToolsRequest) {
    srv.lazyInjectSessionTools(ctx)
})
hooks.AddBeforeCallTool(func(ctx context.Context, _ any, _ *mcp.CallToolRequest) {
    srv.lazyInjectSessionTools(ctx)
})

// Session-creation wiring relocated from server.New (memory/Redis backend selection
// stays here; THV_SESSION_REDIS_PASSWORD read inside buildSessionDataStorage).
sessionManager := transportsession.NewManager(ttl, transportsession.NewStreamableSession)
sessionDataStorage, err := buildSessionDataStorage(ctx, cfg) // closeStorageOnErr guard
vmcpSessMgr, optimizerCleanup, err := sessionmanager.New(sessionDataStorage, sessMgrCfg, backendRegistry)

Testing Strategy

Unit Tests

• lazyInjectSessionTools is a no-op when the per-session SDK tool store already has
  tools (normal pod-local case), and injects from the vMCP session manager when the store is
  empty (cross-pod case) — mirror the existing *Server-method coverage after relocation.
• buildSessionDataStorage selects in-process storage for nil/empty/"memory" provider
  and Redis (reading THV_SESSION_REDIS_PASSWORD) for "redis"; returns an error for any
  other provider.

Integration / Behavioral Parity Tests

• session_management_integration_test.go and
  session_management_realbackend_integration_test.go stay green: session registration,
  tools/list and tools/call after initialize, and the cross-pod Redis re-hydration path
  behave identically with hooks + session creation driven from Serve.

Edge Cases

• On an error returned from Serve after the session data storage is built, the
  closeStorageOnErr guard closes the storage so the cleanup goroutine does not leak.
• A session re-hydrated from Redis on a second pod (no initialize, empty SDK tool
  store) gets its tools re-injected by the before-list/before-call hooks.
• Identity binding (MetadataKeyIdentityBinding) and the "fixed at initialize"
  capability set remain enforced by the session layer; none of it crosses the VMCP
  boundary.

Out of Scope

• Replacing the discovery-into-context aggregation path with direct VMCP.ListTools/
  CallTool calls — that is P2.4 Replace discovery-into-context with direct VMCP calls #5442 (which this task blocks).
• Moving the HTTP middleware chain — that is P2.3 Move middleware chain under Serve; remove authz + annotation mw #5441.
• Moving the AS runner, status reporter, optimizer, and health monitor — that is P2.5 Move AS runner, status reporter, optimizer, health monitor under Serve #5443.
• Reducing server.New to the Serve(ctx, New(...), ...) wrapper — that is Phase 3
  (P3.1 deriveCoreConfig/deriveServerConfig config split #5444 / P3.2).
• Removing the legacy MetadataKeyTokenHash/Salt (superseded by
  MetadataKeyIdentityBinding, vMCP terminates MCP session on every legitimate access-token refresh #5306) — out of this epic.

References

• RFC THV-0076: /Users/trey/Documents/GitHub/stacklok/toolhive-rfcs/rfcs/THV-0076-vmcp-core-interface.md
• Parent story: Phase 2: Serve transport helper, re-home transport, replace discovery #5431
• Epic: vMCP interface refactor #5419