Repository navigation
fix(mcp): give concurrent tool calls isolated db sessions - #42629
sadpandajoe merged 8 commits into
Conversation
db.session is scoped by greenlet ident, but async MCP tool calls are asyncio tasks on one event-loop greenlet, so every in-flight call resolves to the same Session. The first call to finish removes that session at app-context teardown, detaching the other calls' instances and surfacing DetachedInstanceError on the next attribute read (seen in generate_chart failing for a chart it had already committed). Scope the session registry to a per-call ContextVar token set while _get_app_context_manager pushes the call's app context, so teardown removes exactly that call's session. Outside MCP tool calls the scope falls back to the greenlet ident, leaving web, CLI, and Celery paths unchanged. Installed from init_fastmcp_server.
Code Review Agent Run #3bef32Actionable 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 |
![codeant-ai-for-open-source[bot]](https://avatars.githubusercontent.com/u/151810221?s=60&v=4){kind=small}
|
The flagged issue is correct. The current implementation of MCP tool execution in the pull request already addresses this race condition by introducing Since the fix is already implemented in the provided diff, no further action is required for this specific issue. I have checked the PR comments, and there are no other pending review comments to address. |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #42629 +/- ##
==========================================
- Coverage 66.65% 66.62% -0.03%
==========================================
Files 2866 2867 +1
Lines 162911 162890 -21
Branches 37530 37505 -25
==========================================
- Hits 108584 108532 -52
- Misses 52202 52231 +29
- Partials 2125 2127 +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:
|
The full unit-test run leaks a request context into the worker from an earlier test (sql_lab_test.py pushes one and never pops it), so _get_app_context_manager() takes its nullcontext branch and the per-call session token is never set — the three scoping tests then fail on suite ordering rather than on the code under test. Enter the contexts through _mcp_tool_call_context() directly, which always sets the token and pushes a fresh app context.
✅ Deploy Preview for superset-docs-preview ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
Code Review Agent Run #de0d0eActionable 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 |
run_server(use_factory_config=True) builds the app via create_mcp_app() directly, bypassing init_fastmcp_server() where the per-call scoping is installed. Concurrent tool calls in factory-configured deployments would still share the greenlet-scoped db.session. Install the scoping in this branch too.
|
Good catch. The factory-configured path in |
Code Review Agent Run #7ca966Actionable 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 |
Callers that serve the app straight from the factory never pass through init_fastmcp_server() or run_server(), so they kept the greenlet-scoped default and the shared-session race came back on that path. The lazy import is guarded for the module-import-time case where superset.extensions is still mid-initialization; the server entry points install the scopefunc there instead.
|
Addressed in f3bd3f7. create_mcp_app() now installs the per-call session scoping itself, so deployments that serve the app directly from the factory get the same isolation as the init_fastmcp_server() and run_server() paths. The lazy import is guarded for the module-import-time case where superset.extensions is still mid-initialization, with the server entry points as the deterministic fallback. Added a test covering the factory-only path. |
Code Review Agent Run #c1c099Actionable 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.
Reviewed the session-isolation change at HEAD f3bd3f7. The mechanism is sound for the path it targets:
- Concurrency model / scope key. Async tool calls are asyncio tasks multiplexed on one event-loop greenlet, so the stock greenlet-keyed
scoped_sessioncollapses them onto oneSession. Keying the registry on a per-callContextVartoken (("mcp_tool_call", id(token))) is correct: each asyncio task copies its context, so_mcp_session_token.set(object())yields a distinct token per concurrent call and two in-flight calls cannot resolve to the same key. Outside tool calls the scopefunc falls back togreenlet.getcurrent(), matching the flask-sqlalchemy default, so web/CLI/Celery paths are unchanged. - Teardown (success + exception). The token is reset in
finallyafter thewith app_context()block pops, so the app-context teardown'sdb.session.remove()still resolves to this call's key and removes exactly that session; the exception path exits the samewith(plus_cleanup_session_on_error()). Ordering is right and no stale registry entry is left keyed by a to-be-reusedid(). - Identity binding. The token is set before the app context is pushed and
_setup_user_context()runs inside it, so identity resolution and the subsequent queries share the same isolated session — no session-X-identity / session-Y-query drift. - Thread-safety. No new shared mutable registry is introduced;
mcp_session_scopefunconly reads aContextVar, andinstall_mcp_session_scoping()is a one-time idempotent scopefunc swap. - Wiring. Scoping is installed on all three entry paths (
create_mcp_app,init_fastmcp_server,run_server) and the scope is entered per dispatch via_get_app_context_manager()in both the sync and asyncmcp_auth_hookwrappers. - Tests.
test_concurrent_tool_calls_get_isolated_sessionsproves isolation non-vacuously (realasyncio.gather, asserts distinctSessionobjects and that call A's session survives call B's teardown and still executes SQL); the counterfactual with the greenlet scopefunc restored demonstrates the collision the change prevents.
On the three open codeant-ai threads:
- app.py:1037 / server.py:977 (factory path) — both flag that
create_mcp_app(**factory_config)bypasses the install. Addressed at HEAD:create_mcp_app()now installs the scopefunc itself (app.py:707-726), in addition toinit_fastmcp_server()andrun_server(). These are moot and can be resolved. - auth.py:1051 (request-backed path) — valid but scoped. When
has_request_context()is true the call takes thenullcontext()branch, no token is set, and concurrent request-backed calls on one greenlet again share the greenlet-keyed session. This is not a regression (it is the pre-existing behavior for that path, and the default OSS streamable-http transport pushes an app context rather than a Flask request context, so it takes the fixed path), but the isolation guarantee does not extend to deployments where external middleware pushes a per-request Flask context. Worth either extending the per-call token to that branch (while preserving the externally suppliedg.user) or documenting the per-request-scoping requirement for request-backed deployments.
Nit (optional): keying on id(token) is safe here only because remove() clears the entry while the token is still referenced; keying the registry on the token object directly would remove even the theoretical id()-reuse concern at no cost.
CI is green across the suite.
|
just pending the unresolved codeant comments |
|
@goingforstudying-ctrl One thing before this can go in. Codeant's thread on |
When a request context is active (e.g. Preset's WorkspaceContextMiddleware path), _get_app_context_manager() returned nullcontext() and never set _mcp_session_token, so concurrent request-backed tool calls still shared the greenlet-scoped session — the same race the rest of this change eliminates. Reuse the request context but tag the call, and deregister the per-call session on exit since no app-context teardown runs here.
|
Good catch, thanks. The request-context branch now tags the call with a per-call session token as well (new _request_tool_call_context in auth.py), so request-backed calls get isolated sessions like every other path instead of sharing the greenlet-scoped one. Since no app context is pushed in that branch, the per-call session is deregistered explicitly on exit, leaving the request's own session to the normal request lifecycle. Pushed as 9b73ab4, with a test covering concurrent tool calls under an active request context plus registry cleanup on exit. The full mcp_service unit test module passes locally. |
Code Review Agent Run #82e3beActionable 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 |
|
Hi @goingforstudying-ctrl — heads up so you hear it from me rather than find it: I opened #43007 for #42622 without having seen this PR, which predates mine by ten days and reaches the same diagnosis. @rusackas asked us to compare notes before either lands, so I backported your approach onto a Superset 6.1 harness and measured both under identical conditions. Full write-up is in this comment; the short version:
So the pool consequence is a property of isolating sessions at all, not of how — whichever lands needs a concurrency guard, and I will write one that applies to either. Two things I flagged as costs of your approach, in case they are useful rather than annoying: the install in One caveat: backporting yours to 6.1 needed an import fix because 6.1 imports No stake in which one lands — happy to close mine and bring the measurements and the guard here if the maintainers prefer yours. |
|
@goingforstudying-ctrl LGTM, thanks for chasing the request-context gap down, that thread's closed out now. @verdier appreciate the head-to-head, that's exactly what I was hoping for before picking one. This one got here first and the threads are resolved, so I'm inclined to land this and have you bring the concurrency guard over as a follow-up. I'll close #43007 in favor of this, and thanks for opening it! |
rusackas
added
the
merge-if-green
|
Bito Automatic Review Skipped – PR Already Merged |
Co-authored-by: goingforstudying-ctrl <goingforstudying-ctrl@users.noreply.github.com> Co-authored-by: Evan Rusackas <evan@preset.io> Co-authored-by: Joe Li <joe@preset.io> (cherry picked from commit d0658ba)
5 participants
SUMMARY
Ran into #42622 while looking at the MCP concurrency reports: fire two async tool calls at the same time and one of them can come back with
DetachedInstanceErroreven though each call pushes its own app context. Dug into why, and the session scope is the culprit —db.session(flask-sqlalchemy 2.5.1) keys its registry by greenlet ident, and asyncio tasks on the event loop all live on the same greenlet, so N concurrent tool calls end up sharing ONE Session. Whichever call finishes first pops its app context, the teardown handler callsdb.session.remove(), and every other in-flight call is suddenly holding detached instances. That's also what makes #42567'sgenerate_chartreport failure for a chart it had already committed, and why retrying agents end up creating duplicates.The fix keys the registry on a per-call token instead.
_get_app_context_manager()sets a fresh ContextVar token when it pushes the call's app context, and the MCP scopefunc resolvesdb.sessionfrom that token — each call gets its own Session, and the app-context teardown removes exactly that call's. Anywhere outside an MCP tool call the scopefunc falls back to the plain greenlet ident, so web, CLI, and Celery paths behave exactly as before. It's wired up frominit_fastmcp_server()so every serving mode gets it.The new tests include an interleaved two-task race that reproduces the shared-session teardown (fails without the fix, passes with it), plus a counterfactual that pins down the old behavior: with the greenlet scope restored, both calls really do resolve to the same Session and the first teardown removes it out from under the survivor. Not 100% sure the ContextVar-token approach is what you'd pick long-term vs. upgrading flask-sqlalchemy, but it's contained to the MCP service and doesn't move anything else.
TESTING INSTRUCTIONS
pytest tests/unit_tests/mcp_service/test_session_scope.py(5 new tests)test_g_user_race_condition,test_auth_*,test_middleware,test_rbac_tool_enforcement, the top-levelmcp_servicetests, and the fullchart/tool suite (~1900 tests total) — all greenpre-commit run(mypy, ruff, pylint) passes on the changed filesADDITIONAL INFORMATION