perf: parallelize toolset startup and prefetch MCP catalog#3688
Conversation
docker-agent
left a comment
There was a problem hiding this comment.
Assessment: 🟢 APPROVE
The parallelisation and prefetch changes are well-structured. concurrent.MapSlice/ForEach are correct — Go 1.26.5 gives per-iteration loop variable semantics so there is no closure capture race. The two-wave start ordering (independent → peer-dependent) correctly preserves the invariant that the deferred aggregator starts only after its source toolsets have settled. gateway.Prefetch is safe: loader.load immediately wraps its fetch in context.WithoutCancel, so the background goroutine completes even if the caller's context is cancelled. The interactiveOAuthMu locking order (per-transport oauthFlowMu → global interactiveOAuthMu) is consistent at every call-site, ruling out deadlock.
Lower-confidence findings (not posted inline)
- [medium] pkg/tools/mcp/oauth.go:492 —
interactiveOAuthMu.Lock()does not honour context cancellation (confidence: 🟠 weak 52/100): In the interactive-prompt path (ensureToolSetsAreStarted), multiple toolsets starting concurrently could each queue oninteractiveOAuthMu. Async.Mutexlock is not cancellable, so if the agent run context is cancelled while a goroutine waits for the mutex, that goroutine remains blocked until the in-progress OAuth flow finishes. In practice this is the intended serialisation behaviour and the goroutine is not permanently leaked (it will proceed once the mutex is released and then return quickly on the cancelled context); the non-interactive startup path (emitToolsProgressively) already guards this withWithoutInteractivePromptsso the mutex is never contended there. Mentioning for awareness — a context-aware mutex (or aTryLockwith ctx-done select) would make cancellation instantaneous, but it is not required for correctness.
docker-agent
left a comment
There was a problem hiding this comment.
Assessment: 🟡 NEEDS ATTENTION
This PR's concurrent toolset startup is well-structured — the two-wave approach (independents first, then PeerDependent toolsets) is sound, concurrent.MapSlice/ForEach is correctly implemented, and the wrapper chain for PeerDependent detection works because all wrapper types implement Unwrapper. One liveness concern was found in the new OAuth serialization mutex.
| // Queue behind other transports' flows only after the cheap decline | ||
| // short-circuit; flows on this transport are already serialized by | ||
| // oauthFlowMu, so the latch cannot change while we wait here. | ||
| interactiveOAuthMu.Lock() |
There was a problem hiding this comment.
[medium] interactiveOAuthMu.Lock() is not context-aware: a cancelled toolset blocks indefinitely waiting behind an active OAuth flow
Trigger: This PR introduces concurrent toolset startup. If toolset A is in an interactive OAuth flow (waiting for user browser interaction, which can take minutes), toolset B — started concurrently — can also reach startInteractiveFlowLocked and block at interactiveOAuthMu.Lock() on line 492. If toolset B's ctx is cancelled before it acquires the mutex (e.g. the agent session is torn down, a startup timeout fires, or the user cancels), sync.Mutex.Lock() is not context-aware and cannot be interrupted. Toolset B's goroutine is stuck behind the user interaction for toolset A.
Impact: While blocked at interactiveOAuthMu.Lock(), toolset B's goroutine still holds oauthFlowMu (acquired by the caller authorizeOnce). Every concurrent RPC to toolset B's transport that also needs oauthFlowMu is therefore also stalled — potentially for minutes — even though toolset B's context is already cancelled. Once toolset B finally acquires interactiveOAuthMu, it calls handleOAuthFlow(ctx, ...) and returns immediately (cancelled context), but the delay was already paid. Before this PR sequential startup meant only one toolset could be doing OAuth at a time, so this scenario could not arise.
Suggested fix: Replace the bare sync.Mutex with a context-aware acquire that selects on ctx.Done():
acquired := make(chan struct{}, 1)
go func() { interactiveOAuthMu.Lock(); acquired <- struct{}{} }()
select {
case <-acquired:
defer interactiveOAuthMu.Unlock()
case <-ctx.Done():
return ctx.Err()
}Or use a sync.Cond-based approach, or a channel-based semaphore, to allow waiting goroutines to be preempted on context cancellation.
| Confidence | Score |
|---|---|
| 🟢 strong | 97/100 |
The first ref: toolset creation paid a blocking network round-trip to resolve the MCP catalog in the middle of LoadTeam. Kick the memoized fetch off in the background as soon as the config shows a catalog ref, so it overlaps model and environment resolution.
The startup sidebar pass started each toolset sequentially, so an agent with several MCP servers paid the sum of their handshakes before its tools appeared. Start them all concurrently (the context is already non-interactive, so no start can block on a user-driven flow), then process results in configuration order to keep progress events and once-per-streak warnings deterministic.
ensureToolSetsAreStarted ran Start() sequentially, so the first turn of an agent with several MCP toolsets (e.g. a sub-agent reached via transfer_task) paid the sum of the handshakes. Start them concurrently: Start() is single-flight per toolset, and warnings keep being recorded in configuration order after all starts settle.
A remote MCP tools/list is a network round-trip; listing sequentially made a turn pay the sum of them. List all started toolsets concurrently and merge in configuration order so duplicate handling (first toolset in config wins), collision reporting and once-per-streak warnings are unchanged.
Five call sites (toolset start/list, startup sidebar pass, hook dispatch, tool-call batches) hand-rolled the same pattern: a results slice indexed by i, one goroutine per element, WaitGroup. MapSlice captures it once — apply f to every element concurrently and return results in input order.
The deferred aggregator's Start lists its source toolsets' tools, so the concurrent fan-outs could race it against the very toolsets it depends on, yielding a spurious warning and missing deferred tools on the first turn. Mark such toolsets with tools.PeerDependent and start them in a second wave in both ensureToolSetsAreStarted (via new concurrent.ForEach) and emitToolsProgressively. emitToolsProgressively also consumes start outcomes in configuration order again instead of a barrier, so an early fast toolset is counted in the sidebar while a slow later one is still starting.
Toolsets now start concurrently, so two OAuth-requiring MCP servers could prompt the user simultaneously, race to bind a fixed callback port, and break legacy clients that answer elicitations without an elicitation_id (resolveSingle refuses ambiguous pending requests). A process-wide mutex around the interactive flow restores the one-at-a- time behavior of sequential startup; the per-transport sticky-decline short-circuit stays ahead of the queue so declined toolsets never wait.
f5ea16a to
babb9e1
Compare
docker-agent
left a comment
There was a problem hiding this comment.
Assessment: 🟡 NEEDS ATTENTION
Two LIKELY findings in the concurrency changes introduced by this PR. Neither is a crash or data-loss bug, but both represent new regressions vs. the sequential-startup behaviour this PR replaces.
| // Queue behind other transports' flows only after the cheap decline | ||
| // short-circuit; flows on this transport are already serialized by | ||
| // oauthFlowMu, so the latch cannot change while we wait here. | ||
| interactiveOAuthMu.Lock() |
There was a problem hiding this comment.
[medium] Holding per-transport oauthFlowMu while blocking on global interactiveOAuthMu can stall concurrent RPC calls for the duration of another transport's interactive OAuth flow
authorizeOnce acquires t.oauthFlowMu, then calls startInteractiveFlowLocked, which blocks on the process-wide interactiveOAuthMu. Before this PR, toolsets started sequentially, so Transport A's OAuth would never begin while Transport B's user-facing flow was in progress. Now that toolsets start concurrently, Transport A can be holding its own per-transport t.oauthFlowMu while waiting for Transport B's user to click "Authorize" in the browser — a wait measured in minutes.
During that entire window, any concurrent RoundTrip call for Transport A that encounters a 401 and calls authorizeOnce will block on t.oauthFlowMu with its request context ticking down, potentially timing out. The comment in the code acknowledges this lock ordering but does not address the extended hold duration introduced by concurrency.
A low-risk mitigation: acquire interactiveOAuthMu before t.oauthFlowMu (i.e. restructure authorizeOnce to try the cheap token-check first without holding oauthFlowMu, and only take t.oauthFlowMu for the minimal critical section). Alternatively, document that callers must use a context with a timeout long enough to survive a full interactive flow.
| Confidence | Score |
|---|---|
| 🟡 moderate | 75/100 |
| // instead of paying the network round-trip on the caller's critical path. | ||
| func Prefetch(ctx context.Context) { | ||
| loader := loaderFrom(ctx) | ||
| go func() { _, _ = loader.load(ctx) }() |
There was a problem hiding this comment.
[low] Prefetch winning the sync.Once race with a transient error permanently prevents catalog loading for this process lifetime
Prefetch fires go func() { _, _ = loader.load(ctx) }() in the background. Inside load, sync.Once memoizes the first result — including errors — so a transient failure during the prefetch window permanently caches the error; no subsequent ServerSpec call can retry.
Before this PR, Prefetch did not exist, so the once.Do was triggered only by a real ServerSpec call (at which point the caller could propagate the error and the user could retry by restarting). Now the once.Do fires at load time, in the background, when the caller has no visibility.
fetchAndCache does have a disk-cache fallback, so a pure network blip is unlikely to produce an error (it would return the cached catalog instead). The residual scenario is: no network AND no on-disk cache (cold start on first run with no connectivity). In that case, the process must be restarted to retry, whereas before this PR the first real call would fail immediately and the user could react.
Consider returning the error through a channel or logging it at WARN level so the failure is visible, and/or not entering once.Do from the background goroutine (only trigger the fetch eagerly without committing the error result — keep once.Do solely on the real call path).
| Confidence | Score |
|---|---|
| 🟡 moderate | 62/100 |
Profiling
examples/gopher.yamlon Apple Silicon revealed that agent startup blocked sequentially on MCP handshakes: every toolset was started one after another, so total startup time was the sum of all handshake latencies. For a sub-agent likelibrarianthat reaches three MCP servers, that summed to around 5.2 s on a warm cache.This PR attacks that sequential work at every layer.
gateway.Prefetch()kicks off the memoized catalog fetch in the background as soon as a config withref:toolsets is loaded, so the network round-trip overlaps model and environment resolution instead of blockingLoadTeam. InsideEmitStartupInfo,agent.Run, andtools/list, toolsets are now started and queried concurrently using a newconcurrent.MapSlicefan-out helper that replaces five hand-rolled results-slice +WaitGrouppatterns (including pre-existing ones inpkg/hooks/executor.goandpkg/runtime/toolexec/dispatcher.go). Merged tool lists remain in configuration order so dedup, collision detection, and warning semantics are unchanged.Two correctness fixes accompany the parallelism. The
deferredaggregator implements a newtools.PeerDependentmarker so it is started only after all its peer toolsets have settled — preserving the ordering invariant that existed before when startup was sequential. A process-wide mutex inpkg/tools/mcp/oauth.goserializes interactive OAuth flows so concurrent starts cannot race on the fixed callback port, issue simultaneous prompts, or confuse legacy MCP clients that answer elicitations without anelicitation_id.Measured on Apple Silicon, warm caches:
LoadTeamdrops from ~150–190 ms to ~130–145 ms; gopher root-agent tools ready in 1.34 s (from 1.47 s, now bound by the inherent ~1.3 sgoplsMCP handshake); librarian sub-agent first turn drops from 5.2 s to ~3.4 s. The single-flight guarantee onStartis preserved, configuration-order warnings and dedup are preserved, and the deferred toolset ordering invariant is preserved.