fix(mcp-oauth): anchor 401 handler task to prevent GC mid-flight - #12228
haileymarshall wants to merge 1 commit into
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 (NousResearch#11997, NousResearch#11998, NousResearch#12000, NousResearch#12001, NousResearch#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.
|
Salvaged onto current main as #51757 with full credit to you, @haileymarshall — the GC-eligible-task defect is still live (the call site moved to |
|
Merged via PR #55952. Your fix commit was cherry-picked onto current main with your authorship preserved in git log (commit 9f22f36). I added a small test follow-up on top: the two new inflight-task tests asserted the discard done-callback had already run by the time handle_401 returned, but that callback fires on a later event-loop tick via call_soon, so I added an await asyncio.sleep(0) before the cleanup assertions. Thanks for catching the GC-anchor gap in mcp_oauth_manager.py — same class as the adapter PRs. #55952 |
Summary
MCPOAuthManager.handle_401()spawned its dedup'd recovery coroutine viaasyncio.create_task(_do_handle())and dropped the returned task handle. Python's event loop only keeps weak references to tasks, so the coroutine can be garbage-collected between itsawaitcheckpoints. When that happens:pending.set_result(...)never runs → every concurrent caller stuck onawait pendinghangs forever.finally: entry.pending_401.pop(key, None)cleanup never runs → subsequent 401s for the same access token latch onto the dead future, so the manager is poisoned for that key until process restart.This is the exact pattern the adapter-side PRs are cleaning up (#11997 dingtalk, #11998 weixin, #12000 qqbot, #12001 rl_training, #12006 bluebubbles).
mcp_oauth_manager.pywas missed.Fix
self._inflight_tasks: set[asyncio.Task]toMCPOAuthManager.__init__.handle_401, keep the task alive viaself._inflight_tasks.add(task)and drop it on completion viatask.add_done_callback(self._inflight_tasks.discard).Test plan
test_handle_401_tracks_inflight_task_to_prevent_gc— asserts the task is added to_inflight_taskswhile running and removed via the done-callback once finished.test_handle_401_dedup_survives_even_if_task_reference_dropped— fans out 8 concurrent callers onto the same dedup'd handler, forcesgc.collect()between scheduling and completion, and verifies all callers resolve within a 5s timeout. Without the fix this test can intermittently hang.python3 -m py_compile tools/mcp_oauth_manager.py tests/tools/test_mcp_oauth_manager.pyWhy
invalidate_if_disk_changedwasn't affectedIt runs on the caller's stack (directly
await-ed fromHermesMCPOAuthProvider.async_auth_flow), so no orphaned task can exist. Only thecreate_taskpath inhandle_401has this exposure.