From e0463a38fffdeb32dec9297039ede25a113ae543 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 11 Jul 2026 16:10:39 -0700 Subject: [PATCH 01/41] fix(completion): forward aws credential kwargs into litellm_params so the responses bridge keeps WIF auth Chat-completions requests to responses-only Bedrock Mantle models are bridged to the Responses API, but completion() forwarded only aws_bedrock_project_id into get_litellm_params, so aws_role_name, aws_web_identity_token, aws_session_name and the other SigV4 credential kwargs never reached sign_request and botocore fell back to the default credential chain ("Bedrock Mantle auth failed: no Bearer token and no usable AWS credentials"). Forward the whole AWS credential kwarg family, extracted from the OPTIONAL_KWARGS_KEYS set get_litellm_params already supports. --- .../litellm_core_utils/get_litellm_params.py | 56 ++++++++++------- litellm/main.py | 7 ++- tests/test_litellm/test_main.py | 62 +++++++++++++++++++ 3 files changed, 99 insertions(+), 26 deletions(-) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 352e55e9c23..b8ef9d8cca7 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -2,26 +2,8 @@ from typing import Optional from litellm.llms.openai.data_residency import infer_openai_data_residency -# Pre-define optional kwargs keys as frozenset for O(1) lookups -# These are extracted from kwargs only if present, avoiding unnecessary .get() calls -OPTIONAL_KWARGS_KEYS = frozenset( +AWS_CREDENTIAL_KWARGS_KEYS = frozenset( { - "azure_ad_token", - "tenant_id", - "client_id", - "client_secret", - "azure_username", - "azure_password", - "azure_scope", - "timeout", - "gcs_bucket_name", - "bucket_name", - "vertex_credentials", - "vertex_project", - "vertex_location", - "vertex_ai_project", - "vertex_ai_location", - "vertex_ai_credentials", "aws_region_name", "aws_access_key_id", "aws_secret_access_key", @@ -34,14 +16,40 @@ OPTIONAL_KWARGS_KEYS = frozenset( "aws_external_id", "aws_bedrock_runtime_endpoint", "aws_bedrock_project_id", - "tpm", - "rpm", - "itpm", - "otpm", - "use_xai_oauth", } ) +# Pre-define optional kwargs keys as frozenset for O(1) lookups +# These are extracted from kwargs only if present, avoiding unnecessary .get() calls +OPTIONAL_KWARGS_KEYS = ( + frozenset( + { + "azure_ad_token", + "tenant_id", + "client_id", + "client_secret", + "azure_username", + "azure_password", + "azure_scope", + "timeout", + "gcs_bucket_name", + "bucket_name", + "vertex_credentials", + "vertex_project", + "vertex_location", + "vertex_ai_project", + "vertex_ai_location", + "vertex_ai_credentials", + "tpm", + "rpm", + "itpm", + "otpm", + "use_xai_oauth", + } + ) + | AWS_CREDENTIAL_KWARGS_KEYS +) + # Backward-compatible alias for existing imports/tests. _OPTIONAL_KWARGS_KEYS = OPTIONAL_KWARGS_KEYS diff --git a/litellm/main.py b/litellm/main.py index 7d457d9cdd1..6fd68921fb0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -92,7 +92,10 @@ from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) -from litellm.litellm_core_utils.get_litellm_params import OPTIONAL_KWARGS_KEYS +from litellm.litellm_core_utils.get_litellm_params import ( + AWS_CREDENTIAL_KWARGS_KEYS, + OPTIONAL_KWARGS_KEYS, +) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, @@ -5322,7 +5325,7 @@ def completion( # type: ignore tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), use_xai_oauth=kwargs.get("use_xai_oauth", False), - aws_bedrock_project_id=kwargs.get("aws_bedrock_project_id"), + **{key: kwargs[key] for key in AWS_CREDENTIAL_KWARGS_KEYS if key in kwargs}, ) cast(LiteLLMLoggingObj, logging).update_environment_variables( model=model, diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 28cf4fa0744..c5e70e0fabf 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2081,3 +2081,65 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): assert response.usage.prompt_tokens > 0 assert response.usage.completion_tokens > 0 assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + +@pytest.mark.asyncio +async def test_acompletion_forwards_aws_credentials_through_responses_bridge( + respx_mock: respx.MockRouter, monkeypatch +): + from botocore.credentials import Credentials + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + original_disable_aiohttp = litellm.disable_aiohttp_transport + try: + litellm.disable_aiohttp_transport = True + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + get_credentials_mock = MagicMock(return_value=Credentials("fake-key", "fake-secret")) + monkeypatch.setattr(BaseAWSLLM, "get_credentials", get_credentials_mock) + + respx_mock.post("https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses").respond( + json={ + "id": "resp_123", + "object": "response", + "created_at": 1760144904, + "status": "completed", + "model": "openai.gpt-5.4", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + } + ) + + response = await litellm.acompletion( + model="bedrock_mantle/openai.gpt-5.4", + messages=[{"role": "user", "content": "hi"}], + api_base="https://bedrock-mantle.us-east-2.api.aws/v1", + aws_region_name="us-east-2", + aws_session_name="litellm-gcp", + aws_role_name="arn:aws:iam::123456789012:role/litellm-bedrock-role", + aws_web_identity_token="oidc/google/108963886734710037768", + num_retries=0, + ) + + assert response.choices[0].message.content == "ok" + credential_kwargs = get_credentials_mock.call_args.kwargs + assert credential_kwargs["aws_role_name"] == "arn:aws:iam::123456789012:role/litellm-bedrock-role" + assert credential_kwargs["aws_web_identity_token"] == "oidc/google/108963886734710037768" + assert credential_kwargs["aws_session_name"] == "litellm-gcp" + authorization = respx_mock.calls.last.request.headers["Authorization"] + assert authorization.startswith("AWS4-HMAC-SHA256") + assert "fake-key" in authorization + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() From 201730efe573efb7f41941173b74d08e55b77b80 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 11 Jul 2026 16:35:13 -0700 Subject: [PATCH 02/41] fix(bedrock): allow bedrock-mantle:CreateInference in the web identity session policy --- litellm/llms/bedrock/base_aws_llm.py | 9 ++++ .../test_web_identity_session_policy.py | 48 +++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index f449851b76f..df811f8d262 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -877,6 +877,15 @@ class BaseAWSLLM: "Resource": "*", "Condition": {"Bool": {"aws:SecureTransport": "true"}}, }, + { + "Sid": "BedrockMantleLiteLLM", + "Effect": "Allow", + "Action": [ + "bedrock-mantle:CreateInference", + ], + "Resource": "*", + "Condition": {"Bool": {"aws:SecureTransport": "true"}}, + }, ], } assume_role_params = { diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py index 7e9c8a273ae..0cbdc518cc2 100644 --- a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -158,6 +158,54 @@ class TestClaudePlatformActionsCovered: ) +class TestBedrockMantleActionsCovered: + """LIT-3859: bedrock_mantle inference authorizes against the + ``bedrock-mantle`` action namespace, so the session-policy ceiling + must include it or every Mantle request via OIDC/WIF auth denies + with "no session policy allows the bedrock-mantle:CreateInference + action" even when the role's identity policy grants it.""" + + def test_bedrock_mantle_create_inference_present(self): + policy = _captured_policy() + all_actions: set = set() + for stmt in policy["Statement"]: + stmt_actions = stmt.get("Action") + if isinstance(stmt_actions, str): + all_actions.add(stmt_actions) + elif isinstance(stmt_actions, list): + all_actions.update(stmt_actions) + assert "bedrock-mantle:CreateInference" in all_actions, ( + "bedrock-mantle:CreateInference missing from session policy — " + "bedrock_mantle/* requests will 403 on OIDC/WIF auth" + ) + + def test_bedrock_mantle_statement_allows(self): + policy = _captured_policy() + stmt = _statement_by_sid(policy, "BedrockMantleLiteLLM") + assert stmt["Effect"] == "Allow" + assert stmt["Resource"] == "*" + + def test_no_bedrock_mantle_wildcard(self): + policy = _captured_policy() + stmt = _statement_by_sid(policy, "BedrockMantleLiteLLM") + actions = stmt["Action"] + if isinstance(actions, str): + actions = [actions] + assert "bedrock-mantle:*" not in actions, ( + "session policy must not grant bedrock-mantle:* — " + "the ceiling should match the documented action set" + ) + + def test_bedrock_mantle_statement_carries_secure_transport_condition(self): + policy = _captured_policy() + stmt = _statement_by_sid(policy, "BedrockMantleLiteLLM") + cond = stmt.get("Condition") or {} + assert cond.get("Bool", {}).get("aws:SecureTransport") == "true", ( + "BedrockMantleLiteLLM must require aws:SecureTransport=true " + "to keep parity with the bedrock statement" + ) + + def _make_jwt(payload: dict) -> str: def _segment(data: dict) -> str: return base64.urlsafe_b64encode(json.dumps(data).encode()).rstrip(b"=").decode() From cc27528d4f248b5b6e908f424caa270a2df84513 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 11 Jul 2026 16:47:00 -0700 Subject: [PATCH 03/41] test(main): assert the responses bridge forwards static aws keys as well as web identity params --- tests/test_litellm/test_main.py | 28 +++++++++++++++++++++------- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index c5e70e0fabf..4611aafa3c1 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2084,8 +2084,24 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): @pytest.mark.asyncio +@pytest.mark.parametrize( + "aws_credential_kwargs", + [ + { + "aws_session_name": "litellm-gcp", + "aws_role_name": "arn:aws:iam::123456789012:role/litellm-bedrock-role", + "aws_web_identity_token": "oidc/google/108963886734710037768", + }, + { + "aws_access_key_id": "AKIASTATICKEYFORTEST", + "aws_secret_access_key": "static-secret-key", + "aws_session_token": "static-session-token", + }, + ], + ids=["web_identity", "static_keys"], +) async def test_acompletion_forwards_aws_credentials_through_responses_bridge( - respx_mock: respx.MockRouter, monkeypatch + respx_mock: respx.MockRouter, monkeypatch, aws_credential_kwargs: dict ): from botocore.credentials import Credentials @@ -2126,17 +2142,15 @@ async def test_acompletion_forwards_aws_credentials_through_responses_bridge( messages=[{"role": "user", "content": "hi"}], api_base="https://bedrock-mantle.us-east-2.api.aws/v1", aws_region_name="us-east-2", - aws_session_name="litellm-gcp", - aws_role_name="arn:aws:iam::123456789012:role/litellm-bedrock-role", - aws_web_identity_token="oidc/google/108963886734710037768", num_retries=0, + **aws_credential_kwargs, ) assert response.choices[0].message.content == "ok" credential_kwargs = get_credentials_mock.call_args.kwargs - assert credential_kwargs["aws_role_name"] == "arn:aws:iam::123456789012:role/litellm-bedrock-role" - assert credential_kwargs["aws_web_identity_token"] == "oidc/google/108963886734710037768" - assert credential_kwargs["aws_session_name"] == "litellm-gcp" + assert credential_kwargs["aws_region_name"] == "us-east-2" + for key, value in aws_credential_kwargs.items(): + assert credential_kwargs[key] == value authorization = respx_mock.calls.last.request.headers["Authorization"] assert authorization.startswith("AWS4-HMAC-SHA256") assert "fake-key" in authorization From 45fed6a50a231822aeb616c99d4b6f17ffd48da0 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 14:44:10 -0700 Subject: [PATCH 04/41] feat(mcp): generalize the bridge envelope identity to a key_hash or user_id subject The scripted two-header client mints under a virtual key it presents at the token endpoint (key_hash), but the interactive DCR client authenticates via SSO at the bridged authorize, which yields a user, not a key. Make EnvelopeIdentity a discriminated subject (subject_type key_hash | user_id) with key_hash_identity / user_identity constructors, and dispatch admission on it: a key_hash reloads the key, a user_id reloads the user and admits them as themselves (user-level budget and SCIM enforced via the same centralized gate; no team bound, since a user belongs to many teams or none). The interactive producer that mints a user_id envelope lands in the follow-up commit. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 57 +++++++++- .../mcp_server/discoverable_endpoints.py | 4 +- .../outbound_credentials/envelope.py | 48 ++++++-- .../auth/test_user_api_key_auth_mcp.py | 105 +++++++++++++++++- .../test_bridge_credentials.py | 7 +- .../outbound_credentials/test_envelope.py | 35 ++++-- .../mcp_server/test_discoverable_endpoints.py | 3 +- 7 files changed, 230 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index e300a22e5db..faec35db41a 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -18,6 +18,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti is_bridge_envelope_shaped, resolve_bridge_envelope, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeIdentity, +) from litellm.proxy._types import ( UI_TEAM_ID, LiteLLM_TeamTable, @@ -543,7 +546,7 @@ class MCPRequestHandler: header_key = server.alias or server.server_name if header_key is None: raise HTTPException(status_code=500, detail="Server misconfigured: MCP server has no routable name") - admitted = await MCPRequestHandler._reload_admitted_key(result.identity.key_hash) + admitted = await MCPRequestHandler._reload_admitted_principal(result.identity) await MCPRequestHandler._enforce_admitted_live_policy(admitted=admitted, request=request, route=route) injected = {header_key: {"Authorization": result.upstream_authorization.get_secret_value()}} new_headers = {**(mcp_server_auth_headers or {}), **injected} @@ -572,6 +575,58 @@ class MCPRequestHandler: route=route, ) + @staticmethod + async def _reload_admitted_principal(identity: EnvelopeIdentity) -> UserAPIKeyAuth: + """Reload the live litellm record the envelope's subject references. + + Dispatches on the sealed subject type: a ``key_hash`` reloads the virtual key that + minted the envelope (the scripted two-header client that presents a litellm key at the + token endpoint), a ``user_id`` reloads the user that authenticated interactively (the + DCR client, whose SSO login at the bridged authorize yields a user, not a key). Both + return a ``UserAPIKeyAuth`` the caller runs through the centralized policy gate, so + team/project/org/budget/SCIM enforcement is identical to the principal presenting + itself directly.""" + match identity.subject_type: + case "key_hash": + return await MCPRequestHandler._reload_admitted_key(identity.subject) + case "user_id": + return await MCPRequestHandler._reload_admitted_user(identity.subject) + case _: + assert_never(identity.subject_type) + + @staticmethod + async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth: + """Reload the live user an interactively-minted envelope references and admit them as + themselves. + + The DCR client authenticates via SSO at the bridged authorize, which yields a user + subject rather than a virtual key, so the envelope admits under the user's own + identity: the reloaded ``user_id`` rides on the returned ``UserAPIKeyAuth`` and the + caller's centralized policy gate then enforces the user's live budget and org state, + and a SCIM-deactivated owner fails closed here exactly as the key path enforces it. No + team is bound; a user may belong to many teams or none, so the envelope grants the + user's own access rather than silently selecting one team's scope. A missing user + fails closed with a 401 rather than admitting an unresolved identity.""" + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Server misconfigured: no database connection") + try: + user_object = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + except (ProxyException, HTTPException): + raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + if user_object is None: + raise HTTPException(status_code=401, detail="Invalid or expired credential") + if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: + raise HTTPException(status_code=401, detail="Invalid or expired credential") + return UserAPIKeyAuth(user_id=user_object.user_id) + @staticmethod async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: """Reload the live key record an admitted envelope references and re-check live policy. diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 8d1713a5911..3e727ce95bb 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -961,15 +961,15 @@ def _finish_bridge_mint( build_bridge_token_response, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import - EnvelopeIdentity, SealedEnvelope, UpstreamTokenGrant, + key_hash_identity, ) grant = _bridge_grant_from_token_response(token_response) if not isinstance(grant, UpstreamTokenGrant): return _upstream_rejection_to_mint_error(grant) - identity = EnvelopeIdentity(server_id=mcp_server.server_id, key_hash=ready.key_hash) + identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=ready.key_hash) sealed = build_bridge_token_response(identity, grant, ready.keys, now) if not isinstance(sealed, SealedEnvelope): return "too_large" diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py index 517c2ef5c8f..783e64d13e2 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -67,20 +67,42 @@ typed error, never truncated.""" _ENVELOPE_JWT_ALGORITHM = "HS256" -class EnvelopeIdentity(BaseModel): - """The litellm identity the envelope binds the inner grant to. +EnvelopeSubjectType: TypeAlias = Literal["key_hash", "user_id"] +"""Discriminator for what litellm principal the envelope binds the grant to. - ``key_hash`` is the hashed litellm key that authorized the mint, never a raw - credential (and the edge rejects a bare hash presented as a bearer). Admission - reloads the live key record by it, so the key's current team/org/object-permission - restrictions and its revocation state are enforced at use time rather than frozen at - mint time. ``server_id`` binds the envelope to one MCP server so it cannot be replayed - across a server boundary. +``key_hash`` is a hashed virtual key (the scripted two-header client mints under the key it +presents at the token endpoint); ``user_id`` is a litellm user subject (the interactive DCR +client mints under the SSO-authenticated user, which is the only identity that browser login +yields). Admission reloads a key record for the first and a user record for the second, then +runs both through the same live-policy gate, so team/org/budget/revocation enforcement is +identical either way.""" + + +class EnvelopeIdentity(BaseModel): + """The litellm principal the envelope binds the inner grant to. + + ``subject`` is the principal identifier and ``subject_type`` says how to resolve it: a + hashed litellm key (``key_hash``) or a litellm user id (``user_id``), never a raw + credential (and the edge rejects a bare hash or id presented as a bearer). Admission + reloads the live record by it, so the principal's current team/org restrictions and its + revocation state are enforced at use time rather than frozen at mint time. ``server_id`` + binds the envelope to one MCP server so it cannot be replayed across a server boundary. """ model_config = ConfigDict(frozen=True) server_id: str = Field(min_length=1) - key_hash: str = Field(min_length=1) + subject_type: EnvelopeSubjectType + subject: str = Field(min_length=1) + + +def key_hash_identity(server_id: str, key_hash: str) -> EnvelopeIdentity: + """The identity for the scripted client that mints under a presented virtual key.""" + return EnvelopeIdentity(server_id=server_id, subject_type="key_hash", subject=key_hash) + + +def user_identity(server_id: str, user_id: str) -> EnvelopeIdentity: + """The identity for the interactive DCR client that mints under its SSO user subject.""" + return EnvelopeIdentity(server_id=server_id, subject_type="user_id", subject=user_id) class UpstreamTokenGrant(BaseModel): @@ -200,7 +222,8 @@ class _EnvelopeClaims(BaseModel): iat: int exp: int server_id: str = Field(min_length=1) - key_hash: str = Field(min_length=1) + subject_type: EnvelopeSubjectType + subject: str = Field(min_length=1) grant: str = Field(min_length=1) @@ -236,7 +259,8 @@ def mint_envelope( iat=int(now.timestamp()), exp=int(expires_at.timestamp()), server_id=identity.server_id, - key_hash=identity.key_hash, + subject_type=identity.subject_type, + subject=identity.subject, grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key), ) token = ENVELOPE_PREFIX + jwt.encode( @@ -281,7 +305,7 @@ def open_envelope( if not isinstance(grant, UpstreamTokenGrant): return grant return OpenedEnvelope( - identity=EnvelopeIdentity(server_id=claims.server_id, key_hash=claims.key_hash), + identity=EnvelopeIdentity(server_id=claims.server_id, subject_type=claims.subject_type, subject=claims.subject), grant=grant, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index c785ac577f7..a6affe5496c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -4910,6 +4910,7 @@ class TestMCPDcrBridgeDelegateAdmission: cls, *, key_hash=None, + user_id=None, server_id="bridge-server-id", access_token="inner-upstream-access-token", token_type="Bearer", @@ -4921,17 +4922,23 @@ class TestMCPDcrBridgeDelegateAdmission: envelope_keys_from_master_key, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( - EnvelopeIdentity, SealedEnvelope, UpstreamTokenGrant, + key_hash_identity, mint_envelope, + user_identity, ) from pydantic import SecretStr + identity = ( + user_identity(server_id=server_id, user_id=user_id) + if user_id is not None + else key_hash_identity(server_id=server_id, key_hash=key_hash or cls._KEY_HASH) + ) keys = envelope_keys_from_master_key(master_key or cls._MASTER_KEY) now = minted_at or datetime.now(timezone.utc) sealed = mint_envelope( - identity=EnvelopeIdentity(server_id=server_id, key_hash=key_hash or cls._KEY_HASH), + identity=identity, grant=UpstreamTokenGrant( access_token=SecretStr(access_token), token_type=token_type, @@ -4999,6 +5006,22 @@ class TestMCPDcrBridgeDelegateAdmission: stack.enter_context(patcher) yield get_key_object + @staticmethod + @contextlib.contextmanager + def _patch_user_reload(*, return_value=None, side_effect=None): + """Patch the user-subject reload path an interactively-minted envelope takes: the + ``get_user_object`` lookup ``_reload_admitted_user`` runs (which also drives the SCIM gate), + plus the ``prisma_client`` / ``user_api_key_cache`` globals. The centralized gate's own + fetches fail-safe to None under the MagicMock prisma, so an unblocked user admits. Yields the + ``get_user_object`` mock so a caller can assert the sealed user_id was the reload key.""" + get_user_object = AsyncMock(return_value=return_value, side_effect=side_effect) + with ( + patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + ): + yield get_user_object + @staticmethod def _mcp_request(path="/mcp/bridge_delegate_server"): """A minimal ``Request`` for direct ``_admit_dcr_bridge_delegate`` calls, mirroring how @@ -5060,6 +5083,84 @@ class TestMCPDcrBridgeDelegateAdmission: "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} } + async def test_user_subject_envelope_admits_under_the_reloaded_user(self): + """An interactively-minted (user_id) envelope admits under the reloaded USER, not a key: the + reload is keyed by the sealed user_id, the admitted auth carries that user_id, the raw-key + pipeline is never invoked, and the inner upstream token is injected for egress. This is the + interactive-DCR admission the whole flow exists for.""" + envelope = self._mint_bridge_envelope(user_id="sso-user-7") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_user_reload( + return_value=MagicMock(user_id="sso-user-7", metadata={"scim_active": True}) + ) as get_user_object, + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + (auth_result, _h, _s, mcp_server_auth_headers, _o, _r) = await MCPRequestHandler.process_mcp_request(scope) + + assert get_user_object.await_args.kwargs["user_id"] == "sso-user-7" + assert auth_result.user_id == "sso-user-7" + mock_auth.assert_not_called() + assert mcp_server_auth_headers == { + "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} + } + + async def test_user_subject_envelope_missing_user_fails_closed_401(self): + """A user_id envelope whose user has since been deleted must fail closed: get_user_object + resolves None, so admission 401s instead of admitting an unresolved identity.""" + envelope = self._mint_bridge_envelope(user_id="ghost-user") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_user_reload(return_value=None), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + + async def test_user_subject_envelope_scim_deactivated_user_fails_closed_401(self): + """SCIM-deactivating the envelope's user revokes it immediately: the reloaded user carries + scim_active False, so admission 401s rather than letting an offboarded user keep tool access + until the envelope expires.""" + envelope = self._mint_bridge_envelope(user_id="offboarded-user") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_user_reload( + return_value=MagicMock(user_id="offboarded-user", metadata={"scim_active": False}) + ), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + async def test_revoked_key_envelope_fails_closed_401(self): """An envelope whose key has since been deleted must fail closed: ``get_key_object`` raises for the missing row, so admission 401s instead of admitting the caller as an unrestricted diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index 82e8e2aae89..ecea86bbed4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -28,13 +28,14 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import EnvelopeTooLarge, SealedEnvelope, UpstreamTokenGrant, + key_hash_identity, mint_envelope, ) _NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) _MASTER_KEY = "sk-master-key-for-derivation-tests-0123456789" _ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" -_IDENTITY = EnvelopeIdentity(server_id="srv-456", key_hash="hashed-key-123") +_IDENTITY = key_hash_identity(server_id="srv-456", key_hash="hashed-key-123") _SERVER_ID = _IDENTITY.server_id @@ -138,7 +139,7 @@ def test_resolve_envelope_minted_for_another_server_is_invalid(): captured or misrouted envelope cannot forward one server's upstream credential to another. The valid access token stays sealed; the mismatch alone fails the resolve.""" keys = envelope_keys_from_master_key(_MASTER_KEY) - other_server_identity = EnvelopeIdentity(server_id="srv-OTHER", key_hash=_IDENTITY.key_hash) + other_server_identity = key_hash_identity(server_id="srv-OTHER", key_hash=_IDENTITY.subject) token = _sealed_token(keys, identity=other_server_identity) result = resolve_bridge_envelope(token, keys, _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeInvalid) @@ -155,7 +156,7 @@ def test_resolve_non_ascii_server_id_stays_total_and_does_not_raise(): unicode server_id); it stays total and returns a typed result. A matching non-ASCII id admits, a mismatching one is BridgeEnvelopeInvalid, and neither raises.""" keys = envelope_keys_from_master_key(_MASTER_KEY) - unicode_identity = EnvelopeIdentity(server_id="srv-café", key_hash=_IDENTITY.key_hash) + unicode_identity = key_hash_identity(server_id="srv-café", key_hash=_IDENTITY.subject) token = _sealed_token(keys, identity=unicode_identity) assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-café"), BridgeEnvelopeAdmitted) assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-cafe"), BridgeEnvelopeInvalid) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py index b44f3f84cc9..7a2b51c2a95 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py @@ -36,8 +36,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import SealedEnvelope, UpstreamTokenGrant, is_envelope, + key_hash_identity, mint_envelope, open_envelope, + user_identity, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value @@ -51,7 +53,7 @@ _WRONG_SIGNING = EnvelopeKeys(signing_key=SecretStr(_OTHER_SIGNING_KEY), encrypt _WRONG_ENCRYPTION = EnvelopeKeys(signing_key=SecretStr(_SIGNING_KEY), encryption_key=SecretStr(_OTHER_ENCRYPTION_KEY)) _ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" _REFRESH_TOKEN = "upstream-refresh-token-do-not-leak-1d0aa4b7" -_IDENTITY = EnvelopeIdentity(server_id="srv-456", key_hash="hashed-key-123") +_IDENTITY = key_hash_identity(server_id="srv-456", key_hash="hashed-key-123") def _full_grant() -> UpstreamTokenGrant: @@ -137,12 +139,13 @@ def test_minimal_grant_round_trips_without_none_leakage_into_claims(): def test_claim_layout_and_no_plaintext_token_in_envelope(): token = _sealed_token(_full_grant()) claims = _unverified_claims(token) - assert set(claims) == {"iss", "iat", "exp", "server_id", "key_hash", "grant"} + assert set(claims) == {"iss", "iat", "exp", "server_id", "subject_type", "subject", "grant"} assert claims["iss"] == ENVELOPE_ISSUER assert claims["iat"] == int(_NOW.timestamp()) assert claims["exp"] == int(_NOW.timestamp()) + 600 assert claims["server_id"] == "srv-456" - assert claims["key_hash"] == "hashed-key-123" + assert claims["subject_type"] == "key_hash" + assert claims["subject"] == "hashed-key-123" assert _ACCESS_TOKEN not in token assert _ACCESS_TOKEN not in json.dumps(claims) assert _REFRESH_TOKEN not in json.dumps(claims) @@ -226,11 +229,11 @@ def test_wrong_issuer_is_malformed_payload(): def test_missing_identity_claim_is_malformed_payload(): claims = _unverified_claims(_sealed_token(_full_grant())) - forged = _forge({key: value for key, value in claims.items() if key != "key_hash"}) + forged = _forge({key: value for key, value in claims.items() if key != "subject"}) assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) -@pytest.mark.parametrize("identity_claim", ["server_id", "key_hash"]) +@pytest.mark.parametrize("identity_claim", ["server_id", "subject"]) def test_signed_empty_identity_claim_is_malformed_payload_not_a_raise(identity_claim): claims = _unverified_claims(_sealed_token(_full_grant())) forged = _forge({**claims, identity_claim: ""}) @@ -463,9 +466,11 @@ def test_non_positive_expires_in_is_rejected_at_construction_without_leaking(): def test_empty_identity_and_key_fields_are_rejected_at_construction(): with pytest.raises(ValidationError): - EnvelopeIdentity(server_id="", key_hash="hashed-key-123") + EnvelopeIdentity(server_id="", subject_type="key_hash", subject="hashed-key-123") with pytest.raises(ValidationError): - EnvelopeIdentity(server_id="srv-456", key_hash="") + EnvelopeIdentity(server_id="srv-456", subject_type="key_hash", subject="") + with pytest.raises(ValidationError): + EnvelopeIdentity(server_id="srv-456", subject_type="not-a-subject-type", subject="x") with pytest.raises(ValidationError): EnvelopeKeys(signing_key=SecretStr(""), encryption_key=SecretStr(_ENCRYPTION_KEY)) with pytest.raises(ValidationError): @@ -474,6 +479,20 @@ def test_empty_identity_and_key_fields_are_rejected_at_construction(): UpstreamTokenGrant(access_token=SecretStr(""), token_type="Bearer") +def test_user_subject_identity_round_trips(): + """The user_id subject variant seals and opens with its discriminator intact, so the edge can + tell an interactively-minted (user) envelope from a scripted (key_hash) one and reload the right + kind of record.""" + identity = user_identity(server_id="srv-456", user_id="user-42") + sealed = mint_envelope(identity, _full_grant(), _KEYS, _NOW) + assert isinstance(sealed, SealedEnvelope) + opened = open_envelope(sealed.token.get_secret_value(), _KEYS, _NOW) + assert isinstance(opened, OpenedEnvelope) + assert opened.identity.server_id == "srv-456" + assert opened.identity.subject_type == "user_id" + assert opened.identity.subject == "user-42" + + def test_public_models_are_frozen(): sealed = mint_envelope(_IDENTITY, _full_grant(), _KEYS, _NOW) assert isinstance(sealed, SealedEnvelope) @@ -484,4 +503,4 @@ def test_public_models_are_frozen(): with pytest.raises(ValidationError): opened.grant = _minimal_grant() with pytest.raises(ValidationError): - _IDENTITY.key_hash = "someone-elses-hash" + _IDENTITY.subject = "someone-elses-hash" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 68466e624ec..e43312e400e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4441,7 +4441,8 @@ async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) opened = resolve_bridge_envelope(token, keys, datetime.now(timezone.utc), server.server_id) assert isinstance(opened, BridgeEnvelopeAdmitted) - assert opened.identity.key_hash == "hashed-litellm-key-77" + assert opened.identity.subject_type == "key_hash" + assert opened.identity.subject == "hashed-litellm-key-77" assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN" From 02e9c5631a88d0bdb52d1d4ccb1de21e9c29bede Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 14:58:52 -0700 Subject: [PATCH 05/41] feat(mcp): interactive SSO sign-in for dcr_bridge oauth_delegate DCR clients Completes the oauth_delegate bridge for real DCR clients (Claude Code, Claude Desktop), which send no litellm key and cannot use the scripted two-header path. On the short-circuit bridge arm the gateway now captures the SSO-authenticated litellm user from the browser session at /authorize and seals it into the OAuth state; at /callback it seals that user plus the upstream code into a gateway authorization code the client echoes back; at /token it recovers the user, exchanges the real upstream code, and mints a user-subject envelope. The user identity captured in the browser thus rides to the back-channel token call with nothing stored server-side, and admission opens the envelope under that user. The scripted key_hash path is unchanged (raw upstream code, key from the request); without a session the browser is sent through login first. --- .../mcp_server/discoverable_endpoints.py | 176 +++++++++++++-- .../mcp_server/test_discoverable_endpoints.py | 200 +++++++++++++++++- 2 files changed, 352 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 3e727ce95bb..bd26abe0b26 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -12,7 +12,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response -from pydantic import BaseModel, SecretStr, ValidationError +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError from typing_extensions import assert_never from litellm._logging import verbose_logger @@ -41,6 +41,7 @@ from litellm.types.mcp_server.mcp_server_manager import MCPServer if TYPE_CHECKING: from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeIdentity, EnvelopeKeys, UpstreamTokenGrant, ) @@ -98,6 +99,8 @@ def encode_state_with_base_url( code_challenge: Optional[str] = None, code_challenge_method: Optional[str] = None, client_redirect_uri: Optional[str] = None, + litellm_user_id: str | None = None, + mcp_server_id: str | None = None, ) -> str: """ Encode the base_url, original state, and PKCE parameters using encryption. @@ -108,6 +111,11 @@ def encode_state_with_base_url( code_challenge: PKCE code challenge from client code_challenge_method: PKCE code challenge method from client client_redirect_uri: Original redirect_uri from client + litellm_user_id: The SSO-authenticated litellm user captured at the bridge authorize + (interactive dcr_bridge oauth_delegate only); the callback seals it into the gateway + authorization code so the token mint can bind the envelope to this user + mcp_server_id: The bridge server the interactive flow targets, sealed alongside + litellm_user_id so the gateway code cannot be replayed against another server Returns: An encrypted string that encodes all values @@ -118,6 +126,8 @@ def encode_state_with_base_url( "code_challenge": code_challenge, "code_challenge_method": code_challenge_method, "client_redirect_uri": client_redirect_uri, + "litellm_user_id": litellm_user_id, + "mcp_server_id": mcp_server_id, } state_json = json.dumps(state_data, sort_keys=True) encrypted_state = encrypt_value_helper(state_json) @@ -145,6 +155,68 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data +_BRIDGE_AUTH_CODE_PREFIX = "llm_bcode_" + + +class _BridgeAuthorizationCode(BaseModel): + """The identity and upstream code the gateway seals into the authorization code it hands a DCR + client for an interactive dcr_bridge oauth_delegate sign-in, recovered at the token endpoint.""" + + model_config = ConfigDict(frozen=True) + upstream_code: str = Field(min_length=1) + litellm_user_id: str = Field(min_length=1) + mcp_server_id: str = Field(min_length=1) + + +def is_bridge_authorization_code(code: str) -> bool: + """Cheap prefix check that ``code`` is a gateway-sealed bridge authorization code rather than a + raw upstream code, so the token endpoint can route without decrypting.""" + return code.startswith(_BRIDGE_AUTH_CODE_PREFIX) + + +def seal_bridge_authorization_code(upstream_code: str, litellm_user_id: str, mcp_server_id: str) -> str: + """Seal the upstream authorization code and the SSO-captured litellm user into a gateway + authorization code. The DCR client only echoes this opaque value back at the token endpoint; the + gateway decrypts it there to recover the user (to bind the envelope) and the upstream code (to + exchange with the upstream), so a litellm identity captured in the browser at authorize survives + to the back-channel token call with nothing stored server-side. Encrypted with the repo's + authenticated symmetric helper (the same family the OAuth state uses), so the client can neither + read nor forge it.""" + payload = json.dumps( + {"upstream_code": upstream_code, "litellm_user_id": litellm_user_id, "mcp_server_id": mcp_server_id}, + sort_keys=True, + ) + return _BRIDGE_AUTH_CODE_PREFIX + encrypt_value_helper(payload) + + +def open_bridge_authorization_code(code: str) -> _BridgeAuthorizationCode | None: + """Recover the sealed identity and upstream code, or ``None`` when ``code`` is not a gateway + bridge code or does not decrypt / validate. Total over hostile input: a raw upstream code (the + scripted two-header path) returns ``None`` and the caller falls through to the existing + behavior.""" + if not is_bridge_authorization_code(code): + return None + decrypted = decrypt_value_helper( + code[len(_BRIDGE_AUTH_CODE_PREFIX) :], "bridge_authorization_code", return_original_value=False + ) + if not isinstance(decrypted, str): + return None + try: + return _BridgeAuthorizationCode.model_validate_json(decrypted) + except ValidationError: + return None + + +def _redirect_to_litellm_login(request: Request) -> RedirectResponse: + """Send an unauthenticated browser through litellm login before the interactive bridge authorize + can capture its identity. The bridge oauth_delegate flow seals the SSO user into the gateway code, + so a session is required; without one there is nothing to bind. After login the user re-initiates + the connection, which then finds the session cookie (the seamless return-to round-trip, which is + origin-validated against the control-plane URL, is a follow-up).""" + base_url = get_request_base_url(request) + return RedirectResponse(f"{base_url}/sso/key/generate") + + # LIT-4197: some upstream authorization servers reject an over-long ``state`` # (the encrypted OAuth session blob routinely exceeds their limit). The upstream # only needs an opaque value it echoes back on ``/callback``, so we forward a @@ -697,12 +769,31 @@ async def authorize_with_server( parsed = urlparse(redirect_uri) base_url = urlunparse(parsed._replace(query="")) request_base_url = get_request_base_url(request) + + # Interactive dcr_bridge oauth_delegate sign-in: this arm runs the gateway /callback and /token in + # the loop, so the gateway can capture the litellm user here (from the browser's UI session) and + # carry it to the back-channel token mint. Seal the SSO user and the target server into the state; + # the callback reads them back to mint the gateway authorization code. A DCR client cannot present a + # litellm key, so the browser session is the only identity source; without one there is nothing to + # bind, so send the user through login first. Every other oauth2 server keeps the identity-less state. + litellm_user_id: str | None = None + if mcp_server.is_dcr_bridge and mcp_server.is_oauth_delegate: + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import + _user_id_from_session_cookie, + ) + + litellm_user_id = _user_id_from_session_cookie(request) + if litellm_user_id is None: + return _redirect_to_litellm_login(request) + encoded_state = encode_state_with_base_url( base_url=base_url, original_state=state, code_challenge=code_challenge, code_challenge_method=code_challenge_method, client_redirect_uri=redirect_uri, + litellm_user_id=litellm_user_id, + mcp_server_id=mcp_server.server_id if litellm_user_id else None, ) relay_state = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES) @@ -824,11 +915,14 @@ _BridgeMintError = Literal[ @dataclass(frozen=True, slots=True) class _BridgeMintReady: - """Everything the seal needs, resolved once before the exchange: the authorizing key hash and the - master-key-derived envelope keys. Passing this forward means identity resolution and key derivation - happen exactly once, and ``_finish_bridge_mint`` has no preconditions left that could fail.""" + """Everything the seal needs, resolved once before the exchange: the identity to bind the envelope + to and the master-key-derived envelope keys. The identity is a key_hash subject for the scripted + two-header client (resolved from the litellm key it presents) or a user_id subject for the + interactive SSO client (the user recovered from the gateway authorization code), so one phase-3 seal + serves both. Resolving identity here means ``_finish_bridge_mint`` has no preconditions left to + fail.""" - key_hash: str + identity: "EnvelopeIdentity" keys: "EnvelopeKeys" @@ -844,8 +938,8 @@ def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse: status, code, desc = ( 400, "invalid_request", - "this server issues a gateway-bound credential; send a litellm credential " - "(x-litellm-api-key or Authorization) on the token request", + "this server issues a gateway-bound credential; complete the interactive sign-in, or " + "send a litellm credential (x-litellm-api-key or Authorization) on the token request", ) case "unsupported_grant": status, code, desc = ( @@ -923,18 +1017,30 @@ def _upstream_rejection_to_mint_error(rejection: _UpstreamGrantRejection) -> _Br assert_never(rejection) -async def _prepare_bridge_mint(request: Request, grant_type: str) -> "_BridgeMintReady | _BridgeMintError": +async def _prepare_bridge_mint( + request: Request, + grant_type: str, + mcp_server: MCPServer, + bridge_identity: _BridgeAuthorizationCode | None = None, +) -> "_BridgeMintReady | _BridgeMintError": """Phase 1, BEFORE the upstream exchange: reject a grant this mint does not support, confirm the gateway can mint (master_key set), resolve the litellm identity, and derive the envelope keys. Returns a ready context or a precise failure value. Running before the exchange is what makes every - failure here fail closed without consuming the single-use code or rotating a refresh token. A bridge - server issues only envelopes and seals no upstream refresh_token, so the client holds none to - present: the refresh_token grant is rejected up front rather than exchanged (which could rotate the - upstream credential) and its result then discarded. Identity-resolution failures keep their origin - so the mapper statuses each truthfully.""" + failure here fail closed without consuming the single-use code. + + Two identity sources, one envelope. The interactive DCR client authenticates via SSO at the bridged + authorize, so its identity arrives as ``bridge_identity`` (the user recovered from the gateway + authorization code) and mints a user subject. The scripted two-header client presents a litellm key + on the token request instead, so its identity is the active key's hash and mints a key_hash subject. + A missing or invalid presented key keeps its resolution origin so the mapper statuses it truthfully; + neither source present is ``no_identity``.""" from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import envelope_keys_from_master_key, ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + key_hash_identity, + user_identity, + ) from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import master_key, ) @@ -943,16 +1049,21 @@ async def _prepare_bridge_mint(request: Request, grant_type: str) -> "_BridgeMin return "unsupported_grant" if not master_key: return "not_configured" + keys = envelope_keys_from_master_key(master_key) + if bridge_identity is not None: + identity = user_identity(server_id=mcp_server.server_id, user_id=bridge_identity.litellm_user_id) + return _BridgeMintReady(identity=identity, keys=keys) resolved = await _resolve_active_litellm_key(request) if not isinstance(resolved, _ResolvedKey): return _key_resolution_failure_to_mint_error(resolved) - return _BridgeMintReady(key_hash=resolved.key_hash, keys=envelope_keys_from_master_key(master_key)) + identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=resolved.key_hash) + return _BridgeMintReady(identity=identity, keys=keys) def _finish_bridge_mint( ready: "_BridgeMintReady", mcp_server: MCPServer, token_response: object, now: datetime ) -> "JSONResponse | _BridgeMintError": - """Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held envelope using + """Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held envelope under the pre-resolved identity and keys, so the client holds one bearer that admits it and forwards the upstream token with nothing stored server-side. The only failures here are properties of the upstream response (no usable token, an already-expired lifetime, or a token too large to seal), @@ -963,14 +1074,12 @@ def _finish_bridge_mint( from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import SealedEnvelope, UpstreamTokenGrant, - key_hash_identity, ) grant = _bridge_grant_from_token_response(token_response) if not isinstance(grant, UpstreamTokenGrant): return _upstream_rejection_to_mint_error(grant) - identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=ready.key_hash) - sealed = build_bridge_token_response(identity, grant, ready.keys, now) + sealed = build_bridge_token_response(ready.identity, grant, ready.keys, now) if not isinstance(sealed, SealedEnvelope): return "too_large" # Report expires_in from the JWT's own second-truncated exp, rounding the elapsed portion up, so the @@ -1014,6 +1123,7 @@ async def exchange_token_with_server( except TokenEndpointAuthConfigError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc + bridge_identity: _BridgeAuthorizationCode | None = None if grant_type == "refresh_token": if not refresh_token: raise HTTPException( @@ -1033,6 +1143,19 @@ async def exchange_token_with_server( status_code=400, detail="code is required for authorization_code grant", ) + # Interactive dcr_bridge oauth_delegate: the client presents the gateway authorization code the + # callback sealed. Recover the SSO user and the real upstream code from it; the upstream exchange + # below uses the upstream code, and the mint binds the envelope to the recovered user. Bind the + # sealed server to this request so a code minted for one bridge server cannot be spent at another. + # A raw upstream code (scripted path) opens to None and the code is used as-is. + bridge_identity = open_bridge_authorization_code(code) + if bridge_identity is not None: + if bridge_identity.mcp_server_id != mcp_server.server_id: + raise HTTPException( + status_code=400, + detail="Authorization code was issued for a different MCP server", + ) + code = bridge_identity.upstream_code bridge_token_relay = _dcr_bridge_relays_client_registration(mcp_server) if bridge_token_relay and not redirect_uri: raise HTTPException( @@ -1058,7 +1181,7 @@ async def exchange_token_with_server( # phase 3. A failure here returns without ever touching the upstream credential. bridge_mint_ready: _BridgeMintReady | None = None if mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge: - prepared = await _prepare_bridge_mint(request, grant_type) + prepared = await _prepare_bridge_mint(request, grant_type, mcp_server, bridge_identity) if not isinstance(prepared, _BridgeMintReady): return _bridge_mint_error_response(prepared) bridge_mint_ready = prepared @@ -1706,7 +1829,20 @@ async def callback( # states while permitting same-origin / allowlisted clients. redirect_uri = _get_validated_client_redirect_uri(request, state_data) - params = {"code": code, "state": original_state} + # Interactive dcr_bridge oauth_delegate: the state carries the litellm user the authorize step + # captured. Instead of forwarding the raw upstream code (which the client would present at the + # token endpoint with no way to prove who signed in), seal the user and the upstream code into a + # gateway authorization code and forward THAT. The token endpoint decrypts it to bind the + # envelope to this user. Every other flow forwards the raw code unchanged. + litellm_user_id = state_data.get("litellm_user_id") + mcp_server_id = state_data.get("mcp_server_id") + forwarded_code = code + if isinstance(litellm_user_id, str) and litellm_user_id and isinstance(mcp_server_id, str) and mcp_server_id: + forwarded_code = seal_bridge_authorization_code( + upstream_code=code, litellm_user_id=litellm_user_id, mcp_server_id=mcp_server_id + ) + + params = {"code": forwarded_code, "state": original_state} complete_returned_url = _append_query_params(redirect_uri, params) response = RedirectResponse(url=complete_returned_url, status_code=302) _clear_oauth_state_cookie(response, request, state) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index e43312e400e..966619ee6b8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4365,7 +4365,7 @@ async def test_register_bridge_relay_never_persists(): _BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef" -async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_client_out=None): +async def _exchange_for_bridge_server(server, upstream_body, key_hash, code="auth-code", fake_client_out=None): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _ResolvedKey, exchange_token_with_server, @@ -4398,13 +4398,17 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_clie request=_bridge_mock_request(), mcp_server=server, grant_type="authorization_code", - code="auth-code", + code=code, redirect_uri="https://claude.ai/api/mcp/auth_callback", client_id="dcr-client-123", client_secret=None, code_verifier="verifier", ) - if server.is_oauth_delegate and server.is_dcr_bridge: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import is_bridge_authorization_code + + # The key_hash path resolves the presented litellm key; the interactive SSO path recovers identity + # from the gateway authorization code instead, so it never awaits the resolver. + if server.is_oauth_delegate and server.is_dcr_bridge and not is_bridge_authorization_code(code): key_resolver.assert_awaited_once() else: key_resolver.assert_not_awaited() @@ -4446,6 +4450,193 @@ async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN" +def test_bridge_authorization_code_round_trips_and_rejects_hostile_input(): + """The gateway authorization code seals and recovers the upstream code and the SSO user, and is + total over hostile input: a raw upstream code (scripted path) opens to None, and a tampered or + non-gateway value opens to None rather than raising.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + is_bridge_authorization_code, + open_bridge_authorization_code, + seal_bridge_authorization_code, + ) + + with patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY): + sealed = seal_bridge_authorization_code( + upstream_code="up-code", litellm_user_id="sso-user-9", mcp_server_id="srv-1" + ) + assert is_bridge_authorization_code(sealed) + opened = open_bridge_authorization_code(sealed) + assert opened is not None + assert opened.upstream_code == "up-code" + assert opened.litellm_user_id == "sso-user-9" + assert opened.mcp_server_id == "srv-1" + assert open_bridge_authorization_code("raw-upstream-code") is None + assert open_bridge_authorization_code(sealed[:-4] + "aaaa") is None + + +@pytest.mark.asyncio +async def test_interactive_bridge_token_exchange_mints_user_subject_envelope(): + """An interactive dcr_bridge oauth_delegate exchange (the client presents the gateway code the + callback sealed, and NO litellm key) mints an envelope bound to the SSO-captured user: it opens + to a user_id subject, and the upstream exchange used the real upstream code recovered from the + gateway code, not the sealed wrapper.""" + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_bridge_authorization_code, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + BridgeEnvelopeAdmitted, + envelope_keys_from_master_key, + resolve_bridge_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + with patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY): + gateway_code = seal_bridge_authorization_code( + upstream_code="REAL-UPSTREAM-CODE", litellm_user_id="sso-user-42", mcp_server_id=server.server_id + ) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + captured: dict = {} + response = await _exchange_for_bridge_server( + server, upstream, key_hash=None, code=gateway_code, fake_client_out=captured + ) + + token = json.loads(response.body)["access_token"] + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = resolve_bridge_envelope(token, keys, datetime.now(timezone.utc), server.server_id) + assert isinstance(opened, BridgeEnvelopeAdmitted) + assert opened.identity.subject_type == "user_id" + assert opened.identity.subject == "sso-user-42" + assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN" + assert captured["client"].post.call_args.kwargs["data"]["code"] == "REAL-UPSTREAM-CODE" + + +@pytest.mark.asyncio +async def test_interactive_bridge_gateway_code_for_another_server_is_rejected_400(): + """A gateway authorization code is bound to the server it was minted for: presenting it at another + server's token endpoint is a 400, so a code cannot be replayed across a server boundary.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_bridge_authorization_code, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + with patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY): + gateway_code = seal_bridge_authorization_code( + upstream_code="up-code", litellm_user_id="sso-user-42", mcp_server_id="a-different-server-id" + ) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + with pytest.raises(HTTPException) as exc: + await _exchange_for_bridge_server(server, upstream, key_hash=None, code=gateway_code) + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_interactive_bridge_authorize_seals_sso_user_into_state(): + """On the short-circuit bridge oauth_delegate arm, authorize captures the SSO user from the UI + session cookie and seals it (and the target server) into the encrypted OAuth state, so the + callback can later mint a user-bound gateway code; it still proceeds to the upstream redirect.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import authorize_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate, client_id="admin-client", registration_url=None) + captured: dict = {} + + def _capture(**kwargs): + captured.update(kwargs) + return "mocked_encrypted_state" + + with ( + patch( + "litellm.proxy._experimental.mcp_server.byok_oauth_endpoints._user_id_from_session_cookie", + return_value="sso-user-42", + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encode_state_with_base_url", + side_effect=_capture, + ), + ): + response = await authorize_with_server( + request=_bridge_mock_request(), + mcp_server=server, + client_id="ignored", + redirect_uri="http://127.0.0.1:60108/callback", + state="s", + code_challenge="chal", + code_challenge_method="S256", + ) + + assert captured["litellm_user_id"] == "sso-user-42" + assert captured["mcp_server_id"] == server.server_id + assert "/sso/key/generate" not in response.headers["location"] + + +@pytest.mark.asyncio +async def test_interactive_bridge_authorize_without_session_redirects_to_login(): + """Without a UI session there is no identity to bind, so the short-circuit bridge oauth_delegate + authorize sends the browser through litellm login instead of proceeding to the upstream.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import authorize_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate, client_id="admin-client", registration_url=None) + with patch( + "litellm.proxy._experimental.mcp_server.byok_oauth_endpoints._user_id_from_session_cookie", + return_value=None, + ): + response = await authorize_with_server( + request=_bridge_mock_request(), + mcp_server=server, + client_id="ignored", + redirect_uri="http://127.0.0.1:60108/callback", + state="s", + code_challenge="chal", + code_challenge_method="S256", + ) + assert "/sso/key/generate" in response.headers["location"] + + +@pytest.mark.asyncio +async def test_interactive_bridge_callback_seals_user_into_gateway_code(): + """When the OAuth state carries the captured SSO user, the callback forwards a gateway + authorization code (sealing the user and upstream code) to the client instead of the raw upstream + code, so the client's later token call can prove who signed in.""" + from urllib.parse import parse_qs, urlparse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + is_bridge_authorization_code, + ) + + state_data = { + "original_state": "client-state", + "client_redirect_uri": "http://127.0.0.1:60108/cb", + "base_url": "http://127.0.0.1:60108/cb", + "litellm_user_id": "sso-user-42", + "mcp_server_id": "bridge_srv", + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_encoded_oauth_state", + return_value="enc", + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash", + return_value=state_data, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._get_validated_client_redirect_uri", + return_value="http://127.0.0.1:60108/cb", + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await callback(request=_bridge_mock_request(), code="REAL-UPSTREAM-CODE", state="relay") + + forwarded_code = parse_qs(urlparse(response.headers["location"]).query)["code"][0] + assert is_bridge_authorization_code(forwarded_code) + + @pytest.mark.asyncio async def test_oauth_delegate_bridge_token_exchange_fails_closed_without_litellm_identity(): """Without a resolvable litellm identity on the token request, the exchange must not mint an @@ -4727,10 +4918,11 @@ def test_bridge_reported_expires_in_can_be_zero_at_jwt_exp_boundary(): from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( envelope_keys_from_master_key, ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity from litellm.types.mcp import MCPAuth ready = _BridgeMintReady( - key_hash="hashed-litellm-key-77", + identity=key_hash_identity(server_id="bridge_srv", key_hash="hashed-litellm-key-77"), keys=envelope_keys_from_master_key(_BRIDGE_MASTER_KEY), ) response = _finish_bridge_mint( From f96899ae2b793f049a01a91db648980cd85021d2 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 13 Jul 2026 11:08:08 -0700 Subject: [PATCH 06/41] fix(mcp): classify the user-subject reload's errors like the key path (503 outage, 401 missing) _reload_admitted_user mirrored only part of _reload_admitted_key's error contract: it caught ProxyException and HTTPException but had no arm for anything else, so a transient DB outage surfaced as an opaque 500 instead of the retryable 503 the key path guarantees, and a missing user surfaced as a 500 too. The missing-user case is the subtle one: get_user_object raises a bare Exception for a deleted user (not a ProxyException like get_key_object does for a missing key), so the ProxyException/HTTPException clause never caught it and the user_object-is-None branch it was supposed to hit is unreachable on the production path. Add the same except-Exception arm the key path uses, with the one deliberate difference the differing get_user_object contract requires: a database-service-unavailable error still raises the retryable 503, while a missing user or any other non-outage resolution failure fails closed as a 401 rather than propagating as a 500. The regression tests now drive the real behavior (get_user_object raising) rather than a None return that never happens in production, and cover both the 503 outage and the 401 missing-user paths. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 13 +++++-- .../auth/test_user_api_key_auth_mcp.py | 34 +++++++++++++++---- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index faec35db41a..01861abbbc0 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -605,8 +605,14 @@ class MCPRequestHandler: caller's centralized policy gate then enforces the user's live budget and org state, and a SCIM-deactivated owner fails closed here exactly as the key path enforces it. No team is bound; a user may belong to many teams or none, so the envelope grants the - user's own access rather than silently selecting one team's scope. A missing user - fails closed with a 401 rather than admitting an unresolved identity.""" + user's own access rather than silently selecting one team's scope. + + Error handling mirrors the key path's retryable-503 contract, with one deliberate + difference: ``get_key_object`` raises a ``ProxyException`` for a missing key, but + ``get_user_object`` raises a bare ``Exception`` for a missing user (it does not surface as a + ``ProxyException``/``HTTPException``). So a transient DB outage still surfaces as a retryable + 503 via ``_raise_503_if_db_unavailable``, while a missing user, or any other non-outage + resolution failure, fails closed as a 401 rather than propagating as an opaque 500.""" from litellm.proxy.auth.auth_checks import get_user_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -621,6 +627,9 @@ class MCPRequestHandler: ) except (ProxyException, HTTPException): raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + except Exception as e: # noqa: BLE001 # DB outage -> retryable 503; a missing user (bare Exception) or any other resolution failure -> fail closed 401, never an opaque 500 + MCPRequestHandler._raise_503_if_db_unavailable(e) + raise HTTPException(status_code=401, detail="Invalid or expired credential") from None if user_object is None: raise HTTPException(status_code=401, detail="Invalid or expired credential") if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index a6affe5496c..a441e6a154b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -5117,8 +5117,10 @@ class TestMCPDcrBridgeDelegateAdmission: } async def test_user_subject_envelope_missing_user_fails_closed_401(self): - """A user_id envelope whose user has since been deleted must fail closed: get_user_object - resolves None, so admission 401s instead of admitting an unresolved identity.""" + """A user_id envelope whose user has since been deleted must fail closed with a 401, not a 500. + get_user_object raises a bare Exception for a missing user (it does not return None on the + production path), so the reload must catch it and fail closed rather than let it propagate as an + opaque 500. Regression for the missing-user path surfacing as a 500.""" envelope = self._mint_bridge_envelope(user_id="ghost-user") scope = { "type": "http", @@ -5129,7 +5131,7 @@ class TestMCPDcrBridgeDelegateAdmission: with ( patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), - self._patch_user_reload(return_value=None), + self._patch_user_reload(side_effect=Exception("user not found")), ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() with pytest.raises(HTTPException) as exc_info: @@ -5137,6 +5139,28 @@ class TestMCPDcrBridgeDelegateAdmission: assert exc_info.value.status_code == 401 + async def test_user_subject_envelope_db_outage_is_retryable_503(self): + """A transient database outage while reloading the envelope's user is a retryable 503, not an + opaque 500, matching the key path's contract so an interactive DCR client retries instead of + treating a live identity as invalid. Regression for the user reload dropping the 503 arm.""" + envelope = self._mint_bridge_envelope(user_id="sso-user-7") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_user_reload(side_effect=ConnectionError("auth database unreachable")), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 503 + async def test_user_subject_envelope_scim_deactivated_user_fails_closed_401(self): """SCIM-deactivating the envelope's user revokes it immediately: the reloaded user carries scim_active False, so admission 401s rather than letting an offboarded user keep tool access @@ -5151,9 +5175,7 @@ class TestMCPDcrBridgeDelegateAdmission: with ( patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), - self._patch_user_reload( - return_value=MagicMock(user_id="offboarded-user", metadata={"scim_active": False}) - ), + self._patch_user_reload(return_value=MagicMock(user_id="offboarded-user", metadata={"scim_active": False})), ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() with pytest.raises(HTTPException) as exc_info: From c46863b0e64a4962b84ddf41dc1a9faf7faac3dd Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 13 Jul 2026 11:18:53 -0700 Subject: [PATCH 07/41] fix(mcp): admit a user-subject envelope with the user's own MCP object permission _reload_admitted_user returned a bare UserAPIKeyAuth(user_id=...), so the shared get_allowed_mcp_servers found no key/team/object-permission grants and an interactive SSO client could admit successfully yet see zero tools on a normal (allow_all_keys=False) server. The key path returns the full key record whose object permission drives that computation; the user path dropped it. Resolve the user's own MCP object permission and put it on the returned auth, so the same get_allowed_mcp_servers the key path uses grants the user their litellm-granted servers and access groups. This reuses get_object_permission (the id-to-grants resolver keys and teams already use) and does not duplicate any permission logic; get_user_object does not load object_permission, so it is resolved from the user's object_permission_id the same way the key and team paths do. Only the user's own object permission is bound. A UserAPIKeyAuth carries a single team_id while a user may belong to many teams, so team-inherited MCP grants for a user are a follow-up: they need a many-teams union get_allowed_mcp_servers does not do off one auth object, and faking one here would be the kind of half-measure that spawns more bugs. Tests cover the user's object permission riding onto the admitted auth, and the existing admit/SCIM/missing-user/503 cases still hold. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 32 ++++++++++--- .../auth/test_user_api_key_auth_mcp.py | 46 ++++++++++++++++++- 2 files changed, 70 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 01861abbbc0..ac3e4439fa1 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -601,11 +601,14 @@ class MCPRequestHandler: The DCR client authenticates via SSO at the bridged authorize, which yields a user subject rather than a virtual key, so the envelope admits under the user's own - identity: the reloaded ``user_id`` rides on the returned ``UserAPIKeyAuth`` and the - caller's centralized policy gate then enforces the user's live budget and org state, - and a SCIM-deactivated owner fails closed here exactly as the key path enforces it. No - team is bound; a user may belong to many teams or none, so the envelope grants the - user's own access rather than silently selecting one team's scope. + identity: the reloaded ``user_id`` and the user's own MCP object permission ride on the + returned ``UserAPIKeyAuth``, and the SAME ``get_allowed_mcp_servers`` the key path uses then + computes which servers the user may reach, so the user's litellm MCP grants and access groups + gate the request exactly as a key's do. Only the user's OWN object permission is bound: a + ``UserAPIKeyAuth`` carries a single ``team_id`` while a user may belong to many teams, so + team-inherited MCP grants for a user are a follow-up (they need a many-teams union + ``get_allowed_mcp_servers`` does not do off one auth object). The caller's centralized policy + gate enforces the user's live budget and org state, and a SCIM-deactivated owner fails closed. Error handling mirrors the key path's retryable-503 contract, with one deliberate difference: ``get_key_object`` raises a ``ProxyException`` for a missing key, but @@ -613,7 +616,7 @@ class MCPRequestHandler: ``ProxyException``/``HTTPException``). So a transient DB outage still surfaces as a retryable 503 via ``_raise_503_if_db_unavailable``, while a missing user, or any other non-outage resolution failure, fails closed as a 401 rather than propagating as an opaque 500.""" - from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.auth.auth_checks import get_object_permission, get_user_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache if prisma_client is None: @@ -634,7 +637,22 @@ class MCPRequestHandler: raise HTTPException(status_code=401, detail="Invalid or expired credential") if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: raise HTTPException(status_code=401, detail="Invalid or expired credential") - return UserAPIKeyAuth(user_id=user_object.user_id) + # Resolve the user's own MCP object permission (get_user_object does not load it) so the shared + # get_allowed_mcp_servers can grant the user their litellm-granted servers. Reuses the same + # get_object_permission resolver the key and team paths use; no permission logic is duplicated. + object_permission = user_object.object_permission + if user_object.object_permission_id and object_permission is None: + object_permission = await get_object_permission( + object_permission_id=user_object.object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return UserAPIKeyAuth( + user_id=user_object.user_id, + user_role=user_object.user_role, + object_permission=object_permission, + object_permission_id=user_object.object_permission_id, + ) @staticmethod async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index a441e6a154b..0b1905689ac 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -5103,7 +5103,13 @@ class TestMCPDcrBridgeDelegateAdmission: patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), self._patch_user_reload( - return_value=MagicMock(user_id="sso-user-7", metadata={"scim_active": True}) + return_value=MagicMock( + user_id="sso-user-7", + metadata={"scim_active": True}, + user_role=None, + object_permission=None, + object_permission_id=None, + ) ) as get_user_object, ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() @@ -5116,6 +5122,44 @@ class TestMCPDcrBridgeDelegateAdmission: "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} } + async def test_user_subject_envelope_carries_the_users_mcp_object_permission(self): + """The admitted user's own MCP object permission rides on the returned auth so the shared + get_allowed_mcp_servers grants the user their litellm-granted servers, rather than admitting a + bare user with no MCP access. Regression for the signed-in SSO client getting zero tools because + the reload dropped the user's object permission.""" + object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="op-user-7", mcp_servers=["bridge_delegate_server"] + ) + envelope = self._mint_bridge_envelope(user_id="sso-user-7") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ), + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_user_reload( + return_value=MagicMock( + user_id="sso-user-7", + metadata={"scim_active": True}, + user_role=None, + object_permission=object_permission, + object_permission_id="op-user-7", + ) + ), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + (auth_result, _h, _s, _headers, _o, _r) = await MCPRequestHandler.process_mcp_request(scope) + + assert auth_result.object_permission is not None + assert auth_result.object_permission.mcp_servers == ["bridge_delegate_server"] + async def test_user_subject_envelope_missing_user_fails_closed_401(self): """A user_id envelope whose user has since been deleted must fail closed with a 401, not a 500. get_user_object raises a bare Exception for a missing user (it does not return None on the From 7fce761cdeca0ec8431e1b8feb1b8192df92521a Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 13 Jul 2026 12:19:38 -0700 Subject: [PATCH 08/41] fix(ui): respect litellm_key_header_name in BYOK credential save and workflow runs fetches (#33103) --- ui/litellm-dashboard/eslint-suppressions.json | 299 +++++++++--------- .../mcp-servers/_components/mcp_servers.tsx | 1 - .../playground/components/chat_ui/ChatUI.tsx | 1 - .../workflows/WorkflowRuns.test.tsx | 19 +- .../(dashboard)/workflows/WorkflowRuns.tsx | 8 +- .../mcp_tools/ByokCredentialModal.test.tsx | 72 +++++ .../mcp_tools/ByokCredentialModal.tsx | 37 +-- .../src/components/networking.test.ts | 64 ++++ .../src/components/networking.tsx | 3 + 9 files changed, 324 insertions(+), 180 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.test.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index ba2e2a375ca..953de4fe480 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -515,6 +515,152 @@ "count": 2 } }, + "src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.test.tsx": { + "unused-imports/no-unused-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx": { + "react-hooks/immutability": { + "count": 2 + } + }, + "src/app/(dashboard)/mcp-servers/_components/MCPSubmissionsTab.tsx": { + "react-hooks/set-state-in-effect": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 1 + }, + "unused-imports/no-unused-imports": { + "count": 2 + } + }, + "src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/OpenAPIQuickPicker.tsx": { + "react-hooks/set-state-in-effect": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.tsx": { + "no-nested-ternary": { + "count": 3 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx": { + "no-nested-ternary": { + "count": 2 + } + }, + "src/app/(dashboard)/mcp-servers/_components/create_mcp_server.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 4 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx": { + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/static-components": { + "count": 4 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_connection_status.tsx": { + "no-nested-ternary": { + "count": 3 + }, + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_discovery.tsx": { + "react-hooks/set-state-in-effect": { + "count": 2 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_config.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_display.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/immutability": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 5 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 2 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/mcp_tools.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 2 + } + }, "src/app/(dashboard)/memory/_components/MemoryView.tsx": { "react-hooks/set-state-in-effect": { "count": 1 @@ -1860,45 +2006,11 @@ "count": 1 } }, - "src/components/mcp_tools/ByokCredentialModal.tsx": { - "no-restricted-syntax": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.test.tsx": { - "unused-imports/no-unused-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx": { - "react-hooks/immutability": { - "count": 2 - } - }, - "src/app/(dashboard)/mcp-servers/_components/MCPSubmissionsTab.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/mcp_tools/MCPToolArgumentsForm.tsx": { "no-nested-ternary": { "count": 5 } }, - "src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - }, - "unused-imports/no-unused-imports": { - "count": 2 - } - }, "src/components/mcp_tools/McpCrudPermissionPanel.tsx": { "no-nested-ternary": { "count": 3 @@ -1907,123 +2019,6 @@ "count": 1 } }, - "src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/OpenAPIQuickPicker.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.tsx": { - "no-nested-ternary": { - "count": 3 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx": { - "no-nested-ternary": { - "count": 2 - } - }, - "src/app/(dashboard)/mcp-servers/_components/create_mcp_server.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 4 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/static-components": { - "count": 4 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_connection_status.tsx": { - "no-nested-ternary": { - "count": 3 - }, - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_discovery.tsx": { - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_config.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_display.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/immutability": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 5 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/mcp-servers/_components/mcp_tools.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, "src/components/model_add/AddCredentialModal.tsx": { "no-restricted-imports": { "count": 1 @@ -2545,4 +2540,4 @@ "count": 1 } } -} +} \ No newline at end of file diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx index 0afb4bd9314..f186fef22da 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx @@ -684,7 +684,6 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) refetch(); setByokModalServer(null); }} - accessToken={accessToken || ""} /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index 689abf66b41..5133fafb4a7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -2186,7 +2186,6 @@ const ChatUI: React.FC = ({ loadMCPServers(); setByokModalServer(null); }} - accessToken={accessToken || ""} /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.test.tsx index e73abfe7cd6..1aee3fcc8ab 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.test.tsx @@ -4,7 +4,10 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import WorkflowRuns from "./WorkflowRuns"; -vi.mock("@/components/networking", () => ({ proxyBaseUrl: "" })); +vi.mock("@/components/networking", () => ({ + proxyBaseUrl: "", + getGlobalLitellmHeaderName: () => "x-litellm-api-key", +})); interface FakeRun { run_id: string; @@ -78,4 +81,18 @@ describe("WorkflowRuns (migrated onto shared DataTable)", () => { expect(await screen.findByText("No workflow runs yet")).toBeInTheDocument(); }); + + it("sends the configured litellm key header on every fetch instead of hardcoding Authorization", async () => { + const user = userEvent.setup(); + const fetchSpy = mockFetch(RUNS); + vi.stubGlobal("fetch", fetchSpy); + render(); + + await user.click(await screen.findByText("First run")); + + await waitFor(() => expect(fetchSpy).toHaveBeenCalledTimes(3)); + for (const [url, init] of fetchSpy.mock.calls as [string, RequestInit][]) { + expect(init.headers, url).toEqual({ "x-litellm-api-key": "Bearer tok" }); + } + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.tsx index 9afa07251c2..7354b7479f4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/workflows/WorkflowRuns.tsx @@ -2,7 +2,7 @@ import React, { useState, useEffect, useCallback, useMemo } from "react"; import { Button, Collapse, Drawer, Empty, Spin, Tooltip, Typography } from "antd"; import { ReloadOutlined } from "@ant-design/icons"; import type { ColumnDef, ColumnFiltersState } from "@tanstack/react-table"; -import { proxyBaseUrl } from "@/components/networking"; +import { getGlobalLitellmHeaderName, proxyBaseUrl } from "@/components/networking"; import { DataTable, DataTableFilterDrawer, @@ -507,7 +507,7 @@ const WorkflowRuns: React.FC = ({ accessToken }) => { setLoadingRuns(true); try { const res = await fetch(`${proxyBaseUrl ?? ""}/v1/workflows/runs?limit=100`, { - headers: { Authorization: `Bearer ${accessToken}` }, + headers: { [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}` }, }); if (!res.ok) throw new Error(`HTTP ${res.status}`); const data = await res.json(); @@ -531,10 +531,10 @@ const WorkflowRuns: React.FC = ({ accessToken }) => { const base = proxyBaseUrl ?? ""; const [evRes, msgRes] = await Promise.all([ fetch(`${base}/v1/workflows/runs/${run.run_id}/events`, { - headers: { Authorization: `Bearer ${accessToken}` }, + headers: { [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}` }, }), fetch(`${base}/v1/workflows/runs/${run.run_id}/messages`, { - headers: { Authorization: `Bearer ${accessToken}` }, + headers: { [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}` }, }), ]); const evData = evRes.ok ? await evRes.json() : { events: [] }; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.test.tsx new file mode 100644 index 00000000000..021aec5f85f --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.test.tsx @@ -0,0 +1,72 @@ +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { registerAuthHeaderNameGetter, registerAuthTokenGetter, registerBaseUrlGetter } from "@/lib/http/runtime"; +import { ByokCredentialModal } from "./ByokCredentialModal"; +import type { MCPServer } from "./types"; + +const fetchSpy = vi.hoisted(() => { + const spy = vi.fn<(request: Request) => Promise>(); + vi.stubGlobal("fetch", spy); + return spy; +}); + +vi.mock("@/components/molecules/message_manager", () => ({ + default: { success: vi.fn(), error: vi.fn() }, +})); + +const SERVER = { server_id: "srv-1", alias: "Linear", server_name: "Linear" } as MCPServer; + +const jsonResponse = (body: unknown, status = 200) => + new Response(JSON.stringify(body), { status, headers: { "Content-Type": "application/json" } }); + +async function fillAndSubmit(user: ReturnType) { + await user.click(screen.getByText("Continue to Authentication")); + await user.type(screen.getByPlaceholderText("Enter your API key"), "linear-key"); + await user.click(screen.getByRole("button", { name: /Connect & Authorize/ })); +} + +beforeEach(() => { + fetchSpy.mockReset(); + registerBaseUrlGetter(() => ""); + registerAuthTokenGetter(() => "sk-session"); +}); + +describe("ByokCredentialModal", () => { + it("saves the credential with the session's configured litellm key header, not a hardcoded Authorization", async () => { + registerAuthHeaderNameGetter(() => "x-litellm-api-key"); + fetchSpy.mockResolvedValue(jsonResponse({ server_id: "srv-1", has_credential: true })); + const onSuccess = vi.fn(); + const user = userEvent.setup(); + render( {}} onSuccess={onSuccess} />); + + await fillAndSubmit(user); + + await waitFor(() => expect(onSuccess).toHaveBeenCalledWith("srv-1")); + const request = fetchSpy.mock.calls[0][0]; + expect(request.method).toBe("POST"); + expect(new URL(request.url).pathname).toBe("/v1/mcp/server/srv-1/user-credential"); + expect(request.headers.get("x-litellm-api-key")).toBe("Bearer sk-session"); + expect(request.headers.get("Authorization")).toBeNull(); + expect(await request.json()).toEqual({ credential: "linear-key", save: true }); + }); + + it("surfaces the backend's detail.error message when the save fails", async () => { + registerAuthHeaderNameGetter(() => "Authorization"); + fetchSpy.mockResolvedValue( + jsonResponse({ detail: { error: "This MCP server does not support BYOK credentials" } }, 400), + ); + const MessageManager = (await import("@/components/molecules/message_manager")).default; + const onSuccess = vi.fn(); + const user = userEvent.setup(); + render( {}} onSuccess={onSuccess} />); + + await fillAndSubmit(user); + + await waitFor(() => + expect(MessageManager.error).toHaveBeenCalledWith("This MCP server does not support BYOK credentials"), + ); + expect(onSuccess).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx index cb07db871fd..f36de019aa5 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx @@ -3,6 +3,8 @@ import React, { useState } from "react"; import { Modal, Input, Switch } from "antd"; import MessageManager from "@/components/molecules/message_manager"; +import { fetchClient } from "@/lib/http/api"; +import { ApiError } from "@/lib/http/client"; import { KeyOutlined, LockOutlined, @@ -14,21 +16,22 @@ import { } from "@ant-design/icons"; import { MCPServer } from "./types"; +const byokSaveErrorMessage = (e: unknown): string => { + if (e instanceof ApiError) { + const detail = (e.body as { detail?: { error?: string } } | null)?.detail?.error; + if (detail) return detail; + } + return e instanceof Error && e.message ? e.message : "Failed to connect"; +}; + interface ByokCredentialModalProps { server: MCPServer; open: boolean; onClose: () => void; onSuccess: (serverId: string) => void; - accessToken: string; } -export const ByokCredentialModal: React.FC = ({ - server, - open, - onClose, - onSuccess, - accessToken, -}) => { +export const ByokCredentialModal: React.FC = ({ server, open, onClose, onSuccess }) => { const [step, setStep] = useState<1 | 2>(1); const [apiKey, setApiKey] = useState(""); const [saveKey, setSaveKey] = useState(true); @@ -52,23 +55,15 @@ export const ByokCredentialModal: React.FC = ({ } setLoading(true); try { - const response = await fetch(`/v1/mcp/server/${server.server_id}/user-credential`, { - method: "POST", - headers: { - "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, - }, - body: JSON.stringify({ credential: apiKey.trim(), save: saveKey }), + await fetchClient.POST("/v1/mcp/server/{server_id}/user-credential", { + params: { path: { server_id: server.server_id } }, + body: { credential: apiKey.trim(), save: saveKey }, }); - if (!response.ok) { - const err = await response.json(); - throw new Error(err?.detail?.error || "Failed to save credential"); - } MessageManager.success(`Connected to ${serverDisplayName}`); onSuccess(server.server_id); handleClose(); - } catch (e: any) { - MessageManager.error(e.message || "Failed to connect"); + } catch (e) { + MessageManager.error(byokSaveErrorMessage(e)); } finally { setLoading(false); } diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index b9d04a00f61..e6ee4de2735 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -530,3 +530,67 @@ describe("buildModelGroupTestRequest", () => { expect(body).toEqual({ model: "text-embedding-3-small", input: "test from litellm" }); }); }); + +describe("testMCPToolsListRequest auth headers", () => { + const originalFetch = global.fetch; + + const captureFetch = () => { + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + status: 200, + headers: { get: () => "application/json" }, + json: vi.fn().mockResolvedValue({ tools: [] }), + } as any); + global.fetch = mockFetch as any; + return mockFetch; + }; + + const sentHeaders = (mockFetch: ReturnType): Record => + (mockFetch.mock.calls[0][1] as RequestInit).headers as Record; + + afterEach(() => { + Networking.setGlobalLitellmHeaderName("Authorization"); + global.fetch = originalFetch; + }); + + it("sends the litellm key under a custom litellm_key_header_name even when an upstream OAuth token uses Authorization", async () => { + Networking.setGlobalLitellmHeaderName("x-litellm-key"); + const mockFetch = captureFetch(); + + await Networking.testMCPToolsListRequest("sk-key", {}, "upstream-oauth-token"); + + const headers = sentHeaders(mockFetch); + expect(headers["x-litellm-key"]).toBe("Bearer sk-key"); + expect(headers["Authorization"]).toBe("Bearer upstream-oauth-token"); + }); + + it("Bearer-prefixes x-litellm-api-key when it is the configured key header (raw values fail _get_bearer_token)", async () => { + Networking.setGlobalLitellmHeaderName("x-litellm-api-key"); + const mockFetch = captureFetch(); + + await Networking.testMCPToolsListRequest("sk-key", {}, "upstream-oauth-token"); + + const headers = sentHeaders(mockFetch); + expect(headers["x-litellm-api-key"]).toBe("Bearer sk-key"); + expect(headers["Authorization"]).toBe("Bearer upstream-oauth-token"); + }); + + it("never clobbers the upstream OAuth token on default deployments", async () => { + const mockFetch = captureFetch(); + + await Networking.testMCPToolsListRequest("sk-key", {}, "upstream-oauth-token"); + + const headers = sentHeaders(mockFetch); + expect(headers["Authorization"]).toBe("Bearer upstream-oauth-token"); + expect(headers["x-litellm-api-key"]).toBe("sk-key"); + }); + + it("sends the litellm key as the bearer on default deployments without an OAuth token", async () => { + const mockFetch = captureFetch(); + + await Networking.testMCPToolsListRequest("sk-key", {}); + + const headers = sentHeaders(mockFetch); + expect(headers["Authorization"]).toBe("Bearer sk-key"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 10d7f12604d..875a345a4c5 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6650,6 +6650,9 @@ export const testMCPToolsListRequest = async ( }; if (accessToken) { headers["x-litellm-api-key"] = accessToken; + if (globalLitellmHeaderName.toLowerCase() !== "authorization") { + headers[globalLitellmHeaderName] = `Bearer ${accessToken}`; + } } if (oauthAccessToken) { headers["Authorization"] = `Bearer ${oauthAccessToken}`; From aa9dcb43cfaee139ad515ea2a66eab0c31f04a6c Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 13 Jul 2026 12:19:57 -0700 Subject: [PATCH 09/41] refactor(ui): standardize debounce waits behind shared DEBOUNCE_WAIT_MS constant (#33040) --- .../usage/_components/components/UsagePageView.tsx | 3 ++- .../src/app/(dashboard)/users/_components/view_users.tsx | 3 ++- .../PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx | 4 ++-- .../ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx | 4 ++-- .../src/components/VirtualKeysPage/VirtualKeysTable.tsx | 3 ++- .../src/components/common_components/team_dropdown.tsx | 4 ++-- .../src/components/common_components/team_multi_select.tsx | 4 ++-- .../src/components/team/TeamVirtualKeysTable.tsx | 3 ++- ui/litellm-dashboard/src/utils/debounceConstants.ts | 1 + 9 files changed, 17 insertions(+), 12 deletions(-) create mode 100644 ui/litellm-dashboard/src/utils/debounceConstants.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index a95d48aa75b..d2f75609d18 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -8,6 +8,7 @@ import { DownOutlined, ExportOutlined, InfoCircleOutlined, LoadingOutlined, RightOutlined } from "@ant-design/icons"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { Card, Col, @@ -94,7 +95,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { // Debounced search for user selector const [userSearchInput, setUserSearchInput] = useState(""); const [debouncedUserSearch, setDebouncedUserSearch] = useDebouncedState("", { - wait: 300, + wait: DEBOUNCE_WAIT_MS, }); const { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx index 18709e05df8..db3b17d6af3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx @@ -16,6 +16,7 @@ import { import OnboardingModal, { InvitationLink } from "@/components/onboarding_link"; import { updateExistingKeys } from "@/utils/dataUtils"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { isAdminRole, isProxyAdminRole } from "@/utils/roles"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; import { useQuery, useQueryClient } from "@tanstack/react-query"; @@ -86,7 +87,7 @@ const ViewUserDashboard: React.FC = ({ const [userToDelete, setUserToDelete] = useState(null); const [activeTab, setActiveTab] = useState("users"); const [filters, setFilters] = useState(initialFilters); - const [debouncedFilters, setDebouncedFilters, debouncer] = useDebouncedState(filters, { wait: 300 }); + const [debouncedFilters, setDebouncedFilters, debouncer] = useDebouncedState(filters, { wait: DEBOUNCE_WAIT_MS }); const [isInvitationLinkModalVisible, setIsInvitationLinkModalVisible] = useState(false); const [invitationLinkData, setInvitationLinkData] = useState(null); const [baseUrl, setBaseUrl] = useState(null); diff --git a/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx b/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx index d42d5ab324a..1d19ba3255d 100644 --- a/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx +++ b/ui/litellm-dashboard/src/components/KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect.tsx @@ -1,4 +1,5 @@ import { useInfiniteKeyAliases } from "@/app/(dashboard)/hooks/keys/useKeyAliases"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { LoadingOutlined } from "@ant-design/icons"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; import { Select } from "antd"; @@ -16,7 +17,6 @@ export interface PaginatedKeyAliasSelectProps { } const SCROLL_THRESHOLD = 0.8; -const DEBOUNCE_MS = 300; export const PaginatedKeyAliasSelect = ({ value, @@ -30,7 +30,7 @@ export const PaginatedKeyAliasSelect = ({ }: PaginatedKeyAliasSelectProps) => { const [searchInput, setSearchInput] = useState(""); const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { - wait: DEBOUNCE_MS, + wait: DEBOUNCE_WAIT_MS, }); const teamId = allFilters?.["Team ID"] || undefined; diff --git a/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx b/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx index da3ecf77ded..a77b2cf561e 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx @@ -1,4 +1,5 @@ import { useInfiniteModelInfo } from "@/app/(dashboard)/hooks/models/useModels"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { LoadingOutlined } from "@ant-design/icons"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; import { Select, Space, Typography } from "antd"; @@ -17,7 +18,6 @@ export interface PaginatedModelSelectProps { } const SCROLL_THRESHOLD = 0.8; -const DEBOUNCE_MS = 300; export const PaginatedModelSelect = ({ value, @@ -30,7 +30,7 @@ export const PaginatedModelSelect = ({ }: PaginatedModelSelectProps) => { const [searchInput, setSearchInput] = useState(""); const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { - wait: DEBOUNCE_MS, + wait: DEBOUNCE_WAIT_MS, }); const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteModelInfo( diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index cae6dc54df5..697318af62f 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -4,6 +4,7 @@ import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrgan import { useAllTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; import { ColumnDef, @@ -64,7 +65,7 @@ export function VirtualKeysTable() { pageSize: 50, }); const [filters, setFilters] = useState(DEFAULT_KEY_FILTERS); - const [debouncedFilters] = useDebouncedValue(filters, { wait: 300 }); + const [debouncedFilters] = useDebouncedValue(filters, { wait: DEBOUNCE_WAIT_MS }); const sortBy = sorting.length > 0 ? sorting[0].id : null; const sortOrder = sorting.length > 0 ? (sorting[0].desc ? "desc" : "asc") : null; diff --git a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx index 84140135db7..7d27886c7f5 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx @@ -3,6 +3,7 @@ import { Select, Typography } from "antd"; import { LoadingOutlined } from "@ant-design/icons"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; import { useInfiniteTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { Team } from "../key_team_helpers/key_list"; const { Text } = Typography; @@ -19,7 +20,6 @@ interface TeamDropdownProps { } const SCROLL_THRESHOLD = 0.8; -const DEBOUNCE_MS = 300; const TeamDropdown: React.FC = ({ value, @@ -31,7 +31,7 @@ const TeamDropdown: React.FC = ({ }) => { const [searchInput, setSearchInput] = useState(""); const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { - wait: DEBOUNCE_MS, + wait: DEBOUNCE_WAIT_MS, }); const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteTeams( diff --git a/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx b/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx index d91f83c589b..3a48b5f7b50 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx @@ -3,6 +3,7 @@ import { Select, Typography } from "antd"; import { LoadingOutlined } from "@ant-design/icons"; import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; import { useInfiniteTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { Team } from "../key_team_helpers/key_list"; const { Text } = Typography; @@ -17,7 +18,6 @@ interface TeamMultiSelectProps { } const SCROLL_THRESHOLD = 0.8; -const DEBOUNCE_MS = 300; const TeamMultiSelect: React.FC = ({ value = [], @@ -29,7 +29,7 @@ const TeamMultiSelect: React.FC = ({ }) => { const [searchInput, setSearchInput] = useState(""); const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { - wait: DEBOUNCE_MS, + wait: DEBOUNCE_WAIT_MS, }); const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteTeams( diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index 207e6f2ccfe..eeccc7482e3 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -9,6 +9,7 @@ import { DataTableToolbar, } from "@/components/shared/DataTable"; import { Input } from "@/components/ui/input"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { ColumnDef, ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; @@ -43,7 +44,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi const [columnFilters, setColumnFilters] = useState([]); const [filtersOpen, setFiltersOpen] = useState(false); const [searchInput, setSearchInput] = useState(""); - const [searchQuery] = useDebouncedValue(searchInput, { wait: 300 }); + const [searchQuery] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); const handleSearchChange = useCallback((value: string) => { setSearchInput(value); diff --git a/ui/litellm-dashboard/src/utils/debounceConstants.ts b/ui/litellm-dashboard/src/utils/debounceConstants.ts new file mode 100644 index 00000000000..bb8a1e4014c --- /dev/null +++ b/ui/litellm-dashboard/src/utils/debounceConstants.ts @@ -0,0 +1 @@ +export const DEBOUNCE_WAIT_MS = 300; From fa09cde3c09b68354e7f11b4654d15ff77f088cc Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 13 Jul 2026 12:49:01 -0700 Subject: [PATCH 10/41] feat(ui): rebuild the Virtual Keys table on the shared DataTable (#32991) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(ui): rebuild the Virtual Keys table on the shared DataTable Replaces the hand-rolled Tremor table and bespoke toolbar/pagination on the admin Virtual Keys page with the shared DataTable: server-side sort, paginate, and filter, a sticky scrolling body, a search plus column-visibility plus filters toolbar, a right-side filter drawer, and a rows-per-page footer. A page header with the existing key icon carries the Create New Key action. Adds reusable, shadcn-default building blocks for the tables migrating onto the DataTable next: shared IdentityCell, ModelsCell, and SpendBudgetCell in shared/table_cells, plus a shared PageHeader. The models cell reveals overflow in a hover tooltip and the spend/budget cell uses the Meter primitive. All data and domain logic is preserved, including the useKeys query, team and org alias resolution, the user popover, and the KeyInfoView detail swap. The rich async Team/Org/Alias filters move into the drawer, and the toolbar search maps to the key-alias substring search. Status now also reflects key expiry alongside blocked and SCIM-blocked. The VirtualKeysTable tests are updated to the new markup and extended with focused coverage for each new shared cell * fix(ui): address Virtual Keys redesign review feedback Fold the status badge into the clickable Key cell and drop the separate Status column so a key's alias, secret, and status read as one unit. The Key cell is now the single click target that opens the key detail; the whole-row click is removed Migrate the filter drawer off AntD to shadcn. A new Combobox composed from Popover and Input backs the Team, Organization, and Key Alias filters, keeping search and the alias infinite-scroll Show $0.00 for zero spend instead of a hyphen, and extend the shared DataTable with badge, chips, and meter skeleton shapes so the loading state matches the loaded cells (status pill, model chips, spend meter) rather than uniform bars Fix key sorting: the Key column sent its column id "key" as sort_by, which /key/list rejects with 400. It now sorts by the backend field key_alias * fix(ui): use the shadcn base combobox and refine the keys filters and skeletons Replace the hand-rolled filter combobox with the supported shadcn Base UI combobox (ui/combobox, added via the CLI and reused through a small SearchSelect wrapper). Its vended input-group and textarea deps are written for React 19 (plain functions with ref-as-prop); this app is on React 18, where those subcomponents drop the refs Base UI passes for focus and anchoring, so InputGroupInput, InputGroupButton, and ComboboxTrigger are adapted to forwardRef. Those ui/ files now diverge from the registry, and a future shadcn add would overwrite the adaptation until the app moves to React 19. Adds class-variance-authority, which input-group needs Give loading skeletons a per-column renderSkeleton escape hatch on the shared DataTable and mirror the Key cell exactly (alias line, secret, status pill), so skeleton rows match the real rows instead of being shorter and simpler Resolve the automated review: the toolbar search and the drawer Key Alias filter both mapped to the key-alias query, so the search silently overrode the drawer value while its chip stayed visible. Consolidate to a single alias search in the toolbar (placeholder now "Search by key alias…") and drop the redundant drawer field. Re-add coverage for the Created By column's alias-over-email display Refine the Team and Organization filters: they match on name and id, so the labels read "Team" and "Organization" rather than "... ID", each option shows the name with the id on a muted second line instead of "name (id)", and the active-filter chip shows the friendly name * chore(ui): drop duplicate class-variance-authority, use the repo cva package in input-group --- ui/litellm-dashboard/eslint-suppressions.json | 8 - .../VirtualKeysPage/VirtualKeysTable.test.tsx | 205 ++-- .../VirtualKeysPage/VirtualKeysTable.tsx | 970 ++++-------------- .../VirtualKeysPage/keyTableColumns.tsx | 353 +++++++ .../shared/DataTable/DataTable.test.tsx | 32 + .../components/shared/DataTable/DataTable.tsx | 26 +- .../components/shared/DataTable/columnMeta.ts | 3 + .../src/components/shared/DataTable/types.ts | 2 +- .../src/components/shared/PageHeader.test.tsx | 31 + .../src/components/shared/PageHeader.tsx | 29 + .../components/shared/SearchSelect.test.tsx | 64 ++ .../src/components/shared/SearchSelect.tsx | 76 ++ .../shared/table_cells/identity_cell.test.tsx | 38 + .../shared/table_cells/identity_cell.tsx | 48 + .../components/shared/table_cells/index.ts | 3 + .../shared/table_cells/models_cell.test.tsx | 45 + .../shared/table_cells/models_cell.tsx | 56 + .../table_cells/spend_budget_cell.test.tsx | 53 + .../shared/table_cells/spend_budget_cell.tsx | 44 + .../src/components/ui/combobox.tsx | 266 +++++ .../src/components/ui/input-group.tsx | 140 +++ .../src/components/ui/textarea.tsx | 18 + .../src/components/user_dashboard.tsx | 27 +- 23 files changed, 1647 insertions(+), 890 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/PageHeader.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/SearchSelect.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/combobox.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/input-group.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/textarea.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 953de4fe480..7c6df4e7bb9 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1690,14 +1690,6 @@ "count": 1 } }, - "src/components/VirtualKeysPage/VirtualKeysTable.tsx": { - "no-nested-ternary": { - "count": 2 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/components/activity_metrics.tsx": { "no-nested-ternary": { "count": 1 diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 02f5d588149..cf4beeed40f 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -63,7 +63,7 @@ const mockKey: KeyResponse = { key_alias: "Test Key Alias", spend: 5.5, max_budget: 100, - expires: "2024-12-31T23:59:59Z", + expires: "2999-12-31T23:59:59Z", models: ["gpt-3.5-turbo", "gpt-4"], aliases: {}, config: {}, @@ -154,6 +154,8 @@ const keysResult = (keys: KeyResponse[], data: Partial = {}, extra ...extra, }) as any; +const openFilters = () => fireEvent.click(screen.getByRole("button", { name: "Filters" })); + beforeEach(() => { vi.clearAllMocks(); @@ -170,6 +172,12 @@ it("should render VirtualKeysTable component", () => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); +it("renders the page header with the create-key action slot", () => { + renderWithProviders(Create New Key} />); + expect(screen.getByRole("heading", { name: "Virtual Keys" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Create New Key" })).toBeInTheDocument(); +}); + it("should display key information correctly", async () => { renderWithProviders(); @@ -177,6 +185,7 @@ it("should display key information correctly", async () => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); expect(screen.getByText("Test Team")).toBeInTheDocument(); expect(screen.getByText("$5.5000")).toBeInTheDocument(); + expect(screen.getByText("of $100")).toBeInTheDocument(); }); }); @@ -188,14 +197,49 @@ it("should display user email correctly", async () => { }); }); -it("should show loading message only on initial load (isPending)", () => { +it("shows the user alias over the email in the visible cell when both exist", async () => { + mockUseKeys.mockReturnValue( + keysResult([{ ...mockKey, user: { user_id: "user-1", user_email: "user@example.com", user_alias: "The User" } }]), + ); + + renderWithProviders(); + + const row = (await screen.findByText("Test Key Alias")).closest("tr") as HTMLElement; + expect(within(row).getByText("The User")).toBeInTheDocument(); + expect(within(row).queryByText("user@example.com")).not.toBeInTheDocument(); +}); + +it("shows created_by_user alias over email in the Created By column when it is enabled", async () => { + mockUseKeys.mockReturnValue( + keysResult([ + { + ...mockKey, + created_by: "some-uuid", + created_by_user: { user_id: "some-uuid", user_email: "creator@example.com", user_alias: "The Creator" }, + }, + ]), + ); + const user = userEvent.setup(); + renderWithProviders(); + + // Created By is hidden by default; turn it on via the Columns menu. + await user.click(screen.getByRole("button", { name: "Columns" })); + await user.click(await screen.findByText("Created By")); + await user.keyboard("{Escape}"); + + const row = (await screen.findByText("Test Key Alias")).closest("tr") as HTMLElement; + expect(within(row).getByText("The Creator")).toBeInTheDocument(); + expect(within(row).queryByText("creator@example.com")).not.toBeInTheDocument(); +}); + +it("should show a loading state on the initial load and hide the data", () => { mockUseKeys.mockReturnValue(keysResult([], {}, { data: null, isPending: true, isFetching: true })); renderWithProviders(); - expect(screen.getByText("🚅 Loading keys...")).toBeInTheDocument(); + expect(screen.getByText("Loading keys...")).toBeInTheDocument(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); - expect(screen.queryByText("Test Team")).not.toBeInTheDocument(); }); it("should show 'No keys found' message when the key list is empty", () => { @@ -206,61 +250,52 @@ it("should show 'No keys found' message when the key list is empty", () => { expect(screen.getByText("No keys found")).toBeInTheDocument(); }); -it("should handle models with more than 3 entries to trigger expansion UI", () => { +it("collapses models beyond the visible limit into a '+N more' badge", () => { mockUseKeys.mockReturnValue( keysResult([{ ...mockKey, models: ["gpt-3.5-turbo", "gpt-4", "gpt-4-turbo", "claude-3", "claude-3-5-sonnet"] }]), ); renderWithProviders(); - expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); + expect(screen.getByText("+2 more")).toBeInTheDocument(); }); -it("should render table headers correctly", () => { +it("should render the redesigned table headers", () => { renderWithProviders(); - expect(screen.getByText("Key ID")).toBeInTheDocument(); - expect(screen.getByText("Key Alias")).toBeInTheDocument(); + expect(screen.getByText("Key")).toBeInTheDocument(); expect(screen.getByText("Team")).toBeInTheDocument(); expect(screen.getByText("Models")).toBeInTheDocument(); - expect(screen.getByText("Spend (USD)")).toBeInTheDocument(); + expect(screen.getByText("Spend / Budget")).toBeInTheDocument(); }); -it("should handle column resizing hover events", () => { +it("sorts by the backend key_alias field (not the column label) when the Key header is clicked", async () => { renderWithProviders(); - const headerCell = document.querySelector("[data-header-id]") as HTMLElement; - expect(headerCell).toBeInTheDocument(); + const keyHeader = screen.getByText("Key").closest("button") as HTMLElement; + fireEvent.click(keyHeader); - const resizer = headerCell?.querySelector(".resizer") as HTMLElement; - expect(resizer).toBeInTheDocument(); - expect(resizer.style.opacity).toBe("0"); - - fireEvent.mouseEnter(headerCell); - expect(resizer.style.opacity).toBe("0.5"); - - fireEvent.mouseLeave(headerCell); - expect(resizer.style.opacity).toBe("0"); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: "key_alias" })); + }); }); -it("should open KeyInfoView when clicking on a key ID button", async () => { +it("should open KeyInfoView when clicking the key cell", async () => { renderWithProviders(); await waitFor(() => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); - expect(screen.getByText(/Showing.*results/)).toBeInTheDocument(); + expect(screen.getByTestId("pagination-range")).toBeInTheDocument(); - const keyIdButton = screen.getByText("sk-1234567890abcdef"); - fireEvent.click(keyIdButton); + fireEvent.click(screen.getByText("Test Key Alias")); await waitFor(() => { expect(screen.getByText("Back to Keys")).toBeInTheDocument(); - expect(screen.getByText("Created At")).toBeInTheDocument(); }); - expect(screen.queryByText(/Showing.*results/)).not.toBeInTheDocument(); + expect(screen.queryByTestId("pagination-range")).not.toBeInTheDocument(); }); it("should display 'Default Proxy Admin' for user_id when value is 'default_user_id'", async () => { @@ -282,44 +317,6 @@ it("should display 'Default Proxy Admin' for user_id when value is 'default_user }); }); -it("should display created_by_user email in 'Created By' column when available", async () => { - mockUseKeys.mockReturnValue( - keysResult([ - { - ...mockKey, - created_by: "some-uuid-1234", - created_by_user: { user_id: "some-uuid-1234", user_email: "creator@example.com", user_alias: null }, - }, - ]), - ); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("creator@example.com")).toBeInTheDocument(); - }); -}); - -it("should display created_by_user alias over email when both are available", async () => { - mockUseKeys.mockReturnValue( - keysResult([ - { - ...mockKey, - created_by: "some-uuid-1234", - created_by_user: { user_id: "some-uuid-1234", user_email: "creator@example.com", user_alias: "The Creator" }, - }, - ]), - ); - - renderWithProviders(); - - // Scope to the key's row so we assert the visible cell value: the hover popover that - // also holds the email is portaled out of the row, not the displayed "Created By" text. - const row = (await screen.findByText("Test Key Alias")).closest("tr") as HTMLElement; - expect(within(row).getByText("The Creator")).toBeInTheDocument(); - expect(within(row).queryByText("creator@example.com")).not.toBeInTheDocument(); -}); - it("should render table without crashing when models is null", async () => { mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, models: null as unknown as string[] }])); @@ -327,6 +324,7 @@ it("should render table without crashing when models is null", async () => { await waitFor(() => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); }); }); @@ -341,13 +339,14 @@ it("should display 'Unknown' for last_active when value is null", async () => { }); describe("server-side filtering – the LIT-4080 regression guard", () => { - it("threads an active User ID filter into the useKeys query so any refetch keeps it", async () => { + it("threads an applied User ID filter into the useKeys query so any refetch keeps it", async () => { renderWithProviders(); - fireEvent.click(screen.getByRole("button", { name: "Filters" })); + openFilters(); - const userIdInput = await screen.findByPlaceholderText("Enter User ID..."); + const userIdInput = await screen.findByPlaceholderText(/Enter User ID/); fireEvent.change(userIdInput, { target: { value: "user-42" } }); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); await waitFor(() => { expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ userID: "user-42" })); @@ -361,18 +360,19 @@ describe("server-side filtering – the LIT-4080 regression guard", () => { expect(lastCall[2] ?? {}).toMatchObject({ userID: undefined, teamID: undefined, keyHash: undefined }); }); - it("drops the filter from the useKeys query when Reset Filters is clicked", async () => { + it("drops the filter from the useKeys query when it is cleared", async () => { renderWithProviders(); - fireEvent.click(screen.getByRole("button", { name: "Filters" })); - const userIdInput = await screen.findByPlaceholderText("Enter User ID..."); + openFilters(); + const userIdInput = await screen.findByPlaceholderText(/Enter User ID/); fireEvent.change(userIdInput, { target: { value: "user-42" } }); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); await waitFor(() => { expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ userID: "user-42" })); }); - fireEvent.click(screen.getByRole("button", { name: "Reset Filters" })); + fireEvent.click(screen.getByTestId("datatable-clear-filters")); await waitFor(() => { const lastCall = mockUseKeys.mock.calls[mockUseKeys.mock.calls.length - 1]; @@ -388,8 +388,8 @@ describe("pagination display – total count comes from useKeys", () => { renderWithProviders(); await waitFor(() => { - expect(screen.getByText("Showing 1 - 50 of 509 results")).toBeInTheDocument(); - expect(screen.getByText("Page 1 of 11")).toBeInTheDocument(); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-50 of 509"); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 11"); }); }); @@ -399,57 +399,44 @@ describe("pagination display – total count comes from useKeys", () => { renderWithProviders(); await waitFor(() => { - expect(screen.getByText("Showing 1 - 1 of 1 results")).toBeInTheDocument(); - expect(screen.getByText("Page 1 of 1")).toBeInTheDocument(); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-1 of 1"); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 1"); }); }); }); -describe("refetch button", () => { - it("should show Fetch button in normal state", () => { +describe("refresh button", () => { + it("renders an enabled refresh control in the normal state", () => { renderWithProviders(); - const fetchButton = screen.getByTitle("Fetch data"); - expect(fetchButton).toBeInTheDocument(); - expect(fetchButton).not.toBeDisabled(); - expect(screen.getByText("Fetch")).toBeInTheDocument(); + const refresh = screen.getByTestId("datatable-refresh"); + expect(refresh).toBeInTheDocument(); + expect(refresh).not.toBeDisabled(); }); - it("should show Fetching state and keep table data visible during refetch", () => { + it("disables the refresh control while a fetch is in flight but keeps data visible", () => { mockUseKeys.mockReturnValue(keysResult([mockKey], {}, { isFetching: true })); renderWithProviders(); - expect(screen.getByText("Fetching")).toBeInTheDocument(); - expect(screen.getByTitle("Fetch data")).toBeDisabled(); + expect(screen.getByTestId("datatable-refresh")).toBeDisabled(); expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); - expect(screen.queryByText("🚅 Loading keys...")).not.toBeInTheDocument(); }); - it("should call refetch when Fetch button is clicked", () => { + it("calls refetch when the refresh control is clicked", () => { const mockRefetch = vi.fn(); mockUseKeys.mockReturnValue(keysResult([mockKey], {}, { refetch: mockRefetch })); renderWithProviders(); - fireEvent.click(screen.getByTitle("Fetch data")); + fireEvent.click(screen.getByTestId("datatable-refresh")); expect(mockRefetch).toHaveBeenCalledTimes(1); }); - - it("should show Fetch button enabled on error so user can retry", () => { - mockUseKeys.mockReturnValue(keysResult([], {}, { data: null, isError: true })); - - renderWithProviders(); - - const fetchButton = screen.getByTitle("Fetch data"); - expect(fetchButton).not.toBeDisabled(); - expect(screen.getByText("Fetch")).toBeInTheDocument(); - }); }); -describe("Status column reflects key.blocked / scim_blocked metadata", () => { - it("should render Active for a non-blocked key", async () => { +describe("Status column reflects blocked / expiry / scim metadata", () => { + it("renders Active for a non-blocked, unexpired key", async () => { mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, blocked: false, metadata: {} }])); renderWithProviders(); @@ -459,7 +446,19 @@ describe("Status column reflects key.blocked / scim_blocked metadata", () => { }); }); - it("should render Blocked when key.blocked is true", async () => { + it("renders Expired when the expiry date has passed", async () => { + mockUseKeys.mockReturnValue( + keysResult([{ ...mockKey, blocked: false, metadata: {}, expires: "2020-01-01T00:00:00Z" }]), + ); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByTestId(`key-status-${mockKey.token_id}`)).toHaveTextContent("Expired"); + }); + }); + + it("renders Blocked when key.blocked is true", async () => { mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, blocked: true, metadata: {} }])); renderWithProviders(); @@ -470,7 +469,7 @@ describe("Status column reflects key.blocked / scim_blocked metadata", () => { expect(screen.queryByText(/Blocked by SCIM/i)).not.toBeInTheDocument(); }); - it("should mark a SCIM-blocked key with the SCIM tooltip reason", async () => { + it("marks a SCIM-blocked key with the SCIM tooltip reason", async () => { mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, blocked: true, metadata: { scim_blocked: true } }])); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 697318af62f..112b6cbcba4 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -1,811 +1,251 @@ "use client"; -import { useKeys, KeyListCallOptions } from "@/app/(dashboard)/hooks/keys/useKeys"; + +import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useAllTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; -import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; -import { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; import { - ColumnDef, - flexRender, - getCoreRowModel, - PaginationState, - SortingState, - useReactTable, -} from "@tanstack/react-table"; -import { Badge, Icon, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Text } from "@tremor/react"; -import { InfoCircleOutlined, SyncOutlined } from "@ant-design/icons"; -import { Button as AntButton, Popover, Skeleton, Typography } from "antd"; -import { DateCell, IdCell, MoneyCell, StatusBadge } from "@/components/shared/table_cells"; -import React, { useDeferredValue, useMemo, useState } from "react"; -import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; -import { PaginatedKeyAliasSelect } from "../KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect"; + DataTable, + DataTableFilterDrawer, + DataTableFilterField, + DataTableToolbar, +} from "@/components/shared/DataTable"; +import { SearchSelect } from "@/components/shared/SearchSelect"; +import { PageHeader } from "@/components/shared/PageHeader"; +import { Input } from "@/components/ui/input"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; +import { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; +import { KeyRound } from "lucide-react"; +import React, { useCallback, useMemo, useState } from "react"; + import { KeyResponse, Team } from "../key_team_helpers/key_list"; -import FilterComponent, { FilterOption } from "../molecules/filter"; -import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; import KeyInfoView from "../templates/key_info_view"; +import { getKeyTableColumns, KEY_TABLE_HIDDEN_COLUMNS } from "./keyTableColumns"; -type KeyFilterState = { - "Team ID": string; - "Organization ID": string; - "Key Alias": string; - "User ID": string; - "Key Hash": string; +interface VirtualKeysTableProps { + headerActions?: React.ReactNode; +} + +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +const toSortOrder = (sorting: SortingState): "asc" | "desc" | undefined => { + const active = sorting[0]; + if (!active) return undefined; + return active.desc ? "desc" : "asc"; }; -const DEFAULT_KEY_FILTERS: KeyFilterState = { - "Team ID": "", - "Organization ID": "", - "Key Alias": "", - "User ID": "", - "Key Hash": "", +const FILTER_LABELS: Record = { + team_id: "Team", + org_id: "Organization", + user_id: "User ID", + key_hash: "Key ID", }; -type KeyListFilterOptions = Pick< - KeyListCallOptions, - "teamID" | "organizationID" | "selectedKeyAlias" | "userID" | "keyHash" ->; +export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { + const { data: fetchedOrganizations } = useOrganizations(); + const organizations = useMemo(() => fetchedOrganizations ?? [], [fetchedOrganizations]); + const { data: fetchedTeams } = useAllTeams(); + const allTeams = useMemo(() => fetchedTeams ?? [], [fetchedTeams]); -const toKeyListFilters = (filters: KeyFilterState): KeyListFilterOptions => ({ - teamID: filters["Team ID"].trim() || undefined, - organizationID: filters["Organization ID"].trim() || undefined, - selectedKeyAlias: filters["Key Alias"].trim() || undefined, - userID: filters["User ID"].trim() || undefined, - keyHash: filters["Key Hash"].trim() || undefined, -}); - -export function VirtualKeysTable() { - const { data: fetchedOrganizations, isLoading: isOrgsLoading } = useOrganizations(); - const resolvedOrganizations = useMemo(() => fetchedOrganizations ?? [], [fetchedOrganizations]); const [selectedKey, setSelectedKey] = useState(null); - const [sorting, setSorting] = React.useState([{ id: "created_at", desc: true }]); - const [tablePagination, setTablePagination] = React.useState({ - pageIndex: 0, - pageSize: 50, - }); - const [filters, setFilters] = useState(DEFAULT_KEY_FILTERS); - const [debouncedFilters] = useDebouncedValue(filters, { wait: DEBOUNCE_WAIT_MS }); + const [sorting, setSorting] = useState(DEFAULT_SORTING); + const [tablePagination, setTablePagination] = useState({ pageIndex: 0, pageSize: 50 }); + const [columnFilters, setColumnFilters] = useState([]); + const [filtersOpen, setFiltersOpen] = useState(false); + const [searchInput, setSearchInput] = useState(""); + const [searchQuery] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); - const sortBy = sorting.length > 0 ? sorting[0].id : null; - const sortOrder = sorting.length > 0 ? (sorting[0].desc ? "desc" : "asc") : null; + const getFilterValue = useCallback( + (columnId: string): string | undefined => { + const entry = columnFilters.find((filter) => filter.id === columnId); + return typeof entry?.value === "string" && entry.value.trim() ? entry.value.trim() : undefined; + }, + [columnFilters], + ); + + const sortBy = sorting[0]?.id; + const sortOrder = toSortOrder(sorting); + + const keyListOptions = { + teamID: getFilterValue("team_id"), + organizationID: getFilterValue("org_id"), + selectedKeyAlias: searchQuery.trim() || undefined, + userID: getFilterValue("user_id"), + keyHash: getFilterValue("key_hash"), + sortBy, + sortOrder, + expand: "user", + }; const { data: keys, isPending: isLoading, isFetching, - isError, refetch, - } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, { - ...toKeyListFilters(debouncedFilters), - sortBy: sortBy || undefined, - sortOrder: sortOrder || undefined, - expand: "user", - }); - const [expandedAccordions, setExpandedAccordions] = useState>({}); + } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, keyListOptions); const keyList = useMemo(() => keys?.keys ?? [], [keys]); + const rowCount = keys?.total_count ?? 0; - const { data: fetchedTeams, isLoading: isTeamsLoading } = useAllTeams(); - const allTeams = useMemo(() => fetchedTeams ?? [], [fetchedTeams]); - - // Defer the transition so the button stays in loading state until the table - // has rendered with the new data (mirrors the spend-logs pattern) - const isFetchingDeferred = useDeferredValue(isFetching); - const isButtonLoading = (isFetching || isFetchingDeferred) && !isError; - - const handleRefresh = () => { - refetch(); - }; - - const handleFilterChange = (newFilters: Record) => { - setFilters({ - "Team ID": newFilters["Team ID"] || "", - "Organization ID": newFilters["Organization ID"] || "", - "Key Alias": newFilters["Key Alias"] || "", - "User ID": newFilters["User ID"] || "", - "Key Hash": newFilters["Key Hash"] || "", - }); + const handleSearchChange = useCallback((value: string) => { + setSearchInput(value); setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); - }; + }, []); - const handleFilterReset = () => { - setFilters(DEFAULT_KEY_FILTERS); + const handleSortingChange = useCallback>((updaterOrValue) => { + setSorting(updaterOrValue); setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); - }; + }, []); - const totalCount = keys?.total_count ?? 0; + const handleColumnFiltersChange = useCallback>((updaterOrValue) => { + setColumnFilters(updaterOrValue); + setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); - const columns: ColumnDef[] = useMemo( - () => [ - { - id: "expander", - header: () => null, - size: 40, - enableSorting: false, - cell: ({ row }) => - row.getCanExpand() ? ( - - ) : null, - }, - { - id: "token", - accessorKey: "token", - header: "Key ID", - size: 100, - enableSorting: true, - cell: (info) => setSelectedKey(info.row.original)} />, - }, - { - id: "key_alias", - accessorKey: "key_alias", - header: "Key Alias", - size: 150, - enableSorting: true, - cell: (info) => { - const value = info.getValue() as string; - const width = info.cell.column.getSize(); - return ( - - {value ?? "-"} - - ); - }, - }, - { - id: "status", - header: "Status", - size: 100, - enableSorting: false, - cell: ({ row }) => { - const key = row.original; - if (key.blocked !== true) { - return ; - } - const isScimBlocked = (key.metadata as Record | null | undefined)?.scim_blocked === true; - const reason = isScimBlocked - ? "Blocked by SCIM (external identity provider deactivated or deleted the owning user)." - : "Blocked. Requests using this key will be rejected with 401."; - return ( - - ); - }, - }, - { - id: "key_name", - accessorKey: "key_name", - header: "Secret Key", - size: 120, - enableSorting: false, - cell: (info) => {info.getValue() as string}, - }, - { - id: "team_alias", - accessorKey: "team_id", - header: "Team", - size: 120, - enableSorting: false, - cell: (info) => { - const teamId = info.getValue() as string | null; - if (!teamId) return "-"; - const team = allTeams.find((t) => t.team_id === teamId); - const displayValue = team?.team_alias || teamId; - const width = info.cell.column.getSize(); - return ( - - {displayValue} - - ); - }, - }, - { - id: "organization_alias", - accessorKey: "org_id", - header: "Organization", - size: 140, - enableSorting: false, - cell: (info) => { - const orgId = info.getValue() as string | null; - if (!orgId) return "-"; - const org = resolvedOrganizations.find((o) => o.organization_id === orgId); - const displayValue = org?.organization_alias || orgId; - const width = info.cell.column.getSize(); - return ( - - {displayValue} - - ); - }, - }, - { - id: "user", - accessorKey: "user", - header: () => ( - - User - - - - - ), - size: 160, - enableSorting: false, - cell: ({ row }) => { - const key = row.original; - const userAlias = key.user?.user_alias ?? null; - const userEmail = key.user?.user_email ?? key.user_email ?? null; - const userId = key.user_id ?? null; - const isDefaultAdmin = userId === "default_user_id"; - const displayValue = userAlias || userEmail || userId; - const width = 160; - - const popoverContent = ( -
- {[ - { label: "User Alias", value: userAlias }, - { label: "User Email", value: userEmail }, - { label: "User ID", value: userId }, - ].map(({ label, value }) => ( -
- {label} - {value ? ( - - {value} - - ) : ( - - - )} -
- ))} -
- ); - - if (isDefaultAdmin && !userAlias && !userEmail) { - return ( - - - - - - ); - } - - return ( - - - {displayValue || "-"} - - - ); - }, - }, - { - id: "created_at", - accessorKey: "created_at", - header: "Created At", - size: 120, - enableSorting: true, - cell: (info) => , - }, - { - id: "created_by", - accessorKey: "created_by", - header: "Created By", - size: 160, - enableSorting: false, - cell: (info) => { - const userId = info.getValue() as string | null; - if (!userId) return "-"; - const key = info.row.original; - const createdByUser = key.created_by_user; - const userAlias = createdByUser?.user_alias ?? null; - const userEmail = createdByUser?.user_email ?? null; - const isDefaultAdmin = userId === "default_user_id"; - const displayValue = userAlias || userEmail || userId; - const width = 160; - - const popoverContent = ( -
- {[ - { label: "User Alias", value: userAlias }, - { label: "User Email", value: userEmail }, - { label: "User ID", value: userId }, - ].map(({ label, value }) => ( -
- {label} - {value ? ( - - {value} - - ) : ( - - - )} -
- ))} -
- ); - - if (isDefaultAdmin && !userAlias && !userEmail) { - return ( - - - - - - ); - } - - return ( - - - {displayValue} - - - ); - }, - }, - { - id: "updated_at", - accessorKey: "updated_at", - header: "Updated At", - size: 120, - enableSorting: true, - cell: (info) => , - }, - { - id: "last_active", - accessorKey: "last_active", - header: () => ( - - Last Active - - - - - ), - size: 130, - enableSorting: false, - cell: (info) => , - }, - { - id: "expires", - accessorKey: "expires", - header: "Expires", - size: 120, - enableSorting: false, - cell: (info) => , - }, - { - id: "spend", - accessorKey: "spend", - header: "Spend (USD)", - size: 100, - enableSorting: true, - cell: (info) => , - }, - { - id: "max_budget", - accessorKey: "max_budget", - header: "Budget (USD)", - size: 110, - enableSorting: true, - cell: (info) => { - const maxBudget = info.getValue() as number | null; - if (maxBudget !== null) { - return `$${formatNumberWithCommas(maxBudget)}`; - } - const teamId = info.row.original.team_id; - const team = allTeams.find((t) => t.team_id === teamId); - if (team?.max_budget != null) { - return `$${formatNumberWithCommas(team.max_budget)} (Team)`; - } - return "Unlimited"; - }, - }, - { - id: "budget_reset_at", - accessorKey: "budget_reset_at", - header: "Budget Reset", - size: 130, - enableSorting: false, - cell: (info) => , - }, - { - id: "models", - accessorKey: "models", - header: "Models", - size: 200, - enableSorting: false, - cell: (info) => { - const models = info.getValue() as string[]; - return ( -
- {Array.isArray(models) ? ( -
- {models.length === 0 ? ( - - All Proxy Models - - ) : ( - <> -
- {models.length > 3 && ( -
- { - setExpandedAccordions((prev) => ({ - ...prev, - [info.row.id]: !prev[info.row.id], - })); - }} - /> -
- )} -
- {models.slice(0, 3).map((model, index) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} - {models.length > 3 && !expandedAccordions[info.row.id] && ( - - - +{models.length - 3} {models.length - 3 === 1 ? "more model" : "more models"} - - - )} - {expandedAccordions[info.row.id] && ( -
- {models.slice(3).map((model, index) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} -
- )} -
-
- - )} -
- ) : null} -
- ); - }, - }, - { - id: "rate_limits", - header: "Rate Limits", - size: 140, - enableSorting: false, - cell: ({ row }) => { - const key = row.original; - return ( -
-
TPM: {key.tpm_limit !== null ? key.tpm_limit : "Unlimited"}
-
RPM: {key.rpm_limit !== null ? key.rpm_limit : "Unlimited"}
-
- ); - }, - }, - ], - [allTeams, resolvedOrganizations], + const columns = useMemo( + () => getKeyTableColumns({ allTeams, organizations, onSelectKey: setSelectedKey }), + [allTeams, organizations], ); - const filterOptions: FilterOption[] = [ - { - name: "Team ID", - label: "Team ID", - isSearchable: true, - loading: isTeamsLoading, - searchFn: async (searchText: string) => { - if (!allTeams || allTeams.length === 0) return []; + const teamOptions = useMemo( + () => + allTeams.map((team) => ({ + label: team.team_alias || team.team_id, + value: team.team_id, + sublabel: team.team_alias ? team.team_id : undefined, + })), + [allTeams], + ); - const filteredTeams = allTeams.filter( - (team) => - team.team_id.toLowerCase().includes(searchText.toLowerCase()) || - (team.team_alias && team.team_alias.toLowerCase().includes(searchText.toLowerCase())), - ); + const orgOptions = useMemo( + () => + organizations + .filter((org) => org.organization_id) + .map((org) => { + const id = org.organization_id as string; + return { label: org.organization_alias || id, value: id, sublabel: org.organization_alias ? id : undefined }; + }), + [organizations], + ); - return filteredTeams.map((team) => ({ - label: `${team.team_alias || team.team_id} (${team.team_id})`, - value: team.team_id, - })); - }, + const formatFilterValue = useCallback( + (columnId: string, value: unknown): string => { + const raw = String(value); + if (columnId === "team_id") { + return allTeams.find((team) => team.team_id === raw)?.team_alias || raw; + } + if (columnId === "org_id") { + return organizations.find((org) => org.organization_id === raw)?.organization_alias || raw; + } + return raw; }, - { - name: "Organization ID", - label: "Organization ID", - isSearchable: true, - loading: isOrgsLoading, - searchFn: async (searchText: string) => { - if (!resolvedOrganizations || resolvedOrganizations.length === 0) return []; + [allTeams, organizations], + ); - const filteredOrgs = resolvedOrganizations.filter( - (org) => org.organization_id?.toLowerCase().includes(searchText.toLowerCase()) ?? false, - ); - - return filteredOrgs - .filter((org) => org.organization_id !== null && org.organization_id !== undefined) - .map((org) => ({ - label: `${org.organization_id || "Unknown"} (${org.organization_id})`, - value: org.organization_id as string, - })); - }, - }, - { - name: "Key Alias", - label: "Key Alias", - customComponent: PaginatedKeyAliasSelect, - }, - { - name: "User ID", - label: "User ID", - isSearchable: false, - }, - { - name: "Key Hash", - label: "Key ID", - isSearchable: false, - }, - ]; - - const table = useReactTable({ - data: keyList, - columns: columns.filter((col) => col.id !== "expander"), - columnResizeMode: "onChange", - columnResizeDirection: "ltr", - state: { - sorting, - pagination: tablePagination, - }, - onSortingChange: (updaterOrValue) => { - const newSorting = typeof updaterOrValue === "function" ? updaterOrValue(sorting) : updaterOrValue; - setSorting(newSorting); - setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); - }, - onPaginationChange: setTablePagination, - getCoreRowModel: getCoreRowModel(), - enableSorting: true, - manualSorting: true, - manualPagination: true, - pageCount: Math.ceil(totalCount / tablePagination.pageSize), - }); - - const { pageIndex, pageSize } = table.getState().pagination; - const start = pageIndex * pageSize + 1; - const end = Math.min((pageIndex + 1) * pageSize, totalCount); - const rangeLabel = `${start} - ${end}`; - return ( -
- {selectedKey ? ( + if (selectedKey) { + return ( +
setSelectedKey(null)} keyData={selectedKey} teams={allTeams} + onDelete={refetch} /> - ) : ( -
-
- + ); + } + + return ( +
+ } + title="Virtual Keys" + subtitle="Every key that authenticates requests to the gateway." + actions={headerActions} + /> + row.token} + defaultColumnVisibility={KEY_TABLE_HIDDEN_COLUMNS} + sortingMode="server" + sorting={sorting} + onSortingChange={handleSortingChange} + paginationMode="server" + pagination={tablePagination} + onPaginationChange={setTablePagination} + rowCount={rowCount} + filterMode="server" + columnFilters={columnFilters} + onColumnFiltersChange={handleColumnFiltersChange} + enableColumnResizing + columnResizeMode="onChange" + isLoading={isLoading} + loadingMessage="Loading keys..." + noDataMessage="No keys found" + maxBodyHeight="calc(75vh - 210px)" + size="compact" + toolbar={(table) => ( + <> + refetch?.()} + isRefreshing={isFetching} + onOpenFilters={() => setFiltersOpen(true)} + filterLabels={FILTER_LABELS} + formatFilterValue={formatFilterValue} /> -
- -
-
- {isLoading ? ( - - ) : ( - - Showing {rangeLabel} of {totalCount} results - + + {({ get, set }) => ( + <> + + set("team_id", value)} + placeholder="Select a team…" + emptyText="No teams found" + /> + + + set("org_id", value)} + placeholder="Select an organization…" + emptyText="No organizations found" + /> + + + set("user_id", event.target.value)} + placeholder="Enter User ID…" + /> + + + set("key_hash", event.target.value)} + placeholder="Enter Key ID…" + /> + + )} - - } - onClick={handleRefresh} - disabled={isButtonLoading} - title="Fetch data" - > - {isButtonLoading ? "Fetching" : "Fetch"} - -
- -
- {isLoading ? ( - - ) : ( - - Page {pageIndex + 1} of {table.getPageCount()} - - )} - - {isLoading ? ( - - ) : ( - - )} - - {isLoading ? ( - - ) : ( - - )} -
-
-
-
-
- - - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - { - const resizer = document.querySelector(`[data-header-id="${header.id}"] .resizer`); - if (resizer) { - (resizer as HTMLElement).style.opacity = "0.5"; - } - }} - onMouseLeave={() => { - const resizer = document.querySelector(`[data-header-id="${header.id}"] .resizer`); - if (resizer && !header.column.getIsResizing()) { - (resizer as HTMLElement).style.opacity = "0"; - } - }} - onClick={header.column.getCanSort() ? header.column.getToggleSortingHandler() : undefined} - > -
-
- {header.isPlaceholder - ? null - : flexRender(header.column.columnDef.header, header.getContext())} -
- {header.id !== "actions" && header.column.getCanSort() && ( -
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
- )} -
header.column.resetSize()} - onMouseDown={header.getResizeHandler()} - onTouchStart={header.getResizeHandler()} - className={`resizer ${table.options.columnResizeDirection} ${header.column.getIsResizing() ? "isResizing" : ""}`} - style={{ - position: "absolute", - right: 0, - top: 0, - height: "100%", - width: "5px", - background: header.column.getIsResizing() ? "#3b82f6" : "transparent", - cursor: "col-resize", - userSelect: "none", - touchAction: "none", - opacity: header.column.getIsResizing() ? 1 : 0, - }} - /> -
- - ))} - - ))} - - - {isLoading ? ( - - -
-

🚅 Loading keys...

-
- - - ) : keyList.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - - {row.getVisibleCells().map((cell) => ( - 3 ? "px-0" : ""}`} - > - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No keys found

-
- - - )} - -
-
-
-
-
- )} + + + )} + />
); } diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx new file mode 100644 index 00000000000..e2dc48fed9d --- /dev/null +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -0,0 +1,353 @@ +"use client"; + +import { InfoCircleOutlined } from "@ant-design/icons"; +import { ColumnDef } from "@tanstack/react-table"; +import { Popover, Typography } from "antd"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { Skeleton } from "@/components/ui/skeleton"; +import { + DateCell, + IdCell, + IdentityCell, + ModelsCell, + SpendBudgetCell, + StatusBadge, + type StatusTone, +} from "@/components/shared/table_cells"; + +import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; +import { KeyResponse, Team } from "../key_team_helpers/key_list"; +import { Organization } from "../networking"; + +interface KeyStatus { + tone: StatusTone; + label: string; + tooltip?: string; +} + +const getKeyStatus = (key: KeyResponse): KeyStatus => { + if (key.blocked === true) { + const isScimBlocked = (key.metadata as Record | null | undefined)?.scim_blocked === true; + return { + tone: "error", + label: "Blocked", + tooltip: isScimBlocked + ? "Blocked by SCIM (external identity provider deactivated or deleted the owning user)." + : "Blocked. Requests using this key will be rejected with 401.", + }; + } + const expiresAt = key.expires ? Date.parse(key.expires) : Number.NaN; + if (!Number.isNaN(expiresAt) && expiresAt < Date.now()) { + return { tone: "warning", label: "Expired", tooltip: "This key has passed its expiry date." }; + } + return { tone: "success", label: "Active" }; +}; + +const UserPopoverCell = ({ + userAlias, + userEmail, + userId, + width, +}: { + userAlias: string | null; + userEmail: string | null; + userId: string | null; + width: number; +}) => { + const displayValue = userAlias || userEmail || userId; + const isDefaultAdmin = userId === "default_user_id"; + + const popoverContent = ( +
+ {[ + { label: "User Alias", value: userAlias }, + { label: "User Email", value: userEmail }, + { label: "User ID", value: userId }, + ].map(({ label, value }) => ( +
+ {label} + {value ? ( + + {value} + + ) : ( + - + )} +
+ ))} +
+ ); + + if (isDefaultAdmin && !userAlias && !userEmail) { + return ( + + + + + + ); + } + + return ( + + + {displayValue || "-"} + + + ); +}; + +const InfoHeader = ({ label, tooltip }: { label: string; tooltip: string }) => ( + + {label} + + + + +); + +interface KeyTableColumnsDeps { + allTeams: Team[]; + organizations: Organization[]; + onSelectKey: (key: KeyResponse) => void; +} + +export const getKeyTableColumns = ({ + allTeams, + organizations, + onSelectKey, +}: KeyTableColumnsDeps): ColumnDef[] => [ + { + id: "key_alias", + accessorKey: "key_alias", + meta: { + title: "Key", + renderSkeleton: () => ( +
+ +
+ + +
+
+ ), + }, + header: ({ column }) => , + size: 260, + enableSorting: true, + cell: ({ row }) => { + const status = getKeyStatus(row.original); + return ( + + } + onClick={() => onSelectKey(row.original)} + /> + ); + }, + }, + { + id: "token", + accessorKey: "token", + meta: { title: "Key ID" }, + header: ({ column }) => , + size: 120, + enableSorting: true, + cell: (info) => onSelectKey(info.row.original)} />, + }, + { + id: "team_alias", + accessorKey: "team_id", + meta: { title: "Team" }, + header: "Team", + size: 120, + enableSorting: false, + cell: (info) => { + const teamId = info.getValue() as string | null; + if (!teamId) return "-"; + const team = allTeams.find((t) => t.team_id === teamId); + const displayValue = team?.team_alias || teamId; + const width = info.cell.column.getSize(); + return ( + + {displayValue} + + ); + }, + }, + { + id: "organization_alias", + accessorKey: "org_id", + meta: { title: "Organization" }, + header: "Organization", + size: 140, + enableSorting: false, + cell: (info) => { + const orgId = info.getValue() as string | null; + if (!orgId) return "-"; + const org = organizations.find((o) => o.organization_id === orgId); + const displayValue = org?.organization_alias || orgId; + const width = info.cell.column.getSize(); + return ( + + {displayValue} + + ); + }, + }, + { + id: "user", + accessorKey: "user", + meta: { title: "User" }, + header: () => ( + + ), + size: 160, + enableSorting: false, + cell: ({ row }) => { + const key = row.original; + return ( + + ); + }, + }, + { + id: "created_at", + accessorKey: "created_at", + meta: { title: "Created At" }, + header: ({ column }) => , + size: 120, + enableSorting: true, + cell: (info) => , + }, + { + id: "created_by", + accessorKey: "created_by", + meta: { title: "Created By" }, + header: "Created By", + size: 160, + enableSorting: false, + cell: (info) => { + const userId = info.getValue() as string | null; + if (!userId) return "-"; + const createdByUser = info.row.original.created_by_user; + return ( + + ); + }, + }, + { + id: "updated_at", + accessorKey: "updated_at", + meta: { title: "Updated At" }, + header: ({ column }) => , + size: 120, + enableSorting: true, + cell: (info) => , + }, + { + id: "last_active", + accessorKey: "last_active", + meta: { title: "Last Active" }, + header: () => ( + + ), + size: 130, + enableSorting: false, + cell: (info) => , + }, + { + id: "expires", + accessorKey: "expires", + meta: { title: "Expires" }, + header: "Expires", + size: 120, + enableSorting: false, + cell: (info) => , + }, + { + id: "spend", + accessorKey: "spend", + meta: { title: "Spend / Budget", skeleton: "meter" }, + header: ({ column }) => , + size: 180, + enableSorting: true, + cell: ({ row }) => { + const teamId = row.original.team_id; + const team = allTeams.find((t) => t.team_id === teamId); + return ( + + ); + }, + }, + { + id: "budget_reset_at", + accessorKey: "budget_reset_at", + meta: { title: "Budget Reset" }, + header: "Budget Reset", + size: 130, + enableSorting: false, + cell: (info) => , + }, + { + id: "models", + accessorKey: "models", + meta: { title: "Models", skeleton: "chips" }, + header: "Models", + size: 220, + enableSorting: false, + cell: (info) => , + }, + { + id: "rate_limits", + meta: { title: "Rate Limits" }, + header: "Rate Limits", + size: 140, + enableSorting: false, + cell: ({ row }) => { + const key = row.original; + return ( +
+
TPM: {key.tpm_limit !== null ? key.tpm_limit : "Unlimited"}
+
RPM: {key.rpm_limit !== null ? key.rpm_limit : "Unlimited"}
+
+ ); + }, + }, +]; + +export const KEY_TABLE_HIDDEN_COLUMNS: Record = { + token: false, + organization_alias: false, + created_by: false, + updated_at: false, + expires: false, + budget_reset_at: false, + rate_limits: false, +}; diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index ef0d842ad3e..d8d4c9392dc 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -303,6 +303,38 @@ describe("DataTable loading", () => { // per-column widths differ instead of every cell sharing one fixed width expect(new Set(bars.map((bar) => bar.className)).size).toBeGreaterThan(1); }); + + it("renders shape-specific skeletons for badge, chips, and meter columns", () => { + const columns: ColumnDef[] = [ + { id: "badge", header: "Badge", meta: { skeleton: "badge" }, cell: () => null }, + { id: "chips", header: "Chips", meta: { skeleton: "chips" }, cell: () => null }, + { id: "meter", header: "Meter", meta: { skeleton: "meter" }, cell: () => null }, + ]; + render(); + + const firstRow = screen.getAllByTestId("skeleton-row").at(0); + const cells = Array.from(firstRow?.querySelectorAll("td") ?? []); + const barsIn = (cell: Element | undefined) => cell?.querySelectorAll('[data-slot="skeleton"]').length ?? 0; + + // badge = a single pill, chips = three pills, meter = value bar + track bar + expect(barsIn(cells[0])).toBe(1); + expect(cells[0]?.querySelector('[data-slot="skeleton"]')?.className).toContain("rounded-full"); + expect(barsIn(cells[1])).toBe(3); + expect(barsIn(cells[2])).toBe(2); + }); + + it("uses a column's renderSkeleton override when provided", () => { + const columns: ColumnDef[] = [ + { + id: "custom", + header: "Custom", + meta: { renderSkeleton: () =>
loading
}, + cell: () => null, + }, + ]; + render(); + expect(screen.getAllByTestId("custom-skeleton").length).toBeGreaterThan(0); + }); }); describe("DataTable column visibility", () => { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index 758ca5a597b..bc13318dd58 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -338,7 +338,11 @@ const SKELETON_WIDTHS = ["w-[58%]", "w-[44%]", "w-[70%]", "w-[50%]", "w-[64%]", function SkeletonCell({ column, index }: { column: Column | undefined; index: number }) { const meta = column?.columnDef.meta; const width = SKELETON_WIDTHS[index % SKELETON_WIDTHS.length]; - if (meta?.skeleton === "twoLine") { + const shape = meta?.skeleton; + if (meta?.renderSkeleton !== undefined) { + return <>{meta.renderSkeleton()}; + } + if (shape === "twoLine") { return (
@@ -346,6 +350,26 @@ function SkeletonCell({ column, index }: { column: Column
); } + if (shape === "badge") { + return ; + } + if (shape === "chips") { + return ( +
+ + + +
+ ); + } + if (shape === "meter") { + return ( +
+ + +
+ ); + } return ; } diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/columnMeta.ts b/ui/litellm-dashboard/src/components/shared/DataTable/columnMeta.ts index 0f14c277c6f..eff4e0cb7db 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/columnMeta.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/columnMeta.ts @@ -1,4 +1,5 @@ import type { RowData } from "@tanstack/react-table"; +import type * as React from "react"; import type { ColumnPinnedSide, DataTableSkeletonShape } from "./types"; @@ -10,5 +11,7 @@ declare module "@tanstack/react-table" { title?: string; pinned?: ColumnPinnedSide; skeleton?: DataTableSkeletonShape; + /** Full control over this column's loading skeleton, for cells the built-in shapes can't mirror. */ + renderSkeleton?: () => React.ReactNode; } } diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts index f5130b4c823..672ab512ef4 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts @@ -18,7 +18,7 @@ export type FilterMode = "none" | "client" | "server"; export type ColumnResizeMode = "onEnd" | "onChange"; export type DataTableSize = "compact" | "default"; export type ColumnPinnedSide = "left" | "right"; -export type DataTableSkeletonShape = "text" | "twoLine"; +export type DataTableSkeletonShape = "text" | "twoLine" | "badge" | "chips" | "meter"; export interface DataTableProps { data: TData[]; diff --git a/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx b/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx new file mode 100644 index 00000000000..f7a313271da --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx @@ -0,0 +1,31 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; + +import { PageHeader } from "./PageHeader"; + +describe("PageHeader", () => { + it("renders the title as a heading", () => { + render(); + expect(screen.getByRole("heading", { name: "Virtual Keys" })).toBeInTheDocument(); + }); + + it("renders the subtitle, icon, and actions when provided", () => { + render( + } + actions={} + />, + ); + expect(screen.getByText("Every key that authenticates requests")).toBeInTheDocument(); + expect(screen.getByTestId("icon")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Create New Key" })).toBeInTheDocument(); + }); + + it("omits the optional slots when not provided", () => { + render(); + expect(screen.queryByRole("button")).not.toBeInTheDocument(); + expect(document.querySelector("p")).toBeNull(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/PageHeader.tsx b/ui/litellm-dashboard/src/components/shared/PageHeader.tsx new file mode 100644 index 00000000000..34d478e1cd9 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/PageHeader.tsx @@ -0,0 +1,29 @@ +"use client"; + +import * as React from "react"; + +interface PageHeaderProps { + title: React.ReactNode; + subtitle?: React.ReactNode; + icon?: React.ReactNode; + actions?: React.ReactNode; +} + +export function PageHeader({ title, subtitle, icon, actions }: PageHeaderProps) { + return ( +
+
+ {icon != null && ( + + {icon} + + )} +
+

{title}

+ {subtitle != null &&

{subtitle}

} +
+
+ {actions != null &&
{actions}
} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx new file mode 100644 index 00000000000..acf50d282b4 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx @@ -0,0 +1,64 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import { SearchSelect } from "./SearchSelect"; + +const OPTIONS = [ + { label: "Acme Prod", value: "team-1" }, + { label: "Growth", value: "team-2" }, + { label: "Data Team", value: "team-3" }, +]; + +describe("SearchSelect", () => { + it("renders the placeholder when nothing is selected", () => { + render(); + expect(screen.getByPlaceholderText("Select Team…")).toBeInTheDocument(); + }); + + it("shows the selected option's label in the field", () => { + render(); + expect(screen.getByRole("combobox")).toHaveValue("Growth"); + }); + + it("shows a clear control only when a value is selected", () => { + const { rerender } = render(); + expect(document.querySelector('[data-slot="combobox-clear"]')).toBeNull(); + rerender(); + expect(document.querySelector('[data-slot="combobox-clear"]')).not.toBeNull(); + }); + + it("filters the options client-side as you type", async () => { + const user = userEvent.setup(); + render(); + const input = screen.getByRole("combobox"); + await user.click(input); + await user.type(input, "grow"); + expect(await screen.findByText("Growth")).toBeInTheDocument(); + expect(screen.queryByText("Acme Prod")).not.toBeInTheDocument(); + }); + + it("renders a muted sublabel and matches it when searching", async () => { + const user = userEvent.setup(); + render( + , + ); + const input = screen.getByRole("combobox"); + await user.click(input); + expect(await screen.findByText("team-abc-123")).toBeInTheDocument(); + await user.type(input, "abc-123"); + expect(await screen.findByText("Acme Prod")).toBeInTheDocument(); + }); + + it("selects an option and reports its value", async () => { + const onValueChange = vi.fn(); + const user = userEvent.setup(); + render(); + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("Growth")); + expect(onValueChange).toHaveBeenCalledWith("team-2"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx new file mode 100644 index 00000000000..c29e099a1c6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx @@ -0,0 +1,76 @@ +"use client"; + +import { + Combobox, + ComboboxContent, + ComboboxEmpty, + ComboboxInput, + ComboboxItem, + ComboboxList, +} from "@/components/ui/combobox"; + +export interface SearchSelectOption { + label: string; + value: string; + /** Optional muted second line (e.g. an id); also matched when searching. */ + sublabel?: string; +} + +interface SearchSelectProps { + options: SearchSelectOption[]; + value?: string; + onValueChange: (value: string) => void; + placeholder?: string; + emptyText?: string; + disabled?: boolean; + className?: string; +} + +export function SearchSelect({ + options, + value, + onValueChange, + placeholder = "Select…", + emptyText = "No results", + disabled = false, + className, +}: SearchSelectProps) { + const selected = options.find((option) => option.value === value) ?? null; + + return ( + onValueChange(item?.value ?? "")} + isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value} + itemToStringLabel={(item: SearchSelectOption) => item.label} + filter={(item: SearchSelectOption, query: string) => { + const q = query.trim().toLowerCase(); + if (!q) return true; + return item.label.toLowerCase().includes(q) || (item.sublabel?.toLowerCase().includes(q) ?? false); + }} + disabled={disabled} + > + + + {emptyText} + + {(item: SearchSelectOption) => ( + + + {item.label} + {item.sublabel != null && item.sublabel !== "" && ( + {item.sublabel} + )} + + + )} + + + + ); +} diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.test.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.test.tsx new file mode 100644 index 00000000000..db4e93c7cb2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.test.tsx @@ -0,0 +1,38 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import { IdentityCell } from "./identity_cell"; + +describe("IdentityCell", () => { + it("renders the title", () => { + render(); + expect(screen.getByText("prod-gateway")).toBeInTheDocument(); + }); + + it("renders the subtitle and an inline badge together", () => { + render(Active} />); + expect(screen.getByText("sk-...v0Pw")).toBeInTheDocument(); + expect(screen.getByText("Active")).toBeInTheDocument(); + }); + + it("omits the subtitle row when there is no subtitle or badge", () => { + render(); + expect(document.querySelector("span.font-mono")).toBeNull(); + }); + + it("renders a static div (no button) when not clickable", () => { + render(); + expect(screen.queryByRole("button")).not.toBeInTheDocument(); + }); + + it("renders a clickable button and fires onClick", async () => { + const onClick = vi.fn(); + const user = userEvent.setup(); + render(); + const button = screen.getByRole("button"); + expect(button.querySelector(".lucide-chevron-right")).not.toBeNull(); + await user.click(button); + expect(onClick).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.tsx new file mode 100644 index 00000000000..4d3e3d8e4dd --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.tsx @@ -0,0 +1,48 @@ +"use client"; + +import { ChevronRight } from "lucide-react"; +import * as React from "react"; + +import { cn } from "@/lib/cva.config"; + +interface IdentityCellProps { + title: React.ReactNode; + subtitle?: React.ReactNode; + badge?: React.ReactNode; + onClick?: () => void; + className?: string; + titleClassName?: string; +} + +export function IdentityCell({ title, subtitle, badge, onClick, className, titleClassName }: IdentityCellProps) { + const hasSubtitleRow = (subtitle != null && subtitle !== "") || badge != null; + + const body = ( +
+ {title} + {hasSubtitleRow && ( + + {subtitle != null && subtitle !== "" && ( + {subtitle} + )} + {badge} + + )} +
+ ); + + if (onClick != null) { + return ( + + ); + } + + return
{body}
; +} diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/index.ts b/ui/litellm-dashboard/src/components/shared/table_cells/index.ts index e189413d43d..9fdd04d169c 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/index.ts +++ b/ui/litellm-dashboard/src/components/shared/table_cells/index.ts @@ -1,5 +1,8 @@ export { CellTooltip } from "./cell_tooltip"; export { DateCell, formatCellDate, formatFullTimestamp, type DatePrecision } from "./date_cell"; export { IdCell, type IdCellVariant } from "./id_cell"; +export { IdentityCell } from "./identity_cell"; +export { ModelsCell } from "./models_cell"; export { MoneyCell } from "./money_cell"; +export { SpendBudgetCell } from "./spend_budget_cell"; export { StatusBadge, type StatusTone } from "./status_badge"; diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx new file mode 100644 index 00000000000..d3fad1d3244 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx @@ -0,0 +1,45 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it } from "vitest"; + +import { ModelsCell } from "./models_cell"; + +describe("ModelsCell", () => { + it("shows 'All Proxy Models' when the list is empty, null, or undefined", () => { + const { rerender } = render(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + rerender(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + rerender(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + }); + + it("renders every model with no overflow badge when at or below the limit", () => { + render(); + expect(screen.getByText("gpt-4o")).toBeInTheDocument(); + expect(screen.getByText("claude-sonnet-4-5")).toBeInTheDocument(); + expect(screen.getByText("o3-mini")).toBeInTheDocument(); + expect(screen.queryByText(/more$/)).not.toBeInTheDocument(); + }); + + it("collapses models beyond the limit into a '+N more' badge", () => { + render(); + expect(screen.getByText("a")).toBeInTheDocument(); + expect(screen.getByText("b")).toBeInTheDocument(); + expect(screen.queryByText("c")).not.toBeInTheDocument(); + expect(screen.getByText("+3 more")).toBeInTheDocument(); + }); + + it("reveals the hidden models in a tooltip on hover", async () => { + const user = userEvent.setup(); + render(); + await user.hover(screen.getByText("+2 more")); + expect(await screen.findByText("c")).toBeInTheDocument(); + expect(await screen.findByText("d")).toBeInTheDocument(); + }); + + it("labels the all-proxy-models wildcard", () => { + render(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx new file mode 100644 index 00000000000..712d8511c78 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx @@ -0,0 +1,56 @@ +"use client"; + +import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; +import { Badge } from "@/components/ui/badge"; + +import { CellTooltip } from "./cell_tooltip"; + +interface ModelsCellProps { + models: string[] | null | undefined; + maxVisible?: number; +} + +const WILDCARD_MODEL = "all-proxy-models"; + +const formatModel = (model: string): string => { + if (model === WILDCARD_MODEL) { + return "All Proxy Models"; + } + const name = getModelDisplayName(model); + return name.length > 30 ? `${name.slice(0, 30)}...` : name; +}; + +export function ModelsCell({ models, maxVisible = 3 }: ModelsCellProps) { + if (!Array.isArray(models) || models.length === 0) { + return All Proxy Models; + } + + const visible = models.slice(0, maxVisible); + const overflow = models.slice(maxVisible); + + return ( +
+ {visible.map((model, index) => ( + + {formatModel(model)} + + ))} + {overflow.length > 0 && ( + + {overflow.map((model, index) => ( + {formatModel(model)} + ))} +
+ } + trigger={ + + +{overflow.length} more + + } + /> + )} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx new file mode 100644 index 00000000000..707441aef1d --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx @@ -0,0 +1,53 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; + +import { SpendBudgetCell } from "./spend_budget_cell"; + +const indicator = (container: HTMLElement) => container.querySelector('[data-slot="meter-indicator"]'); + +describe("SpendBudgetCell", () => { + it("shows Unlimited and renders no meter when there is no budget", () => { + const { container } = render(); + expect(screen.getByText("· Unlimited")).toBeInTheDocument(); + expect(screen.queryByRole("meter")).not.toBeInTheDocument(); + expect(indicator(container)).toBeNull(); + }); + + it("shows $0.00 for zero or undefined spend, never a hyphen", () => { + const { rerender } = render(); + expect(screen.getByText("$0.00")).toBeInTheDocument(); + expect(screen.queryByText("-")).not.toBeInTheDocument(); + rerender(); + expect(screen.getByText("$0.00")).toBeInTheDocument(); + expect(screen.queryByText("-")).not.toBeInTheDocument(); + }); + + it("renders a meter carrying the spend and budget when a budget exists", () => { + render(); + const meter = screen.getByRole("meter"); + expect(meter).toHaveAttribute("aria-valuenow", "25"); + expect(meter).toHaveAttribute("aria-valuemax", "100"); + expect(screen.getByText("of $100")).toBeInTheDocument(); + }); + + it("keeps the default tone below 80% usage", () => { + const { container } = render(); + expect(indicator(container)?.className).toContain("bg-primary"); + }); + + it("switches to the warning tone at 80% usage", () => { + const { container } = render(); + expect(indicator(container)?.className).toContain("bg-amber-500"); + }); + + it("switches to the over tone above 100% usage", () => { + const { container } = render(); + expect(indicator(container)?.className).toContain("bg-destructive"); + }); + + it("falls back to the team budget and labels it", () => { + render(); + expect(screen.getByText("of $200 (Team)")).toBeInTheDocument(); + expect(screen.getByRole("meter")).toHaveAttribute("aria-valuemax", "200"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx new file mode 100644 index 00000000000..10956f23b1c --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx @@ -0,0 +1,44 @@ +"use client"; + +import { Meter, MeterIndicator, MeterTrack } from "@/components/ui/meter"; +import { formatNumberWithCommas, getSpendString } from "@/utils/dataUtils"; + +interface SpendBudgetCellProps { + spend: number | null | undefined; + maxBudget: number | null | undefined; + teamMaxBudget?: number | null; +} + +const meterTone = (pct: number): "default" | "warning" | "over" => { + if (pct > 100) return "over"; + if (pct >= 80) return "warning"; + return "default"; +}; + +export function SpendBudgetCell({ spend, maxBudget, teamMaxBudget }: SpendBudgetCellProps) { + const spendValue = typeof spend === "number" && !Number.isNaN(spend) ? spend : 0; + const budget = maxBudget ?? teamMaxBudget ?? null; + const isTeamBudget = maxBudget == null && teamMaxBudget != null; + const hasBudget = typeof budget === "number" && budget > 0; + const pct = hasBudget ? (spendValue / budget) * 100 : 0; + + const spendText = spendValue > 0 ? getSpendString(spendValue, 4) : "$0.00"; + const budgetLabel = + budget === null ? "· Unlimited" : `of $${formatNumberWithCommas(budget)}${isTeamBudget ? " (Team)" : ""}`; + + return ( +
+
+ {spendText}{" "} + {budgetLabel} +
+ {hasBudget && ( + + + + + + )} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/ui/combobox.tsx b/ui/litellm-dashboard/src/components/ui/combobox.tsx new file mode 100644 index 00000000000..2854928140e --- /dev/null +++ b/ui/litellm-dashboard/src/components/ui/combobox.tsx @@ -0,0 +1,266 @@ +"use client"; + +import * as React from "react"; +import { Combobox as ComboboxPrimitive } from "@base-ui/react"; + +import { cn } from "@/lib/cva.config"; +import { Button } from "@/components/ui/button"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; +import { ChevronDownIcon, XIcon, CheckIcon } from "lucide-react"; + +const Combobox = ComboboxPrimitive.Root; + +function ComboboxValue({ ...props }: ComboboxPrimitive.Value.Props) { + return ; +} + +const ComboboxTrigger = React.forwardRef< + React.ComponentRef, + ComboboxPrimitive.Trigger.Props +>(({ className, children, ...props }, ref) => { + return ( + + {children} + + + ); +}); +ComboboxTrigger.displayName = "ComboboxTrigger"; + +function ComboboxClear({ className, ...props }: ComboboxPrimitive.Clear.Props) { + return ( + } + className={cn(className)} + {...props} + > + + + ); +} + +function ComboboxInput({ + className, + children, + disabled = false, + showTrigger = true, + showClear = false, + ...props +}: ComboboxPrimitive.Input.Props & { + showTrigger?: boolean; + showClear?: boolean; +}) { + return ( + + } {...props} /> + + {showTrigger && ( + } + data-slot="input-group-button" + className="group-has-data-[slot=combobox-clear]/input-group:hidden data-pressed:bg-transparent" + disabled={disabled} + /> + )} + {showClear && } + + {children} + + ); +} + +function ComboboxContent({ + className, + side = "bottom", + sideOffset = 6, + align = "start", + alignOffset = 0, + anchor, + ...props +}: ComboboxPrimitive.Popup.Props & + Pick) { + return ( + + + + + + ); +} + +function ComboboxList({ className, ...props }: ComboboxPrimitive.List.Props) { + return ( + + ); +} + +function ComboboxItem({ className, children, ...props }: ComboboxPrimitive.Item.Props) { + return ( + + {children} + } + > + + + + ); +} + +function ComboboxGroup({ className, ...props }: ComboboxPrimitive.Group.Props) { + return ; +} + +function ComboboxLabel({ className, ...props }: ComboboxPrimitive.GroupLabel.Props) { + return ( + + ); +} + +function ComboboxCollection({ ...props }: ComboboxPrimitive.Collection.Props) { + return ; +} + +function ComboboxEmpty({ className, ...props }: ComboboxPrimitive.Empty.Props) { + return ( + + ); +} + +function ComboboxSeparator({ className, ...props }: ComboboxPrimitive.Separator.Props) { + return ( + + ); +} + +function ComboboxChips({ + className, + ...props +}: React.ComponentPropsWithRef & ComboboxPrimitive.Chips.Props) { + return ( + + ); +} + +function ComboboxChip({ + className, + children, + showRemove = true, + ...props +}: ComboboxPrimitive.Chip.Props & { + showRemove?: boolean; +}) { + return ( + + {children} + {showRemove && ( + } + className="-ml-1 opacity-50 hover:opacity-100" + data-slot="combobox-chip-remove" + > + + + )} + + ); +} + +function ComboboxChipsInput({ className, ...props }: ComboboxPrimitive.Input.Props) { + return ( + + ); +} + +function useComboboxAnchor() { + return React.useRef(null); +} + +export { + Combobox, + ComboboxInput, + ComboboxContent, + ComboboxList, + ComboboxItem, + ComboboxGroup, + ComboboxLabel, + ComboboxCollection, + ComboboxEmpty, + ComboboxSeparator, + ComboboxChips, + ComboboxChip, + ComboboxChipsInput, + ComboboxTrigger, + ComboboxValue, + useComboboxAnchor, +}; diff --git a/ui/litellm-dashboard/src/components/ui/input-group.tsx b/ui/litellm-dashboard/src/components/ui/input-group.tsx new file mode 100644 index 00000000000..8ee9b7f17bd --- /dev/null +++ b/ui/litellm-dashboard/src/components/ui/input-group.tsx @@ -0,0 +1,140 @@ +"use client"; + +import * as React from "react"; +import { type VariantProps } from "cva"; + +import { cn, cva } from "@/lib/cva.config"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Textarea } from "@/components/ui/textarea"; + +function InputGroup({ className, ...props }: React.ComponentProps<"div">) { + return ( +
[data-align=block-end]]:h-auto has-[>[data-align=block-end]]:flex-col has-[>[data-align=block-start]]:h-auto has-[>[data-align=block-start]]:flex-col has-[>textarea]:h-auto dark:bg-input/30 dark:has-[[data-slot][aria-invalid=true]]:ring-destructive/40 has-[>[data-align=block-end]]:[&>input]:pt-3 has-[>[data-align=block-start]]:[&>input]:pb-3 has-[>[data-align=inline-end]]:[&>input]:pr-1.5 has-[>[data-align=inline-start]]:[&>input]:pl-1.5", + className, + )} + {...props} + /> + ); +} + +const inputGroupAddonVariants = cva({ + base: "flex h-auto cursor-text items-center justify-center gap-2 py-1.5 text-sm font-medium text-muted-foreground select-none group-data-[disabled=true]/input-group:opacity-50 [&>kbd]:rounded-[calc(var(--radius)-5px)] [&>svg:not([class*='size-'])]:size-4", + variants: { + align: { + "inline-start": "order-first pl-2 has-[>button]:-ml-1 has-[>kbd]:ml-[-0.15rem]", + "inline-end": "order-last pr-2 has-[>button]:-mr-1 has-[>kbd]:mr-[-0.15rem]", + "block-start": + "order-first w-full justify-start px-2.5 pt-2 group-has-[>input]/input-group:pt-2 [.border-b]:pb-2", + "block-end": "order-last w-full justify-start px-2.5 pb-2 group-has-[>input]/input-group:pb-2 [.border-t]:pt-2", + }, + }, + defaultVariants: { + align: "inline-start", + }, +}); + +function InputGroupAddon({ + className, + align = "inline-start", + ...props +}: React.ComponentProps<"div"> & VariantProps) { + return ( +
{ + if ((e.target as HTMLElement).closest("button")) { + return; + } + e.currentTarget.parentElement?.querySelector("input")?.focus(); + }} + {...props} + /> + ); +} + +const inputGroupButtonVariants = cva({ + base: "flex items-center gap-2 text-sm shadow-none", + variants: { + size: { + xs: "h-6 gap-1 rounded-[calc(var(--radius)-5px)] px-1.5 [&>svg:not([class*='size-'])]:size-3.5", + sm: "", + "icon-xs": "size-6 rounded-[calc(var(--radius)-5px)] p-0 has-[>svg]:p-0", + "icon-sm": "size-8 p-0 has-[>svg]:p-0", + }, + }, + defaultVariants: { + size: "xs", + }, +}); + +const InputGroupButton = React.forwardRef< + React.ComponentRef, + Omit, "size" | "type"> & + VariantProps & { + type?: "button" | "submit" | "reset"; + } +>(({ className, type = "button", variant = "ghost", size = "xs", ...props }, ref) => { + return ( +