Files
litellm/litellm-proxy-extras/litellm_proxy_extras
ishaan-berri 693ad49719 Litellm ishaan march23 - MCP Toolsets + GCP Caching fix (#25146) (#25155)
* Litellm ishaan march23 - MCP Toolsets + GCP Caching fix  (#25146)

* feat(mcp): MCP Toolsets — curated tool subsets from one or more MCP servers (#24335)

* feat(mcp): add LiteLLM_MCPToolsetTable and mcp_toolsets to ObjectPermissionTable

* feat(mcp): add prisma migration for MCPToolset table

* feat(mcp): add MCPToolset Python types

* feat(mcp): add toolset_db.py with CRUD helpers for MCPToolset

* feat(mcp): add toolset CRUD endpoints to mcp_management_endpoints

* fix(mcp): skip allow_all_keys servers when explicit mcp_servers permission is set (toolset scope fix)

* feat(mcp): add _apply_toolset_scope and toolset route handling in server.py

* fix(mcp): resolve toolset names in responses API before fetching tools

* feat(mcp): add mcp_toolsets field to LiteLLM_ObjectPermissionTable type

* feat(mcp): register LiteLLM_MCPToolsetTable in prisma client initialization

* feat(mcp): validate mcp_toolsets in key-vs-team permission check

* feat(mcp): register toolset routes in proxy_server.py

* feat(mcp): add MCPToolset and MCPToolsetTool TypeScript types

* feat(mcp): add fetchMCPToolsets, createMCPToolset, updateMCPToolset, deleteMCPToolset API functions

* feat(mcp): add useMCPToolsets React Query hook

* feat(mcp): add toolsets (purple) as third option type in MCPServerSelector

* feat(mcp): extract toolsets from combined MCP field in key form

* feat(mcp): extract toolsets from combined MCP field in team form

* feat(mcp): show toolsets section in MCPServerPermissions read view

* feat(mcp): pass mcp_toolsets through object_permissions_view

* feat(mcp): add MCPToolsetsTab component for creating and managing toolsets

* feat(mcp): add Toolsets tab to mcp_servers.tsx

* feat(mcp): pass mcpToolsets to playground chat and responses API calls

* feat(mcp): generate correct server_url for toolsets in playground API calls

* docs(mcp): add MCP Toolsets documentation

* docs(mcp): add mcp_toolsets to sidebar

* fix(mcp): replace x-mcp-toolset-id header with ContextVar to prevent client forgery

* fix(mcp): use ContextVar + StreamingResponse for toolset MCP routes (fixes SSE streaming)

* fix(mcp): cache toolset permission lookups to avoid per-request DB calls

* test(mcp): add tests for toolset scope enforcement, ContextVar isolation, and access control

* fix(mcp): cache toolset name lookups in MCPServerManager to avoid per-request DB calls

* fix(mcp): prevent body_iter deadlock + use cached toolset lookup in responses API

- _stream_mcp_asgi_response: add done callback to handler_task that puts
  the EOF sentinel on body_queue when the task exits, preventing body_iter
  from hanging forever if the handler raises after headers are sent.
- litellm_proxy_mcp_handler: replace raw get_mcp_toolset_by_name() DB call
  with global_mcp_server_manager.get_toolset_by_name_cached() so toolset
  resolution uses the 60s TTL cache added for this purpose instead of
  hitting the DB on every responses-API request.

* fix(mcp): toolset access control, asyncio fix, and real unit tests

- server.py: _apply_toolset_scope now enforces that non-admin keys must
  have the requested toolset_id in their mcp_toolsets grant list;
  admin keys always bypass the check.
- mcp_management_endpoints.py: three access-control fixes:
  * fetch_mcp_toolsets: non-admin keys with mcp_toolsets=None now
    return [] instead of all toolsets (only admins get 'all' when
    the field is absent)
  * fetch_mcp_toolset: non-admin keys that haven't been granted the
    requested toolset_id now get 403 instead of the full result
  * add_mcp_toolset: duplicate toolset_name now returns 409 Conflict
    instead of an opaque 500
- proxy_server.py: use asyncio.get_running_loop() instead of
  get_event_loop() inside an already-running coroutine (Python 3.10+).
- test_mcp_toolset_scope.py: replace four hollow tests that only
  asserted local variable properties with real tests that call the
  production fetch_mcp_toolsets() and handle_streamable_http_mcp()
  functions with mocked dependencies.

* fix(mcp): add mcp_toolsets to ObjectPermissionBase, fix multi-toolset overwrite, fix delete 404, allow standalone key toolsets

* fix(mcp): add auth check on toolset resolution in responses API; union mcp_servers in _merge_toolset_permissions

* fix(mcp): handle RecordNotFoundError in update_mcp_toolset; union direct servers with toolset servers

* fix(mcp): use _user_has_admin_view; deny None mcp_toolsets for non-admin; use direct RecordNotFoundError import; fix docstring

* fix(mcp): add @default(now()) to MCPToolsetTable.updated_at; fix test for non-admin toolset access

* fix: use UniqueViolationError import; guard _ensure_eof for error/cancel only

* fix(mcp): preserve mcp_access_groups in toolset scope, use shared Redis cache for toolset perms

- Remove mcp_access_groups=[] from _apply_toolset_scope (server.py) and the
  responses API toolset path (litellm_proxy_mcp_handler.py). A key's access-group
  grants remain valid even when the request is scoped to a single toolset; clearing
  them silently revoked legitimate entitlements.

- Switch resolve_toolset_tool_permissions and get_toolset_by_name_cached to use
  user_api_key_cache (Redis-backed DualCache in production) instead of per-instance
  in-memory dicts. Cache entries are now shared across workers, eliminating the
  per-worker stale-toolset-permission window flagged as a P1 by Greptile.

- Use union merge (set union of tool names per server) when applying toolset
  permissions in the responses API path so direct-server tool restrictions are not
  overwritten by toolset permissions.

* fix(mcp): return 404 when edit_mcp_toolset target does not exist

* fix(mcp): align mcp_toolsets default to None in LiteLLM_ObjectPermissionTable

* fix(mcp): admin toolset visibility, in-place tool name mutation, test helper coercion

* fix(mcp): treat None/[] team mcp_toolsets as no restriction in key validation

* fix(mcp): allow_all_keys backward compat, blocked_tools API write-path, efficient startup query

* fix(mcp): use _mcp_active_toolset_id ContextVar to detect toolset scope, avoiding DB-default false-positive

* fix(mcp): remove dead toolset cache stubs, log invalidation failures, align schema updated_at defaults

* fix(mcp): deserialise MCPToolset from Redis cache hit, replace fastapi import in test

* fix(mcp): evict name-cache on toolset mutation, 409 on rename conflict, warning-level list errors

* fix(redis): regenerate GCP IAM token per connection for async cluster (#24426)

* fix(redis): regenerate GCP IAM token per connection for async cluster clients

Async RedisCluster was generating the IAM token once at startup and
storing it as a static password. After the 1-hour GCP token TTL, any
new connection (including to newly-discovered cluster nodes) would fail
to authenticate.

Fix: introduce GCPIAMCredentialProvider that implements redis-py's
CredentialProvider protocol. It calls _generate_gcp_iam_access_token()
on every new connection, matching what the sync redis_connect_func
already does. async_redis.RedisCluster accepts a credential_provider
kwarg which is invoked per-connection.

* refactor(redis): move GCPIAMCredentialProvider to its own file

Extract GCPIAMCredentialProvider and _generate_gcp_iam_access_token
into litellm/_redis_credential_provider.py. _redis.py imports them
from there, keeping the public API unchanged.

* fix: address Greptile review issues

- GCPIAMCredentialProvider now inherits from redis.credentials.CredentialProvider
  so redis-py's async path calls get_credentials_async() properly
- move _redis_credential_provider import to top of _redis.py (PEP 8)
- remove dead else-branch that silently no-oped (gcp_service_account from
  redis_kwargs.get() was always None since it's popped by _get_redis_client_logic)
- remove mid-function 'from litellm import get_secret_str' inline import
- remove unused 'call' import from test_redis.py

* chore: retrigger CI/review

* chore: sync schema.prisma copies from root

* chore: sync schema.prisma copies from root

* fix(proxy_server): use bounded asyncio.Queue with maxsize to prevent unbounded growth

* fix(a2a/pydantic_ai): make api_base Optional to match base class signature

* fix(a2a/pydantic_ai): make api_base Optional in handler and guard against None

* fix(mcp): remove unused get_all_mcp_servers import

* fix(mcp): remove unused MCPToolset import

* refactor(mcp): extract toolset permission logic to reduce statement count below PLR0915 limit

* fix(tests): update reload_servers_from_database tests to mock prisma directly

---------

Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fix(toolset_db): lazy-import prisma to avoid ImportError when prisma not installed

* fix(tests): update UI tests for toolset tab and updated empty state text

* fix(tests): add get_mcp_server_by_name to fake_manager stub

---------

Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
2026-04-04 16:23:21 -07:00
..