Fix stateful session context propagation and affinity forwarding for appropriate list tools response#2974
Conversation
7df45b6 to
a5748e1
Compare
cde9111 to
9ab8853
Compare
Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com>
Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com>
- Remove debug log statement that leaked auth headers and user context at INFO level (security: token/PII exposure in production logs) - Apply _get_request_context_or_default() to call_tool, read_resource, and list_resource_templates for consistent context recovery in stateful sessions (previously only list_* handlers were updated) - Fix server_id injection to create params dict when absent, ensuring JSON-RPC requests without params still get proper server scoping - Remove orphaned blank line in list_resources left from prior edit - Clean up test imports: remove duplicate pytest, unused asyncio/types/ transport imports, fix isort ordering - Fix test property cleanup in test_list_resources_direct_proxy_meta_ lookup_error to properly save/restore original descriptor - Run black + isort on all changed files Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com>
…tion
Cover all branches in _get_request_context_or_default():
- Fast path when ContextVars are already populated
- Anonymous user string-to-dict conversion
- URL without /servers/{id}/mcp pattern (no server_id match)
- LookupError fallback (no active request context)
- Generic Exception fallback with error logging
- Cookie JWT token forwarding to require_auth_override
Cover server_id injection edge cases:
- JSON-RPC body without params key (params dict created)
- URL without server pattern (no injection occurs)
Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com>
Signed-off-by: Mihai Criveti <crivetimihai@gmail.com>
9ab8853 to
157204a
Compare
_get_request_context_or_default() was returning raw JWT payloads from
require_auth_override() without normalizing to the canonical user
context shape {email, teams, is_admin, is_authenticated} that MCP
handlers expect. This caused handlers to look up missing keys (email,
teams) from the raw JWT which uses different field names (sub,
token_use, nested user.is_admin).
Changes:
- Add _normalize_jwt_payload() that mirrors streamable_http_auth
normalization: extracts email from sub/email, resolves teams via
normalize_token_teams (API) or _resolve_teams_from_db_sync (session),
and detects nested user.is_admin
- call_tool now extracts app_user_email from the already-recovered
user_context instead of reading from stale ContextVar via
get_user_email_from_context()
- Update test mocks to use realistic JWT payloads instead of
pre-normalized dicts
- Add 6 dedicated tests for _normalize_jwt_payload covering API token,
session token (admin/non-admin), nested is_admin, email fallback,
and no-email edge case
Signed-off-by: Mihai Criveti <crivetimihai@gmail.com>
Review Changes SummaryRebased onto Commit 1:
|
Verify call_tool passes app_user_email from the recovered fallback context (via _get_request_context_or_default) rather than reading from the stale user_context_var ContextVar. This prevents regression of the stateful session identity propagation fix. Signed-off-by: Mihai Criveti <crivetimihai@gmail.com>
…DAR) Signed-off-by: Mihai Criveti <crivetimihai@gmail.com>
Pre-existing issues identified during reviewTwo pre-existing gaps were identified during the review of this PR. Neither was introduced by this PR — they exist on 1. [HIGH]
|
… appropriate list tools response (#2974) * maintain server_id context with session affinity Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * test cases for new improvements Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * linting Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * fix: address review findings for stateful session context propagation - Remove debug log statement that leaked auth headers and user context at INFO level (security: token/PII exposure in production logs) - Apply _get_request_context_or_default() to call_tool, read_resource, and list_resource_templates for consistent context recovery in stateful sessions (previously only list_* handlers were updated) - Fix server_id injection to create params dict when absent, ensuring JSON-RPC requests without params still get proper server scoping - Remove orphaned blank line in list_resources left from prior edit - Clean up test imports: remove duplicate pytest, unused asyncio/types/ transport imports, fix isort ordering - Fix test property cleanup in test_list_resources_direct_proxy_meta_ lookup_error to properly save/restore original descriptor - Run black + isort on all changed files Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * test: add missing coverage for context propagation and affinity injection Cover all branches in _get_request_context_or_default(): - Fast path when ContextVars are already populated - Anonymous user string-to-dict conversion - URL without /servers/{id}/mcp pattern (no server_id match) - LookupError fallback (no active request context) - Generic Exception fallback with error logging - Cookie JWT token forwarding to require_auth_override Cover server_id injection edge cases: - JSON-RPC body without params key (params dict created) - URL without server pattern (no injection occurs) Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: normalize raw JWT payload in stateful session context recovery _get_request_context_or_default() was returning raw JWT payloads from require_auth_override() without normalizing to the canonical user context shape {email, teams, is_admin, is_authenticated} that MCP handlers expect. This caused handlers to look up missing keys (email, teams) from the raw JWT which uses different field names (sub, token_use, nested user.is_admin). Changes: - Add _normalize_jwt_payload() that mirrors streamable_http_auth normalization: extracts email from sub/email, resolves teams via normalize_token_teams (API) or _resolve_teams_from_db_sync (session), and detects nested user.is_admin - call_tool now extracts app_user_email from the already-recovered user_context instead of reading from stale ContextVar via get_user_email_from_context() - Update test mocks to use realistic JWT payloads instead of pre-normalized dicts - Add 6 dedicated tests for _normalize_jwt_payload covering API token, session token (admin/non-admin), nested is_admin, email fallback, and no-email edge case Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * test: add regression test for call_tool recovered email propagation Verify call_tool passes app_user_email from the recovered fallback context (via _get_request_context_or_default) rather than reading from the stale user_context_var ContextVar. This prevents regression of the stateful session identity propagation fix. Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: add missing docstring params for _normalize_jwt_payload (flake8 DAR) Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> --------- Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> Co-authored-by: Mihai Criveti <crivetimihai@gmail.com> Signed-off-by: Vishu Bhatnagar <vishu.bhatnagar@ibm.com>
… appropriate list tools response (IBM#2974) * maintain server_id context with session affinity Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * test cases for new improvements Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * linting Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * fix: address review findings for stateful session context propagation - Remove debug log statement that leaked auth headers and user context at INFO level (security: token/PII exposure in production logs) - Apply _get_request_context_or_default() to call_tool, read_resource, and list_resource_templates for consistent context recovery in stateful sessions (previously only list_* handlers were updated) - Fix server_id injection to create params dict when absent, ensuring JSON-RPC requests without params still get proper server scoping - Remove orphaned blank line in list_resources left from prior edit - Clean up test imports: remove duplicate pytest, unused asyncio/types/ transport imports, fix isort ordering - Fix test property cleanup in test_list_resources_direct_proxy_meta_ lookup_error to properly save/restore original descriptor - Run black + isort on all changed files Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * test: add missing coverage for context propagation and affinity injection Cover all branches in _get_request_context_or_default(): - Fast path when ContextVars are already populated - Anonymous user string-to-dict conversion - URL without /servers/{id}/mcp pattern (no server_id match) - LookupError fallback (no active request context) - Generic Exception fallback with error logging - Cookie JWT token forwarding to require_auth_override Cover server_id injection edge cases: - JSON-RPC body without params key (params dict created) - URL without server pattern (no injection occurs) Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: normalize raw JWT payload in stateful session context recovery _get_request_context_or_default() was returning raw JWT payloads from require_auth_override() without normalizing to the canonical user context shape {email, teams, is_admin, is_authenticated} that MCP handlers expect. This caused handlers to look up missing keys (email, teams) from the raw JWT which uses different field names (sub, token_use, nested user.is_admin). Changes: - Add _normalize_jwt_payload() that mirrors streamable_http_auth normalization: extracts email from sub/email, resolves teams via normalize_token_teams (API) or _resolve_teams_from_db_sync (session), and detects nested user.is_admin - call_tool now extracts app_user_email from the already-recovered user_context instead of reading from stale ContextVar via get_user_email_from_context() - Update test mocks to use realistic JWT payloads instead of pre-normalized dicts - Add 6 dedicated tests for _normalize_jwt_payload covering API token, session token (admin/non-admin), nested is_admin, email fallback, and no-email edge case Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * test: add regression test for call_tool recovered email propagation Verify call_tool passes app_user_email from the recovered fallback context (via _get_request_context_or_default) rather than reading from the stale user_context_var ContextVar. This prevents regression of the stateful session identity propagation fix. Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: add missing docstring params for _normalize_jwt_payload (flake8 DAR) Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> --------- Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> Co-authored-by: Mihai Criveti <crivetimihai@gmail.com> Signed-off-by: Nithin Katta <Nithin.Katta@ibm.com>
… appropriate list tools response (#2974) * maintain server_id context with session affinity Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * test cases for new improvements Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * linting Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * fix: address review findings for stateful session context propagation - Remove debug log statement that leaked auth headers and user context at INFO level (security: token/PII exposure in production logs) - Apply _get_request_context_or_default() to call_tool, read_resource, and list_resource_templates for consistent context recovery in stateful sessions (previously only list_* handlers were updated) - Fix server_id injection to create params dict when absent, ensuring JSON-RPC requests without params still get proper server scoping - Remove orphaned blank line in list_resources left from prior edit - Clean up test imports: remove duplicate pytest, unused asyncio/types/ transport imports, fix isort ordering - Fix test property cleanup in test_list_resources_direct_proxy_meta_ lookup_error to properly save/restore original descriptor - Run black + isort on all changed files Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * test: add missing coverage for context propagation and affinity injection Cover all branches in _get_request_context_or_default(): - Fast path when ContextVars are already populated - Anonymous user string-to-dict conversion - URL without /servers/{id}/mcp pattern (no server_id match) - LookupError fallback (no active request context) - Generic Exception fallback with error logging - Cookie JWT token forwarding to require_auth_override Cover server_id injection edge cases: - JSON-RPC body without params key (params dict created) - URL without server pattern (no injection occurs) Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: normalize raw JWT payload in stateful session context recovery _get_request_context_or_default() was returning raw JWT payloads from require_auth_override() without normalizing to the canonical user context shape {email, teams, is_admin, is_authenticated} that MCP handlers expect. This caused handlers to look up missing keys (email, teams) from the raw JWT which uses different field names (sub, token_use, nested user.is_admin). Changes: - Add _normalize_jwt_payload() that mirrors streamable_http_auth normalization: extracts email from sub/email, resolves teams via normalize_token_teams (API) or _resolve_teams_from_db_sync (session), and detects nested user.is_admin - call_tool now extracts app_user_email from the already-recovered user_context instead of reading from stale ContextVar via get_user_email_from_context() - Update test mocks to use realistic JWT payloads instead of pre-normalized dicts - Add 6 dedicated tests for _normalize_jwt_payload covering API token, session token (admin/non-admin), nested is_admin, email fallback, and no-email edge case Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * test: add regression test for call_tool recovered email propagation Verify call_tool passes app_user_email from the recovered fallback context (via _get_request_context_or_default) rather than reading from the stale user_context_var ContextVar. This prevents regression of the stateful session identity propagation fix. Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: add missing docstring params for _normalize_jwt_payload (flake8 DAR) Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> --------- Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> Co-authored-by: Mihai Criveti <crivetimihai@gmail.com> Signed-off-by: Nithin Katta <Nithin.Katta@ibm.com>
… appropriate list tools response (#2974) * maintain server_id context with session affinity Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * test cases for new improvements Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * linting Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> * fix: address review findings for stateful session context propagation - Remove debug log statement that leaked auth headers and user context at INFO level (security: token/PII exposure in production logs) - Apply _get_request_context_or_default() to call_tool, read_resource, and list_resource_templates for consistent context recovery in stateful sessions (previously only list_* handlers were updated) - Fix server_id injection to create params dict when absent, ensuring JSON-RPC requests without params still get proper server scoping - Remove orphaned blank line in list_resources left from prior edit - Clean up test imports: remove duplicate pytest, unused asyncio/types/ transport imports, fix isort ordering - Fix test property cleanup in test_list_resources_direct_proxy_meta_ lookup_error to properly save/restore original descriptor - Run black + isort on all changed files Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * test: add missing coverage for context propagation and affinity injection Cover all branches in _get_request_context_or_default(): - Fast path when ContextVars are already populated - Anonymous user string-to-dict conversion - URL without /servers/{id}/mcp pattern (no server_id match) - LookupError fallback (no active request context) - Generic Exception fallback with error logging - Cookie JWT token forwarding to require_auth_override Cover server_id injection edge cases: - JSON-RPC body without params key (params dict created) - URL without server pattern (no injection occurs) Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: normalize raw JWT payload in stateful session context recovery _get_request_context_or_default() was returning raw JWT payloads from require_auth_override() without normalizing to the canonical user context shape {email, teams, is_admin, is_authenticated} that MCP handlers expect. This caused handlers to look up missing keys (email, teams) from the raw JWT which uses different field names (sub, token_use, nested user.is_admin). Changes: - Add _normalize_jwt_payload() that mirrors streamable_http_auth normalization: extracts email from sub/email, resolves teams via normalize_token_teams (API) or _resolve_teams_from_db_sync (session), and detects nested user.is_admin - call_tool now extracts app_user_email from the already-recovered user_context instead of reading from stale ContextVar via get_user_email_from_context() - Update test mocks to use realistic JWT payloads instead of pre-normalized dicts - Add 6 dedicated tests for _normalize_jwt_payload covering API token, session token (admin/non-admin), nested is_admin, email fallback, and no-email edge case Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * test: add regression test for call_tool recovered email propagation Verify call_tool passes app_user_email from the recovered fallback context (via _get_request_context_or_default) rather than reading from the stale user_context_var ContextVar. This prevents regression of the stateful session identity propagation fix. Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: add missing docstring params for _normalize_jwt_payload (flake8 DAR) Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> --------- Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> Co-authored-by: Mihai Criveti <crivetimihai@gmail.com>
🐛 Bug-fix PR
Closes #2973
📌 Summary
Fixed a bug where server_id and other request context (headers, user auth) were lost when using stateful sessions and server_id was lost when using session_affinity. It ensures that:
Long-lived session tasks can dynamically recover context from the active request.
Internal session affinity forwarding for POST requests correctly preserves the server_id.
🔁 Reproduction Steps
/servers/server-id/mcpand list the toolsBefore fix: Returns tools for "default_server_id" (empty/all).
After fix: Returns tools strictly for
<id>.🐞 Root Cause
Two separate issues contributed to this bug:
ContextVar Loss: in stateful sessions, the run_server task is long-lived and spawned separately. The contextvars (server_id_var) set during the initial connection or previous requests do not propagate to this isolated task context, causing handlers to see the default value.
Forwarding Data Loss: When session affinity is enabled, stateful POST requests are forwarded internally to the /rpc endpoint in main.py. The original streamablehttp_transport extracted server_id from the URL, but
main.py's /rpc endpoint expects it in the JSON body (params). The forwarding logic was not injecting the server_id from the URL into the JSON body, so main.py received server_id=None.
💡 Fix Description
Dynamic Context Recovery: Introduced
_get_request_context_or_defaulthelper in streamablehttp_transport.py. This function first checks contextvars (fast path). If missing (stateful case), it inspects mcp_app.request_context to retrieve the underlying starlette.requests.Request object and extracts server_id, headers, and re-verifies user auth viaverify_credentials.
Affinity Forwarding Injection: Updated handle_streamable_http to extract the server_id from the URL regex and inject it into the params of the JSON-RPC body before forwarding the request to /rpc. This ensures correct routing in main.py.
🧪 Verification
make lintmake testmake coverage📐 MCP Compliance (if relevant)
✅ Checklist
make black isort pre-commit)