litellm

tiennm99/litellm

Fork 0

mirror of https://github.com/tiennm99/litellm.git synced 2026-07-29 18:21:26 +00:00

Files

T

History

c792df64d2 feat(mcp): support stateless and stateful clients via session-id routing (#26857 )

* feat(mcp): support stateless and stateful clients via session-id routing

- Add session_manager_stateful (stateless=False) alongside stateless
- Route by mcp-session-id: has ID → stateful, initialize (no ID) → stateful, else → stateless
- Peek POST body to detect initialize for routing; replay via wrapped receive
- Handle stale session IDs for both managers
- Add test_mcp_routing_initialize_to_stateful_no_session_to_stateless
- Update test_valid_mcp_session_id_is_preserved, test_concurrent_initialize_session_managers

Made-with: Cursor

* fix(mcp): respect stateful routing and harden initialize detection

Ensure streamable MCP requests are dispatched via the computed target session manager, and guard initialize detection against non-object JSON bodies. Update stale-session test patches to target the stateful manager so routing assertions remain correct.

Made-with: Cursor

* test(mcp): patch stateless/stateful managers in concurrency init test

Update concurrent session-manager initialization test to patch session_manager_stateless and session_manager_stateful directly, matching initialize_session_managers() behavior and preventing NameError from undefined mocks.

Made-with: Cursor

* Fix tests

* Fix tests

* Fix MCP stateful routing edge cases

* Fix stateful MCP auth context refresh

* Fix MCP stateful session cleanup

* fix(mcp): bind stateful sessions to creator and reject hijacks

Stateful mcp-session-id was usable by any authenticated proxy caller. Track
the session creator's hashed API key (or user_id) when a new session is
issued and reject mismatched callers with 403 before _set_or_update_auth_context
overwrites the stored MCPAuthenticatedUser. Also formats nested with-statements
in test_mcp_stale_session.py and fixes a pre-existing AsyncMock mismatch in
test_stale_mcp_session_id_is_stripped.

* fix(mcp): serialize concurrent requests on same stateful session

Bugbot's 'Concurrent requests share context' finding: _update_auth_context
mutates the single MCPAuthenticatedUser stored per session in place on
every request, so two requests sharing one mcp-session-id can overwrite
each other's mcp_servers / auth headers / oauth state / client_ip while
in-flight callbacks are still reading the same object.

Owner-binding alone narrows this to same-principal racing, but the
in-place mutation race remains. Add a per-session asyncio.Lock around
handle_request so concurrent same-session requests run sequentially. The
lock is allocated on demand and torn down with the rest of the session
state on DELETE / idle expiry.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* fix(mcp): include OAuth2 bearer in stateful session owner fingerprint

UserAPIKeyAuth() for OAuth2 passthrough has no api_key/user_id, so every
OAuth caller fingerprinted to "anonymous" and could hijack another OAuth
caller's mcp-session-id. Hash the upstream Authorization header into the
fingerprint as oauth:<sha256>.

* fix(mcp): don't hold stateful session lock for streaming GETs

The per-session lock wraps handle_request, so a long-lived GET (SSE
stream held open for the life of the session) would block every
subsequent POST on the same mcp-session-id. Only POST/DELETE mutate the
shared MCPAuthenticatedUser, so it's sufficient to serialize those —
GETs run lock-free and stream concurrently.

* fix(mcp): allow None user_api_key_auth in MCPAuthenticatedUser

The set_auth_context / _set_or_update_auth_context / _update_auth_context
helpers in server.py all accept Optional[UserAPIKeyAuth] and pass it
straight into MCPAuthenticatedUser, but the dataclass-style constructor
typed user_api_key_auth as required UserAPIKeyAuth. Mypy flagged this on
the stateful-routing branch:

  server.py:3227: error: Incompatible types in assignment (expression
    has type "UserAPIKeyAuth | None", variable has type "UserAPIKeyAuth")
  server.py:3255: error: Argument "user_api_key_auth" to
    "MCPAuthenticatedUser" has incompatible type "UserAPIKeyAuth | None";
    expected "UserAPIKeyAuth"

Widen the parameter type to Optional[UserAPIKeyAuth] to match the call
sites. Runtime behavior is unchanged.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* style: replace with new alias

* fix(mcp): fall back to client_ip in stateful session owner fingerprint

Addresses Greptile review on PR #26857: when no API key, user_id, or
OAuth bearer is available (e.g. unauthenticated/passthrough callers),
the owner fingerprint collapsed to a single 'anonymous' value, allowing
two unrelated callers to drive each other's stateful MCP sessions.

Fold client IP into the fingerprint as a fallback identity signal so
distinct anonymous sources do not share an owner identity.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Fix active stateful MCP session cleanup

* test(mcp): cancel leaked stateful auth-context cleanup task

initialize_session_managers() spawns a real asyncio.create_task running
_cleanup_expired_stateful_session_auth_contexts(). The
test_concurrent_initialize_session_managers test was saving and
restoring the session-manager context-manager globals but did not save,
cancel, or restore _stateful_auth_context_cleanup_task.

Because pyproject.toml sets asyncio_default_fixture_loop_scope=session,
the event loop is shared across tests in the same session, so the
leaked task kept running against module-level dicts for the rest of the
test run. Save and cancel the task in the finally block so the test
fully cleans up after itself.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Fix stateful MCP session fingerprinting

* Hash MCP session user owner fingerprints

* Fix stale MCP session DELETE cleanup

* fix(mcp): harden owner fingerprint hashing for non-str api keys

_owner_fingerprint_for assumed api_key/user_id supported .encode();
MagicMock-based tests (and any non-str truthy values) crashed with
TypeError before routing. Only hash str/bytes secrets; fall through
otherwise so MCP routing and session tests behave correctly.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Fix MCP stateful cleanup loop resilience

* Fix stateful MCP initialize auth capture

* fix(mcp): drop orphan per-session lock when auth context absent

Defensive cleanup for _stateful_session_locks entries created on
sessions that never enter _stateful_session_auth_contexts. The
periodic cleanup loop only iterates auth_context_last_seen, so such
locks would otherwise live forever. Add a test that reproduces the
leak and verifies the request finalizer pops the lock.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* chore(mcp): trim verbose comment on lock cleanup

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Fix stateful MCP delete failure tracking

* fix test

* fix(mcp): cap routing-peek body size to bound pre-dispatch memory

Authenticated clients that POST without an mcp-session-id forced the proxy
to buffer the entire request body before routing, since the peek loop
drained every body chunk to decide whether the JSON-RPC method was
'initialize'. Cap the peek at 4 KB (more than enough for an initialize
envelope) and let the remainder stream through wrapped_receive into the
downstream handler.

* test: replace dall-e-3 with gpt-image-1 in health check and router tests (#27813)

OpenAI returns 'The model dall-e-3 does not exist' for the test account,
breaking test_openai_img_gen_health_check and test_image_generation.
Switch to gpt-image-1, matching the existing TestOpenAIGPTImage1 pattern.

* fix(tests): drop dall-e-only test classes; route live image tests via gpt-image-1

Second wave of failures from the 2026-05-12 DALL-E shutdown:
- tests/image_gen_tests/test_image_edits.py::TestOpenAIImageEditDallE2
  and tests/image_gen_tests/test_image_generation.py::TestOpenAIDalle3
  are explicitly named for the deprecated models and can't pass; remove.
  gpt-image-1 coverage already exists in sibling classes.
- tests/local_testing/test_router.py image gen tests use dall-e-3 only
  as a routing example; swap to gpt-image-1.
- tests/local_testing/test_custom_callback_input.py image_generation
  success/failure paths swapped to gpt-image-1.

* Fix MCP initialize session active tracking

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* Fix MCP reinitialize session tracking

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* Fix MCP reinitialize auth context aliasing

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* Apply black formatting after merge

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Run owner-binding 403 before consuming POST body

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Harden MCP routing peek bound and stateful purge race

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Remove inadvertently committed Next.js build artifacts

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Run owner check before stale MCP session cleanup

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* fix(mcp): reverse cleanup ordering to terminate transport before clearing owner

Reverses _purge_expired_stateful_session_auth_contexts so the transport
is popped from server_instances and terminated BEFORE owner/auth tracking
is cleared. The previous order left a window where _stateful_session_owners
was already empty but server_instances still served the session, so a
concurrent request would observe expected_owner is None and bypass the
owner-binding check. Addresses Greptile review on PR #26857.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* test(mcp): fully reset stateful session tracking in auth-context refresh test

Use _remove_stateful_session_tracking in teardown so the test no longer
leaks _stateful_session_auth_context_last_seen and _stateful_session_locks
between tests, matching the cleanup used by the sibling stateful tests.

* fix(mcp): cap concurrent stateful sessions per caller to bound memory

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
Co-authored-by: Sameerlite <sameerlite@users.noreply.github.com>
Co-authored-by: yuneng-jiang <yuneng@berri.ai>
Co-authored-by: Yassin Kortam <yassin@berri.ai>
Co-authored-by: Claude Babysitter <claude@anthropic.com>
Co-authored-by: mateo-berri <mateo@berri.ai>