fix(mcp): cache StatefulHTTP JWT admin lookups

[Lang-Akshay]

Fixes #4988

Problem

The StreamableHTTP stateful-session fallback called _normalize_jwt_payload() for requests without propagated context. Each call opened a new SessionLocal() for is_user_admin(), defeating the session-scoped memoization and adding a synchronous database lookup to the request path.

Fix

• Check the existing AuthCache before opening a database session.
• Warm the cache on misses through the existing batched auth-context lookup.
• Preserve the platform-admin fast path and direct DB fallback when caching is unavailable.
• Enforce token revocation, active-user, strict user-in-DB, and API/legacy team-membership checks in the fallback.
• Keep JWT team normalization and session-token policy unchanged.

[msureshkumar88]

Thanks for tackling the cold-cache latency regression from #4988 — the AuthCache/batched-lookup path mirrors the primary _auth_jwt flow nicely and should close the ~1s per-request hit described in the issue.

While going through the added enforcement logic I ran into two spots where the new checks don't actually take effect at runtime, plus a couple of smaller things worth a look before merge.

Blocking

1. The new revocation/disabled/team-membership rejections get silently swallowed by the caller.
_normalize_jwt_payload now raises HTTPException(401/403) for a revoked token, a disabled account, strict user-in-DB misses, and stale team membership (all inside this PR's new code). But its only caller, _get_request_context_or_default (streamablehttp_transport.py:2182-2194, unchanged by this PR), wraps the call in:

except Exception as e:
    logger.warning("Failed to recover user context in stateful session: %s", e)
    user_ctx = {}

HTTPException is an Exception, so it gets caught here too — the request falls through to an anonymous user_ctx = {} instead of being rejected. Worse, _should_enforce_streamable_rbac({}) then skips RBAC entirely (no is_authenticated key), so the request proceeds with no permission check rather than the intended 401/403. Could you catch HTTPException separately in the caller and propagate it, so this path actually enforces what it's designed to enforce?

2. require_user_in_db isn't re-checked when the cache already has a user.

if cached_user is None and settings.require_user_in_db and email != platform_admin_email:
    raise HTTPException(...)

This only fires when cached_user is None. The primary auth path (mcpgateway/auth.py:1714) re-verifies against the DB even when cached_ctx.user is populated:

if cached_ctx.user:
    if settings.require_user_in_db:
        db_user = await asyncio.to_thread(_get_user_by_email_sync, email)
        if db_user is None:
            raise HTTPException(...)

Would it make sense to mirror that here? Otherwise a deleted user with a still-cached record stays trusted for the cache TTL on this path, diverging from the primary path's guarantee.

Suggestions

3. Team-membership revalidation logic is now duplicated a third time. The same cache-check/DB-query/cache-write pattern already exists in token_scoping.py::_validate_team_membership and in the primary path (auth.py ~5399-5439). This new copy (streamablehttp_transport.py:2344-2369) uses a slightly different error path (HTTPException vs. the primary path's _send_error) even though the message text matches — a small refactor into a shared helper would keep these from drifting apart.

4. Test coverage gap around the caller. All the new tests call _normalize_jwt_payload directly and assert it raises — none exercise it through the actual caller, _get_request_context_or_default, so the swallow-bug in #1 wasn't caught. Worth adding a test that goes through the real fallback path and asserts the request is actually denied, plus a case for #2 (cached user present + require_user_in_db=True + user removed from DB).

Minor

5. .secrets.baseline references tests/performance/benchmark_issue_4988.sh:14, but that file isn't part of this diff (and isn't on main). Looks like a benchmark script that didn't get git add-ed — either include it or drop the stray baseline entry.

Scope-wise, this bundles the perf fix with new security enforcement (revocation/active-user/team-membership checks) that #4988 doesn't ask for — not a blocker, but calling it out since it's what's driving the extra surface area under review here.

No schema/UI changes, nothing else jumped out.

[Lang-Akshay]

Finding	Status	Resolution
1. _get_request_context_or_default catches _normalize_jwt_payload’s HTTPException at lines 2182–2194, replaces the context with {}, and downstream _should_enforce_streamable_rbac skips enforcement.	valid	Addressed
2. With require_user_in_db=True, the cache-hit branch only checks cached_user is None; a stale cached user remains trusted. The primary auth path explicitly rechecks the database for cached users.	valid	Addressed
3. Team-membership validation is duplicated in the fallback normalizer, token_scoping.py, and the primary Streamable HTTP auth path. The implementations and error handling can drift.	valid	Addressed
4. Added tests exercise _normalize_jwt_payload directly, but no test verifies propagation through _get_request_context_or_default for revoked, disabled, missing, or stale-membership users. Existing caller coverage only tests generic auth failure.	valid	Addressed
5. .secrets.baseline adds tests/performance/benchmark_issue_4988.sh:14, but the file is absent both from the branch and main.	valid	Addressed
6. The issue requests a performance fix for repeated database lookups. Revocation, disabled-user, and membership enforcement expand the PR beyond that scope.	out of scope	Out of scope

All valid findings from the reviewers were addressed; finding 6 was left out of scope as requested.

[gandhipratik203]

token_scoping.py:888-901: next(get_db()) now runs unconditionally, so every team-scoped request pays a session checkout even on a cache hit (previously only on a miss). validate_token_team_membership opens its own session when db is None, so the pre-acquisition can be dropped.

streamablehttp_transport.py:2310+: the direct-lookup branch never calls set_auth_context, so with auth_cache_batch_queries=False the cache stays permanently cold. Mirror the write the primary path does at 5231+.

[Lang-Akshay]

#	File	Finding	Resolution	Commit
1	token_scoping.py:888-901	_check_team_membership created a DB session via next(get_db()) unconditionally before calling validate_token_team_membership, which checks cache first and may return without touching the DB. Every team-scoped request paid a connection pool checkout even on cache hits.	Dropped the 16-line session wrapper. The shared helper already manages its own SessionLocal() on cache miss; the caller just passes db through (or None).	9c334e5ba
2	streamablehttp_transport.py:2310+	The direct-lookup branch in _normalize_jwt_payload (cache enabled, auth_cache_batch_queries=False, cache miss) never called set_auth_context, so the cache stayed permanently cold. The primary path at 5328+ writes after the same lookup.	Added set_auth_context after successful user lookup, mirroring the primary path. Subsequent fallback-path requests now hit cache instead of DB.	9c334e5ba

[gandhipratik203]

LGTM! Thanks for the changes!

[msureshkumar88]

Thanks for this — the hardening here is genuinely valuable. Bringing the stateful-session fallback in line with the primary _auth_jwt path closes a real gap: before this PR, _normalize_jwt_payload() did nothing but an is_user_admin() lookup, with no revocation check, no is_active check, and no user-existence check at all. Routing team-membership validation through the new shared validate_token_team_membership() helper is also a nice consolidation — the old inline cache-then-DB block in _StreamableHttpAuthHandler was duplicated logic waiting to drift.

One blocking item, then a few optional notes.

Blocking

Cache warm writes a partial user dict into the shared AuthCache

mcpgateway/transports/streamablehttp_transport.py:2344-2360

The new cache-warm in the direct-fallback branch stores a three-field user dict:

CachedAuthContext(
    user={
        "email": getattr(user_record, "email", email),
        "is_admin": bool(getattr(user_record, "is_admin", False)),
        "is_active": bool(getattr(user_record, "is_active", True)),
    },
    ...
)

Every other writer into this cache stores the full ten-field shape — _get_auth_context_batched_sync() (mcpgateway/auth.py:1202-1213) includes password_hash, full_name, auth_provider, password_change_required, email_verified_at, created_at, and updated_at, and the batched branch a few lines above in this same function (line ~2298) passes that full dict straight through.

Because get_auth_cache() is a process-wide singleton keyed on (email, jti), a subsequent REST request carrying the same token hits the cached-context branch in get_current_user() (mcpgateway/auth.py:1772), calls _user_from_cached_dict() (mcpgateway/auth.py:1289), and gets an EmailUser built from .get(..., default) fallbacks for the seven missing fields:

• password_change_required → False, so a user flagged for a forced password change silently isn't
• auth_provider → "local", masking an SSO-provisioned identity
• email_verified_at → None
• password_hash → "", created_at/updated_at → now()

That holds for the remainder of the cache TTL, and which shape wins is a race between whichever path warms the entry first.

Two ways to resolve it, whichever fits better:

1. Build the same full dict the batched path does — either by reusing _get_auth_context_batched_sync() here instead of _get_user_by_email_sync(), or by extracting the dict construction in auth.py into a small shared helper (_user_to_cache_dict(user_record)) and calling it from both sites. This keeps the perf win.
2. Skip the cache warm on this branch entirely and leave the DB round trip. Loses the optimization, but no risk of a partial entry.

A regression test asserting that an entry written by the fallback path round-trips through _user_from_cached_dict() with auth_provider and password_change_required intact would lock this down.

Suggestions (non-blocking)

resolve_session_teams() isn't given preresolved_db_teams in the batched branch

mcpgateway/transports/streamablehttp_transport.py:2374

_get_auth_context_batched_sync() already returns team_ids/team_names one call earlier, but the session-token branch calls:

final_teams = await resolve_session_teams(payload, email, {"is_admin": effective_is_admin})

The sibling _StreamableHttpAuthHandler in this same file (lines ~5395-5401) passes preresolved_db_teams=preresolved for exactly this case. Without it, resolve_session_teams() → _resolve_teams_from_db() re-fetches membership that was just retrieved, which cuts against the round-trip reduction this PR is aiming for. Threading the batched team_ids through (where the branch has them) would close that.

Similarly at line 2387: for API/legacy tokens, validate_token_team_membership() does its own cache-then-DB membership resolution, duplicating what the batched context already fetched. Lower priority since it only affects the batched branch, but worth a look if you're optimizing this path.

Log messages lost some triage detail

mcpgateway/transports/streamablehttp_transport.py:5427

The refactor collapsed three distinct log lines — a cached rejection (with the (cached) suffix), a DB rejection naming the specific missing_teams, and a separate cache-hit event — into one generic message. An operator triaging a spike in 403s can no longer tell stale-cache rejections from genuine membership changes, or see which team IDs mismatched. If validate_token_team_membership() could return or surface the missing set (or the cache outcome already flowing through on_cache_event), the caller could restore that detail.

Notes — no change requested

A couple of things that looked suspicious on first read but check out, recorded so they don't get re-raised:

• if email == platform_admin_email: db_user_is_admin = True (line 2363) reads like an unconditional admin override that ignores a DB is_admin=False, and it does diverge from get_current_user(), which only bootstraps platform-admin when no DB row exists. But it faithfully reproduces the behavior of the is_user_admin() call it replaces — mcpgateway/utils/admin_check.py:83-85 returns True on the platform-admin email with no DB check at all. No behavior change here; if that semantic is worth revisiting it belongs in its own change against admin_check.py. Same reasoning covers the getattr(settings, "platform_admin_email", "") empty-string default at line 2250 — it matches the helper rather than the other call sites in this file.

• The require_user_in_db gating at lines 2266, 2289, and 2334 doesn't fully match the primary path, which rejects an unknown user with 401 regardless of that setting (mcpgateway/auth.py:2132-2137), so a valid JWT for a deleted account still authenticates here when REQUIRE_USER_IN_DB=false. Since the pre-PR code had no existence check whatsoever, this is a clear improvement rather than a regression — flagging it only so the remaining parity gap is a deliberate choice rather than an oversight. Happy to see it as follow-up work.

[Lang-Akshay]

Finding: Cache warm writes partial user dict into shared AuthCache

Status: Fixed

Problem

The direct-fallback branch in streamablehttp_transport.py (line ~2344) warmed the process-wide AuthCache with a 3-field user dict (email, is_admin, is_active), while every other cache writer stores the full 10-field shape. Because the cache is a singleton keyed on (email, jti), a subsequent request hitting the cached-context path in get_current_user() would reconstruct an EmailUser via _user_from_cached_dict() with wrong defaults for the 7 missing fields:

• password_change_required defaulted to False -- silently bypassing forced password changes
• auth_provider defaulted to "local" -- masking SSO-provisioned identities
• password_hash defaulted to "", created_at/updated_at to now(), etc.

Which shape won the cache was a race between whichever code path warmed the entry first for a given (email, jti) pair.

Fix

Expanded the 3-field dict to the full 10-field shape, matching the batched-path writer in auth.py:1202-1213. The user_record returned by _get_user_by_email_sync() already had all fields -- they were just being discarded. No new helpers, no new files.

Test

Added TestUserFromCachedDictFullShape in test_auth.py with two cases:

1. Full 10-field dict round-trips through _user_from_cached_dict with auth_provider and password_change_required intact.
2. Partial 3-field dict documents the wrong defaults it would produce (the old bug shape).

[msureshkumar88]

Thanks for turning this around — the blocking item from the last round is resolved and I'm happy with where this landed.

Verified fixed

Partial user dict in the cache warm (streamablehttp_transport.py:2338-2362) — the fallback branch now writes the full ten-field shape (password_hash, full_name, auth_provider, password_change_required, email_verified_at, created_at, updated_at all present), matching _get_auth_context_batched_sync(). A REST request that later hits the cached-context branch and rebuilds via _user_from_cached_dict() no longer silently picks up password_change_required=False / auth_provider="local" defaults. Whichever path warms the entry first, the shape is now the same.

Also confirming the earlier round is fully addressed:

• except HTTPException: raise added at both _get_request_context_or_default handlers (lines 2189, 2200), so the new 401/403 rejections actually propagate instead of falling through to an anonymous user_ctx = {}.
• require_user_in_db is now re-verified against the DB even on a cache hit, matching auth.py:1714.
• Team-membership revalidation consolidated into the shared validate_token_team_membership() helper — three copies down to one, and token_scoping.py sheds 56 lines for it.
• The stray tests/performance/benchmark_issue_4988.sh entry is gone from .secrets.baseline.

Test suites for the three touched modules pass clean, and CI is green (31 success / 5 skipped).

Re-checked and cleared — recording so these don't get re-raised

I took another pass specifically to confirm a few things that read as suspicious but aren't:

• if email == platform_admin_email: db_user_is_admin = True (line ~2371) still looks like an unconditional override of a DB is_admin=False. I checked this empirically rather than by inspection this time: I ran the same payload (platform-admin email, DB record with is_admin=False) through _normalize_jwt_payload on both this branch and a clean main worktree, and both return is_admin=True. The reason is is_user_admin() — the helper this line replaces — has an explicit no-DB-round-trip fast path at mcpgateway/utils/admin_check.py:83-85. So this is behavior-preserving, not new. Worth noting the natural comparison to auth.py:1899/2125 doesn't apply: those are the virtual-user-creation bootstrap sites, a different concern from the visibility fast path. If that semantic deserves revisiting, it's a change against admin_check.py and affects every caller, not just this path.

• _check_token_revoked_sync at line 2314 isn't individually wrapped, so a transient DB error there unwinds to the outer handler and yields an empty context. Pre-PR the same position had an unwrapped is_user_admin() call under the same handler, so the shape is unchanged — and with the except HTTPException: raise added above, genuine 401/403s now propagate correctly, which is a net improvement. Failing closed to anonymous on a DB error seems defensible.

Optional follow-ups (not blocking)

• resolve_session_teams() at line ~2381 still doesn't receive preresolved_db_teams even though the batched branch has team_ids in hand one call earlier (auth.py:1269), and the sibling path passes it (auth.py:1826-1827). Costs a redundant membership query on that branch.
• The generic "no longer member of teams" warning still drops the specific missing team IDs the old code logged.
• The fix above is in place but there's no test asserting a fallback-warmed entry round-trips through _user_from_cached_dict() with auth_provider / password_change_required intact — would be a cheap guard against regression.

Approving. Nice cleanup on the shared helper.