From 693ad49719f147e103479d4024063d39374eab3d Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Sat, 4 Apr 2026 16:23:21 -0700 Subject: [PATCH] Litellm ishaan march23 - MCP Toolsets + GCP Caching fix (#25146) (#25155) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 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 Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/my-website/docs/mcp_toolsets.md | 231 ++++++++ docs/my-website/sidebars.js | 1 + .../migration.sql | 19 + .../litellm_proxy_extras/schema.prisma | 13 + litellm/_redis.py | 71 +-- litellm/_redis_credential_provider.py | 53 ++ .../providers/pydantic_ai_agents/handler.py | 10 +- .../_experimental/mcp_server/mcp_context.py | 16 + .../mcp_server/mcp_server_manager.py | 184 +++++- .../proxy/_experimental/mcp_server/server.py | 154 ++++- .../_experimental/mcp_server/toolset_db.py | 117 ++++ litellm/proxy/_types.py | 100 ++-- litellm/proxy/db/prisma_client.py | 8 +- .../mcp_management_endpoints.py | 210 ++++++- .../object_permission_utils.py | 46 +- litellm/proxy/proxy_server.py | 170 +++++- litellm/proxy/schema.prisma | 13 + .../mcp/litellm_proxy_mcp_handler.py | 129 ++++- litellm/types/mcp_server/mcp_toolset.py | 34 ++ schema.prisma | 13 + .../mcp_server/test_mcp_server.py | 20 +- .../mcp_server/test_mcp_toolset_scope.py | 288 ++++++++++ .../mcp/test_litellm_proxy_mcp_handler.py | 1 + tests/test_litellm/test_redis.py | 199 +++++-- .../hooks/mcpServers/useMCPToolsets.ts | 16 + .../MCPServerSelector.tsx | 58 +- .../components/mcp_tools/MCPToolsetsTab.tsx | 524 ++++++++++++++++++ .../components/mcp_tools/mcp_servers.test.tsx | 6 +- .../src/components/mcp_tools/mcp_servers.tsx | 5 + .../src/components/mcp_tools/types.tsx | 14 + .../src/components/networking.tsx | 93 ++++ .../components/object_permissions_view.tsx | 11 +- .../permissions/MCPServerPermissions.test.tsx | 2 +- .../permissions/MCPServerPermissions.tsx | 128 ++++- .../components/playground/chat_ui/ChatUI.tsx | 270 +++++++-- .../playground/llm_calls/chat_completion.tsx | 37 +- .../playground/llm_calls/responses_api.tsx | 42 +- .../src/components/team/TeamInfo.tsx | 8 +- .../components/templates/key_info_view.tsx | 3 +- 39 files changed, 2962 insertions(+), 355 deletions(-) create mode 100644 docs/my-website/docs/mcp_toolsets.md create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260321000000_add_mcp_toolsets/migration.sql create mode 100644 litellm/_redis_credential_provider.py create mode 100644 litellm/proxy/_experimental/mcp_server/mcp_context.py create mode 100644 litellm/proxy/_experimental/mcp_server/toolset_db.py create mode 100644 litellm/types/mcp_server/mcp_toolset.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPToolsets.ts create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/MCPToolsetsTab.tsx diff --git a/docs/my-website/docs/mcp_toolsets.md b/docs/my-website/docs/mcp_toolsets.md new file mode 100644 index 00000000000..5f27cdcc0fc --- /dev/null +++ b/docs/my-website/docs/mcp_toolsets.md @@ -0,0 +1,231 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# MCP Toolsets + +A **Toolset** is a named collection of specific tools drawn from one or more MCP servers. Instead of giving an agent access to every tool on every server, you pick exactly which tools it needs — from whichever servers they live on — and bundle them under a single name. + +## How it works + +``` + ┌─────────────────────────────────┐ + │ MCP Toolset │ + │ "devtooling-prod" │ + └────────────┬────────────────────┘ + │ + ┌──────────────────┴──────────────────┐ + │ │ + ┌────────▼────────┐ ┌────────▼────────┐ + │ CircleCI MCP │ │ DeepWiki MCP │ + │ (10+ tools) │ │ (3 tools) │ + └────────┬────────┘ └────────┬────────┘ + │ │ + ┌─────────┴──────────┐ ┌──────────┴──────────┐ + │ ✓ get_build_logs │ │ ✓ read_wiki_structure│ + │ ✓ find_flaky_tests │ │ ✓ read_wiki_contents │ + │ ✓ get_pipeline_ │ │ ✗ ask_question │ + │ status │ └─────────────────────┘ + │ ✓ run_pipeline │ + │ ✗ list_followed_ │ + │ projects │ + └────────────────────┘ + + Agent sees exactly 6 tools, nothing more. +``` + +Instead of 13+ tools across two servers, the agent gets 6 — the ones it actually needs. + +**Why this matters:** +- Smaller tool lists → fewer tokens, faster responses, less hallucination +- Combine tools from GitHub + Linear + CircleCI into one named grant +- Assign to keys and teams the same way you assign MCP servers today + +--- + +## Create a toolset + +### 1. Go to the MCP page + +Navigate to **MCP** in the left sidebar. + +![Navigate to MCP](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/1a96c713-6a37-4f96-92f1-07bd58c1973c/ascreenshot_23515f386ccc4597b0633987667fe01f_text_export.jpeg) + +### 2. Open the Toolsets tab + +Click the **Toolsets** tab on the MCP page. + +![Click Toolsets tab](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/65b6986b-595a-4b28-8fdc-a7b36bc76e59/ascreenshot_ca70c18fe7ec415486f96a6b405bf550_text_export.jpeg) + +### 3. Click "New Toolset" + +![New Toolset button](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/798c55c4-5d6b-4815-a642-70ac9f34f102/ascreenshot_3f144f54a1a944e28454239c837b4e6d_text_export.jpeg) + +### 4. Enter a name + +Type a name for the toolset. Pick something descriptive — this is what agents will reference. + +![Enter toolset name](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/62b412e0-d38f-44c3-99e4-3693f1512f6a/ascreenshot_b678c7c988a04f8b887b0f54c4dd95a7_text_export.jpeg) + +![Toolset name field](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/ba5ebc95-cab7-470b-a7c9-21f12b9b01a3/ascreenshot_a602e982a2a44890a83dca64d61c38eb_text_export.jpeg) + +### 5. Add the first tool + +Select an MCP server from the dropdown, then choose the tool you want to include from that server. + +![Select MCP server](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/2aa5bcba-6414-42e3-9813-efb0a9078e32/ascreenshot_58fbff35ba654210a1b4dc5452aa6bd9_text_export.jpeg) + +![Choose server from dropdown](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/4fd9cffb-d3ba-461a-8679-89f278bf67ad/ascreenshot_b61e9e85a51b494a8d09fe61198d63e1_text_export.jpeg) + +![Select tool from server](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/60718e72-2062-494b-9a23-456992c88cbd/ascreenshot_7a1f8eeab30a4a05ba39c450e5458b78_text_export.jpeg) + +### 6. Add tools from a second server + +Click **Add Tool**, pick a different MCP server, and select another tool. Repeat for as many tools as you need — they can come from any number of servers. + +![Add tool from second server](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/f34e0600-cc74-4b18-8794-88d45f326144/ascreenshot_98834b14ab9343e39fb503e458d72b7c_text_export.jpeg) + +![Select second server](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/75150368-2202-4da1-99f1-6f0620e9b133/ascreenshot_f94d0bc08ea147348a9cf021cce7d854_text_export.jpeg) + +![Select tool from second server](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/ed2cdf6e-025d-4d50-8b12-ed68745d5c51/ascreenshot_0c1c7f76524b46c5a056fda5e6956e2b_text_export.jpeg) + +### 7. Create the toolset + +Click **Create Toolset** to save. + +![Create Toolset](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/021ca7b3-2d9a-49a0-8758-dae3dc3bcb4d/ascreenshot_14c6434e71114a6091e359a996f20e12_text_export.jpeg) + +--- + +## Use a toolset in the Playground + +Once created, your toolset appears alongside MCP servers in the **MCP Servers** dropdown in the Playground — it's selectable the same way. + +### 1. Go to the Playground + +![Navigate to Playground](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/f9d4aa4c-d98e-4767-b98e-aad2890e97ca/ascreenshot_d84239c441bb4e828f229d0c9e079e3f_text_export.jpeg) + +![Click Playground](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/d8a07563-97fe-453a-b974-88da46c87294/ascreenshot_ea494300a536400abb2ea6bf3bdfd5ab_text_export.jpeg) + +### 2. Select your toolset from MCP Servers + +In the left panel under **MCP Servers**, open the dropdown and pick your toolset. The model will only see the tools you included in it. + +![Select MCP servers dropdown](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/ee8cb38c-c4ff-4b4b-844c-22f2e40832ae/ascreenshot_e300fb39cea0434fb5e3986e912a2b8d_text_export.jpeg) + +![Open MCP server picker](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/8672070c-5d07-4f63-878c-6fc7dcbc9b65/ascreenshot_326ddd0868224c99a6fa5dab2d144f1f_text_export.jpeg) + +![Select toolset](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/955826ad-2bbb-403e-ab26-c1ac03ec2675/ascreenshot_13f837ad53574535986ca7ca5998d34a_text_export.jpeg) + +![Toolset selected and active](https://colony-recorder.s3.amazonaws.com/files/2026-03-22/9a59c3b9-1563-4731-838f-1c35d636ddc9/ascreenshot_c05d8fa5f37a4b3093fc46e26f293b4d_text_export.jpeg) + +The model now has access to exactly the tools in your toolset and nothing else. + +--- + +## Use a toolset via API + +Pass the toolset's route as the `server_url` in your tools list. LiteLLM resolves it server-side — no public URL needed. + + + + +```python +import openai + +client = openai.OpenAI( + api_key="your-litellm-key", + base_url="http://your-proxy/v1", +) + +response = client.responses.create( + model="gpt-4o", + input="What CI/CD tools do you have?", + tools=[ + { + "type": "mcp", + "server_label": "devtooling-prod", + "server_url": "litellm_proxy/mcp/devtooling-prod", + "require_approval": "never", + } + ], +) +print(response.output_text) +``` + + + + +```python +import openai + +client = openai.OpenAI( + api_key="your-litellm-key", + base_url="http://your-proxy/v1", +) + +response = client.chat.completions.create( + model="gpt-4o", + messages=[{"role": "user", "content": "What CI/CD tools do you have?"}], + tools=[ + { + "type": "mcp", + "server_label": "devtooling-prod", + "server_url": "litellm_proxy/mcp/devtooling-prod", + "require_approval": "never", + } + ], +) +print(response.choices[0].message.content) +``` + + + + +```bash +curl http://your-proxy/v1/responses \ + -H "Authorization: Bearer your-litellm-key" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4o", + "input": "What CI/CD tools do you have?", + "tools": [ + { + "type": "mcp", + "server_label": "devtooling-prod", + "server_url": "litellm_proxy/mcp/devtooling-prod", + "require_approval": "never" + } + ] + }' +``` + + + + +--- + +## Manage toolsets via API + +```bash +# List all toolsets +curl http://your-proxy/v1/mcp/toolset \ + -H "Authorization: Bearer your-litellm-key" + +# Create a toolset +curl -X POST http://your-proxy/v1/mcp/toolset \ + -H "Authorization: Bearer your-litellm-key" \ + -H "Content-Type: application/json" \ + -d '{ + "toolset_name": "devtooling-prod", + "description": "CircleCI + DeepWiki tools for the dev team", + "tools": [ + {"server_id": "", "tool_name": "get_build_failure_logs"}, + {"server_id": "", "tool_name": "run_pipeline"}, + {"server_id": "", "tool_name": "read_wiki_structure"} + ] + }' + +# Delete a toolset +curl -X DELETE http://your-proxy/v1/mcp/toolset/ \ + -H "Authorization: Bearer your-litellm-key" +``` diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index ab4ae46ec3b..ab8f257c7d5 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -325,6 +325,7 @@ const sidebars = { "mcp_control", "mcp_cost", "mcp_guardrail", + "mcp_toolsets", { type: "link", label: "MCP Troubleshooting Guide", diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260321000000_add_mcp_toolsets/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260321000000_add_mcp_toolsets/migration.sql new file mode 100644 index 00000000000..eb9fd29499f --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260321000000_add_mcp_toolsets/migration.sql @@ -0,0 +1,19 @@ +-- CreateTable: LiteLLM_MCPToolsetTable +CREATE TABLE IF NOT EXISTS "LiteLLM_MCPToolsetTable" ( + "toolset_id" TEXT NOT NULL, + "toolset_name" TEXT NOT NULL, + "description" TEXT, + "tools" JSONB NOT NULL DEFAULT '[]', + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT, + + CONSTRAINT "LiteLLM_MCPToolsetTable_pkey" PRIMARY KEY ("toolset_id") +); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MCPToolsetTable_toolset_name_key" ON "LiteLLM_MCPToolsetTable"("toolset_name"); + +-- AlterTable: add mcp_toolsets to ObjectPermissionTable +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "mcp_toolsets" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 5d804d75b6a..fce95465b55 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -273,6 +273,7 @@ model LiteLLM_ObjectPermissionTable { agent_access_groups String[] @default([]) models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission + mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -331,6 +332,18 @@ model LiteLLM_MCPServerTable { @@index([approval_status]) } +// Named collection of {server_id, tool_name} pairs that can be granted to keys/teams +model LiteLLM_MCPToolsetTable { + toolset_id String @id @default(uuid()) + toolset_name String @unique + description String? + tools Json @default("[]") // [{server_id: string, tool_name: string}] + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} + // Per-user BYOK credentials for MCP servers model LiteLLM_MCPUserCredentials { id String @id @default(uuid()) diff --git a/litellm/_redis.py b/litellm/_redis.py index 2bf32d71b21..f12afbac297 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -18,6 +18,10 @@ import redis # type: ignore import redis.asyncio as async_redis # type: ignore from litellm import get_secret, get_secret_str +from litellm._redis_credential_provider import ( + GCPIAMCredentialProvider, + _generate_gcp_iam_access_token, +) from litellm.constants import REDIS_CONNECTION_POOL_TIMEOUT, REDIS_SOCKET_TIMEOUT from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker @@ -107,33 +111,6 @@ def _redis_kwargs_from_environment(): return return_dict -def _generate_gcp_iam_access_token(service_account: str) -> str: - """ - Generate GCP IAM access token for Redis authentication. - - Args: - service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com' - - Returns: - Access token string for GCP IAM authentication - """ - try: - from google.cloud import iam_credentials_v1 - except ImportError: - raise ImportError( - "google-cloud-iam is required for GCP IAM Redis authentication. " - "Install it with: pip install google-cloud-iam" - ) - - client = iam_credentials_v1.IAMCredentialsClient() - request = iam_credentials_v1.GenerateAccessTokenRequest( - name=service_account, - scope=["https://www.googleapis.com/auth/cloud-platform"], - ) - response = client.generate_access_token(request=request) - return str(response.access_token) - - def create_gcp_iam_redis_connect_func( service_account: str, ssl_ca_certs: Optional[str] = None, @@ -266,7 +243,7 @@ def _get_redis_client_logic(**env_overrides): service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs ) # Store GCP service account in redis_connect_func for async cluster access - redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account + redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account # type: ignore[attr-defined] # Remove GCP-specific kwargs that shouldn't be passed to Redis client redis_kwargs.pop("gcp_service_account", None) @@ -413,41 +390,13 @@ def get_redis_async_client( # Handle GCP IAM authentication for async clusters redis_connect_func = cluster_kwargs.pop("redis_connect_func", None) - from litellm import get_secret_str - # Get GCP service account - first try from redis_connect_func, then from environment - gcp_service_account = None + # Use a CredentialProvider so the IAM token is regenerated on every new + # connection — mirrors the sync path where redis_connect_func is invoked + # per connection. Without this, the token would expire after ~1 hour. if redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"): - gcp_service_account = redis_connect_func._gcp_service_account - else: - gcp_service_account = redis_kwargs.get( - "gcp_service_account" - ) or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT") - - verbose_logger.debug( - f"DEBUG: Redis cluster kwargs: redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}" - ) - - # If GCP IAM is configured (indicated by redis_connect_func), generate access token and use as password - if redis_connect_func and gcp_service_account: - verbose_logger.debug( - "DEBUG: Generating IAM token for service account (value not logged for security reasons)" - ) - try: - # Generate IAM access token using the helper function - access_token = _generate_gcp_iam_access_token(gcp_service_account) - cluster_kwargs["password"] = access_token - verbose_logger.debug( - "DEBUG: Successfully generated GCP IAM access token for async Redis cluster" - ) - except Exception as e: - verbose_logger.error(f"Failed to generate GCP IAM access token: {e}") - from redis.exceptions import AuthenticationError - - raise AuthenticationError("Failed to generate GCP IAM access token") - else: - verbose_logger.debug( - f"DEBUG: Not using GCP IAM auth - redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}" + cluster_kwargs["credential_provider"] = GCPIAMCredentialProvider( + redis_connect_func._gcp_service_account ) new_startup_nodes: List[ClusterNode] = [] diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py new file mode 100644 index 00000000000..495d2a879bd --- /dev/null +++ b/litellm/_redis_credential_provider.py @@ -0,0 +1,53 @@ +import asyncio +from typing import Tuple + +from redis.credentials import CredentialProvider # type: ignore[attr-defined] + + +def _generate_gcp_iam_access_token(service_account: str) -> str: + """ + Generate GCP IAM access token for Redis authentication. + + Args: + service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com' + + Returns: + Access token string for GCP IAM authentication + """ + try: + from google.cloud import iam_credentials_v1 + except ImportError: + raise ImportError( + "google-cloud-iam is required for GCP IAM Redis authentication. " + "Install it with: pip install google-cloud-iam" + ) + + client = iam_credentials_v1.IAMCredentialsClient() + request = iam_credentials_v1.GenerateAccessTokenRequest( + name=service_account, + scope=["https://www.googleapis.com/auth/cloud-platform"], + ) + response = client.generate_access_token(request=request) + return str(response.access_token) + + +class GCPIAMCredentialProvider(CredentialProvider): + """ + redis.credentials.CredentialProvider implementation that generates a fresh GCP IAM + token on every new connection. This fixes the 1-hour token expiry issue for async + Redis cluster clients, which previously generated the token once at startup and + cached it as a static password. + """ + + def __init__(self, gcp_service_account: str) -> None: + self._gcp_service_account = gcp_service_account + + def get_credentials(self) -> Tuple[str]: + token = _generate_gcp_iam_access_token(self._gcp_service_account) + return (token,) + + async def get_credentials_async(self) -> Tuple[str]: + token = await asyncio.to_thread( + _generate_gcp_iam_access_token, self._gcp_service_account + ) + return (token,) diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py index 7d4167752f8..5b8d6b94ff2 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py @@ -5,7 +5,7 @@ Pydantic AI agents follow A2A protocol but don't support streaming natively. This handler provides fake streaming by converting non-streaming responses into streaming chunks. """ -from typing import Any, AsyncIterator, Dict +from typing import Any, AsyncIterator, Dict, Optional from litellm._logging import verbose_logger from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( @@ -26,7 +26,7 @@ class PydanticAIHandler: async def handle_non_streaming( request_id: str, params: Dict[str, Any], - api_base: str, + api_base: Optional[str] = None, timeout: float = 60.0, ) -> Dict[str, Any]: """ @@ -41,6 +41,8 @@ class PydanticAIHandler: Returns: A2A SendMessageResponse dict """ + if api_base is None: + raise ValueError("api_base is required for Pydantic AI agents") verbose_logger.info(f"Pydantic AI: Routing to Pydantic AI agent at {api_base}") # Send request directly to Pydantic AI agent @@ -57,7 +59,7 @@ class PydanticAIHandler: async def handle_streaming( request_id: str, params: Dict[str, Any], - api_base: str, + api_base: Optional[str] = None, timeout: float = 60.0, chunk_size: int = 50, delay_ms: int = 10, @@ -80,6 +82,8 @@ class PydanticAIHandler: Yields: A2A streaming response events """ + if api_base is None: + raise ValueError("api_base is required for Pydantic AI agents") verbose_logger.info( f"Pydantic AI: Faking streaming for Pydantic AI agent at {api_base}" ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_context.py b/litellm/proxy/_experimental/mcp_server/mcp_context.py new file mode 100644 index 00000000000..12830db1d6a --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/mcp_context.py @@ -0,0 +1,16 @@ +""" +Shared ContextVars for the MCP server layer. + +Lives in its own module to avoid circular imports between +mcp_server_manager.py and server.py. +""" + +from contextvars import ContextVar +from typing import Optional + +# Set server-side in proxy_server.py route handlers when a request arrives via +# /toolset/{name}/mcp or the toolset fallback in dynamic_mcp_route. +# Never populated from client-supplied headers. +_mcp_active_toolset_id: ContextVar[Optional[str]] = ContextVar( + "_mcp_active_toolset_id", default=None +) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index cfbc2c437b0..7b87e7e7e61 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -503,12 +503,12 @@ class MCPServerManager: ) # Update tool name to server name mapping (for both prefixed and base names) - self.tool_name_to_mcp_server_name_mapping[ - base_tool_name - ] = server_prefix - self.tool_name_to_mcp_server_name_mapping[ - prefixed_tool_name - ] = server_prefix + self.tool_name_to_mcp_server_name_mapping[base_tool_name] = ( + server_prefix + ) + self.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = ( + server_prefix + ) registered_count += 1 verbose_logger.debug( @@ -790,7 +790,18 @@ class MCPServerManager: f"Allowed MCP Servers for user api key auth: {allowed_mcp_servers}" ) combined_servers = set(allowed_mcp_servers) - combined_servers.update(allow_all_server_ids) + # Only skip allow_all_keys servers when the request is inside a toolset + # scope. toolset_mcp_route / dynamic_mcp_route set _mcp_active_toolset_id + # before calling the handler — that ContextVar is the reliable signal. + # Using op.mcp_toolsets==[] would false-positive on DB-default rows where + # Postgres initialises the column to ARRAY[]::TEXT[]. + from litellm.proxy._experimental.mcp_server.mcp_context import ( # noqa: PLC0415 + _mcp_active_toolset_id, + ) + + in_toolset_scope = _mcp_active_toolset_id.get() is not None + if not in_toolset_scope: + combined_servers.update(allow_all_server_ids) if len(combined_servers) == 0: verbose_logger.debug( @@ -801,6 +812,132 @@ class MCPServerManager: verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}.") return allow_all_server_ids + async def resolve_toolset_tool_permissions( + self, + toolset_ids: List[str], + ) -> Dict[str, List[str]]: + """ + Resolve a list of toolset IDs into a mcp_tool_permissions dict. + + Returns: {server_id: [tool_name, ...]} — the union of all tools across + the given toolsets. Results are cached via ``user_api_key_cache`` (a + Redis-backed ``DualCache`` in production) so that cache entries are + shared across workers and cold-cache DB hits are minimised. + """ + from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL + from litellm.proxy._experimental.mcp_server.toolset_db import list_mcp_toolsets + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if not toolset_ids or prisma_client is None: + return {} + + cache_key = "toolset_perms:" + ",".join(sorted(toolset_ids)) + cached = await user_api_key_cache.async_get_cache(key=cache_key) + if cached is not None: + return cached + + try: + toolsets = await list_mcp_toolsets(prisma_client, toolset_ids=toolset_ids) + tool_permissions: Dict[str, List[str]] = {} + for toolset in toolsets: + for tool in toolset.tools: + raw_name = tool["tool_name"] + unprefixed, _ = split_server_prefix_from_name(raw_name) + tool_permissions.setdefault(tool["server_id"], []) + if unprefixed not in tool_permissions[tool["server_id"]]: + tool_permissions[tool["server_id"]].append(unprefixed) + await user_api_key_cache.async_set_cache( + key=cache_key, + value=tool_permissions, + ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + ) + return tool_permissions + except Exception as e: + verbose_logger.warning(f"Failed to resolve toolset permissions: {str(e)}") + return {} + + def invalidate_toolset_cache(self, toolset_id: Optional[str] = None) -> None: + """Evict cached toolset permission entries. + + Called after create/update/delete of a toolset so stale data is not served. + The in-memory layer of ``user_api_key_cache`` is cleared immediately; + Redis entries expire naturally after the configured TTL. + Pass toolset_id to evict only entries containing that ID, or None to clear all. + """ + # Clear the in-memory layer of the shared DualCache for affected keys. + # We can't enumerate Redis keys by pattern, so Redis entries expire via TTL. + try: + from litellm.proxy.proxy_server import user_api_key_cache + + in_mem = getattr(user_api_key_cache, "in_memory_cache", None) + if in_mem is None: + return + cache_dict = getattr(in_mem, "cache_dict", {}) + if toolset_id is None: + keys_to_remove = [k for k in cache_dict if k.startswith("toolset_")] + else: + # Evict permission-cache entries that reference this toolset ID. + # Also evict ALL name-cache entries (toolset_name:*): we can't map + # toolset_id → toolset_name without a DB call, and the name may have + # changed in an update anyway. + keys_to_remove = [ + k + for k in cache_dict + if (k.startswith("toolset_perms:") and toolset_id in k) + or k.startswith("toolset_name:") + ] + for k in keys_to_remove: + cache_dict.pop(k, None) + except Exception as e: + verbose_logger.warning( + f"invalidate_toolset_cache: failed to evict in-memory entries: {e}" + ) + + async def get_toolset_by_name_cached( + self, + prisma_client: Any, + toolset_name: str, + ) -> Optional[Any]: + """Return a toolset by name, cached in ``user_api_key_cache`` (Redis-backed + ``DualCache`` in production) to avoid a DB hit on every routed request. + + Serialisation note: the cache value is stored as a plain JSON-safe dict via + ``model_dump(mode="json")`` so that Redis round-trips correctly in multi-worker + deployments. On a cache hit we reconstruct the ``MCPToolset`` Pydantic object + so callers can always use attribute access (e.g. ``toolset.toolset_id``). + """ + from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL + from litellm.proxy.proxy_server import user_api_key_cache + from litellm.types.mcp_server.mcp_toolset import MCPToolset + + cache_key = f"toolset_name:{toolset_name}" + cached = await user_api_key_cache.async_get_cache(key=cache_key) + if cached is not None: + # Sentinel value used to cache "not found" so we don't re-query for + # names that don't exist. + if cached == "__not_found__": + return None + # Redis deserialises JSON back as a plain dict — reconstruct the model. + if isinstance(cached, dict): + return MCPToolset(**cached) + return cached + + from litellm.proxy._experimental.mcp_server.toolset_db import ( + get_mcp_toolset_by_name, + ) + + toolset = await get_mcp_toolset_by_name(prisma_client, toolset_name) + await user_api_key_cache.async_set_cache( + key=cache_key, + value=( + toolset.model_dump(mode="json") + if toolset is not None + else "__not_found__" + ), + ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + ) + return toolset + def filter_server_ids_by_ip( self, server_ids: List[str], client_ip: Optional[str] ) -> List[str]: @@ -1077,6 +1214,24 @@ class MCPServerManager: tools = global_mcp_tool_registry.convert_tools_to_mcp_sdk_tool_type( _tools ) + # OpenAPI tools are stored in the registry with their prefix already + # applied (e.g. "test_petstore-getinventory"). Do NOT pass them + # through _create_prefixed_tools — that would add the prefix a second + # time producing "test_petstore-test_petstore-getinventory". + if not add_prefix: + prefix = get_server_prefix(server) + sep = MCP_TOOL_PREFIX_SEPARATOR + tools = [ + ( + t.model_copy( + update={"name": t.name[len(prefix) + len(sep) :]} + ) + if t.name.startswith(f"{prefix}{sep}") + else t + ) + for t in tools + ] + return tools else: tools = await self._fetch_tools_with_timeout(client, server.name) @@ -2418,7 +2573,6 @@ class MCPServerManager: async def reload_servers_from_database(self): """Re-synchronize the in-memory MCP server registry with the database.""" - from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers from litellm.proxy.management_endpoints.mcp_management_endpoints import ( get_prisma_client_or_throw, ) @@ -2429,9 +2583,19 @@ class MCPServerManager: prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to your proxy" ) - db_mcp_servers = await get_all_mcp_servers( - prisma_client, approval_status="active" + # Load only "active", legacy "approved", and NULL (no approval workflow) rows. + # Pending/rejected servers are excluded at the DB level so we never load them. + from litellm.proxy._experimental.mcp_server.db import LiteLLM_MCPServerTable + + raw_rows = await prisma_client.db.litellm_mcpservertable.find_many( + where={ + "OR": [ + {"approval_status": None}, + {"approval_status": {"in": ["active", "approved"]}}, + ] + } ) + db_mcp_servers = [LiteLLM_MCPServerTable(**r.model_dump()) for r in raw_rows] verbose_logger.info(f"Found {len(db_mcp_servers)} MCP servers in database") previous_registry = self.registry diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index cd06de2a2df..166be712d5a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1,6 +1,7 @@ """ LiteLLM MCP Server Routes """ + # pyright: reportInvalidTypeForm=false, reportArgumentType=false, reportOptionalCall=false import asyncio @@ -36,6 +37,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, ) +from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_active_toolset_id from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_DESCRIPTION, @@ -1455,6 +1457,49 @@ if MCP_AVAILABLE: return filtered_tools + async def _merge_toolset_permissions( + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> Optional[UserAPIKeyAuth]: + """ + Resolve mcp_toolsets on the key's object_permission into tool-level permissions + and merge them (union) into object_permission.mcp_tool_permissions. + + Returns the (possibly mutated copy of) user_api_key_auth. + """ + if user_api_key_auth is None: + return None + op = user_api_key_auth.object_permission + if op is None: + return user_api_key_auth + toolset_ids = getattr(op, "mcp_toolsets", None) or [] + if not toolset_ids: + return user_api_key_auth + + toolset_perms = ( + await global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=toolset_ids + ) + ) + if not toolset_perms: + return user_api_key_auth + + # Merge toolset_perms into existing mcp_tool_permissions (union) + existing = dict(op.mcp_tool_permissions or {}) + for server_id, tool_names in toolset_perms.items(): + existing_tools = existing.get(server_id, []) + merged = list(set(existing_tools) | set(tool_names)) + existing[server_id] = merged + + # Build updated object_permission with merged tool permissions and server IDs. + # Union the toolset's server IDs into mcp_servers so downstream server-level + # filtering doesn't silently drop servers that the toolset references but that + # aren't already in the key's explicit mcp_servers list. + merged_servers = list(set(op.mcp_servers or []) | set(existing.keys())) + updated_op = op.model_copy( + update={"mcp_servers": merged_servers, "mcp_tool_permissions": existing} + ) + return user_api_key_auth.model_copy(update={"object_permission": updated_op}) + async def _list_mcp_tools( user_api_key_auth: Optional[UserAPIKeyAuth] = None, mcp_auth_header: Optional[str] = None, @@ -1479,6 +1524,11 @@ if MCP_AVAILABLE: """ if not MCP_AVAILABLE: return [] + + # Resolve toolset permissions and merge into the key's object_permission + # so that the existing filter_tools_by_key_team_permissions logic picks them up. + user_api_key_auth = await _merge_toolset_permissions(user_api_key_auth) + # Get tools from managed MCP servers with error handling managed_tools = [] try: @@ -1822,9 +1872,9 @@ if MCP_AVAILABLE: "litellm_logging_obj", None ) if litellm_logging_obj: - litellm_logging_obj.model_call_details[ - "mcp_tool_call_metadata" - ] = standard_logging_mcp_tool_call + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( + standard_logging_mcp_tool_call + ) litellm_logging_obj.model = f"MCP: {name}" # Resolve the MCP server early so BYOK checks and credential injection # apply to ALL dispatch paths (local tool registry AND managed MCP server). @@ -1836,9 +1886,9 @@ if MCP_AVAILABLE: mcp_server.mcp_info or {} ).get("mcp_server_cost_info") if litellm_logging_obj: - litellm_logging_obj.model_call_details[ - "mcp_tool_call_metadata" - ] = standard_logging_mcp_tool_call + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( + standard_logging_mcp_tool_call + ) # BYOK: retrieve the stored per-user credential. A single DB call # both checks existence and fetches the value, avoiding a double query. @@ -2358,6 +2408,63 @@ if MCP_AVAILABLE: ] return False + async def _apply_toolset_scope( + user_api_key_auth: UserAPIKeyAuth, + toolset_id: str, + ) -> UserAPIKeyAuth: + """ + Restrict a key's MCP permissions to a single toolset. + + When a request arrives via /toolset/{name}/mcp we override the key's + object_permission so that only the toolset's tools are visible. + + Raises HTTPException(403) if the key has an explicit toolset grant list + that does not include toolset_id (i.e. mcp_toolsets is set but empty, + or set to a list that omits this toolset). Admin keys always pass. + """ + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view + + # Access control: non-admin keys must have this toolset in their grant list. + # Use _user_has_admin_view so that PROXY_ADMIN_VIEW_ONLY is also treated as admin. + is_admin = _user_has_admin_view(user_api_key_auth) + if not is_admin: + op = user_api_key_auth.object_permission + granted = getattr(op, "mcp_toolsets", None) if op else None + # granted=None → key has no explicit toolset grants → deny (same semantics as + # fetch_mcp_toolsets which returns [] for non-admin keys with no grants configured). + # granted=[] or list without toolset_id → also deny. + if granted is None or toolset_id not in granted: + raise HTTPException( + status_code=403, + detail=f"API key does not have access to toolset '{toolset_id}'.", + ) + + tool_permissions = ( + await global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=[toolset_id] + ) + ) + server_ids = list(tool_permissions.keys()) + existing_op = user_api_key_auth.object_permission + if existing_op is not None: + updated_op = existing_op.model_copy( + update={ + "mcp_servers": server_ids, + "mcp_tool_permissions": tool_permissions, + "mcp_toolsets": [], + # mcp_access_groups is preserved: a key's access-group grants + # remain valid even when the request is scoped to a single toolset. + } + ) + else: + updated_op = LiteLLM_ObjectPermissionTable( + object_permission_id="toolset-scope", + mcp_servers=server_ids, + mcp_tool_permissions=tool_permissions, + ) + return user_api_key_auth.model_copy(update={"object_permission": updated_op}) + async def handle_streamable_http_mcp( scope: Scope, receive: Receive, send: Send ) -> None: @@ -2402,6 +2509,21 @@ if MCP_AVAILABLE: headers={"www-authenticate": authorization_uri}, ) + # Strip any client-supplied x-mcp-toolset-id to prevent forgery. + scope["headers"] = [ + (k, v) + for k, v in scope.get("headers", []) + if k.lower() != b"x-mcp-toolset-id" + ] + + # Apply toolset scope if set server-side via ContextVar (set by + # /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py). + active_toolset_id = _mcp_active_toolset_id.get() + if active_toolset_id and user_api_key_auth is not None: + user_api_key_auth = await _apply_toolset_scope( + user_api_key_auth, active_toolset_id + ) + # Inject masked debug headers when client sends x-litellm-mcp-debug: true _debug_headers = MCPDebug.maybe_build_debug_headers( raw_headers=raw_headers, @@ -2580,17 +2702,15 @@ if MCP_AVAILABLE: ) auth_context_var.set(auth_user) - def get_auth_context() -> ( - Tuple[ - Optional[UserAPIKeyAuth], - Optional[str], - Optional[List[str]], - Optional[Dict[str, Dict[str, str]]], - Optional[Dict[str, str]], - Optional[Dict[str, str]], - Optional[str], - ] - ): + def get_auth_context() -> Tuple[ + Optional[UserAPIKeyAuth], + Optional[str], + Optional[List[str]], + Optional[Dict[str, Dict[str, str]]], + Optional[Dict[str, str]], + Optional[Dict[str, str]], + Optional[str], + ]: """ Get the UserAPIKeyAuth from the auth context variable. diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py new file mode 100644 index 00000000000..08ac7dbd33b --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -0,0 +1,117 @@ +import json +from typing import List, Optional + +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.proxy.utils import PrismaClient +from litellm.types.mcp_server.mcp_toolset import ( + MCPToolset, + NewMCPToolsetRequest, + UpdateMCPToolsetRequest, +) + + +def _toolset_from_row(row) -> MCPToolset: + data = row.model_dump() + tools = data.get("tools") or [] + if isinstance(tools, str): + tools = json.loads(tools) + data["tools"] = tools + return MCPToolset(**data) + + +async def create_mcp_toolset( + prisma_client: PrismaClient, + data: NewMCPToolsetRequest, + touched_by: str, +) -> MCPToolset: + data_dict = data.model_dump(exclude_none=True) + data_dict["toolset_id"] = str(uuid.uuid4()) + data_dict["tools"] = json.dumps(data_dict.get("tools", [])) + data_dict["created_by"] = touched_by + data_dict["updated_by"] = touched_by + row = await prisma_client.db.litellm_mcptoolsettable.create(data=data_dict) + return _toolset_from_row(row) + + +async def get_mcp_toolset( + prisma_client: PrismaClient, + toolset_id: str, +) -> Optional[MCPToolset]: + row = await prisma_client.db.litellm_mcptoolsettable.find_unique( + where={"toolset_id": toolset_id} + ) + if row is None: + return None + return _toolset_from_row(row) + + +async def list_mcp_toolsets( + prisma_client: PrismaClient, + toolset_ids: Optional[List[str]] = None, +) -> List[MCPToolset]: + try: + where = {} + if toolset_ids is not None: + where = {"toolset_id": {"in": toolset_ids}} + rows = await prisma_client.db.litellm_mcptoolsettable.find_many(where=where) + return [_toolset_from_row(r) for r in rows] + except Exception as e: + verbose_proxy_logger.warning( + "litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - {}".format( + str(e) + ) + ) + return [] + + +async def get_mcp_toolset_by_name( + prisma_client: PrismaClient, + toolset_name: str, +) -> Optional[MCPToolset]: + row = await prisma_client.db.litellm_mcptoolsettable.find_first( + where={"toolset_name": toolset_name} + ) + if row is None: + return None + return _toolset_from_row(row) + + +async def update_mcp_toolset( + prisma_client: PrismaClient, + data: UpdateMCPToolsetRequest, + touched_by: str, +) -> Optional[MCPToolset]: + data_dict = data.model_dump(exclude_none=True, exclude={"toolset_id"}) + if "tools" in data_dict: + data_dict["tools"] = json.dumps(data_dict["tools"]) + data_dict["updated_by"] = touched_by + try: + row = await prisma_client.db.litellm_mcptoolsettable.update( + where={"toolset_id": data.toolset_id}, + data=data_dict, + ) + except Exception as e: + from prisma.errors import RecordNotFoundError + + if isinstance(e, RecordNotFoundError): + return None + raise + return _toolset_from_row(row) + + +async def delete_mcp_toolset( + prisma_client: PrismaClient, + toolset_id: str, +) -> Optional[MCPToolset]: + try: + row = await prisma_client.db.litellm_mcptoolsettable.delete( + where={"toolset_id": toolset_id} + ) + except Exception as e: + from prisma.errors import RecordNotFoundError + + if isinstance(e, RecordNotFoundError): + return None + raise + return _toolset_from_row(row) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 07f3ee64d25..140bcefc1a9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -857,6 +857,8 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): mcp_servers: Optional[List[str]] = None mcp_access_groups: Optional[List[str]] = None mcp_tool_permissions: Optional[Dict[str, List[str]]] = None + mcp_toolsets: Optional[List[str]] = None + blocked_tools: Optional[List[str]] = None vector_stores: Optional[List[str]] = None agents: Optional[List[str]] = None agent_access_groups: Optional[List[str]] = None @@ -885,9 +887,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[ - dict - ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[dict] = ( + {} + ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -1029,9 +1031,9 @@ class RegenerateKeyRequest(GenerateKeyRequest): spend: Optional[float] = None metadata: Optional[dict] = None new_master_key: Optional[str] = None - grace_period: Optional[ - str - ] = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke + grace_period: Optional[str] = ( + None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke + ) class ResetSpendRequest(LiteLLMPydanticObjectBase): @@ -1541,12 +1543,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[ - AllowedModelRegion - ] = None # require all user requests to use models in this specific region - default_model: Optional[ - str - ] = None # if no equivalent model in allowed region - default all requests to this model + allowed_model_region: Optional[AllowedModelRegion] = ( + None # require all user requests to use models in this specific region + ) + default_model: Optional[str] = ( + None # if no equivalent model in allowed region - default all requests to this model + ) object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @model_validator(mode="before") @@ -1569,12 +1571,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[ - AllowedModelRegion - ] = None # require all user requests to use models in this specific region - default_model: Optional[ - str - ] = None # if no equivalent model in allowed region - default all requests to this model + allowed_model_region: Optional[AllowedModelRegion] = ( + None # require all user requests to use models in this specific region + ) + default_model: Optional[str] = ( + None # if no equivalent model in allowed region - default all requests to this model + ) object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @@ -1664,15 +1666,15 @@ class NewTeamRequest(TeamBase): ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm model_tpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[ - float - ] = None # allow user to set a budget for all team members - team_member_rpm_limit: Optional[ - int - ] = None # allow user to set RPM limit for all team members - team_member_tpm_limit: Optional[ - int - ] = None # allow user to set TPM limit for all team members + team_member_budget: Optional[float] = ( + None # allow user to set a budget for all team members + ) + team_member_rpm_limit: Optional[int] = ( + None # allow user to set RPM limit for all team members + ) + team_member_tpm_limit: Optional[int] = ( + None # allow user to set TPM limit for all team members + ) team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None @@ -1769,9 +1771,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase): class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str - callback_type: Optional[ - Literal["success", "failure", "success_and_failure"] - ] = "success_and_failure" + callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = ( + "success_and_failure" + ) callback_vars: Dict[str, str] @model_validator(mode="before") @@ -1848,6 +1850,8 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): vector_stores: Optional[List[str]] = [] agents: Optional[List[str]] = [] agent_access_groups: Optional[List[str]] = [] + mcp_toolsets: Optional[List[str]] = None + blocked_tools: Optional[List[str]] = [] class LiteLLM_TeamTable(TeamBase): @@ -2111,9 +2115,9 @@ class ConfigList(LiteLLMPydanticObjectBase): stored_in_db: Optional[bool] field_default_value: Any premium_field: bool = False - nested_fields: Optional[ - List[FieldDetail] - ] = None # For nested dictionary or Pydantic fields + nested_fields: Optional[List[FieldDetail]] = ( + None # For nested dictionary or Pydantic fields + ) class UserHeaderMapping(LiteLLMPydanticObjectBase): @@ -2471,9 +2475,9 @@ class UserAPIKeyAuth( user_max_budget: Optional[float] = None request_route: Optional[str] = None user: Optional[Any] = None # Expanded user object when expand=user is used - created_by_user: Optional[ - Any - ] = None # Expanded created_by user when expand=user is used + created_by_user: Optional[Any] = ( + None # Expanded created_by user when expand=user is used + ) end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None # Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery # and forwarded into outbound tokens by guardrails such as MCPJWTSigner. @@ -2612,9 +2616,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: Optional[str] = None created_at: datetime updated_at: datetime - user: Optional[ - Any - ] = None # You might want to replace 'Any' with a more specific type if available + user: Optional[Any] = ( + None # You might want to replace 'Any' with a more specific type if available + ) litellm_budget_table: Optional[LiteLLM_BudgetTable] = None user_email: Optional[str] = None @@ -3769,9 +3773,9 @@ class TeamModelDeleteRequest(BaseModel): # Organization Member Requests class OrganizationMemberAddRequest(OrgMemberAddRequest): organization_id: str - max_budget_in_organization: Optional[ - float - ] = None # Users max budget within the organization + max_budget_in_organization: Optional[float] = ( + None # Users max budget within the organization + ) class OrganizationMemberDeleteRequest(MemberDeleteRequest): @@ -4026,9 +4030,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase): Maps provider names to their budget configs. """ - providers: Dict[ - str, ProviderBudgetResponseObject - ] = {} # Dictionary mapping provider names to their budget configurations + providers: Dict[str, ProviderBudgetResponseObject] = ( + {} + ) # Dictionary mapping provider names to their budget configurations class ProxyStateVariables(TypedDict): @@ -4190,9 +4194,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): enforce_rbac: bool = False roles_jwt_field: Optional[str] = None # v2 on role mappings role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[ - str - ] = None # can be either user / team, inferred from the role mapping + object_id_jwt_field: Optional[str] = ( + None # can be either user / team, inferred from the role mapping + ) scope_mappings: Optional[List[ScopeMapping]] = None enforce_scope_based_access: bool = False enforce_team_based_model_access: bool = False diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 82ee11a0f42..114103508ea 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -431,7 +431,13 @@ class PrismaManager: else: # Use prisma db push with increased timeout subprocess.run( - ["prisma", "db", "push", "--accept-data-loss"], + [ + "prisma", + "db", + "push", + "--accept-data-loss", + "--skip-generate", + ], timeout=60, check=True, ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index e4bb288cda9..40d00adeb0b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -37,9 +37,10 @@ from fastapi import ( from fastapi.responses import JSONResponse try: - from prisma.errors import RecordNotFoundError + from prisma.errors import RecordNotFoundError, UniqueViolationError except ImportError: RecordNotFoundError = Exception # type: ignore + UniqueViolationError = Exception # type: ignore import litellm from litellm._logging import verbose_logger, verbose_proxy_logger @@ -424,17 +425,17 @@ if MCP_AVAILABLE: inherited_credentials["scopes"] = existing_server.scopes # AWS SigV4 fields if existing_server.aws_access_key_id: - inherited_credentials[ - "aws_access_key_id" - ] = existing_server.aws_access_key_id + inherited_credentials["aws_access_key_id"] = ( + existing_server.aws_access_key_id + ) if existing_server.aws_secret_access_key: - inherited_credentials[ - "aws_secret_access_key" - ] = existing_server.aws_secret_access_key + inherited_credentials["aws_secret_access_key"] = ( + existing_server.aws_secret_access_key + ) if existing_server.aws_session_token: - inherited_credentials[ - "aws_session_token" - ] = existing_server.aws_session_token + inherited_credentials["aws_session_token"] = ( + existing_server.aws_session_token + ) if existing_server.aws_region_name: inherited_credentials["aws_region_name"] = existing_server.aws_region_name if existing_server.aws_service_name: @@ -2031,3 +2032,192 @@ if MCP_AVAILABLE: f"Failed to load OpenAPI registry from {_OPENAPI_REGISTRY_PATH}: {e}" ) return {"apis": []} + + # --------------------------------------------------------------------------- + # MCP Toolset endpoints + # --------------------------------------------------------------------------- + + from litellm.proxy._experimental.mcp_server.toolset_db import ( + create_mcp_toolset, + delete_mcp_toolset, + get_mcp_toolset, + list_mcp_toolsets, + update_mcp_toolset, + ) + from litellm.types.mcp_server.mcp_toolset import ( + NewMCPToolsetRequest, + UpdateMCPToolsetRequest, + ) + + @router.post( + "/toolset", + description="Create a new MCP toolset (admin only)", + status_code=status.HTTP_201_CREATED, + ) + @management_endpoint_wrapper + async def add_mcp_toolset( + payload: NewMCPToolsetRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header(None), + ): + """Create a named toolset — a curated selection of {server_id, tool_name} pairs.""" + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "Only proxy admins can create MCP toolsets."}, + ) + touched_by = ( + litellm_changed_by or user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME + ) + try: + result = await create_mcp_toolset(prisma_client, payload, touched_by) + except UniqueViolationError: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "error": f"A toolset named '{payload.toolset_name}' already exists." + }, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.invalidate_toolset_cache() + return result + + @router.get( + "/toolset", + description="List MCP toolsets accessible to the calling key", + ) + @management_endpoint_wrapper + async def fetch_mcp_toolsets( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """Return toolsets the calling key is allowed to access.""" + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + is_admin = _user_has_admin_view(user_api_key_dict) + op = user_api_key_dict.object_permission + # mcp_toolsets=None or [] both mean "not restricted by toolsets". + # For admins: either value → no restriction → return all. + # For non-admins: either value → no toolsets explicitly granted → return nothing. + # (An admin whose DB row has mcp_toolsets=[] should still see all toolsets.) + raw_toolsets = getattr(op, "mcp_toolsets", None) if op else None + if not raw_toolsets: + if is_admin: + return await list_mcp_toolsets(prisma_client) + return [] + return await list_mcp_toolsets(prisma_client, toolset_ids=raw_toolsets) + + @router.get( + "/toolset/{toolset_id}", + description="Get a specific MCP toolset by ID", + ) + @management_endpoint_wrapper + async def fetch_mcp_toolset( + toolset_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + # Non-admin keys may only fetch toolsets they've been explicitly granted. + if not _user_has_admin_view(user_api_key_dict): + op = user_api_key_dict.object_permission + granted = getattr(op, "mcp_toolsets", None) if op else None + if granted is None or toolset_id not in granted: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "API key does not have access to this toolset."}, + ) + toolset = await get_mcp_toolset(prisma_client, toolset_id) + if toolset is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"Toolset '{toolset_id}' not found."}, + ) + return toolset + + @router.put( + "/toolset", + description="Update an existing MCP toolset (admin only)", + ) + @management_endpoint_wrapper + async def edit_mcp_toolset( + payload: UpdateMCPToolsetRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header(None), + ): + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "Only proxy admins can update MCP toolsets."}, + ) + touched_by = ( + litellm_changed_by or user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME + ) + try: + result = await update_mcp_toolset(prisma_client, payload, touched_by) + except UniqueViolationError: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "error": ( + f"A toolset named '{payload.toolset_name}' already exists." + if payload.toolset_name + else "A toolset with that name already exists." + ) + }, + ) + if result is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"Toolset '{payload.toolset_id}' not found."}, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.invalidate_toolset_cache( + getattr(payload, "toolset_id", None) + ) + return result + + @router.delete( + "/toolset/{toolset_id}", + description="Delete an MCP toolset (admin only)", + status_code=status.HTTP_202_ACCEPTED, + ) + @management_endpoint_wrapper + async def remove_mcp_toolset( + toolset_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header(None), + ): + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "Only proxy admins can delete MCP toolsets."}, + ) + deleted = await delete_mcp_toolset(prisma_client, toolset_id) + if deleted is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"Toolset '{toolset_id}' not found."}, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.invalidate_toolset_cache(toolset_id) + return Response(status_code=status.HTTP_202_ACCEPTED) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 8aba8307b9d..410f636693f 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -208,10 +208,10 @@ async def _resolve_team_allowed_mcp_servers( ) direct_servers: List[str] = team_object_permission.mcp_servers or [] - access_group_servers: List[ - str - ] = await MCPRequestHandler._get_mcp_servers_from_access_groups( - team_object_permission.mcp_access_groups or [] + access_group_servers: List[str] = ( + await MCPRequestHandler._get_mcp_servers_from_access_groups( + team_object_permission.mcp_access_groups or [] + ) ) raw_tool_perms = team_object_permission.mcp_tool_permissions or {} if isinstance(raw_tool_perms, str): @@ -286,6 +286,19 @@ def _extract_requested_mcp_access_groups( return set() +def _extract_requested_mcp_toolsets( + object_permission: Optional[dict], +) -> Set[str]: + """Extract MCP toolset IDs from a key's object_permission dict.""" + if not object_permission or not isinstance(object_permission, dict): + return set() + + toolsets = object_permission.get("mcp_toolsets") + if isinstance(toolsets, list): + return set(toolsets) + return set() + + async def validate_key_mcp_servers_against_team( object_permission: Optional[dict], team_obj: Optional["LiteLLM_TeamTableCachedObj"], @@ -305,8 +318,10 @@ async def validate_key_mcp_servers_against_team( requested_servers = _extract_requested_mcp_server_ids(object_permission) requested_access_groups = _extract_requested_mcp_access_groups(object_permission) + requested_toolsets = _extract_requested_mcp_toolsets(object_permission) + # Nothing to validate - if not requested_servers and not requested_access_groups: + if not requested_servers and not requested_access_groups and not requested_toolsets: return allow_all_keys_servers = _get_allow_all_keys_server_ids() @@ -364,3 +379,24 @@ async def validate_key_mcp_servers_against_team( status_code=status.HTTP_403_FORBIDDEN, detail={"error": detail}, ) + + # Validate requested toolsets against team's allowed toolsets. + # Only enforce the team-based restriction when a team is present — standalone + # keys (no team) can freely be granted any toolset by an admin. + if requested_toolsets and team_obj is not None: + team_op = team_obj.object_permission + team_mcp_toolsets = team_op.mcp_toolsets if team_op is not None else None + # None or [] means the team has no toolset restriction — allow any toolsets. + if team_mcp_toolsets: + disallowed_toolsets = requested_toolsets - set(team_mcp_toolsets) + if disallowed_toolsets: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + f"Key requests MCP toolsets not allowed by team '{team_obj.team_id}': " + f"{sorted(disallowed_toolsets)}. " + f"Team allows: {sorted(team_mcp_toolsets)}." + ) + }, + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fd88f44fbcf..1b98d1ac848 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -194,6 +194,7 @@ def generate_feedback_box(): print() # noqa +import contextlib from collections import defaultdict from contextlib import asynccontextmanager from functools import lru_cache @@ -2159,11 +2160,9 @@ def run_ollama_serve(): with open(os.devnull, "w") as devnull: subprocess.Popen(command, stdout=devnull, stderr=devnull) except Exception as e: - verbose_proxy_logger.debug( - f""" + verbose_proxy_logger.debug(f""" LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve` - """ - ) + """) def _get_process_rss_mb() -> Optional[float]: @@ -6019,9 +6018,9 @@ class ProxyStartupEvent: """ from litellm.secret_managers.main import str_to_bool - _use_redis_transaction_buffer: Optional[ - Union[bool, str] - ] = general_settings.get("use_redis_transaction_buffer", False) + _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( + general_settings.get("use_redis_transaction_buffer", False) + ) if isinstance(_use_redis_transaction_buffer, str): _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) @@ -13908,18 +13907,147 @@ app.include_router(agent_endpoints_router) app.include_router(compliance_router) app.include_router(a2a_router) app.include_router(access_group_router) + + +async def _stream_mcp_asgi_response( + handle_fn, scope: dict, receive +) -> "StreamingResponse": + """ + Call an ASGI MCP handler and return a StreamingResponse so SSE/streaming works. + + asyncio.create_task copies the current context, so any ContextVar set before + this call (e.g. _mcp_active_toolset_id) is visible inside the handler task. + """ + from starlette.responses import StreamingResponse + + headers_ready: asyncio.Future = asyncio.get_running_loop().create_future() + body_queue: asyncio.Queue = asyncio.Queue(maxsize=1024) + + async def bridging_send(message): + if message["type"] == "http.response.start": + if not headers_ready.done(): + headers_ready.set_result( + (message.get("status", 200), message.get("headers", [])) + ) + elif message["type"] == "http.response.body": + chunk = message.get("body", b"") + if chunk: + await body_queue.put(chunk) + if not message.get("more_body", False): + await body_queue.put(None) # EOF sentinel + + handler_task = asyncio.create_task(handle_fn(scope, receive, bridging_send)) + + # If the handler task dies (exception or cancellation) without sending the EOF + # sentinel, body_iter() would block forever on body_queue.get(). The callback + # below guarantees the queue gets unblocked regardless of how the task ends. + def _ensure_eof(task: asyncio.Task) -> None: + if task.cancelled() or task.exception() is not None: + body_queue.put_nowait(None) + + handler_task.add_done_callback(_ensure_eof) + + try: + status, raw_headers = await asyncio.wait_for( + asyncio.shield(headers_ready), timeout=30.0 + ) + except asyncio.TimeoutError: + handler_task.cancel() + raise HTTPException( + status_code=504, detail="MCP handler did not respond in time" + ) + + headers_dict = {k.decode("latin-1"): v.decode("latin-1") for k, v in raw_headers} + + async def body_iter(): + try: + while True: + chunk = await body_queue.get() + if chunk is None: + break + yield chunk + finally: + if not handler_task.done(): + handler_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await handler_task + + return StreamingResponse( + body_iter(), + status_code=status, + headers=headers_dict, + media_type=headers_dict.get("content-type"), + ) + + ######################################################## # MCP Server ######################################################## +# Toolset-namespaced MCP routes - handle /toolset/{toolset_name}/mcp +# Must be declared BEFORE /{mcp_server_name}/mcp to avoid being swallowed by the catchall. +@app.api_route( + "/toolset/{toolset_name}/mcp", + methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], +) +async def toolset_mcp_route(toolset_name: str, request: Request): + """ + Namespace a toolset as its own MCP endpoint. + + Connecting to /toolset//mcp exposes exactly the tools defined in + the toolset. Access is enforced: non-admin API keys must have the toolset + listed in their object_permission.mcp_toolsets grant list, or the request + will be rejected with a 403. + """ + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.server import ( + _mcp_active_toolset_id, + handle_streamable_http_mcp, + ) + + if prisma_client is None: + raise HTTPException(status_code=503, detail="Database not available") + + toolset = await global_mcp_server_manager.get_toolset_by_name_cached( + prisma_client, toolset_name + ) + if toolset is None: + raise HTTPException( + status_code=404, + detail=f"Toolset '{toolset_name}' not found", + ) + + scope = dict(request.scope) + scope["path"] = "/mcp" + + token = _mcp_active_toolset_id.set(toolset.toolset_id) + try: + return await _stream_mcp_asgi_response( + handle_streamable_http_mcp, scope, request.receive + ) + finally: + _mcp_active_toolset_id.reset(token) + + except HTTPException as e: + raise e + except Exception as e: + verbose_proxy_logger.error( + f"Error handling toolset MCP route for {toolset_name}: {str(e)}" + ) + raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") + + # Dynamic MCP server routes - handle /{mcp_server_name}/mcp @app.api_route( "/{mcp_server_name}/mcp", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], ) async def dynamic_mcp_route(mcp_server_name: str, request: Request): - """Handle dynamic MCP server routes like /github_mcp/mcp""" + """Handle dynamic MCP server routes like /github_mcp/mcp and toolset routes like /devtooling-prod/mcp""" try: # Validate that the MCP server exists from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -13933,6 +14061,32 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): mcp_server_name, client_ip=client_ip ) if mcp_server is None: + # Check if this is a toolset name — toolsets are accessible at /{name}/mcp + # the same way individual servers are, no separate /toolset/ prefix needed. + if prisma_client is not None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.server import ( + _mcp_active_toolset_id, + handle_streamable_http_mcp, + ) + + toolset = await global_mcp_server_manager.get_toolset_by_name_cached( + prisma_client, mcp_server_name + ) + if toolset is not None: + scope = dict(request.scope) + scope["path"] = "/mcp" + + token = _mcp_active_toolset_id.set(toolset.toolset_id) + try: + return await _stream_mcp_asgi_response( + handle_streamable_http_mcp, scope, request.receive + ) + finally: + _mcp_active_toolset_id.reset(token) + raise HTTPException( status_code=404, detail=f"MCP server '{mcp_server_name}' not found" ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 5d804d75b6a..fce95465b55 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -273,6 +273,7 @@ model LiteLLM_ObjectPermissionTable { agent_access_groups String[] @default([]) models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission + mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -331,6 +332,18 @@ model LiteLLM_MCPServerTable { @@index([approval_status]) } +// Named collection of {server_id, tool_name} pairs that can be granted to keys/teams +model LiteLLM_MCPToolsetTable { + toolset_id String @id @default(uuid()) + toolset_name String @unique + description String? + tools Json @default("[]") // [{server_id: string, tool_name: string}] + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} + // Per-user BYOK credentials for MCP servers model LiteLLM_MCPUserCredentials { id String @id @default(uuid()) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 7a3934ffdaa..b729cdb92f2 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -119,6 +119,54 @@ class LiteLLM_Proxy_MCP_Handler: return mcp_tools_with_litellm_proxy, other_tools + @staticmethod + async def _apply_toolset_permissions( + resolved_toolset_ids: List[str], + resolved_mcp_servers: List[str], + user_api_key_auth: Any, + ) -> Any: + """Apply resolved toolset permissions to user_api_key_auth and return updated auth.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + tool_permissions = ( + await global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=resolved_toolset_ids + ) + ) + all_server_ids = list( + set(tool_permissions.keys()) | set(resolved_mcp_servers) + ) + existing_op = user_api_key_auth.object_permission + if existing_op is not None: + merged_tool_perms = dict(existing_op.mcp_tool_permissions or {}) + for server_id, tool_names in tool_permissions.items(): + existing_tools = merged_tool_perms.get(server_id, []) + merged_tool_perms[server_id] = list( + set(existing_tools) | set(tool_names) + ) + updated_op = existing_op.model_copy( + update={ + "mcp_servers": all_server_ids, + "mcp_tool_permissions": merged_tool_perms, + "mcp_toolsets": [], + } + ) + else: + updated_op = LiteLLM_ObjectPermissionTable( + object_permission_id="toolset-scope", + mcp_servers=all_server_ids, + mcp_tool_permissions=tool_permissions, + ) + return user_api_key_auth.model_copy(update={"object_permission": updated_op}) + except Exception as _e: + verbose_logger.debug(f"Could not apply toolset permissions: {_e}") + return user_api_key_auth + @staticmethod async def _get_mcp_tools_from_manager( user_api_key_auth: Any, @@ -160,10 +208,75 @@ class LiteLLM_Proxy_MCP_Handler: ): mcp_servers.append(server_url.split("/")[-1]) + # Resolve toolset names: collect all toolset IDs first, then apply their + # combined permissions in a single pass so multiple toolsets are unioned + # rather than the last one overwriting the others. + resolved_mcp_servers: List[str] = [] + resolved_toolset_ids: List[str] = [] + for name in mcp_servers: + if not global_mcp_server_manager.get_mcp_server_by_name(name): + try: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is not None: + toolset = ( + await global_mcp_server_manager.get_toolset_by_name_cached( + prisma_client, name + ) + ) + if toolset is not None: + # Access control: only allow if the key explicitly grants this toolset. + if user_api_key_auth is not None: + from litellm.proxy.management_endpoints.common_utils import ( + _user_has_admin_view, + ) + + is_admin = _user_has_admin_view(user_api_key_auth) + if not is_admin: + op = user_api_key_auth.object_permission + granted = ( + getattr(op, "mcp_toolsets", None) + if op + else None + ) + # None means no grants configured → deny (consistent with + # fetch_mcp_toolsets which returns [] for unconfigured keys) + if ( + granted is None + or toolset.toolset_id not in granted + ): + verbose_logger.debug( + f"Key does not have access to toolset '{name}', skipping." + ) + continue + resolved_toolset_ids.append(toolset.toolset_id) + # Don't add to resolved_mcp_servers — toolset scope + # restricts via object_permission, not server name filter. + continue + except Exception as _e: + verbose_logger.debug(f"Could not resolve '{name}' as toolset: {_e}") + resolved_mcp_servers.append(name) + + # Apply all resolved toolsets at once (union), avoiding permission overwrite. + if resolved_toolset_ids and user_api_key_auth is not None: + user_api_key_auth = await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions( + resolved_toolset_ids=resolved_toolset_ids, + resolved_mcp_servers=resolved_mcp_servers, + user_api_key_auth=user_api_key_auth, + ) + + # When toolsets were resolved we updated object_permission.mcp_servers to the + # full union (toolset server IDs + direct server names). Passing a name-based + # filter here would exclude those toolset server IDs (which are UUIDs, not + # names), so use None and let the auth object's mcp_servers do the filtering. + effective_server_filter = ( + None if resolved_toolset_ids else (resolved_mcp_servers or None) + ) + tools = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, + mcp_servers=effective_server_filter, mcp_server_auth_headers=mcp_server_auth_headers, log_list_tools_to_spendlogs=True, list_tools_log_source="responses", @@ -178,7 +291,7 @@ class LiteLLM_Proxy_MCP_Handler: ) allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( - mcp_servers=mcp_servers, + mcp_servers=effective_server_filter, allowed_mcp_servers=allowed_mcp_servers, ) @@ -682,14 +795,14 @@ class LiteLLM_Proxy_MCP_Handler: standard_logging_mcp_tool_call["mcp_server_logo_url"] = logo_url cost_info = mcp_info.get("mcp_server_cost_info") if cost_info: - standard_logging_mcp_tool_call[ - "mcp_server_cost_info" - ] = cost_info + standard_logging_mcp_tool_call["mcp_server_cost_info"] = ( + cost_info + ) if litellm_logging_obj: - litellm_logging_obj.model_call_details[ - "mcp_tool_call_metadata" - ] = standard_logging_mcp_tool_call + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( + standard_logging_mcp_tool_call + ) litellm_logging_obj.model = f"MCP: {tool_name}" litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value diff --git a/litellm/types/mcp_server/mcp_toolset.py b/litellm/types/mcp_server/mcp_toolset.py new file mode 100644 index 00000000000..7f78a22bfe2 --- /dev/null +++ b/litellm/types/mcp_server/mcp_toolset.py @@ -0,0 +1,34 @@ +from datetime import datetime +from typing import List, Optional + +from pydantic import BaseModel +from typing_extensions import TypedDict + + +class MCPToolsetTool(TypedDict): + server_id: str + tool_name: str + + +class MCPToolset(BaseModel): + toolset_id: str + toolset_name: str + description: Optional[str] = None + tools: List[MCPToolsetTool] = [] + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None + + +class NewMCPToolsetRequest(BaseModel): + toolset_name: str + description: Optional[str] = None + tools: List[MCPToolsetTool] = [] + + +class UpdateMCPToolsetRequest(BaseModel): + toolset_id: str + toolset_name: Optional[str] = None + description: Optional[str] = None + tools: Optional[List[MCPToolsetTool]] = None diff --git a/schema.prisma b/schema.prisma index 5d804d75b6a..fce95465b55 100644 --- a/schema.prisma +++ b/schema.prisma @@ -273,6 +273,7 @@ model LiteLLM_ObjectPermissionTable { agent_access_groups String[] @default([]) models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission + mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -331,6 +332,18 @@ model LiteLLM_MCPServerTable { @@index([approval_status]) } +// Named collection of {server_id, tool_name} pairs that can be granted to keys/teams +model LiteLLM_MCPToolsetTable { + toolset_id String @id @default(uuid()) + toolset_name String @unique + description String? + tools Json @default("[]") // [{server_id: string, tool_name: string}] + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} + // Per-user BYOK credentials for MCP servers model LiteLLM_MCPUserCredentials { id String @id @default(uuid()) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index edc69ad6a4f..8cab5689568 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1897,18 +1897,18 @@ class TestMCPServerManagerReload: db_row = _make_db_mcp_server("server-1", timestamp) + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock( + return_value=[db_row] + ) with patch( - "litellm.proxy._experimental.mcp_server.db.get_all_mcp_servers", - new=AsyncMock(return_value=[db_row]), - ) as mock_get_all, patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=object(), + return_value=mock_prisma, ), patch.object( manager, "build_mcp_server_from_table", AsyncMock() ) as mock_build: await manager.reload_servers_from_database() - mock_get_all.assert_awaited_once() mock_build.assert_not_awaited() assert manager.registry["server-1"] is existing_server @@ -1940,12 +1940,13 @@ class TestMCPServerManagerReload: updated_at=new_timestamp, ) + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock( + return_value=[db_row] + ) with patch( - "litellm.proxy._experimental.mcp_server.db.get_all_mcp_servers", - new=AsyncMock(return_value=[db_row]), - ) as mock_get_all, patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=object(), + return_value=mock_prisma, ), patch.object( manager, "build_mcp_server_from_table", @@ -1953,7 +1954,6 @@ class TestMCPServerManagerReload: ) as mock_build: await manager.reload_servers_from_database() - mock_get_all.assert_awaited_once() mock_build.assert_awaited_once_with(db_row) assert manager.registry["server-1"] is rebuilt_server diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py new file mode 100644 index 00000000000..2ffa997bdde --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py @@ -0,0 +1,288 @@ +"""Tests for MCP toolset scope enforcement.""" + +import asyncio +from typing import Dict, List, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LitellmUserRoles, + UserAPIKeyAuth, +) + + +def _make_auth( + mcp_servers: Optional[List[str]] = None, + mcp_tool_permissions: Optional[Dict[str, List[str]]] = None, + mcp_toolsets: Optional[List[str]] = None, +) -> UserAPIKeyAuth: + op = LiteLLM_ObjectPermissionTable( + object_permission_id="test", + mcp_servers=mcp_servers, + mcp_tool_permissions=mcp_tool_permissions or {}, + mcp_toolsets=mcp_toolsets, + ) + return UserAPIKeyAuth( + api_key="sk-test", + object_permission=op, + ) + + +class TestApplyToolsetScope: + """Tests for _apply_toolset_scope helper.""" + + @pytest.mark.asyncio + async def test_restricts_to_toolset_servers_and_tools(self): + from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope + + toolset_perms = { + "server-a": ["tool1", "tool2"], + "server-b": ["tool3"], + } + with patch( + "litellm.proxy._experimental.mcp_server.server." + "global_mcp_server_manager.resolve_toolset_tool_permissions", + new=AsyncMock(return_value=toolset_perms), + ): + # Key has been explicitly granted toolset-123 — access check passes. + auth = _make_auth( + mcp_servers=["server-a", "server-b", "server-c"], + mcp_toolsets=["toolset-123"], + ) + result = await _apply_toolset_scope(auth, "toolset-123") + + op = result.object_permission + assert op is not None + assert set(op.mcp_servers or []) == {"server-a", "server-b"} + assert op.mcp_tool_permissions == toolset_perms + + @pytest.mark.asyncio + async def test_admin_creates_object_permission_when_none(self): + """Admin key with object_permission=None can access any toolset.""" + from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope + + toolset_perms = {"server-a": ["tool1"]} + with patch( + "litellm.proxy._experimental.mcp_server.server." + "global_mcp_server_manager.resolve_toolset_tool_permissions", + new=AsyncMock(return_value=toolset_perms), + ): + auth = UserAPIKeyAuth( + api_key="sk-test", + user_role=LitellmUserRoles.PROXY_ADMIN, + object_permission=None, + ) + result = await _apply_toolset_scope(auth, "toolset-123") + + op = result.object_permission + assert op is not None + assert op.mcp_servers == ["server-a"] + assert op.mcp_tool_permissions == toolset_perms + + @pytest.mark.asyncio + async def test_non_admin_no_object_permission_raises_403(self): + """Non-admin key with object_permission=None is denied (no grants configured).""" + from starlette.exceptions import HTTPException + + from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope + + auth = UserAPIKeyAuth(api_key="sk-test", object_permission=None) + with pytest.raises(HTTPException) as exc_info: + await _apply_toolset_scope(auth, "toolset-123") + assert exc_info.value.status_code == 403 + + +class TestFetchMCPToolsetsAccess: + """Tests for GET /v1/mcp/toolset access control.""" + + @pytest.mark.asyncio + async def test_non_admin_empty_grants_returns_empty(self): + """Non-admin key with mcp_toolsets=[] must not see any toolsets.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_toolsets, + ) + + auth = _make_auth(mcp_toolsets=[]) + mock_client = MagicMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets", + new=AsyncMock(return_value=[]), + ) as mock_list, + ): + result = await fetch_mcp_toolsets(user_api_key_dict=auth) + + assert result == [] + mock_list.assert_not_called() + + @pytest.mark.asyncio + async def test_admin_unrestricted_returns_all(self): + """Admin key with mcp_toolsets absent (None) gets all toolsets.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_toolsets, + ) + + auth = UserAPIKeyAuth( + api_key="sk-test", + user_role=LitellmUserRoles.PROXY_ADMIN, + object_permission=None, + ) + fake_toolsets = [MagicMock(), MagicMock()] + mock_client = MagicMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets", + new=AsyncMock(return_value=fake_toolsets), + ) as mock_list, + ): + result = await fetch_mcp_toolsets(user_api_key_dict=auth) + + assert result == fake_toolsets + mock_list.assert_called_once_with(mock_client) + + @pytest.mark.asyncio + async def test_non_admin_none_grants_returns_empty(self): + """Non-admin key with no object_permission (field absent) gets no toolsets.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_toolsets, + ) + + auth = UserAPIKeyAuth(api_key="sk-test", object_permission=None) + mock_client = MagicMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets", + new=AsyncMock(return_value=[]), + ) as mock_list, + ): + result = await fetch_mcp_toolsets(user_api_key_dict=auth) + + assert result == [] + mock_list.assert_not_called() + + @pytest.mark.asyncio + async def test_populated_grants_filters_toolsets(self): + """Key with explicit toolset IDs fetches only those IDs from the DB.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_toolsets, + ) + + auth = _make_auth(mcp_toolsets=["ts-1", "ts-2"]) + fake_toolsets = [MagicMock(toolset_id="ts-1"), MagicMock(toolset_id="ts-2")] + mock_client = MagicMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets", + new=AsyncMock(return_value=fake_toolsets), + ) as mock_list, + ): + result = await fetch_mcp_toolsets(user_api_key_dict=auth) + + assert len(result) == 2 + mock_list.assert_called_once_with(mock_client, toolset_ids=["ts-1", "ts-2"]) + + +class TestMCPActiveToolsetContextVar: + """Tests for _mcp_active_toolset_id ContextVar — clients cannot inject it.""" + + def test_contextvar_default_is_none(self): + from litellm.proxy._experimental.mcp_server.server import _mcp_active_toolset_id + + assert _mcp_active_toolset_id.get() is None + + def test_contextvar_set_and_reset(self): + from litellm.proxy._experimental.mcp_server.server import _mcp_active_toolset_id + + token = _mcp_active_toolset_id.set("toolset-abc") + assert _mcp_active_toolset_id.get() == "toolset-abc" + _mcp_active_toolset_id.reset(token) + assert _mcp_active_toolset_id.get() is None + + @pytest.mark.asyncio + async def test_client_header_is_stripped_in_scope(self): + """handle_streamable_http_mcp strips x-mcp-toolset-id from scope before passing to session manager.""" + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + ) + + scope = { + "type": "http", + "path": "/mcp", + "method": "GET", + "query_string": b"", + "headers": [ + (b"authorization", b"Bearer sk-test"), + (b"x-mcp-toolset-id", b"evil-toolset"), + (b"content-type", b"application/json"), + ], + } + mock_auth = UserAPIKeyAuth(api_key="sk-test") + + async def fake_receive(): + return {"type": "http.disconnect"} + + async def fake_send(msg): + pass + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new=AsyncMock( + return_value=(mock_auth, None, [], {}, {}, scope["headers"]) + ), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.IPAddressUtils", + MagicMock(get_mcp_client_ip=MagicMock(return_value="127.0.0.1")), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + MagicMock(get_mcp_server_by_name=MagicMock(return_value=None)), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.MCPDebug", + MagicMock( + maybe_build_debug_headers=MagicMock(return_value=None), + ), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.set_auth_context", + MagicMock(), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session", + new=AsyncMock(return_value=True), + ), + ): + await handle_streamable_http_mcp(scope, fake_receive, fake_send) + + header_keys = [k for k, _ in scope["headers"]] + assert b"x-mcp-toolset-id" not in header_keys + assert b"authorization" in header_keys + assert b"content-type" in header_keys diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index 15fdc7bd0c4..f706883f384 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -385,6 +385,7 @@ async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch fake_manager = types.SimpleNamespace( get_allowed_mcp_servers=AsyncMock(return_value=[]), get_mcp_servers_from_ids=MagicMock(return_value=[]), + get_mcp_server_by_name=MagicMock(return_value=None), ) monkeypatch.setattr( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 15907190998..3ee82699eb8 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -1,16 +1,20 @@ +import json +import os +from unittest.mock import MagicMock, patch + +import pytest +import redis +import redis.asyncio as async_redis + from litellm._redis import ( - get_redis_url_from_environment, _get_redis_cluster_kwargs, get_redis_async_client, get_redis_client, get_redis_connection_pool, + get_redis_url_from_environment, ) -import json -import os -import pytest -from unittest.mock import MagicMock, patch -import redis -import redis.asyncio as async_redis +from litellm._redis_credential_provider import GCPIAMCredentialProvider + def test_get_redis_url_from_environment_single_url(monkeypatch): """Test when REDIS_URL is directly provided""" @@ -23,6 +27,7 @@ def test_get_redis_url_from_environment_single_url(monkeypatch): # Assert that the returned URL matches the expected value assert redis_url == "redis://redis-server:6379/0" + def test_get_redis_url_from_environment_host_port(monkeypatch): """Test when REDIS_HOST and REDIS_PORT are provided""" # Set the environment variables @@ -39,6 +44,7 @@ def test_get_redis_url_from_environment_host_port(monkeypatch): # Assert that the returned URL matches the expected value assert redis_url == "redis://redis-server:6379" + def test_get_redis_url_from_environment_with_ssl(monkeypatch): """Test when SSL is enabled""" # Set the environment variables @@ -55,6 +61,7 @@ def test_get_redis_url_from_environment_with_ssl(monkeypatch): # Assert that the returned URL uses rediss:// protocol assert redis_url == "rediss://redis-server:6379" + def test_get_redis_url_from_environment_with_username_password(monkeypatch): """Test when username and password are provided""" # Set the environment variables @@ -69,6 +76,7 @@ def test_get_redis_url_from_environment_with_username_password(monkeypatch): # Assert that the returned URL includes username:password@ assert redis_url == "redis://user:password@redis-server:6379" + def test_get_redis_url_from_environment_with_password_only(monkeypatch): """Test when only password is provided""" # Set the environment variables @@ -85,6 +93,7 @@ def test_get_redis_url_from_environment_with_password_only(monkeypatch): # Assert that the returned URL includes :password@ assert redis_url == "redis://password@redis-server:6379" + def test_get_redis_url_from_environment_with_all_options(monkeypatch): """Test when all options are provided""" # Set the environment variables @@ -100,6 +109,7 @@ def test_get_redis_url_from_environment_with_all_options(monkeypatch): # Assert that the returned URL includes all components assert redis_url == "rediss://user:password@redis-server:6379" + def test_get_redis_url_from_environment_missing_host_port(monkeypatch): """Test error when required variables are missing""" # Make sure these environment variables don't exist @@ -110,9 +120,13 @@ def test_get_redis_url_from_environment_missing_host_port(monkeypatch): # Call the function and expect a ValueError with pytest.raises(ValueError) as excinfo: get_redis_url_from_environment() - + # Check the error message - assert "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" in str(excinfo.value) + assert ( + "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" + in str(excinfo.value) + ) + def test_get_redis_url_from_environment_missing_port(monkeypatch): """Test error when only REDIS_HOST is provided but REDIS_PORT is missing""" @@ -124,57 +138,135 @@ def test_get_redis_url_from_environment_missing_port(monkeypatch): # Call the function and expect a ValueError with pytest.raises(ValueError) as excinfo: get_redis_url_from_environment() - + # Check the error message - assert "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" in str(excinfo.value) + assert ( + "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" + in str(excinfo.value) + ) + def test_max_connections_in_cluster_kwargs(): """Test that max_connections is included in Redis cluster kwargs""" kwargs = _get_redis_cluster_kwargs() - assert "max_connections" in kwargs, "max_connections should be in available Redis cluster kwargs" + assert ( + "max_connections" in kwargs + ), "max_connections should be in available Redis cluster kwargs" + def test_get_redis_async_client_with_connection_pool(): """Test that connection_pool parameter is properly passed to Redis client""" # Create a mock connection pool mock_pool = MagicMock(spec=async_redis.BlockingConnectionPool) - + # Mock the Redis client creation - with patch('litellm._redis.async_redis.Redis') as mock_redis, \ - patch('litellm._redis._get_redis_client_logic') as mock_logic: - + with patch("litellm._redis.async_redis.Redis") as mock_redis, patch( + "litellm._redis._get_redis_client_logic" + ) as mock_logic: + # Configure mock to return basic redis kwargs - mock_logic.return_value = { - "host": "localhost", - "port": 6379, - "db": 0 - } - + mock_logic.return_value = {"host": "localhost", "port": 6379, "db": 0} + # Call get_redis_async_client with connection_pool get_redis_async_client(connection_pool=mock_pool) - + # Verify Redis was called with connection_pool in kwargs call_kwargs = mock_redis.call_args[1] - assert "connection_pool" in call_kwargs, "connection_pool should be passed to Redis client" - assert call_kwargs["connection_pool"] == mock_pool, "connection_pool should match the provided pool" + assert ( + "connection_pool" in call_kwargs + ), "connection_pool should be passed to Redis client" + assert ( + call_kwargs["connection_pool"] == mock_pool + ), "connection_pool should match the provided pool" + def test_get_redis_async_client_without_connection_pool(): """Test that Redis client works without connection_pool parameter""" - with patch('litellm._redis.async_redis.Redis') as mock_redis, \ - patch('litellm._redis._get_redis_client_logic') as mock_logic: - + with patch("litellm._redis.async_redis.Redis") as mock_redis, patch( + "litellm._redis._get_redis_client_logic" + ) as mock_logic: + # Configure mock to return basic redis kwargs - mock_logic.return_value = { - "host": "localhost", - "port": 6379, - "db": 0 - } - + mock_logic.return_value = {"host": "localhost", "port": 6379, "db": 0} + # Call get_redis_async_client without connection_pool get_redis_async_client() - + # Verify Redis was called without connection_pool in kwargs call_kwargs = mock_redis.call_args[1] - assert "connection_pool" not in call_kwargs, "connection_pool should not be in kwargs when not provided" + assert ( + "connection_pool" not in call_kwargs + ), "connection_pool should not be in kwargs when not provided" + + +def test_gcp_iam_credential_provider_get_credentials(): + """GCPIAMCredentialProvider.get_credentials() returns a fresh token tuple on every call.""" + service_account = "projects/-/serviceAccounts/test@project.iam.gserviceaccount.com" + + with patch( + "litellm._redis_credential_provider._generate_gcp_iam_access_token", + return_value="tok-1", + ) as mock_gen: + provider = GCPIAMCredentialProvider(service_account) + creds = provider.get_credentials() + + assert creds == ("tok-1",) + mock_gen.assert_called_once_with(service_account) + + +def test_gcp_iam_credential_provider_regenerates_token_on_each_call(): + """Each call to get_credentials() generates a new token (no caching).""" + service_account = "projects/-/serviceAccounts/test@project.iam.gserviceaccount.com" + tokens = ["tok-1", "tok-2", "tok-3"] + + with patch( + "litellm._redis_credential_provider._generate_gcp_iam_access_token", + side_effect=tokens, + ) as mock_gen: + provider = GCPIAMCredentialProvider(service_account) + results = [provider.get_credentials() for _ in range(3)] + + assert results == [("tok-1",), ("tok-2",), ("tok-3",)] + assert mock_gen.call_count == 3 + + +def test_get_redis_async_client_gcp_cluster_uses_credential_provider(): + """ + When startup_nodes + gcp_service_account are provided, the async cluster client + must be constructed with a GCPIAMCredentialProvider — not a static password. + This ensures that the 1-hour IAM token expiry does not cause auth failures. + """ + startup_nodes = [{"host": "redis-node-1", "port": 6379}] + + mock_connect_func = MagicMock() + mock_connect_func._gcp_service_account = ( + "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com" + ) + + redis_kwargs = { + "startup_nodes": startup_nodes, + "redis_connect_func": mock_connect_func, + } + + with patch("litellm._redis.async_redis.RedisCluster") as mock_cluster, patch( + "litellm._redis._get_redis_client_logic", return_value=redis_kwargs + ): + get_redis_async_client() + + assert mock_cluster.called + cluster_call_kwargs = mock_cluster.call_args[1] + + # Must use credential_provider, not a static password + assert ( + "credential_provider" in cluster_call_kwargs + ), "async GCP cluster must use credential_provider for per-connection token refresh" + assert isinstance( + cluster_call_kwargs["credential_provider"], GCPIAMCredentialProvider + ) + assert ( + "password" not in cluster_call_kwargs + ), "async GCP cluster must not use a static password (expires after 1h)" + @patch("litellm._redis.init_redis_cluster") def test_sync_client_prefers_cluster_over_url(mock_init_cluster, monkeypatch): @@ -194,6 +286,7 @@ def test_sync_client_prefers_cluster_over_url(mock_init_cluster, monkeypatch): "startup_nodes" in call_kwargs ), "startup_nodes must be forwarded to init_redis_cluster" + @patch("litellm._redis.async_redis.RedisCluster") def test_async_client_prefers_cluster_over_url(mock_cluster_cls, monkeypatch): """ @@ -207,12 +300,18 @@ def test_async_client_prefers_cluster_over_url(mock_cluster_cls, monkeypatch): mock_cluster_cls.assert_called_once() call_kwargs = mock_cluster_cls.call_args[1] - assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to async RedisCluster" - assert len(call_kwargs["startup_nodes"]) == 1, "should forward exactly 1 cluster node" + assert ( + "startup_nodes" in call_kwargs + ), "startup_nodes must be forwarded to async RedisCluster" + assert ( + len(call_kwargs["startup_nodes"]) == 1 + ), "should forward exactly 1 cluster node" @patch("litellm._redis.async_redis.RedisCluster") -def test_async_client_prefers_cluster_over_url_via_env_var(mock_cluster_cls, monkeypatch): +def test_async_client_prefers_cluster_over_url_via_env_var( + mock_cluster_cls, monkeypatch +): """ Test get_redis_async_client returns async RedisCluster when REDIS_CLUSTER_NODES is set even if REDIS_URL is also set. @@ -227,10 +326,15 @@ def test_async_client_prefers_cluster_over_url_via_env_var(mock_cluster_cls, mon mock_cluster_cls.assert_called_once() call_kwargs = mock_cluster_cls.call_args[1] - assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to async RedisCluster" + assert ( + "startup_nodes" in call_kwargs + ), "startup_nodes must be forwarded to async RedisCluster" + @patch("litellm._redis.init_redis_cluster") -def test_sync_client_prefers_cluster_over_url_via_env_var(mock_init_cluster, monkeypatch): +def test_sync_client_prefers_cluster_over_url_via_env_var( + mock_init_cluster, monkeypatch +): """ Test get_redis_client returns RedisCluster when REDIS_CLUSTER_NODES is set even if REDIS_URL is also set. @@ -246,11 +350,16 @@ def test_sync_client_prefers_cluster_over_url_via_env_var(mock_init_cluster, mon mock_init_cluster.assert_called_once() call_kwargs = mock_init_cluster.call_args[0][0] - assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to init_redis_cluster" + assert ( + "startup_nodes" in call_kwargs + ), "startup_nodes must be forwarded to init_redis_cluster" assert len(call_kwargs["startup_nodes"]) == 1 + @patch("litellm._redis.init_redis_cluster") -def test_sync_client_preserves_password_for_cluster_when_url_also_set(mock_init_cluster, monkeypatch): +def test_sync_client_preserves_password_for_cluster_when_url_also_set( + mock_init_cluster, monkeypatch +): """ Test _get_redis_client_logic does not strip password from redis_kwargs when startup_nodes is present even if REDIS_URL is also set. @@ -264,7 +373,9 @@ def test_sync_client_preserves_password_for_cluster_when_url_also_set(mock_init_ mock_init_cluster.assert_called_once() call_kwargs = mock_init_cluster.call_args[0][0] - assert "password" in call_kwargs, "password must not be stripped when routing to cluster" + assert ( + "password" in call_kwargs + ), "password must not be stripped when routing to cluster" assert call_kwargs["password"] == "secret" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPToolsets.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPToolsets.ts new file mode 100644 index 00000000000..aacb446d4bc --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPToolsets.ts @@ -0,0 +1,16 @@ +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { fetchMCPToolsets } from "@/components/networking"; +import { MCPToolset } from "@/components/mcp_tools/types"; +import useAuthorized from "../useAuthorized"; + +const mcpToolsetKeys = createQueryKeys("mcpToolsets"); + +export const useMCPToolsets = () => { + const { accessToken } = useAuthorized(); + return useQuery({ + queryKey: mcpToolsetKeys.list(), + queryFn: async () => await fetchMCPToolsets(accessToken!), + enabled: !!accessToken, + }); +}; diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx index dc4ed25786b..fc4b20517cf 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx @@ -1,13 +1,15 @@ import { useMCPAccessGroups } from "@/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { Select } from "antd"; import React from "react"; interface MCPServerSelectorProps { - onChange: (selected: { servers: string[]; accessGroups: string[] }) => void; + onChange: (selected: { servers: string[]; accessGroups: string[]; toolsets: string[] }) => void; value?: { servers: string[]; accessGroups: string[]; + toolsets?: string[]; }; className?: string; accessToken: string; @@ -16,6 +18,8 @@ interface MCPServerSelectorProps { teamId?: string | null; } +const TOOLSET_PREFIX = "toolset:"; + const MCPServerSelector: React.FC = ({ onChange, value, @@ -27,33 +31,61 @@ const MCPServerSelector: React.FC = ({ }) => { const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers(teamId); const { data: accessGroups = [], isLoading: groupsLoading } = useMCPAccessGroups(); + const { data: toolsets = [], isLoading: toolsetsLoading } = useMCPToolsets(); - const loading = serversLoading || groupsLoading; + const loading = serversLoading || groupsLoading || toolsetsLoading; - // Combine options, access groups first + const accessGroupSet = new Set(accessGroups); + + // Combine options: access groups (green) + servers (blue) + toolsets (purple) const options = [ ...accessGroups.map((group) => ({ label: group, value: group, - isAccessGroup: true, + type: "accessGroup" as const, searchText: `${group} Access Group`, })), ...mcpServers.map((server) => ({ label: `${server.server_name || server.server_id} (${server.server_id})`, value: server.server_id, - isAccessGroup: false, + type: "server" as const, searchText: `${server.server_name || server.server_id} ${server.server_id} MCP Server`, })), + ...toolsets.map((toolset) => ({ + label: toolset.toolset_name, + value: `${TOOLSET_PREFIX}${toolset.toolset_id}`, + type: "toolset" as const, + searchText: `${toolset.toolset_name} ${toolset.toolset_id} Toolset`, + })), ]; - // Flatten value for Select - const selectedValues = [...(value?.servers || []), ...(value?.accessGroups || [])]; + const colorByType: Record = { + accessGroup: "#52c41a", + server: "#1890ff", + toolset: "#722ed1", + }; + const labelByType: Record = { + accessGroup: "Access Group", + server: "MCP Server", + toolset: "Toolset", + }; + + // Flatten value for Select — prefix toolset IDs + const selectedValues = [ + ...(value?.servers || []), + ...(value?.accessGroups || []), + ...(value?.toolsets || []).map((id) => `${TOOLSET_PREFIX}${id}`), + ]; // Handle selection const handleChange = (selected: string[]) => { - const servers = selected.filter((v) => !accessGroups.includes(v)); - const accessGroupsSelected = selected.filter((v) => accessGroups.includes(v)); - onChange({ servers, accessGroups: accessGroupsSelected }); + const toolsetsSelected = selected + .filter((v) => v.startsWith(TOOLSET_PREFIX)) + .map((v) => v.slice(TOOLSET_PREFIX.length)); + const rest = selected.filter((v) => !v.startsWith(TOOLSET_PREFIX)); + const servers = rest.filter((v) => !accessGroupSet.has(v)); + const accessGroupsSelected = rest.filter((v) => accessGroupSet.has(v)); + onChange({ servers, accessGroups: accessGroupsSelected, toolsets: toolsetsSelected }); }; return ( @@ -83,20 +115,20 @@ const MCPServerSelector: React.FC = ({ width: 8, height: 8, borderRadius: "50%", - background: opt.isAccessGroup ? "#52c41a" : "#1890ff", + background: colorByType[opt.type], flexShrink: 0, }} /> {opt.label} - {opt.isAccessGroup ? "Access Group" : "MCP Server"} + {labelByType[opt.type]} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolsetsTab.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolsetsTab.tsx new file mode 100644 index 00000000000..d6f91f99957 --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolsetsTab.tsx @@ -0,0 +1,524 @@ +import React, { useState, useCallback } from "react"; +import { Button, Text, Title } from "@tremor/react"; +import { Modal, Form, Input, message, Spin, Card, Typography, Space } from "antd"; +import { PlusIcon, PencilIcon, TrashIcon } from "@heroicons/react/outline"; +import { ColumnDef } from "@tanstack/react-table"; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useQueryClient } from "@tanstack/react-query"; +import { DataTable } from "../view_logs/table"; +import { + createMCPToolset, + updateMCPToolset, + deleteMCPToolset, + listMCPTools, + getProxyBaseUrl, +} from "../networking"; +import { MCPToolset, MCPToolsetTool } from "./types"; + +const { Text: AntdText } = Typography; + +interface MCPToolsetsTabProps { + accessToken: string | null; + userRole: string | null; +} + +interface ToolsetFormValues { + toolset_name: string; + description?: string; +} + +interface MCPToolListProps { + serverId: string; + serverName: string; + accessToken: string | null; + selectedTools: MCPToolsetTool[]; + onToggle: (tool: MCPToolsetTool) => void; +} + +interface ToolEntry { + name: string; + description?: string; +} + +function MCPToolList({ serverId, serverName, accessToken, selectedTools, onToggle }: MCPToolListProps) { + const [tools, setTools] = useState([]); + const [loading, setLoading] = useState(false); + const [expanded, setExpanded] = useState(false); + + const selectedSet = new Set(selectedTools.filter((t) => t.server_id === serverId).map((t) => t.tool_name)); + + const fetchTools = useCallback(async () => { + if (!accessToken || tools.length > 0) return; + setLoading(true); + try { + const result = await listMCPTools(accessToken, serverId); + const toolList = Array.isArray(result) ? result : result?.tools ?? []; + setTools(toolList.map((t: any) => ({ name: t.name ?? t.tool_name ?? t, description: t.description ?? "" }))); + } catch { + setTools([]); + } finally { + setLoading(false); + } + }, [accessToken, serverId, tools.length]); + + const handleToggle = () => { + if (!expanded) fetchTools(); + setExpanded(!expanded); + }; + + return ( +
+ + {expanded && ( +
+ {loading ? ( +
+ ) : tools.length === 0 ? ( +

No tools found for this server.

+ ) : ( +
+ {tools.map((tool) => { + const selected = selectedSet.has(tool.name); + return ( + + ); + })} +
+ )} +
+ )} +
+ ); +} + +interface CreateToolsetModalProps { + open: boolean; + onClose: () => void; + onSave: (name: string, description: string | undefined, tools: MCPToolsetTool[]) => Promise; + accessToken: string | null; + initialToolset?: MCPToolset; +} + +function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset }: CreateToolsetModalProps) { + const [form] = Form.useForm(); + const [selectedTools, setSelectedTools] = useState(initialToolset?.tools || []); + const [saving, setSaving] = useState(false); + const [serverSearch, setServerSearch] = useState(""); + const { data: mcpServers = [] } = useMCPServers(); + + React.useEffect(() => { + if (open) { + form.setFieldsValue({ + toolset_name: initialToolset?.toolset_name || "", + description: initialToolset?.description || "", + }); + setSelectedTools(initialToolset?.tools || []); + setServerSearch(""); + } + }, [open, initialToolset]); + + const handleToggleTool = (tool: MCPToolsetTool) => { + setSelectedTools((prev) => { + const exists = prev.some((t) => t.server_id === tool.server_id && t.tool_name === tool.tool_name); + return exists + ? prev.filter((t) => !(t.server_id === tool.server_id && t.tool_name === tool.tool_name)) + : [...prev, tool]; + }); + }; + + const handleSubmit = async () => { + const values = await form.validateFields(); + setSaving(true); + try { + await onSave(values.toolset_name, values.description, selectedTools); + onClose(); + } finally { + setSaving(false); + } + }; + + const filteredServers = mcpServers.filter((s) => { + const q = serverSearch.toLowerCase(); + return ( + !q || + (s.alias || "").toLowerCase().includes(q) || + (s.server_name || "").toLowerCase().includes(q) + ); + }); + + return ( + +
+
+ + + + + + +
+
+ +
+ {/* Left panel: Available Tools */} +
+
+ Available Tools +
+ setServerSearch(e.target.value)} + className="mb-2" + allowClear + /> +
+ {filteredServers.length === 0 ? ( + {mcpServers.length === 0 ? "No MCP servers configured" : "No servers match your search"} + ) : ( + filteredServers.map((server) => ( + + )) + )} +
+
+ + {/* Divider */} +
+ + {/* Right panel: Your Toolset */} +
+ + Your Toolset{" "} + ({selectedTools.length} tools) + +
+ {selectedTools.length === 0 ? ( + No tools added yet + ) : ( + selectedTools.map((tool, idx) => ( + + )) + )} +
+
+
+ +
+ + +
+ + ); +} + +function toolsetColumns( + isAdmin: boolean, + onEdit: (t: MCPToolset) => void, + onDelete: (id: string) => void, + proxyBaseUrl: string, +): ColumnDef[] { + return [ + { + header: "Toolset ID", + accessorKey: "toolset_id", + cell: ({ row }) => ( + + {row.original.toolset_id.slice(0, 8)}… + + ), + }, + { + header: "Name", + accessorKey: "toolset_name", + cell: ({ row }) => { + const url = `${proxyBaseUrl}/toolset/${row.original.toolset_name}/mcp`; + return ( +
+
+ + {row.original.toolset_name} +
+ +
+ ); + }, + }, + { + header: "Description", + accessorKey: "description", + cell: ({ row }) => ( + {row.original.description || "—"} + ), + }, + { + header: "Tools", + accessorKey: "tools", + cell: ({ row }) => { + const tools = row.original.tools; + return ( +
+ {tools.slice(0, 4).map((t, i) => ( + + {t.tool_name} + + ))} + {tools.length > 4 && ( + +{tools.length - 4} more + )} +
+ ); + }, + }, + { + header: "Created", + accessorKey: "created_at", + cell: ({ row }) => ( + + {row.original.created_at ? new Date(row.original.created_at).toLocaleDateString() : "—"} + + ), + }, + ...(isAdmin ? [{ + header: "", + id: "actions", + cell: ({ row }: { row: { original: MCPToolset } }) => ( +
+ + +
+ ), + } as ColumnDef] : []), + ]; +} + +function ToolsetUsageGuide() { + const [copied, setCopied] = useState(false); + const proxyBaseUrl = getProxyBaseUrl(); + + const snippet = `{ + "mcpServers": { + "my-toolset": { + "url": "${proxyBaseUrl}/toolset//mcp", + "headers": { "x-litellm-api-key": "Bearer " } + } + } +}`; + + const copy = async () => { + try { + await navigator.clipboard.writeText(snippet); + setCopied(true); + setTimeout(() => setCopied(false), 1500); + } catch { + // ignore + } + }; + + return ( +
+

How toolsets work

+

+ Create a toolset, assign it to a key via API Keys → Edit Key → MCP Servers, then point your MCP client at the toolset URL. The client only sees the tools you picked. +

+
Claude Code / Cursor config
+
+ 
+          {snippet}
+
+ +
+
+ ); +} + +export function MCPToolsetsTab({ accessToken, userRole }: MCPToolsetsTabProps) { + const queryClient = useQueryClient(); + const { data: toolsets = [], isLoading } = useMCPToolsets(); + const [createOpen, setCreateOpen] = useState(false); + const [editToolset, setEditToolset] = useState(null); + const [deleteId, setDeleteId] = useState(null); + const [deleting, setDeleting] = useState(false); + + const isAdmin = userRole === "Admin" || userRole === "proxy_admin"; + + const handleCreate = async (name: string, description: string | undefined, tools: MCPToolsetTool[]) => { + if (!accessToken) return; + await createMCPToolset(accessToken, { toolset_name: name, description, tools }); + message.success("Toolset created"); + queryClient.invalidateQueries({ queryKey: ["mcpToolsets"] }); + }; + + const handleUpdate = async (name: string, description: string | undefined, tools: MCPToolsetTool[]) => { + if (!accessToken || !editToolset) return; + await updateMCPToolset(accessToken, { toolset_id: editToolset.toolset_id, toolset_name: name, description, tools }); + message.success("Toolset updated"); + queryClient.invalidateQueries({ queryKey: ["mcpToolsets"] }); + setEditToolset(null); + }; + + const handleDelete = async () => { + if (!accessToken || !deleteId) return; + setDeleting(true); + try { + await deleteMCPToolset(accessToken, deleteId); + message.success("Toolset deleted"); + queryClient.invalidateQueries({ queryKey: ["mcpToolsets"] }); + setDeleteId(null); + } finally { + setDeleting(false); + } + }; + + const proxyBaseUrl = getProxyBaseUrl(); + const columns = toolsetColumns(isAdmin, setEditToolset, setDeleteId, proxyBaseUrl); + + return ( +
+
+
+ MCP Toolsets + + Curated collections of tools from one or more MCP servers. Assign toolsets to keys and teams via the MCP permissions dropdown. + +
+ {isAdmin && ( + + )} +
+ + + +
} + getRowCanExpand={() => false} + isLoading={isLoading} + noDataMessage="No toolsets yet. Click 'New Toolset' to create one." + loadingMessage="Loading toolsets..." + enableSorting={true} + /> + + setCreateOpen(false)} + onSave={handleCreate} + accessToken={accessToken} + /> + + {editToolset && ( + setEditToolset(null)} + onSave={handleUpdate} + accessToken={accessToken} + initialToolset={editToolset} + /> + )} + + setDeleteId(null)} + onOk={handleDelete} + okText="Delete" + okButtonProps={{ danger: true, loading: deleting }} + title="Delete Toolset" + > +

