fix(mcp): give concurrent tool calls isolated db sessions - #42629
fix(mcp): give concurrent tool calls isolated db sessions#42629goingforstudying-ctrl wants to merge 4 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}
| if has_request_context(): | ||
| return nullcontext() |
There was a problem hiding this comment.
Suggestion: Request-backed MCP calls bypass _mcp_tool_call_context() and therefore never set _mcp_session_token. When concurrent async requests execute on the same greenlet, their db.session lookups still use the shared greenlet key, so one request's teardown can remove the session used by another and detach its ORM objects. Apply the per-call session token to this path as well, while preserving the externally supplied g.user context. [race condition]
Severity Level: Critical 🚨
- ❌ Concurrent request-backed MCP tools can lose ORM session state.
- ❌ Long-running tools may fail after another request completes.
- ⚠️ Affected deployments include streamable HTTP with external request middleware.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/mcp_service/auth.py
**Line:** 1050:1051
**Comment:**
*Race Condition: Request-backed MCP calls bypass `_mcp_tool_call_context()` and therefore never set `_mcp_session_token`. When concurrent async requests execute on the same greenlet, their `db.session` lookups still use the shared greenlet key, so one request's teardown can remove the session used by another and detach its ORM objects. Apply the per-call session token to this path as well, while preserving the externally supplied `g.user` context.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
|
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 65.44% 65.42% -0.02%
==========================================
Files 2810 2811 +1
Lines 159362 159455 +93
Branches 36372 36381 +9
==========================================
+ Hits 104294 104324 +30
- Misses 53026 53087 +61
- Partials 2042 2044 +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. |
|
|
||
| # Give each tool call its own SQLAlchemy session; the greenlet-scoped | ||
| # default lets concurrent async calls share (and tear down) one session. | ||
| install_mcp_session_scoping() |
There was a problem hiding this comment.
Suggestion: The factory-configured server path in server.py calls create_mcp_app(**factory_config) directly and never reaches init_fastmcp_server(). Consequently, deployments using use_factory_config=True do not install the MCP-aware scope function, so _mcp_tool_call_context() still creates per-call tokens that db.session ignores and concurrent tool calls continue sharing the greenlet-scoped session. Install the scoping during every MCP server creation path, including the factory path. [api mismatch]
Severity Level: Major ⚠️
- ❌ Factory-configured concurrent tools can share SQLAlchemy sessions.
- ❌ One call teardown can detach another call’s ORM instances.
- ⚠️ Default MCP server initialization is not affected.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/mcp_service/app.py
**Line:** 1016:1016
**Comment:**
*Api Mismatch: The factory-configured server path in `server.py` calls `create_mcp_app(**factory_config)` directly and never reaches `init_fastmcp_server()`. Consequently, deployments using `use_factory_config=True` do not install the MCP-aware scope function, so `_mcp_tool_call_context()` still creates per-call tokens that `db.session` ignores and concurrent tool calls continue sharing the greenlet-scoped session. Install the scoping during every MCP server creation path, including the factory path.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
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 |
| mcp_instance = create_mcp_app(**factory_config) | ||
| # The factory path bypasses init_fastmcp_server(), so install the | ||
| # per-tool-call session scoping here as well; without it concurrent | ||
| # tool calls share the greenlet-scoped db.session. | ||
| # Lazy import mirrors init_fastmcp_server() to avoid the circular | ||
| # import through superset.extensions during startup. | ||
| from superset.mcp_service.session_scope import ( # noqa: PLC0415 | ||
| install_mcp_session_scoping, | ||
| ) | ||
|
|
||
| install_mcp_session_scoping() |
There was a problem hiding this comment.
Suggestion: Session scoping is installed only in run_server() after the factory returns. Callers that use the public create_mcp_app() factory directly receive a FastMCP instance whose protected tools still execute with the original greenlet-scoped db.session, so concurrent calls through that supported factory path retain the DetachedInstanceError/session-sharing failure. Install the MCP session scope as part of factory initialization or otherwise guarantee it before the returned instance can serve requests. [api mismatch]
Severity Level: Critical 🚨
- ❌ Direct factory deployments lack concurrent session isolation.
- ❌ Concurrent tool calls can produce detached ORM instances.
- ⚠️ Public factory behavior differs from `run_server()`.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/mcp_service/server.py
**Line:** 967:977
**Comment:**
*Api Mismatch: Session scoping is installed only in `run_server()` after the factory returns. Callers that use the public `create_mcp_app()` factory directly receive a FastMCP instance whose protected tools still execute with the original greenlet-scoped `db.session`, so concurrent calls through that supported factory path retain the DetachedInstanceError/session-sharing failure. Install the MCP session scope as part of factory initialization or otherwise guarantee it before the returned instance can serve requests.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
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 |
2 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