fix(mcp-oauth): anchor 401 handler task to prevent GC mid-flight - #55952
Merged
Conversation
`handle_401` spawned a dedup'd recovery coroutine via `asyncio.create_task(_do_handle())` and discarded the returned task reference. Python's event loop only keeps weak references to tasks, so the coroutine could be garbage-collected before it called `pending.set_result(...)`. Every concurrent caller awaiting that future then hangs forever, and the `finally: entry.pending_401.pop(...)` cleanup never runs — so subsequent 401s for the same key latch onto the dead future too. Same pattern the adapter-side fixes address (#11997, #11998, #12000, #12001, #12006). Hold the task in a process-wide set on the manager and discard it via `add_done_callback` once it completes. Regression test covers both the structural invariant (task tracked, then removed on completion) and a concurrent dedup path with a forced `gc.collect()` between the handler's await points.
The discard done-callback added via task.add_done_callback runs on a later event-loop iteration (call_soon) than the one that resolves `pending` and lets handle_401 return. Both inflight-task tests asserted the live set was empty immediately after the await returned, racing the callback. Add a single `await asyncio.sleep(0)` before the cleanup assertions.
Collaborator
Duplicate of #51757 — byte-identical salvage of #12228 (same |
3 tasks
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
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
Salvage of #12228 (@haileymarshall) onto current
main. Anchors the MCP OAuth 401-recovery task to a strong-reference set so the event loop can't garbage-collect it mid-run.Root cause
MCPOAuthManager.handle_401()launched its deduped recovery coroutine viaasyncio.create_task(_do_handle())and dropped the handle. The event loop only weakly references tasks, so_do_handlecould be GC'd between itsawaitcheckpoints. When that happened:pending.set_result(...)never ran → every concurrent 401 caller stuck onawait pendinghung forever.finally: entry.pending_401.pop(key, None)cleanup never ran → subsequent 401s for the same access token latched onto the dead future, poisoning the manager for that key until process restart.Same fire-and-forget pattern already fixed on the adapter side (#11997, #11998, #12000, #12001, #12006);
mcp_oauth_manager.pywas missed.Changes
tools/mcp_oauth_manager.py: addself._inflight_tasks: set[asyncio.Task]in__init__; inhandle_401,self._inflight_tasks.add(task)on launch andtask.add_done_callback(self._inflight_tasks.discard)on completion.tests/tools/test_mcp_oauth_manager.py: two new tests (inflight tracking + 8-caller dedup-survival withgc.collect()mid-flight).Follow-up fix on top of the contributor commit
The two new tests asserted
_inflight_taskswas empty immediately afterhandle_401returned, but thediscarddone-callback is scheduled vialoop.call_soonand runs on a later loop iteration than the one that resolvespending. Added a singleawait asyncio.sleep(0)before the cleanup assertions in both tests to let the callback fire.Validation
tests/tools/test_mcp_oauth_manager.pyInfographic