feat(mcp): per-user connection lifecycle — pool + routing seam + materialized type=user servers (#317 increments 1–3)

[initializ-mk]

Stacked on #327 (base: feat/mcp-delegated-user-resolver — it uses delegatedSubject from there). Rebase onto main once #327 merges.

First increment of #317's per-user connection lifecycle (design recorded on the issue). #327 landed per-user tokens on a shared connection; a session-stateful MCP server binds identity at initialize, so full per-user isolation needs one connection per user. The investigation settled that per-subject sessions require per-subject Client instances (initialize is Client-driven) — i.e. a pool.

What

subjectConnPool — one lazily-established MCP Client per requesting-user subject (keyed off auth.IdentityFromContext(ctx)), implementing the ClientResolver seam the tool-call path will use to route each call to the right user's connection.

• Lazy — nothing connects until a user's first call; no user in ctx → ErrNoToken (never at startup).
• Identity-preserving, cancel-decoupled establish — the connect ctx carries the subject's identity (so authFn resolves their token at initialize) but is background-derived with a hard timeout, so one caller's cancel can't tear down the shared establish (the B2 lesson from the OAuth refresh path).
• Single-flighted per subject — a burst of a user's first calls opens exactly one connection; distinct subjects never block each other.
• Lifecycle — Evict (drop+close on a connection error → next call re-establishes) and Close (tear down all).

Scope

Self-contained library component, no hot-path wiring yet — deliberately, because the tool-call path is core and a bug there breaks every MCP tool. The documented increments after this:
2. ClientResolver routing seam — ToolHandle/MCPTool resolve the client per-call from the ctx subject (backward-compatible: a static resolver preserves today's shared-connection behavior for non-per-user servers).
3. Materialized tool schemas — so a type: user server (no connection at startup) registers its tools from platform-materialized config, then connects per-user lazily only for calls.

Tests

Per-subject lazy establish + reuse (distinct users → distinct connections, one connect each), no-subject → ErrNoToken, single-flight (20 concurrent first-calls → one connect), evict-reconnects, close-teardown, and identity-preserved-on-the-connect-ctx. Full forge-core mcp suite passes; golangci-lint 0 issues (the component is exercised by tests, not flagged unused).

[initializ-mk]

Review: per-subject connection pool (#317 connection lifecycle, increment 1)

This is the follow-up the #327 review's finding-1 pointed at — connection-level isolation for session-stateful MCP servers — and it's a well-built concurrency primitive. I traced the tricky parts and they're correct:

• Per-subject isolation is sound. Both conns and inFly are keyed by subject; distinct users get distinct connections, no shared-connection path.
• The B2-lesson connect ctx is exactly right (the crux of this PR): the establish ctx is context.Background-derived + hard-timeout-bounded (so a caller's cancel can't tear down the shared establish) AND carries the subject's identity via auth.WithIdentity (so authFn resolves their token at initialize). Those two requirements pull in opposite directions and the implementation threads them correctly. And since subject != "" already implies a non-nil identity, the identity is always preserved.
• Single-flight per subject is correct — concurrent first-calls find the same inFly[subject] flight and wait on one fl.done; one connect. The channel close provides the happens-before barrier, so reading fl.client/fl.err after <-fl.done is race-free even though they're written outside the lock.
• Caller-cancel is decoupled from the establish — a waiter that cancels returns ctx.Err() while the goroutine keeps going and still caches the connection for the next caller (no lost work, no leak).
• Close-during-establish is handled — the goroutine checks p.closed and tears down a connection it just made, so no leak.
• Staging discipline is good: no hot-path wiring yet, exactly because a bug on the tool-call path breaks every MCP tool.

No blockers. Two robustness minors, a test gap, and a process note:

1. Minor (robustness): a panic in connect poisons the subject's flight

The establish goroutine has no defer guaranteeing close(fl.done) and delete(p.inFly, subject). If p.connect (factory + Initialize) panics — a nil-deref in a transport, say — close(fl.done) never runs (waiters hang until their own ctx cancels) and inFly[subject] is never cleared, so every future call for that subject also hangs on the dead flight. Most singleflight implementations guard this. Recommend a deferred cleanup: recover() → set fl.err to a panic error, delete(inFly, subject) under lock, close(fl.done) — so a panicking factory fails that one establish cleanly instead of wedging the slot. Inline.

2. Minor: Close doesn't join in-flight establishes

A connection being established concurrently with Close is torn down by its own goroutine (it checks p.closed) — so no leak — but Close returns before that happens, so "all connections closed" isn't guaranteed synchronously at Close return. For deterministic shutdown/resource release, a sync.WaitGroup over the establish goroutines would close the gap. Inline.

3. Test gap: the connect-error and panic paths are untested

Coverage is good for reuse / no-subject / single-flight / evict / close, but the connect returns error branch ("leave nothing cached" — ClientFor returns the err, len()==0, a retry re-establishes) has no test, and neither does the panic case (finding 1). The error path is a cheap add and pins that a failed establish doesn't cache a bad connection.

4. Process: no CI runs on this branch

gh pr checks reports no checks — because the base is feat/mcp-delegated-user-resolver (a feature branch), and the CI workflow triggers on PRs targeting main/develop. So "full mcp suite passes / lint clean" is local-only, same situation as #276. Retargeting to main after #327 merges (as the PR plans) is what puts CI on record — worth doing before merge so the concurrency tests actually run in CI, given this is a concurrency primitive where CI's -race matters.

Verdict

Merge-ready after finding 1 — the panic guard is the one I'd fold in, since a wedged per-subject slot (hanging every future call for that user) is a nastier failure than the establish error it's meant to survive, and it's a few lines. Finding 2 and the test/CI items are follow-ups, though I'd get the connect-error test + real CI (-race) before this lands, since it's concurrency-critical. The core design — per-subject isolation, single-flight, and the identity-preserving/cancel-decoupled establish — is correct and cleanly implements the connection-level isolation #327 deferred.

[initializ-mk]

Minor (robustness): guard this goroutine against a connect panic.

If p.connect(connCtx) panics, close(fl.done) and delete(p.inFly, subject) below never run — so waiters hang on fl.done (until their own ctx cancels) and inFly[subject] stays populated forever, wedging every future call for that subject on a dead flight. A factory + Initialize over a transport is exactly the kind of code that can nil-deref.

Suggest a deferred cleanup so a panic fails just this one establish cleanly:

defer func() {
    if r := recover(); r != nil {
        p.mu.Lock(); delete(p.inFly, subject); p.mu.Unlock()
        fl.client, fl.err = nil, fmt.Errorf("%w: connect panicked: %v", ErrProtocolError, r)
        close(fl.done)
    }
}()

(or restructure so delete(inFly) + close(fl.done) are unconditionally deferred). Most singleflight impls guard this.

[initializ-mk]

Fixed. The establish goroutine now defers all cleanup (clear inFly, cache-or-close, close fl.done) behind a recover, so a panicking connect fails just that one establish with an error instead of wedging the subject's flight. Test TestSubjectConnPool_ConnectPanicDoesNotWedge asserts the panic surfaces as an error (not a hang) and the subject re-establishes on the next call.

[initializ-mk]

Minor: Close tears down conns but doesn't join in-flight establish goroutines. A connection being established concurrently with Close is closed by its own goroutine (it checks p.closed), so there's no leak — but Close returns before that happens, so callers can't assume all connections are closed when Close returns. If deterministic teardown matters (resource release before process exit), track the establish goroutines with a sync.WaitGroup and wg.Wait() here. Non-blocking; just noting the shutdown isn't fully synchronous.

[initializ-mk]

Done — added a sync.WaitGroup over establishes. Close marks the pool closed, wg.Wait()s the in-flight establishes (each self-closes on seeing closed), then closes the cached connections — so teardown is synchronous: nothing open and no establish still running when Close returns.

[initializ-mk]

Folded in the robustness findings (8aeb3e0):

• Finding 1 (panic guard) — the establish goroutine's cleanup (clear inFly[subject], cache-or-close, close(fl.done)) is now unconditionally deferred behind a recover, so a panicking connect fails that one establish cleanly instead of wedging the subject's flight and hanging every future call. Test asserts panic → error (not hang) + subject re-establishes.
• Finding 2 (synchronous Close) — Close now wg.Wait()s the in-flight establishes (they self-close on seeing closed) before closing cached conns, so teardown is deterministic.
• Finding 3 (tests) — added the connect-error path (surfaces error, caches nothing, retry re-establishes) and the panic path. Whole pool suite passes under -race.
• Finding 4 (CI) — can't run real CI yet: base is feat/mcp-delegated-user-resolver (feat(mcp): delegated per-user token resolution for type=user (#317) #327), which is still open, and CI triggers on main. I'll retarget feat(mcp): per-user connection lifecycle — pool + routing seam + materialized type=user servers (#317 increments 1–3) #329 to main once feat(mcp): delegated per-user token resolution for type=user (#317) #327 merges so the concurrency tests run under CI's -race before this lands — noting it so it's not forgotten.

[initializ-mk]

Retargeted to main (rebased off #327, now merged) and added increment 2 — the ClientResolver routing seam — to this PR.

Increment 2 (ab752fd) — the tool-call path now resolves the MCP Client per call from a ClientResolver instead of holding one fixed Client. Behavior-preserving: Manager.Tools() hands every ToolHandle a StaticResolver over the shared client, so existing servers call the exact same connection as before. The per-subject pool (increment 1) is the other ClientResolver implementation — increment 3 wires it in for type: user servers (with materialized-schema tool registration) with no further tool-adapter change.

• MCPTool.Execute resolves per call; a resolver error surfaces as a tool error (not a nil-deref); nil resolver falls back to the fixed client.
• Tests: two users' calls route to their own clients; resolver-error surfaces. Full mcp suite green under -race.

So #329 now = increment 1 (pool) + increment 2 (routing seam). Increment 3 (materialized schemas + lazy type: user server + pool-as-resolver) is the next PR.

[initializ-mk]

Added increment 3 — the per-user connection lifecycle now works end-to-end. This is the increment that closes the #327 finding-1 caveat (per-user connection isolation, not just per-call token).

Increment 3 (d42686f):

• Materialized schemas (tools.schemas) — a type: user server has no user (and no connection) at startup, so it can't tools/list. The platform materializes the tool schemas into config; Forge registers them without a live connection (validated the same way discovered tools are).
• Lazy type: user server — a new runMaterialized path reaches Ready without eager-connecting (Configured → Ready); the server exposes the increment-1 subjectConnPool as its per-call ClientResolver. Each user's first call establishes their own connection whose initialize runs under their token.
• Manager.Tools() includes a pool server (which has no shared client) and hands its ToolHandle the pool resolver — via the increment-2 seam, no adapter change.

The vertical, proven under -race: a type: user server registers its tools with zero startup connections, then a call for alice@corp.com establishes a per-user connection whose initialize carries Bearer tok-alice@corp.com — the per-user platform token. Two users → two sessions.

So #329 now = increments 1 (pool) + 2 (routing seam) + 3 (materialized + lazy server) — the complete per-user connection lifecycle. Retitled accordingly.

[initializ-mk]

Re-review — the PR now delivers the full #317 per-user connection lifecycle end-to-end

Since the first review, the earlier findings are resolved and two new increments landed (routing seam + materialized/lazy type=user). I re-traced the security-critical property and the whole vertical is correct.

Prior findings — all resolved

• 1 (panic guard): the establish cleanup is now one unconditionally-deferred func with recover() — a panicking connect fails that one establish cleanly, fl.done always closes, inFly[subject] is always cleared, wg.Done() always fires. Correct.
• 2 (synchronous Close): WaitGroup — mark closed, release lock, wg.Wait(), then close cached conns. Deterministic teardown, no deadlock (lock released before Wait).
• 3 (test gap): connect-error + panic tests added.
• 4 (no CI): base retargeted to main — CI now runs, all 10 checks green (the -race vertical is on record, which is what mattered for concurrency-critical code).

Increment 2 (routing seam) — behavior-preserving, and slightly safer

MCPTool.Execute resolves the client per-call; Manager.Tools() hands ordinary servers a StaticResolver over the shared client, whose ClientFor returns (sharedClient, nil) — identical behavior for bearer/static/oauth/agent-principal. Its only new error path is a nil client, which was previously a nil-deref in CallTool, so the change turns a latent crash into a clean ErrTransportUnavailable. Tested both ways.

Increment 3 (materialized + lazy type=user) — the completion, correct and secure

• Identity is bound per connection — the crux. establish(ctx) runs factory + Initialize under the pool's per-subject ctx (which carries the user's identity), so each user's initialize resolves their token and the remote server binds the session to them. This closes the shared-connection caveat #327 deferred — end-to-end.
• required:true still rejected for type=user (the check precedes pool creation), so the materialized path only reaches Ready lazily — Ready never lies about a connection.
• Schema validation is defensively good — empty-name and __-separator rejected (prevents a materialized name spoofing the <server>__<tool> namespace), input_schema validated like a real tools/list.
• Trust model is sound — a divergent/malicious materialized schema causes a runtime CallTool failure at the real per-user connection, not privilege escalation.

Three minor findings on increment 3 (two inline):

A. Docs: the materialized tool surface is a global declaration, not per-user-filtered. All users see the same tools.schemas; a user whose platform grant lacks a tool still sees it registered and gets a runtime auth error when the LLM calls it. The docs cover the global allow/deny filter and the lazy establish, but not this per-user divergence — worth a line so operators (and the LLM's expectations) know the schema is what the server can expose, filtered per user only at call time. Inline.

B. Robustness: go r.Run(context.Background()) starts before Initialize. On an Initialize failure the code calls cli.Close(), which must stop the just-started Run goroutine or it leaks per failed establish. Low risk (the Client contract), but worth a test that an establish failure after Run started leaves no goroutine behind. Inline.

C. Lifecycle (nit): materialized-schema staleness. If the real server's tools change, the forge.yaml schema is stale until re-materialized — a platform responsibility, worth a one-line note on the re-materialization trigger in the docs.

Verdict

Merge-ready. The complete arc is correct: per-subject isolation, single-flight, the identity-preserving/cancel-decoupled establish (B2 lesson), behavior-preserving routing seam, and per-connection identity binding that closes the #327 caveat — all under green -race CI. Findings A–C are a docs caveat, a leak-safety test, and a lifecycle note; none blocks. Clean, well-staged completion of #317.

[initializ-mk]

Minor: worth caveating that tools.schemas is a global declaration, not per-user-filtered. allow/deny filter the set the same way for every user, but the actual per-user tool access is enforced only at call time by that user's own connection — so a user whose platform grant lacks a materialized tool still sees it registered, and the LLM may attempt it and get a runtime auth error from the real connection. One line here ("the schema set is what the server can expose; a given user may not be granted all of them — ungranted calls fail at the per-user connection, not at registration") sets the right expectation.

[initializ-mk]

Added both caveats to the materialized-schemas section: (1) the schema set is a global declaration — a user whose grant lacks a tool still sees it registered and gets a runtime auth error at their own connection (the per-user gate is the connection, not registration); (2) a staleness note — schemas are a snapshot until the platform re-materializes (redeploy), same as allow:["*"].

[initializ-mk]

Minor (leak-safety): Run is started here before Initialize, and on an Initialize failure below the code relies on cli.Close() to stop this goroutine. If there's any window where Close() doesn't stop a just-started Run, each failed establish leaks a goroutine — and establish failures are expected (a user without a grant hits 401 at initialize). Low risk given the Client contract, but worth a test: force an Initialize failure after Run has started and assert no goroutine remains (or that Close joins Run).

[initializ-mk]

Added TestServer_Establish_InitializeFailure_StopsRun: a fake client whose Run blocks until Close and whose Initialize fails — establish starts Run, Initialize fails, cli.Close() fires, and the test asserts Run returned (no leaked goroutine). Pins the Client contract on the expected-failure path (a user without a grant 401s at initialize).

[initializ-mk]

Folded in the three minor re-review findings (36d29b7) — thanks for the thorough trace of the vertical:

• B (leak-safety) — new TestServer_Establish_InitializeFailure_StopsRun: forces an Initialize failure after the demux Run started and asserts Close stops it (no leaked goroutine). This is the expected-failure path — a user without a grant 401s at initialize — so it's worth pinning.
• A (docs) — the materialized schema set is now documented as a global declaration: a user whose grant lacks a tool still sees it registered and gets a runtime auth error at their connection; the per-user gate is the connection, not registration.
• C (docs) — added the staleness note: schemas are a snapshot until the platform re-materializes (redeploy).

All under green -race CI. Nothing changed in the security-critical paths you validated.