Repository navigation
fix(oauth): narrow context.lock scope in async_auth_flow (rebase of #2660 by @peisuke, closes #2847) - #2858
fix(oauth): narrow context.lock scope in async_auth_flow (rebase of #2660 by @peisuke, closes #2847)#2858Bartok9 wants to merge 6 commits into
Conversation
425161d to
a9be0b9
Compare
|
I reproduced this downstream with MCP SDK 1.28.1 against a production Streamable HTTP OAuth server: OAuth and I opened the
The backport also serializes refresh and authorization transitions under one semaphore, re-checks tokens after acquiring it, and restores each request's MCP protocol version before RFC 8707 request construction. |
|
Thanks @ctaylor86 — really helpful production repro and the v1.x work in #3243. Agree on both points for
This PR is a conflict-free salvage of @peisuke's #2660 onto Happy to fold (1)/(2) into this branch if that's the preferred vehicle for |
a9be0b9 to
16e003c
Compare
A long-poll GET SSE no longer holds OAuthContext.lock. Token refresh is single-flight via refresh_lock. A 401 whose access token changed while the request was in flight retries once instead of starting a second authorization. The 401 discovery path stays inside the lock.
16e003c to
614b033
Compare
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
3 issues found across 1 file
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="src/mcp/client/auth/oauth2.py">
<violation number="1" location="src/mcp/client/auth/oauth2.py:611">
P2: `_unauthorized_token_changed()` returns `False` for every 403, even after a concurrent step-up replaced the token. A second request sent with the same under-scoped token then starts another interactive authorization after the first granted that scope. Check the replacement token against the challenged scope before reauthorizing.</violation>
<violation number="2" location="src/mcp/client/auth/oauth2.py:644">
P1: `context.lock` still encloses HTTP yields in the 401/step-up branch, and `refresh_lock` encloses the refresh yield. AnyIO locks can only be released by the task that acquired them, so cross-task resumption still raises the OAuth lock error. Move yielded HTTP requests outside these task-owned locks.</violation>
<violation number="3" location="src/mcp/client/auth/oauth2.py:669">
P2: This retry is yielded while `context.lock` is held, so a long-lived SSE response can block every other request using this provider. Update the header under the lock, then release it before yielding the retry.</violation>
</file>
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
| if needs_refresh: | ||
| async with self.context.refresh_lock: | ||
| refresh_request = None | ||
| async with self.context.lock: |
There was a problem hiding this comment.
P1: context.lock still encloses HTTP yields in the 401/step-up branch, and refresh_lock encloses the refresh yield. AnyIO locks can only be released by the task that acquired them, so cross-task resumption still raises the OAuth lock error. Move yielded HTTP requests outside these task-owned locks.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At src/mcp/client/auth/oauth2.py, line 644:
<comment>`context.lock` still encloses HTTP yields in the 401/step-up branch, and `refresh_lock` encloses the refresh yield. AnyIO locks can only be released by the task that acquired them, so cross-task resumption still raises the OAuth lock error. Move yielded HTTP requests outside these task-owned locks.</comment>
<file context>
@@ -595,40 +603,72 @@ def _select_authorization_server(self, advertised: list[str]) -> str:
+ if needs_refresh:
+ async with self.context.refresh_lock:
+ refresh_request = None
+ async with self.context.lock:
+ if self._needs_token_refresh():
+ refresh_request = await self._refresh_token()
</file context>
|
|
||
| def _unauthorized_token_changed(self, response: httpx2.Response, sent_access_token: str | None) -> bool: | ||
| """True when a 401 is stale because another request already replaced the access token.""" | ||
| if response.status_code != 401: |
There was a problem hiding this comment.
P2: _unauthorized_token_changed() returns False for every 403, even after a concurrent step-up replaced the token. A second request sent with the same under-scoped token then starts another interactive authorization after the first granted that scope. Check the replacement token against the challenged scope before reauthorizing.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At src/mcp/client/auth/oauth2.py, line 611:
<comment>`_unauthorized_token_changed()` returns `False` for every 403, even after a concurrent step-up replaced the token. A second request sent with the same under-scoped token then starts another interactive authorization after the first granted that scope. Check the replacement token against the challenged scope before reauthorizing.</comment>
<file context>
@@ -595,40 +603,72 @@ def _select_authorization_server(self, advertised: list[str]) -> str:
+
+ def _unauthorized_token_changed(self, response: httpx2.Response, sent_access_token: str | None) -> bool:
+ """True when a 401 is stale because another request already replaced the access token."""
+ if response.status_code != 401:
+ return False
+ current = self.context.current_tokens.access_token if self.context.current_tokens else None
</file context>
| if self._unauthorized_token_changed(response, sent_access_token): | ||
| # Another request refreshed or re-authorized while this one was in flight. | ||
| self._add_auth_header(request) | ||
| yield request |
There was a problem hiding this comment.
P2: This retry is yielded while context.lock is held, so a long-lived SSE response can block every other request using this provider. Update the header under the lock, then release it before yielding the retry.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At src/mcp/client/auth/oauth2.py, line 669:
<comment>This retry is yielded while `context.lock` is held, so a long-lived SSE response can block every other request using this provider. Update the header under the lock, then release it before yielding the retry.</comment>
<file context>
@@ -595,40 +603,72 @@ def _select_authorization_server(self, advertised: list[str]) -> str:
+ if self._unauthorized_token_changed(response, sent_access_token):
+ # Another request refreshed or re-authorized while this one was in flight.
+ self._add_auth_header(request)
+ yield request
+ return
</file context>
pre-commit ruff-format wants this condition on one line.
The refresh re-check and the failure flag now sit outside the lock's exit so branch coverage can see them. A peer that stores a valid token while this caller waits does not start a second refresh. A non-200 refresh clears the token and sends the original request unauthenticated.
There was a problem hiding this comment.
2 issues found across 2 files (changes from recent commits).
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="tests/client/test_auth.py">
<violation number="1" location="tests/client/test_auth.py:3784">
P2: This assigns the unrelated `HoldableLock` class to a field typed as `anyio.Lock`, so strict Pyright rejects the test. Type the test double to the lock interface or cast it at the assignment.</violation>
</file>
<file name="src/mcp/client/auth/oauth2.py">
<violation number="1" location="src/mcp/client/auth/oauth2.py:651">
P2: A concurrent 401 flow can replace `current_tokens` after this boolean is captured but before this lock is reacquired. Re-check `_needs_token_refresh()` while holding this lock; otherwise `_refresh_token()` may raise when the replacement token lacks a refresh token.</violation>
</file>
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
| return self._inner.locked() | ||
|
|
||
| lock = HoldableLock() | ||
| oauth_provider.context.refresh_lock = lock |
There was a problem hiding this comment.
P2: This assigns the unrelated HoldableLock class to a field typed as anyio.Lock, so strict Pyright rejects the test. Type the test double to the lock interface or cast it at the assignment.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At tests/client/test_auth.py, line 3784:
<comment>This assigns the unrelated `HoldableLock` class to a field typed as `anyio.Lock`, so strict Pyright rejects the test. Type the test double to the lock interface or cast it at the assignment.</comment>
<file context>
@@ -3667,3 +3668,144 @@ async def echo_callback() -> AuthorizationCodeResult:
+ return self._inner.locked()
+
+ lock = HoldableLock()
+ oauth_provider.context.refresh_lock = lock
+ await lock.acquire()
+ lock.arm = True
</file context>
| async with self.context.lock: | ||
| refresh_request = await self._refresh_token() |
There was a problem hiding this comment.
P2: A concurrent 401 flow can replace current_tokens after this boolean is captured but before this lock is reacquired. Re-check _needs_token_refresh() while holding this lock; otherwise _refresh_token() may raise when the replacement token lacks a refresh token.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. When an issue isn't valid or won't be fixed in this PR, reply in its thread with the reason and then resolve the thread. At src/mcp/client/auth/oauth2.py, line 651:
<comment>A concurrent 401 flow can replace `current_tokens` after this boolean is captured but before this lock is reacquired. Re-check `_needs_token_refresh()` while holding this lock; otherwise `_refresh_token()` may raise when the replacement token lacks a refresh token.</comment>
<file context>
@@ -640,14 +640,23 @@ async def _auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Req
+ still_needs_refresh = self._needs_token_refresh()
+ refresh_request = None
+ if still_needs_refresh:
+ async with self.context.lock:
refresh_request = await self._refresh_token()
if refresh_request is not None:
</file context>
| async with self.context.lock: | |
| refresh_request = await self._refresh_token() | |
| async with self.context.lock: | |
| refresh_request = ( | |
| await self._refresh_token() if self._needs_token_refresh() else None | |
| ) |
An unauthenticated request holds context.lock until that 401 login finishes, so a peer waits on the lock instead of starting a second login. A request that already has a token still yields outside the lock.
The waiter test enters the stand-in once before it is armed, and drops the unused lock methods that coverage could not reach.
This branch has not been deployed
2 participants
Summary
Conflict-free salvage of #2660 by @peisuke onto current
main, opened as a fresh, mergeable PR so the fix can land without waiting on the original draft to be un-conflicted. All credit for the fix goes to @peisuke — the commit is authored by them; I only resolved the rebase againstmainand squashed to a single commit.Closes #2847. Addresses the same
RuntimeError: The current task is not holding this lockreported in #2644 and #2847.Why a new PR
#2660 has been a
CONFLICTINGdraft since May 22 with no maintainer review. Issue #2847 (June 12) is a second, independent production repro of the exact bug, and the reporter explicitly asks for #2660 to be reviewed/merged. Rather than let the fix keep sitting, this PR provides a ready-to-merge version. If maintainers prefer to merge #2660 directly, or @peisuke wants to push the rebase themselves, this PR can be closed — the goal is just to unblock the fix.The bug
OAuthContext.lockis ananyio.Lock, which records task identity atacquire()and enforces same-taskrelease().async_auth_flowholds this lock acrossyieldpoints. When httpx drives the generator from a different task during concurrent OAuth connections (e.g. Notion's frequent token refresh behind a gateway),release()raisesRuntimeError: The current task is not holding this lock.The fix (@peisuke's approach)
Narrow the lock scope so no HTTP
yield(including the long-poll GET SSE and the token-refresh round trips) runs while holdingcontext.lock, plus a single-flightrefresh_lockwith a re-check under the lock. Keeps trio portability (unlike swapping toasyncio.Lock).Verification
Rebased onto current
main. Single commit (authored by @peisuke), touches onlysrc/mcp/client/auth/oauth2.pyandtests/client/test_auth.py(+348/-18).