feat: retry shared-connection MCP auth failures after bounded reconnect wait with concurrent dedup - #5726
Merged
Pratham-Mishra04 merged 1 commit intoAug 8, 2026
Conversation
This was referenced Jul 30, 2026
Merged
refactor: migrate OAuth flow state/PKCE fields from
oauth_configs to new mcp_oauth_flows table
#5709
Merged
|
|
This was referenced Jul 30, 2026
Merged
Merged
Merged
Collaborator
Author
18 tasks
This was referenced Aug 6, 2026
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Pratham-Mishra04
force-pushed
the
07-30-feat_heal_shared_mcp_connections_through_token_expiry_without_failed_calls
branch
from
August 6, 2026 21:53
4cfec68 to
e340bd6
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@core/mcp/toolmanager.go`:
- Around line 934-940: Update the retry logic around GetClientForTool to
reacquire the original MCP client via GetClientByName(executionConfig.Name)
after ReconnectClient. Reject the retry when the client is missing,
executionConfig.ID differs from the original, or the tool is no longer
available; preserve provider isolation. Add coverage for routing changes during
reconnect and verify the second client receives no retry.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 49db4186-949b-48e5-801e-958146a91c9b
📒 Files selected for processing (14)
core/mcp/agent_test.gocore/mcp/auth_retry_test.gocore/mcp/clientmanager.gocore/mcp/codemode/starlark/starlark_test.gocore/mcp/mcp.gocore/mcp/toolmanager.gocore/mcp/toolmanager_test.goframework/configstore/rdb.goframework/configstore/store.goframework/oauth2/sync.goframework/oauth2/sync_test.goframework/oauth2/tokenexchange.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/server/server.go
🚧 Files skipped from review as they are similar to previous changes (13)
- framework/configstore/store.go
- framework/oauth2/tokenexchange.go
- framework/configstore/rdb.go
- core/mcp/codemode/starlark/starlark_test.go
- core/mcp/toolmanager_test.go
- transports/bifrost-http/server/server.go
- core/mcp/mcp.go
- framework/oauth2/sync.go
- core/mcp/agent_test.go
- framework/oauth2/sync_test.go
- core/mcp/clientmanager.go
- transports/bifrost-http/lib/config.go
- core/mcp/auth_retry_test.go
Pratham-Mishra04
force-pushed
the
07-30-feat_heal_shared_mcp_connections_through_token_expiry_without_failed_calls
branch
from
August 8, 2026 08:43
e340bd6 to
40ffed8
Compare
Pratham-Mishra04
force-pushed
the
07-31-feat_ui_support_for_token_exchange_mcp_auth_type
branch
from
August 8, 2026 08:43
01d9624 to
0c63dfc
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@core/mcp/auth_retry_test.go`:
- Around line 266-276: Remove the unused authRetryClientManager.resetInflight
method and update the nearby comment to describe retained inflight-record
behavior without referencing that removed helper.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 421f0d89-a9a4-4628-97f5-93a3ed097acb
📒 Files selected for processing (6)
core/mcp/auth_retry_test.gocore/mcp/toolmanager.goframework/configstore/rdb.goframework/configstore/store.goframework/oauth2/sync.gotransports/bifrost-http/lib/config.go
🚧 Files skipped from review as they are similar to previous changes (4)
- framework/configstore/store.go
- framework/configstore/rdb.go
- transports/bifrost-http/lib/config.go
- core/mcp/toolmanager.go
Comment on lines
+266
to
+276
| // resetInflight clears any completed inflight record so the next | ||
| // ReconnectClient call starts from a clean no-op state. Only needed by tests | ||
| // that reuse the same authRetryClientManager across multiple reconnect | ||
| // phases and require no stale op to be observable via AwaitReconnect; | ||
| // ReconnectClient's own CompareAndSwap-style replace makes this unnecessary | ||
| // for tests that just call ReconnectClient again. | ||
| func (m *authRetryClientManager) resetInflight() { | ||
| m.inflightMu.Lock() | ||
| m.inflight = nil | ||
| m.inflightMu.Unlock() | ||
| } |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
resetInflight is unused and fails the unused linter.
golangci-lint reports func (*authRetryClientManager).resetInflight is unused. No test in this file calls it. Remove it, or add the test that needs the clean no-op state. The comment block at lines 253-261 can keep documenting the retained-record behavior without referencing a method that does not exist.
🧹 Proposed fix: drop the dead helper
-// resetInflight clears any completed inflight record so the next
-// ReconnectClient call starts from a clean no-op state. Only needed by tests
-// that reuse the same authRetryClientManager across multiple reconnect
-// phases and require no stale op to be observable via AwaitReconnect;
-// ReconnectClient's own CompareAndSwap-style replace makes this unnecessary
-// for tests that just call ReconnectClient again.
-func (m *authRetryClientManager) resetInflight() {
- m.inflightMu.Lock()
- m.inflight = nil
- m.inflightMu.Unlock()
-}
-Also update the trailing sentence at line 260-261 to drop the resetInflight reference.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| // resetInflight clears any completed inflight record so the next | |
| // ReconnectClient call starts from a clean no-op state. Only needed by tests | |
| // that reuse the same authRetryClientManager across multiple reconnect | |
| // phases and require no stale op to be observable via AwaitReconnect; | |
| // ReconnectClient's own CompareAndSwap-style replace makes this unnecessary | |
| // for tests that just call ReconnectClient again. | |
| func (m *authRetryClientManager) resetInflight() { | |
| m.inflightMu.Lock() | |
| m.inflight = nil | |
| m.inflightMu.Unlock() | |
| } |
🧰 Tools
🪛 golangci-lint (2.12.2)
[error] 272-272: func (*authRetryClientManager).resetInflight is unused
(unused)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@core/mcp/auth_retry_test.go` around lines 266 - 276, Remove the unused
authRetryClientManager.resetInflight method and update the nearby comment to
describe retained inflight-record behavior without referencing that removed
helper.
Source: Linters/SAST tools
Collaborator
Author
Merge activity
|
Pratham-Mishra04
changed the base branch from
07-31-feat_ui_support_for_token_exchange_mcp_auth_type
to
graphite-base/5726
August 8, 2026 10:02
…-hint gap in budget test authRetryClientManager.ReconnectClient cleared m.inflight to nil in the same critical section that closed the op's done channel, unlike the real beginExclusiveClientOp it mirrors (which deliberately leaves a completed op in place, replacing it only lazily on the next call). A caller that lost the race and got 'already in progress' could poll AwaitReconnect a moment after the winner cleared inflight, observe nil, and wrongly conclude nothing was ever in flight — the TestExecuteTool_AuthFailureRetry_Shared_Concurrent401sJoinOneReconnect errs[1] == nil assertion depended on winning that race. Mirror the real CompareAndSwap-on-a-done-op replace behavior instead, and add resetInflight for tests that need an explicit clean-slate reset. TestExecuteTool_AuthFailureRetry_Shared_FallsBackWhenReconnectExceedsBudget built its client state with no idempotent hint, so its 'no retry' assertions could pass for the wrong reason (annotation fail-closed) instead of proving the budget give-up path it's named for. Pass an explicit idempotent hint, matching the sibling reconnect-mechanics tests.
Pratham-Mishra04
force-pushed
the
07-30-feat_heal_shared_mcp_connections_through_token_expiry_without_failed_calls
branch
from
August 8, 2026 10:05
40ffed8 to
8b7af9a
Compare
Pratham-Mishra04
deleted the
07-30-feat_heal_shared_mcp_connections_through_token_expiry_without_failed_calls
branch
August 8, 2026 10:08
This was referenced Aug 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
On shared MCP connections, a 401 auth failure previously triggered a background reconnect but returned the original error immediately with no retry. This PR upgrades that path to a bounded-wait retry: the caller waits up to
MCPSharedAuthRetryReconnectBudget(10 s, further capped by the request's own remaining deadline) for the reconnect to finish, then retries the same call once on the healed connection. If the reconnect exceeds the budget it keeps running in the background and the original error surfaces as before. Concurrent 401s on the same client now deduplicate: the loser of the reconnect race joins the winner's in-flight attempt via a newAwaitReconnectinterface method rather than failing outright.Additionally,
OAuthTokenRefreshWorker(renamed fromTokenRefreshWorker) gains aSetOnTokenRefreshedcallback. The HTTP server installs this callback after the boot dial to proactively recycle shared MCP connections whenever the background token refresh worker lands a fresh credential, instead of waiting for the next call to fail with a 401.Changes
inflightClientOpstruct andbeginExclusiveClientOp: Replaced thesync.Map[string]boolsentinel inMCPManagerwith async.Map[string]*inflightClientOpthat carries adonechannel and final error. All exclusive-operation entry points (ReconnectClient,DisableClient,EnableClient,UpdateClient,UpdateClientConnection) now usebeginExclusiveClientOpand propagate their return error to waiters via a deferredfinish(retErr).AwaitReconnect: New method onMCPManager(andClientManagerinterface) that lets a caller block up to a given budget for an in-flight exclusive operation to complete, returning the operation's final error. A timed-out wait never cancels the underlying operation.recoverSharedConnection: New method extracted fromattemptAuthFailureRecoverythat implements the bounded-wait retry for shared connections. It triggers the background reconnect, waits for the result channel, falls back toAwaitReconnectif the trigger lost the race to a concurrent reconnect, then re-acquires the healed connection and retries the call once.triggerBackgroundReconnect: Now returns a<-chan error(buffered, capacity 1) so the caller can observe the reconnect outcome without blocking the goroutine.Retry opt-out ordering: The destructive/non-idempotent tool gate now sets a
retryOptedOutflag rather than returning early, so connection healing always runs on the shared path even when the retry itself is suppressed.OAuthTokenRefreshWorkerrename andSetOnTokenRefreshedcallback:TokenRefreshWorkeris renamed toOAuthTokenRefreshWorkerthroughout. A newSetOnTokenRefreshed(func(mcpClientID, authMode string))method (backed by anatomic.Pointer) lets callers register a hook invoked after each successful proactive refresh. The HTTP server uses this to trigger a backgroundReconnectMCPClientfor shared-auth clients immediately after a token is refreshed.Test coverage: Replaced the single
TestExecuteTool_AuthFailureRetry_Shared_FailsFastAndTriggersBackgroundReconnecttest with four focused tests covering: successful bounded-wait retry, budget-exceeded fallback, reconnect-failed fallback, destructive-tool reconnect-without-retry, and concurrent 401 deduplication. TheauthRetryClientManagermock now implements the fullinflightClientOpdedup logic to faithfully simulate the real manager. NewOAuthTokenRefreshWorkertests cover the callback firing, non-firing on failure, non-firing for tokens with no MCP client ID, and nil-safety.Type of change
Affected areas
How to test
go test ./core/mcp/... ./framework/oauth2/...Key scenarios exercised by the new tests:
CallToolinvocations and a successful response.CallToolinvocation, the original error, and the reconnect still running.ReconnectClientcall, 1 rejected duplicate, 4 totalCallToolinvocations, and both callers succeeding.Breaking changes
TokenRefreshWorkeris renamed toOAuthTokenRefreshWorkerand theConfigfieldTokenRefreshWorkeris renamed toOAuthTokenRefreshWorker. Any code referencing these names directly must be updated. TheClientManagerinterface gains a new required methodAwaitReconnect(clientID string, budget time.Duration) (bool, error); all implementations (including mocks) must add this method.Security considerations
The bounded-wait retry re-uses the caller's existing request context deadline to cap the reconnect wait, preventing a malicious or slow upstream from holding a request open indefinitely beyond its own timeout. The
SetOnTokenRefreshedcallback is stored as anatomic.Pointerto avoid data races between the worker goroutine and the serving layer installing the callback after startup.Checklist
docs/contributing/README.mdand followed the guidelines