Are you sure you want to delete this toolset? Keys and teams using it will lose access to the scoped tools.

+ +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx index 4d67456b78a..fc1d941c177 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx @@ -102,7 +102,7 @@ describe("MCPServers", () => { vi.mocked(networking.fetchMCPServers).mockResolvedValue(mockServers); const queryClient = createQueryClient(); - const { getByText } = render( + const { getByText, getAllByText } = render( , @@ -121,8 +121,8 @@ describe("MCPServers", () => { // Verify the mocked server data is rendered in the table expect(getByText("Test Server 1")).toBeInTheDocument(); expect(getByText("Test Server 2")).toBeInTheDocument(); - expect(getByText("test-server-1")).toBeInTheDocument(); - expect(getByText("test-server-2")).toBeInTheDocument(); + expect(getAllByText("test-server-1").length).toBeGreaterThan(0); + expect(getAllByText("test-server-2").length).toBeGreaterThan(0); // Verify the API was called // Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index 34eeb6b1e86..f0fed7c7fec 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -9,6 +9,7 @@ import { useMCPServerHealth } from "../../app/(dashboard)/hooks/mcpServers/useMC import NotificationsManager from "../molecules/notifications_manager"; import { deleteMCPServer } from "../networking"; import { MCPSubmissionsTab } from "./MCPSubmissionsTab"; +import { MCPToolsetsTab } from "./MCPToolsetsTab"; import { DataTable } from "../view_logs/table"; import CreateMCPServer from "./create_mcp_server"; import MCPConnect from "./mcp_connect"; @@ -348,6 +349,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID })
All Servers + Toolsets Connect Semantic Filter Network Settings @@ -426,6 +428,9 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID })
)} + + + diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 67526ca2b07..acf8310c803 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -231,6 +231,20 @@ export interface MCPServerProps { userID: string | null; } +export interface MCPToolsetTool { + server_id: string; + tool_name: string; +} + +export interface MCPToolset { + toolset_id: string; + toolset_name: string; + description?: string; + tools: MCPToolsetTool[]; + created_at?: string; + created_by?: string; +} + // Discoverable MCP server from the curated registry export interface DiscoverableMCPServer { name: string; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 7f8dea4cbfe..28f8d308de7 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6675,6 +6675,99 @@ export const deleteMCPServer = async (accessToken: string, serverId: string) => } }; +export const fetchMCPToolsets = async (accessToken: string): Promise => { + try { + const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/toolset`; + const response = await fetch(url, { + method: HTTP_REQUEST.GET, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + return await response.json(); + } catch (error) { + console.error("Failed to fetch MCP toolsets:", error); + throw error; + } +}; + +export const createMCPToolset = async (accessToken: string, formValues: Record) => { + try { + const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/toolset`; + const response = await fetch(url, { + method: HTTP_REQUEST.POST, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(formValues), + }); + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + return await response.json(); + } catch (error) { + console.error("Failed to create MCP toolset:", error); + throw error; + } +}; + +export const updateMCPToolset = async (accessToken: string, formValues: Record) => { + try { + const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/toolset`; + const response = await fetch(url, { + method: HTTP_REQUEST.PUT, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(formValues), + }); + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + return await response.json(); + } catch (error) { + console.error("Failed to update MCP toolset:", error); + throw error; + } +}; + +export const deleteMCPToolset = async (accessToken: string, toolsetId: string) => { + try { + const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/toolset/${toolsetId}`; + const response = await fetch(url, { + method: HTTP_REQUEST.DELETE, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + } catch (error) { + console.error("Failed to delete MCP toolset:", error); + throw error; + } +}; + export const registerMCPServer = async (accessToken: string, formValues: Record) => { try { const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/server/register`; diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx index ac55a57c7c3..685467e1d3e 100644 --- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx +++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx @@ -9,6 +9,7 @@ interface ObjectPermission { mcp_servers: string[]; mcp_access_groups?: string[]; mcp_tool_permissions?: Record; + mcp_toolsets?: string[]; vector_stores: string[]; agents?: string[]; agent_access_groups?: string[]; @@ -31,17 +32,19 @@ export function ObjectPermissionsView({ const mcpServers = objectPermission?.mcp_servers || []; const mcpAccessGroups = objectPermission?.mcp_access_groups || []; const mcpToolPermissions = objectPermission?.mcp_tool_permissions || {}; + const mcpToolsets = objectPermission?.mcp_toolsets || []; const agents = objectPermission?.agents || []; const agentAccessGroups = objectPermission?.agent_access_groups || []; const content = (
- { ); // Verify empty state message - expect(screen.getByText("No MCP servers or access groups configured")).toBeInTheDocument(); + expect(screen.getByText("No MCP servers, access groups, or toolsets configured")).toBeInTheDocument(); // Verify count badge shows 0 expect(screen.getByText("0")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx index 962fc00670a..649e8efdf4a 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx @@ -2,25 +2,28 @@ import React, { useState, useEffect } from "react"; import { Text, Badge } from "@tremor/react"; import { ServerIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; import { Tooltip } from "antd"; -import { fetchMCPServers } from "../networking"; -import { MCPServer } from "../mcp_tools/types"; +import { fetchMCPServers, fetchMCPToolsets } from "../networking"; +import { MCPServer, MCPToolset } from "../mcp_tools/types"; interface MCPServerPermissionsProps { mcpServers: string[]; mcpAccessGroups?: string[]; mcpToolPermissions?: Record; + mcpToolsets?: string[]; accessToken?: string | null; } -export function MCPServerPermissions({ - mcpServers, - mcpAccessGroups = [], +export function MCPServerPermissions({ + mcpServers, + mcpAccessGroups = [], mcpToolPermissions = {}, - accessToken + mcpToolsets = [], + accessToken }: MCPServerPermissionsProps) { const [mcpServerDetails, setMCPServerDetails] = useState([]); - const [accessGroupNames, setAccessGroupNames] = useState([]); + const [toolsetDetails, setToolsetDetails] = useState([]); const [expandedServers, setExpandedServers] = useState>(new Set()); + const [expandedToolsets, setExpandedToolsets] = useState>(new Set()); const toggleServerExpansion = (serverId: string) => { setExpandedServers((prev) => { @@ -34,6 +37,18 @@ export function MCPServerPermissions({ }); }; + const toggleToolsetExpansion = (toolsetId: string) => { + setExpandedToolsets((prev) => { + const newSet = new Set(prev); + if (newSet.has(toolsetId)) { + newSet.delete(toolsetId); + } else { + newSet.add(toolsetId); + } + return newSet; + }); + }; + // Fetch MCP server details when component mounts useEffect(() => { const fetchMCPServerDetails = async () => { @@ -53,20 +68,23 @@ export function MCPServerPermissions({ fetchMCPServerDetails(); }, [accessToken, mcpServers.length]); - // Fetch MCP access group names + // Fetch toolset details useEffect(() => { - const fetchGroups = async () => { - if (accessToken && mcpAccessGroups.length > 0) { + const fetchToolsets = async () => { + if (accessToken && mcpToolsets.length > 0) { try { - const groups = await import("../networking").then((m) => m.fetchMCPAccessGroups(accessToken)); - setAccessGroupNames(Array.isArray(groups) ? groups : groups.data || []); + const all = await fetchMCPToolsets(accessToken); + const filtered = Array.isArray(all) + ? all.filter((t: MCPToolset) => mcpToolsets.includes(t.toolset_id)) + : []; + setToolsetDetails(filtered); } catch (error) { - console.error("Error fetching MCP access groups:", error); + console.error("Error fetching toolsets:", error); } } }; - fetchGroups(); - }, [accessToken, mcpAccessGroups.length]); + fetchToolsets(); + }, [accessToken, mcpToolsets.length]); // Function to get display name for MCP server const getMCPServerDisplayName = (serverId: string) => { @@ -78,17 +96,12 @@ export function MCPServerPermissions({ return serverId; }; - // Function to get display name for access group - const getAccessGroupDisplayName = (group: string) => { - return group; - }; - // Merge servers and access groups into one list const mergedItems = [ ...mcpServers.map((server) => ({ type: "server", value: server })), ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group })), ]; - const totalCount = mergedItems.length; + const totalCount = mergedItems.length + mcpToolsets.length; return (
@@ -99,21 +112,21 @@ export function MCPServerPermissions({ {totalCount}
- + {totalCount > 0 ? (
{mergedItems.map((item, index) => { const toolsForServer = item.type === "server" ? mcpToolPermissions[item.value] : undefined; const hasToolRestrictions = toolsForServer && toolsForServer.length > 0; const isExpanded = expandedServers.has(item.value); - + return (
-
hasToolRestrictions && toggleServerExpansion(item.value)} className={`flex items-center gap-3 py-2 px-3 rounded-lg border border-gray-200 transition-all ${ - hasToolRestrictions - ? 'cursor-pointer hover:bg-gray-50 hover:border-gray-300' + hasToolRestrictions + ? 'cursor-pointer hover:bg-gray-50 hover:border-gray-300' : 'bg-white' }`} > @@ -128,14 +141,14 @@ export function MCPServerPermissions({ ) : (
- {getAccessGroupDisplayName(item.value)} + {item.value} Group
)}
- + {hasToolRestrictions && (
{toolsForServer.length} @@ -148,7 +161,7 @@ export function MCPServerPermissions({
)}
- + {/* Show tool permissions if expanded */} {hasToolRestrictions && isExpanded && (
@@ -167,11 +180,66 @@ export function MCPServerPermissions({
); })} + + {/* Toolsets section */} + {mcpToolsets.length > 0 && mcpToolsets.map((toolsetId, index) => { + const detail = toolsetDetails.find((t) => t.toolset_id === toolsetId); + const isExpanded = expandedToolsets.has(toolsetId); + const toolCount = detail?.tools.length ?? 0; + + return ( +
+
toolCount > 0 && toggleToolsetExpansion(toolsetId)} + className={`flex items-center gap-3 py-2 px-3 rounded-lg border border-purple-200 transition-all ${ + toolCount > 0 ? 'cursor-pointer hover:bg-purple-50 hover:border-purple-300' : 'bg-white' + }`} + > +
+ + + {detail?.toolset_name ?? toolsetId} + + + Toolset + +
+ {toolCount > 0 && ( +
+ {toolCount} + {toolCount === 1 ? "tool" : "tools"} + {isExpanded ? ( + + ) : ( + + )} +
+ )} +
+ + {toolCount > 0 && isExpanded && detail && ( +
+
+ {detail.tools.map((tool, toolIndex) => ( + + {tool.server_id.slice(0, 6)}… + {tool.tool_name} + + ))} +
+
+ )} +
+ ); + })}
) : (
- No MCP servers or access groups configured + No MCP servers, access groups, or toolsets configured
)}
diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index e80c2b2f742..2bfe08efae8 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -34,7 +34,8 @@ import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "../../mcp_tools/M import { MCPServer } from "../../mcp_tools/types"; import { ByokCredentialModal } from "../../mcp_tools/ByokCredentialModal"; import NotificationsManager from "../../molecules/notifications_manager"; -import { callMCPTool, fetchMCPServers, listMCPTools } from "../../networking"; +import { callMCPTool, fetchMCPServers, fetchMCPToolsets, listMCPTools } from "../../networking"; +import { MCPToolset } from "../../mcp_tools/types"; import TagSelector from "../../tag_management/TagSelector"; import VectorStoreSelector from "../../vector_store_management/VectorStoreSelector"; import { makeA2ASendMessageRequest } from "../llm_calls/a2a_send_message"; @@ -111,6 +112,8 @@ const ChatUI: React.FC = ({ fixedModel, }) => { const [mcpServers, setMCPServers] = useState([]); + const [mcpToolsets, setMCPToolsets] = useState([]); + const [isToolsetsInfoModalVisible, setIsToolsetsInfoModalVisible] = useState(false); const [byokModalServer, setByokModalServer] = useState(null); const [selectedMCPServers, setSelectedMCPServers] = useState(() => { const saved = sessionStorage.getItem("selectedMCPServers"); @@ -257,15 +260,19 @@ const ChatUI: React.FC = ({ const chatEndRef = useRef(null); - // Fetch MCP servers + // Fetch MCP servers and toolsets const loadMCPServers = async () => { const userApiKey = apiKeySource === "session" ? accessToken : apiKey; if (!userApiKey) return; setIsLoadingMCPServers(true); try { - const servers = await fetchMCPServers(userApiKey); + const [servers, toolsets] = await Promise.all([ + fetchMCPServers(userApiKey), + fetchMCPToolsets(userApiKey).catch(() => []), + ]); setMCPServers(Array.isArray(servers) ? servers : servers.data || []); + setMCPToolsets(Array.isArray(toolsets) ? toolsets : []); } catch (error) { console.error("Error fetching MCP servers:", error); } finally { @@ -416,17 +423,29 @@ const ChatUI: React.FC = ({ loadMCPServers(); }, [accessToken, userID, userRole, apiKeySource, apiKey, token, simplified]); - // Load tools when MCP direct mode has a server selected + // Load tools when MCP direct mode has a server (or toolset) selected useEffect(() => { if ( endpointType === EndpointType.MCP && selectedMCPServers.length === 1 && - selectedMCPServers[0] !== "__all__" && - !serverToolsMap[selectedMCPServers[0]] + selectedMCPServers[0] !== "__all__" ) { - loadServerTools(selectedMCPServers[0]); + const selected = selectedMCPServers[0]; + if (selected.startsWith("toolset:")) { + // For a toolset, load tools for each server in it + const toolsetId = selected.slice("toolset:".length); + const toolset = mcpToolsets.find((t) => t.toolset_id === toolsetId); + if (toolset) { + const uniqueServerIds = [...new Set(toolset.tools.map((t) => t.server_id))]; + uniqueServerIds.forEach((sid) => { + if (!serverToolsMap[sid]) loadServerTools(sid); + }); + } + } else if (!serverToolsMap[selected]) { + loadServerTools(selected); + } } - }, [endpointType, selectedMCPServers, serverToolsMap]); + }, [endpointType, selectedMCPServers, serverToolsMap, mcpToolsets]); // Fetch agents when A2A endpoint is selected useEffect(() => { @@ -572,19 +591,34 @@ const ChatUI: React.FC = ({ // For MCP direct mode, require server and tool selection, and get form values early let mcpToolArguments: Record = {}; if (endpointType === EndpointType.MCP) { - const mcpServerId = + const rawSelected = selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" ? selectedMCPServers[0] : null; - if (!mcpServerId) { + if (!rawSelected) { NotificationsManager.fromBackend("Please select an MCP server to test"); return; } + // Resolve the real server ID (toolsets use toolset: prefix) + const mcpServerId = rawSelected.startsWith("toolset:") ? rawSelected : rawSelected; if (!selectedMCPDirectTool) { NotificationsManager.fromBackend("Please select an MCP tool to call"); return; } - const mcpTool = (serverToolsMap[selectedMCPServers[0]] || []).find( + // For toolsets, find the tool in the servers that back this toolset + const toolsetForSelected = rawSelected.startsWith("toolset:") + ? mcpToolsets.find((t) => t.toolset_id === rawSelected.slice("toolset:".length)) + : null; + let searchPool: any[] = []; + if (toolsetForSelected) { + const uniqueServerIds = [...new Set(toolsetForSelected.tools.map((t) => t.server_id))]; + uniqueServerIds.forEach((sid) => { + searchPool = searchPool.concat(serverToolsMap[sid] || []); + }); + } else { + searchPool = serverToolsMap[rawSelected] || []; + } + const mcpTool = searchPool.find( (t: any) => t.name === selectedMCPDirectTool, ); if (!mcpTool) { @@ -742,6 +776,7 @@ const ChatUI: React.FC = ({ mcpServerToolRestrictions, handleMCPEvent, mockTestFallbacks, + mcpToolsets, ); } else if (endpointType === EndpointType.IMAGE) { // For image generation @@ -822,6 +857,7 @@ const ChatUI: React.FC = ({ customProxyBaseUrl || undefined, mcpServers, mcpServerToolRestrictions, + mcpToolsets, ); } else if (endpointType === EndpointType.ANTHROPIC_MESSAGES) { const apiChatHistory = [ @@ -879,14 +915,22 @@ const ChatUI: React.FC = ({ // Handle MCP direct tool calls (no chat completions) if (endpointType === EndpointType.MCP) { - const mcpServerId = + const rawSelected = selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" ? selectedMCPServers[0] : null; - if (mcpServerId && selectedMCPDirectTool) { + // For toolsets, resolve the real server_id from the toolset's tool list + let resolvedServerId = rawSelected; + if (rawSelected?.startsWith("toolset:")) { + const toolsetId = rawSelected.slice("toolset:".length); + const toolset = mcpToolsets.find((t) => t.toolset_id === toolsetId); + const toolEntry = toolset?.tools.find((t) => t.tool_name === selectedMCPDirectTool); + resolvedServerId = toolEntry?.server_id ?? rawSelected; + } + if (resolvedServerId && !resolvedServerId.startsWith("toolset:") && selectedMCPDirectTool) { const result = await callMCPTool( effectiveApiKey, - mcpServerId, + resolvedServerId, selectedMCPDirectTool, mcpToolArguments, selectedGuardrails.length > 0 ? { guardrails: selectedGuardrails } : undefined, @@ -1301,11 +1345,14 @@ const ChatUI: React.FC = ({ className="ml-1" title={ endpointType === EndpointType.MCP - ? "Select an MCP server to test tools directly." - : "Select MCP servers to use in your conversation." + ? "Select an MCP server or toolset to test tools directly." + : "Select MCP servers or toolsets to use in your conversation." } > - + setIsToolsetsInfoModalVisible(true)} + /> {/* MCP Tool selector - only for MCP direct mode */} {endpointType === EndpointType.MCP && selectedMCPServers.length === 1 && - selectedMCPServers[0] !== "__all__" && ( -
- Select Tool - setSelectedMCPDirectTool(value)} + options={toolOptions} + allowClear + className="rounded-md" + /> +
+ ); + })()} {/* Tool restrictions UI (optional) - hidden for MCP direct mode */} {selectedMCPServers.length > 0 && @@ -1924,7 +2037,21 @@ const ChatUI: React.FC = ({ selectedMCPDirectTool ? (
{(() => { - const mcpTool = (serverToolsMap[selectedMCPServers[0]] || []).find( + const rawSel = selectedMCPServers[0]; + let toolPool: any[] = []; + if (rawSel.startsWith("toolset:")) { + const toolsetId = rawSel.slice("toolset:".length); + const toolset = mcpToolsets.find((t) => t.toolset_id === toolsetId); + if (toolset) { + const uniqueServerIds = [...new Set(toolset.tools.map((t) => t.server_id))]; + uniqueServerIds.forEach((sid) => { + toolPool = toolPool.concat(serverToolsMap[sid] || []); + }); + } + } else { + toolPool = serverToolsMap[rawSel] || []; + } + const mcpTool = toolPool.find( (t: any) => t.name === selectedMCPDirectTool, ); return mcpTool ? ( @@ -2070,6 +2197,47 @@ const ChatUI: React.FC = ({ accessToken={accessToken || ""} /> )} + + {/* Toolsets info modal */} + setIsToolsetsInfoModalVisible(false)} + footer={[ + , + ]} + width={600} + > +
+

+ Toolsets are named collections of specific tools from one or more MCP servers. + Instead of exposing all tools from a server, a toolset gives an agent exactly the tools it needs. +

+
+

How to use a toolset:

+
    +
  1. Select a Toolset (purple badge) from the MCP Servers dropdown.
    2. +
  2. The tool picker will show only the tools included in that toolset.
    3. +
  3. Select a tool and fill in its parameters, then send.
    4. +
  4. The tool call is routed to the correct underlying MCP server automatically.
    5. +
+
+
+

+ Example: A "GitHub Read-only" toolset might include only list_repos and get_file from a GitHub MCP server — preventing agents from making writes. +

+
+
+

Creating toolsets:

+

+ Admins can create and manage toolsets from the MCP page → Toolsets tab. + Toolsets can then be assigned to keys and teams to scope their tool access. +

+
+
+
); }; diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx index 3197c9409ce..b2281eaadce 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx @@ -3,7 +3,7 @@ import { ChatCompletionMessageParam } from "openai/resources/chat/completions"; import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { VectorStoreSearchResponse } from "../chat_ui/types"; import { getProxyBaseUrl } from "@/components/networking"; -import { MCPServer, type MCPEvent } from "../../mcp_tools/types"; +import { MCPServer, MCPToolset, type MCPEvent } from "../../mcp_tools/types"; export async function makeOpenAIChatCompletionRequest( chatHistory: { role: string; content: string | any[] }[], @@ -30,6 +30,7 @@ export async function makeOpenAIChatCompletionRequest( mcpServerToolRestrictions?: Record, onMCPEvent?: (event: MCPEvent) => void, mockTestFallbacks?: boolean, + mcpToolsets?: MCPToolset[], ) { // base url should be the current base_url const isLocal = process.env.NODE_ENV === "development"; @@ -82,19 +83,31 @@ export async function makeOpenAIChatCompletionRequest( require_approval: "never", }); } else { - // Individual servers selected - create one entry per server + // Individual servers/toolsets selected - create one entry per item selectedMCPServers.forEach((serverId) => { - const server = mcpServers?.find((s) => s.server_id === serverId); - const serverName = server?.alias || server?.server_name || serverId; - const allowedTools = mcpServerToolRestrictions?.[serverId] || []; + if (serverId.startsWith("toolset:")) { + const toolsetId = serverId.slice("toolset:".length); + const toolset = mcpToolsets?.find((t) => t.toolset_id === toolsetId); + const toolsetName = toolset?.toolset_name || toolsetId; + tools.push({ + type: "mcp", + server_label: toolsetName, + server_url: `litellm_proxy/mcp/${encodeURIComponent(toolsetName)}`, + require_approval: "never", + }); + } else { + const server = mcpServers?.find((s) => s.server_id === serverId); + const serverName = server?.alias || server?.server_name || serverId; + const allowedTools = mcpServerToolRestrictions?.[serverId] || []; - tools.push({ - type: "mcp", - server_label: "litellm", - server_url: `litellm_proxy/mcp/${serverName}`, - require_approval: "never", - ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), - }); + tools.push({ + type: "mcp", + server_label: "litellm", + server_url: `litellm_proxy/mcp/${serverName}`, + require_approval: "never", + ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), + }); + } }); } } diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx index 48d0efca6ee..4e88a356cf3 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx @@ -4,7 +4,7 @@ import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { getProxyBaseUrl } from "@/components/networking"; import NotificationManager from "@/components/molecules/notifications_manager"; import type { MCPEvent } from "../../mcp_tools/types"; -import { MCPServer } from "../../mcp_tools/types"; +import { MCPServer, MCPToolset } from "../../mcp_tools/types"; import { CodeInterpreterResult, CodeInterpreterState, @@ -37,6 +37,7 @@ export async function makeOpenAIResponsesRequest( customBaseUrl?: string, mcpServers?: MCPServer[], mcpServerToolRestrictions?: Record, + mcpToolsets?: MCPToolset[], ) { if (!accessToken) { throw new Error("Virtual Key is required"); @@ -102,21 +103,34 @@ export async function makeOpenAIResponsesRequest( require_approval: "never", }); } else { - // Individual servers selected - create one entry per server + // Individual servers/toolsets selected - create one entry per item selectedMCPServers.forEach((serverId) => { - const server = mcpServers?.find((s) => s.server_id === serverId); - // Use server_name for both routing and labelling. server_name is the - // unique registered identifier; aliases can collide across servers. - const routeName = server?.server_name || serverId; - const allowedTools = mcpServerToolRestrictions?.[serverId] || []; + if (serverId.startsWith("toolset:")) { + // Toolset: same /{name}/mcp pattern as individual servers + const toolsetId = serverId.slice("toolset:".length); + const toolset = mcpToolsets?.find((t) => t.toolset_id === toolsetId); + const toolsetName = toolset?.toolset_name || toolsetId; + tools.push({ + type: "mcp", + server_label: toolsetName, + server_url: `${proxyBaseUrl}/mcp/${encodeURIComponent(toolsetName)}`, + require_approval: "never", + }); + } else { + const server = mcpServers?.find((s) => s.server_id === serverId); + // Use server_name for both routing and labelling. server_name is the + // unique registered identifier; aliases can collide across servers. + const routeName = server?.server_name || serverId; + const allowedTools = mcpServerToolRestrictions?.[serverId] || []; - tools.push({ - type: "mcp", - server_label: routeName, // unique per request — collisions cause silent tool-routing failures - server_url: `${proxyBaseUrl}/mcp/${encodeURIComponent(routeName)}`, - require_approval: "never", - ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), - }); + tools.push({ + type: "mcp", + server_label: routeName, // unique per request — collisions cause silent tool-routing failures + server_url: `${proxyBaseUrl}/mcp/${encodeURIComponent(routeName)}`, + require_approval: "never", + ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), + }); + } }); } } diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index abc967bb984..e59af82e97b 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -105,6 +105,7 @@ export interface TeamData { mcp_servers: string[]; mcp_access_groups?: string[]; mcp_tool_permissions?: Record; + mcp_toolsets?: string[]; vector_stores: string[]; agents?: string[]; agent_access_groups?: string[]; @@ -516,9 +517,10 @@ const TeamInfoView: React.FC = ({ } // Handle object_permission updates - const { servers, accessGroups } = values.mcp_servers_and_groups || { + const { servers, accessGroups, toolsets } = values.mcp_servers_and_groups || { servers: [], accessGroups: [], + toolsets: [], }; const serverIds = new Set(servers || []); const mcpToolPermissions = Object.fromEntries( @@ -535,6 +537,9 @@ const TeamInfoView: React.FC = ({ if (mcpToolPermissions) { updateData.object_permission.mcp_tool_permissions = mcpToolPermissions; } + if (toolsets) { + updateData.object_permission.mcp_toolsets = toolsets; + } delete values.mcp_servers_and_groups; delete values.mcp_tool_permissions; @@ -838,6 +843,7 @@ const TeamInfoView: React.FC = ({ mcp_servers_and_groups: { servers: info.object_permission?.mcp_servers || [], accessGroups: info.object_permission?.mcp_access_groups || [], + toolsets: info.object_permission?.mcp_toolsets || [], }, mcp_tool_permissions: info.object_permission?.mcp_tool_permissions || {}, agents_and_groups: { diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 5d00ab3d0b9..24e3e18b93c 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -160,11 +160,12 @@ export default function KeyInfoView({ } if (formValues.mcp_servers_and_groups !== undefined) { - const { servers, accessGroups } = formValues.mcp_servers_and_groups || { servers: [], accessGroups: [] }; + const { servers, accessGroups, toolsets } = formValues.mcp_servers_and_groups || { servers: [], accessGroups: [], toolsets: [] }; formValues.object_permission = { ...currentKeyData.object_permission, mcp_servers: servers || [], mcp_access_groups: accessGroups || [], + mcp_toolsets: toolsets || [], }; // Remove mcp_servers_and_groups from the top level as it should be in object_permission delete formValues.mcp_servers_and_groups;