fix(mcp): give each concurrent tool call its own SQLAlchemy session - #43007
fix(mcp): give each concurrent tool call its own SQLAlchemy session#43007verdier wants to merge 2 commits into
Conversation
`flask-sqlalchemy` scopes `db.session` on `greenlet.getcurrent()`, which does not tell asyncio tasks apart: MCP tool calls all run in the same greenlet on the event loop thread, so they resolve to one shared `Session`. Every call runs in its own Flask app context, and `flask-sqlalchemy` removes the session on app context teardown — so the first call to finish removes the session out from under every call still in flight. They are left holding detached instances, and the next attribute read raises `DetachedInstanceError`, for a chart the tool has already committed. Scope the session on the running asyncio task instead. Off the event loop the greenlet identity is kept, so the WSGI web tier, Celery workers and the MCP thread pool behave exactly as before. Measured on apache/superset:6.1.0-py311, `superset mcp run`, concurrent `generate_chart` calls asserted against rows in the metadata database: | concurrency | before | after | |---|---|---| | 10 | 2/10, 8 DetachedInstanceError | 10/10, 0 | | 20 | 6/20, 14 DetachedInstanceError | 20/20, 0 | A probe on the auth hook counts the sessions actually handed out: one session for ten concurrent calls before, ten after. Sessions are no longer shared, so N concurrent tool calls now hold N connections: the SQLAlchemy pool has to be sized for the concurrency the MCP server is meant to serve, the way the web tier is sized for its workers. Fixes apache#42622
Code Review Agent Run #85ce88Actionable Suggestions - 0Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #43007 +/- ##
==========================================
- Coverage 66.40% 66.40% -0.01%
==========================================
Files 2860 2860
Lines 161514 161529 +15
Branches 37201 37203 +2
==========================================
+ Hits 107252 107259 +7
- Misses 52222 52228 +6
- Partials 2040 2042 +2
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Rewrote the description rather than appending to it — the connection-demand part deserved numbers instead of a caveat, and it turned up something worth your eye. The cliff is exactly at What makes it a cliff rather than a queue is not the pool size. Oversubscription is already the norm and drains fine everywhere else: the web tier is Which leads to the part I had not appreciated until I measured it: the tools are So the fix is correct but not self-limiting, and the two ways to make it so are a semaphore sized from the pool (small, turns the stall into a wait) or moving the blocking sections onto a worker thread while keeping Glad to add the semaphore to this PR and open a separate issue for the thread work, or to keep this one to the correctness fix alone. Whichever you prefer. |
Code Review Agent Run #010e78Actionable Suggestions - 0Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
aminghadersohi
left a comment
aminghadersohi
left a comment
There was a problem hiding this comment.
Thanks for this — it's a careful fix with an honest write-up, and the test that reproduces the DetachedInstanceError from #42567 by driving two real concurrent tasks through a sibling teardown is the right way to pin the bug. A few notes, and a direct answer to the trade-off you raised.
The correctness fix is sound, and the parity claim holds
I verified the "off the event loop, behaves exactly as before" claim rather than taking the docstring's word:
- Superset pins
flask-sqlalchemy==2.5.1(requirements/base.txt), whose default scope function is the werkzeug local ident —greenlet.getcurrentwhen greenlet is importable, elsethreading.get_ident. Your_greenlet_identguard mirrors that import ordering exactly, so the fallback is a faithful reproduction of the current default. asyncio.current_task()raisesRuntimeErroronly with no running loop and returnsNoneunder a loop with no current task; both fall back to_greenlet_ident(). Correct on both edges.- I grepped the backend for any event-loop entry point outside the MCP service (
asyncio.run,new_event_loop,run_until_complete,get_event_loop,ensure_future) — zero hits. The WSGI web tier, Celery workers, andasync_query_managerare fully synchronous, and the MCP sync-tool path runs in worker threads with no running loop. So althoughdbis global, the only place where a loop is running with a live task is the MCP async path — there's no pre-existing code where the old shared-session behavior was load-bearing under a live task. The global blast radius is real in principle but empty in practice on the pinned stack.
One forward-looking caveat: pyproject.toml allows flask-sqlalchemy < 4.0, and in 3.x the default scope switched to id(app_ctx) (per-app-context). If Superset ever moves to 3.x, this fallback would pin the web tier back to greenlet scoping instead of the library's new default — worth a comment noting the coupling to 2.5.1 semantics.
Registry teardown — I went looking for the leak, and it's covered on the dominant path
A per-task scopefunc means scoped_session keeps one registry entry per task, each pinning a Session and a strong ref to the Task. What removes it: mcp_auth_hook wraps every tool call (both async_wrapper and sync_wrapper) in with _get_app_context_manager():, which pushes a new app.app_context() per call whenever only a bare app context exists. On __exit__, flask-sqlalchemy's teardown_appcontext fires db.session.remove(), which runs synchronously in the same task — so scopefunc() returns that task and the correct entry is popped. with/__exit__ guarantees this on normal return, on exception, and on CancelledError, so the async path doesn't grow the registry unboundedly.
The one path where teardown does not fire per call is the has_request_context() → nullcontext() branch (embedded-under-WSGI): there, per-task entries created under a shared request context are only reclaimed when that request context tears down. If multiple concurrent tasks ever ran under a single shared request context, sibling entries would linger until the request ends — bounded by request lifetime, not unbounded, so minor, but worth being aware of if embedded concurrency is ever in play.
The pool cliff — you asked directly, so: net improvement, but it should ship with the guard
Your self-assessment is correct and the mechanism is exactly as you describe. Before this change, every concurrent MCP task shared one greenlet-scoped session and therefore effectively one connection — DB access was serialized, which masked both the bug and any pool pressure. That is unambiguously wrong: SQLAlchemy Session objects aren't safe to share across concurrent tasks, so the shared-session behavior is a genuine data-integrity/crash bug at any concurrency ≥ 2. Fixing it is not optional.
But because the tools are async def doing blocking DB work and holding a checked-out connection across await ctx.* points, giving each task its own session is precisely what makes pool_size + max_overflow reachable — the exact cliff you measured (15 clean, 16 → 30s QueuePool timeout, pool_size=40 → 20 clean). So the fix trades a silent correctness bug for a loud, bounded, tunable availability limit. That's a strictly better failure mode, so yes — net improvement, and it should land.
My one substantive ask: because this PR is what makes the pool exhaustible (before it, you couldn't hold more than ~one connection across the whole tool call set), I'd pair the isolation fix with a concurrency guard in the same change rather than as a follow-up — a semaphore sized from pool_size + max_overflow is the minimal, low-risk option and turns a 30s stall into fast, predictable backpressure. The worker-thread offload you mention is the more complete fix (it restores real concurrency instead of just capping it), but it's a larger change and fine as a follow-up. Shipping isolation alone silently lowers the effective concurrency ceiling operators hit, so at minimum that new limit deserves to be enforced or loudly documented here.
Minor
test_greenlet_scoping_detaches_the_instanceasserts the bug viapytest.raises(DetachedInstanceError)on the library default andtest_task_scoping_keeps_the_instance_aliveasserts the fix on observable session identity (not on the scopefunc directly) — both genuinely drive two concurrent tasks through a real teardown, so neither is a false-pass. Nicely done.
Context, not a request
There's an open PR, #42629, targeting the same problem. Flagging it only so the maintainers can coordinate — this review is on #43007's own merits and takes no position on which should land.
I'm not a committer, so this is a comment rather than an approving review; a maintainer will need to give the formal approval.
|
@verdier this is the direction I asked for on the issue, and the write-up holds up. Couple of things before it merges though. @aminghadersohi's ask for a semaphore in this PR rather than a follow-up seems right to me, that 30s stall is a full freeze on the event loop, not just backpressure. Also #42629 is chasing the same bug with a ContextVar scoped to the MCP call instead of the global scopefunc. Worth comparing notes with that PR before we pick one to land. |
3 participants
SUMMARY
Fixes #42622, taking up @rusackas' preference in the issue thread: scope the session per asyncio task, mirroring what the thread pool already does for sync calls.
flask-sqlalchemy2.5.1 scopesdb.sessionwithscopefunc=_ident_func, and_ident_funcisgreenlet.getcurrent. Async MCP tool calls are asyncio tasks, not greenlets: they all run in the same greenlet on the event loop thread, so they all resolve to the sameSession. Each call still runs in its own Flask app context — deliberately, sog.userdoes not race — andflask-sqlalchemyregistersteardown_appcontext → session.remove(). So the first call to finish removes the session shared by every call still running, leaving them holding detached instances. The next attribute read raisesDetachedInstanceError: for a chart the tool has already committed (#42567), or on aUserlazy-load before any write at all.This scopes on the running task, and keeps the greenlet identity off the event loop so the WSGI web tier, Celery workers and the MCP thread pool are untouched:
Measured on
apache/superset:6.1.0-py311,superset mcp run,stateless_http=True, N concurrentgenerate_chartcalls withsave_chart=true, asserted against rows in the metadata database — the response envelope alone is not trustworthy here, the rows get written either way and it is the responses that lie:DetachedInstanceErrorDetachedInstanceErrorA probe on the auth hook logging
id(db.session())per call counts what is actually handed out: one session for ten concurrent calls before, ten after — ten distinct tasks in both runs.Sessions are no longer shared, so N concurrent calls hold N connections instead of one. On
masterthe async path never contends for a connection — one session, one connection, however many calls — which is precisely why the corruption was silent. This surfaces the real demand, and there is a cliff:pool_size=5, max_overflow=10)QueuePool limit ... timeout 30.00pool_size=40The cliff is exactly at
pool_size + max_overflow, and it is a cliff rather than a queue because of where the wait happens. Oversubscription itself is already normal here and copes fine: the web tier runsgunicorn --worker-class gthread --threads 20against that same 15-connection pool (docker/entrypoints/run-server.sh), and the MCP sync path runs tools through anyio's default thread limiter — 40 threads against the same 15. Both are oversubscribed, both drain, because each thread waits on its own.On the event loop there is one thread, and the task waiting for a connection is that thread. The tasks holding the connections cannot reach their app context teardown to release them. Nothing moves until the 30s timeout fires.
Worth noting where that leaves the async path: the tools are declared
async def, but they perform no async I/O. Ingenerate_chart, all 34awaits are onctx.*—report_progress,info,debug, i.e. talking to the MCP client. Every database call underneath is blocking. So the tools are async in order to report progress, not in order to yield the loop: they occupy it without ever giving it back. That is the property that turns a pool wait into a stall, and it predates this PR — this change is simply the first thing to depend on it.Sizing
SQLALCHEMY_ENGINE_OPTIONSfor the intended concurrency is enough to make this behave today (the 20/20 row above). Making it self-limiting is a design decision I would rather not fold into a bug fix:await ctx.*progress reporting on the loop — the right shape, and the same footing as the sync path, but it touches every toolFor the avoidance of doubt, I am not suggesting a move to async drivers: that would be Superset's whole data layer, not the MCP tools, and it would buy nothing for work that is short queries and CPU. The question is only whether blocking work belongs on the loop.
Happy to add (1) here, and to open a separate issue for (2) — or to leave both out if you would rather ship the correctness fix alone. Your call.
TESTING INSTRUCTIONS
tests/unit_tests/extensions/test_session_scope.pycovers both directions, no database or MCP server needed:DetachedInstanceError— the bug, pinnedEnd to end: run
superset mcp runwith JWT auth, fire 10 concurrentgenerate_chartcalls withsave_chart=trueagainst a valid dataset, and assert on rows inslicesrather than on the response envelope.ADDITIONAL INFORMATION