diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 00000000000..1a29c0b6691 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,2 @@ +/ui/ @yuneng-jiang @ryan-crabbe-berri +/litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri diff --git a/.github/actions/setup-uv-with-retries/action.yml b/.github/actions/setup-uv-with-retries/action.yml new file mode 100644 index 00000000000..1627038dc3d --- /dev/null +++ b/.github/actions/setup-uv-with-retries/action.yml @@ -0,0 +1,47 @@ +name: "Set up uv with retries" +description: >- + Install uv via astral-sh/setup-uv, retrying on transient failures. Even with + an exact pinned version, the action resolves the artifact URL by fetching + https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a + single request with no retry, timeout, or fallback, so one connection-level + network error ("fetch failed") fails the whole job before any test runs. + Retrying the full step covers the manifest fetch and the binary download. + +inputs: + version: + description: "uv version to install" + required: true + +runs: + using: composite + steps: + - name: Set up uv (attempt 1) + id: attempt-1 + continue-on-error: true + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + with: + version: ${{ inputs.version }} + + - name: Wait before attempt 2 + if: steps.attempt-1.outcome == 'failure' + shell: bash + run: sleep 15 + + - name: Set up uv (attempt 2) + id: attempt-2 + if: steps.attempt-1.outcome == 'failure' + continue-on-error: true + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + with: + version: ${{ inputs.version }} + + - name: Wait before attempt 3 + if: steps.attempt-2.outcome == 'failure' + shell: bash + run: sleep 30 + + - name: Set up uv (attempt 3) + if: steps.attempt-2.outcome == 'failure' + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + with: + version: ${{ inputs.version }} diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index bd9fc2285d1..d7e80b32749 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -41,3 +41,27 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac ✅ Test ## Changes + +## QA runbook + + + +### Final Attestation + +- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 9fd81b27f3b..92230fc8892 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -63,7 +63,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/auto_update_price_and_context_window.yml b/.github/workflows/auto_update_price_and_context_window.yml index 1c6c318c717..1a638a4a331 100644 --- a/.github/workflows/auto_update_price_and_context_window.yml +++ b/.github/workflows/auto_update_price_and_context_window.yml @@ -18,7 +18,7 @@ jobs: with: persist-credentials: false - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" - name: Update JSON Data diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index 439126aa1ee..9c24bad00f1 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -31,7 +31,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index a1772102b89..54a8e53d7a3 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -37,7 +37,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/mutation-test.yml b/.github/workflows/mutation-test.yml index 183f12f969c..6684952b998 100644 --- a/.github/workflows/mutation-test.yml +++ b/.github/workflows/mutation-test.yml @@ -39,7 +39,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/oss_daily_guardrails.yml b/.github/workflows/oss_daily_guardrails.yml index 950c51c9b60..f9dc746ee05 100644 --- a/.github/workflows/oss_daily_guardrails.yml +++ b/.github/workflows/oss_daily_guardrails.yml @@ -35,7 +35,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 872a1799d98..9d28ca211cf 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -38,7 +38,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 0b5b0e9b976..09406d77634 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -33,7 +33,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" @@ -172,7 +172,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-litellm-ui-lint.yml b/.github/workflows/test-litellm-ui-lint.yml index 5a5c4709ca2..804894b1e50 100644 --- a/.github/workflows/test-litellm-ui-lint.yml +++ b/.github/workflows/test-litellm-ui-lint.yml @@ -89,4 +89,4 @@ jobs: - name: Check for dead code (knip) if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }} - run: npm run knip + run: npm run knip:ci diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml index 5b5290880c1..a5a4e722133 100644 --- a/.github/workflows/test-mcp.yml +++ b/.github/workflows/test-mcp.yml @@ -32,7 +32,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-semgrep.yml b/.github/workflows/test-semgrep.yml index f0dcb9887be..6e9f5e42fa2 100644 --- a/.github/workflows/test-semgrep.yml +++ b/.github/workflows/test-semgrep.yml @@ -31,7 +31,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-terraform-provider.yml b/.github/workflows/test-terraform-provider.yml index 03d8ff3461c..058a2538c15 100644 --- a/.github/workflows/test-terraform-provider.yml +++ b/.github/workflows/test-terraform-provider.yml @@ -74,7 +74,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-unit-documentation.yml b/.github/workflows/test-unit-documentation.yml index 03f9f0a510b..c12a289ce9f 100644 --- a/.github/workflows/test-unit-documentation.yml +++ b/.github/workflows/test-unit-documentation.yml @@ -42,7 +42,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/.github/workflows/test-unit-proxy-legacy.yml b/.github/workflows/test-unit-proxy-legacy.yml index 0068e80e584..bcbf365babf 100644 --- a/.github/workflows/test-unit-proxy-legacy.yml +++ b/.github/workflows/test-unit-proxy-legacy.yml @@ -59,7 +59,7 @@ jobs: python-version: "3.12" - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 + uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" diff --git a/CLAUDE.md b/CLAUDE.md index 78da2c65d96..9f708716c6d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -39,6 +39,8 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a Python max line length is 120, not 88 +On a fresh worktree or clone, run `make bootstrap` before anything else. It provisions everything tests, `make pre-commit`, and a local proxy need + Run tests before you commit. Also, run `make pre-commit` right before each commit, which generates types (as needed) and formats/lints your code. Any errors found must be fixed. It only runs when there are staged frontend and/or backend changes and calculates violations, generates types, etc. based on the worktree, so stage what you need or stash/delete unwanted files in litellm/ or ui/ (where backend and frontend lint run, respectively) before running it. If it fails because dashboard api types are stale, it already regenerated them for you. You just need to stage the schema.d.ts, re-run `make pre-commit` to confirm it passes, and commit When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing diff --git a/Makefile b/Makefile index 965a3254616..8b657dcb465 100644 --- a/Makefile +++ b/Makefile @@ -9,11 +9,12 @@ lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \ install-dev install-proxy-dev install-test-deps install-hooks \ install-helm-unittest check-circular-imports check-import-safety pre-commit \ - lint-install lint-fetch-base + lint-install lint-fetch-base bootstrap # Default target help: @echo "Available commands:" + @echo " make bootstrap - Provision a fresh clone/worktree" @echo " make install-dev - Install development dependencies" @echo " make install-proxy-dev - Install proxy development dependencies" @echo " make install-dev-ci - Install dev dependencies (CI-compatible, pins OpenAI)" @@ -71,6 +72,18 @@ info: install-dev: $(UV) sync --inexact --frozen +bootstrap: + $(UV) sync --inexact --frozen --extra proxy --group proxy-dev --group e2e-dev + $(UV_RUN) python scripts/prisma_generate_if_needed.py + cd ui/litellm-dashboard && npm ci --no-audit --no-fund + @main_root=$$(git worktree list --porcelain | head -1 | sed 's/^worktree //'); \ + if [ "$$main_root" != "$$(git rev-parse --show-toplevel)" ] && [ -f "$$main_root/.env" ] && [ ! -f .env ]; then \ + cp "$$main_root/.env" .env && echo "bootstrap: copied .env from $$main_root"; \ + else \ + echo "bootstrap: .env left untouched"; \ + fi + @echo "bootstrap: done" + install-proxy-dev: $(UV) sync --frozen --group proxy-dev --extra proxy diff --git a/README.md b/README.md index 90d3e944fcc..32b0160dbaa 100644 --- a/README.md +++ b/README.md @@ -552,17 +552,12 @@ The Terraform modules live at [`terraform/litellm/aws/`](./terraform/litellm/aws 2. Run dependent services `docker-compose up db prometheus` #### Backend -1. (In root) create virtual environment `python -m venv .venv` -2. Activate virtual environment `source .venv/bin/activate` -3. Install dependencies `uv sync --all-extras --group proxy-dev` -4. `uv run prisma generate` -5. `prisma generate` -6. Start proxy backend `python litellm/proxy/proxy_cli.py` +1. Run `make bootstrap` +2. Start proxy backend: `uv run python litellm/proxy/proxy_cli.py` #### Frontend -1. Navigate to `ui/litellm-dashboard` -2. Install dependencies `npm install` -3. Run `npm run dev` to start the dashboard +1. Navigate to `ui/litellm-dashboard` (dependencies were already installed w/ `make bootstrap`) +2. Start dashboard: `npm run dev` ### Verify Docker Image Signatures diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index b9ac98f515c..f209ab54f64 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -29,7 +29,7 @@ class CheckBatchCost: proxy_logging_obj: "ProxyLogging", prisma_client: "PrismaClient", llm_router: "Router", - track_unmanaged_vertex_batch_cost: bool = False, + track_unmanaged_batch_cost: bool = False, ): from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -37,7 +37,7 @@ class CheckBatchCost: self.proxy_logging_obj: ProxyLogging = proxy_logging_obj self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router - self._track_unmanaged_vertex_batch_cost = track_unmanaged_vertex_batch_cost + self._track_unmanaged_batch_cost = track_unmanaged_batch_cost # Cached after the first poll cycle. Once we know the column is absent we skip # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True @@ -118,11 +118,11 @@ class CheckBatchCost: Resolve (model_id, batch_id) for a managed-object row, where model_id is a router deployment id and batch_id is the raw provider batch id. - Managed batches encode both in a base64 unified id. Unmanaged Vertex batches, created with - a raw gs:// input_file_id, store the raw provider job id as unified_object_id; when - track_unmanaged_vertex_batch_cost is enabled the model is derived from the gs:// path and - mapped to a configured vertex_ai deployment. Returns None (recording a metric) when the row - can't be routed. + Managed batches encode both in a base64 unified id. Unmanaged batches (created outside + LiteLLM's own /v1/batches with a raw input_file_id) store the raw provider job id as + unified_object_id instead; when track_unmanaged_batch_cost is enabled the model is derived + from the provider-specific input_file_id layout (Vertex gs:// or Bedrock s3://) and mapped + to a matching deployment. Returns None (recording a metric) when the row can't be routed. """ from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, @@ -142,8 +142,43 @@ class CheckBatchCost: return None return model_id, get_batch_id_from_unified_batch_id(decoded) - if self._track_unmanaged_vertex_batch_cost: - return self._resolve_unmanaged_vertex_routing(job, prom_logger) + if self._track_unmanaged_batch_cost: + from litellm.llms.bedrock.batches.transformation import ( + BedrockBatchesConfig, + ) + from litellm.llms.vertex_ai.batches.transformation import ( + VertexAIBatchTransformation, + ) + + input_file_id = self._get_input_file_id(job) + if VertexAIBatchTransformation.is_unmanaged_gcs_batch_input_file_id( + input_file_id + ): + assert input_file_id is not None # narrowed by is_unmanaged_gcs_batch_input_file_id + return self._resolve_unmanaged_provider_routing( + job=job, + prom_logger=prom_logger, + llm_provider="vertex_ai", + bare_model_name=VertexAIBatchTransformation.get_bare_model_name_from_gcs_file( + input_file_id + ), + ) + if BedrockBatchesConfig.is_unmanaged_s3_batch_input_file_id(input_file_id): + assert input_file_id is not None # narrowed by is_unmanaged_s3_batch_input_file_id + return self._resolve_unmanaged_provider_routing( + job=job, + prom_logger=prom_logger, + llm_provider="bedrock", + bare_model_name=BedrockBatchesConfig.get_bare_model_name_from_s3_file( + input_file_id + ), + ) + verbose_proxy_logger.info( + f"Skipping job {unified_object_id}: not a recognized unmanaged batch " + "(no gs:// or s3:// input_file_id with an embedded model)" + ) + self._record_error(prom_logger, "invalid_unified_id") + return None verbose_proxy_logger.info( f"Skipping job {unified_object_id} because it is not a valid unified object id" @@ -151,36 +186,17 @@ class CheckBatchCost: self._record_error(prom_logger, "invalid_unified_id") return None - def _resolve_unmanaged_vertex_routing( + def _resolve_unmanaged_provider_routing( self, job: "LiteLLM_ManagedObjectTable", prom_logger: Optional["PrometheusLogger"], + llm_provider: str, + bare_model_name: str, ) -> Optional[Tuple[str, str]]: - from litellm.llms.vertex_ai.batches.transformation import ( - VertexAIBatchTransformation, - ) - - input_file_id = self._get_input_file_id(job) - if not VertexAIBatchTransformation.is_unmanaged_gcs_batch_input_file_id( - input_file_id - ): - verbose_proxy_logger.info( - f"Skipping job {job.unified_object_id}: not an unmanaged vertex batch " - "(no gs:// input_file_id with a publishers/ model path)" - ) - self._record_error(prom_logger, "invalid_unified_id") - return None - assert input_file_id is not None # narrowed by is_unmanaged_gcs_batch_input_file_id - - bare_model_name = VertexAIBatchTransformation.get_bare_model_name_from_gcs_file( - input_file_id - ) - deployment_id = self._get_vertex_ai_deployment_id_for_bare_model( - bare_model_name - ) + deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider) if deployment_id is None: verbose_proxy_logger.info( - f"Skipping unmanaged vertex batch {job.unified_object_id}: no vertex_ai " + f"Skipping unmanaged {llm_provider} batch {job.unified_object_id}: no {llm_provider} " f"deployment configured for model {bare_model_name}" ) self._record_error(prom_logger, "unmanaged_no_matching_deployment") @@ -188,22 +204,22 @@ class CheckBatchCost: return deployment_id, job.unified_object_id - def _get_vertex_ai_deployment_id_for_bare_model( - self, bare_model_name: str + def _get_deployment_id_for_bare_model( + self, bare_model_name: str, llm_provider: str ) -> Optional[str]: model_group = self.llm_router.resolve_model_name_from_model_id(bare_model_name) deployment_id = ( - self._get_vertex_ai_deployment_id(model_group) if model_group else None + self._get_deployment_id_for_provider(model_group, llm_provider) if model_group else None ) if deployment_id is not None: return deployment_id - return self._get_vertex_ai_deployment_id_from_matching_deployments( - bare_model_name + return self._get_deployment_id_from_matching_deployments( + bare_model_name, llm_provider ) - def _get_vertex_ai_deployment_id_from_matching_deployments( - self, bare_model_name: str + def _get_deployment_id_from_matching_deployments( + self, bare_model_name: str, llm_provider: str ) -> Optional[str]: from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -215,13 +231,13 @@ class CheckBatchCost: if not self._is_bare_model_match(actual_model, bare_model_name): continue try: - _, llm_provider, _, _ = get_llm_provider( + _, deployment_llm_provider, _, _ = get_llm_provider( model=actual_model, custom_llm_provider=litellm_params.get("custom_llm_provider"), ) except Exception: continue - if llm_provider != "vertex_ai": + if deployment_llm_provider != llm_provider: continue model_info = deployment.get("model_info") or {} deployment_id = model_info.get("id") @@ -231,15 +247,21 @@ class CheckBatchCost: @staticmethod def _is_bare_model_match(actual_model: str, bare_model_name: str) -> bool: + # Bedrock model ids may have ":" replaced with "-" in the S3 object key (see + # BedrockBatchesConfig.get_bare_model_name_from_s3_file), so normalize both sides; + # a no-op for providers like vertex_ai whose model ids never contain a colon. + normalized_actual = actual_model.replace(":", "-") + normalized_bare = bare_model_name.replace(":", "-") return ( - actual_model == bare_model_name - or actual_model.endswith(f"/{bare_model_name}") - or actual_model.endswith(f":{bare_model_name}") + normalized_actual == normalized_bare + or normalized_actual.endswith(f"/{normalized_bare}") ) - def _get_vertex_ai_deployment_id(self, model_group: str) -> Optional[str]: + def _get_deployment_id_for_provider( + self, model_group: str, llm_provider: str + ) -> Optional[str]: """ - Returns the first deployment id for `model_group` whose provider is vertex_ai, + Returns the first deployment id for `model_group` whose provider is `llm_provider`, skipping deployments from other providers that happen to share the model group name. """ from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -249,13 +271,13 @@ class CheckBatchCost: if deployment_info is None: continue try: - _, llm_provider, _, _ = get_llm_provider( + _, deployment_llm_provider, _, _ = get_llm_provider( model=deployment_info.litellm_params.model, custom_llm_provider=deployment_info.litellm_params.custom_llm_provider, ) except Exception: continue - if llm_provider == "vertex_ai": + if deployment_llm_provider == llm_provider: return deployment_id return None diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 85ccbef752f..97571a4576d 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.49" +version = "0.1.50" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.49" +version = "0.1.50" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713230852_add_key_type_to_litellm_verification_token/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713230852_add_key_type_to_litellm_verification_token/migration.sql new file mode 100644 index 00000000000..708b7601346 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713230852_add_key_type_to_litellm_verification_token/migration.sql @@ -0,0 +1,6 @@ +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "key_type" TEXT; + +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "key_type" TEXT; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index fb4d8d0b5a3..a23cecc3911 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -422,6 +422,7 @@ model LiteLLM_VerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") @@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index a54db2ace65..b67d9d8570a 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.76" +version = "0.4.77" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.76" +version = "0.4.77" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 11b07d39981..2fcb8455e90 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -118,6 +118,7 @@ def _batch_cost_calculator( total_cost = _get_batch_job_cost_from_file_content( file_content_dictionary=file_content_dictionary, custom_llm_provider=custom_llm_provider, + model_name=model_name, model_info=model_info, ) verbose_logger.debug("total_cost=%s", total_cost) @@ -363,6 +364,7 @@ def _count_entry_tokens( def _get_batch_job_cost_from_file_content( file_content_dictionary: List[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + model_name: Optional[str] = None, model_info: Optional[ModelInfo] = None, ) -> float: """ @@ -377,9 +379,15 @@ def _get_batch_job_cost_from_file_content( for _item in file_content_dictionary: if _batch_response_was_successful(_item, custom_llm_provider): _response_body = _get_response_from_batch_job_output_file(_item, custom_llm_provider) - if model_info is not None or custom_llm_provider == "anthropic": + if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"): usage = _get_batch_job_usage_from_response_body(_response_body, custom_llm_provider) - model = _response_body.get("model", "") + # Bedrock batch output lines report a short internal model id + # (e.g. "claude-sonnet-4-6") that is not in the cost map; use the + # deployment model name for pricing when available. + if custom_llm_provider == "bedrock" and model_name: + model = model_name + else: + model = _response_body.get("model") or model_name or "" prompt_cost, completion_cost = batch_cost_calculator( usage=usage, model=model, @@ -485,7 +493,7 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov """ Get the tokens of a batch job from the response body """ - if custom_llm_provider == "anthropic": + if custom_llm_provider in ("anthropic", "bedrock"): from litellm.llms.anthropic.chat.transformation import AnthropicConfig return AnthropicConfig().calculate_usage( @@ -513,6 +521,8 @@ def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom """ if custom_llm_provider == "anthropic": return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("message", None) or {} + if custom_llm_provider == "bedrock": + return batch_job_output_file.get("modelOutput", None) or {} _response: dict = batch_job_output_file.get("response", None) or {} _response_body = _response.get("body", None) or {} return _response_body @@ -523,9 +533,12 @@ def _batch_response_was_successful(batch_job_output_file: dict, custom_llm_provi Check if the batch job response was successful OpenAI-shaped output rows report ``response.status_code == 200``; Anthropic - message batch results lines report ``result.type == "succeeded"``. + message batch results lines report ``result.type == "succeeded"``; Bedrock + batch output lines report ``modelOutput`` (and no ``error``). """ if custom_llm_provider == "anthropic": return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("type") == "succeeded" + if custom_llm_provider == "bedrock": + return batch_job_output_file.get("modelOutput") is not None and batch_job_output_file.get("error") is None _response: dict = batch_job_output_file.get("response", None) or {} return _response.get("status_code", None) == 200 diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index c8bfcabc64e..856556f7c56 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -515,6 +515,22 @@ class CustomGuardrail(CustomLogger): return True return False + def uses_apply_guardrail_interface(self) -> bool: + return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail + + def _deployment_pre_call_target(self) -> "CustomLogger": + if not self.uses_apply_guardrail_interface(): + return self + try: + from litellm.proxy.utils import unified_guardrail + except ImportError as e: + raise ImportError( + f"Guardrail {self.guardrail_name or type(self).__name__} implements apply_guardrail, which needs " + "the litellm proxy dependencies to run at the deployment level. " + "Install them with: pip install 'litellm[proxy]'" + ) from e + return unified_guardrail + async def async_pre_call_deployment_hook( self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] ) -> Optional[dict]: @@ -533,7 +549,10 @@ class CustomGuardrail(CustomLogger): # CHECK IF GUARDRAIL REJECTS THE REQUEST if call_type == CallTypes.completion or call_type == CallTypes.acompletion: - result = await self.async_pre_call_hook( + target = self._deployment_pre_call_target() + if target is not self: + kwargs["guardrail_to_apply"] = self + result = await target.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( user_id=kwargs.get("user_api_key_user_id"), team_id=kwargs.get("user_api_key_team_id"), @@ -543,7 +562,7 @@ class CustomGuardrail(CustomLogger): ), cache=dc, data=kwargs, - call_type=call_type.value or "acompletion", # type: ignore + call_type="completion" if call_type == CallTypes.completion else "acompletion", ) if result is not None and isinstance(result, dict): diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f575372fc3d..64d4dd578b2 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -239,6 +239,18 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_output_audio_tokens_metric"), ) + self.litellm_video_duration_seconds_metric = self._counter_factory( + "litellm_video_duration_seconds_metric", + "Seconds of video generated, from usage.duration_seconds on video generation calls", + labelnames=self.get_labels_for_metric("litellm_video_duration_seconds_metric"), + ) + + self.litellm_images_generated_metric = self._counter_factory( + "litellm_images_generated_metric", + "Number of images generated, from the image generation response", + labelnames=self.get_labels_for_metric("litellm_images_generated_metric"), + ) + # Remaining Budget for Team self.litellm_remaining_team_budget_metric = self._gauge_factory( "litellm_remaining_team_budget_metric", @@ -1336,6 +1348,12 @@ class PrometheusLogger(CustomLogger): label_context=label_context, ) + self._increment_media_generation_metrics( + standard_logging_payload=standard_logging_payload, + enum_values=enum_values, + label_context=label_context, + ) + # MCP tool call metrics self._increment_mcp_tool_call_metrics( standard_logging_payload=standard_logging_payload, @@ -1459,8 +1477,65 @@ class PrometheusLogger(CustomLogger): ), ] - for counter, metric_name, value in detail_metrics: - if not isinstance(value, (int, float)) or value <= 0: + PrometheusLogger._inc_sparse_usage_counters( + self, + detail_metrics, + enum_values=enum_values, + label_context=label_context, + ) + + def _increment_media_generation_metrics( + self, + standard_logging_payload: StandardLoggingPayload, + enum_values: UserAPIKeyLabelValues, + label_context: PrometheusLabelFactoryContext | None = None, + ) -> None: + """ + Increment video-seconds and images-generated counters from + ``standard_logging_payload["metadata"]["usage_object"]``. Video + providers report ``duration_seconds`` there; image generation calls + report ``output_image_count``. Both are sparse: only emitted when the + value is present and > 0, so token-only call types are unaffected. + """ + metadata = standard_logging_payload.get("metadata") or {} + usage_object = metadata.get("usage_object") if isinstance(metadata, dict) else None + if not isinstance(usage_object, dict): + return + + media_metrics: list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]] = [ + ( + self.litellm_video_duration_seconds_metric, + "litellm_video_duration_seconds_metric", + usage_object.get("duration_seconds"), + ), + ( + self.litellm_images_generated_metric, + "litellm_images_generated_metric", + usage_object.get("output_image_count"), + ), + ] + + PrometheusLogger._inc_sparse_usage_counters( + self, + media_metrics, + enum_values=enum_values, + label_context=label_context, + ) + + def _inc_sparse_usage_counters( + self, + counters_with_values: list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]], + enum_values: UserAPIKeyLabelValues, + label_context: PrometheusLabelFactoryContext | None = None, + ) -> None: + """ + Increment each ``(counter, metric_name, value)`` entry whose value is + a positive number. Non-numeric values (including booleans from + malformed provider usage dicts) and values <= 0 are skipped, keeping + scrape output sparse. + """ + for counter, metric_name, value in counters_with_values: + if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0: continue PrometheusLogger._inc_labeled_counter( self, @@ -1716,6 +1791,35 @@ class PrometheusLogger(CustomLogger): amount=float(response_cost), ) + @staticmethod + def _get_remaining_from_v3_rate_limit_headers( + standard_logging_payload: StandardLoggingPayload | None, + rate_limit_type: Literal["requests", "tokens"], + ) -> int | None: + """ + Read the per-(key, model) remaining value emitted by the v3 rate + limiter (``parallel_request_limiter_v3.py``), which writes + ``x-ratelimit-model_per_key-remaining-{requests,tokens}`` into + ``standard_logging_object.hidden_params.additional_headers`` instead + of the ``litellm-key-remaining-*`` metadata keys the legacy limiter + sets. The header carries no model group; it always refers to this + request's model group, which is what the gauges are labeled with. + Values are written in-process as plain ints (never HTTP-serialized + strings), so anything else is rejected rather than coerced. + """ + if standard_logging_payload is None: + return None + hidden_params = standard_logging_payload.get("hidden_params") + if hidden_params is None: + return None + additional_headers = hidden_params.get("additional_headers") + if additional_headers is None: + return None + value = dict(additional_headers).get(f"x-ratelimit-model_per_key-remaining-{rate_limit_type}") + if isinstance(value, bool) or not isinstance(value, int): + return None + return value + def _set_virtual_key_rate_limit_metrics( self, user_api_key: Optional[str], @@ -1733,11 +1837,20 @@ class PrometheusLogger(CustomLogger): model_group = get_model_group_from_litellm_kwargs(kwargs) remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}" remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}" + standard_logging_payload: StandardLoggingPayload | None = kwargs.get("standard_logging_object") remaining_requests = metadata.get(remaining_requests_variable_name) + if remaining_requests is None: + remaining_requests = self._get_remaining_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, rate_limit_type="requests" + ) if remaining_requests is None: remaining_requests = sys.maxsize remaining_tokens = metadata.get(remaining_tokens_variable_name) + if remaining_tokens is None: + remaining_tokens = self._get_remaining_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, rate_limit_type="tokens" + ) if remaining_tokens is None: remaining_tokens = sys.maxsize diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index 53a982cd2c4..e8252d87572 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -161,8 +161,13 @@ def get_s3_object_key( start_time: datetime, s3_file_name: str, ) -> str: + sanitized_s3_file_name = s3_file_name.replace("/", "_") s3_object_key = ( - (s3_path.rstrip("/") + "/" if s3_path else "") + prefix + start_time.strftime("%Y-%m-%d") + "/" + s3_file_name + (s3_path.rstrip("/") + "/" if s3_path else "") + + prefix + + start_time.strftime("%Y-%m-%d") + + "/" + + sanitized_s3_file_name ) # we need the s3 key to include the time, so we log cache hits too s3_object_key += ".json" return s3_object_key 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/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 936d79b22d6..461ab62b815 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -73,6 +73,7 @@ from litellm.litellm_core_utils.model_param_helper import ModelParamHelper from litellm.litellm_core_utils.redact_messages import ( redact_message_input_output_from_custom_logger, redact_message_input_output_from_logging, + redact_streaming_responses_for_custom_logger, ) from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.llms.base_llm.search.transformation import SearchResponse @@ -2576,6 +2577,9 @@ class Logging(LiteLLMLoggingBaseClass): model_call_details = callback.redact_standard_logging_payload_from_model_call_details( model_call_details=model_call_details ) + model_call_details = redact_streaming_responses_for_custom_logger( + model_call_details=model_call_details, custom_logger=callback + ) ################################## if self.stream is True: if "async_complete_streaming_response" in model_call_details: @@ -5208,10 +5212,15 @@ def get_standard_logging_object_payload( call_type = kwargs.get("call_type") cache_hit = kwargs.get("cache_hit", False) # Extract usage as a plain dict, avoiding Pydantic round-trip - usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict( + raw_usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict( response_obj=response_obj, combined_usage_object=cast(Optional[Usage], kwargs.get("combined_usage_object")), ) + usage_dict = ( + {**raw_usage_dict, "output_image_count": len(init_response_obj.data)} + if isinstance(init_response_obj, ImageResponse) and init_response_obj.data + else raw_usage_dict + ) id = response_obj.get("id", kwargs.get("litellm_call_id")) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index c039f0f43ee..33bf546c239 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -445,6 +445,7 @@ class PromptTokensDetailsResult(TypedDict): text_tokens: int audio_tokens: int image_tokens: int + video_tokens: int character_count: int image_count: int video_length_seconds: float @@ -473,6 +474,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) audio_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0 image_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0 + video_tokens = _coerce_token_count(getattr(usage.prompt_tokens_details, "video_tokens", 0)) character_count = ( cast( Optional[int], @@ -503,6 +505,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: text_tokens=text_tokens, audio_tokens=audio_tokens, image_tokens=image_tokens, + video_tokens=video_tokens, character_count=character_count, image_count=image_count, video_length_seconds=float(video_length_seconds), @@ -515,6 +518,7 @@ class CompletionTokensDetailsResult(TypedDict): text_tokens: int reasoning_tokens: int image_tokens: int + video_tokens: int def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult: @@ -546,12 +550,14 @@ def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsRes ) or 0 ) + video_tokens = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0)) return CompletionTokensDetailsResult( audio_tokens=audio_tokens, text_tokens=text_tokens, reasoning_tokens=reasoning_tokens, image_tokens=image_tokens, + video_tokens=video_tokens, ) @@ -586,6 +592,13 @@ def _calculate_input_cost( image_token_cost_key = "input_cost_per_token" prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"]) + ### VIDEO TOKEN COST + if prompt_tokens_details["video_tokens"]: + video_token_cost_key = "input_cost_per_video_token" + if model_info.get(video_token_cost_key) is None: + video_token_cost_key = "input_cost_per_token" + prompt_cost += calculate_cost_component(model_info, video_token_cost_key, prompt_tokens_details["video_tokens"]) + ### CACHE WRITING COST - Now uses tiered pricing if ( prompt_tokens_details["cache_creation_tokens"] @@ -698,6 +711,7 @@ def generic_cost_per_token( text_tokens=usage.prompt_tokens, audio_tokens=0, image_tokens=0, + video_tokens=0, character_count=0, image_count=0, video_length_seconds=0.0, @@ -716,13 +730,14 @@ def generic_cost_per_token( audio_tokens = prompt_tokens_details["audio_tokens"] cache_creation = prompt_tokens_details["cache_creation_tokens"] image_tokens = prompt_tokens_details["image_tokens"] + video_tokens = prompt_tokens_details["video_tokens"] # Check for double-counting: sum of details > prompt_tokens means overlap - total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + video_tokens has_double_counting = cache_hit > 0 and total_details > usage.prompt_tokens if (text_tokens == 0 and prompt_tokens_details["image_count"] == 0) or has_double_counting: - text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens + text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens # Clamp to zero: inconsistent streaming usage if text_tokens < 0: text_tokens = 0 @@ -751,6 +766,7 @@ def generic_cost_per_token( audio_tokens = 0 reasoning_tokens = 0 image_tokens = 0 + video_tokens = 0 is_text_tokens_total = False if usage.completion_tokens_details is not None: completion_tokens_details = _parse_completion_tokens_details(usage) @@ -758,19 +774,20 @@ def generic_cost_per_token( text_tokens = completion_tokens_details["text_tokens"] reasoning_tokens = completion_tokens_details["reasoning_tokens"] image_tokens = completion_tokens_details["image_tokens"] + video_tokens = completion_tokens_details["video_tokens"] # Handle text_tokens calculation: # 1. If text_tokens is explicitly provided and > 0, use it - # 2. If there's a breakdown (reasoning/audio/image tokens), calculate text_tokens as the remainder + # 2. If there's a breakdown (reasoning/audio/image/video tokens), calculate text_tokens as the remainder # 3. If no breakdown at all, assume all completion_tokens are text_tokens - has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0 + has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0 or video_tokens > 0 if text_tokens == 0: if has_token_breakdown: # Calculate text tokens as remainder when we have a breakdown # This handles cases like OpenAI's reasoning models where text_tokens isn't provided text_tokens = max( 0, - usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens, + usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens - video_tokens, ) else: # No breakdown at all, all tokens are text tokens @@ -803,6 +820,14 @@ def generic_cost_per_token( ) completion_cost += float(image_tokens) * _output_cost_per_image_token + ## VIDEO COST + if not is_text_tokens_total and video_tokens and video_tokens > 0: + _output_cost_per_video_token = _get_cost_per_unit(model_info, "output_cost_per_video_token", None) + _output_cost_per_video_token = ( + _output_cost_per_video_token if _output_cost_per_video_token is not None else completion_base_cost + ) + completion_cost += float(video_tokens) * _output_cost_per_video_token + ## REGIONAL DATA-RESIDENCY UPLIFT # Applied as a flat multiplier across all token costs for the request # when the upstream is a regionalized OpenAI host (eu./us.api.openai.com). diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 8bb0e12905e..f7ff4d6b16f 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5494,3 +5494,56 @@ def has_tool_with_name(tools: Any, tool_name: str) -> bool: elif tool.get("name") == tool_name: return True return False + + +def resolve_structured_messages( + messages: list[dict[str, Any]] | None, + request_kwargs: dict[str, Any], +) -> list[dict[str, Any]] | None: + """ + Normalize a request's messages to OpenAI-spec chat-completions shape, + regardless of which API surface produced them (chat completions, + Anthropic /v1/messages, Responses API ``input``, etc). + + Returns ``messages`` unchanged if already present. Otherwise dispatches + through the guardrail translation handlers (the same per-surface + conversion logic guardrails use) to convert e.g. Responses API ``input`` + into a message list. Returns ``None`` if no messages could be resolved. + """ + if messages: + return messages + + from litellm.litellm_core_utils.api_route_to_call_types import ( + get_call_types_for_route, + ) + from litellm.llms import load_guardrail_translation_mappings + from litellm.types.utils import CallTypes + + mappings = load_guardrail_translation_mappings() + call_type: CallTypes | None = None + + # 1. Try route-based inference from proxy metadata + route = request_kwargs.get("litellm_metadata", {}).get("user_api_key_request_route") + if route: + call_types_list = get_call_types_for_route(route) + if call_types_list: + for ct in call_types_list: + if ct in mappings: + call_type = ct + break + + # 2. Fallback: try each mapped handler until one produces messages + handlers_to_try: list[Any] = [] + if call_type is not None and call_type in mappings: + handlers_to_try.append(mappings[call_type]()) + else: + handlers_to_try.extend(handler_cls() for handler_cls in mappings.values()) + + for handler in handlers_to_try: + structured = handler.get_structured_messages(request_kwargs) + if structured: + return [ + msg if isinstance(msg, dict) else msg.model_dump() # type: ignore + for msg in structured + ] + return None diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index cc9264e93f8..6e8429839ad 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -38,10 +38,45 @@ def redact_message_input_output_from_custom_logger( litellm_logging_obj: LiteLLMLoggingObject, result, custom_logger: CustomLogger ): if hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True: - return perform_redaction(litellm_logging_obj.model_call_details, result) + return perform_redaction(litellm_logging_obj.model_call_details, result, redact_streaming_responses=False) return result +def redact_streaming_responses_for_custom_logger(model_call_details: dict, custom_logger: CustomLogger) -> dict: + """ + Returns a copy of model_call_details whose streaming response entries are redacted deepcopies + when the custom logger has opted out of message logging. The shared model_call_details is left + untouched so other callbacks still receive the unredacted response. + """ + if not (hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True): + return model_call_details + redacted_entries = { + streaming_key: _redacted_streaming_response_copy(model_call_details[streaming_key]) + for streaming_key in ("complete_streaming_response", "async_complete_streaming_response") + if model_call_details.get(streaming_key) is not None + } + if not redacted_entries: + return model_call_details + return {**model_call_details, **redacted_entries} + + +def _redacted_streaming_response_copy(streaming_response): + redacted_response = copy.deepcopy(streaming_response) + _redact_streaming_response(redacted_response) + return redacted_response + + +def _redact_streaming_response(streaming_response): + if hasattr(streaming_response, "choices"): + for choice in streaming_response.choices: + _redact_choice_content(choice) + redact_vertex_ai_metadata_from_logged_object(streaming_response) + elif hasattr(streaming_response, "output"): + _redact_responses_api_output(streaming_response.output) + if hasattr(streaming_response, "reasoning") and streaming_response.reasoning is not None: + streaming_response.reasoning = None + + def _redact_choice_content(choice): """Helper to redact content in a choice (message or delta).""" if isinstance(choice, litellm.Choices): @@ -150,9 +185,13 @@ def _redact_model_response_dict_choices(choices, redacted_str: str): _redact_choice_content(choice) -def perform_redaction(model_call_details: dict, result): +def perform_redaction(model_call_details: dict, result, redact_streaming_responses: bool = True): """ Performs the actual redaction on the logging object and result. + + redact_streaming_responses=False skips the in-place redaction of the shared streaming + response entries; per-callback redaction hands each opted-out callback its own redacted + copy via redact_streaming_responses_for_custom_logger instead. """ # Redact model_call_details model_call_details["messages"] = [{"role": "user", "content": "redacted-by-litellm"}] @@ -162,17 +201,9 @@ def perform_redaction(model_call_details: dict, result): redact_vertex_ai_metadata_from_litellm_params(model_call_details) # Redact streaming response - if model_call_details.get("stream", False) is True and "complete_streaming_response" in model_call_details: - _streaming_response = model_call_details["complete_streaming_response"] - if hasattr(_streaming_response, "choices"): - for choice in _streaming_response.choices: - _redact_choice_content(choice) - redact_vertex_ai_metadata_from_logged_object(_streaming_response) - elif hasattr(_streaming_response, "output"): - _redact_responses_api_output(_streaming_response.output) - # Redact reasoning field in ResponsesAPIResponse - if hasattr(_streaming_response, "reasoning") and _streaming_response.reasoning is not None: - _streaming_response.reasoning = None + if redact_streaming_responses and model_call_details.get("stream", False) is True: + for _streaming_key in ("complete_streaming_response", "async_complete_streaming_response"): + _redact_streaming_response(model_call_details.get(_streaming_key)) # Redact result if result is not None: diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 6033e54fb77..0ec1f3eae13 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -227,6 +227,10 @@ DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING = ( "Sonnet 4.6+, and Mythos Preview." ) +DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING = ( + "Dropping adaptive `thinking` for model=%s: max_tokens is too small to fit the minimum thinking budget." +) + DROP_UNSUPPORTED_SPEED_WARNING = ( "Dropping unsupported `speed` for model=%s (drop_params=True). Fast mode is only supported on select Opus models." ) @@ -1220,6 +1224,23 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): llm_provider=llm_provider, ) + @staticmethod + def _cap_thinking_budget_to_max_tokens( + thinking: AnthropicThinkingParam, max_tokens: Optional[int] + ) -> Optional[AnthropicThinkingParam]: + """Cap a legacy ``thinking.budget_tokens`` below ``max_tokens`` (Anthropic + requires ``max_tokens > budget_tokens``). Returns the (possibly capped) + thinking dict, or ``None`` when ``max_tokens`` is too small to fit even the + minimum thinking budget and thinking should be dropped.""" + budget = thinking.get("budget_tokens") + if max_tokens is None or not isinstance(budget, int): + return thinking + if max_tokens <= ANTHROPIC_MIN_THINKING_BUDGET_TOKENS: + return None + if budget < max_tokens: + return thinking + return AnthropicThinkingParam(type=thinking.get("type", "enabled"), budget_tokens=max_tokens - 1) + def _extract_json_schema_from_response_format(self, value: Optional[dict]) -> Optional[dict]: if value is None: return None @@ -1420,24 +1441,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): output_key=param, ) elif param == "response_format" and isinstance(value, dict): - if any( - substring in model - for substring in { - "sonnet-4.5", - "sonnet-4-5", - "opus-4.1", - "opus-4-1", - "opus-4.5", - "opus-4-5", - "opus-4.6", - "opus-4-6", - "opus-4.7", - "opus-4-7", - "sonnet-4.6", - "sonnet-4-6", - "sonnet_4.6", - "sonnet_4_6", - } + if AnthropicConfig._supports_model_capability( + model, + "supports_native_structured_output", + self._resolved_provider, ): _output_format = self.map_response_format_to_anthropic_output_format(value) if _output_format is not None: @@ -1463,7 +1470,38 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ): optional_params["metadata"] = {"user_id": value} elif param == "thinking": - optional_params["thinking"] = value + if ( + isinstance(value, dict) + and value.get("type") == "adaptive" + and not AnthropicConfig._is_adaptive_thinking_model(model, self._resolved_provider) + ): + # Callers (e.g. Claude Code) send adaptive thinking + # unconditionally; translate it down to the legacy + # `thinking={type: enabled, budget_tokens}` interface a + # pre-4.6 model actually supports instead of forwarding a + # shape the model will reject. + max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens") + legacy_thinking = AnthropicConfig._map_reasoning_effort( + reasoning_effort="medium", + model=model, + custom_llm_provider=self._resolved_provider, + llm_provider=self._resolved_provider, + ) + capped_thinking = ( + AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens) + if legacy_thinking is not None + else None + ) + if capped_thinking is not None: + optional_params["thinking"] = capped_thinking + else: + litellm.verbose_logger.warning( + DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, + model, + ) + optional_params.pop("thinking", None) + else: + optional_params["thinking"] = value elif param == "reasoning_effort": # Accept both string ("low") and dict ({"effort": "low", # "summary": "concise"}). The Responses->Chat parser keeps the diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 0bcf34a45d6..e006662ec4d 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -340,11 +340,15 @@ class AnthropicModelInfo(BaseLLMModelInfo): def _get_model_capability(model: str, key: str) -> Optional[bool]: """Read boolean capability ``key`` from the model map, or None when no entry declares it.""" + from litellm.utils import _get_bundled_model_cost_map + try: - for cand in AnthropicModelInfo._model_map_lookup_candidates(model): - value = litellm.model_cost.get(cand, {}).get(key) - if isinstance(value, bool): - return value + candidates = AnthropicModelInfo._model_map_lookup_candidates(model) + for model_cost in (litellm.model_cost, _get_bundled_model_cost_map()): + for cand in candidates: + value = model_cost.get(cand, {}).get(key) + if isinstance(value, bool): + return value except Exception: pass return None diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 01412d3c8b2..05679bf39ab 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -3,7 +3,6 @@ from typing import Any, AsyncIterator, Dict, List, Optional, Tuple import httpx from litellm.constants import ( - ANTHROPIC_MIN_THINKING_BUDGET_TOKENS, DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, @@ -358,7 +357,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): except _BadRequestError as e: raise AnthropicError(message=str(e.message), status_code=400) capped_thinking = ( - AnthropicMessagesConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens) + AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens) if legacy_thinking is not None else None ) @@ -377,19 +376,34 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params.pop("output_config", None) @staticmethod - def _cap_thinking_budget_to_max_tokens(thinking: Dict, max_tokens: Optional[int]) -> Optional[Dict]: - """Cap a legacy ``thinking.budget_tokens`` below ``max_tokens`` (Anthropic - requires ``max_tokens > budget_tokens``). Returns the (possibly capped) - thinking dict, or ``None`` when ``max_tokens`` is too small to fit even the - minimum thinking budget and thinking should be dropped.""" - budget = thinking.get("budget_tokens") - if max_tokens is None or not isinstance(budget, int): - return thinking - if max_tokens <= ANTHROPIC_MIN_THINKING_BUDGET_TOKENS: - return None - if budget < max_tokens: - return thinking - return {**thinking, "budget_tokens": max_tokens - 1} + def _drop_incompatible_temperature_for_thinking( + model: str, optional_params: dict, custom_llm_provider: str + ) -> None: + """Anthropic rejects any ``temperature`` other than 1 while extended thinking + is enabled ("temperature may only be set to 1 when thinking is enabled"). + + Clients like Claude Code send ``thinking``/``output_config.effort`` together + with a pinned ``temperature`` (e.g. the safety classifier uses ``temperature=0`` + for determinism). When the request lands on a non-adaptive model, the effort + interface is reshaped above into legacy ``thinking={type: enabled}`` (or kept + as ``output_config.effort`` on Opus 4.5), and the leftover ``temperature`` would + 400. Preserving the thinking the caller asked for wins over an unhonorable + sampling value (Anthropic forces ``temperature=1`` under thinking regardless), + so drop it and let the API default apply. + + Adaptive models (4.6+) own this natively and are left untouched. + """ + if AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider): + return + temperature = optional_params.get("temperature") + if temperature is None or temperature == 1: + return + thinking = optional_params.get("thinking") + output_config = optional_params.get("output_config") + thinking_enabled = isinstance(thinking, dict) and thinking.get("type") == "enabled" + effort_enabled = isinstance(output_config, dict) and output_config.get("effort") is not None + if thinking_enabled or effort_enabled: + optional_params.pop("temperature", None) def transform_anthropic_messages_request( self, @@ -431,6 +445,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): custom_llm_provider=self._resolved_provider, ) + self._drop_incompatible_temperature_for_thinking( + model=model, + optional_params=anthropic_messages_optional_request_params, + custom_llm_provider=self._resolved_provider, + ) + system_param = anthropic_messages_optional_request_params.get("system") if self.should_strip_billing_metadata() and system_param is not None: filtered_system = self._filter_billing_headers_from_system(system_param) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 1a052f457c5..172e54de98e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -198,14 +198,16 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @staticmethod def translate_tool_choice_to_responses_api( tool_choice: AnthropicMessagesToolChoice, - ) -> Dict[str, Any]: + ) -> Union[str, dict[str, Any]]: """Convert Anthropic tool_choice to Responses API tool_choice.""" tc_type = tool_choice.get("type") if tc_type == "any": - return {"type": "required"} + return "required" elif tc_type == "tool": return {"type": "function", "name": tool_choice.get("name", "")} - return {"type": "auto"} + elif tc_type == "none": + return "none" + return "auto" @staticmethod def translate_context_management_to_responses_api( 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/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index b0e28b6ba90..4fcf7cf91cb 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -5,6 +5,9 @@ from typing import Any, Dict, List, Literal, Optional, Union, cast from httpx import Headers, Response +from litellm.litellm_core_utils.cloud_storage_security import ( + BEDROCK_MANAGED_S3_BATCH_PREFIX, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -26,6 +29,15 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders from ..base_aws_llm import BaseAWSLLM from ..common_utils import CommonBatchFilesUtils +# Bedrock batch input files are uploaded as +# s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see +# BedrockFilesTransformation._get_s3_object_name). A uuid4 is always 36 hex/dash +# characters, so it can be stripped off the end unambiguously even though the +# model name itself may contain dashes. +_S3_BATCH_FILE_UUID_SUFFIX_PATTERN = re.compile( + r"-[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.jsonl$" +) + class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): """ @@ -40,6 +52,41 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK + @classmethod + def _get_bare_model_name_from_s3_key(cls, object_key: str) -> Optional[str]: + if not object_key.startswith(BEDROCK_MANAGED_S3_BATCH_PREFIX): + return None + model_part = object_key[len(BEDROCK_MANAGED_S3_BATCH_PREFIX) :] + match = _S3_BATCH_FILE_UUID_SUFFIX_PATTERN.search(model_part) + if not match or match.start() == 0: + return None + return model_part[: match.start()] + + @classmethod + def is_unmanaged_s3_batch_input_file_id(cls, input_file_id: Optional[str]) -> bool: + """ + Returns True if `input_file_id` is a raw s3:// Bedrock batch input file (i.e. not a + LiteLLM-managed unified file id) whose object key embeds the model name in the + `litellm-bedrock-files-{model}-{uuid}.jsonl` layout. + """ + if input_file_id is None or not input_file_id.startswith("s3://"): + return False + object_key = input_file_id.rsplit("/", 1)[-1] + return cls._get_bare_model_name_from_s3_key(object_key) is not None + + @classmethod + def get_bare_model_name_from_s3_file(cls, input_file_id: str) -> str: + """ + Extracts the bare model name (e.g. "us.anthropic.claude-sonnet-4-20250514-v1-0") from + an unmanaged batch's s3:// input file id. Note any ":" in the original model id was + replaced with "-" at upload time, so callers must fuzzy-match against configured + deployments rather than expect an exact string match. + """ + object_key = input_file_id.rsplit("/", 1)[-1] + bare_model_name = cls._get_bare_model_name_from_s3_key(object_key) + assert bare_model_name is not None # narrowed by is_unmanaged_s3_batch_input_file_id + return bare_model_name + def validate_environment( self, headers: dict, diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index be904fb27be..c38b3593465 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -33,6 +33,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( make_valid_bedrock_tool_name, ) from litellm.llms.anthropic.chat.transformation import ( + DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT, AnthropicConfig, @@ -899,7 +900,28 @@ class AmazonConverseConfig(BaseConfig): "tool_choice": {"disable_parallel_tool_use": disable_parallel} } if param == "thinking": - optional_params["thinking"] = value + if ( + isinstance(value, dict) + and value.get("type") == "adaptive" + and not AnthropicConfig._is_adaptive_thinking_model(model, "bedrock") + ): + max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens") + legacy_thinking = AnthropicConfig._map_reasoning_effort( + reasoning_effort="medium", + model=model, + custom_llm_provider="bedrock", + ) + capped = ( + AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens) + if legacy_thinking is not None + else None + ) + if capped is not None: + optional_params["thinking"] = capped + else: + litellm.verbose_logger.warning(DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, model) + else: + optional_params["thinking"] = value elif param == "reasoning_effort" and isinstance(value, str): self._handle_reasoning_effort_parameter( model=model, reasoning_effort=value, optional_params=optional_params diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index d107ca7a0d7..3c2ae238a0b 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -20,6 +20,8 @@ from litellm.types.utils import LlmProviders from ..common_utils import OpenAIError +OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS = 16 + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -59,6 +61,19 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): key="supports_none_reasoning_effort", ) + @staticmethod + def _enforce_min_max_output_tokens(max_output_tokens: "int | None") -> "int | None": + """Raise sub-minimum max_output_tokens up to the OpenAI Responses API minimum. + + OpenAI's Responses API rejects max_output_tokens below 16 for every model + (not gpt-5 specific), so a client like Claude Code that sends a max_tokens=1 + warmup probe on model switch would otherwise 400. Values that are None or + already at/above the minimum are returned unchanged. + """ + if isinstance(max_output_tokens, int) and max_output_tokens < OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS: + return OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS + return max_output_tokens + def get_supported_openai_params(self, model: str) -> list: """ All OpenAI Responses API params are supported @@ -92,6 +107,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): """ params = dict(response_api_optional_params) + if "max_output_tokens" in params: + params["max_output_tokens"] = self._enforce_min_max_output_tokens(params.get("max_output_tokens")) + if self._is_gpt_5_model(model=model): temperature = params.get("temperature") if temperature is not None and temperature != 1: diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index f9ed8cea9b5..8c4bb1aa0c5 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -998,6 +998,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): response_modalities.append("IMAGE") elif modality == "audio": response_modalities.append("AUDIO") + elif modality == "video": + response_modalities.append("VIDEO") else: response_modalities.append("MODALITY_UNSPECIFIED") return response_modalities 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/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b4ae842c4e6..77ca423b866 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11331,6 +11331,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11362,6 +11363,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true @@ -11424,6 +11426,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, @@ -11479,6 +11482,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11506,6 +11510,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11559,6 +11564,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true @@ -11586,6 +11592,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true @@ -11614,6 +11621,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "provider_specific_entry": { @@ -11648,6 +11656,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "provider_specific_entry": { @@ -11682,6 +11691,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, @@ -11718,6 +11728,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, @@ -11788,6 +11799,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, @@ -18747,6 +18759,49 @@ }, "supports_image_size": false }, + "gemini/gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "rpm": 1000, + "tpm": 4000000, + "output_cost_per_token_batches": 6e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -18778,6 +18833,49 @@ ], "supports_function_calling": false, "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, + "gemini/gemini-3.1-flash-image": { + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.045, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, "supports_response_schema": true, "supports_system_messages": true, "supports_vision": true, @@ -18819,6 +18917,7 @@ ], "supports_function_calling": false, "supports_prompt_caching": true, + "supports_reasoning": false, "supports_response_schema": true, "supports_system_messages": true, "supports_vision": true, @@ -19513,6 +19612,39 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-omni-flash-preview": { + "input_cost_per_audio_token": 1.5e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true, + "tpm": 800000 + }, "gemini/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -19677,6 +19809,37 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-omni-flash-preview": { + "input_cost_per_audio_token": 1.5e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/omni-flash-preview", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.5-flash": { "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, @@ -36777,6 +36940,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -36786,6 +36950,22 @@ "tpm": 8000000, "supports_image_size": false }, + "vertex_ai/gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "supports_reasoning": false, + "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" + }, "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -36799,8 +36979,23 @@ "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, + "supports_reasoning": false, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, + "vertex_ai/gemini-3.1-flash-image": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "supports_reasoning": false, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" + }, "vertex_ai/gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, @@ -36812,6 +37007,7 @@ "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, "output_cost_per_token": 3e-06, + "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, "vertex_ai/gemini-3.1-flash-lite-preview": { @@ -45051,8 +45247,8 @@ "rules": [ { "name": "bedrock-claude-ids", - "pattern": "anthropic\\.claude-", - "description": "Any Bedrock-syntax Claude id: the dotted anthropic.claude- segment appears in bare (anthropic.claude-...), region-prefixed (us./eu./au./jp./apac.) and global.-prefixed ids, for every version. Routes these to bedrock before the bare-id Anthropic rule is consulted.", + "pattern": "^(?:[a-z-]+\\.)?anthropic\\.claude-", + "description": "A Bedrock-syntax Claude id, for every version: anthropic.claude- at the start of the name, optionally behind a single dotted geo segment (us./eu./au./jp./apac./global./us-gov.). Anchored to the start because routing rules see the raw request string and provider inference feeds the proxy's provider/* wildcard access checks: an id under an unrecognized namespace such as bedrockz/anthropic.claude-... must stay unroutable rather than resolve to bedrock and slip through a bedrock/* key. Routes to bedrock before the bare-id Anthropic rule is consulted.", "model_info": { "litellm_provider": "bedrock" } diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index d67726be584..519066b8266 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -36,6 +36,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): budget_reset_at: Optional[datetime] = None allowed_cache_controls: Optional[list] = [] allowed_routes: Optional[list] = [] + key_type: str | None = None permissions: Dict = {} model_spend: Dict = {} model_max_budget: Dict = {} 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..421f1dcfbea 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,89 @@ 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`` 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, but ``get_user_object`` defeats a + type-based check: where ``get_key_object`` raises a typed ``ProxyException`` for a missing key + and lets a DB outage propagate raw, ``get_user_object`` catches every DB failure and re-raises a + bare ``ValueError``, so a missing user and a real outage look identical and the original error + survives only as ``__context__``. ``_raise_503_if_db_unavailable`` therefore walks the cause + chain: a transient DB outage still surfaces as a retryable 503, while a missing user, or any + other non-outage resolution failure, fails closed as a 401 rather than an opaque 500. The + object-permission load shares this one boundary, so an outage there is classified the same + way (``get_object_permission`` itself swallows a failed load to ``None``, matching how + ``get_key_object`` best-effort-loads a key's object permission).""" + 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: + 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, + ) + # 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 is not None else None + if user_object is not None and object_permission is None and user_object.object_permission_id: + 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, + ) + except (ProxyException, HTTPException): + raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + except Exception as e: # noqa: BLE001 # a DB outage anywhere in the resolution is a retryable 503, not an opaque 500; anything else fails closed as 401 + 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: + raise HTTPException(status_code=401, detail="Invalid or expired credential") + 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: """Reload the live key record an admitted envelope references and re-check live policy. @@ -615,10 +701,14 @@ class MCPRequestHandler: """Raise a retryable 503 when ``e`` means the auth database is unreachable, else return so the caller applies its own fail-closed mapping. A DB outage must not masquerade as an auth failure (401) or surface as an opaque 500; the caller retries. Mirrors ``UserAPIKeyAuthExceptionHandler``, - which renders a service-unavailable database error as 503 on the standard pipeline.""" + which renders a service-unavailable database error as 503 on the standard pipeline. + + Classifies across the ``__cause__``/``__context__`` chain, not just ``e`` itself: ``get_user_object`` + re-raises every DB failure as a bare ``ValueError``, so a type-based check on the top exception + would miss a real outage wrapped inside it.""" from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler - if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): raise HTTPException( status_code=503, detail="Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py new file mode 100644 index 00000000000..19048e2eb7c --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -0,0 +1,694 @@ +"""Bridge token flow: litellm identity resolution and the DCR-bridge oauth_delegate mint/refresh pipeline.""" + +import math +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Literal, Optional + +from fastapi import HTTPException, Request +from fastapi.responses import JSONResponse +from pydantic import SecretStr +from typing_extensions import assert_never + +from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +if TYPE_CHECKING: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _BridgeAuthorizationCode + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeIdentity, + EnvelopeKeys, + RefreshCredential, + UpstreamTokenGrant, + ) + from litellm.proxy._types import UserAPIKeyAuth + + +def _litellm_key_from_request(request: Request) -> Optional[str]: + """Return the LiteLLM API key presented on the request, or ``None``. + + Accepts the key from ``x-litellm-api-key`` (what MCP clients such as Claude Desktop/Code + send) as well as ``Authorization``; either may carry a bare token or ``Bearer ``. + ``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry + an OAuth/upstream bearer. + """ + for header_value in ( + request.headers.get("x-litellm-api-key"), + request.headers.get("Authorization") or request.headers.get("authorization"), + ): + if not header_value: + continue + value = header_value.strip() + if value.lower().startswith("bearer "): + value = value[7:].strip() + if value: + return value + return None + + +def _key_is_active(key_obj: "UserAPIKeyAuth") -> bool: + """``True`` when the presented key is neither blocked nor past its expiry. + + The OAuth token endpoint is unauthenticated, so the presented key is validated here before it is + trusted; a revoked or expired key must not mint a bridge envelope or write a stored credential. + ``get_key_object`` resolves a row without these checks (the main ``user_api_key_auth`` pipeline + enforces them downstream, which this endpoint bypasses), so they are applied here. Deleted keys + are already rejected upstream, where ``get_key_object`` raises on a row that no longer exists. + + This is an active-state gate only; it deliberately does not require a ``user_id``. A valid + team-scoped or service-account key has no ``user_id`` yet is a legitimate credential, so gating + on ``user_id`` presence would wrongly reject it. Callers that need the user (the per-user token + store) derive it separately via :func:`_active_key_user_id`. + + Total by design: ``expires`` is typed ``str | datetime``, and an unparseable string would make + ``datetime.fromisoformat`` raise. Since the callers run this outside their key-resolution + ``try``, an uncaught parse error would surface as a 500 instead of the endpoint's fail-closed + behavior, so a malformed expiry is treated as inactive (return ``False``) rather than raising. + """ + if key_obj.blocked is True: + return False + expires = key_obj.expires + if expires is not None: + if isinstance(expires, datetime): + expiry = expires + else: + try: + expiry = datetime.fromisoformat(expires) + except (ValueError, TypeError): + return False + if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None: + expiry = expiry.replace(tzinfo=timezone.utc) + if expiry < datetime.now(timezone.utc): + return False + return True + + +def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> str | None: + """The active key's ``user_id``, or ``None`` when the key is blocked/expired or simply has no + ``user_id`` (a team-scoped or service-account key). Used only by the per-user token store, which + needs a user to key the stored credential; the bridge mint uses the key hash and does not.""" + return key_obj.user_id if _key_is_active(key_obj) else None + + +@dataclass(frozen=True, slots=True) +class _ResolvedKey: + """An active litellm key resolved from the token request: its hash (the value ``get_key_object`` + and the cache/DB layer key the record by) and the live record.""" + + key_hash: str + key: "UserAPIKeyAuth" + + +_KeyResolutionFailure = Literal["no_active_key", "unavailable", "unresolvable"] +"""Why a token request yielded no active litellm key, kept distinct so a caller statuses each truthfully +instead of blaming the client for a gateway problem: +- ``no_active_key``: none was presented, or the presented key is unknown / blocked / expired (the + caller's request is at fault) +- ``unavailable``: the auth database was transiently unreachable while resolving (retryable) +- ``unresolvable``: the gateway cannot resolve identity right now (no DB connection, or an unexpected + error) -- a gateway fault, not the caller's +The classification mirrors admission's ``_reload_admitted_key`` so the mint (ingress) and admission +(egress) never disagree on the status of the same outage.""" + + +async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyResolutionFailure": + """Resolve the presented litellm key to an active key record, or say precisely why not. + + Single resolution path the OAuth token endpoint reuses, resolving authoritatively via + ``get_key_object`` (cache first, then DB). The failure is a value, not a bare ``None``, so a caller + can tell "the client sent no usable credential" (a request error) apart from "the gateway could not + check" (an infrastructure error) and status each truthfully; collapsing both to ``None`` is what let + a DB outage read as a 400. A resolved key is still gated by ``_key_is_active``, so a blocked or + expired key is ``no_active_key`` while a valid team-scoped or service-account key (no ``user_id``) + resolves. Classification mirrors admission's ``_reload_admitted_key``: no DB connection is a gateway + fault, a ``ProxyException`` / ``HTTPException`` from ``get_key_object`` is an unknown or invalid key, + a database-service-unavailable error is a retryable outage, and anything else is an unexpected + gateway fault.""" + token = _litellm_key_from_request(request) + if not token: + return "no_active_key" + from litellm.proxy._types import hash_token # noqa: PLC0415 # inline import avoids a module-load circular import + + return await _reload_active_key_by_hash(hash_token(token)) + + +async def _reload_active_key_by_hash(key_hash: str) -> "_ResolvedKey | _KeyResolutionFailure": + """Reload the live key record for ``key_hash`` (cache first, then DB) and gate it on active state, + returning the resolved key or a precise failure. Shared by the token request's presented-key + resolution (:func:`_resolve_active_litellm_key`, which hashes the presented key) and the refresh + path (which already holds the hash sealed in the refresh envelope), so both re-validate identity + through one active-key gate and one failure classification. Classification mirrors admission's + ``_reload_admitted_key``: no DB connection is a gateway fault, a ``ProxyException`` / ``HTTPException`` + from ``get_key_object`` is an unknown or invalid key, a database-service-unavailable error is a + retryable outage, and anything else is an unexpected gateway fault. A blocked or expired key is + ``no_active_key``, so a revoked key can neither mint nor refresh a bridge envelope.""" + from litellm.proxy._types import ( + ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import + ) + from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import + get_key_object, + ) + from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import + PrismaDBExceptionHandler, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + return "unresolvable" + try: + key_obj = await get_key_object( + hashed_token=key_hash, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + except (ProxyException, HTTPException): + return "no_active_key" + except Exception as exc: # noqa: BLE001 # classify: a DB outage is retryable, anything else is an opaque gateway fault + if PrismaDBExceptionHandler.is_database_service_unavailable_error(exc): + return "unavailable" + verbose_logger.debug( + "_reload_active_key_by_hash: unexpected key-resolution error (%s)", + type(exc).__name__, + ) + return "unresolvable" + if not _key_is_active(key_obj): + return "no_active_key" + return _ResolvedKey(key_hash=key_hash, key=key_obj) + + +async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | None": + """Re-validate a live litellm user by id, returning ``None`` when the user is active or a precise + failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a + user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a + deactivated user cannot keep refreshing, mirroring how admission re-validates the same user subject on + the egress side. No DB connection is a gateway fault (``unresolvable``) and a + database-service-unavailable error is a retryable outage (``unavailable``). Everything else fails + closed as ``no_active_key`` (the caller maps it to invalid_grant): a ``ProxyException`` / + ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object`` + catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look + identical, the original error surviving only as ``__context__``), so the outage check walks the cause + chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault.""" + from litellm.proxy._types import ( + ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import + ) + from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import + get_user_object, + ) + from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import + PrismaDBExceptionHandler, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + return "unresolvable" + 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): + return "no_active_key" + except Exception as exc: # noqa: BLE001 # a DB outage is retryable; a missing user (get_user_object's wrapped ValueError) or any other resolution failure fails closed as no_active_key, never a 500 + if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(exc): + return "unavailable" + verbose_logger.debug("_reload_active_user_by_id: user-resolution error (%s)", type(exc).__name__) + return "no_active_key" + if user_object is None: + return "no_active_key" + if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: + return "no_active_key" + return None + + +async def _key_owner_scim_deactivated(key: "UserAPIKeyAuth") -> bool: + """True only when the key's owning user was explicitly SCIM-deactivated, so a refresh revokes an + offboarded owner's key exactly as admission does via ``_reject_if_admitted_owner_scim_deactivated``. + A key with no owner, a missing owner record, or a failed lookup fails OPEN (returns ``False``), + matching admission and the standard builder: a key may outlive its owner record, and a transient DB + blip must not revoke a live key. Only an explicit ``scim_active`` of ``False`` gates renewal.""" + if key.user_id is None: + return False + from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import + get_user_object, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + return False + try: + owner = await get_user_object( + user_id=key.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + except Exception as exc: # noqa: BLE001 # fail open: a missing owner (get_user_object's wrapped ValueError) or a DB blip must not revoke a live key + verbose_logger.debug("refresh: key-owner SCIM lookup failed, not revoking (%s)", type(exc).__name__) + return False + return owner is not None and isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False + + +async def _revalidate_active_subject(identity: "EnvelopeIdentity") -> "_KeyResolutionFailure | None": + """Re-validate that the subject sealed in a refresh envelope is still live, dispatching on its type: + a key_hash reloads the virtual key, a user_id reloads the user. Returns ``None`` when the subject is + active or a precise failure otherwise, so revocation gates renewal for either identity source the same + way admission gates the egress: a blocked or expired key, a SCIM-deactivated key owner (mirroring + admission's owner check, so an offboarded user cannot keep renewing a still-active key), and a + deactivated or deleted user all fail closed to ``no_active_key``.""" + match identity.subject_type: + case "key_hash": + reloaded = await _reload_active_key_by_hash(identity.subject) + if not isinstance(reloaded, _ResolvedKey): + return reloaded + if await _key_owner_scim_deactivated(reloaded.key): + return "no_active_key" + return None + case "user_id": + return await _reload_active_user_by_id(identity.subject) + case _: + assert_never(identity.subject_type) + + +async def _extract_user_id_from_request(request: Request) -> str | None: + """The litellm ``user_id`` for the token request, so a per-user token is stored under the same + identity the egress later reads it by. Storage is best-effort, so every non-resolved outcome + (including a transient DB outage) collapses to ``None`` here and the caller simply skips the store; + the bridge mint, which must status those outcomes differently, consumes + :func:`_resolve_active_litellm_key` directly.""" + resolved = await _resolve_active_litellm_key(request) + if not isinstance(resolved, _ResolvedKey): + return None + return _active_key_user_id(resolved.key) + + +_UpstreamGrantRejection = Literal["no_access_token", "expired_lifetime"] +"""Why an upstream token response cannot back a bridge envelope: +- ``no_access_token``: the response carries no usable ``access_token`` +- ``expired_lifetime``: the response reports a parseable, non-positive ``expires_in``, i.e. an upstream + token that is already dead, so sealing it would forward a bearer the edge cannot use +An absent or unparseable ``expires_in`` is NOT a rejection; the lifetime is merely unknown and the +envelope caps it, the by-design behaviour for an upstream that omits the field.""" + + +def _classify_upstream_lifetime(raw_expires_in: object) -> "int | Literal['unspecified', 'expired']": + """Classify an upstream ``expires_in`` into a positive number of seconds, ``"unspecified"`` (absent + or unparseable, so the envelope caps it), or ``"expired"`` (a non-positive value the upstream reports + as already elapsed). Telling "we do not know the lifetime" apart from "the upstream says it is + already dead" is what stops an explicitly-expired token from silently receiving the envelope's 1h + cap. The expired decision is made on the parsed numeric value, not on ``int(...)`` of it, so a + positive sub-second lifetime in ``(0, 1)`` is not truncated to ``0`` and misread as elapsed; the + envelope works in whole seconds, so such a lifetime clamps up to its 1s floor. ``bool`` is excluded + (an ``int`` subclass but never a real lifetime), and the conversions can raise on ``NaN`` / + ``Infinity`` / oversized input, which reads as unparseable rather than surfacing as a 500.""" + if raw_expires_in is None or isinstance(raw_expires_in, bool) or not isinstance(raw_expires_in, (int, float, str)): + return "unspecified" + try: + numeric = float(raw_expires_in) + seconds = int(numeric) + except (ValueError, TypeError, OverflowError): + return "unspecified" + if numeric <= 0: + return "expired" + return max(1, seconds) + + +def _bridge_grant_from_token_response(token_response: object) -> "UpstreamTokenGrant | _UpstreamGrantRejection": + """Validate an upstream OAuth token response into a typed grant, or say why it cannot back an + envelope. Each field is isinstance-checked so nothing untyped from ``response.json()`` reaches the + grant. ``expires_in`` is read three ways (see :func:`_classify_upstream_lifetime`): an unknown + lifetime leaves the grant ``expires_in`` ``None`` for the envelope to cap, a positive value is + honoured, and an explicit already-elapsed value is a rejection rather than a silent fall-through to + the cap.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + UpstreamTokenGrant, + ) + + if not isinstance(token_response, dict): + return "no_access_token" + access = token_response.get("access_token") + if not isinstance(access, str) or not access: + return "no_access_token" + lifetime = _classify_upstream_lifetime(token_response.get("expires_in")) + if lifetime == "expired": + return "expired_lifetime" + token_type = token_response.get("token_type") + scope = token_response.get("scope") + return UpstreamTokenGrant( + access_token=SecretStr(access), + token_type=token_type if isinstance(token_type, str) and token_type else "Bearer", + # The upstream refresh_token is deliberately NOT sealed: the edge never consumes it (it forwards + # only token_type + access_token), so it would be dead weight embedding a long-lived upstream + # credential in the client-held bearer, and it enlarges the envelope. Refresh support is a + # follow-up (a dedicated refresh-envelope); the client re-runs authorization_code at the cap. + refresh_token=None, + scope=scope if isinstance(scope, str) and scope else None, + expires_in=lifetime if isinstance(lifetime, int) else None, + ) + + +# --------------------------------------------------------------------------- +# DCR-bridge oauth_delegate mint: a three-phase pipeline whose failures are values. +# +# prepare (before the upstream exchange) -> validate every precondition and resolve identity+keys +# exchange (the single-use upstream code is consumed here, in exchange_token_with_server) +# finish (after the exchange) -> seal the upstream grant into the client-held envelope +# +# Every precondition lives in ``prepare``, which runs BEFORE the exchange, so no failure can burn the +# single-use code or rotate a refresh token, for either grant type -- that whole class of bug is gone +# by construction rather than guarded case by case. Failures are values mapped to an OAuth-shaped +# response in one place (``_bridge_mint_error_response``), so status codes and the RFC 6749 §5.2 body +# shape are uniform. Adding a failure mode is a new literal plus a match arm the type checker forces. +# --------------------------------------------------------------------------- + +_BridgeMintError = Literal[ + "no_identity", + "invalid_refresh", + "identity_unavailable", + "identity_unresolvable", + "not_configured", + "no_upstream_token", + "upstream_token_expired", + "too_large", +] + + +@dataclass(frozen=True, slots=True) +class _BridgeMintReady: + """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.""" + + identity: "EnvelopeIdentity" + keys: "EnvelopeKeys" + + +def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse: + """Map a bridge-mint failure value to its token-endpoint response: one place, RFC 6749 §5.2 shape + (top-level ``error``, no-store headers) for every case, with a status truthful about where the + failure is. The caller's request is 400, a transient gateway outage is 503, a gateway + misconfiguration is 500, and an upstream problem is 502. The identity-resolution statuses match how + admission statuses the same conditions on the egress side, so mint and admit never disagree under + one outage.""" + match error: + case "no_identity": + status, code, desc = ( + 400, + "invalid_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 "invalid_refresh": + status, code, desc = ( + 400, + "invalid_grant", + "the refresh credential is not a valid, live refresh envelope for this server; " + "re-run authorization_code to obtain a new one", + ) + case "identity_unavailable": + status, code, desc = ( + 503, + "temporarily_unavailable", + "the authentication database is temporarily unreachable; retry shortly", + ) + case "identity_unresolvable": + status, code, desc = ( + 500, + "server_error", + "the gateway could not resolve the litellm identity for this request", + ) + case "not_configured": + status, code, desc = ( + 500, + "server_error", + "the gateway is not configured to mint a gateway-bound credential (master_key is not set)", + ) + case "no_upstream_token": + status, code, desc = ( + 502, + "server_error", + "the upstream token response has no usable access_token", + ) + case "upstream_token_expired": + status, code, desc = ( + 502, + "server_error", + "the upstream token response reports an already-expired lifetime", + ) + case "too_large": + status, code, desc = ( + 502, + "server_error", + "the upstream token is too large to seal into a gateway-bound credential", + ) + case _: + assert_never(error) + return JSONResponse( + status_code=status, content={"error": code, "error_description": desc}, headers=TOKEN_NO_CACHE_HEADERS + ) + + +def _key_resolution_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError: + """Lift an identity-resolution failure into the mint taxonomy, preserving origin so the status stays + truthful: the caller's missing credential is 400, a transient DB outage is 503, and a gateway that + cannot resolve identity is 500.""" + match failure: + case "no_active_key": + return "no_identity" + case "unavailable": + return "identity_unavailable" + case "unresolvable": + return "identity_unresolvable" + case _: + assert_never(failure) + + +def _upstream_rejection_to_mint_error(rejection: _UpstreamGrantRejection) -> _BridgeMintError: + """Lift an upstream-response rejection into the mint taxonomy; both are upstream faults (502).""" + match rejection: + case "no_access_token": + return "no_upstream_token" + case "expired_lifetime": + return "upstream_token_expired" + case _: + assert_never(rejection) + + +async def _prepare_bridge_mint( + request: Request, + mcp_server: MCPServer, + bridge_identity: "_BridgeAuthorizationCode | None" = None, +) -> "_BridgeMintReady | _BridgeMintError": + """Phase 1 for the authorization_code grant, BEFORE the upstream exchange: 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. + + 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``. The refresh_token grant has its own phase-1 + (:func:`_prepare_bridge_refresh`), which recovers identity from the presented refresh envelope.""" + 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, + ) + + 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) + identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=resolved.key_hash) + return _BridgeMintReady(identity=identity, keys=keys) + + +@dataclass(frozen=True, slots=True) +class _BridgeRefreshReady: + """A validated refresh request: the identity+keys to mint the renewed pair under, the upstream refresh + token (unwrapped from the client's refresh envelope) to exchange with the upstream IdP, and the scope + sealed alongside it at mint. The upstream refresh token is a ``SecretStr`` like every other credential + in this layer, so a repr or a traceback that captures this value never exposes the raw upstream refresh + token in plaintext. ``upstream_scope`` carries the originally-granted scope so the renewal re-requests + it when the client (a DCR/MCP client that typically omits scope on refresh) sends none, keeping the + renewed token's scope stable against an upstream that would otherwise narrow or drop it.""" + + ready: "_BridgeMintReady" + upstream_refresh_token: SecretStr + upstream_scope: str | None = None + + +def _refresh_key_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError: + """Lift an identity-resolution failure on the refresh path into the mint taxonomy. Unlike the mint + path, a resolved-but-inactive (or unknown) key is ``invalid_grant`` rather than ``invalid_request``: + the client did present an identity (sealed in the refresh envelope), but it is no longer live, so the + refresh is invalid and the client must re-authenticate. A transient outage is still 503 and a gateway + fault still 500, matching the mint path and admission.""" + match failure: + case "no_active_key": + return "invalid_refresh" + case "unavailable": + return "identity_unavailable" + case "unresolvable": + return "identity_unresolvable" + case _: + assert_never(failure) + + +async def _prepare_bridge_refresh( + mcp_server: MCPServer, refresh_value: str | None +) -> "_BridgeRefreshReady | _BridgeMintError": + """Phase 1 for the refresh_token grant, BEFORE the upstream exchange: open the client's refresh + envelope, re-validate the sealed litellm identity so a revoked key cannot keep refreshing, and + recover the upstream refresh token to exchange. Identity comes entirely from the sealed envelope, not + the HTTP request, so the request object is not needed here. The client presents a refresh envelope, + never a raw upstream refresh token, so a missing value, a non-envelope, an unopenable envelope, or one + minted for another server is ``invalid_grant``. Running before the exchange means a rejected refresh + never consumes or rotates the upstream refresh token.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import + BridgeRefreshOpened, + envelope_keys_from_master_key, + open_bridge_refresh_envelope, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + master_key, + ) + + if not master_key: + return "not_configured" + if not refresh_value: + return "invalid_refresh" + keys = envelope_keys_from_master_key(master_key) + opened = open_bridge_refresh_envelope(refresh_value, keys, datetime.now(timezone.utc), mcp_server.server_id) + if not isinstance(opened, BridgeRefreshOpened): + return "invalid_refresh" + failure = await _revalidate_active_subject(opened.identity) + if failure is not None: + return _refresh_key_failure_to_mint_error(failure) + return _BridgeRefreshReady( + ready=_BridgeMintReady(identity=opened.identity, keys=keys), + upstream_refresh_token=opened.refresh.refresh_token, + upstream_scope=opened.refresh.scope, + ) + + +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 access envelope + using the pre-resolved identity and keys, and, when the upstream returned a refresh token, seal a + long-lived refresh envelope alongside it so the client can renew without re-authenticating. Shared by + the authorization_code and refresh_token paths, so a renewal that the upstream rotates re-issues a + fresh refresh envelope. The only hard failures here are properties of the upstream access token (no + usable token, an already-expired lifetime, or a token too large to seal); a refresh token that cannot + be sealed degrades to an access-only response rather than failing the whole exchange.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import + build_bridge_token_response, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + SealedEnvelope, + UpstreamTokenGrant, + ) + + grant = _bridge_grant_from_token_response(token_response) + if not isinstance(grant, UpstreamTokenGrant): + return _upstream_rejection_to_mint_error(grant) + 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 + # client is never told the bearer lives past the point admission (which uses that exp) rejects it. + expires_in = max(0, int(sealed.expires_at.timestamp()) - math.ceil(now.timestamp())) + refresh_envelope = _mint_refresh_envelope_value(ready.identity, token_response, ready.keys, now, mcp_server) + body = { + "access_token": sealed.token.get_secret_value(), + "token_type": "Bearer", + "expires_in": expires_in, + # A refresh envelope rides along only when the upstream returned a refresh token to seal; when it + # rotates on renewal, the client receives the new one and the old envelope's upstream token dies. + **({"refresh_token": refresh_envelope} if refresh_envelope is not None else {}), + } + return JSONResponse(body, headers=TOKEN_NO_CACHE_HEADERS) + + +def _upstream_refresh_credential(token_response: object) -> "RefreshCredential | None": + """Extract the upstream refresh grant from a token response, or ``None`` when there is none to seal. + Each field is isinstance-checked so nothing untyped reaches the refresh envelope; ``refresh_expires_in`` + (the refresh token's own lifetime, when the upstream reports it) is classified like ``expires_in`` and + bounds the refresh envelope's TTL. An upstream that reports the refresh token itself as already elapsed + (``refresh_expires_in`` non-positive) yields ``None`` rather than a refresh envelope: sealing a dead + token would hand the client a full-TTL-capped envelope the IdP will reject, so the exchange degrades to + an access-only response (the client re-authenticates at access expiry), mirroring how + :func:`_bridge_grant_from_token_response` refuses an already-elapsed access token instead of capping it.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + RefreshCredential, + ) + + if not isinstance(token_response, dict): + return None + refresh = token_response.get("refresh_token") + if not isinstance(refresh, str) or not refresh: + return None + lifetime = _classify_upstream_lifetime(token_response.get("refresh_expires_in")) + if lifetime == "expired": + return None + scope = token_response.get("scope") + return RefreshCredential( + refresh_token=SecretStr(refresh), + scope=scope if isinstance(scope, str) and scope else None, + expires_in=lifetime if isinstance(lifetime, int) else None, + ) + + +def _mint_refresh_envelope_value( + identity: "EnvelopeIdentity", token_response: object, keys: "EnvelopeKeys", now: datetime, mcp_server: MCPServer +) -> str | None: + """Seal the upstream refresh grant (if any) into a refresh envelope and return its bearer string, or + ``None`` when the upstream returned no refresh token or the refresh token is too large to seal. A + too-large refresh token degrades to an access-only response (logged) rather than failing an exchange + that already succeeded upstream: the client simply re-authenticates when the access envelope expires.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import + build_bridge_refresh_token_response, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + SealedEnvelope, + ) + + refresh_credential = _upstream_refresh_credential(token_response) + if refresh_credential is None: + return None + sealed = build_bridge_refresh_token_response(identity, refresh_credential, keys, now) + if isinstance(sealed, SealedEnvelope): + return sealed.token.get_secret_value() + verbose_logger.warning( + "bridge mint: the upstream refresh token is too large to seal into a refresh envelope for " + "server=%s; issuing an access-only response, so the client re-authenticates at access expiry", + mcp_server.server_id, + ) + return None diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ebceb320906..54aff86aab2 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -10,7 +10,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, ValidationError +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -21,6 +21,24 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( TokenEndpointAuthConfigError, build_token_endpoint_client_auth, ) +from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _bridge_mint_error_response, + _BridgeMintReady, + _BridgeRefreshReady, + _extract_user_id_from_request, + _finish_bridge_mint, + _prepare_bridge_mint, + _prepare_bridge_refresh, +) +from litellm.proxy._experimental.mcp_server.faults import ( + CallerRejected, + CredentialSource, + UpstreamProtocolFault, + classify_upstream_dcr_rejection, + classify_upstream_token_rejection, + dcr_fault_detail, + render_token_fault, +) from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, @@ -37,7 +55,7 @@ from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer if TYPE_CHECKING: - from litellm.proxy._types import LiteLLM_MCPServerTable, UserAPIKeyAuth + from litellm.proxy._types import LiteLLM_MCPServerTable # TTL cache for upstream OAuth metadata fetched from pass-through MCP servers. # Keeps us from hammering the upstream IdP on each discovery request. @@ -91,6 +109,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. @@ -101,6 +121,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 @@ -111,6 +136,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) @@ -138,6 +165,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 @@ -304,90 +393,6 @@ def _validate_token_response( ) -def _litellm_key_from_request(request: Request) -> Optional[str]: - """Return the LiteLLM API key presented on the request, or ``None``. - - Accepts the key from ``x-litellm-api-key`` (what MCP clients such as Claude Desktop/Code - send) as well as ``Authorization``; either may carry a bare token or ``Bearer ``. - ``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry - an OAuth/upstream bearer. - """ - for header_value in ( - request.headers.get("x-litellm-api-key"), - request.headers.get("Authorization") or request.headers.get("authorization"), - ): - if not header_value: - continue - value = header_value.strip() - if value.lower().startswith("bearer "): - value = value[7:].strip() - if value: - return value - return None - - -def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> Optional[str]: - """The key's ``user_id``, or ``None`` if the key is blocked or expired. - - The OAuth token endpoint is unauthenticated, so the presented key is validated here before its - identity is trusted to key a stored credential; a revoked or expired key must not be able to - write or overwrite the per-user OAuth token. ``get_key_object`` resolves a row without these - checks (the main ``user_api_key_auth`` pipeline enforces them downstream, which this endpoint - bypasses), so they are applied here. Deleted keys are already rejected upstream, where - ``get_key_object`` raises on a row that no longer exists. - """ - if key_obj.blocked is True: - return None - expires = key_obj.expires - if expires is not None: - expiry = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires) - if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None: - expiry = expiry.replace(tzinfo=timezone.utc) - if expiry < datetime.now(timezone.utc): - return None - return key_obj.user_id - - -async def _extract_user_id_from_request(request: Request) -> Optional[str]: - """Resolve the LiteLLM ``user_id`` at the OAuth token endpoint so a per-user token is stored - under the same identity the egress later reads it by (``user_api_key_auth.user_id``). - - Resolves authoritatively via ``get_key_object`` (cache first, then DB) instead of a raw cache - peek. On a multi-replica gateway the token-exchange request can land on a worker whose in-memory - cache never saw the key, and a cross-replica Redis hit deserializes to a plain ``dict`` rather - than a ``UserAPIKeyAuth``; the previous code read only ``Authorization`` and did - ``getattr(cached, "user_id")`` with no ``model_type`` rehydration and no DB fallback, so it - silently returned ``None`` and the token was never persisted, which makes the egress 401 on every - reconnect. The resolved key is validated (``_active_key_user_id``) before its identity is trusted, - so a blocked or expired key cannot write. Returns ``None`` when no key is present, the key cannot - be resolved, or it is blocked/expired. - """ - token = _litellm_key_from_request(request) - if not token: - return None - try: - from litellm.proxy._types import hash_token # noqa: PLC0415 - from litellm.proxy.auth.auth_checks import get_key_object # noqa: PLC0415 - from litellm.proxy.proxy_server import ( # noqa: PLC0415 - prisma_client, - user_api_key_cache, - ) - - key_obj = await get_key_object( - hashed_token=hash_token(token), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - return _active_key_user_id(key_obj) - except Exception as exc: - verbose_logger.debug( - "_extract_user_id_from_request: could not resolve a LiteLLM user_id for the presented " - "key (%s); per-user token will not be stored server-side.", - type(exc).__name__, - ) - return None - - async def _store_per_user_token_server_side( server: MCPServer, user_id: str, @@ -620,12 +625,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) @@ -654,6 +678,13 @@ async def authorize_with_server( return response +def _token_credential_source(mcp_server: MCPServer) -> CredentialSource: + """Mirrors the resolved-client rule in :func:`exchange_token_with_server`: when the server has a + stored client_id the gateway presents its own credentials upstream, so a credential rejection is + the operator's fault, not the caller's.""" + return "gateway_stored" if mcp_server.client_id else "caller_supplied" + + async def exchange_token_with_server( request: Request, mcp_server: MCPServer, @@ -688,25 +719,61 @@ 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 + bridge_mint_ready: _BridgeMintReady | None = None + bridge_upstream_refresh: SecretStr | None = None + bridge_upstream_scope: str | None = None + refresh_request_scope: str | None = None + is_bridge = mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge + if grant_type == "refresh_token": - if not refresh_token: + # Phase 1 for a bridge refresh: open the client's refresh envelope, re-validate the sealed + # identity, and unwrap the real upstream refresh token BEFORE building token_data, so the exchange + # sends the upstream token and never the envelope. A failure returns without touching the upstream. + if is_bridge: + prepared_refresh = await _prepare_bridge_refresh(mcp_server, refresh_token) + if not isinstance(prepared_refresh, _BridgeRefreshReady): + return _bridge_mint_error_response(prepared_refresh) + bridge_mint_ready = prepared_refresh.ready + bridge_upstream_refresh = prepared_refresh.upstream_refresh_token + bridge_upstream_scope = prepared_refresh.upstream_scope + # A bridge server sends the unwrapped upstream refresh token recovered from the client's refresh + # envelope above; every other server sends the client's own refresh token verbatim. + upstream_refresh_token = ( + bridge_upstream_refresh.get_secret_value() if bridge_upstream_refresh is not None else refresh_token + ) + if not upstream_refresh_token: raise HTTPException( status_code=400, detail="refresh_token is required for refresh_token grant", ) token_data: dict = { "grant_type": "refresh_token", - "refresh_token": refresh_token, + "refresh_token": upstream_refresh_token, **client_auth.body, } - if scope: - token_data["scope"] = scope + refresh_request_scope = scope or bridge_upstream_scope + if refresh_request_scope: + token_data["scope"] = refresh_request_scope else: if not code: raise HTTPException( 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( @@ -726,32 +793,49 @@ async def exchange_token_with_server( } if code_verifier: token_data["code_verifier"] = code_verifier - + # Phase 1 for a bridge authorization_code mint: resolve identity (the SSO user recovered above, or + # the presented litellm key) and the envelope keys BEFORE the exchange consumes the single-use code. + if is_bridge: + prepared = await _prepare_bridge_mint(request, mcp_server, bridge_identity) + if not isinstance(prepared, _BridgeMintReady): + return _bridge_mint_error_response(prepared) + bridge_mint_ready = prepared async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) - response = await async_client.post( - mcp_server.token_url, - headers={"Accept": "application/json", **client_auth.headers}, - data=token_data, - ) + try: + response = await async_client.post( + mcp_server.token_url, + headers={"Accept": "application/json", **client_auth.headers}, + data=token_data, + ) + if response is not None: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + fault = classify_upstream_token_rejection( + exc.response, + credential_source=_token_credential_source(mcp_server), + log_context=mcp_server.server_id, + ) + upstream_rejected_bridge_refresh = ( + is_bridge + and grant_type == "refresh_token" + and isinstance(fault, CallerRejected) + and fault.code == "invalid_grant" + ) + if upstream_rejected_bridge_refresh: + verbose_logger.info( + "bridge refresh: the upstream rejected the sealed refresh token for server=%s with " + "invalid_grant (revoked or expired at the IdP); returning invalid_grant so the client " + "re-runs authorization_code rather than an opaque upstream error", + mcp_server.server_id, + ) + return _bridge_mint_error_response("invalid_refresh") + return render_token_fault(fault) if response is None: raise HTTPException( status_code=502, detail="MCP upstream token endpoint returned no response", ) - - try: - response.raise_for_status() - except httpx.HTTPStatusError as exc: - if "invalid_target" in exc.response.text: - verbose_logger.warning( - "MCP server %s: the upstream authorization server rejected the token request with " - "invalid_target; it may require RFC 8707 resource indicators, which the gateway " - "does not send yet (tracked as LIT-4339)", - mcp_server.server_id, - ) - raise token_response = response.json() - access_token = token_response["access_token"] # Validate token response against server-configured rules before any storage. # This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc. @@ -791,8 +875,23 @@ async def exchange_token_with_server( mcp_server.server_id, ) + # A DCR-bridge oauth_delegate server hands the client a gateway-bound envelope (identity plus the + # upstream token) instead of the raw upstream token, so the one bearer both admits the caller and + # forwards the upstream credential. Only this mode mints; every other server returns the raw token. + if bridge_mint_ready is not None: + if refresh_request_scope and isinstance(token_response, dict) and not token_response.get("scope"): + token_response = {**token_response, "scope": refresh_request_scope} + # Phase 3: seal the upstream grant into the client-held envelope; failures map through the same + # OAuth-shaped response as the phase-1 preconditions. + minted = _finish_bridge_mint(bridge_mint_ready, mcp_server, token_response, datetime.now(timezone.utc)) + return minted if isinstance(minted, JSONResponse) else _bridge_mint_error_response(minted) + + raw_access_token = token_response.get("access_token") if isinstance(token_response, dict) else None + if not isinstance(raw_access_token, str) or not raw_access_token: + return render_token_fault(UpstreamProtocolFault(note="the upstream token response has no usable access_token")) + result = { - "access_token": access_token, + "access_token": raw_access_token, "token_type": token_response.get("token_type", "Bearer"), } @@ -1048,21 +1147,6 @@ async def _persist_dcr_client_registration( return "failed" -_MAX_UPSTREAM_ERROR_CHARS = 500 - - -def _safe_upstream_error_detail(response: httpx.Response) -> str: - """Bounded plaintext summary of an upstream registration failure for the client. - - RFC 7591 error bodies are small JSON objects (``error`` / ``error_description``); relaying the - text lets the client read the real reason instead of a bare 500, and the length bound keeps a - hostile or oversized upstream body from bloating the gateway response.""" - body = response.text - if not body: - return response.reason_phrase or "upstream registration failed" - return body[:_MAX_UPSTREAM_ERROR_CHARS] - - async def register_client_with_server( request: Request, mcp_server: MCPServer, @@ -1122,19 +1206,24 @@ async def register_client_with_server( } async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Register) - response = await async_client.post( - mcp_server.registration_url, - headers=headers, - json=register_data, - ) + try: + response = await async_client.post( + mcp_server.registration_url, + headers=headers, + json=register_data, + ) + if response is not None: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + status_code, detail = dcr_fault_detail( + classify_upstream_dcr_rejection(exc.response, log_context=mcp_server.server_id) + ) + raise HTTPException(status_code=status_code, detail=detail) from exc if response is None: raise HTTPException( status_code=502, detail="MCP upstream registration endpoint returned no response", ) - if bridge_relay and response.status_code >= 400: - raise HTTPException(status_code=response.status_code, detail=_safe_upstream_error_detail(response)) - response.raise_for_status() token_response = response.json() @@ -1362,7 +1451,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/litellm/proxy/_experimental/mcp_server/faults/__init__.py b/litellm/proxy/_experimental/mcp_server/faults/__init__.py new file mode 100644 index 00000000000..da078f0e242 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/faults/__init__.py @@ -0,0 +1,38 @@ +"""Typed fault values for upstream OAuth/DCR failures (phase 1 of the MCP error-handling framework). + +The invariant this package exists to enforce: an upstream failure is classified ONCE into a single +fault value, and the response status, wire error code, and prose are all derived from that value. +Deriving all three from one classification makes contradictory pairings (a caller-fault error code on +a server-fault status) unrepresentable, and gives the trust-boundary rule one enforcement point: +spec-defined machine fields may cross to callers, upstream prose and raw bodies go to server logs. +""" + +from litellm.proxy._experimental.mcp_server.faults.classify import ( + classify_upstream_dcr_rejection, + classify_upstream_token_rejection, +) +from litellm.proxy._experimental.mcp_server.faults.render_oauth import ( + dcr_fault_detail, + render_token_fault, +) +from litellm.proxy._experimental.mcp_server.faults.types import ( + CallerRejected, + CredentialSource, + GatewayRejected, + UpstreamOAuthFault, + UpstreamProtocolFault, + UpstreamReportedFault, +) + +__all__ = [ + "CallerRejected", + "CredentialSource", + "GatewayRejected", + "UpstreamOAuthFault", + "UpstreamProtocolFault", + "UpstreamReportedFault", + "classify_upstream_dcr_rejection", + "classify_upstream_token_rejection", + "dcr_fault_detail", + "render_token_fault", +] diff --git a/litellm/proxy/_experimental/mcp_server/faults/classify.py b/litellm/proxy/_experimental/mcp_server/faults/classify.py new file mode 100644 index 00000000000..8b3a09f8d8d --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/faults/classify.py @@ -0,0 +1,133 @@ +"""The single place that reads upstream OAuth/DCR failure responses. + +Every accessor here is total: an upstream that lies about its content encoding, sends an undecodable +body, or omits the spec fields yields a classified fault, never an exception. Nothing outside this +module should touch a failed upstream response's body. +""" + +from __future__ import annotations + +import httpx + +from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.faults.types import ( + GATEWAY_CAPABILITY_CODES, + GATEWAY_CREDENTIAL_CODES, + MAX_WIRE_FIELD_CHARS, + CallerRejected, + CredentialSource, + GatewayRejected, + UpstreamOAuthFault, + UpstreamProtocolFault, + UpstreamReportedFault, +) + + +def _safe_text(response: httpx.Response) -> str: + try: + return response.text + except Exception: + return "" + + +def _safe_json(response: httpx.Response) -> object: + try: + return response.json() + except Exception: + return None + + +def _bounded_field(value: object) -> str | None: + if not isinstance(value, str) or not value: + return None + return value[:MAX_WIRE_FIELD_CHARS] + + +def _log_out_of_contract(endpoint_kind: str, response: httpx.Response, log_context: str) -> None: + verbose_logger.warning( + "MCP upstream %s endpoint (%s) returned HTTP %s outside the OAuth error contract (first %s chars): %s", + endpoint_kind, + log_context, + response.status_code, + MAX_WIRE_FIELD_CHARS, + _safe_text(response)[:MAX_WIRE_FIELD_CHARS], + ) + + +def _classify_oauth_error_code( + code: str, + description: str | None, + error_uri: str | None, + credential_source: CredentialSource, + log_context: str, +) -> UpstreamOAuthFault: + """Blame assignment for a contract-conformant OAuth error code, shared by the token and DCR + classifiers. Codes by which the upstream blames itself keep that blame; ``invalid_target`` is a + gateway capability gap (RFC 8707 resource indicators, LIT-4339) no matter whose credentials were + presented; credential-indicting codes follow the credential source; everything else, including + codes we do not recognize, is the caller's to act on. The upstream's HTTP status is deliberately + never consulted: status derives from this classification at render time, which is what keeps + status and code from contradicting each other.""" + if code == "server_error" or code == "temporarily_unavailable": + return UpstreamReportedFault(code=code) + if code in GATEWAY_CAPABILITY_CODES: + verbose_logger.warning( + "MCP server %s: the upstream authorization server rejected the request with " + "invalid_target; it may require RFC 8707 resource indicators, which the gateway " + "does not send yet (tracked as LIT-4339)", + log_context, + ) + return GatewayRejected(code=code) + if credential_source == "gateway_stored" and code in GATEWAY_CREDENTIAL_CODES: + verbose_logger.warning( + "MCP server %s: upstream authorization server rejected the gateway's configured client " + "credentials (%s): %s", + log_context, + code, + description or "", + ) + return GatewayRejected(code=code) + return CallerRejected(code=code, description=description, error_uri=error_uri) + + +def classify_upstream_token_rejection( + response: httpx.Response, + credential_source: CredentialSource, + log_context: str, +) -> UpstreamOAuthFault: + """Classify a token-endpoint rejection into exactly one fault: a body with an RFC 6749 §5.2 + ``error`` field goes through blame assignment (:func:`_classify_oauth_error_code`); anything + without a usable ``error`` field is an upstream protocol fault.""" + parsed = _safe_json(response) + fields = parsed if isinstance(parsed, dict) else {} + code = _bounded_field(fields.get("error")) + if code is None: + _log_out_of_contract("token", response, log_context) + return UpstreamProtocolFault(note=f"upstream token endpoint returned HTTP {response.status_code}") + return _classify_oauth_error_code( + code, + description=_bounded_field(fields.get("error_description")), + error_uri=_bounded_field(fields.get("error_uri")), + credential_source=credential_source, + log_context=log_context, + ) + + +def classify_upstream_dcr_rejection(response: httpx.Response, log_context: str) -> UpstreamOAuthFault: + """Classify a dynamic-client-registration rejection. RFC 7591 §3.2.2 errors carry + ``error`` / ``error_description`` and go through the same blame assignment as token errors + (registration sends no client credentials, so credential codes stay caller-actionable); anything + without a usable ``error`` field is an upstream protocol fault.""" + parsed = _safe_json(response) + fields = parsed if isinstance(parsed, dict) else {} + code = _bounded_field(fields.get("error")) + if code is None: + _log_out_of_contract("registration", response, log_context) + return UpstreamProtocolFault(note=f"upstream registration failed with HTTP {response.status_code}") + return _classify_oauth_error_code( + code, + description=_bounded_field(fields.get("error_description")), + error_uri=None, + credential_source="caller_supplied", + log_context=log_context, + ) diff --git a/litellm/proxy/_experimental/mcp_server/faults/render_oauth.py b/litellm/proxy/_experimental/mcp_server/faults/render_oauth.py new file mode 100644 index 00000000000..89ce5011830 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/faults/render_oauth.py @@ -0,0 +1,89 @@ +"""Render upstream OAuth/DCR faults onto the wire. The only place that chooses statuses and bodies +for these faults, so every consumer emits the same contract: RFC 6749 §5.2-shaped JSON with the §5.1 +no-store headers on token endpoints, HTTPException details on registration. Status, code, and prose +all derive from the fault tag; exhaustive matches keep a new fault arm from shipping unrendered. +""" + +from __future__ import annotations + +from fastapi.responses import JSONResponse +from typing_extensions import assert_never + +from litellm.proxy._experimental.mcp_server.faults.types import UpstreamOAuthFault +from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS + + +def _gateway_rejected_description(code: str) -> str: + if code == "invalid_target": + return ( + "the upstream authorization server rejected the request (invalid_target); " + "it may require RFC 8707 resource indicators, which the gateway does not send yet" + ) + return ( + f"the upstream authorization server rejected the gateway's configured client credentials " + f"({code}); verify the MCP server's client_id and client_secret" + ) + + +def _upstream_reported_status_and_description(code: str) -> tuple[int, str]: + if code == "temporarily_unavailable": + return 503, "the upstream authorization server is temporarily unavailable; retry shortly" + return 502, "the upstream authorization server reported an internal error" + + +def render_token_fault(fault: UpstreamOAuthFault) -> JSONResponse: + """RFC 6749 §5.2 response for a token-endpoint fault. Caller-actionable rejections relay the + upstream's code on the status that code implies (401 for invalid_client per §5.2, else 400); + gateway-side faults are 502 ``server_error`` with gateway-authored prose so a caller is never + blamed for, or shown the internals of, a failure only the operator can fix.""" + match fault.tag: + case "caller_rejected": + content = { + "error": fault.code, + **({"error_description": fault.description} if fault.description else {}), + **({"error_uri": fault.error_uri} if fault.error_uri else {}), + } + status_code = 401 if fault.code == "invalid_client" else 400 + return JSONResponse(status_code=status_code, content=content, headers=TOKEN_NO_CACHE_HEADERS) + case "gateway_rejected": + return JSONResponse( + status_code=502, + content={ + "error": "server_error", + "error_description": _gateway_rejected_description(fault.code), + }, + headers=TOKEN_NO_CACHE_HEADERS, + ) + case "upstream_reported_fault": + status_code, description = _upstream_reported_status_and_description(fault.code) + return JSONResponse( + status_code=status_code, + content={"error": fault.code, "error_description": description}, + headers=TOKEN_NO_CACHE_HEADERS, + ) + case "upstream_protocol_fault": + return JSONResponse( + status_code=502, + content={"error": "server_error", "error_description": fault.note}, + headers=TOKEN_NO_CACHE_HEADERS, + ) + case _: + assert_never(fault.tag) + + +def dcr_fault_detail(fault: UpstreamOAuthFault) -> tuple[int, str]: + """Status and detail string for a registration fault, raised as HTTPException by the caller. + RFC 7591 §3.2.2 defines registration errors as 400, so a contract-conformant rejection is 400 + regardless of the status the upstream chose; everything else is a 502 upstream fault.""" + match fault.tag: + case "caller_rejected": + detail = f"{fault.code}: {fault.description}" if fault.description else fault.code + return 400, detail + case "gateway_rejected": + return 502, _gateway_rejected_description(fault.code) + case "upstream_reported_fault": + return _upstream_reported_status_and_description(fault.code) + case "upstream_protocol_fault": + return 502, fault.note + case _: + assert_never(fault.tag) diff --git a/litellm/proxy/_experimental/mcp_server/faults/types.py b/litellm/proxy/_experimental/mcp_server/faults/types.py new file mode 100644 index 00000000000..128b5e3e6cf --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/faults/types.py @@ -0,0 +1,79 @@ +"""Fault taxonomy for upstream OAuth token and DCR registration failures. + +Each fault is a frozen model on a ``tag`` literal. The tag alone decides the HTTP status, the wire +error code, and whose prose the caller sees, so those three facts can never disagree the way they can +when an upstream's status and error code are relayed independently. +""" + +from __future__ import annotations + +from typing import Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict + +MAX_WIRE_FIELD_CHARS = 500 +"""Bound on every upstream-derived string that crosses to a caller or into a log line.""" + +CredentialSource: TypeAlias = Literal["gateway_stored", "caller_supplied"] +"""Whose client credentials the gateway presented upstream: the MCP server's stored configuration or +credentials the caller supplied on the request. Decides whether a credential rejection is the +caller's problem to fix or the gateway operator's.""" + +GATEWAY_CREDENTIAL_CODES: frozenset[str] = frozenset({"invalid_client", "unauthorized_client"}) +"""RFC 6749 error codes that indict the OAuth client's credentials or grant authorization. When the +gateway presented its own stored credentials, these are gateway-side faults the caller cannot act on; +when the caller supplied the credentials, they are the caller's to fix.""" + +GATEWAY_CAPABILITY_CODES: frozenset[str] = frozenset({"invalid_target"}) +"""Codes that indict a gateway capability regardless of whose credentials were presented: +``invalid_target`` means the upstream wants RFC 8707 resource indicators, which the gateway does not +send yet (LIT-4339). Never the caller's fault.""" + +UPSTREAM_FAULT_CODES: frozenset[str] = frozenset({"server_error", "temporarily_unavailable"}) +"""Codes by which the upstream blames itself. Relaying them as caller faults would invert blame, so +they classify as upstream-reported faults and render on the 5xx their meaning implies.""" + + +class CallerRejected(BaseModel): + """The upstream spoke the OAuth error contract and the failure is actionable by our caller + (e.g. ``invalid_grant``: re-run authorization). The code and its bounded prose relay on the + 4xx status the code itself implies.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["caller_rejected"] = "caller_rejected" + code: str + description: str | None = None + error_uri: str | None = None + + +class GatewayRejected(BaseModel): + """The upstream rejected the request for a cause only the gateway operator can address: the + server's stored client credentials or a gateway capability gap. Not actionable by the caller: + rendered as 502 with gateway-authored prose naming the code; the upstream's prose goes to + server logs only.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["gateway_rejected"] = "gateway_rejected" + code: str + + +class UpstreamReportedFault(BaseModel): + """The upstream blamed itself in the OAuth vocabulary. Rendered on the 5xx the code implies + (``server_error`` 502, ``temporarily_unavailable`` 503) so blame and status agree.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["upstream_reported_fault"] = "upstream_reported_fault" + code: Literal["server_error", "temporarily_unavailable"] + + +class UpstreamProtocolFault(BaseModel): + """The upstream broke the error contract: no JSON ``error`` field, an undecodable body, or a + success response without a usable token. Rendered as 502 with a gateway-authored note; the + upstream body never crosses to the caller.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["upstream_protocol_fault"] = "upstream_protocol_fault" + note: str + + +UpstreamOAuthFault: TypeAlias = CallerRejected | GatewayRejected | UpstreamReportedFault | UpstreamProtocolFault diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py index 5530fbc46fd..c352f3a683e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py @@ -21,11 +21,16 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import EnvelopeKeys, EnvelopeMintError, OpenedEnvelope, + OpenedRefreshEnvelope, + RefreshCredential, SealedEnvelope, UpstreamTokenGrant, is_envelope, + is_refresh_envelope, mint_envelope, + mint_refresh_envelope, open_envelope, + open_refresh_envelope, ) _SIGNING_KEY_DOMAIN = b"litellm-mcp-bridge:envelope-signing:" @@ -92,6 +97,67 @@ def build_bridge_token_response( return mint_envelope(identity, grant, keys, now) +def build_bridge_refresh_token_response( + identity: EnvelopeIdentity, + refresh: RefreshCredential, + keys: EnvelopeKeys, + now: datetime, +) -> SealedEnvelope | EnvelopeMintError: + """Seal ``refresh`` for ``identity`` into the long-lived refresh envelope the token endpoint returns + alongside the access envelope, so the client can renew without re-authenticating. A thin, pure + wrapper over :func:`mint_refresh_envelope`; returns the mint error as a value for the caller to map. + """ + return mint_refresh_envelope(identity, refresh, keys, now) + + +class BridgeRefreshOpened(BaseModel): + """A valid refresh envelope presented to the token endpoint: the identity to re-validate and renew + under, and the upstream refresh grant to exchange.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["opened"] = "opened" + identity: EnvelopeIdentity + refresh: RefreshCredential + + +class BridgeRefreshInvalid(BaseModel): + """The presented refresh grant is not a valid refresh envelope for this server (not refresh-shaped, + will not open, or minted for a different server); the token endpoint fails the refresh closed.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["invalid"] = "invalid" + + +BridgeRefreshResult: TypeAlias = BridgeRefreshOpened | BridgeRefreshInvalid + + +def open_bridge_refresh_envelope( + refresh_value: str, + keys: EnvelopeKeys, + now: datetime, + expected_server_id: str, +) -> BridgeRefreshResult: + """Open a refresh envelope a bridge ``oauth_delegate`` client presented on a refresh_token grant. + + The token-endpoint mirror of :func:`resolve_bridge_envelope`: strips an optional ``Bearer`` scheme, + then returns ``BridgeRefreshOpened`` with the recovered identity and upstream refresh grant, or + ``BridgeRefreshInvalid`` for anything that is not a valid refresh envelope for this server. Never + raises; total over hostile input via :func:`open_refresh_envelope`. ``expected_server_id`` binds the + envelope to the server the request targets, so a refresh envelope minted for one server cannot renew + against another. A raw upstream refresh token (not envelope-shaped) is ``BridgeRefreshInvalid``: this + mode never hands the client a bare upstream refresh token, so it must never accept one. + """ + candidate = _strip_bearer(refresh_value) + if not is_refresh_envelope(candidate): + return BridgeRefreshInvalid() + opened = open_refresh_envelope(candidate, keys, now) + if not isinstance(opened, OpenedRefreshEnvelope): + return BridgeRefreshInvalid() + if opened.identity.server_id != expected_server_id: + return BridgeRefreshInvalid() + return BridgeRefreshOpened(identity=opened.identity, refresh=opened.refresh) + + class NotBridgeEnvelope(BaseModel): """The bearer is not an envelope; admission continues on its normal path.""" @@ -128,10 +194,12 @@ def _strip_bearer(value: str) -> str: def is_bridge_envelope_shaped(authorization_value: str) -> bool: - """Cheap, keyless test that an ``Authorization`` value carries an envelope (optional - ``Bearer`` scheme stripped). The admission edge engages the bridge arm only for an - envelope, so a plain upstream bearer falls through to normal oauth2 admission.""" - return is_envelope(_strip_bearer(authorization_value)) + """Cheap, keyless test that an ``Authorization`` value carries an envelope of either kind (optional + ``Bearer`` scheme stripped). The admission edge engages the bridge arm for an access envelope (to + admit) and for a refresh envelope (to reject it explicitly, since a refresh credential is never + usable at the tool-call edge); a plain upstream bearer falls through to normal oauth2 admission.""" + candidate = _strip_bearer(authorization_value) + return is_envelope(candidate) or is_refresh_envelope(candidate) def resolve_bridge_envelope( @@ -148,6 +216,10 @@ def resolve_bridge_envelope( envelope, and ``BridgeEnvelopeInvalid`` for an envelope-shaped bearer that will not open. Never raises: it is total over hostile input via :func:`open_envelope`. + A refresh envelope is ``BridgeEnvelopeInvalid`` here: it is a valid gateway credential but only ever + presented back to the token endpoint, never usable to authenticate a tool call, so admission must + fail it closed rather than let it fall through to another arm. + ``expected_server_id`` is the ``server_id`` of the MCP server the request targets; an opened envelope whose sealed ``server_id`` does not match is rejected as ``BridgeEnvelopeInvalid``. Binding here (rather than leaving it to the caller) prevents @@ -157,6 +229,8 @@ def resolve_bridge_envelope( unlike ``hmac.compare_digest`` on ``str``, does not raise on a non-ASCII server_id. """ candidate = _strip_bearer(authorization_value) + if is_refresh_envelope(candidate): + return BridgeEnvelopeInvalid() if not is_envelope(candidate): return NotBridgeEnvelope() opened = open_envelope(candidate, keys, now) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py index 517c2ef5c8f..9118a3e129d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -44,18 +44,33 @@ from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value ENVELOPE_PREFIX = "llm_env_" -"""Marker prefix on every serialized envelope so the edge can cheaply tell an envelope +"""Marker prefix on every serialized ACCESS envelope so the edge can cheaply tell an envelope from a raw upstream token before doing any cryptography.""" +REFRESH_ENVELOPE_PREFIX = "llm_refresh_" +"""Marker prefix on every serialized REFRESH envelope. A distinct prefix keeps the two credentials +routable without crypto and, together with the signed ``kind`` claim, stops one from being presented +where the other is expected: a refresh envelope carries a long-lived upstream refresh token and is only +ever presented back to the token endpoint, never forwarded upstream on a tool call.""" + ENVELOPE_ISSUER = "litellm-mcp-bridge" """``iss`` claim stamped into every envelope and required back on open.""" MAX_ENVELOPE_TTL_SECONDS = 3600 -"""Hard ceiling on envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)`` +"""Hard ceiling on ACCESS envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)`` (the cap alone when the upstream omits ``expires_in``), matching the 1h lifetime of the BYOK session bearer this module's signing approach is borrowed from: a client-held credential should never outlive a bounded window even when the upstream token does.""" +MAX_REFRESH_ENVELOPE_TTL_SECONDS = 1209600 +"""Hard ceiling on REFRESH envelope lifetime (14 days). A refresh envelope only renews the short-lived +access envelope, and each renewal re-validates the sealed litellm key (revocation gates it) and is +re-minted with a fresh window, so the practical bound is idle time, not a fixed session. ``exp`` is +``min(upstream refresh_expires_in, this cap)`` (the cap alone when the upstream omits it); if the +upstream refresh token dies first, the next renewal simply fails at the upstream and the client +re-authenticates. The value is deliberately far shorter than a typical upstream refresh-token lifetime +so a leaked refresh envelope is bounded even if the upstream would have honoured it for longer.""" + MAX_ENVELOPE_BYTES = 12288 """Size cap on the final serialized envelope (prefix + JWT, in bytes). Upstream JWTs commonly run 2-4KB; base64 plus encryption overhead roughly doubles that inside the @@ -66,21 +81,48 @@ typed error, never truncated.""" _ENVELOPE_JWT_ALGORITHM = "HS256" +EnvelopeKind = Literal["access", "refresh"] +"""Which credential an envelope is. Stamped into the signed claims and required to match on open, so a +signature-valid envelope of one kind cannot be replayed as the other even if its wire prefix is swapped +(the prefix is not part of the signed payload; this claim is).""" + + +EnvelopeSubjectType: TypeAlias = Literal["key_hash", "user_id"] +"""Discriminator for what litellm principal the envelope binds the grant to. + +``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 identity the envelope binds the inner grant to. + """The litellm principal the envelope binds the inner 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. + ``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): @@ -99,6 +141,21 @@ class UpstreamTokenGrant(BaseModel): expires_in: int | None = Field(default=None, gt=0) +class RefreshCredential(BaseModel): + """The upstream refresh grant sealed inside a refresh envelope. + + Only the refresh token (plus the scope to re-request and the refresh token's own lifetime, when the + upstream reports it) is sealed; the access token is never in a refresh envelope. ``refresh_token`` is + a ``SecretStr`` so reprs never leak it, and ``expires_in`` (the refresh token's lifetime, not the + access token's) must be positive when present. + """ + + model_config = ConfigDict(frozen=True) + refresh_token: SecretStr = Field(min_length=1) + scope: str | None = None + expires_in: int | None = Field(default=None, gt=0) + + class EnvelopeKeys(BaseModel): """Injected key material: the HS256 signing key and the symmetric encryption key. @@ -121,13 +178,21 @@ class SealedEnvelope(BaseModel): class OpenedEnvelope(BaseModel): - """A validated envelope: the identity it was minted for and the recovered grant.""" + """A validated access envelope: the identity it was minted for and the recovered grant.""" model_config = ConfigDict(frozen=True) identity: EnvelopeIdentity grant: UpstreamTokenGrant +class OpenedRefreshEnvelope(BaseModel): + """A validated refresh envelope: the identity it was minted for and the recovered refresh grant.""" + + model_config = ConfigDict(frozen=True) + identity: EnvelopeIdentity + refresh: RefreshCredential + + class EnvelopeTooLarge(BaseModel): """The serialized envelope exceeded ``MAX_ENVELOPE_BYTES``; carries sizes only.""" @@ -199,8 +264,10 @@ class _EnvelopeClaims(BaseModel): iss: str iat: int exp: int + kind: EnvelopeKind 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) @@ -213,11 +280,25 @@ class _GrantWire(BaseModel): expires_in: int | None = None +class _RefreshWire(BaseModel): + model_config = ConfigDict(frozen=True) + refresh_token: str + scope: str | None = None + expires_in: int | None = None + + def is_envelope(candidate: str) -> bool: - """Cheap prefix check so the edge can route envelopes vs raw tokens without crypto.""" + """Cheap prefix check for an ACCESS envelope so the edge can route envelopes vs raw tokens without + crypto. A refresh envelope has a different prefix and is not an access envelope.""" return candidate.startswith(ENVELOPE_PREFIX) +def is_refresh_envelope(candidate: str) -> bool: + """Cheap prefix check for a REFRESH envelope so the token endpoint can route a refresh grant that + carries an envelope vs a raw upstream refresh token without crypto.""" + return candidate.startswith(REFRESH_ENVELOPE_PREFIX) + + def mint_envelope( identity: EnvelopeIdentity, grant: UpstreamTokenGrant, @@ -231,23 +312,15 @@ def mint_envelope( serialized envelope exceeds ``MAX_ENVELOPE_BYTES``. """ expires_at = now + timedelta(seconds=_envelope_ttl_seconds(grant.expires_in)) - claims = _EnvelopeClaims( - iss=ENVELOPE_ISSUER, - iat=int(now.timestamp()), - exp=int(expires_at.timestamp()), - server_id=identity.server_id, - key_hash=identity.key_hash, - grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key), + return _seal( + kind="access", + prefix=ENVELOPE_PREFIX, + identity=identity, + grant_blob=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key), + expires_at=expires_at, + signing_key=keys.signing_key, + now=now, ) - token = ENVELOPE_PREFIX + jwt.encode( - claims.model_dump(), - keys.signing_key.get_secret_value(), - algorithm=_ENVELOPE_JWT_ALGORITHM, - ) - size_bytes = len(token.encode("utf-8")) - if size_bytes > MAX_ENVELOPE_BYTES: - return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES) - return SealedEnvelope(token=SecretStr(token), expires_at=expires_at) def open_envelope( @@ -263,35 +336,136 @@ def open_envelope( re-derived, so it is stale by up to the envelope's lifetime; callers that need a live remaining lifetime should use ``now`` against the upstream, not this field. """ - if not is_envelope(candidate): - return NotAnEnvelope() - # UTF-8 byte length is never below character length, so a character count already over the - # cap rejects an oversize candidate in O(1) without encoding it; the exact byte check then - # runs only on candidates already bounded to <= MAX_ENVELOPE_BYTES characters. - if len(candidate) > MAX_ENVELOPE_BYTES: - return MalformedPayload() - if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES: - return MalformedPayload() - claims = _decode_claims(candidate.removeprefix(ENVELOPE_PREFIX), keys.signing_key) + claims = _open_claims(candidate, prefix=ENVELOPE_PREFIX, expected_kind="access", keys=keys, now=now) if not isinstance(claims, _EnvelopeClaims): return claims - if now.timestamp() >= claims.exp: - return Expired() grant = _decrypt_grant(claims.grant, keys.encryption_key) 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, ) +def mint_refresh_envelope( + identity: EnvelopeIdentity, + refresh: RefreshCredential, + keys: EnvelopeKeys, + now: datetime, +) -> SealedEnvelope | EnvelopeMintError: + """Seal ``refresh`` for ``identity`` into a long-lived, client-held refresh envelope. + + ``exp`` is ``min(refresh.expires_in, MAX_REFRESH_ENVELOPE_TTL_SECONDS)`` seconds from ``now`` (the + cap alone when the upstream omits the refresh lifetime). Sealing a distinct ``kind="refresh"`` claim + is what keeps a refresh envelope from ever opening as an access credential at the MCP edge. Returns + ``EnvelopeTooLarge`` when the serialized envelope exceeds ``MAX_ENVELOPE_BYTES``. + """ + expires_at = now + timedelta(seconds=_refresh_ttl_seconds(refresh.expires_in)) + return _seal( + kind="refresh", + prefix=REFRESH_ENVELOPE_PREFIX, + identity=identity, + grant_blob=_encrypt_grant_blob(_refresh_plaintext(refresh), keys.encryption_key), + expires_at=expires_at, + signing_key=keys.signing_key, + now=now, + ) + + +def open_refresh_envelope( + candidate: str, + keys: EnvelopeKeys, + now: datetime, +) -> OpenedRefreshEnvelope | EnvelopeOpenError: + """Validate a refresh ``candidate`` and recover the identity and inner refresh grant. + + Total over hostile input exactly like :func:`open_envelope`: every invalid, expired, tampered, + wrong-kind, or undecryptable candidate maps to a distinct ``EnvelopeOpenError`` variant, never a + raise. The ``kind="refresh"`` claim is required, so an access envelope re-prefixed as a refresh one + is rejected as ``MalformedPayload``. + """ + claims = _open_claims(candidate, prefix=REFRESH_ENVELOPE_PREFIX, expected_kind="refresh", keys=keys, now=now) + if not isinstance(claims, _EnvelopeClaims): + return claims + refresh = _decrypt_refresh(claims.grant, keys.encryption_key) + if not isinstance(refresh, RefreshCredential): + return refresh + return OpenedRefreshEnvelope( + identity=EnvelopeIdentity(server_id=claims.server_id, subject_type=claims.subject_type, subject=claims.subject), + refresh=refresh, + ) + + +def _seal( + kind: EnvelopeKind, + prefix: str, + identity: EnvelopeIdentity, + grant_blob: str, + expires_at: datetime, + signing_key: SecretStr, + now: datetime, +) -> SealedEnvelope | EnvelopeTooLarge: + """Sign the claims for either envelope kind and enforce the size cap. Shared by both mints so the + JWT shape, issuer, and size guard cannot drift between access and refresh envelopes.""" + claims = _EnvelopeClaims( + iss=ENVELOPE_ISSUER, + iat=int(now.timestamp()), + exp=int(expires_at.timestamp()), + kind=kind, + server_id=identity.server_id, + subject_type=identity.subject_type, + subject=identity.subject, + grant=grant_blob, + ) + token = prefix + jwt.encode(claims.model_dump(), signing_key.get_secret_value(), algorithm=_ENVELOPE_JWT_ALGORITHM) + size_bytes = len(token.encode("utf-8")) + if size_bytes > MAX_ENVELOPE_BYTES: + return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES) + return SealedEnvelope(token=SecretStr(token), expires_at=expires_at) + + +def _open_claims( + candidate: str, + prefix: str, + expected_kind: EnvelopeKind, + keys: EnvelopeKeys, + now: datetime, +) -> _EnvelopeClaims | EnvelopeOpenError: + """Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an attacker-controlled + candidate, shared by both openers so the security gate is identical for access and refresh. Returns + the validated claims or a distinct ``EnvelopeOpenError``; never raises.""" + if not candidate.startswith(prefix): + return NotAnEnvelope() + # UTF-8 byte length is never below character length, so a character count already over the cap + # rejects an oversize candidate in O(1) without encoding it; the exact byte check then runs only on + # candidates already bounded to <= MAX_ENVELOPE_BYTES characters. + if len(candidate) > MAX_ENVELOPE_BYTES: + return MalformedPayload() + if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES: + return MalformedPayload() + claims = _decode_claims(candidate.removeprefix(prefix), keys.signing_key) + if not isinstance(claims, _EnvelopeClaims): + return claims + if claims.kind != expected_kind: + return MalformedPayload() + if now.timestamp() >= claims.exp: + return Expired() + return claims + + def _envelope_ttl_seconds(upstream_expires_in: int | None) -> int: if upstream_expires_in is None: return MAX_ENVELOPE_TTL_SECONDS return min(upstream_expires_in, MAX_ENVELOPE_TTL_SECONDS) +def _refresh_ttl_seconds(upstream_refresh_expires_in: int | None) -> int: + if upstream_refresh_expires_in is None: + return MAX_REFRESH_ENVELOPE_TTL_SECONDS + return min(upstream_refresh_expires_in, MAX_REFRESH_ENVELOPE_TTL_SECONDS) + + def _grant_plaintext(grant: UpstreamTokenGrant) -> str: wire = _GrantWire( access_token=grant.access_token.get_secret_value(), @@ -303,6 +477,15 @@ def _grant_plaintext(grant: UpstreamTokenGrant) -> str: return wire.model_dump_json(exclude_none=True) +def _refresh_plaintext(refresh: RefreshCredential) -> str: + wire = _RefreshWire( + refresh_token=refresh.refresh_token.get_secret_value(), + scope=refresh.scope, + expires_in=refresh.expires_in, + ) + return wire.model_dump_json(exclude_none=True) + + def _decode_claims( compact: str, signing_key: SecretStr, @@ -364,3 +547,22 @@ def _decrypt_grant( return UpstreamTokenGrant.model_validate_json(plaintext) except ValidationError: return MalformedPayload() + + +def _decrypt_refresh( + blob: str, + encryption_key: SecretStr, +) -> RefreshCredential | DecryptFailed | MalformedPayload: + from nacl.exceptions import CryptoError + + try: + plaintext = decrypt_value( + value=base64.urlsafe_b64decode(blob), + signing_key=encryption_key.get_secret_value(), + ) + except (CryptoError, ValueError): + return DecryptFailed() + try: + return RefreshCredential.model_validate_json(plaintext) + except ValidationError: + return MalformedPayload() diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5090aa7d7d5..68a61b85175 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3719,8 +3719,15 @@ if MCP_AVAILABLE: headers={"www-authenticate": upstream_www_authenticate}, ) + def _get_authorization_header_from_scope(scope: Scope) -> Optional[str]: + """First ``Authorization`` header value in the ASGI scope, or None.""" + for key, value in scope.get("headers", []): + if key.lower() == b"authorization": + return value.decode("latin-1") + return None + def _scope_has_authorization_header(scope: Scope) -> bool: - return any(key.lower() == b"authorization" for key, _ in scope.get("headers", [])) + return _get_authorization_header_from_scope(scope) is not None def _get_forwarded_auth_from_scope(scope: Scope) -> Optional[str]: """Return the upstream-bound ``Authorization`` header value, or None. @@ -3733,17 +3740,24 @@ if MCP_AVAILABLE: ``MCPRequestHandler.process_mcp_request``), and forwarding it upstream would leak the proxy key to a third-party MCP server. """ - authorization = None - has_litellm_key_header = False - for key, value in scope.get("headers", []): - key_lower = key.lower() - if key_lower == b"authorization": - authorization = value.decode("latin-1") - elif key_lower == b"x-litellm-api-key": - has_litellm_key_header = True + has_litellm_key_header = any(key.lower() == b"x-litellm-api-key" for key, _ in scope.get("headers", [])) if not has_litellm_key_header: return None - return authorization + return _get_authorization_header_from_scope(scope) + + def _is_delegate_upstream_probe_target(server: MCPServer) -> bool: + """Whether ``server`` is an interactive delegate-auth server whose client-supplied + token should be preflighted upstream. + + Mirrors the anonymous-delegate gate in ``get_allowed_mcp_servers``: the flow is + resolved via ``effective_oauth2_flow`` so an unstamped M2M-shape row fails closed + (its stored client credentials drive egress; the caller's bearer is irrelevant). + """ + return ( + server.auth_type == MCPAuth.oauth2 + and server.delegate_auth_to_upstream is True + and MCPServerManager.effective_oauth2_flow(server) != "client_credentials" + ) async def _probe_upstream_auth( url: str, @@ -3805,7 +3819,7 @@ if MCP_AVAILABLE: mcp_servers: Optional[List[str]], client_ip: Optional[str], ) -> None: - """Probe pass-through upstream servers in parallel before the MCP session starts. + """Probe pass-through and delegate-auth upstream servers in parallel before the MCP session starts. Only servers the caller's key is already authorized to reach are probed — the list is derived from _get_allowed_mcp_servers so that a user cannot @@ -3813,11 +3827,42 @@ if MCP_AVAILABLE: The MCP SDK commits HTTP 200 headers before invoking handlers, so a 401 can only be returned before that point. This function raises HTTPException(401) - with a WWW-Authenticate header if any upstream rejects the client token. + with a WWW-Authenticate header if any upstream rejects the client token, or 403 + if the upstream accepts it but forbids the caller. Fails-open: network errors are logged and the request is allowed through. + + Delegate-auth servers (``auth_type=oauth2`` + ``delegate_auth_to_upstream``) + are probed with the caller's bare ``Authorization`` bearer. That bearer is only + an upstream token (never a LiteLLM key) when admission took the delegate bypass, + so the delegate target is resolved through ``get_mcp_server_by_name`` -- the same + resolver admission used -- rather than the wider allowed-server prefix/access-group + matching. A name that only reaches a delegate server via server_id or an access + group would have been admitted as a real LiteLLM key, so probing it would leak that + key upstream; requiring the admission-resolver match closes that gap. Without the + probe a rejected token is absorbed by the tools/list handler and masked as an empty + tool list. Gated to single-server routes so one rejected token cannot 401 a + multi-server aggregate connect, matching the OBO preflight gating; the challenge + echoes the requested name so aliased routes get the same resource_metadata URL as + the tokenless preemptive challenge. """ forwarded_auth = _get_forwarded_auth_from_scope(scope) - if not forwarded_auth: + requested_single_target = mcp_servers[0] if mcp_servers is not None and len(mcp_servers) == 1 else None + # The bare Authorization header (no x-litellm-api-key) is a valid upstream token + # only when admission classified it as one, i.e. the single requested name resolves + # to a delegate server under admission's own resolver. Resolve it the same way here + # so a server_id- or access-group-named delegate (which admission would have treated + # as a LiteLLM key) is never probed with that key. + delegate_server = ( + global_mcp_server_manager.get_mcp_server_by_name(requested_single_target, client_ip=client_ip) + if requested_single_target + else None + ) + delegate_auth = ( + _get_authorization_header_from_scope(scope) + if delegate_server is not None and _is_delegate_upstream_probe_target(delegate_server) + else None + ) + if not forwarded_auth and not delegate_auth: return # Use the authorized server set, not the raw user-supplied names, so that @@ -3827,33 +3872,49 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, client_ip=client_ip, ) - passthrough_servers = [ - srv - for srv in allowed_servers - # Restrict to genuine OAuth pass-through servers (auth_type none + - # Authorization in extra_headers). Gateway-managed OAuth2 servers - # must not receive the ``resource_metadata=`` challenge emitted - # below — they require ``authorization_uri=`` pointing at the - # gateway AS metadata. ``is_oauth_passthrough`` already requires - # ``auth_type in (None, MCPAuth.none)``, which is mutually - # exclusive with ``has_client_credentials`` (oauth2 + M2M flow), - # so M2M servers are implicitly excluded here. - if srv.is_oauth_passthrough - ] - if not passthrough_servers: + passthrough_targets: Tuple[Tuple[MCPServer, str, str], ...] = ( + tuple( + (srv, forwarded_auth, srv.name) + for srv in allowed_servers + # Restrict to genuine OAuth pass-through servers (auth_type none + + # Authorization in extra_headers). Gateway-managed OAuth2 servers + # must not receive the ``resource_metadata=`` challenge emitted + # below — they require ``authorization_uri=`` pointing at the + # gateway AS metadata. ``is_oauth_passthrough`` already requires + # ``auth_type in (None, MCPAuth.none)``, which is mutually + # exclusive with ``has_client_credentials`` (oauth2 + M2M flow), + # so M2M servers are implicitly excluded here. + if srv.is_oauth_passthrough + ) + if forwarded_auth + else () + ) + # Probe the admission-resolved delegate server only when the caller is actually + # authorized for it (present in the IP-filtered allowed set), keyed by server_id. + delegate_targets: Tuple[Tuple[MCPServer, str, str], ...] = ( + tuple( + (srv, delegate_auth, requested_single_target) + for srv in allowed_servers + if delegate_server is not None and srv.server_id == delegate_server.server_id + ) + if delegate_auth and requested_single_target + else () + ) + probe_targets = passthrough_targets + delegate_targets + if not probe_targets: return probe_results = await asyncio.gather( - *[_probe_upstream_auth(srv.url or "", forwarded_auth) for srv in passthrough_servers] + *[_probe_upstream_auth(srv.url or "", auth_header) for srv, auth_header, _ in probe_targets] ) - for srv, (probe_status, _) in zip(passthrough_servers, probe_results): + for (srv, _, challenge_server_name), (probe_status, _) in zip(probe_targets, probe_results): if probe_status == 401: # Token is missing or expired: keep pass-through clients on the # protected-resource discovery flow so they re-authorize against # the upstream IdP metadata proxied by LiteLLM. www_authenticate = _get_passthrough_www_authenticate( scope=scope, - server_name=srv.name, + server_name=challenge_server_name, invalid_token=True, ) raise HTTPException( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 23fe7730c17..5e3ea4b7dcb 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1118,6 +1118,7 @@ class GenerateKeyRequest(KeyRequestBase): class GenerateKeyResponse(KeyRequestBase): key: str # type: ignore key_name: Optional[str] = None + key_type: str | None = None expires: Optional[datetime] = None user_id: Optional[str] = None token_id: Optional[str] = None @@ -2421,6 +2422,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "is active as a reminder that hard enforcement is relaxed." ), ) + skip_user_budget_on_team_key: bool | None = Field( + None, + description=( + "If True, restores the legacy behavior where a user's personal " + "max_budget is NOT enforced when their key belongs to a team; only " + "the team (and team-member) budgets apply. Defaults to False, meaning " + "the user's personal max_budget is always enforced regardless of " + "whether the key belongs to a team (see GitHub issue #12905)." + ), + ) user_url_validation: Optional[bool] = Field( None, description=( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 230b9b70ff0..00f6d44e25a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -626,26 +626,29 @@ async def common_checks( ) async def _user_max_budget_check() -> None: - # 4.1 personal budget, if personal key - if ( - (team_object is None or team_object.team_id is None) - and user_object is not None - and user_object.max_budget is not None - ): - from litellm.proxy.proxy_server import get_current_spend + if user_object is None or user_object.max_budget is None: + return + skip_for_team = ( + general_settings.get("skip_user_budget_on_team_key") is True + and team_object is not None + and team_object.team_id is not None + ) + if skip_for_team: + return + from litellm.proxy.proxy_server import get_current_spend - user_budget = user_object.max_budget - user_spend = await get_current_spend( - counter_key=f"spend:user:{user_object.user_id}", - fallback_spend=user_object.spend or 0.0, + user_budget = user_object.max_budget + user_spend = await get_current_spend( + counter_key=f"spend:user:{user_object.user_id}", + fallback_spend=user_object.spend or 0.0, + max_budget=user_budget, + ) + if math.isfinite(user_budget) and user_spend >= user_budget: + raise litellm.BudgetExceededError( + current_cost=user_spend, max_budget=user_budget, + message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", ) - if math.isfinite(user_budget) and user_spend >= user_budget: - raise litellm.BudgetExceededError( - current_cost=user_spend, - max_budget=user_budget, - message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", - ) # Each scope reads a distinct counter key with no cross-scope ordering # dependency, so the per-scope Redis-first reads run concurrently instead @@ -4383,14 +4386,23 @@ def _model_custom_llm_provider_matches_wildcard_pattern(model: str, allowed_mode or - `model=claude-3-5-sonnet-20240620` - `allowed_model_pattern=anthropic/*` + + A model that already carries a namespace get_llm_provider did not consume + (e.g. `bedrockz/anthropic.claude-...`) is never granted here: its provider was + inferred from a fragment of the full string, so rebuilding + `{provider}/{model}` would produce `bedrock/bedrockz/...` and slip an + unrecognized namespace through a `bedrock/*` key. """ try: - model, custom_llm_provider, _, _ = get_llm_provider(model=model) + stripped_model, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: return False + if stripped_model == model and "/" in model: + return False + return is_model_allowed_by_pattern( - model=f"{custom_llm_provider}/{model}", + model=f"{custom_llm_provider}/{stripped_model}", allowed_model_pattern=allowed_model_pattern, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2613510bd0c..b402212fb2e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1191,13 +1191,15 @@ async def _user_api_key_auth_builder( return await handle_oauth2_proxy_request(request=request) if general_settings.get("enable_jwt_auth", False) is True: - from litellm.proxy.proxy_server import premium_user - - if premium_user is not True: - raise ValueError(f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}") is_jwt = jwt_handler.is_jwt(token=api_key) verbose_proxy_logger.debug("is_jwt: %s", is_jwt) if is_jwt: + from litellm.proxy.proxy_server import premium_user + + if premium_user is not True: + raise ValueError( + f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}" + ) # Try JWT-to-Virtual-Key mapping first to avoid # unnecessary DB queries in auth_builder do_standard_jwt_auth = True @@ -2442,6 +2444,7 @@ async def _reserve_budget_after_common_checks( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, + skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True, ) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 3a93896a206..e4c565e4464 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -8,6 +8,10 @@ from litellm.proxy._types import ( ) from litellm.secret_managers.main import str_to_bool +# Bounds the __cause__/__context__ walk in is_database_service_unavailable_error_in_chain. +# Real exception chains are a few links deep; the cap also makes the walk cycle-safe. +_MAX_EXCEPTION_CHAIN_DEPTH = 20 + class PrismaDBExceptionHandler: """ @@ -218,6 +222,32 @@ class PrismaDBExceptionHandler: ), ) + @staticmethod + def is_database_service_unavailable_error_in_chain(e: BaseException) -> bool: + """Like ``is_database_service_unavailable_error`` but also walks the + ``__cause__`` / ``__context__`` chain. + + ``is_database_service_unavailable_error`` classifies a single exception + by type, which a caller that catches a raw DB failure and re-raises a + domain exception of a different type defeats. ``get_user_object`` in + ``litellm/proxy/auth/auth_checks.py`` is the concrete case: it wraps + every DB error, a genuine outage included, in a bare ``ValueError`` + whose original error survives only as ``__context__``. A type check on + the ``ValueError`` misses the outage, so the caller would mistake an + infrastructure fault for an auth failure. Walking the chain recovers the + real signal, which is the PEP 3134 way to inspect a wrapped cause. + + The walk is depth-bounded, which also makes it cycle-safe. + """ + current: BaseException | None = e + for _ in range(_MAX_EXCEPTION_CHAIN_DEPTH): + if not isinstance(current, Exception): + return False + if PrismaDBExceptionHandler.is_database_service_unavailable_error(current): + return True + current = current.__cause__ or current.__context__ + return False + @staticmethod def handle_db_exception(e: Exception): """ diff --git a/litellm/proxy/dev_config.yaml b/litellm/proxy/dev_config.yaml index 4e9c710d446..f78431f694b 100644 --- a/litellm/proxy/dev_config.yaml +++ b/litellm/proxy/dev_config.yaml @@ -199,9 +199,10 @@ model_list: general_settings: master_key: sk-1234 - # Opt-in: let CheckBatchCost track cost for unmanaged Vertex batches created with a raw gs:// input_file_id. - # Requires a vertex_ai deployment configured for the batched model. Defaults to false. - # track_unmanaged_vertex_batch_cost: true + # Opt-in: let CheckBatchCost track cost for unmanaged batches created with a raw + # gs:// (Vertex) or s3:// (Bedrock) input_file_id. Requires a matching deployment + # configured for the batched model. Defaults to false. + # track_unmanaged_batch_cost: true sandbox_tools: - sandbox_tool_name: e2b_sandbox diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 6cdf49818e6..8861a9e3bbc 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -34,6 +34,12 @@ def is_text_content_call_type(call_type: str) -> bool: TEXT_PART_TYPES: FrozenSet[str] = frozenset({"text", "input_text", "output_text"}) +# Responses-API item types whose ``output`` field carries user/tool text +# that guardrails should inspect. ``function_call_output`` is the +# built-in shape; ``custom_tool_call_output`` is the custom-tool +# counterpart (see ``ChatCompletionCustomToolCallOutput``). +_OUTPUT_ITEM_TYPES: frozenset[str] = frozenset({"function_call_output", "custom_tool_call_output"}) + def _iter_text_parts_in_content(content: Any) -> Iterator[str]: """Yield text fragments from a ``message.content`` value (string or @@ -72,7 +78,7 @@ def _coerce_input_to_messages(input_value: Any) -> List[Dict[str, Any]]: messages.append({"role": item.get("role") or "user", "content": [item]}) elif "content" in item: messages.append({"role": item.get("role") or "user", "content": item["content"]}) - elif item.get("type") == "function_call_output" and "output" in item: + elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: messages.append({"role": item.get("role") or "tool", "content": item["output"]}) return messages @@ -157,7 +163,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: input_value[idx] = {**item, "text": visit(item["text"])} elif "content" in item: item["content"] = _rewrite_content(item["content"]) - elif item.get("type") == "function_call_output" and "output" in item: + elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: item["output"] = _rewrite_content(item["output"]) return visited diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 9c4cef2f06b..31ce0cc2214 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -772,7 +772,13 @@ class LassoGuardrail(CustomGuardrail): data: Request data (used for conversation_id generation and tools extraction) cache: Cache instance for storing conversation_id (optional for post-call) """ - payload: Dict[str, Any] = {"messages": messages, "messageType": message_type} + payload: Dict[str, Any] = { + "messages": messages, + "messageType": message_type, + # Drives the "Used By" badge on Lasso Application API Keys: every call from this + # integration is attributed as "litellm" on the keys list. + "source": {"type": "litellm"}, + } # Add optional parameters if available if self.user_id: diff --git a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py index 1b663c16d5b..7fe942bcb38 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py @@ -44,12 +44,18 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys +from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus +from litellm.types.utils import ( + GenericGuardrailAPIInputs, + GuardrailStatus, + StandardLoggingGuardrailInformation, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( @@ -60,6 +66,13 @@ if TYPE_CHECKING: ) +def _sanitize_scan_result_for_logging(scan_result: dict) -> dict: + without_secrets = {key: value for key, value in scan_result.items() if key != "secret_fields"} + redacted = redact_nested_match_and_regex_keys(without_secrets) + masked = mask_credentials_in_payload(redacted if isinstance(redacted, dict) else without_secrets) + return masked if isinstance(masked, dict) else without_secrets + + _DEFAULT_API_BASE = "https://api-xecguard.cycraft.ai" _SCAN_ENDPOINT = "/xecguard/v1/scan" _GROUNDING_ENDPOINT = "/xecguard/v1/grounding" @@ -246,16 +259,21 @@ class XecGuardGuardrail(CustomGuardrail): "guardrail_intervened" if scan_result.get("decision") == "UNSAFE" else "success" ) end_time = datetime.now() - kwargs["standard_logging_object"]["guardrail_information"] = { - "duration": (end_time - start_time).total_seconds(), - "end_time": end_time.timestamp(), - "guardrail_mode": "logging_only", - "guardrail_name": "xecguard", - "guardrail_response": scan_result, - "guardrail_status": guardrail_status, - "masked_entity_count": None, - "start_time": start_time.timestamp(), - } + slg = StandardLoggingGuardrailInformation( + guardrail_name=self.guardrail_name or "xecguard", + guardrail_mode=GuardrailEventHooks.logging_only, + guardrail_response=_sanitize_scan_result_for_logging(scan_result), + guardrail_status=guardrail_status, + start_time=start_time.timestamp(), + end_time=end_time.timestamp(), + duration=(end_time - start_time).total_seconds(), + masked_entity_count=None, + ) + existing = kwargs["standard_logging_object"].get("guardrail_information") + if isinstance(existing, list): + existing.append(slg) + else: + kwargs["standard_logging_object"]["guardrail_information"] = [slg] except Exception as exc: verbose_proxy_logger.debug( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b6342f4fa1a..b839426fcda 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -28,11 +28,20 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( + CallTypes, StandardLoggingPayload, StandardLoggingPayloadErrorInformation, ) from litellm.utils import get_end_user_id_for_cost_tracking +_PASS_THROUGH_CALL_TYPES: frozenset[str] = frozenset( + { + CallTypes.pass_through.value, + CallTypes.llm_passthrough_route.value, + CallTypes.allm_passthrough_route.value, + } +) + class _ProxyDBLogger(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -219,11 +228,13 @@ class _ProxyDBLogger(CustomLogger): verbose_proxy_logger.debug( f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}" ) + call_type: Optional[str] = kwargs.get("call_type") if _should_track_cost_callback( user_api_key=user_api_key, user_id=user_id, team_id=team_id, end_user_id=end_user_id, + call_type=call_type, ): ## UPDATE DATABASE await _update_database_and_spend_counters( @@ -412,9 +423,15 @@ def _should_track_cost_callback( user_id: Optional[str], team_id: Optional[str], end_user_id: Optional[str], + call_type: Optional[str] = None, ) -> bool: """ Determine if the cost callback should be tracked based on the kwargs + + Pass-through endpoints can be configured with ``auth=false``, which leaves + the request with no key/user/team/end-user to attribute spend to. Those + requests still forward real provider traffic that operators expect to see + in request/usage logs, so they are tracked even when unauthenticated. """ # don't run track cost callback if user opted into disabling spend @@ -423,7 +440,7 @@ def _should_track_cost_callback( if user_api_key is not None or user_id is not None or team_id is not None or end_user_id is not None: return True - return False + return call_type in _PASS_THROUGH_CALL_TYPES def _get_budget_reservation_from_metadata(metadata: dict) -> Optional[dict]: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b128b0ea57e..df311bed7b2 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -42,7 +42,7 @@ from litellm.proxy._experimental.mcp_server.db import ( rotate_mcp_user_env_vars_master_key, ) from litellm.proxy._types import * -from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.proxy._types import LiteLLM_VerificationToken, hash_token from litellm.proxy.auth.auth_checks import ( _delete_cache_key_object, can_team_access_model, @@ -468,7 +468,10 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: Handle the key type. """ key_type = data.key_type - data_json.pop("key_type", None) + if key_type is None: + data_json.pop("key_type", None) + return data_json + data_json["key_type"] = key_type.value if key_type == LiteLLMKeyType.LLM_API: data_json["allowed_routes"] = ["llm_api_routes"] elif key_type == LiteLLMKeyType.MANAGEMENT: @@ -3566,6 +3569,7 @@ async def generate_key_helper_fn( created_by: Optional[str] = None, updated_by: Optional[str] = None, allowed_routes: Optional[list] = None, + key_type: str | None = None, sso_user_id: Optional[str] = None, object_permission_id: Optional[str] = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable object_permission: Optional[LiteLLM_ObjectPermissionBase] = None, @@ -3706,6 +3710,7 @@ async def generate_key_helper_fn( "created_by": created_by, "updated_by": updated_by, "allowed_routes": allowed_routes or [], + "key_type": key_type, "object_permission_id": object_permission_id, "router_settings": router_settings_json, "access_group_ids": access_group_ids or [], @@ -3772,7 +3777,10 @@ async def generate_key_helper_fn( return user_data ## CREATE KEY - verbose_proxy_logger.debug("prisma_client: Creating Key= %s", key_data) + verbose_proxy_logger.debug( + "prisma_client: Creating Key= %s", + {**key_data, "token": hash_token(token=token)}, + ) create_key_response = await prisma_client.insert_data(data=key_data, table_name="key") key_data["token_id"] = getattr(create_key_response, "token", None) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f468f5ec30b..797f4600857 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3550,7 +3550,7 @@ async def team_info( try: team_info: Optional[BaseModel] = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id}, - include={"object_permission": True}, + include={"litellm_model_table": True, "object_permission": True}, ) if team_info is None: raise Exception diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 065464aa565..d8015bb8031 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -4034,30 +4034,41 @@ class MicrosoftSSOHandler: base_url = MicrosoftSSOHandler.get_graph_api_base_url() # Endpoint to get app role assignments for the given service principal endpoint = f"/servicePrincipals/{service_principal_id}/appRoleAssignedTo" - url = base_url + endpoint + next_link: str | None = base_url + endpoint headers = { "Authorization": f"Bearer {access_token}", "Content-Type": "application/json", } - response = await async_client.get(url, headers=headers) - response_json = response.json() - verbose_proxy_logger.debug(f"Response from service principal app role assigned to: {response_json}") group_ids: List[str] = [] service_principal_teams: List[MicrosoftServicePrincipalTeam] = [] + page_count = 0 - for _object in response_json.get("value", []): - if _object.get("principalType") == "Group": - # Append the group ID to the list - group_ids.append(_object.get("principalId")) - # Append the service principal team to the list - service_principal_teams.append( - MicrosoftServicePrincipalTeam( - principalDisplayName=_object.get("principalDisplayName"), - principalId=_object.get("principalId"), + while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: + response = await async_client.get(next_link, headers=headers) + response_json = response.json() + verbose_proxy_logger.debug(f"Response from service principal app role assigned to: {response_json}") + + for _object in response_json.get("value", []): + if _object.get("principalType") == "Group": + # Append the group ID to the list + group_ids.append(_object.get("principalId")) + # Append the service principal team to the list + service_principal_teams.append( + MicrosoftServicePrincipalTeam( + principalDisplayName=_object.get("principalDisplayName"), + principalId=_object.get("principalId"), + ) ) - ) + + next_link = response_json.get("@odata.nextLink") + page_count += 1 + + if next_link is not None and page_count >= MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: + verbose_proxy_logger.warning( + f"Reached maximum page limit of {MicrosoftSSOHandler.MAX_GRAPH_API_PAGES}. Some service principal group assignments may not be included." + ) return group_ids, service_principal_teams diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f0ca1f6396f..a0b07ece4dd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7882,7 +7882,7 @@ class ProxyStartupEvent: proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client, llm_router=llm_router, - track_unmanaged_vertex_batch_cost=general_settings.get("track_unmanaged_vertex_batch_cost", False), + track_unmanaged_batch_cost=general_settings.get("track_unmanaged_batch_cost", False), ) scheduler.add_job( check_batch_cost_job.check_batch_cost, @@ -14805,6 +14805,7 @@ async def get_config_list( "forward_client_headers_to_llm_api": {"type": "Boolean"}, "mcp_required_fields": {"type": "List"}, "cancel_on_disconnect": {"type": "Boolean"}, + "skip_user_budget_on_team_key": {"type": "Boolean"}, } return_val = [] diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index fb4d8d0b5a3..a23cecc3911 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -422,6 +422,7 @@ model LiteLLM_VerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") @@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 38fd0d3f343..e1a093a4a48 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -123,6 +123,7 @@ async def reserve_budget_for_request( proxy_logging_obj: ProxyLogging, end_user_id: Optional[str] = None, end_user_object: Optional[Any] = None, + skip_user_budget_on_team_key: bool = False, ) -> Optional[dict]: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None @@ -141,6 +142,7 @@ async def reserve_budget_for_request( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, + skip_user_budget_on_team_key=skip_user_budget_on_team_key, ) if not counters: return None @@ -296,6 +298,7 @@ async def _get_budget_counters( proxy_logging_obj: ProxyLogging, end_user_id: Optional[str] = None, end_user_object: Optional[Any] = None, + skip_user_budget_on_team_key: bool = False, ) -> List[_BudgetCounter]: counters: List[_BudgetCounter] = [] @@ -344,8 +347,9 @@ async def _get_budget_counters( ) ) + is_team_key = team_object is not None and team_object.team_id is not None if ( - (team_object is None or team_object.team_id is None) + not (is_team_key and skip_user_budget_on_team_key) and user_object is not None and user_object.user_id is not None and user_object.max_budget is not None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d7649b524aa..62fc28256cd 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3592,7 +3592,10 @@ class PrismaClient: """ start_time = time.time() try: - verbose_proxy_logger.debug("PrismaClient: insert_data: %s", data) + verbose_proxy_logger.debug( + "PrismaClient: insert_data: %s", + {**data, "token": self.hash_token(token=data["token"])} if data.get("token") is not None else data, + ) if table_name == "key": token = data["token"] hashed_token = self.hash_token(token=token) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 92dd5a513b9..12f9be970c7 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -261,7 +261,7 @@ async def aresponses_api_with_mcp( pre_processed_mcp_tools=original_mcp_tools, ) - return LiteLLM_Proxy_MCP_Handler._create_mcp_streaming_response( + mcp_streaming_response = LiteLLM_Proxy_MCP_Handler._create_mcp_streaming_response( input=input, model=model, all_tools=all_tools, @@ -272,6 +272,10 @@ async def aresponses_api_with_mcp( tool_server_map=tool_server_map, **kwargs, ) + await mcp_streaming_response._create_initial_response_iterator() + if mcp_streaming_response._initial_creation_error is not None: + raise mcp_streaming_response._initial_creation_error + return mcp_streaming_response # Determine if we should auto-execute tools should_auto_execute = bool(mcp_tools_with_litellm_proxy) and LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index c705a04963c..d6b45855be5 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -5,6 +5,8 @@ from litellm._uuid import uuid from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, + ErrorEvent, + ErrorEventError, MCPCallArgumentsDeltaEvent, MCPCallArgumentsDoneEvent, MCPCallCompletedEvent, @@ -316,6 +318,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Cache the response ID to ensure consistency across all events self._cached_response_id: Optional[str] = None + self._initial_creation_error: Exception | None = None + self._stream_error: Exception | None = None + self._error_event_emitted = False + self._last_sequence_number = 0 + def _extract_mcp_headers_from_params(self) -> None: """Extract MCP headers from original request params to pass to tool calls""" from typing import Dict, Optional @@ -380,10 +387,31 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): return LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(self.mcp_tools_with_litellm_proxy) + def _make_stream_error_event(self) -> ResponsesAPIStreamingResponse: + err = self._stream_error + status_code = getattr(err, "status_code", None) + return ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, + sequence_number=self._last_sequence_number + 1, + error=ErrorEventError( + type="mcp_gateway_error", + code=str(status_code) if status_code is not None else "internal_error", + message=str(err) if err is not None else "MCP gateway stream failed", + param=None, + ), + ) + def __aiter__(self): return self async def __anext__(self) -> ResponsesAPIStreamingResponse: + chunk = await self._anext_impl() + sequence_number = getattr(chunk, "sequence_number", None) + if isinstance(sequence_number, int) and sequence_number > self._last_sequence_number: + self._last_sequence_number = sequence_number + return chunk + + async def _anext_impl(self) -> ResponsesAPIStreamingResponse: """ Phase-based streaming: 1. initial_response - Stream the first LLM response (includes response.created, response.in_progress, response.output_item.added) @@ -438,10 +466,16 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.phase = "continue_initial_response" return await self.__anext__() self.phase = "finished" + if self._stream_error is not None and not self._error_event_emitted: + self._error_event_emitted = True + return self._make_stream_error_event() raise StopAsyncIteration # Phase 6: Finished if self.phase == "finished": + if self._stream_error is not None and not self._error_event_emitted: + self._error_event_emitted = True + return self._make_stream_error_event() raise StopAsyncIteration # Should not reach here @@ -530,6 +564,12 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined] + if self._cached_response_id is None and hasattr(chunk, "response"): + new_response = getattr(chunk, "response", None) + new_response_id = getattr(new_response, "id", None) if new_response is not None else None + if new_response_id: + self._cached_response_id = new_response_id + # Ensure response ID consistency - update chunk if needed if self._cached_response_id and hasattr(chunk, "response"): response_obj = getattr(chunk, "response", None) @@ -589,6 +629,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): traceback.print_exc() self.base_iterator = None + self._initial_creation_error = e + self._stream_error = e # Don't set phase to "finished" here — let __anext__ emit any # pre-generated MCP discovery events before ending the iteration. @@ -761,6 +803,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if hasattr(follow_up_response, "__aiter__"): self.base_iterator = follow_up_response self.collected_response = None + self._cached_response_id = None except Exception as e: verbose_logger.error(f"Error creating follow-up iterator: {e}") @@ -768,6 +811,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): traceback.print_exc() self.base_iterator = None + self._stream_error = e def __iter__(self): return self diff --git a/litellm/router.py b/litellm/router.py index 6e773a06c7f..6e8127110cc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -179,7 +179,9 @@ from litellm.types.router import ( RouterModelGroupAliasItem, RouterRateLimitError, RouterRateLimitErrorBasic, + RoutingContext, RoutingGroup, + RoutingPlugin, RoutingStrategy, SearchToolTypedDict, ) @@ -299,6 +301,7 @@ class Router: enable_pre_call_checks: bool = False, enable_tag_filtering: bool = False, tag_filtering_match_any: bool = True, + plugins: list[RoutingPlugin] | None = None, retry_after: int = 0, # min time to wait before retrying a failed request retry_policy: Optional[Union[RetryPolicy, dict]] = None, # set custom retries for different exceptions model_group_retry_policy: Dict[str, RetryPolicy] = {}, # set custom retry policies based on model group @@ -477,6 +480,7 @@ class Router: self.complexity_routers: Dict[str, "ComplexityRouter"] = {} self.adaptive_routers: Dict[str, "AdaptiveRouter"] = {} self.quality_routers: Dict[str, "QualityRouter"] = {} + self.routing_plugins: list[RoutingPlugin] = list(plugins) if plugins else [] # Initialize model_group_alias early since it's used in set_model_list self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = ( @@ -7552,7 +7556,11 @@ class Router: if default_model is None and complexity_router_config: tiers = complexity_router_config.get("tiers", {}) # Use MEDIUM tier as fallback default - default_model = tiers.get("MEDIUM") or tiers.get("SIMPLE") + medium = tiers.get("MEDIUM") or tiers.get("SIMPLE") + if isinstance(medium, list): + default_model = medium[0] if medium else None + else: + default_model = medium if default_model is None: raise ValueError( @@ -7589,15 +7597,6 @@ class Router: AdaptiveRouterPostCallHook, ) - for _cb_list in ( - litellm.callbacks, - litellm.success_callback, - litellm.failure_callback, - litellm._async_success_callback, - litellm._async_failure_callback, - ): - litellm.logging_callback_manager.remove_callbacks_by_type(_cb_list, AdaptiveRouterPostCallHook) - for entry in self.model_list or []: lp = entry.get("litellm_params") if isinstance(entry, dict) else entry.litellm_params lp_model = (lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None @@ -7606,15 +7605,29 @@ class Router: model_name = entry.get("model_name") if isinstance(entry, dict) else entry.model_name if not model_name or not lp: continue - if model_name in self.adaptive_routers: - continue deployment = Deployment( model_name=model_name, litellm_params=(lp if not isinstance(lp, dict) else LiteLLM_Params(**lp)), model_info=(entry.get("model_info") if isinstance(entry, dict) else entry.model_info), ) + if model_name in self.adaptive_routers: + continue self.init_adaptive_router_deployment(deployment=deployment) + for model_name, complexity_router in self.complexity_routers.items(): + if not complexity_router.config.adaptive or model_name in self.adaptive_routers: + continue + adaptive_router = complexity_router._ensure_adaptive_router() + if adaptive_router is not None: + self.adaptive_routers[model_name] = adaptive_router + + for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): + litellm.logging_callback_manager.remove_callback_from_all_lists(callback) + for adaptive_router in self.adaptive_routers.values(): + litellm.logging_callback_manager.add_litellm_callback( + AdaptiveRouterPostCallHook(adaptive_router=adaptive_router) + ) + def init_adaptive_router_deployment(self, deployment: Deployment) -> None: """ Build an AdaptiveRouter instance for this deployment and register its @@ -10321,6 +10334,12 @@ class Router: metadata_variable_name=self._get_metadata_variable_name_from_kwargs(request_kwargs), ) + # narrow to whatever `self.routing_plugins` left in candidate_models + healthy_deployments = self._filter_by_routing_plugin_candidates( + healthy_deployments=healthy_deployments, + request_kwargs=request_kwargs, + ) + ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order = (request_kwargs or {}).pop("_target_order", None) healthy_deployments = litellm.utils._get_order_filtered_deployments( @@ -10596,6 +10615,76 @@ class Router: ) raise e + async def _run_routing_plugins( + self, + model: str, + request_kwargs: dict, + messages: list[dict[str, Any]] | None, + ) -> RoutingContext: + """ + Build a RoutingContext for `model`, run it through `self.routing_plugins` + in order, then stash the narrowed candidate list and accumulated signals + onto `request_kwargs["metadata"]` so `_filter_by_routing_plugin_candidates` + (called later, during healthy-deployment filtering) can consume them. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + resolve_structured_messages, + ) + + deployments = self.get_model_list(model_name=model) or [] + candidate_models = [ + d["litellm_params"]["model"] for d in deployments if d.get("litellm_params", {}).get("model") + ] + + metadata_key = self._get_metadata_variable_name_from_kwargs(request_kwargs) + metadata = request_kwargs.setdefault(metadata_key, {}) + + context = RoutingContext( + raw_messages=messages or [], + structured_messages=resolve_structured_messages(messages=messages, request_kwargs=request_kwargs) or [], + candidate_models=candidate_models, + metadata=metadata, + ) + + for plugin in self.routing_plugins: + context = await plugin.run(context) + + metadata["routing_plugin_signals"] = context.signals + if len(context.candidate_models) < len(candidate_models): + metadata["_routing_plugin_candidate_models"] = context.candidate_models + + return context + + def _filter_by_routing_plugin_candidates( + self, + healthy_deployments: Union[list[dict], dict], + request_kwargs: dict, + ) -> Union[list[dict], dict]: + """ + Narrow `healthy_deployments` to whatever `self.routing_plugins` left in + `context.candidate_models`. Raises rather than silently falling back to + the unfiltered pool -- a plugin narrowing to nothing is a policy decision + (e.g. no model this tenant's budget allows), not something to bypass. + """ + if not self.routing_plugins or not isinstance(healthy_deployments, list): + return healthy_deployments + + metadata_key = self._get_metadata_variable_name_from_kwargs(request_kwargs) + candidate_models = (request_kwargs.get(metadata_key) or {}).get("_routing_plugin_candidate_models") + # `is None` (not falsy-check): a plugin narrowing to an empty list must + # still hit the "no deployments left" raise below, not be treated the + # same as "no plugin ever set this key". + if candidate_models is None: + return healthy_deployments + + candidate_set = set(candidate_models) + filtered = [d for d in healthy_deployments if d.get("litellm_params", {}).get("model") in candidate_set] + + if not filtered: + raise ValueError(f"No deployments left after routing-plugin filtering. candidate_models={candidate_models}") + + return filtered + async def async_pre_routing_hook( self, model: str, @@ -10610,55 +10699,47 @@ class Router: Used for the litellm auto-router to modify the request before the routing decision is made. """ ######################################################### - # Check if any auto-router should be used + # Run the routing-plugin pipeline, if any plugins are configured. + # Plugins narrow the candidate deployment pool (consumed later by + # `_filter_by_routing_plugin_candidates`) and may attach signals for + # downstream strategies (auto-router, complexity-router, ...) to read. ######################################################### - if model in self.auto_routers: - return await self.auto_routers[model].async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) + if self.routing_plugins: + await self._run_routing_plugins(model=model, request_kwargs=request_kwargs, messages=messages) - ######################################################### - # Check if any complexity-router should be used - ######################################################### - if model in self.complexity_routers: - return await self.complexity_routers[model].async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) + router_strategy = ( + self.auto_routers.get(model) + or self.complexity_routers.get(model) + or self.adaptive_routers.get(model) + or self.quality_routers.get(model) + ) + if router_strategy is None: + return None - ######################################################### - # Check if an adaptive-router should be used - ######################################################### - adaptive_router = self.adaptive_routers.get(model) - if adaptive_router is not None: - return await adaptive_router.async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) + pre_routing_hook_response = await router_strategy.async_pre_routing_hook( + model=model, + request_kwargs=request_kwargs, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) - ######################################################### - # Check if any quality-router should be used - ######################################################### - if model in self.quality_routers: - return await self.quality_routers[model].async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) + # `model` (the alias, e.g. "smart-router") is never the deployment actually + # called - apply the alias's own litellm_params (besides `model` itself, + # which is just the alias marker) to the request, since the tier/route + # deployment the hook selected won't have them. Router-only fields + # (tpm, rpm, weight, complexity_router_config, ...) are excluded from the + # actual outbound LLM call downstream by litellm.types.utils.all_litellm_params, + # not here. + if pre_routing_hook_response is not None: + alias_index = self.model_name_to_deployment_indices.get(model, []) + if alias_index: + alias_litellm_params = self.model_list[alias_index[0]].get("litellm_params", {}) + for key, value in alias_litellm_params.items(): + if key != "model" and value is not None: + request_kwargs.setdefault(key, value) - return None + return pre_routing_hook_response def get_available_deployment( self, @@ -10671,6 +10752,18 @@ class Router: """ Returns the deployment based on routing strategy """ + if self.routing_plugins: + raise ValueError( + "Router(plugins=[...]) is configured, but this call resolved to the synchronous " + "deployment-selection path, which never runs the routing-plugin pipeline. This " + "happens for sync Router methods (e.g. Router.completion()) and for async calls " + "with a routing_strategy that has no async-native selector (e.g. legacy " + "'usage-based-routing', v1). Silently skipping " + "configured plugins would let a policy plugin (e.g. a deny-all rule) be bypassed. " + "Use an async Router method with a supported routing_strategy (simple-shuffle, " + "usage-based-routing-v2, cost-based-routing, latency-based-routing, least-busy), " + "or remove `plugins` from the Router config." + ) # users need to explicitly call a specific deployment, by setting `specific_deployment = True` as completion()/embedding() kwarg # When this was no explicit we had several issues with fallbacks timing out diff --git a/litellm/router_strategy/adaptive_router/README.md b/litellm/router_strategy/adaptive_router/README.md index 7f5d7aa21d0..09420a8dd9d 100644 --- a/litellm/router_strategy/adaptive_router/README.md +++ b/litellm/router_strategy/adaptive_router/README.md @@ -56,11 +56,10 @@ Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key - **Per-request decision.** Sample once per eligible model, score with `quality_weight·sample + cost_weight·normalized_cost`, pick the argmax. Routing is stateless per-turn — no sticky lookup. Each call resamples. -- **Owner-cache attribution.** Post-call, the conversation's first picked - model claims an "owner slot" for `OWNER_CACHE_TTL_SECONDS` (24h). Later - turns of the same conversation only fire bandit/state updates if the - same model handled them — mismatches are dropped (no attribution) and - counted in `skipped_updates_total`. Conversation identity is the +- **Previous-response attribution.** Post-call, feedback from the current user + message is attributed to the model that produced the previous response, while + response signals are attributed to the current model. Contexts expire after + 24 hours and the in-memory cache is capped at 1,024 sessions. Conversation identity is the client-supplied `litellm_session_id` if present, otherwise a sha256 over caller identity (api key hash, team, user, end-user) + the first message. - **Per-turn updates.** `satisfaction → +α`. `misalignment, stagnation, @@ -76,12 +75,6 @@ Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key model can still be picked. - **Hard sample cap at 200.** Once `α + β > 200`, deltas are silently dropped. No rescaling — drift is a v1 concern. -- **24h owner-cache TTL.** No explicit eviction below TTL. The in-memory map - can grow if traffic patterns produce many one-shot sessions. -- **Owner-recovery skew.** If model A "owns" a conversation but is then - dethroned in the bandit, later turns served by model B are dropped — so - bandit updates for that conversation flatline until A's TTL expires. - Tracked via `skipped_updates_total`. - **Signals are regex + tool-call only.** No LLM-judge, no embedding similarity, no exemplar storage. Signals are best-effort and biased toward English. - **One AdaptiveRouter per `Router`.** Multiple `adaptive_router/*` deployments diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 69d6a019e68..ec84eb1decf 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -3,25 +3,21 @@ Main adaptive router strategy. See README.md for design overview. One AdaptiveRouter instance per router_name. Holds in-memory caches: - _cells: Beta(alpha, beta) bandit posteriors per (request_type, model) -- _owner_cache: session_key -> (owner_model, expires_at) — the first model - picked for a conversation owns its bandit-update slot - _session_states: (session_key, model) -> SessionState for incremental signal updates Owns the AdaptiveRouterUpdateQueue used by the proxy's flusher to persist state and session snapshots back to Postgres. -Routing is stateless per-turn (Thompson sample fresh on every call). The -owner cache is consulted only at post-call time to decide whether a turn's -signals should fire a bandit update — turns served by a different model than -the conversation's owner are skipped to avoid cross-model misattribution. +Routing is stateless per-turn (Thompson sample fresh on every call). """ from __future__ import annotations import asyncio import time -from dataclasses import asdict -from typing import Any, Dict, List, Optional, Tuple, Union, cast +from collections import OrderedDict +from dataclasses import asdict, dataclass +from typing import Any, Union, cast from litellm._logging import verbose_router_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -38,13 +34,18 @@ from litellm.router_strategy.adaptive_router.config import ( ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, MIN_QUALITY_TIER_HEADER, MIN_QUALITY_TIER_METADATA_KEY, + MIN_TURNS_FOR_CLEAN_CREDIT, OWNER_CACHE_TTL_SECONDS, ) from litellm.router_strategy.adaptive_router.signals import ( SessionState, SignalDelta, Turn, - apply_turn, + advance_session_state, + apply_signal_delta, + detect_response_signals, + detect_user_feedback, + merge_signal_deltas, ) from litellm.router_strategy.adaptive_router.update_queue import ( AdaptiveRouterUpdateQueue, @@ -53,8 +54,7 @@ from litellm.router_strategy.adaptive_router.update_queue import ( # Sweep session-state cache when it exceeds this many live entries. Expired # entries are dropped in bulk; amortizes to O(1) per insert. _SESSION_STATE_SWEEP_THRESHOLD: int = 1024 -# Same pattern for the owner cache. -_OWNER_CACHE_SWEEP_THRESHOLD: int = 1024 +_FEEDBACK_CONTEXT_MAX_ENTRIES: int = 1024 from litellm.repositories.table_repositories import AdaptiveRouterStateRepository from litellm.types.llms.openai import AllMessageValues from litellm.types.router import ( @@ -70,6 +70,17 @@ def _default_prefs() -> AdaptiveRouterPreferences: return AdaptiveRouterPreferences(quality_tier=2, strengths=[]) +@dataclass(frozen=True, slots=True) +class _FeedbackContext: + model_name: str + request_type: RequestType + user_content: str | None + assistant_content: str | None + turn_count: int + clean_credit_awarded: bool + expires_at: float + + class AdaptiveRouter: """One instance per router_name. Holds in-memory caches + the update queue.""" @@ -77,8 +88,8 @@ class AdaptiveRouter: self, router_name: str, config: AdaptiveRouterConfig, - model_to_prefs: Dict[str, AdaptiveRouterPreferences], - model_to_cost: Dict[str, float], + model_to_prefs: dict[str, AdaptiveRouterPreferences], + model_to_cost: dict[str, float], ) -> None: self.router_name = router_name self.config = config @@ -86,13 +97,14 @@ class AdaptiveRouter: self.model_to_cost = model_to_cost self.queue = AdaptiveRouterUpdateQueue() - self._cells: Dict[Tuple[RequestType, str], BanditCell] = {} - self._owner_cache: Dict[str, Tuple[str, float]] = {} - self._session_states: Dict[Tuple[str, str], SessionState] = {} - # Parallel expiry map for _session_states, same TTL as _owner_cache. - # Evicted opportunistically in `get_or_create_session_state`. - self._session_states_expiry: Dict[Tuple[str, str], float] = {} - self._skipped_updates_total: int = 0 + self._cells: dict[tuple[RequestType, str], BanditCell] = {} + self._session_states: dict[tuple[str, str], SessionState] = {} + self._feedback_contexts: OrderedDict[str, _FeedbackContext] = OrderedDict() + self._session_states_expiry: dict[tuple[str, str], float] = {} + self._feedback_attributed_total: int = 0 + self._feedback_without_context_total: int = 0 + self._cross_model_feedback_total: int = 0 + self._response_signal_updates_total: int = 0 # Set to True once the proxy flusher has loaded persisted priors from # Postgres. Checked to support lazy-load on hot-reloaded routers. self._state_loaded: bool = False @@ -145,11 +157,11 @@ class AdaptiveRouter: async def async_pre_routing_hook( self, model: str, - request_kwargs: Dict[str, Any], - messages: Optional[List[Dict[str, Any]]] = None, - input: Optional[Union[str, List]] = None, - specific_deployment: Optional[bool] = False, - ) -> Optional[PreRoutingHookResponse]: + request_kwargs: dict[str, Any], + messages: list[dict[str, Any]] | None = None, + input: Union[str, list] | None = None, + specific_deployment: bool | None = False, + ) -> PreRoutingHookResponse | None: """ Plugin entry point invoked by `Router.async_pre_routing_hook` when the inbound `model` matches this adaptive router's `router_name`. @@ -159,11 +171,9 @@ class AdaptiveRouter: post-call hook can surface it as a response header. Routing is stateless per-turn: every call Thompson-samples fresh, - regardless of any prior pick for the same session. Cross-turn - attribution is enforced post-call via the owner cache (see - `claim_or_check_owner`). + regardless of any prior pick for the same session. """ - user_text = get_last_user_message(cast(List[AllMessageValues], messages or [])) or "" + user_text = get_last_user_message(cast(list[AllMessageValues], messages or [])) or "" request_type = classify_prompt(user_text) min_quality_tier = self._extract_min_quality_tier(request_kwargs) @@ -190,7 +200,7 @@ class AdaptiveRouter: async def pick_model( self, request_type: RequestType, - min_quality_tier: Optional[int] = None, + min_quality_tier: int | None = None, ) -> str: """Thompson-sample across eligible models. Stateless per-turn.""" eligible = self._eligible_models(min_quality_tier) @@ -206,44 +216,7 @@ class AdaptiveRouter: cost_weight=self.config.weights.cost, ) - def claim_or_check_owner(self, session_key: str, current_model: str) -> bool: - """Resolve attribution for a turn under stateless routing. - - Returns True iff this turn should fire a bandit/state update. The - first call for a `session_key` claims ownership for `current_model` - and returns True. Subsequent calls return True only if the owner is - still live AND matches `current_model`. Mismatches (a different - model handled this turn) and expired owners both increment - `_skipped_updates_total` and return False — no attribution. - """ - now = time.time() - existing = self._owner_cache.get(session_key) - if existing is not None and existing[1] > now: - owner_model, _ = existing - if owner_model == current_model: - return True - self._skipped_updates_total += 1 - return False - - # Opportunistic bulk sweep — sessions that never come back would - # otherwise pile up here forever. Same threshold pattern as the - # session-state cache. - if len(self._owner_cache) >= _OWNER_CACHE_SWEEP_THRESHOLD: - self._evict_expired_owner_cache(now) - - # No live owner -> claim for current_model. - self._owner_cache[session_key] = ( - current_model, - now + OWNER_CACHE_TTL_SECONDS, - ) - return True - - def _evict_expired_owner_cache(self, now: float) -> None: - expired = [k for k, (_, exp) in self._owner_cache.items() if exp <= now] - for k in expired: - self._owner_cache.pop(k, None) - - async def get_state_snapshot(self) -> Dict[str, Any]: + async def get_state_snapshot(self) -> dict[str, Any]: """In-memory snapshot for the introspection endpoint. Cheap; no DB hit.""" cells = [] for (rt, model), cell in sorted(self._cells.items(), key=lambda kv: (kv[0][0].value, kv[0][1])): @@ -264,7 +237,7 @@ class AdaptiveRouter: ) queue = await self.queue.queue_size() now = time.time() - owner_cache_live = sum(1 for _, exp in self._owner_cache.values() if exp > now) + feedback_contexts_live = sum(1 for context in self._feedback_contexts.values() if context.expires_at > now) return { "router_name": self.router_name, "available_models": list(self.config.available_models), @@ -274,15 +247,18 @@ class AdaptiveRouter: }, "model_costs": dict(self.model_to_cost), "cells": cells, - "owner_cache_live": owner_cache_live, - "skipped_updates_total": self._skipped_updates_total, + "feedback_contexts_live": feedback_contexts_live, + "feedback_attributed_total": self._feedback_attributed_total, + "feedback_without_context_total": self._feedback_without_context_total, + "cross_model_feedback_total": self._cross_model_feedback_total, + "response_signal_updates_total": self._response_signal_updates_total, "queue": queue, } @staticmethod def _extract_min_quality_tier( - request_kwargs: Dict[str, Any], - ) -> Optional[int]: + request_kwargs: dict[str, Any], + ) -> int | None: """Pull `min_quality_tier` from request headers or metadata. Precedence: headers (`x-litellm-min-quality-tier`) over metadata @@ -310,7 +286,7 @@ class AdaptiveRouter: return None return None - def _eligible_models(self, min_quality_tier: Optional[int]) -> List[str]: + def _eligible_models(self, min_quality_tier: int | None) -> list[str]: if min_quality_tier is None: return list(self.config.available_models) return [ @@ -363,17 +339,131 @@ class AdaptiveRouter: request_type: RequestType, turn: Turn, ) -> SignalDelta: - """Apply one turn, push session snapshot + bandit deltas to the queue.""" - state = self.get_or_create_session_state(session_id, model_name, request_type) - delta = apply_turn(state, turn) - verbose_router_logger.debug("AdaptiveRouter[%s]: record_turn delta=%s", self.router_name, delta) + """Attribute feedback to the previous response and response signals to the current model.""" + async with self._lock: + now = time.time() + while self._feedback_contexts: + oldest_context = next(iter(self._feedback_contexts.values())) + if oldest_context.expires_at > now: + break + self._feedback_contexts.popitem(last=False) + previous = self._feedback_contexts.pop(session_id, None) - # Strip the raw conversation content before persisting. The - # last_user/assistant_content and tool_call_history fields are only - # needed in-memory for the next turn's incremental signal detection; - # writing user prompts and tool payloads to the DB would store PII - # for every adaptive-router conversation. Counts + bookkeeping is - # all the persisted row needs. + effective_request_type = ( + previous.request_type if previous is not None and request_type == RequestType.GENERAL else request_type + ) + current_state = self.get_or_create_session_state( + session_id, + model_name, + effective_request_type, + ) + feedback_delta = detect_user_feedback( + previous.user_content if previous else None, + turn.user_content, + turn.tool_results, + allow_satisfaction=( + previous is not None + and not previous.clean_credit_awarded + and previous.turn_count + 1 >= MIN_TURNS_FOR_CLEAN_CREDIT + ), + ) + previous_assistant = previous.assistant_content if previous else None + response_delta = detect_response_signals( + previous_assistant, + turn.assistant_content, + current_state.tool_call_history, + turn.tool_calls, + turn.tool_results, + turn.response_status, + ) + states_to_persist: dict[str, SessionState] = {model_name: current_state} + bandit_deltas: dict[tuple[RequestType, str], SignalDelta] = {} + + if previous is not None: + feedback_state = self.get_or_create_session_state( + session_id, + previous.model_name, + previous.request_type, + ) + apply_signal_delta(feedback_state, feedback_delta) + if feedback_delta.satisfaction: + feedback_state.clean_credit_awarded = True + states_to_persist[previous.model_name] = feedback_state + if feedback_delta.any_fired(): + self._feedback_attributed_total += 1 + if previous.model_name != model_name: + self._cross_model_feedback_total += 1 + bandit_deltas[(previous.request_type, previous.model_name)] = feedback_delta + else: + if feedback_delta.any_fired(): + self._feedback_without_context_total += 1 + initial_failure = SignalDelta(failure=feedback_delta.failure) + apply_signal_delta(current_state, initial_failure) + bandit_deltas[(effective_request_type, model_name)] = initial_failure + + apply_signal_delta(current_state, response_delta) + if self._compute_bandit_delta(response_delta) != (0.0, 0.0): + self._response_signal_updates_total += 1 + current_key = (effective_request_type, model_name) + bandit_deltas[current_key] = merge_signal_deltas( + bandit_deltas.get(current_key, SignalDelta()), + response_delta, + ) + advance_session_state(current_state, turn) + + next_turn_count = (previous.turn_count if previous else 0) + 1 + clean_credit_awarded = bool((previous and previous.clean_credit_awarded) or feedback_delta.satisfaction) + if len(self._feedback_contexts) >= _FEEDBACK_CONTEXT_MAX_ENTRIES: + self._feedback_contexts.popitem(last=False) + self._feedback_contexts[session_id] = _FeedbackContext( + model_name=model_name, + request_type=effective_request_type, + user_content=turn.user_content, + assistant_content=turn.assistant_content, + turn_count=next_turn_count, + clean_credit_awarded=clean_credit_awarded, + expires_at=now + OWNER_CACHE_TTL_SECONDS, + ) + + for state_model, state in states_to_persist.items(): + await self.queue.add_session_state( + session_id, + self.router_name, + state_model, + self._persistable_session_snapshot(state), + ) + + combined_delta = SignalDelta() + for (attribution_type, target_model), delta in bandit_deltas.items(): + combined_delta = merge_signal_deltas(combined_delta, delta) + d_alpha, d_beta = self._compute_bandit_delta(delta) + if d_alpha == 0 and d_beta == 0: + continue + cell_key = (attribution_type, target_model) + self._cells[cell_key] = apply_delta( + self._cells[cell_key], + d_alpha, + d_beta, + ) + await self.queue.add_state_delta( + self.router_name, + attribution_type.value, + target_model, + d_alpha, + d_beta, + ) + + verbose_router_logger.debug( + "AdaptiveRouter[%s]: feedback_target=%s current_model=%s delta=%s", + self.router_name, + previous.model_name if previous else None, + model_name, + combined_delta, + ) + return combined_delta + + @staticmethod + def _persistable_session_snapshot(state: SessionState) -> dict[str, Any]: snapshot = asdict(state) for sensitive in ( "last_user_content", @@ -382,38 +472,10 @@ class AdaptiveRouter: "pending_tool_calls", ): snapshot.pop(sensitive, None) - await self.queue.add_session_state(session_id, self.router_name, model_name, snapshot) - - d_alpha, d_beta = self._compute_bandit_delta(delta) - verbose_router_logger.debug( - "AdaptiveRouter[%s]: bandit delta alpha=%.2f beta=%.2f", - self.router_name, - d_alpha, - d_beta, - ) - if d_alpha != 0 or d_beta != 0: - # For non-GENERAL turns, attribute to the current-turn classification - # so genuine mid-session topic shifts (e.g. code → math) update the - # correct cell. For GENERAL turns ("thanks!", "ok", "sounds good"), fall - # back to the session's original type so closing pleasantries don't - # misattribute the reward. - attribution_type = ( - request_type if request_type != RequestType.GENERAL else RequestType(state.classified_type) - ) - cell_key = (attribution_type, model_name) - self._cells[cell_key] = apply_delta(self._cells[cell_key], d_alpha, d_beta) - await self.queue.add_state_delta( - self.router_name, - attribution_type.value, - model_name, - d_alpha, - d_beta, - ) - - return delta + return snapshot @staticmethod - def _compute_bandit_delta(delta: SignalDelta) -> Tuple[float, float]: + def _compute_bandit_delta(delta: SignalDelta) -> tuple[float, float]: """ Translate per-turn signal deltas into bandit-cell deltas. diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index c3e3f8ca74a..89ae28be227 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -214,10 +214,6 @@ class AdaptiveRouterPostCallHook(CustomLogger): ) -> None: try: messages = kwargs.get("messages") or [] - if len(messages) < SIGNAL_GATE_MIN_MESSAGES: - # Too few turns for any signal to be meaningful — skip. - return - session_key = _resolve_session_key(kwargs) if not session_key: return @@ -233,10 +229,6 @@ class AdaptiveRouterPostCallHook(CustomLogger): if not current_model: return - if not self.adaptive_router.claim_or_check_owner(session_key, current_model): - # A different model owns this conversation — skip attribution. - return - user_text = _last_user_content(messages) assistant_text, tool_calls = _assistant_content_and_tool_calls(response_obj) tool_results = _recent_tool_results(messages) diff --git a/litellm/router_strategy/adaptive_router/signals.py b/litellm/router_strategy/adaptive_router/signals.py index 2fd1d24fbbe..74fa8936098 100644 --- a/litellm/router_strategy/adaptive_router/signals.py +++ b/litellm/router_strategy/adaptive_router/signals.py @@ -14,7 +14,7 @@ from __future__ import annotations import re from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Set +from typing import Any from litellm.router_strategy.adaptive_router.config import ( LOOP_REPEAT_THRESHOLD, @@ -74,26 +74,26 @@ class SessionState: loop_count: int = 0 exhaustion_count: int = 0 - last_user_content: Optional[str] = None - last_assistant_content: Optional[str] = None - tool_call_history: List[str] = field(default_factory=list) - pending_tool_calls: Dict[str, str] = field(default_factory=dict) + last_user_content: str | None = None + last_assistant_content: str | None = None + tool_call_history: list[str] = field(default_factory=list) + pending_tool_calls: dict[str, str] = field(default_factory=dict) turn_count: int = 0 last_processed_turn: int = -1 clean_credit_awarded: bool = False - terminal_status: Optional[int] = None + terminal_status: int | None = None @dataclass class Turn: """One turn of input. Caller assembles this from the request/response.""" - user_content: Optional[str] = None - assistant_content: Optional[str] = None - tool_calls: List[Dict[str, Any]] = field(default_factory=list) - tool_results: List[Dict[str, Any]] = field(default_factory=list) - response_status: Optional[int] = None + user_content: str | None = None + assistant_content: str | None = None + tool_calls: list[dict[str, Any]] = field(default_factory=list) + tool_results: list[dict[str, Any]] = field(default_factory=list) + response_status: int | None = None # ---- Detection helpers ---------------------------------------------------- @@ -101,13 +101,13 @@ class Turn: _TOKEN_RE = re.compile(r"[A-Za-z0-9]+") -def _tokens(text: Optional[str]) -> Set[str]: +def _tokens(text: str | None) -> set[str]: if not text: return set() return {t.lower() for t in _TOKEN_RE.findall(text)} -def _jaccard(a: Set[str], b: Set[str]) -> float: +def _jaccard(a: set[str], b: set[str]) -> float: union = a | b if not union: return 0.0 @@ -130,7 +130,7 @@ _SATISFACTION_PATTERNS = [ ] -def _detect_misalignment(prev_user: Optional[str], curr_user: Optional[str]) -> bool: +def _detect_misalignment(prev_user: str | None, curr_user: str | None) -> bool: """Fires when consecutive user messages share *some* topic (jaccard > 0) but are sufficiently different (jaccard < threshold) — i.e. user is rephrasing, not changing topic, not repeating.""" @@ -140,7 +140,7 @@ def _detect_misalignment(prev_user: Optional[str], curr_user: Optional[str]) -> return 0.0 < j < MISALIGNMENT_JACCARD_THRESHOLD -def _detect_stagnation(prev_asst: Optional[str], curr_asst: Optional[str]) -> bool: +def _detect_stagnation(prev_asst: str | None, curr_asst: str | None) -> bool: """Fires when consecutive assistant messages are near-duplicates.""" if not prev_asst or not curr_asst: return False @@ -148,19 +148,19 @@ def _detect_stagnation(prev_asst: Optional[str], curr_asst: Optional[str]) -> bo return j >= STAGNATION_JACCARD_NEAR_DUP -def _detect_disengagement(curr_user: Optional[str]) -> bool: +def _detect_disengagement(curr_user: str | None) -> bool: if not curr_user: return False return any(p.search(curr_user) for p in _DISENGAGEMENT_PATTERNS) -def _detect_satisfaction(curr_user: Optional[str]) -> bool: +def _detect_satisfaction(curr_user: str | None) -> bool: if not curr_user: return False return any(p.search(curr_user) for p in _SATISFACTION_PATTERNS) -def _detect_failure(tool_results: List[Dict[str, Any]]) -> bool: +def _detect_failure(tool_results: list[dict[str, Any]]) -> bool: """Any tool result explicitly flagged as an error. We do NOT treat empty content as failure — many tools legitimately return @@ -173,7 +173,7 @@ def _detect_failure(tool_results: List[Dict[str, Any]]) -> bool: return False -def _signature(call: Dict[str, Any]) -> str: +def _signature(call: dict[str, Any]) -> str: """Stable signature for loop detection: name + sorted JSON-ish args.""" name = call.get("name") or call.get("function", {}).get("name", "") call_args = call.get("arguments") @@ -184,7 +184,7 @@ def _signature(call: Dict[str, Any]) -> str: return f"{name}({call_args})" -def _detect_loop(history: List[str], new_calls: List[Dict[str, Any]]) -> bool: +def _detect_loop(history: list[str], new_calls: list[dict[str, Any]]) -> bool: """Fires if any new call's signature appears >= LOOP_REPEAT_THRESHOLD-1 times in recent history (so this call would be the Nth).""" if not new_calls: @@ -209,7 +209,7 @@ _EXHAUSTION_KEYWORDS = ( ) -def _detect_exhaustion(status: Optional[int], tool_results: List[Dict[str, Any]]) -> bool: +def _detect_exhaustion(status: int | None, tool_results: list[dict[str, Any]]) -> bool: if status is not None and status in _EXHAUSTION_STATUSES: return True for r in tool_results: @@ -219,39 +219,53 @@ def _detect_exhaustion(status: Optional[int], tool_results: List[Dict[str, Any]] return False -# ---- Public entrypoint ---------------------------------------------------- +def detect_user_feedback( + previous_user_content: str | None, + current_user_content: str | None, + tool_results: list[dict[str, Any]], + allow_satisfaction: bool, +) -> SignalDelta: + return SignalDelta( + misalignment=int(_detect_misalignment(previous_user_content, current_user_content)), + disengagement=int(_detect_disengagement(current_user_content)), + satisfaction=int(allow_satisfaction and _detect_satisfaction(current_user_content)), + failure=int(_detect_failure(tool_results)), + ) -def apply_turn(state: SessionState, turn: Turn) -> SignalDelta: - """ - Detect signals on this turn, mutate state, return the delta. +def detect_response_signals( + previous_assistant_content: str | None, + current_assistant_content: str | None, + tool_call_history: list[str], + tool_calls: list[dict[str, Any]], + tool_results: list[dict[str, Any]], + response_status: int | None, +) -> SignalDelta: + return SignalDelta( + stagnation=int( + _detect_stagnation( + previous_assistant_content, + current_assistant_content, + ) + ), + loop=int(_detect_loop(tool_call_history, tool_calls)), + exhaustion=int(_detect_exhaustion(response_status, tool_results)), + ) - O(1) per turn (no full-history rescan). Only inspects last_*, recent tool history - (which is bounded at TOOL_CALL_HISTORY_MAX), and the new turn payload. - """ - delta = SignalDelta() - if _detect_misalignment(state.last_user_content, turn.user_content): - delta.misalignment = 1 - if _detect_stagnation(state.last_assistant_content, turn.assistant_content): - delta.stagnation = 1 - if _detect_disengagement(turn.user_content): - delta.disengagement = 1 - if _detect_satisfaction(turn.user_content): - # Gate: only award satisfaction credit once per session, and only - # after MIN_TURNS_FOR_CLEAN_CREDIT turns of context. Early "thanks" - # on turn 1-2 is noise, not a validated quality signal. - current_turn_index = state.turn_count + 1 - if not state.clean_credit_awarded and current_turn_index >= MIN_TURNS_FOR_CLEAN_CREDIT: - delta.satisfaction = 1 - state.clean_credit_awarded = True - if _detect_failure(turn.tool_results): - delta.failure = 1 - if _detect_loop(state.tool_call_history, turn.tool_calls): - delta.loop = 1 - if _detect_exhaustion(turn.response_status, turn.tool_results): - delta.exhaustion = 1 +def merge_signal_deltas(*deltas: SignalDelta) -> SignalDelta: + return SignalDelta( + misalignment=sum(delta.misalignment for delta in deltas), + stagnation=sum(delta.stagnation for delta in deltas), + disengagement=sum(delta.disengagement for delta in deltas), + satisfaction=sum(delta.satisfaction for delta in deltas), + failure=sum(delta.failure for delta in deltas), + loop=sum(delta.loop for delta in deltas), + exhaustion=sum(delta.exhaustion for delta in deltas), + ) + +def apply_signal_delta(state: SessionState, delta: SignalDelta) -> None: state.misalignment_count += delta.misalignment state.stagnation_count += delta.stagnation state.disengagement_count += delta.disengagement @@ -260,6 +274,8 @@ def apply_turn(state: SessionState, turn: Turn) -> SignalDelta: state.loop_count += delta.loop state.exhaustion_count += delta.exhaustion + +def advance_session_state(state: SessionState, turn: Turn) -> None: if turn.user_content: state.last_user_content = turn.user_content if turn.assistant_content: @@ -276,4 +292,38 @@ def apply_turn(state: SessionState, turn: Turn) -> SignalDelta: state.turn_count += 1 state.last_processed_turn = state.turn_count + +# ---- Public entrypoint ---------------------------------------------------- + + +def apply_turn(state: SessionState, turn: Turn) -> SignalDelta: + """ + Detect signals on this turn, mutate state, return the delta. + + O(1) per turn (no full-history rescan). Only inspects last_*, recent tool history + (which is bounded at TOOL_CALL_HISTORY_MAX), and the new turn payload. + """ + feedback_delta = detect_user_feedback( + state.last_user_content, + turn.user_content, + turn.tool_results, + allow_satisfaction=(not state.clean_credit_awarded and state.turn_count + 1 >= MIN_TURNS_FOR_CLEAN_CREDIT), + ) + response_delta = detect_response_signals( + state.last_assistant_content, + turn.assistant_content, + state.tool_call_history, + turn.tool_calls, + turn.tool_results, + turn.response_status, + ) + delta = merge_signal_deltas( + feedback_delta, + response_delta, + ) + apply_signal_delta(state, delta) + if delta.satisfaction: + state.clean_credit_awarded = True + advance_session_state(state, turn) + return delta diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 11719b8a18f..bd3b300b558 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -13,9 +13,12 @@ evaluated before either classification strategy and force a tier outright when m Inspired by ClawRouter: https://github.com/BlockRunAI/ClawRouter """ +from __future__ import annotations + import asyncio +import random import re -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, Literal, Union, cast from pydantic import BaseModel @@ -37,6 +40,7 @@ if TYPE_CHECKING: from semantic_router.routers import SemanticRouter from litellm.router import Router + from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.types.router import PreRoutingHookResponse else: Router = Any @@ -62,7 +66,7 @@ Tiers: {prompt}""" -def _append_custom_keywords(base_keywords: list[str], custom_keywords: Optional[list[str]]) -> list[str]: +def _append_custom_keywords(base_keywords: list[str], custom_keywords: list[str] | None) -> list[str]: if not custom_keywords: return base_keywords base_lowered = frozenset(keyword.lower() for keyword in base_keywords) @@ -94,7 +98,7 @@ def _sanitize_user_api_key_auth(auth: Any) -> Any: return auth -def _classifier_call_metadata(metadata: Optional[dict[str, Any]]) -> Optional[dict[str, Any]]: +def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any] | None: if not metadata: return metadata return { @@ -109,7 +113,7 @@ class DimensionScore: __slots__ = ("name", "score", "signal") - def __init__(self, name: str, score: float, signal: Optional[str] = None): + def __init__(self, name: str, score: float, signal: str | None = None): self.name = name self.score = score self.signal = signal @@ -133,9 +137,9 @@ class ComplexityRouter(CustomLogger): def __init__( self, model_name: str, - litellm_router_instance: "Router", - complexity_router_config: Optional[Dict[str, Any]] = None, - default_model: Optional[str] = None, + litellm_router_instance: Router, + complexity_router_config: dict[str, Any] | None = None, + default_model: str | None = None, ): """ Initialize ComplexityRouter. @@ -172,7 +176,7 @@ class ComplexityRouter(CustomLogger): # embeddings are static, only the prompt is embedded per request). The lock # serializes the one-time build so concurrent cold-start requests don't each # construct the index and fire duplicate embedding calls. - self._semantic_routelayer: Optional[SemanticRouter] = None + self._semantic_routelayer: SemanticRouter | None = None self._semantic_routelayer_lock = asyncio.Lock() # Pre-compile regex patterns for efficiency @@ -184,6 +188,10 @@ class ComplexityRouter(CustomLogger): re.compile(r"[a-z]\)\s", re.IGNORECASE), ] + self.adaptive_router: AdaptiveRouter | None = None + self._model_tiers: dict[str, tuple[ComplexityTier, ...]] = {} + self._adaptive_init_attempted = False + verbose_router_logger.debug(f"ComplexityRouter initialized for {model_name} with tiers: {self.config.tiers}") def _estimate_tokens(self, text: str) -> int: @@ -227,12 +235,12 @@ class ComplexityRouter(CustomLogger): def _score_keyword_match( self, text: str, - keywords: List[str], + keywords: list[str], name: str, signal_label: str, - thresholds: Tuple[int, int], # (low, high) - scores: Tuple[float, float, float], # (none, low, high) - ) -> Tuple[DimensionScore, int]: + thresholds: tuple[int, int], # (low, high) + scores: tuple[float, float, float], # (none, low, high) + ) -> tuple[DimensionScore, int]: """Score based on keyword matches using word boundary matching. Returns: @@ -270,7 +278,7 @@ class ComplexityRouter(CustomLogger): return DimensionScore("questionComplexity", 0.5, f"{count} questions") return DimensionScore("questionComplexity", 0, None) - def classify(self, prompt: str, system_prompt: Optional[str] = None) -> Tuple[ComplexityTier, float, List[str]]: + def classify(self, prompt: str, system_prompt: str | None = None) -> tuple[ComplexityTier, float, list[str]]: """ Classify a prompt by complexity. @@ -329,7 +337,7 @@ class ComplexityRouter(CustomLogger): (0, -1.0, -1.0), ) - dimensions: List[DimensionScore] = [ + dimensions: list[DimensionScore] = [ self._score_token_count(estimated_tokens), code_score, reasoning_score, @@ -371,8 +379,8 @@ class ComplexityRouter(CustomLogger): async def aclassify( self, prompt: str, - system_prompt: Optional[str] = None, - request_kwargs: Optional[dict[str, Any]] = None, + system_prompt: str | None = None, + request_kwargs: dict[str, Any] | None = None, ) -> tuple[ComplexityTier, float, list[str]]: """ Classify a prompt by complexity, using the LLM classifier when configured. @@ -395,8 +403,8 @@ class ComplexityRouter(CustomLogger): async def _classify_with_llm( self, prompt: str, - system_prompt: Optional[str] = None, - request_kwargs: Optional[dict[str, Any]] = None, + system_prompt: str | None = None, + request_kwargs: dict[str, Any] | None = None, ) -> ComplexityTier: """Call the configured classifier model and parse its structured tier response.""" llm_config = self.config.classifier_llm_config @@ -437,23 +445,196 @@ class ComplexityRouter(CustomLogger): """ tier_key = tier.value if isinstance(tier, ComplexityTier) else tier - # Check config tiers mapping - model = self.config.tiers.get(tier_key) - if model: - return model + if tier_key in self.config.tiers: + return self._pick_from_tier_value(self.config.tiers[tier_key], tier_key) - # Fallback to default model if configured if self.config.default_model: return self.config.default_model - # Last resort: return MEDIUM tier model or error - medium_model = self.config.tiers.get(ComplexityTier.MEDIUM.value) - if medium_model: - return medium_model + medium_key = ComplexityTier.MEDIUM.value + if medium_key in self.config.tiers: + return self._pick_from_tier_value(self.config.tiers[medium_key], medium_key) raise ValueError(f"No model configured for tier {tier_key} and no default_model set") - def _lexical_tier_override(self, user_message: str) -> Optional[ComplexityTier]: + @staticmethod + def _pick_from_tier_value(model: str | list[str], tier_key: str) -> str: + if isinstance(model, str): + return model + if not model: + raise ValueError(f"Empty model pool for tier {tier_key}") + return random.choice(model) + + def _tier_pools(self) -> dict[str, list[str]]: + return {tier: (models if isinstance(models, list) else [models]) for tier, models in self.config.tiers.items()} + + def _ensure_adaptive_router(self) -> Any | None: + if not self.config.adaptive: + return None + if self.adaptive_router is not None: + return self.adaptive_router + if self._adaptive_init_attempted: + return self.adaptive_router + self._adaptive_init_attempted = True + + from litellm.router_strategy.adaptive_router.adaptive_router import ( + AdaptiveRouter, + ) + from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + ) + from litellm.types.router import ( + AdaptiveRouterConfig, + AdaptiveRouterPreferences, + ) + + pools = self._tier_pools() + available_models = list(dict.fromkeys(model for models in pools.values() for model in models)) + self._model_tiers = { + model: tuple(ComplexityTier(tier_name) for tier_name, models in pools.items() if model in models) + for model in available_models + } + + model_to_prefs: dict[str, AdaptiveRouterPreferences] = {} + model_to_cost: dict[str, float] = {} + model_list = getattr(self.litellm_router_instance, "model_list", None) or [] + name_to_indices = getattr(self.litellm_router_instance, "model_name_to_deployment_indices", {}) or {} + for name in available_models: + indices = name_to_indices.get(name, []) + if not indices: + model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[]) + model_to_cost[name] = 0.0 + continue + deployment = model_list[indices[0]] + mi = deployment.get("model_info") if isinstance(deployment, dict) else deployment.model_info + mi_dict: dict[str, Any] = mi if isinstance(mi, dict) else (mi.model_dump() if mi else {}) + prefs_raw = mi_dict.get("adaptive_router_preferences") + if prefs_raw is not None: + model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) + else: + model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[]) + + lp = deployment.get("litellm_params") if isinstance(deployment, dict) else deployment.litellm_params + lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {}) + cost = lp_dict.get("input_cost_per_token") + model_to_cost[name] = float(cost) if cost is not None else 0.0 + + self.adaptive_router = AdaptiveRouter( + router_name=self.model_name, + config=AdaptiveRouterConfig( + available_models=available_models, + weights=self.config.adaptive_weights, + ), + model_to_prefs=model_to_prefs, + model_to_cost=model_to_cost, + ) + self._adaptive_chosen_model_key = ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY + return self.adaptive_router + + def _soft_floor_pick( + self, + classified_tier: ComplexityTier, + user_message: str, + request_kwargs: dict[str, Any] | None = None, + ) -> str: + from litellm.router_strategy.adaptive_router.bandit import ( + normalized_cost, + thompson_sample, + ) + from litellm.router_strategy.adaptive_router.classifier import classify_prompt + + adaptive = self._ensure_adaptive_router() + if adaptive is None: + return self.get_model_for_tier(classified_tier) + + request_type = classify_prompt(user_message) + classified_idx = TIER_SEVERITY_ORDER.index(classified_tier) + pools = self._tier_pools() + classified_candidates = tuple(pools.get(classified_tier.value, ())) + cold_start_candidates = tuple( + model for model in classified_candidates if adaptive._cells[(request_type, model)].total_samples == 0 + ) + if cold_start_candidates: + chosen_model = random.choice(cold_start_candidates) + if request_kwargs is not None: + metadata = request_kwargs.setdefault("metadata", {}) + if isinstance(metadata, dict): + metadata["adaptive_router_decision"] = { + "phase": "cold_start", + "classified_tier": classified_tier.value, + "request_type": request_type.value, + "eligible_mode": "classified_tier", + "quality_weight": self.config.adaptive_weights.quality, + "cost_weight": self.config.adaptive_weights.cost, + "tier_distance_penalty": self.config.tier_distance_penalty, + "chosen_model": chosen_model, + "candidates": [ + { + "model": model, + "total_samples": adaptive._cells[(request_type, model)].total_samples, + } + for model in cold_start_candidates + ], + } + return chosen_model + if self.config.adaptive_eligible == "classified_tier": + candidates = list(classified_candidates) + if not candidates: + return self.get_model_for_tier(classified_tier) + else: + candidates = list(adaptive.config.available_models) + + all_costs = [adaptive.model_to_cost.get(m, 0.0) for m in candidates] + quality_weight = self.config.adaptive_weights.quality + cost_weight = self.config.adaptive_weights.cost + penalty_weight = self.config.tier_distance_penalty + + best_model: str | None = None + best_score = float("-inf") + candidate_scores: list[dict[str, Any]] = [] + for model in candidates: + cell = adaptive._cells[(request_type, model)] + quality_sample = thompson_sample(cell) + cost_score = normalized_cost(adaptive.model_to_cost.get(model, 0.0), all_costs) + if self.config.adaptive_eligible == "classified_tier": + distance = 0 + else: + model_tiers = self._model_tiers.get(model, (classified_tier,)) + distance = min( + abs(TIER_SEVERITY_ORDER.index(model_tier) - classified_idx) for model_tier in model_tiers + ) + score = quality_weight * quality_sample + cost_weight * cost_score - penalty_weight * distance + candidate_scores.append( + { + "model": model, + "quality_sample": quality_sample, + "cost_score": cost_score, + "tier_distance": distance, + "score": score, + } + ) + if score > best_score: + best_score = score + best_model = model + if best_model is None: + return self.get_model_for_tier(classified_tier) + if request_kwargs is not None: + metadata = request_kwargs.setdefault("metadata", {}) + if isinstance(metadata, dict): + metadata["adaptive_router_decision"] = { + "phase": "adaptive", + "classified_tier": classified_tier.value, + "request_type": request_type.value, + "eligible_mode": self.config.adaptive_eligible, + "quality_weight": quality_weight, + "cost_weight": cost_weight, + "tier_distance_penalty": penalty_weight, + "chosen_model": best_model, + "candidates": candidate_scores, + } + return best_model + + def _lexical_tier_override(self, user_message: str) -> ComplexityTier | None: """When keyword_tier_rules match literally, the most-severe matched tier wins. Escalating to the highest tier (rather than the first rule in the list) keeps @@ -471,7 +652,7 @@ class ComplexityRouter(CustomLogger): return None return max(matched_tiers, key=TIER_SEVERITY_ORDER.index) - def _get_or_create_semantic_routelayer(self) -> "SemanticRouter": + def _get_or_create_semantic_routelayer(self) -> SemanticRouter: """Build (once) a SemanticRouter with one route per tier, utterances = that tier's keywords.""" if self._semantic_routelayer is not None: return self._semantic_routelayer @@ -510,7 +691,7 @@ class ComplexityRouter(CustomLogger): self._semantic_routelayer = routelayer return routelayer - async def _ensure_semantic_routelayer(self) -> "SemanticRouter": + async def _ensure_semantic_routelayer(self) -> SemanticRouter: """Return the cached route layer, building it once under a lock if needed. The build embeds the static route utterances via the encoder's synchronous path, @@ -526,7 +707,7 @@ class ComplexityRouter(CustomLogger): routelayer = await asyncio.to_thread(self._get_or_create_semantic_routelayer) return routelayer - async def _semantic_tier_override(self, user_message: str, request_kwargs: Dict) -> Optional[ComplexityTier]: + async def _semantic_tier_override(self, user_message: str, request_kwargs: dict) -> ComplexityTier | None: """Match the prompt against keyword_tier_rules by embedding similarity. Embeds the query ourselves (instead of letting SemanticRouter.acall embed it @@ -566,7 +747,7 @@ class ComplexityRouter(CustomLogger): except ValueError: return None - async def _resolve_keyword_tier_override(self, user_message: str, request_kwargs: Dict) -> Optional[ComplexityTier]: + async def _resolve_keyword_tier_override(self, user_message: str, request_kwargs: dict) -> ComplexityTier | None: """Resolve a keyword_tier_rule override, semantically or lexically per config. Returns None (no override -> fall through to the scorer) not only when no rule @@ -587,60 +768,28 @@ class ComplexityRouter(CustomLogger): def _resolve_messages( self, - messages: Optional[List[Dict[str, Any]]], - request_kwargs: Dict, - ) -> Optional[List[Dict[str, Any]]]: + messages: list[dict[str, Any]] | None, + request_kwargs: dict, + ) -> list[dict[str, Any]] | None: """ Resolve messages from the request, converting from other formats if needed. Uses the guardrail translation handler dispatch to convert Responses API ``input`` (or other non-chat-completions formats) into OpenAI-spec messages. """ - if messages: - return messages - - from litellm.litellm_core_utils.api_route_to_call_types import ( - get_call_types_for_route, + from litellm.litellm_core_utils.prompt_templates.factory import ( + resolve_structured_messages, ) - from litellm.llms import load_guardrail_translation_mappings - from litellm.types.utils import CallTypes - mappings = load_guardrail_translation_mappings() - call_type: Optional[CallTypes] = None - - # 1. Try route-based inference from proxy metadata - route = request_kwargs.get("litellm_metadata", {}).get("user_api_key_request_route") - if route: - call_types_list = get_call_types_for_route(route) - if call_types_list: - for ct in call_types_list: - if ct in mappings: - call_type = ct - break - - # 2. Fallback: try each mapped handler until one produces messages - handlers_to_try: List[Any] = [] - if call_type is not None and call_type in mappings: - handlers_to_try.append(mappings[call_type]()) - else: - handlers_to_try.extend(handler_cls() for handler_cls in mappings.values()) - - for handler in handlers_to_try: - structured = handler.get_structured_messages(request_kwargs) - if structured: - return [ - msg if isinstance(msg, dict) else msg.model_dump() # type: ignore - for msg in structured - ] - return None + return resolve_structured_messages(messages=messages, request_kwargs=request_kwargs) @staticmethod def _extract_user_message_and_system_prompt( - messages: List[Dict[str, Any]], - ) -> Tuple[Optional[str], Optional[str]]: + messages: list[dict[str, Any]], + ) -> tuple[str | None, str | None]: """Extract the last user message text and last system prompt from messages.""" - user_message: Optional[str] = None - system_prompt: Optional[str] = None + user_message: str | None = None + system_prompt: str | None = None for msg in reversed(messages): role = msg.get("role", "") @@ -660,17 +809,115 @@ class ComplexityRouter(CustomLogger): return user_message, system_prompt + @staticmethod + def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]: + """Metadata may land on `metadata` or `litellm_metadata` depending on the + endpoint, mirroring DeploymentAffinityCheck's precedence.""" + return [ + metadata + for metadata_key in ("litellm_metadata", "metadata") + if isinstance(metadata := request_kwargs.get(metadata_key), dict) + ] + + @staticmethod + def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: + """Resolve a client-supplied session_id.""" + for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): + session_id = metadata.get("session_id") + if session_id is not None: + return str(session_id) + return None + + @staticmethod + def _get_user_api_key_hash_from_request_kwargs(request_kwargs: dict) -> str | None: + """Resolve the proxy-derived API key hash, the same trust boundary + DeploymentAffinityCheck uses for its own key-based affinity (not the + client-supplied OpenAI `user` param, which isn't authenticated).""" + for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): + user_key = metadata.get("user_api_key_hash") + if user_key is not None: + return str(user_key) + return None + + def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: dict) -> str: + # Namespace by the caller's API key hash so two different callers reusing the + # same client-supplied session_id can't poison each other's routing pin. Falls + # back to "unscoped" only when there's no authenticated caller to scope by + # (e.g. direct Router usage without the proxy layer). + caller_scope = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped" + return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}" + async def async_pre_routing_hook( self, model: str, - request_kwargs: Dict, - messages: Optional[List[Dict[str, Any]]] = None, - input: Optional[Union[str, List]] = None, - specific_deployment: Optional[bool] = False, - ) -> Optional["PreRoutingHookResponse"]: + request_kwargs: dict, + messages: list[dict[str, Any]] | None = None, + input: Union[str, list] | None = None, + specific_deployment: bool | None = False, + ) -> PreRoutingHookResponse | None: """ Pre-routing hook called before the routing decision. + When `session_affinity` is enabled and a session_id is resolvable on the request, + pins the model chosen on the session's first turn and reuses it for every later + turn, skipping classification entirely. Otherwise delegates to `_classify_and_route`. + """ + from litellm.types.router import PreRoutingHookResponse + + session_id = self._get_session_id_from_request_kwargs(request_kwargs) if self.config.session_affinity else None + cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None + + if cache_key is not None: + pinned_model = await self.litellm_router_instance.cache.async_get_cache(key=cache_key) + if isinstance(pinned_model, str): + # Refresh the TTL on every hit so an active session doesn't lose its + # pin mid-conversation just because it outlives the original write. + await self.litellm_router_instance.cache.async_set_cache( + key=cache_key, + value=pinned_model, + ttl=self.config.session_affinity_ttl_seconds, + ) + if self.config.adaptive: + from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + ) + + kwargs_metadata = request_kwargs.setdefault("metadata", {}) + if isinstance(kwargs_metadata, dict): + kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = pinned_model + verbose_router_logger.info( + f"ComplexityRouter: routing decision cause=session_affinity_pin, routed_model={pinned_model}" + ) + has_original_messages = messages is not None and len(messages) > 0 + return PreRoutingHookResponse( + model=pinned_model, + messages=messages if has_original_messages else None, + ) + + response = await self._classify_and_route( + model=model, + request_kwargs=request_kwargs, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) + if cache_key is not None and response is not None: + await self.litellm_router_instance.cache.async_set_cache( + key=cache_key, + value=response.model, + ttl=self.config.session_affinity_ttl_seconds, + ) + return response + + async def _classify_and_route( + self, + model: str, + request_kwargs: dict, + messages: list[dict[str, Any]] | None = None, + input: Union[str, list] | None = None, + specific_deployment: bool | None = False, + ) -> PreRoutingHookResponse | None: + """ Classifies the request by complexity and returns the appropriate model. Supports chat completions (messages), Responses API (input), and other formats via the guardrail translation handler dispatch. @@ -719,12 +966,25 @@ class ComplexityRouter(CustomLogger): ) tier, score, signals = await self.aclassify(user_message, system_prompt, request_kwargs) - routed_model = self.get_model_for_tier(tier) - - verbose_router_logger.info( - f"ComplexityRouter: routing decision cause=complexity_scorer, tier={tier.value}, " - f"score={score:.3f}, signals={signals}, routed_model={routed_model}" - ) + if self.config.adaptive: + routed_model = self._soft_floor_pick(tier, user_message, request_kwargs) + adaptive = self._ensure_adaptive_router() + if adaptive is not None: + kwargs_metadata = request_kwargs.setdefault("metadata", {}) + if isinstance(kwargs_metadata, dict): + chosen_key = getattr(self, "_adaptive_chosen_model_key", "adaptive_router_chosen_model") + kwargs_metadata[chosen_key] = routed_model + verbose_router_logger.info( + f"ComplexityRouter[adaptive]: routing decision cause=complexity_scorer, " + f"tier={tier.value}, score={score:.3f}, " + f"signals={signals}, routed_model={routed_model}" + ) + else: + routed_model = self.get_model_for_tier(tier) + verbose_router_logger.info( + f"ComplexityRouter: routing decision cause=complexity_scorer, tier={tier.value}, " + f"score={score:.3f}, signals={signals}, routed_model={routed_model}" + ) return PreRoutingHookResponse( model=routed_model, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 125de6f7489..e4bd36505e6 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -6,9 +6,11 @@ All values are configurable via proxy config.yaml. """ from enum import Enum -from typing import Dict, List, Literal, Optional +from typing import Literal -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + +from litellm.types.router import AdaptiveRouterWeights class ComplexityTier(str, Enum): @@ -27,11 +29,13 @@ TIER_SEVERITY_ORDER: tuple[ComplexityTier, ...] = ( ComplexityTier.REASONING, ) +DEFAULT_TIER_DISTANCE_PENALTY: float = 0.5 + class KeywordTierRule(BaseModel): """A deterministic override: if any keyword matches, route to this tier.""" - keywords: List[str] = Field( + keywords: list[str] = Field( min_length=1, description="Keywords/phrases that trigger this rule (lexical or semantic match)", ) @@ -56,7 +60,7 @@ class KeywordTierRule(BaseModel): # Note: Keywords should be full words/phrases to avoid substring false positives. # The matching logic uses word boundary detection for single-word keywords. -DEFAULT_CODE_KEYWORDS: List[str] = [ +DEFAULT_CODE_KEYWORDS: list[str] = [ "function", "class", "def", @@ -104,7 +108,7 @@ DEFAULT_CODE_KEYWORDS: List[str] = [ "pull request", ] -DEFAULT_REASONING_KEYWORDS: List[str] = [ +DEFAULT_REASONING_KEYWORDS: list[str] = [ "step by step", "think through", "let's think", @@ -126,7 +130,7 @@ DEFAULT_REASONING_KEYWORDS: List[str] = [ "conclude", ] -DEFAULT_TECHNICAL_KEYWORDS: List[str] = [ +DEFAULT_TECHNICAL_KEYWORDS: list[str] = [ "architecture", "distributed", "scalable", @@ -158,7 +162,7 @@ DEFAULT_TECHNICAL_KEYWORDS: List[str] = [ # Note: "async", "kubernetes", "docker" are in DEFAULT_CODE_KEYWORDS ] -DEFAULT_SIMPLE_KEYWORDS: List[str] = [ +DEFAULT_SIMPLE_KEYWORDS: list[str] = [ "what is", "what's", "define", @@ -191,7 +195,7 @@ DEFAULT_SIMPLE_KEYWORDS: List[str] = [ # ─── Default Dimension Weights ─── -DEFAULT_DIMENSION_WEIGHTS: Dict[str, float] = { +DEFAULT_DIMENSION_WEIGHTS: dict[str, float] = { "tokenCount": 0.10, # Reduced - length is less important than content "codePresence": 0.30, # High - code requests need capable models "reasoningMarkers": 0.25, # High - explicit reasoning requests @@ -204,7 +208,7 @@ DEFAULT_DIMENSION_WEIGHTS: Dict[str, float] = { # ─── Default Tier Boundaries ─── -DEFAULT_TIER_BOUNDARIES: Dict[str, float] = { +DEFAULT_TIER_BOUNDARIES: dict[str, float] = { "simple_medium": 0.15, # Lower threshold to catch more MEDIUM cases "medium_complex": 0.35, # Lower threshold to catch technical COMPLEX cases "complex_reasoning": 0.60, # Reasoning tier reserved for explicit reasoning markers @@ -213,7 +217,7 @@ DEFAULT_TIER_BOUNDARIES: Dict[str, float] = { # ─── Default Token Thresholds ─── -DEFAULT_TOKEN_THRESHOLDS: Dict[str, int] = { +DEFAULT_TOKEN_THRESHOLDS: dict[str, int] = { "simple": 15, # Only very short prompts (<15 tokens) are penalized "complex": 400, # Long prompts (>400 tokens) get complexity boost } @@ -221,7 +225,7 @@ DEFAULT_TOKEN_THRESHOLDS: Dict[str, int] = { # ─── Default Tier to Model Mapping ─── -DEFAULT_TIER_MODELS: Dict[str, str] = { +DEFAULT_TIER_MODELS: dict[str, str] = { "SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "claude-sonnet-4-20250514", @@ -244,44 +248,47 @@ class ClassifierLLMConfig(BaseModel): class ComplexityRouterConfig(BaseModel): """Configuration for the ComplexityRouter.""" - # Tier to model mapping - tiers: Dict[str, str] = Field( + # string = pin; list = random pick when adaptive=False, soft-floor home pool when adaptive=True + tiers: dict[str, str | list[str]] = Field( default_factory=lambda: DEFAULT_TIER_MODELS.copy(), - description="Mapping of complexity tiers to model names", + description=( + "Mapping of complexity tiers to a model or model pool. " + "A list is randomly picked from when adaptive=False, and used as a soft-floor home pool when adaptive=True" + ), ) # Tier boundaries (normalized scores) - tier_boundaries: Dict[str, float] = Field( + tier_boundaries: dict[str, float] = Field( default_factory=lambda: DEFAULT_TIER_BOUNDARIES.copy(), description="Score boundaries between tiers", ) # Token count thresholds - token_thresholds: Dict[str, int] = Field( + token_thresholds: dict[str, int] = Field( default_factory=lambda: DEFAULT_TOKEN_THRESHOLDS.copy(), description="Token count thresholds for simple/complex classification", ) # Dimension weights - dimension_weights: Dict[str, float] = Field( + dimension_weights: dict[str, float] = Field( default_factory=lambda: DEFAULT_DIMENSION_WEIGHTS.copy(), description="Weights for each scoring dimension", ) # Keyword lists (overridable) - code_keywords: Optional[List[str]] = Field( + code_keywords: list[str] | None = Field( default=None, description="Keywords indicating code-related content", ) - reasoning_keywords: Optional[List[str]] = Field( + reasoning_keywords: list[str] | None = Field( default=None, description="Keywords indicating reasoning-required content", ) - technical_keywords: Optional[List[str]] = Field( + technical_keywords: list[str] | None = Field( default=None, description="Keywords indicating technical content", ) - custom_technical_keywords: Optional[list[str]] = Field( + custom_technical_keywords: list[str] | None = Field( default=None, description=( "Domain-specific technical keywords appended to the effective base list " @@ -290,13 +297,13 @@ class ComplexityRouterConfig(BaseModel): "the base list and within this list." ), ) - simple_keywords: Optional[List[str]] = Field( + simple_keywords: list[str] | None = Field( default=None, description="Keywords indicating simple/basic queries", ) # Default model if scoring fails - default_model: Optional[str] = Field( + default_model: str | None = Field( default=None, description="Default model to use if tier cannot be determined", ) @@ -306,13 +313,34 @@ class ComplexityRouterConfig(BaseModel): default="heuristic", description="Classification strategy: local regex/keyword scoring, or an LLM call", ) - classifier_llm_config: Optional[ClassifierLLMConfig] = Field( + classifier_llm_config: ClassifierLLMConfig | None = Field( default=None, description="Configuration for the LLM classifier; required when classifier_type is 'llm'", ) + adaptive: bool = Field( + default=False, + description="Enable adaptive bandit selection with soft complexity floors", + ) + adaptive_weights: AdaptiveRouterWeights = Field( + default_factory=lambda: AdaptiveRouterWeights(quality=0.3, cost=0.7), + description="Quality vs cost weights for adaptive selection (used when adaptive=True)", + ) + tier_distance_penalty: float = Field( + default=DEFAULT_TIER_DISTANCE_PENALTY, + ge=0.0, + description="Score penalty per tier-step away from the classified tier when adaptive=True", + ) + adaptive_eligible: Literal["all", "classified_tier"] = Field( + default="all", + description=( + "When adaptive=True: 'all' scores every pool model with a tier-distance penalty (soft floors); " + "'classified_tier' Thompson-samples only inside the classified tier's pool" + ), + ) + # Deterministic keyword -> tier overrides, evaluated before weighted scoring - keyword_tier_rules: Optional[List[KeywordTierRule]] = Field( + keyword_tier_rules: list[KeywordTierRule] | None = Field( default=None, description="Rules that force a specific tier when their keywords match the prompt", ) @@ -322,7 +350,7 @@ class ComplexityRouterConfig(BaseModel): default=False, description="Match keyword_tier_rules by embedding similarity instead of literal text", ) - embedding_model: Optional[str] = Field( + embedding_model: str | None = Field( default=None, description="Embedding model (LiteLLM model name) used when semantic_keyword_matching is enabled", ) @@ -333,14 +361,56 @@ class ComplexityRouterConfig(BaseModel): description="Minimum cosine similarity for a semantic keyword match", ) + # Session affinity: pin the first turn's routed model for the rest of the session + session_affinity: bool = Field( + default=False, + description=( + "When True and a session_id is resolvable on the request, pin the model chosen on the " + "session's first turn and reuse it for every later turn, skipping re-classification." + ), + ) + session_affinity_ttl_seconds: int = Field( + default=3600, + gt=0, + description="TTL for the session affinity pin; refreshed on every cache hit", + ) + model_config = ConfigDict(extra="allow") # Allow additional fields + @field_validator("tiers", mode="before") + @classmethod + def _coerce_tier_values(cls, value: object) -> object: + if not isinstance(value, dict): + return value + coerced: dict[str, object] = {} + for key, item in value.items(): + if isinstance(item, str): + coerced[key] = item + elif isinstance(item, (list, tuple)): + coerced[key] = list(item) + else: + coerced[key] = item + return coerced + @model_validator(mode="after") def _validate_llm_classifier_config(self) -> "ComplexityRouterConfig": if self.classifier_type == "llm" and self.classifier_llm_config is None: raise ValueError("classifier_llm_config is required when classifier_type is 'llm'") return self + @model_validator(mode="after") + def _validate_adaptive_pools(self) -> "ComplexityRouterConfig": + if not self.adaptive: + return self + normalized = {tier: (models if isinstance(models, list) else [models]) for tier, models in self.tiers.items()} + if not any(normalized.values()): + raise ValueError("adaptive=True requires at least one non-empty tier pool") + empty = [tier for tier, models in normalized.items() if not models] + if empty: + raise ValueError(f"adaptive=True tier pools must be non-empty; empty tiers: {empty}") + self.tiers = normalized + return self + @model_validator(mode="after") def _validate_semantic_matching(self) -> "ComplexityRouterConfig": if not self.semantic_keyword_matching: diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 69bccff701f..318ba1f5956 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -213,6 +213,8 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_input_audio_tokens_metric", "litellm_output_reasoning_tokens_metric", "litellm_output_audio_tokens_metric", + "litellm_video_duration_seconds_metric", + "litellm_images_generated_metric", "litellm_deployment_successful_fallbacks", "litellm_deployment_failed_fallbacks", "litellm_remaining_team_budget_metric", @@ -506,6 +508,9 @@ class PrometheusMetricLabels: litellm_output_reasoning_tokens_metric = litellm_output_tokens_metric litellm_output_audio_tokens_metric = litellm_output_tokens_metric + litellm_video_duration_seconds_metric = litellm_output_tokens_metric + litellm_images_generated_metric = litellm_output_tokens_metric + litellm_deployment_state = [ UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, UserAPIKeyLabelNames.MODEL_ID.value, @@ -717,6 +722,8 @@ class PrometheusMetricLabels: "litellm_input_tokens_metric", "litellm_total_tokens_metric", "litellm_output_tokens_metric", + "litellm_video_duration_seconds_metric", + "litellm_images_generated_metric", } ) # Managed batch metrics diff --git a/litellm/types/router.py b/litellm/types/router.py index 4bac9358392..3bedd97c20c 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -9,7 +9,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Union, get_type_hi import httpx from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator -from typing_extensions import Required, TypedDict +from typing_extensions import Protocol, Required, TypedDict from litellm._uuid import uuid @@ -829,6 +829,35 @@ class PreRoutingHookResponse(BaseModel): messages: Optional[List[Dict[str, Any]]] +class RoutingContext(BaseModel): + """ + Passed through a Router's `plugins` pipeline before the routing decision is made. + + Each plugin reads and mutates this object; the next plugin sees the previous + plugin's changes. `candidate_models` narrows as the pipeline runs -- Router + only selects a deployment whose `litellm_params.model` survives the pipeline. + + `raw_messages` and `structured_messages` mirror the pattern + `CustomGuardrail.apply_guardrail` uses: the message shape differs by API + surface (chat completions, Anthropic /v1/messages, Responses API `input`, + ...), so plugins that need a stable, provider-agnostic shape should read + `structured_messages` (normalized to OpenAI chat-completions format); + plugins that need the exact original payload can read `raw_messages`. + """ + + raw_messages: list[dict[str, Any]] + structured_messages: list[dict[str, Any]] + candidate_models: list[str] + metadata: dict[str, Any] = Field(default_factory=dict) + signals: dict[str, Any] = Field(default_factory=dict) + + +class RoutingPlugin(Protocol): + """Interface a custom routing plugin must implement to run in `Router(plugins=[...])`.""" + + async def run(self, context: RoutingContext) -> RoutingContext: ... + + class RequestType(str, enum.Enum): """Fixed v0 taxonomy. User-extensible types come in v1.""" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e33e2335525..8799b622ad1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -209,6 +209,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_query: Optional[float] # only for rerank models input_cost_per_image: Optional[float] # only for vertex ai models input_cost_per_image_token: Optional[float] # for gpt-image-1 and similar models + input_cost_per_video_token: Optional[float] # for gemini omni models with video input input_cost_per_audio_per_second: Optional[float] # only for vertex ai models input_cost_per_video_per_second: Optional[float] # only for vertex ai models input_cost_per_second: Optional[float] # for OpenAI Speech models @@ -234,6 +235,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_character_above_128k_tokens: Optional[float] # only for vertex ai models output_cost_per_image: Optional[float] output_cost_per_image_token: Optional[float] + output_cost_per_video_token: Optional[float] # for gemini omni models with video output output_vector_size: Optional[int] output_cost_per_reasoning_token: Optional[float] output_cost_per_video_per_second: Optional[float] # only for vertex ai models @@ -3046,6 +3048,7 @@ class CustomPricingLiteLLMParams(BaseModel): output_cost_per_character_above_128k_tokens: Optional[float] = None output_cost_per_image: Optional[float] = None output_cost_per_image_token: Optional[float] = None + output_cost_per_video_token: Optional[float] = None output_cost_per_reasoning_token: Optional[float] = None output_cost_per_video_per_second: Optional[float] = None output_cost_per_audio_per_second: Optional[float] = None @@ -3055,6 +3058,7 @@ class CustomPricingLiteLLMParams(BaseModel): cache_read_input_token_cost_above_272k_tokens: Optional[float] = None cache_read_input_token_cost_above_512k_tokens: Optional[float] = None input_cost_per_image_token: Optional[float] = None + input_cost_per_video_token: Optional[float] = None input_cost_per_token_above_272k_tokens: Optional[float] = None input_cost_per_token_above_512k_tokens: Optional[float] = None output_cost_per_token_above_272k_tokens: Optional[float] = None @@ -3210,6 +3214,16 @@ all_litellm_params = ( "_litellm_tpm_reserved_model", "_litellm_tpm_reserved_scopes", "_litellm_tpm_reservation_released", + "auto_router_config_path", + "auto_router_config", + "auto_router_default_model", + "auto_router_embedding_model", + "complexity_router_config", + "complexity_router_default_model", + "adaptive_router_config", + "adaptive_router_default_model", + "quality_router_config", + "quality_router_default_model", ] + list(StandardCallbackDynamicParams.__annotations__.keys()) + list(CustomPricingLiteLLMParams.model_fields.keys()) diff --git a/litellm/utils.py b/litellm/utils.py index 18b89ee0d13..0636d3683b7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5437,6 +5437,7 @@ def _get_model_info_helper( input_cost_per_second=_model_info.get("input_cost_per_second", None), input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None), input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None), + input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None), input_cost_per_image=_model_info.get("input_cost_per_image", None), input_cost_per_audio_per_second=_model_info.get("input_cost_per_audio_per_second", None), input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None), @@ -5480,6 +5481,7 @@ def _get_model_info_helper( output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None), + output_cost_per_video_token=_model_info.get("output_cost_per_video_token", None), output_vector_size=_model_info.get("output_vector_size", None), citation_cost_per_token=_model_info.get("citation_cost_per_token", None), tiered_pricing=_model_info.get("tiered_pricing", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 94d9f6496bd..df6e8c992ef 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11331,6 +11331,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11362,6 +11363,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true @@ -11424,6 +11426,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, @@ -11479,6 +11482,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11506,6 +11510,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true }, @@ -11559,6 +11564,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true @@ -11586,6 +11592,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true @@ -11614,6 +11621,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "provider_specific_entry": { @@ -11648,6 +11656,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, "provider_specific_entry": { @@ -11682,6 +11691,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, @@ -11718,6 +11728,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, @@ -11788,6 +11799,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_native_structured_output": true, "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, @@ -18865,7 +18877,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "supports_reasoning": false }, "gemini/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, @@ -18898,6 +18911,7 @@ ], "supports_function_calling": false, "supports_prompt_caching": true, + "supports_reasoning": false, "supports_response_schema": true, "supports_system_messages": true, "supports_vision": true, @@ -18939,6 +18953,7 @@ ], "supports_function_calling": false, "supports_prompt_caching": true, + "supports_reasoning": false, "supports_response_schema": true, "supports_system_messages": true, "supports_vision": true, @@ -18980,6 +18995,7 @@ ], "supports_function_calling": false, "supports_prompt_caching": true, + "supports_reasoning": false, "supports_response_schema": true, "supports_system_messages": true, "supports_vision": true, @@ -19674,6 +19690,39 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-omni-flash-preview": { + "input_cost_per_audio_token": 1.5e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true, + "tpm": 800000 + }, "gemini/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -19838,6 +19887,37 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-omni-flash-preview": { + "input_cost_per_audio_token": 1.5e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/omni-flash-preview", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.5-flash": { "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, @@ -36951,6 +37031,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -36973,6 +37054,7 @@ "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, + "supports_reasoning": false, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, "vertex_ai/gemini-3-pro-image-preview": { @@ -36988,6 +37070,7 @@ "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, + "supports_reasoning": false, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, "vertex_ai/gemini-3.1-flash-image": { @@ -37001,6 +37084,7 @@ "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, "output_cost_per_token": 3e-06, + "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, "vertex_ai/gemini-3.1-flash-image-preview": { @@ -37014,6 +37098,7 @@ "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, "output_cost_per_token": 3e-06, + "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, "vertex_ai/gemini-3.1-flash-lite-preview": { @@ -45284,8 +45369,8 @@ "rules": [ { "name": "bedrock-claude-ids", - "pattern": "anthropic\\.claude-", - "description": "Any Bedrock-syntax Claude id: the dotted anthropic.claude- segment appears in bare (anthropic.claude-...), region-prefixed (us./eu./au./jp./apac.) and global.-prefixed ids, for every version. Routes these to bedrock before the bare-id Anthropic rule is consulted.", + "pattern": "^(?:[a-z-]+\\.)?anthropic\\.claude-", + "description": "A Bedrock-syntax Claude id, for every version: anthropic.claude- at the start of the name, optionally behind a single dotted geo segment (us./eu./au./jp./apac./global./us-gov.). Anchored to the start because routing rules see the raw request string and provider inference feeds the proxy's provider/* wildcard access checks: an id under an unrecognized namespace such as bedrockz/anthropic.claude-... must stay unroutable rather than resolve to bedrock and slip through a bedrock/* key. Routes to bedrock before the bare-id Anthropic rule is consulted.", "model_info": { "litellm_provider": "bedrock" } diff --git a/pyproject.toml b/pyproject.toml index c2a0e85d13a..f890cc976f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.93.0" +version = "1.94.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -62,8 +62,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.26.0,<2.0", - "litellm-proxy-extras==0.4.76", - "litellm-enterprise==0.1.49", + "litellm-proxy-extras==0.4.77", + "litellm-enterprise==0.1.50", "RestrictedPython>=8.1,<9.0", "rich>=13.9.4,<14.0", "polars>=1.38.1,<2.0", @@ -205,7 +205,7 @@ ci = [ # protobuf, Pillow is a compiled C extension). "tenacity==8.5.0", "google-generativeai==0.8.6", - "Pillow==12.2.0", + "Pillow==12.3.0", # Azure batch E2E tests still import psycopg2 directly. "psycopg2-binary==2.9.11", "pytest-codspeed==4.3.0", @@ -264,6 +264,8 @@ constraint-dependencies = [ "aiohttp>=3.14.1,<4.0", "packaging>=24.0", "soupsieve>=2.8.4", + "httplib2>=0.32.0", + "setuptools>=83.0.0", ] override-dependencies = [ # a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0. @@ -284,7 +286,7 @@ members = ["enterprise", "litellm-proxy-extras"] profile = "black" [tool.commitizen] -version = "1.93.0" +version = "1.94.0" version_files = [ "pyproject.toml:^version", ] diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 7750ac6628a..dcde6fd1641 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -306,7 +306,7 @@ "limit": 9 }, "TID251": { - "limit": 2710 + "limit": 2701 }, "TRY002": { "limit": 548 @@ -324,7 +324,7 @@ "limit": 883 }, "UP006": { - "limit": 12869 + "limit": 12792 }, "UP007": { "limit": 2570 @@ -354,7 +354,7 @@ "limit": 4 }, "UP035": { - "limit": 2295 + "limit": 2284 }, "UP036": { "limit": 4 @@ -363,6 +363,6 @@ "limit": 105 }, "UP045": { - "limit": 18517 + "limit": 18462 } } diff --git a/schema.prisma b/schema.prisma index fb4d8d0b5a3..a23cecc3911 100644 --- a/schema.prisma +++ b/schema.prisma @@ -422,6 +422,7 @@ model LiteLLM_VerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") @@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken { budget_reset_at DateTime? allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) + key_type String? policies String[] @default([]) access_group_ids String[] @default([]) model_spend Json @default("{}") diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index d7d560ce947..cce0cb61c1e 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -89,6 +89,11 @@ EOF status=0 +bootstrap_hint() { + echo " This checkout looks unprovisioned (fresh worktree or clone)." >&2 + echo " Fix: make bootstrap" >&2 +} + if [ -n "$litellm_py_files" ]; then echo "pre-commit: linting Python (make lint)" make lint || { echo "✗ Python lint failed. Fix the reds above, then re-run make pre-commit." >&2; status=1; } @@ -109,7 +114,13 @@ fi if [ -n "$ui_prettier_files" ] || [ -n "$ui_eslint_files" ]; then echo "pre-commit: linting dashboard (prettier + eslint + lint budgets)" - lint_dashboard || { echo "✗ Dashboard lint failed. See above; format with: (cd ui/litellm-dashboard && npm run format)." >&2; status=1; } + if [ ! -d ui/litellm-dashboard/node_modules ]; then + echo "✗ ui/litellm-dashboard/node_modules is missing; dashboard lint cannot run." >&2 + bootstrap_hint + status=1 + else + lint_dashboard || { echo "✗ Dashboard lint failed. See above; format with: (cd ui/litellm-dashboard && npm run format)." >&2; status=1; } + fi fi if [ -n "$spec_files" ]; then @@ -118,7 +129,15 @@ if [ -n "$spec_files" ]; then # and an up-to-date Prisma client; check-ui-api-types.yml installs those and runs # prisma generate before gen:api, so mirror that here or a stale client can mask # drift that CI will still flag. - if ! uv run --no-sync python scripts/prisma_generate_if_needed.py; then + if [ ! -d ui/litellm-dashboard/node_modules ]; then + echo "✗ ui/litellm-dashboard/node_modules is missing; the gen:api sync check cannot run." >&2 + bootstrap_hint + status=1 + elif ! uv run --no-sync python -c "import orjson, prisma" 2>/dev/null; then + echo "✗ The Python env lacks the proxy deps (orjson/prisma) that gen:api needs." >&2 + bootstrap_hint + status=1 + elif ! uv run --no-sync python scripts/prisma_generate_if_needed.py; then echo "✗ Could not regenerate Prisma client (prisma generate failed)." >&2 status=1 elif ( cd ui/litellm-dashboard && LITELLM_PYTHON="uv run --no-sync python" npm run gen:api ); then diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 88992038cb7..6b1b9950656 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -11,7 +11,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `embeddings/` - the `/embeddings` endpoint across providers - `batches/` - the `/batches` endpoint (placeholder until the first test lands) - `realtime/` - realtime websocket sessions, including the pipecat audio path -- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window) and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) +- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials; also the dashboard UI behavior on top of them, driven through the proxy-served UI at /ui with playwright (optional dep behind importorskip) - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 7d54f05656e..85d9315b8c6 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -17,6 +17,7 @@ from __future__ import annotations import json import time +from datetime import datetime, timedelta, timezone from typing import Callable import pytest @@ -49,7 +50,7 @@ from e2e_http import ( unwrap, ) from lifecycle import ResourceManager -from models import KeyGenerateBody, SpendLogRow, SpendLogsParams +from models import KeyGenerateBody, SpendLogRow pytestmark = pytest.mark.e2e @@ -349,6 +350,10 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( the file-read path fires while the batch itself is not blocked. ``resources.key()`` cannot set limits, so the key is minted on the gateway directly and its delete deferred. + + Snapshots read /spend/logs/v2 over a bounded window around the test instead + of the unpaginated /spend/logs whole-table read, which grows with the + environment and OOMed the e2e runner on stage. """ user_id = f"e2e-batch-rl-{unique_marker()}" key = client.gateway.generate_key( @@ -356,8 +361,13 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( ) resources.defer(lambda: client.gateway.delete_key(key)) + window_start = datetime.now(timezone.utc) - timedelta(hours=1) + window_end = window_start + timedelta(hours=2) before = frozenset( - row.request_id for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams())) + row.request_id + for row in unattributed_rows( + client.gateway.spend_logs_window(start=window_start, end=window_end) + ) ) file = unwrap( @@ -379,7 +389,9 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( new_orphans = [ row - for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams())) + for row in unattributed_rows( + client.gateway.spend_logs_window(start=window_start, end=window_end) + ) if row.request_id not in before ] assert not new_orphans, ( diff --git a/tests/e2e/coverage_registry/logging.yaml b/tests/e2e/coverage_registry/logging.yaml index 65ab8f0096f..c4b26387c13 100644 --- a/tests/e2e/coverage_registry/logging.yaml +++ b/tests/e2e/coverage_registry/logging.yaml @@ -8,7 +8,7 @@ - {id: logging.datadog.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/datadog/datadog.py", rationale: "Powers dashboards/alerts; cardinality regressions common"} - {id: logging.datadog.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions], source: "integrations/datadog/datadog.py", rationale: "Failure metrics for alerting/SLO"} - {id: logging.prometheus.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/prometheus.py", rationale: "Standard OSS metrics; per-key cardinality (existing e2e)"} -- {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"} +- {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"} - {id: logging.otel.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions, messages], source: "integrations/otel/logger.py", rationale: "Error spans for observability continuity"} - {id: logging.braintrust.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/braintrust_logging.py", rationale: "Evals platform spend"} - {id: logging.langsmith.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/langsmith.py", rationale: "LangChain ecosystem"} diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml index 0cfbb5b0b66..12f47331dbb 100644 --- a/tests/e2e/docker-compose.yml +++ b/tests/e2e/docker-compose.yml @@ -18,6 +18,12 @@ configs: type: redis host: redis port: 6379 + # OTEL v2 trace destination for the logging suite's trace-completeness + # tests: the arize_phoenix preset is OTLP with a configurable endpoint + # (PHOENIX_COLLECTOR_HTTP_ENDPOINT below points it at the jaeger service), + # so gen-AI spans export through a preset-owned provider - the code path + # where trace splits actually happen - with no cloud credentials needed. + callbacks: ["arize_phoenix"] router_settings: routing_strategy: simple-shuffle @@ -68,9 +74,14 @@ services: condition: service_healthy redis: condition: service_healthy + jaeger: + condition: service_healthy env_file: .env environment: LITELLM_MASTER_KEY: sk-1234 + LITELLM_OTEL_V2: "true" + PHOENIX_COLLECTOR_HTTP_ENDPOINT: http://jaeger:4318/v1/traces + PHOENIX_API_KEY: local-jaeger-noauth DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm UI_USERNAME: admin UI_PASSWORD: sk-1234 @@ -114,3 +125,15 @@ services: interval: 3s timeout: 3s retries: 20 + +# throwaway OTEL trace destination (OTLP ingest on 4318 inside the network, +# query API on host 16686 for test read-back; see E2E_OTEL_QUERY_URL) + jaeger: + image: jaegertracing/all-in-one:1.62.0 + ports: + - "16686:16686" + healthcheck: + test: ["CMD", "wget", "-qO-", "http://localhost:14269/"] + interval: 3s + timeout: 3s + retries: 20 diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 75bd715a23a..6e6c30709de 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -24,6 +24,14 @@ CONTROL_PLANE_BASE_URL = os.environ.get( UI_USERNAME = os.environ.get("E2E_UI_USERNAME", "admin") UI_PASSWORD = os.environ.get("E2E_UI_PASSWORD", MASTER_KEY) +CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haiku-4-5") +CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5") + +# Jaeger query API of the compose stack's OTEL trace destination (the `jaeger` +# service in docker-compose.yml maps it to host 16686). Trace-completeness tests +# read exported spans back through it. +OTEL_QUERY_URL = os.environ.get("E2E_OTEL_QUERY_URL", "http://localhost:16686").rstrip("/") + # Writes on the proxy are eventually consistent (e.g. spend rows flush on # proxy_batch_write_at, ~60s). Read-backs poll to this deadline, never sleep-once. POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120")) diff --git a/tests/e2e/e2e_gateway.py b/tests/e2e/e2e_gateway.py index 05f83ecc085..d40b96d60fa 100644 --- a/tests/e2e/e2e_gateway.py +++ b/tests/e2e/e2e_gateway.py @@ -12,6 +12,7 @@ import time import warnings from collections.abc import Callable from dataclasses import dataclass +from datetime import datetime from e2e_http import ( NoBody, @@ -50,6 +51,8 @@ from models import ( OcrResponse, SpendLogRow, SpendLogs, + SpendLogsPage, + SpendLogsPageParams, SpendLogsParams, ) from e2e_config import ( @@ -255,6 +258,28 @@ class Gateway: case _: return [] + def spend_logs_window(self, *, start: datetime, end: datetime) -> list[SpendLogRow]: + def fetch(page: int) -> SpendLogsPage: + return unwrap( + self.transport.get( + "/spend/logs/v2", + headers=self.transport.master, + params=SpendLogsPageParams( + start_date=start.strftime("%Y-%m-%d %H:%M:%S"), + end_date=end.strftime("%Y-%m-%d %H:%M:%S"), + page=page, + page_size=100, + ), + response_type=SpendLogsPage, + ) + ) + + first = fetch(1) + return [ + *first.data, + *(row for page in range(2, first.total_pages + 1) for row in fetch(page).data), + ] + def poll_logs_for_key( self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None ) -> list[SpendLogRow]: diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 32005faed4b..ad53b2b4aa2 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -109,14 +109,16 @@ class StreamingResponse(BaseModel): """Raw outcome for calls whose body is provider-native or streamed: status, the x-litellm-call-id header, the x-litellm-response-cost header (StandardLogging response_cost), the content-type (which tells streaming `text/event-stream` from - non-streaming `application/json`), and the body. SpendLogs.request_id is the - completion body id, not call_id. Used by passthrough and streaming, where one - validated JSON model does not fit.""" + non-streaming `application/json`), the response headers (lowercased names, e.g. + the x-ratelimit-* pacing headers and retry-after on a 429), and the body. + SpendLogs.request_id is the completion body id, not call_id. Used by passthrough + and streaming, where one validated JSON model does not fit.""" status_code: int call_id: str | None = None # x-litellm-call-id header response_cost: float | None = None # x-litellm-response-cost header content_type: str | None = None + headers: dict[str, str] = {} body: str chunks: int = 0 # streamed events (0 for non-streaming) @@ -276,12 +278,14 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon call_id = _hdr(resp, "x-litellm-call-id") response_cost = _parse_response_cost(resp) content_type = _hdr(resp, "content-type") + headers = {name.lower(): value for name, value in resp.headers.items()} if not stream or not (200 <= resp.status_code < 300): return StreamingResponse( status_code=resp.status_code, call_id=call_id, response_cost=response_cost, content_type=content_type, + headers=headers, body=resp.text, ) lines = cast("Iterator[bytes]", resp.iter_lines()) @@ -291,6 +295,7 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon call_id=call_id, response_cost=response_cost, content_type=content_type, + headers=headers, body="", chunks=chunks, ) diff --git a/tests/e2e/logging/conftest.py b/tests/e2e/logging/conftest.py index 567355f23ac..43d279602ef 100644 --- a/tests/e2e/logging/conftest.py +++ b/tests/e2e/logging/conftest.py @@ -11,6 +11,7 @@ import os import pytest from logging_client import LangfuseCreds, LoggingClient, build_logging_client, load_langfuse_creds +from otel_client import OtelReader, build_otel_reader def pytest_configure(config: pytest.Config) -> None: @@ -28,6 +29,12 @@ def client() -> LoggingClient: return build_logging_client() +@pytest.fixture(scope="session") +def otel_reader() -> OtelReader: + """Read-back client for the compose stack's Jaeger trace destination.""" + return build_otel_reader() + + @pytest.fixture def datadog_creds() -> None: """Require Datadog shipping credentials. Hard-fail when absent; never skip.""" diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index 6071e657bd3..bffdf71ed80 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -35,6 +35,7 @@ from e2e_http import ( unwrap, ) from models import ( + AnthropicMessagesBody, ChatBody, ChatMessage, ChatResponse, @@ -75,6 +76,14 @@ WEATHER_TOOL = ChatTool( ) +class ResponsesRequestBody(BaseModel): + """OpenAI Responses API /v1/responses request (non-streaming).""" + + model: str + input: str + max_output_tokens: int + + class TeamCallbackBody(BaseModel): callback_name: Literal["langfuse_otel", "langfuse", "langsmith", "gcs"] callback_type: Literal["success", "failure", "success_and_failure"] @@ -455,6 +464,32 @@ class LoggingClient: json=body, ) + def messages_raw(self, key: str, model: str, text: str, *, max_tokens: int = 16) -> StreamingResponse: + """Non-streaming POST /v1/messages (Anthropic-native body): raw outcome + judged by status/body/headers, for tests that need x-litellm-call-id.""" + return self.gateway.transport.send( + "/v1/messages", + headers=self.gateway.transport.bearer(key), + json=AnthropicMessagesBody( + model=model, + max_tokens=max_tokens, + messages=[ChatMessage(role="user", content=text)], + ), + ) + + def responses_raw( + self, key: str, model: str, text: str, *, max_output_tokens: int = 64 + ) -> StreamingResponse: + """Non-streaming POST /v1/responses (OpenAI Responses API): raw outcome + judged by status/body/headers, for tests that need x-litellm-call-id. + max_output_tokens caps reasoning-model output cost; a capped response is + still a 200 and still exports the trace.""" + return self.gateway.transport.send( + "/v1/responses", + headers=self.gateway.transport.bearer(key), + json=ResponsesRequestBody(model=model, input=text, max_output_tokens=max_output_tokens), + ) + def scrape_metrics(self) -> str: return self.gateway.probe("/metrics", params=NoBody()).body diff --git a/tests/e2e/logging/otel_client.py b/tests/e2e/logging/otel_client.py new file mode 100644 index 00000000000..f4a0e4fe102 --- /dev/null +++ b/tests/e2e/logging/otel_client.py @@ -0,0 +1,138 @@ +"""Jaeger read-back for the OTEL trace-completeness tests: typed models over the +Jaeger query API (the destination's own API - completeness is judged on what the +backend actually holds, never on "export succeeded" proxy-side). + +Traces are fetched server-side by the ``litellm.call_id`` tag the gen-AI span +carries (the request's x-litellm-call-id response header), so read-back is +immune to the query page filling up with unrelated traffic (background jobs, +other suites sharing the stack). Jaeger returns every span of a matching trace, +so the completeness assertions see the whole tree. A failed query is a hard +failure, never an empty result - an unreachable destination must not read as +"the trace never arrived". + +External reads go through ``e2e_http`` (the only module allowed to call +``requests.*``). +""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass + +import pytest +from pydantic import BaseModel, ConfigDict, Field + +from e2e_config import OTEL_QUERY_URL, POLL_INTERVAL, POLL_TIMEOUT +from e2e_http import URL, NoBody, Success, get + +#: OTEL resource service.name the proxy exports under (OTEL_SERVICE_NAME default). +JAEGER_SERVICE = "litellm" +#: Span tag carrying the request's x-litellm-call-id (stamped on the gen-AI span). +CALL_ID_TAG = "litellm.call_id" + + +class JaegerTag(BaseModel): + model_config = ConfigDict(extra="ignore") + + key: str + value: str | int | float | bool | None = None + + +class JaegerReference(BaseModel): + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + ref_type: str = Field(alias="refType") + trace_id: str = Field(alias="traceID") + span_id: str = Field(alias="spanID") + + +class JaegerSpan(BaseModel): + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + span_id: str = Field(alias="spanID") + operation_name: str = Field(alias="operationName") + start_time: int = Field(default=0, alias="startTime") + references: list[JaegerReference] = [] + tags: list[JaegerTag] = [] + + @property + def kind(self) -> str: + for tag in self.tags: + if tag.key == "span.kind": + return str(tag.value) + return "" + + +class JaegerTrace(BaseModel): + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + trace_id: str = Field(alias="traceID") + spans: list[JaegerSpan] = [] + + def span_names(self) -> list[str]: + return sorted(span.operation_name for span in self.spans) + + +class JaegerTracesPage(BaseModel): + model_config = ConfigDict(extra="ignore") + + data: list[JaegerTrace] = [] + + +class _TracesQuery(BaseModel): + service: str + tags: str + limit: int = 20 + lookback: str = "1h" + + +def _settled(trace: JaegerTrace, names: set[str], prefixes: set[str]) -> bool: + present = set(trace.span_names()) + return names.issubset(present) and all( + any(name.startswith(prefix) for name in present) for prefix in prefixes + ) + + +@dataclass(frozen=True, slots=True) +class OtelReader: + query_url: str + + def traces_for_call(self, call_id: str) -> list[JaegerTrace]: + """Every trace holding a span tagged with this call id. Jaeger matches + spans server-side and returns their full traces; more than one hit for + one call IS the split-trace bug, so this never collapses to one.""" + result = get( + URL(f"{self.query_url}/api/traces"), + headers=NoBody(), + params=_TracesQuery(service=JAEGER_SERVICE, tags=json.dumps({CALL_ID_TAG: call_id})), + response_type=JaegerTracesPage, + timeout=30.0, + ) + match result: + case Success(data=page): + return page.data + case failure: + pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}") + + def poll_traces_for_call( + self, *, call_id: str, settled_names: set[str], settled_prefixes: set[str] + ) -> list[JaegerTrace]: + """Poll until exactly one trace holds the call and it carries every span + name in ``settled_names`` plus at least one name per prefix in + ``settled_prefixes`` (spans flush in batches, the cost write lands after + the response), then return the hits. At the deadline the last hits are + returned as-is so the caller's assertions report the real final state - + on a split trace this never settles and the orphan comes back.""" + deadline = time.monotonic() + POLL_TIMEOUT + hits: list[JaegerTrace] = [] + while time.monotonic() < deadline: + hits = self.traces_for_call(call_id) + if len(hits) == 1 and _settled(hits[0], settled_names, settled_prefixes): + return hits + time.sleep(POLL_INTERVAL) + return hits + + +def build_otel_reader() -> OtelReader: + return OtelReader(query_url=OTEL_QUERY_URL) diff --git a/tests/e2e/logging/test_otel_trace_e2e.py b/tests/e2e/logging/test_otel_trace_e2e.py new file mode 100644 index 00000000000..90887ea5510 --- /dev/null +++ b/tests/e2e/logging/test_otel_trace_e2e.py @@ -0,0 +1,259 @@ +"""Live e2e: OTEL trace completeness on the admin-owned destination (LIT-3787). + +Covers logging.otel.success.exports_metric: a successful non-streaming call must +land at the OTEL destination as ONE connected trace - a single root SERVER span +with the auth phase, db lookups, and cost write under it, and the gen-AI CLIENT +span parented into the same tree. The regression this pins: the proxy publishing +the global TracerProvider before callbacks init made server spans export through +a different provider than the preset's gen-AI spans, so the destination received +the gen-AI span alone, dangling (fixed in #30590; verified failing at its parent +commit 1bd603d1ac). + +Both halves of the contract are asserted: the recorded state (the proxy reports +the OTEL v2 logger active via /health/readiness/details) and the enforced +behavior (the complete span tree at the destination, read back through the +destination's own query API - never proxy-side "export succeeded" logs). +""" + +from __future__ import annotations + +import time +from collections.abc import Callable + +import pytest +from pydantic import BaseModel, ConfigDict + +from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, unique_marker +from e2e_http import NoBody, StreamingResponse, require_successful_call +from lifecycle import ResourceManager +from logging_client import LoggingClient +from otel_client import JaegerTrace, OtelReader + +pytestmark = pytest.mark.e2e + +MODEL = CHEAP_ANTHROPIC_MODEL +COST_SPAN = "batch_write_to_db _PROXY_track_cost_callback" +DB_SPAN_PREFIX = "postgres " +#: The active OTEL v2 logger's name in /health/readiness/details success_callbacks. +OTEL_V2_LOGGER_NAME = "OpenTelemetryV2" + + +class _ReadinessDetails(BaseModel): + model_config = ConfigDict(extra="ignore") + + success_callbacks: list[str] = [] + + +def _assert_otel_destination_configured(client: LoggingClient) -> None: + """Recorded state: the proxy reports the OTEL v2 logger among its active + callbacks, so a missing/failed destination config fails here, before any + traffic-based assertion can time out confusingly.""" + result = client.gateway.probe("/health/readiness/details", params=NoBody()) + assert result.status_code == 200, ( + f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}" + ) + details = _ReadinessDetails.model_validate_json(result.body) + assert OTEL_V2_LOGGER_NAME in details.success_callbacks, ( + f"the proxy must report the {OTEL_V2_LOGGER_NAME} callback active " + f"(LITELLM_OTEL_V2 + arize_phoenix preset in the compose config); got: {details.success_callbacks}" + ) + + +def _first_ok(client: LoggingClient, send: Callable[[], StreamingResponse]) -> StreamingResponse: + """First successful call on a fresh key. A fresh key may briefly 401 until + the data plane's auth cache picks it up, so retry on 401 to a deadline; a + 401 is rejected before the LLM call so it exports no gen-AI span and cannot + contaminate the trace assertions. Any other failure is behavior under test + and fails hard.""" + deadline = time.monotonic() + client.gateway.poll_timeout + while True: + outcome = send() + if outcome.ok: + return outcome + if outcome.status_code != 401 or time.monotonic() >= deadline: + require_successful_call(outcome) + time.sleep(client.gateway.poll_interval) + + +def _parent_ids(span_id: str, trace: JaegerTrace) -> list[str]: + span = next(s for s in trace.spans if s.span_id == span_id) + return [ref.span_id for ref in span.references if ref.ref_type == "CHILD_OF"] + + +def _chain_reaches(span_id: str, root_id: str, trace: JaegerTrace) -> bool: + """Walk parent references (within the trace) from span_id up to root_id.""" + seen: set[str] = set() + in_trace = {s.span_id for s in trace.spans} + current = span_id + while current not in seen: + if current == root_id: + return True + seen.add(current) + parents = [p for p in _parent_ids(current, trace) if p in in_trace] + if not parents: + return False + current = parents[0] + return False + + +def _assert_complete_trace(hits: list[JaegerTrace], *, route: str, genai_span: str) -> None: + """The enforced behavior: the destination holds exactly one trace for the + call, rooted at the SERVER span, with auth/db/cost children and the gen-AI + span all connected into that one tree - no dangling parent references.""" + assert hits, ( + "no trace for this call arrived at the destination within the deadline " + "(nothing tagged with its call id was found)" + ) + assert len(hits) == 1, ( + f"expected exactly ONE trace for the call, got {len(hits)}: " + f"{[(t.trace_id, t.span_names()) for t in hits]} - more than one trace for " + "one call is the split-trace bug (gen-AI span exported away from its root)" + ) + trace = hits[0] + names = trace.span_names() + in_trace = {span.span_id for span in trace.spans} + + dangling = [ + span.operation_name + for span in trace.spans + if span.references and not any(ref.span_id in in_trace for ref in span.references) + ] + assert not dangling, ( + f"span(s) {dangling} reference a parent that never reached the destination " + f"(orphaned trace); spans present: {names}" + ) + + roots = [span for span in trace.spans if not span.references] + assert len(roots) == 1, f"expected exactly one root span, got {[s.operation_name for s in roots]}; spans: {names}" + root = roots[0] + assert root.operation_name == f"POST {route}", ( + f"the root must be the SERVER span 'POST {route}', got {root.operation_name!r}" + ) + assert root.kind == "server", f"the root span must have kind=server, got {root.kind!r}" + + assert f"auth {route}" in names, f"auth phase span 'auth {route}' missing; spans: {names}" + assert any(name.startswith(DB_SPAN_PREFIX) for name in names), ( + f"no db ('{DB_SPAN_PREFIX}*') span in the trace; spans: {names}" + ) + assert COST_SPAN in names, f"cost write span {COST_SPAN!r} missing; spans: {names}" + + genai = next((span for span in trace.spans if span.operation_name == genai_span), None) + assert genai is not None, f"gen-AI span {genai_span!r} missing; spans: {names}" + assert genai.kind == "client", f"gen-AI span must have kind=client, got {genai.kind!r}" + assert _chain_reaches(genai.span_id, root.span_id, trace), ( + f"gen-AI span {genai_span!r} is in the trace but its parent chain does not " + f"reach the root SERVER span; spans: {names}" + ) + + +def _settled_names(*, route: str, genai_span: str) -> set[str]: + return {f"POST {route}", f"auth {route}", COST_SPAN, genai_span} + + +class TestOtelTraceCompleteness: + @pytest.mark.covers("logging.otel.success.exports_metric", exercised_on=["chat_completions"]) + def test_chat_completions_exports_complete_trace( + self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager + ) -> None: + """This test verifies that a successful non-streaming + /chat/completions request produces one complete OTEL trace. + + The trace should have a single server root span for the incoming request, with + the authentication, database, and cost-recording work beneath it. The span for + the actual model call must also belong to that same trace, rather than being + exported separately with a missing parent. + + This matters because a split trace is easy to miss: all of the spans may still + arrive, but the model call appears without the surrounding request context. + That makes it difficult to understand where time was spent, connect the model + cost to the original request, or investigate a slow or failed call. + + /chat/completions is the main OpenAI-compatible route used by most customers, + so it is important that trace parenting works correctly on this path. + """ + route = "/chat/completions" + _assert_otel_destination_configured(client) + + key = client.key_with_alias(f"otel-trace-chat-{unique_marker()}", models=[MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = _first_ok( + client, lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16) + ) + assert outcome.call_id is not None, "success response must carry x-litellm-call-id" + + hits = otel_reader.poll_traces_for_call( + call_id=outcome.call_id, + settled_names=_settled_names(route=route, genai_span=f"chat {MODEL}"), + settled_prefixes={DB_SPAN_PREFIX}, + ) + _assert_complete_trace(hits, route=route, genai_span=f"chat {MODEL}") + + @pytest.mark.covers("logging.otel.success.exports_metric", exercised_on=["messages"]) + def test_messages_exports_complete_trace( + self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager + ) -> None: + """This test verifies that one successful non-streaming /v1/messages request + produces exactly one complete OTEL trace. + + The trace must have a single root span named "POST /v1/messages". The + authentication, database, cost-writing, and model-call spans must all belong to + the same trace and have valid parent relationships leading back to that root. + + The model-call span is expected to be named "chat ". The test fails if + the request is split across multiple traces, if any span references a missing + parent, or if the model-call span cannot be connected back to the root.""" + route = "/v1/messages" + _assert_otel_destination_configured(client) + + key = client.key_with_alias(f"otel-trace-messages-{unique_marker()}", models=[MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = _first_ok( + client, lambda: client.messages_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16) + ) + assert outcome.call_id is not None, "success response must carry x-litellm-call-id" + + hits = otel_reader.poll_traces_for_call( + call_id=outcome.call_id, + settled_names=_settled_names(route=route, genai_span=f"chat {MODEL}"), + settled_prefixes={DB_SPAN_PREFIX}, + ) + _assert_complete_trace(hits, route=route, genai_span=f"chat {MODEL}") + + @pytest.mark.covers("logging.otel.success.exports_metric", exercised_on=["responses"]) + def test_responses_exports_complete_trace( + self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager + ) -> None: + """This test verifies that one successful non-streaming /v1/responses request + produces exactly one complete OTEL trace. + + The trace must have a single root span named "POST /v1/responses". The + authentication, database, cost-writing, and model-call spans must all belong to + the same trace and have valid parent relationships leading back to that root. + + The model-call span is expected to be named "chat ". The test fails if + the request is split across multiple traces, if any span references a missing + parent, or if the model-call span cannot be connected back to the root.""" + route = "/v1/responses" + _assert_otel_destination_configured(client) + + key = client.key_with_alias(f"otel-trace-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = _first_ok( + client, + lambda: client.responses_raw(key, CHEAP_OPENAI_MODEL, f"reply with one word {marker}"), + ) + assert outcome.call_id is not None, "success response must carry x-litellm-call-id" + + genai_span = f"chat {CHEAP_OPENAI_MODEL}" + hits = otel_reader.poll_traces_for_call( + call_id=outcome.call_id, + settled_names=_settled_names(route=route, genai_span=genai_span), + settled_prefixes={DB_SPAN_PREFIX}, + ) + _assert_complete_trace(hits, route=route, genai_span=genai_span) diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index cbd5db0d59f..7a05a5c520f 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -111,7 +111,7 @@ class TestKeyRoutes: key = _generate_key( client, resources, - KeyGenerateBody(models=["gemini-2.5-flash"], key_alias=alias, tpm_limit=424242), + KeyGenerateBody(models=["gemini-2.5-flash"], key_alias=alias, tpm_limit=424242, rpm_limit=424243), ) info = client.gateway.key_info(key) @@ -122,6 +122,9 @@ class TestKeyRoutes: assert info.tpm_limit == 424242, ( f"/key/info reports tpm_limit {info.tpm_limit}, configured 424242" ) + assert info.rpm_limit == 424243, ( + f"/key/info reports rpm_limit {info.rpm_limit}, configured 424243" + ) _poll_chat_ok(client, key, "gemini-2.5-flash") _assert_model_denied( diff --git a/tests/e2e/models.py b/tests/e2e/models.py index e32f2709181..bf90426188f 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -8,7 +8,7 @@ from __future__ import annotations from typing import Literal -from pydantic import BaseModel, ConfigDict, RootModel +from pydantic import BaseModel, ConfigDict, RootModel, model_validator # ---------- keys ---------- @@ -82,6 +82,7 @@ class KeyInfo(BaseModel): key_alias: str | None = None models: list[str] = [] tpm_limit: int | None = None + rpm_limit: int | None = None team_id: str | None = None spend: float | None = None max_budget: float | None = None @@ -255,6 +256,16 @@ class SpendLogsParams(BaseModel): request_id: str | None = None api_key: str | None = None + @model_validator(mode="after") + def require_filter(self) -> SpendLogsParams: + if self.request_id is None and self.api_key is None: + raise ValueError( + "unfiltered /spend/logs returns the entire spend table and OOMs the " + "runner on long-lived environments; filter by request_id or api_key, " + "or use Gateway.spend_logs_window for a bounded /spend/logs/v2 read" + ) + return self + class SpendLogsPageParams(BaseModel): """Query for /spend/logs/v2, which requires an explicit date window and diff --git a/tests/e2e/quota_management/ratelimit/conftest.py b/tests/e2e/quota_management/ratelimit/conftest.py new file mode 100644 index 00000000000..4a5a73bb5e4 --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/conftest.py @@ -0,0 +1,15 @@ +"""Quota-management suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +live in the parent tests/e2e/conftest.py. QuotaClient holds the shared Gateway, +so the `resources` fixture cleans up keys through it. +""" + +import pytest + +from quota_client import QuotaClient, build_client + + +@pytest.fixture(scope="session") +def client() -> QuotaClient: + return build_client() diff --git a/tests/e2e/quota_management/ratelimit/quota_client.py b/tests/e2e/quota_management/ratelimit/quota_client.py new file mode 100644 index 00000000000..806ab1d1557 --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/quota_client.py @@ -0,0 +1,32 @@ +"""Client for the quota-management suite: the shared Gateway plus raw chat +calls judged by HTTP status, body, and headers (a rate-limit block is a 429 +whose body and retry-after header carry the contract, not a typed success +model).""" + +from __future__ import annotations + +from dataclasses import dataclass + +from e2e_gateway import Gateway, build_gateway +from e2e_http import StreamingResponse +from models import ChatBody, ChatMessage + + +@dataclass(frozen=True, slots=True) +class QuotaClient: + gateway: Gateway + + def chat(self, key: str, model: str, content: str, *, max_tokens: int = 16) -> StreamingResponse: + return self.gateway.transport.send( + "/chat/completions", + headers=self.gateway.transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + ), + ) + + +def build_client() -> QuotaClient: + return QuotaClient(gateway=build_gateway()) diff --git a/tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py b/tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py new file mode 100644 index 00000000000..a6c15b79bb1 --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py @@ -0,0 +1,254 @@ +"""Live e2e: key-level rpm/tpm rate limits on the gateway. + +Covers quota_management.ratelimit.*: a key generated with rpm_limit/tpm_limit gets a +429 once the limit is crossed inside one window (blocks_over_limit), serves +again once the window rolls and no sooner (resets_after_window), and successful responses +report x-ratelimit-* limit/remaining headers so clients can pace +(headers_report_remaining). Each test asserts both halves of the contract: the +recorded state (/key/info echoes the configured limit) and the enforced +behavior (the 429, the recovery, or the headers on live traffic). + +The v3 limiter counts a request against the rpm budget at the pre-call hook, +before model routing, so every call that clears auth consumes budget whether or +not it ultimately succeeds. The tpm budget is reserved pre-call from an estimate +(message chars // 4 + max_tokens) and reconciled to the body's actual +usage.total_tokens after the call, so a block may legitimately fire before the +actual spend crosses the limit; the tpm test asserts the exact contract on both +sides (a 429 only once the blocked call's reservation exceeds the remaining +budget, and no later than the first call after actual spend reaches the limit). +All calls of one test must land inside a single window +(LITELLM_RATE_LIMIT_WINDOW_SIZE, 60s default), which real chat latency +comfortably allows. + +The window opens at the pre-call hook of the first counted call, which happens +after that call is sent, so the send timestamp of the winning first call is a +lower bound on the window start. The reset test uses it to reject an early +reset: recovery must not arrive before the full window has elapsed from that +send, less a small tolerance for the limiter's integer-second window +arithmetic. +""" + +from __future__ import annotations + +import re +import time +from dataclasses import dataclass + +import pytest +from pydantic import BaseModel, ConfigDict, ValidationError + +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from e2e_http import StreamingResponse, require_successful_call +from lifecycle import ResourceManager +from models import KeyGenerateBody +from quota_client import QuotaClient + +pytestmark = pytest.mark.e2e + +MODEL = CHEAP_ANTHROPIC_MODEL +TPM_LIMIT = 60 +CHAT_MAX_TOKENS = 16 +RESERVATION_CHARS_PER_TOKEN = 4 +WINDOW_SECONDS = 60 +RESET_TOLERANCE_SECONDS = 5 +LAST_CALL_LATENCY_MARGIN_SECONDS = 10 + + +@dataclass(frozen=True, slots=True) +class _FirstOk: + sent_at: float + response: StreamingResponse + + +class _ChatUsage(BaseModel): + model_config = ConfigDict(extra="ignore") + + total_tokens: int + + +class _ChatBodyWithUsage(BaseModel): + model_config = ConfigDict(extra="ignore") + + usage: _ChatUsage + + +def _total_tokens(outcome: StreamingResponse) -> int: + try: + return _ChatBodyWithUsage.model_validate_json(outcome.body).usage.total_tokens + except ValidationError: + pytest.fail(f"successful chat body must report usage.total_tokens, got: {outcome.body[:300]}") + + +def _reserved_tokens(content: str) -> int: + return max(1, len(content) // RESERVATION_CHARS_PER_TOKEN) + CHAT_MAX_TOKENS + + +def _remaining_from_429(body: str) -> int: + found = re.search(r"Remaining: (\d+)", body) + if found is None: + pytest.fail(f"429 body must report the remaining budget, got: {body[:300]}") + return int(found.group(1)) + + +@dataclass(frozen=True, slots=True) +class _BlockedByReservation: + outcome: StreamingResponse + content: str + spent: int + + +@dataclass(frozen=True, slots=True) +class _CrossedLimit: + spent: int + + +def _spend_until_blocked_or_crossed( + client: QuotaClient, key: str, first: _FirstOk +) -> _BlockedByReservation | _CrossedLimit: + """Drive chat traffic, summing each body's actual usage.total_tokens, until + the limiter blocks (which the reservation may do before the actual spend + crosses the limit) or the actual spend reaches the limit.""" + window_deadline = first.sent_at + WINDOW_SECONDS - LAST_CALL_LATENCY_MARGIN_SECONDS + spent = _total_tokens(first.response) + while spent < TPM_LIMIT: + assert time.monotonic() < window_deadline, ( + f"spent only {spent} of {TPM_LIMIT} tokens before the {WINDOW_SECONDS}s window could roll; " + "the exact-crossing assertion needs every call inside one window" + ) + content = f"reply with one word {unique_marker()}" + outcome = client.chat(key, MODEL, content, max_tokens=CHAT_MAX_TOKENS) + if outcome.status_code == 429: + return _BlockedByReservation(outcome=outcome, content=content, spent=spent) + require_successful_call(outcome) + spent += _total_tokens(outcome) + return _CrossedLimit(spent=spent) + + +def _limited_key( + client: QuotaClient, + resources: ResourceManager, + *, + rpm_limit: int | None = None, + tpm_limit: int | None = None, +) -> str: + key = client.gateway.generate_key(KeyGenerateBody(models=[MODEL], rpm_limit=rpm_limit, tpm_limit=tpm_limit)) + resources.defer(lambda: client.gateway.delete_key(key)) + return key + + +def _chat(client: QuotaClient, key: str) -> StreamingResponse: + return client.chat(key, MODEL, f"reply with one word {unique_marker()}") + + +def _first_ok(client: QuotaClient, key: str) -> _FirstOk: + """First successful call on a fresh key, which opens the rate-limit window; + `sent_at` is captured just before the winning send, so the window opened no + earlier than it. A fresh key may briefly 401 until the data plane's auth + cache picks it up, so retry on 401 to a deadline; a 401 never reaches the + rate limiter, so only the successful call consumes budget. Any other failure + is behavior under test and fails hard.""" + deadline = time.monotonic() + client.gateway.poll_timeout + while True: + sent_at = time.monotonic() + outcome = _chat(client, key) + if outcome.ok: + return _FirstOk(sent_at=sent_at, response=outcome) + if outcome.status_code != 401 or time.monotonic() >= deadline: + require_successful_call(outcome) + time.sleep(client.gateway.poll_interval) + + +def _assert_rate_limited(outcome: StreamingResponse, limit_type: str) -> None: + assert outcome.status_code == 429, ( + f"expected a 429 {limit_type} rate-limit block, got {outcome.status_code}: {outcome.body[:300]}" + ) + assert "Rate limit exceeded for api_key" in outcome.body, ( + f"429 body must name the api_key scope, got: {outcome.body[:300]}" + ) + assert f"Limit type: {limit_type}" in outcome.body, ( + f"429 body must carry 'Limit type: {limit_type}', got: {outcome.body[:300]}" + ) + retry_after = outcome.headers.get("retry-after") + assert retry_after is not None and retry_after.isdigit() and int(retry_after) > 0, ( + f"429 must carry a positive integer retry-after header, got {retry_after!r}" + ) + + +class TestKeyRateLimits: + @pytest.mark.covers("quota_management.ratelimit.rpm.blocks_over_limit") + def test_rpm_limit_blocks_over_limit(self, client: QuotaClient, resources: ResourceManager) -> None: + key = _limited_key(client, resources, rpm_limit=3) + info = client.gateway.key_info(key) + assert info.rpm_limit == 3, f"/key/info reports rpm_limit {info.rpm_limit}, configured 3" + + _ = _first_ok(client, key) + for _ in range(2): + require_successful_call(_chat(client, key)) + + _assert_rate_limited(_chat(client, key), "requests") + + @pytest.mark.covers("quota_management.ratelimit.tpm.blocks_over_limit") + def test_tpm_limit_blocks_over_limit(self, client: QuotaClient, resources: ResourceManager) -> None: + key = _limited_key(client, resources, tpm_limit=TPM_LIMIT) + info = client.gateway.key_info(key) + assert info.tpm_limit == TPM_LIMIT, f"/key/info reports tpm_limit {info.tpm_limit}, configured {TPM_LIMIT}" + + first = _first_ok(client, key) + match _spend_until_blocked_or_crossed(client, key, first): + case _BlockedByReservation(outcome=outcome, content=content, spent=spent): + _assert_rate_limited(outcome, "tokens") + remaining = _remaining_from_429(outcome.body) + reserved = _reserved_tokens(content) + assert reserved > remaining, ( + f"blocked while the call still fit: {remaining} of {TPM_LIMIT} tokens remained but the call " + f"reserved only {reserved} ({spent} actual tokens spent so far)" + ) + case _CrossedLimit(): + _assert_rate_limited(_chat(client, key), "tokens") + + @pytest.mark.covers("quota_management.ratelimit.rpm.resets_after_window") + def test_rpm_limit_resets_after_window(self, client: QuotaClient, resources: ResourceManager) -> None: + key = _limited_key(client, resources, rpm_limit=1) + + first = _first_ok(client, key) + _assert_rate_limited(_chat(client, key), "requests") + + deadline = time.monotonic() + client.gateway.poll_timeout + while time.monotonic() < deadline: + attempt_sent_at = time.monotonic() + outcome = _chat(client, key) + if outcome.ok: + window_age = attempt_sent_at - first.sent_at + assert window_age >= WINDOW_SECONDS - RESET_TOLERANCE_SECONDS, ( + f"the key recovered {window_age:.1f}s after the window opened, before the " + f"{WINDOW_SECONDS}s window (less {RESET_TOLERANCE_SECONDS}s tolerance) elapsed; " + "the limiter reset early instead of after the window" + ) + return + assert outcome.status_code == 429, ( + f"while the window drains only 429s are acceptable, got {outcome.status_code}: {outcome.body[:300]}" + ) + time.sleep(client.gateway.poll_interval) + pytest.fail("a blocked key never recovered after the rate-limit window elapsed") + + @pytest.mark.covers("quota_management.ratelimit.rpm.headers_report_remaining") + def test_headers_report_limit_and_remaining(self, client: QuotaClient, resources: ResourceManager) -> None: + key = _limited_key(client, resources, rpm_limit=5, tpm_limit=100000) + + first = _first_ok(client, key).response + assert first.headers.get("x-ratelimit-api_key-limit-requests") == "5", ( + f"success response must report the key's request limit, headers: " + f"{ {k: v for k, v in first.headers.items() if 'ratelimit' in k} }" + ) + assert first.headers.get("x-ratelimit-api_key-remaining-requests") == str(5 - 1), ( + f"first call against rpm_limit=5 must leave {5 - 1} remaining, got " + f"{first.headers.get('x-ratelimit-api_key-remaining-requests')!r}" + ) + assert first.headers.get("x-ratelimit-api_key-limit-tokens") == "100000", ( + f"success response must report the key's token limit, got " + f"{first.headers.get('x-ratelimit-api_key-limit-tokens')!r}" + ) + remaining_tokens = first.headers.get("x-ratelimit-api_key-remaining-tokens") + assert remaining_tokens is not None and remaining_tokens.isdigit() and int(remaining_tokens) < 100000, ( + f"one call must leave remaining tokens reported and below the limit, got {remaining_tokens!r}" + ) diff --git a/tests/e2e/test_e2e_gateway.py b/tests/e2e/test_e2e_gateway.py index 9a9aa2fd2cc..70e846ad8d5 100644 --- a/tests/e2e/test_e2e_gateway.py +++ b/tests/e2e/test_e2e_gateway.py @@ -1,17 +1,23 @@ """Unit coverage for the Gateway model-management surface (create_model / -delete_model). +delete_model) and the bounded spend read-back (spend_logs_window). The batches conftest and several llm_translation tests register deployments at runtime through gateway.create_model; when that method went missing, every batch test errored at fixture setup (AttributeError) before a single request reached the proxy. This pins the surface with a typed fake Transport so a rename or signature drift fails here instead of in a live stage run. + +spend_logs_window exists because the unpaginated /spend/logs whole-table read +grew past the e2e runner's memory limit on stage and OOMKilled every run; these +tests pin its /spend/logs/v2 pagination and that SpendLogsParams can no longer +express the unfiltered read. """ from dataclasses import dataclass, field +from datetime import datetime, timezone import pytest -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from batches.batch_client import BatchClient from e2e_gateway import Gateway @@ -30,6 +36,9 @@ from models import ( ModelNewBody, ModelNewResponse, ModelsListResponse, + SpendLogsPage, + SpendLogsPageParams, + SpendLogsParams, ) @@ -46,6 +55,8 @@ class _RecordingTransport: servable_after_gets: int = 0 models_error: UnknownApiError | None = None model_gets: int = 0 + spend_total: int = 0 + spend_gets: list[SpendLogsPageParams] = field(default_factory=list) _created: list[str] = field(default_factory=list) def post[R: BaseModel]( @@ -91,6 +102,22 @@ class _RecordingTransport: return Success( data=response_type.model_validate({"data": [{"id": name} for name in visible]}) ) + if path == "/spend/logs/v2" and response_type is SpendLogsPage: + assert isinstance(params, SpendLogsPageParams) + self.spend_gets.append(params) + offset = (params.page - 1) * params.page_size + count = min(params.page_size, max(self.spend_total - offset, 0)) + return Success( + data=response_type.model_validate( + { + "data": [{"request_id": f"req-{offset + i}"} for i in range(count)], + "total": self.spend_total, + "page": params.page, + "page_size": params.page_size, + "total_pages": (self.spend_total + params.page_size - 1) // params.page_size, + } + ) + ) raise AssertionError(f"unexpected get: {path}") def delete[R: BaseModel]( @@ -202,3 +229,45 @@ def test_gateway_delete_model_posts_the_model_id() -> None: assert path == "/model/delete" assert isinstance(body, ModelDeleteBody) assert body.id == "registered-id" + + +WINDOW_START = datetime(2026, 7, 14, 12, 0, 0, tzinfo=timezone.utc) +WINDOW_END = datetime(2026, 7, 14, 14, 0, 0, tzinfo=timezone.utc) + + +def test_gateway_spend_logs_window_pages_through_every_row_in_the_window() -> None: + transport = _RecordingTransport(spend_total=250) + gateway = Gateway(transport=transport) + + rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END) + + assert len(rows) == 250 + assert len({row.request_id for row in rows}) == 250 + assert [params.page for params in transport.spend_gets] == [1, 2, 3] + assert all(params.start_date == "2026-07-14 12:00:00" for params in transport.spend_gets) + assert all(params.end_date == "2026-07-14 14:00:00" for params in transport.spend_gets) + + +def test_gateway_spend_logs_window_stops_at_an_exact_page_boundary() -> None: + transport = _RecordingTransport(spend_total=200) + gateway = Gateway(transport=transport) + + rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END) + + assert len(rows) == 200 + assert [params.page for params in transport.spend_gets] == [1, 2] + + +def test_gateway_spend_logs_window_returns_empty_for_an_empty_window() -> None: + transport = _RecordingTransport(spend_total=0) + gateway = Gateway(transport=transport) + + rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END) + + assert rows == [] + assert [params.page for params in transport.spend_gets] == [1] + + +def test_spend_logs_params_rejects_the_unfiltered_whole_table_read() -> None: + with pytest.raises(ValidationError, match="spend_logs_window"): + SpendLogsParams() diff --git a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py index 891e5020f37..d56d5e51b04 100644 --- a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py +++ b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py @@ -2,7 +2,7 @@ import io import os import sys -from typing import Optional +from typing import Optional, Union sys.path.insert(0, os.path.abspath("../..")) @@ -12,6 +12,7 @@ import json import logging import time from unittest.mock import AsyncMock, patch +from datetime import datetime import httpx import pytest @@ -20,17 +21,24 @@ import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.responses.main import mock_responses_api_response -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import ( + ModelResponse, + ResponsesAPIResponse, + StandardLoggingPayload, + TextCompletionResponse, +) class TestCustomLogger(CustomLogger): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None + self.response_obj: Optional[Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse]] = None async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): standard_logging_payload = kwargs.get("standard_logging_object", None) self.logged_standard_logging_payload = standard_logging_payload + self.response_obj = response_obj @pytest.mark.asyncio @@ -108,6 +116,78 @@ async def test_dynamic_turn_off_message_logging_overrides_global_off(dynamic_tur assert standard_logging_payload["messages"][0]["content"] == expected_message_content +@pytest.mark.asyncio +async def test_redaction_with_custom_logger_streaming(): + """Test redaction of responses for custom logger callbacks""" + from litellm.litellm_core_utils.litellm_logging import Logging + + class LoggingWithoutSyncSuccessHandler(Logging): + def success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs): + pass + + litellm.turn_off_message_logging = True + test_custom_logger = TestCustomLogger() + + try: + litellm_logging_obj = LoggingWithoutSyncSuccessHandler( + model="gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + litellm_call_id="1234", + start_time=datetime.now(), + function_id="1234", + dynamic_async_success_callbacks=[test_custom_logger], + ) + + response = await litellm.acompletion( + model="gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello", + stream=True, + litellm_logging_obj=litellm_logging_obj, + ) + + # Consume the stream to trigger logging + chunks = [] + async for chunk in response: + chunks.append(chunk) + + await asyncio.sleep(1) + async_complete_streaming_response = test_custom_logger.response_obj + assert async_complete_streaming_response is not None + assert async_complete_streaming_response.choices[0].message.content == "redacted-by-litellm" + finally: + litellm.turn_off_message_logging = False + + +@pytest.mark.asyncio +async def test_streaming_redaction_scoped_to_opted_out_logger(): + """One logger opting out of message logging must not blank the response for other loggers""" + litellm.turn_off_message_logging = False + opted_out_logger = TestCustomLogger(message_logging=False) + compliant_logger = TestCustomLogger() + litellm.callbacks = [opted_out_logger, compliant_logger] + + try: + response = await litellm.acompletion( + model="gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello", + stream=True, + ) + async for _ in response: + pass + + await asyncio.sleep(1) + assert opted_out_logger.response_obj is not None + assert opted_out_logger.response_obj.choices[0].message.content == "redacted-by-litellm" + assert compliant_logger.response_obj is not None + assert compliant_logger.response_obj.choices[0].message.content == "hello" + finally: + litellm.callbacks = [] + + @pytest.mark.asyncio async def test_redaction_responses_api(): """Test redaction with ResponsesAPIResponse format""" diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py index e7136ecb195..e58e6c9694b 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -236,6 +236,8 @@ async def test_can_team_call_model(model, expect_to_work): (["bedrock/*"], "bedrock/anthropic.claude-3-5-sonnet-20240620", True), (["bedrock/*"], "bedrockz/anthropic.claude-3-5-sonnet-20240620", False), (["bedrock/us.*"], "bedrock/us.amazon.nova-micro-v1:0", True), + (["openai/*"], "ft:gpt-4-0613", True), + (["openai/*"], "bedrockz/ft:gpt-4-0613", False), ], ) @pytest.mark.asyncio diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index f8f15f2e008..b822799fb40 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -2,7 +2,8 @@ Unit tests for CheckBatchCost class. Covers: stale-row cleanup (file_purpose scoping), paginated find_many, the batch_processed-column fallback query, and routing of unmanaged -Vertex batches (raw gs:// input_file_id, no managed unified id). +Vertex (raw gs:// input_file_id) and Bedrock (raw s3:// input_file_id, +ARN unified_object_id) batches with no managed unified id. """ from unittest.mock import AsyncMock, MagicMock, patch @@ -31,6 +32,28 @@ def _unmanaged_vertex_file_object( ).model_dump_json() +def _unmanaged_bedrock_file_object( + input_file_id=( + "s3://bucket/litellm-bedrock-files-us.anthropic.claude-sonnet-4-20250514-v1-0" + "-74b61828-9191-4d80-addb-5a0f9ab0ec6a.jsonl" + ), + status="validating", +): + """A LiteLLMBatch JSON blob shaped like what gets stored for an unmanaged Bedrock + batch (raw s3:// input_file_id, ARN unified_object_id).""" + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id="arn:aws:bedrock:us-east-1:298249409318:model-invocation-job/1ofb47x17jua", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id=input_file_id, + object="batch", + status=status, + ).model_dump_json() + + class TestCheckBatchCost: """Test suite for CheckBatchCost class""" @@ -684,7 +707,7 @@ class TestUnmanagedVertexRouting: proxy_logging_obj=MagicMock(), prisma_client=MagicMock(), llm_router=router, - track_unmanaged_vertex_batch_cost=track_unmanaged, + track_unmanaged_batch_cost=track_unmanaged, ) def _job(self, file_object=None): @@ -916,3 +939,285 @@ class TestUnmanagedVertexRouting: update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" + + +class TestUnmanagedBedrockRouting: + """Routing of unmanaged Bedrock batches whose unified_object_id is a raw model-invocation-job ARN.""" + + _ARN = "arn:aws:bedrock:us-east-1:298249409318:model-invocation-job/1ofb47x17jua" + + def _instance(self, track_unmanaged, router): + from litellm_enterprise.proxy.common_utils.check_batch_cost import ( + CheckBatchCost, + ) + + return CheckBatchCost( + proxy_logging_obj=MagicMock(), + prisma_client=MagicMock(), + llm_router=router, + track_unmanaged_batch_cost=track_unmanaged, + ) + + def _job(self, file_object=None): + job = MagicMock() + job.unified_object_id = self._ARN + job.file_object = ( + file_object if file_object is not None else _unmanaged_bedrock_file_object() + ) + return job + + def _bedrock_deployment(self): + deployment = MagicMock() + deployment.litellm_params.custom_llm_provider = "bedrock" + deployment.litellm_params.model = "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0" + return deployment + + def test_flag_off_skips_arn_unified_id_unchanged(self): + """Default (flag off): a raw ARN unified_object_id is skipped exactly as before; no + model derivation or router lookup happens.""" + router = MagicMock() + instance = self._instance(track_unmanaged=False, router=router) + prom = MagicMock() + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(self._job(), prom) + + assert result is None + prom.record_check_batch_cost_error.assert_called_once_with("invalid_unified_id") + router.resolve_model_name_from_model_id.assert_not_called() + router.get_model_ids.assert_not_called() + + def test_flag_on_routes_to_bedrock_deployment(self): + """Flag on: derive the bare model from the s3:// object key (":" restored to "-" is + matched fuzzily), resolve it to a deployment id, and use the raw ARN as the batch id.""" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = "claude-sonnet-4" + router.get_model_ids.return_value = ["deploy-1"] + router.get_deployment = MagicMock(return_value=self._bedrock_deployment()) + instance = self._instance(track_unmanaged=True, router=router) + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(self._job(), MagicMock()) + + assert result == ("deploy-1", self._ARN) + + def test_flag_on_skips_non_bedrock_deployment_sharing_model_group(self): + """Flag on, but the only deployment for the model group is a non-bedrock provider: + must not be selected, even though the model group name matches.""" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = "claude-sonnet-4" + router.get_model_ids.return_value = ["deploy-anthropic"] + non_bedrock_deployment = MagicMock() + non_bedrock_deployment.litellm_params.custom_llm_provider = "anthropic" + non_bedrock_deployment.litellm_params.model = "claude-sonnet-4-20250514" + router.get_deployment = MagicMock(return_value=non_bedrock_deployment) + instance = self._instance(track_unmanaged=True, router=router) + prom = MagicMock() + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(self._job(), prom) + + assert result is None + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) + + def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self): + """The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured + deployment's actual bedrock model id uses ':' (e.g. 'v1:0'). Routing must still match.""" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = None + router.get_model_ids.return_value = [] + router.get_model_list.return_value = [ + { + "model_name": "claude-sonnet-4", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0", + "custom_llm_provider": "bedrock", + }, + "model_info": {"id": "deploy-bedrock"}, + } + ] + instance = self._instance(track_unmanaged=True, router=router) + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(self._job(), MagicMock()) + + assert result == ("deploy-bedrock", self._ARN) + + def test_flag_on_no_matching_deployment_records_metric(self): + """Flag on but no bedrock deployment for the model: skip with a distinct metric.""" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = None + router.get_model_ids.return_value = [] + router.get_model_list.return_value = [] + instance = self._instance(track_unmanaged=True, router=router) + prom = MagicMock() + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(self._job(), prom) + + assert result is None + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) + + def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self): + """Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as + unroutable, do not attempt model derivation.""" + router = MagicMock() + instance = self._instance(track_unmanaged=True, router=router) + prom = MagicMock() + job = self._job( + file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123") + ) + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(job, prom) + + assert result is None + prom.record_check_batch_cost_error.assert_called_once_with("invalid_unified_id") + router.resolve_model_name_from_model_id.assert_not_called() + + @pytest.mark.asyncio + async def test_end_to_end_costs_unmanaged_batch(self): + """Flag on, completed unmanaged batch: the poller polls Bedrock with the raw ARN, + computes cost, and marks batch_processed=True.""" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = "claude-sonnet-4" + router.get_model_ids.return_value = ["deploy-1"] + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = "s3://bucket/out/predictions.jsonl" + mock_response.error_file_id = None + mock_response.completed_at = None + mock_response.created_at = None + mock_response.model_dump_json.return_value = ( + f'{{"id":"{self._ARN}","status":"completed"}}' + ) + router.aretrieve_batch = AsyncMock(return_value=mock_response) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"aws_region_name": "us-east-1"} + ) + + deployment = self._bedrock_deployment() + deployment.model_name = "claude-sonnet-4" + deployment.model_info.model_dump.return_value = {} + router.get_deployment = MagicMock(return_value=deployment) + + instance = self._instance(track_unmanaged=True, router=router) + instance.proxy_logging_obj.get_proxy_hook.return_value = None + instance._has_batch_processed_column = True + + prisma = instance.prisma_client + prisma.db = MagicMock() + prisma.db.litellm_managedobjecttable = MagicMock() + prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + prisma.db.litellm_managedobjecttable.update = AsyncMock() + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) + prisma.db.litellm_usertable = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + mock_file_content = MagicMock() + mock_file_content.content = b'{"id":"req-1"}' + + with ( + patch(_IS_B64, side_effect=[False, None]), + patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + return_value=mock_file_content, + ), + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary", + return_value=[{"id": "req-1"}], + ), + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + return_value=( + 0.02, + {"prompt_tokens": 10, "completion_tokens": 5}, + ["claude-sonnet-4"], + ), + ), + patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", + return_value=("claude-sonnet-4", "bedrock", None, None), + ), + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, + ): + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + mock_logging_cls.return_value = mock_logging_obj + + await instance.check_batch_cost() + + router.aretrieve_batch.assert_awaited_once() + assert router.aretrieve_batch.call_args[1]["model"] == "deploy-1" + assert router.aretrieve_batch.call_args[1]["batch_id"] == self._ARN + + mock_logging_obj.async_success_handler.assert_awaited_once() + assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.02 + + assert prisma.db.litellm_managedobjecttable.update.call_count == 1 + update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert update_data["batch_processed"] is True + assert update_data["status"] == "complete" + + +class TestUnmanagedBatchCostFlagIsGeneralized: + """The single track_unmanaged_batch_cost flag must cover both Vertex and Bedrock, not + just the provider it was originally added for.""" + + def test_one_flag_routes_both_vertex_and_bedrock_jobs(self): + from litellm_enterprise.proxy.common_utils.check_batch_cost import ( + CheckBatchCost, + ) + + router = MagicMock() + router.resolve_model_name_from_model_id.side_effect = [ + "gemini-2.5-flash", + "claude-sonnet-4", + ] + router.get_model_ids.side_effect = [["deploy-vertex"], ["deploy-bedrock"]] + + def _get_deployment(model_id): + if model_id == "deploy-vertex": + deployment = MagicMock() + deployment.litellm_params.custom_llm_provider = "vertex_ai" + deployment.litellm_params.model = "vertex_ai/gemini-2.5-flash" + return deployment + deployment = MagicMock() + deployment.litellm_params.custom_llm_provider = "bedrock" + deployment.litellm_params.model = "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0" + return deployment + + router.get_deployment = MagicMock(side_effect=_get_deployment) + + instance = CheckBatchCost( + proxy_logging_obj=MagicMock(), + prisma_client=MagicMock(), + llm_router=router, + track_unmanaged_batch_cost=True, + ) + + vertex_job = MagicMock() + vertex_job.unified_object_id = "8823717160934178816" + vertex_job.file_object = _unmanaged_vertex_file_object() + + bedrock_job = MagicMock() + bedrock_job.unified_object_id = TestUnmanagedBedrockRouting._ARN + bedrock_job.file_object = _unmanaged_bedrock_file_object() + + with patch(_IS_B64, return_value=False): + vertex_result = instance._resolve_job_routing(vertex_job, MagicMock()) + bedrock_result = instance._resolve_job_routing(bedrock_job, MagicMock()) + + assert vertex_result == ("deploy-vertex", "8823717160934178816") + assert bedrock_result == ("deploy-bedrock", TestUnmanagedBedrockRouting._ARN) diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 958b028c542..5471d2668e4 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -219,8 +219,8 @@ async def test_aaauser_personal_budgets(key_ownership): """ Set a personal budget on a user - - have it only apply when key belongs to user -> raises BudgetExceededError - - if key belongs to team, have key respect team budget -> allows call to go through + User budget is enforced regardless of key ownership (personal or team). + Both cases should raise BudgetExceededError when the user is over budget. """ import asyncio import time @@ -229,7 +229,12 @@ async def test_aaauser_personal_budgets(key_ownership): from starlette.datastructures import URL import litellm - from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth + from litellm.proxy._types import ( + LiteLLM_UserTable, + ProxyErrorTypes, + ProxyException, + UserAPIKeyAuth, + ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import hash_token, user_api_key_cache @@ -273,14 +278,9 @@ async def test_aaauser_personal_budgets(key_ownership): == valid_token ) - try: + with pytest.raises(ProxyException) as exc_info: await user_api_key_auth(request=request, api_key="Bearer " + user_key) - - if key_ownership == "user_key": - pytest.fail("Expected this call to fail. User is over limit.") - except Exception: - if key_ownership == "team_key": - pytest.fail("Expected this call to work. Key is below team budget.") + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded @pytest.mark.asyncio diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 9de5cd69b1e..c7aecac477e 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -818,6 +818,29 @@ def test_anthropic_usage_conversion_includes_cache_tokens(): assert usage.prompt_tokens_details.cache_creation_tokens == 2000 +def test_bedrock_model_output_line_success_check(): + row = { + "recordId": "1", + "modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}}, + } + assert bu._batch_response_was_successful(row, custom_llm_provider="bedrock") is True + assert bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6" + + +def test_bedrock_cost_uses_deployment_model_name(): + row = { + "recordId": "1", + "modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}}, + } + cost = bu._get_batch_job_cost_from_file_content( + file_content_dictionary=[row], + custom_llm_provider="bedrock", + model_name="us.anthropic.claude-sonnet-4-6", + model_info={}, + ) + assert cost > 0 + + def test_anthropic_total_usage_sums_succeeded_only(): rows = [ _anthropic_succeeded_row(usage=_anthropic_usage(10, 5)), diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index d300f326b9e..9289dece83f 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1614,3 +1614,105 @@ class TestGuardrailInterventionClassification: slg = request_data["metadata"]["standard_logging_guardrail_information"][0] assert slg["guardrail_status"] == "guardrail_intervened" + + +class _ApplyStyleGuardrail(CustomGuardrail): + """Overrides only apply_guardrail, like openai_moderation; async_pre_call_hook stays the CustomLogger no-op.""" + + def __init__(self, block: bool): + from litellm.types.guardrails import GuardrailEventHooks + + super().__init__( + guardrail_name="apply-style-guardrail", + event_hook=GuardrailEventHooks.pre_call, + default_on=False, + ) + self.block = block + self.apply_called = False + self.seen_texts = None + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + from fastapi import HTTPException + + self.apply_called = True + self.seen_texts = inputs.get("texts") + if self.block: + raise HTTPException(status_code=400, detail={"error": "Violated moderation policy"}) + return inputs + + +class TestApplyGuardrailStyleDeploymentDispatch: + """LIT-4217 regression: model-level guardrails that implement only the + unified apply_guardrail interface must execute in + async_pre_call_deployment_hook instead of silently hitting the + async_pre_call_hook no-op.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", [CallTypes.completion, CallTypes.acompletion]) + async def test_blocks_when_requested_via_model_level_guardrails(self, call_type): + from fastapi import HTTPException + + guardrail = _ApplyStyleGuardrail(block=True) + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "flagged content"}], + "guardrails": ["apply-style-guardrail"], + "metadata": {}, + } + + with pytest.raises(HTTPException): + await guardrail.async_pre_call_deployment_hook(kwargs, call_type) + + assert guardrail.apply_called is True + assert guardrail.seen_texts == ["flagged content"] + + @pytest.mark.asyncio + async def test_pass_path_runs_guardrail_and_strips_dispatch_key(self): + guardrail = _ApplyStyleGuardrail(block=False) + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "guardrails": ["apply-style-guardrail"], + "metadata": {}, + } + + result = await guardrail.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert guardrail.apply_called is True + assert result is not None + assert "guardrail_to_apply" not in result + assert result["messages"] == [{"role": "user", "content": "hello"}] + + @pytest.mark.asyncio + async def test_skips_when_not_requested(self): + guardrail = _ApplyStyleGuardrail(block=True) + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "guardrails": ["some-other-guardrail"], + "metadata": {}, + } + + result = await guardrail.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert guardrail.apply_called is False + assert result is not None + + @pytest.mark.asyncio + async def test_fails_closed_when_proxy_extras_missing(self): + import sys + from unittest.mock import patch + + guardrail = _ApplyStyleGuardrail(block=True) + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "flagged content"}], + "guardrails": ["apply-style-guardrail"], + "metadata": {}, + } + + with patch.dict(sys.modules, {"litellm.proxy.utils": None}): + with pytest.raises(ImportError, match="litellm\\[proxy\\]"): + await guardrail.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert guardrail.apply_called is False diff --git a/tests/test_litellm/integrations/test_prometheus_media_generation_metrics.py b/tests/test_litellm/integrations/test_prometheus_media_generation_metrics.py new file mode 100644 index 00000000000..2ce92f08ef0 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_media_generation_metrics.py @@ -0,0 +1,180 @@ +""" +Unit tests for the video-seconds and images-generated Prometheus counters (LIT-4254). + +Video providers report ``duration_seconds`` inside the usage object that lands +on ``standard_logging_payload["metadata"]["usage_object"]``; image generation +calls report ``output_image_count`` there. Both counters are sparse: only +incremented when the value is present and > 0. +""" + +from typing import get_args +from unittest.mock import MagicMock + +import pytest + +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import ( + DEFINED_PROMETHEUS_METRICS, + PrometheusMetricLabels, + UserAPIKeyLabelValues, +) + +MEDIA_GENERATION_METRICS = [ + "litellm_video_duration_seconds_metric", + "litellm_images_generated_metric", +] + + +@pytest.fixture +def sample_enum_values(): + return UserAPIKeyLabelValues( + end_user="test-end-user", + hashed_api_key="test-key-hash", + api_key_alias="test-key-alias", + team="test-team", + team_alias="test-team-alias", + user="test-user", + model="sora-2", + ) + + +def _make_mock_logger(): + logger = MagicMock() + for name in MEDIA_GENERATION_METRICS: + setattr(logger, name, MagicMock()) + logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + return logger + + +class TestMediaGenerationMetricsRegistration: + def test_metrics_in_defined_prometheus_metrics(self): + defined = get_args(DEFINED_PROMETHEUS_METRICS) + for name in MEDIA_GENERATION_METRICS: + assert name in defined, f"{name} missing from DEFINED_PROMETHEUS_METRICS" + + def test_metric_labels_defined(self): + for name in MEDIA_GENERATION_METRICS: + assert hasattr(PrometheusMetricLabels, name), f"{name} missing from PrometheusMetricLabels" + + def test_metrics_share_output_token_label_set(self): + assert ( + PrometheusMetricLabels.litellm_video_duration_seconds_metric + == PrometheusMetricLabels.litellm_output_tokens_metric + ) + assert ( + PrometheusMetricLabels.litellm_images_generated_metric + == PrometheusMetricLabels.litellm_output_tokens_metric + ) + + def test_runtime_label_set_matches_output_tokens_metric(self): + """Full parity with litellm_output_tokens_metric, including the org labels + appended via _org_label_metrics, so existing token dashboards can be cloned.""" + expected = PrometheusMetricLabels.get_labels("litellm_output_tokens_metric") + for name in MEDIA_GENERATION_METRICS: + assert PrometheusMetricLabels.get_labels(name) == expected + + +class TestIncrementMediaGenerationMetrics: + def test_video_duration_incremented(self, sample_enum_values): + logger = _make_mock_logger() + payload = {"metadata": {"usage_object": {"duration_seconds": 8.0}}} + + PrometheusLogger._increment_media_generation_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + logger.litellm_video_duration_seconds_metric.labels().inc.assert_called_once_with(8.0) + logger.litellm_images_generated_metric.labels.assert_not_called() + + def test_image_count_incremented(self, sample_enum_values): + logger = _make_mock_logger() + payload = { + "metadata": { + "usage_object": { + "prompt_tokens": 18, + "completion_tokens": 391, + "total_tokens": 409, + "output_image_count": 2, + } + } + } + + PrometheusLogger._increment_media_generation_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + logger.litellm_images_generated_metric.labels().inc.assert_called_once_with(2.0) + logger.litellm_video_duration_seconds_metric.labels.assert_not_called() + + def test_token_only_usage_is_a_noop(self, sample_enum_values): + logger = _make_mock_logger() + payload = { + "metadata": { + "usage_object": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30, + } + } + } + + PrometheusLogger._increment_media_generation_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + for name in MEDIA_GENERATION_METRICS: + getattr(logger, name).labels.assert_not_called() + + @pytest.mark.parametrize("bad_value", [0, 0.0, None, -4.0, "4", True]) + def test_non_positive_or_non_numeric_values_are_ignored(self, sample_enum_values, bad_value): + logger = _make_mock_logger() + payload = { + "metadata": { + "usage_object": { + "duration_seconds": bad_value, + "output_image_count": bad_value, + } + } + } + + PrometheusLogger._increment_media_generation_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + for name in MEDIA_GENERATION_METRICS: + getattr(logger, name).labels.assert_not_called() + + def test_missing_usage_object_is_a_noop(self, sample_enum_values): + logger = _make_mock_logger() + + for payload in ({"metadata": {}}, {"metadata": None}, {"metadata": {"usage_object": "redacted"}}): + PrometheusLogger._increment_media_generation_metrics( + logger, + standard_logging_payload=payload, + enum_values=sample_enum_values, + ) + + for name in MEDIA_GENERATION_METRICS: + getattr(logger, name).labels.assert_not_called() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py index bb035c4c3ee..9c6d2e018ff 100644 --- a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py @@ -326,3 +326,148 @@ async def test_should_leave_rate_limit_labels_blank_for_non_rate_limit_failure() assert isinstance(enum_values, UserAPIKeyLabelValues) assert enum_values.rate_limit_category is None assert enum_values.rate_limit_type is None + + +def _logger_with_mock_virtual_key_gauges() -> PrometheusLogger: + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_remaining_api_key_requests_for_model = MagicMock() + logger.litellm_remaining_api_key_tokens_for_model = MagicMock() + logger.get_labels_for_metric = MagicMock(return_value=[]) + return logger + + +def _kwargs_with_v3_rate_limit_headers(additional_headers: dict) -> dict: + return { + "litellm_params": {"metadata": {"model_group": "gpt-4o-mini"}}, + "standard_logging_object": { + "metadata": {}, + "hidden_params": {"additional_headers": additional_headers}, + }, + } + + +def _set_virtual_key_metrics(logger: PrometheusLogger, kwargs: dict) -> None: + logger._set_virtual_key_rate_limit_metrics( + user_api_key="test-hash", + user_api_key_alias="test-alias", + kwargs=kwargs, + metadata=kwargs["litellm_params"]["metadata"], + model_id="model-123", + ) + + +def test_should_read_v3_remaining_headers_when_metadata_keys_absent(): + """ + Regression for LIT-2577: the default v3 rate limiter writes remaining + per-(key, model) values into + ``standard_logging_object.hidden_params.additional_headers`` as + ``x-ratelimit-model_per_key-remaining-{requests,tokens}`` and never sets + the legacy ``litellm-key-remaining-*`` metadata keys, so the gauges were + pinned to ``sys.maxsize``. + """ + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": 42, + "x-ratelimit-model_per_key-remaining-tokens": 900, + "x-ratelimit-model_per_key-limit-requests": 100, + "x-ratelimit-model_per_key-limit-tokens": 1000, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + 42 + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + 900 + ) + + +def test_should_prefer_legacy_metadata_keys_over_v3_headers(): + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": 42, + "x-ratelimit-model_per_key-remaining-tokens": 900, + } + ) + kwargs["litellm_params"]["metadata"].update( + { + "litellm-key-remaining-requests-gpt-4o-mini": 3, + "litellm-key-remaining-tokens-gpt-4o-mini": 200, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + 3 + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + 200 + ) + + +def test_should_treat_zero_v3_remaining_as_zero(): + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": 0, + "x-ratelimit-model_per_key-remaining-tokens": 0, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + 0 + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + 0 + ) + + +def test_should_keep_maxsize_sentinel_when_no_rate_limit_source_present(): + import sys + + logger = _logger_with_mock_virtual_key_gauges() + kwargs = { + "litellm_params": {"metadata": {"model_group": "gpt-4o-mini"}}, + "standard_logging_object": {"metadata": {}, "hidden_params": {}}, + } + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + ) + + +@pytest.mark.parametrize("bad_value", ["not-a-number", None, True]) +def test_should_ignore_non_int_v3_header_values(bad_value): + import sys + + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": bad_value, + "x-ratelimit-model_per_key-remaining-tokens": bad_value, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + ) diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 246b378c982..f0a33f2ebfc 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1125,6 +1125,47 @@ async def test_combined_prefix_reflects_in_s3_object_key(): assert "myteam/apikey/" in key, f"Expected both prefixes in key: {key}" +def test_s3_object_key_sanitizes_slashes_in_file_name(): + """Response ids containing slashes (e.g. bedrock batch job ARNs) must not + create nested S3 folders; only path/prefix/date slashes are separators.""" + from litellm.integrations.s3 import get_s3_object_key + + start_time = datetime(2026, 2, 11, 0, 35, 18, 391582) + file_name = "time-00-35-18-391582_arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/gl18r6skk9yy" + + key = get_s3_object_key( + s3_path="LiteLLMAPPLogs", + prefix="myteam/", + start_time=start_time, + s3_file_name=file_name, + ) + + assert key == ( + "LiteLLMAPPLogs/myteam/2026-02-11/" + "time-00-35-18-391582_arn:aws:bedrock:us-east-1:123456789012:model-invocation-job_gl18r6skk9yy.json" + ) + + +def test_create_s3_batch_logging_element_flat_key_for_arn_response_id(): + """End-to-end through the s3_v2 element builder: an ARN response id must + yield a flat file directly under the date segment.""" + logger = S3Logger(s3_use_team_prefix=False, s3_use_key_prefix=False) + payload = StandardLoggingPayload( + id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/gl18r6skk9yy", + metadata={}, + messages=[], + ) + + start_time = datetime(2026, 2, 11, 0, 35, 18, 391582) + result = logger.create_s3_batch_logging_element(start_time, payload) + + assert result is not None + date_segment = "2026-02-11/" + file_segment = result.s3_object_key.split(date_segment, 1)[1] + assert "/" not in file_segment, f"Expected flat file under date segment, got: {result.s3_object_key}" + assert file_segment.endswith("model-invocation-job_gl18r6skk9yy.json") + + # -------------------------------------------------------------- # params_source / s3_callback_params_override (audit-log decoupling) # -------------------------------------------------------------- diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index ca23e61352e..b156faf3ea6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -270,6 +270,105 @@ def test_image_tokens_fallback_to_base_cost(): assert round(completion_cost, 12) == round(expected_completion_cost, 12) +def test_video_output_tokens_gemini_omni_flash_preview(): + """Video output tokens are billed at output_cost_per_video_token, not the text rate and not zero.""" + model = "gemini-omni-flash-preview" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + text_tokens = 100 + video_tokens = 46336 + usage = Usage( + completion_tokens=text_tokens + video_tokens, + prompt_tokens=20, + total_tokens=20 + text_tokens + video_tokens, + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=text_tokens, + video_tokens=video_tokens, + ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=20), + ) + model_cost_map = litellm.model_cost[f"gemini/{model}"] + assert model_cost_map["input_cost_per_token"] == 1.5e-06 + assert model_cost_map["output_cost_per_token"] == 9e-06 + assert model_cost_map["output_cost_per_video_token"] == 1.75e-05 + + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="gemini", + ) + + assert round(prompt_cost, 10) == round( + model_cost_map["input_cost_per_token"] * usage.prompt_tokens, + 10, + ) + assert round(completion_cost, 10) == round( + (model_cost_map["output_cost_per_token"] * text_tokens) + + (model_cost_map["output_cost_per_video_token"] * video_tokens), + 10, + ) + + +def test_video_input_tokens_gemini_omni_flash_preview(): + """Video input tokens are billed at the standard input rate instead of being dropped.""" + model = "gemini-omni-flash-preview" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + usage = Usage( + completion_tokens=10, + prompt_tokens=10050, + total_tokens=10060, + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=10), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=50, video_tokens=10000), + ) + model_cost_map = litellm.model_cost[f"gemini/{model}"] + + prompt_cost, _ = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="gemini", + ) + + assert round(prompt_cost, 10) == round( + model_cost_map["input_cost_per_token"] * usage.prompt_tokens, + 10, + ) + + +def test_video_tokens_fallback_to_base_cost(): + """Video output tokens fall back to the base output rate when output_cost_per_video_token is not set.""" + from unittest.mock import patch + + mock_model_info = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + } + + usage = Usage( + completion_tokens=1720, + prompt_tokens=14, + total_tokens=1734, + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=600, + video_tokens=1120, + ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=14), + ) + + with patch( + "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", + return_value=mock_model_info, + ): + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", usage=usage, custom_llm_provider="gemini" + ) + + assert round(prompt_cost, 12) == round(14 * 1e-6, 12) + assert round(completion_cost, 12) == round((600 + 1120) * 2e-6, 12) + + def test_generic_cost_per_token_above_200k_tokens(): # gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing model = "gemini-2.5-pro" @@ -1086,6 +1185,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): "text_tokens": 0, "audio_tokens": 0, "image_tokens": 0, + "video_tokens": 0, "character_count": 0, "image_count": 0, "video_length_seconds": 0.0, diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index bdd71f28b1a..1a38b5dc769 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -138,6 +138,42 @@ def test_shipped_backup_carries_the_claude_routing_rules(): set_fallback_generalizations(previous) +def test_shipped_routing_rules_never_match_through_an_unrecognized_namespace(): + """Routing rules decide ``litellm_provider`` for otherwise-unknown ids, and the + proxy's wildcard access check (``can_key_call_model`` with a ``bedrock/*`` key) + trusts that inference: it rebuilds ``{provider}/{model}`` and matches it against + the key's patterns. A routing pattern that matches as a substring lets + ``bedrockz/anthropic.claude-...`` resolve to bedrock and slip through a + ``bedrock/*`` key, so every shipped routing rule must anchor to the start of + the name and never match an id carrying an unrecognized namespace prefix.""" + backup = GetModelCostMap.load_local_model_cost_map() + rules = backup[FALLBACK_GENERALIZATIONS_KEY]["rules"] + + routing_rules = [r for r in rules if "litellm_provider" in r["model_info"]] + assert routing_rules + assert all(r["pattern"].startswith("^") for r in routing_rules) + + previous = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations(rules) + for bedrock_id in [ + "anthropic.claude-3-5-sonnet-20240620-v1:0", + "anthropic.claude-v2:1", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "us-gov.anthropic.claude-3-5-sonnet-20240620-v1:0", + "global.anthropic.claude-fable-5-20260120-v1:0", + ]: + assert match_routing_generalization(bedrock_id) == "bedrock", bedrock_id + for namespaced in [ + "bedrockz/anthropic.claude-3-5-sonnet-20240620", + "bedrockz/us.anthropic.claude-3-5-sonnet-20240620-v1:0", + "bedrockz/claude-3-5-sonnet-20240620", + ]: + assert match_routing_generalization(namespaced) is None, namespaced + finally: + set_fallback_generalizations(previous) + + def test_shipped_backup_marks_claude_4_6_plus_adaptive_not_4_0(): """Adaptive thinking is data, not code. The bundled backup must carry supports_adaptive_thinking on genuine Claude >= 4.6 entries (every provider diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 0523ed7ecb1..ade2c677745 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3707,3 +3707,69 @@ def test_set_cost_breakdown_stores_reasoning_cost(): cost_for_built_in_tools_cost_usd_dollar=0.0, ) assert "reasoning_cost" not in no_reasoning.cost_breakdown + + +def _build_payload_for_media_response(logging_obj, init_response_obj, kwargs=None): + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.datetime.now() + return get_standard_logging_object_payload( + kwargs=kwargs or {"litellm_call_id": "media-call-id", "model": "test-model", "messages": []}, + init_response_obj=init_response_obj, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + +def test_image_response_sets_output_image_count_on_usage_object(logging_obj): + """Generated-image count must land on metadata.usage_object for callbacks (e.g. Prometheus).""" + from litellm.types.utils import ImageResponse + + response = ImageResponse(created=1, data=[{"url": "https://img/1"}, {"url": "https://img/2"}]) + + payload = _build_payload_for_media_response(logging_obj, response) + + assert payload is not None + assert payload["metadata"]["usage_object"]["output_image_count"] == 2 + + +def test_output_image_count_survives_message_redaction(logging_obj, monkeypatch): + """Redaction replaces the ImageResponse body, so the count must be captured pre-redaction.""" + import litellm + from litellm.types.utils import ImageResponse + + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + response = ImageResponse(created=1, data=[{"url": "https://img/1"}]) + + payload = _build_payload_for_media_response(logging_obj, response) + + assert payload is not None + assert payload["response"] == {"text": "redacted-by-litellm"} + assert payload["metadata"]["usage_object"]["output_image_count"] == 1 + + +def test_non_image_response_has_no_output_image_count(logging_obj): + payload = _build_payload_for_media_response( + logging_obj, {"id": "chatcmpl-1", "usage": {"prompt_tokens": 1, "completion_tokens": 2}} + ) + + assert payload is not None + assert "output_image_count" not in payload["metadata"]["usage_object"] + + +def test_zero_token_video_usage_preserves_duration_seconds(logging_obj): + """Video usage bills by duration; the payload must keep duration_seconds even with zero tokens.""" + payload = _build_payload_for_media_response( + logging_obj, {"id": "video-1", "usage": {"duration_seconds": 4.0}} + ) + + assert payload is not None + assert payload["metadata"]["usage_object"]["duration_seconds"] == 4.0 + assert payload["total_tokens"] == 0 + assert payload["completion_tokens"] == 0 diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 36f220f9a2c..0f7f492ddb6 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -10,9 +10,11 @@ from types import SimpleNamespace import pytest import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.redact_messages import ( _redact_responses_api_output, perform_redaction, + redact_streaming_responses_for_custom_logger, should_redact_message_logging, ) from litellm.responses.main import mock_responses_api_response @@ -442,3 +444,109 @@ class TestPerformRedaction: assert "vertex_ai_url_context_metadata" not in hidden_params assert "vertex_ai_safety_ratings" not in hidden_params assert "vertex_ai_citation_metadata" not in hidden_params + + def test_redact_async_complete_streaming_response(self): + """Test that async_complete_streaming_response is properly redacted.""" + response_obj = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret content", role="assistant") + ) + ] + ) + + model_call_details = { + "messages": [{"role": "user", "content": "hi"}], + "prompt": "hi", + "input": "hi", + "stream": True, + "async_complete_streaming_response": response_obj, + } + + perform_redaction(model_call_details, result=None) + + redacted_response = model_call_details["async_complete_streaming_response"] + assert redacted_response.choices[0].message.content == "redacted-by-litellm" + + def test_redact_complete_streaming_response(self): + """Test that complete_streaming_response is properly redacted.""" + response_obj = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret content", role="assistant") + ) + ] + ) + + model_call_details = { + "messages": [{"role": "user", "content": "hi"}], + "prompt": "hi", + "input": "hi", + "stream": True, + "complete_streaming_response": response_obj, + } + + perform_redaction(model_call_details, result=None) + + redacted_response = model_call_details["complete_streaming_response"] + assert redacted_response.choices[0].message.content == "redacted-by-litellm" + + def test_streaming_responses_untouched_when_disabled(self): + response_obj = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret content", role="assistant") + ) + ] + ) + + model_call_details = { + "messages": [{"role": "user", "content": "hi"}], + "prompt": "hi", + "input": "hi", + "stream": True, + "async_complete_streaming_response": response_obj, + } + + perform_redaction(model_call_details, result=None, redact_streaming_responses=False) + + assert response_obj.choices[0].message.content == "secret content" + + +class TestRedactStreamingResponsesForCustomLogger: + def _model_call_details(self): + response_obj = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret content", role="assistant") + ) + ] + ) + return { + "stream": True, + "async_complete_streaming_response": response_obj, + }, response_obj + + def test_opted_out_logger_gets_redacted_copy(self): + model_call_details, response_obj = self._model_call_details() + opted_out_logger = CustomLogger(message_logging=False) + + redacted_details = redact_streaming_responses_for_custom_logger( + model_call_details=model_call_details, custom_logger=opted_out_logger + ) + + redacted_response = redacted_details["async_complete_streaming_response"] + assert redacted_response.choices[0].message.content == "redacted-by-litellm" + assert response_obj.choices[0].message.content == "secret content" + assert model_call_details["async_complete_streaming_response"] is response_obj + + def test_compliant_logger_gets_shared_response(self): + model_call_details, response_obj = self._model_call_details() + compliant_logger = CustomLogger() + + result_details = redact_streaming_responses_for_custom_logger( + model_call_details=model_call_details, custom_logger=compliant_logger + ) + + assert result_details is model_call_details + assert response_obj.choices[0].message.content == "secret content" diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 7fb38544c52..94a4a3fc945 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -10,6 +10,7 @@ from unittest.mock import MagicMock, patch import litellm from litellm.constants import ( + ANTHROPIC_MIN_THINKING_BUDGET_TOKENS, DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET, @@ -956,15 +957,15 @@ def test_anthropic_structured_output_beta_header(): @pytest.mark.parametrize( "model_name", [ - "claude-opus-4-6-20250918", - "claude-opus-4.6-20250918", + "claude-opus-4-8", + "claude-opus-4-6-20260205", "claude-opus-4-5-20251101", "claude-opus-4.5-20251101", ], ) def test_opus_uses_native_structured_output(model_name): """ - Test that Opus 4.5 and 4.6 models use native Anthropic structured outputs + Test that supported Opus models use native Anthropic structured outputs (output_format) rather than the tool-based workaround. """ config = AnthropicConfig() @@ -1004,6 +1005,43 @@ def test_opus_uses_native_structured_output(model_name): assert optional_params.get("json_mode") is True +def test_native_structured_output_uses_bundled_capability_when_remote_map_lags( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = "claude-opus-4-8" + monkeypatch.setattr( + litellm, + "model_cost", + {model: {"supports_response_schema": True}}, + ) + litellm.get_model_info.cache_clear() + + try: + optional_params = AnthropicConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "answer", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + }, + }, + } + }, + optional_params={}, + model=model, + drop_params=False, + ) + finally: + litellm.get_model_info.cache_clear() + + assert "output_format" in optional_params + assert "tools" not in optional_params + + def test_non_structured_output_model_uses_tool_workaround(): """ Test that models NOT in the native structured output list still use the @@ -2443,6 +2481,78 @@ def test_reasoning_effort_maps_to_adaptive_thinking_for_claude_4_6_models(): assert result["output_config"]["effort"] == effort_map[effort] +def test_raw_adaptive_thinking_translates_to_legacy_for_pre_46_model(): + """Clients like Claude Code send ``thinking={"type": "adaptive"}`` directly + (not via ``reasoning_effort``) on every request, regardless of which model + the request routes to. For a pre-4.6 model that doesn't understand + adaptive thinking, this must be translated to the legacy + ``thinking={type: enabled, budget_tokens}`` interface instead of being + forwarded raw, which Anthropic would reject.""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"thinking": {"type": "adaptive"}, "max_tokens": 8192}, + optional_params={}, + model="claude-haiku-4-5-20251001", + drop_params=False, + ) + + assert result["thinking"]["type"] == "enabled" + assert result["thinking"]["budget_tokens"] == DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET + + +def test_raw_adaptive_thinking_budget_capped_below_max_tokens(): + """Anthropic requires ``max_tokens > thinking.budget_tokens``. When the + default medium budget wouldn't fit, it must be capped below max_tokens + rather than forwarded as an invalid combination.""" + config = AnthropicConfig() + + max_tokens = DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET - 100 + result = config.map_openai_params( + non_default_params={"thinking": {"type": "adaptive"}, "max_tokens": max_tokens}, + optional_params={}, + model="claude-haiku-4-5-20251001", + drop_params=False, + ) + + assert result["thinking"]["type"] == "enabled" + assert result["thinking"]["budget_tokens"] == max_tokens - 1 + + +def test_raw_adaptive_thinking_dropped_when_max_tokens_too_small(): + """When max_tokens can't fit even the minimum thinking budget, thinking + must be dropped entirely so the request still succeeds, matching how the + native /v1/messages passthrough already handles this.""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={ + "thinking": {"type": "adaptive"}, + "max_tokens": ANTHROPIC_MIN_THINKING_BUDGET_TOKENS, + }, + optional_params={}, + model="claude-haiku-4-5-20251001", + drop_params=False, + ) + + assert "thinking" not in result + + +def test_raw_adaptive_thinking_untouched_for_46_plus_model(): + """Adaptive-thinking models understand ``thinking={"type": "adaptive"}`` + natively, so it must pass through unmodified.""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"thinking": {"type": "adaptive"}, "max_tokens": 8192}, + optional_params={}, + model="claude-sonnet-4-6-20260219", + drop_params=False, + ) + + assert result["thinking"] == {"type": "adaptive"} + + @pytest.fixture def local_model_cost_map(monkeypatch): original_model_cost = litellm.model_cost diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py index 06d3effcfbb..5254808e315 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py @@ -2,6 +2,7 @@ import pytest from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, ) @@ -174,6 +175,81 @@ def test_unrecognized_effort_raises_clean_400(): assert exc_info.value.status_code == 400 +def test_pinned_temperature_dropped_when_adaptive_downgraded_to_enabled(): + """Regression (#33203): Claude Code's safety classifier sends adaptive thinking + + temperature=0 to Haiku 4.5. The adaptive interface is downgraded to legacy enabled + thinking, but Anthropic rejects "temperature may only be set to 1 when thinking is + enabled". The pinned temperature must be dropped so the request succeeds while the + downgraded thinking is preserved.""" + params = _claude_code_payload(effort="medium") + params["temperature"] = 0 + result = _transform("claude-haiku-4-5", params) + + assert result["thinking"] == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + } + assert "temperature" not in result + + +def test_temperature_one_preserved_with_enabled_thinking(): + """temperature=1 is compatible with extended thinking, so it must be kept.""" + params = _claude_code_payload(effort="medium") + params["temperature"] = 1 + result = _transform("claude-haiku-4-5", params) + + assert result["thinking"]["type"] == "enabled" + assert result["temperature"] == 1 + + +def test_pinned_temperature_preserved_when_thinking_dropped(): + """When thinking is dropped entirely (non-reasoning model), there is no thinking + conflict, so a pinned temperature must survive untouched.""" + params = _claude_code_payload(effort="medium") + params["temperature"] = 0 + result = _transform("claude-3-5-haiku-latest", params) + + assert "thinking" not in result + assert result["temperature"] == 0 + + +def test_pinned_temperature_preserved_for_adaptive_model(): + """Adaptive models (4.6+) own the thinking/temperature relationship natively, so + the passthrough must not strip a pinned temperature for them.""" + params = _claude_code_payload(effort="high") + params["temperature"] = 0 + result = _transform("claude-sonnet-4-6", params) + + assert result["thinking"] == {"type": "adaptive"} + assert result["temperature"] == 0 + + +def test_pinned_temperature_dropped_for_opus_4_5_effort(): + """Opus 4.5 keeps native output_config.effort (extended thinking), which is equally + incompatible with a pinned non-1 temperature, so the temperature must be dropped.""" + params = _claude_code_payload(effort="medium") + params["temperature"] = 0 + result = _transform("claude-opus-4-5", params) + + assert result["output_config"] == {"effort": "medium"} + assert "temperature" not in result + + +def test_reasoning_effort_with_pinned_temperature_drops_temperature(): + """The reasoning_effort alias synthesizes legacy enabled thinking on a non-adaptive + model; a co-pinned non-1 temperature must be dropped to avoid the Anthropic 400.""" + result = _transform( + "claude-haiku-4-5", + {"max_tokens": 8192, "reasoning_effort": "low", "temperature": 0}, + ) + + assert result["thinking"] == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, + } + assert "temperature" not in result + + def test_non_adaptive_request_without_effort_is_untouched(): """A non-adaptive model receiving a request with no adaptive interface (no effort, no adaptive thinking) must pass through untouched.""" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 14e872485ec..606ff39b35e 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -99,17 +99,11 @@ class TestContextManagementConversion: } ) kwargs = _ADAPTER.translate_request(req) - assert kwargs["context_management"] == [ - {"type": "compaction", "compact_threshold": 100000} - ] + assert kwargs["context_management"] == [{"type": "compaction", "compact_threshold": 100000}] def test_translate_request_drops_anthropic_only_context_management(self): """context_management with only unknown edit types is omitted from kwargs.""" - req = _make_request( - context_management={ - "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] - } - ) + req = _make_request(context_management={"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}) kwargs = _ADAPTER.translate_request(req) assert "context_management" not in kwargs @@ -134,9 +128,7 @@ class TestOutputConfigStructuredOutput: def test_output_config_format_json_schema_converted(self): """output_config.format.json_schema is converted to OpenAI text.format.""" - req = _make_request( - output_config={"format": {"type": "json_schema", "schema": self._SCHEMA}} - ) + req = _make_request(output_config={"format": {"type": "json_schema", "schema": self._SCHEMA}}) kwargs = _ADAPTER.translate_request(req) assert "text" in kwargs fmt = kwargs["text"]["format"] @@ -153,9 +145,7 @@ class TestOutputConfigStructuredOutput: def test_output_format_still_works(self): """The original output_format field still takes precedence when present.""" - req = _make_request( - output_format={"type": "json_schema", "schema": self._SCHEMA} - ) + req = _make_request(output_format={"type": "json_schema", "schema": self._SCHEMA}) kwargs = _ADAPTER.translate_request(req) assert "text" in kwargs assert kwargs["text"]["format"]["type"] == "json_schema" @@ -250,9 +240,7 @@ class TestTranslateMessagesToResponsesInput: ] result = _translate_messages(messages) assert len(result) == 1 - assert result[0]["content"] == [ - {"type": "input_image", "image_url": "data:image/png;base64,abc123"} - ] + assert result[0]["content"] == [{"type": "input_image", "image_url": "data:image/png;base64,abc123"}] def test_user_url_image(self): """User message with URL image source becomes input_image with the URL.""" @@ -268,9 +256,7 @@ class TestTranslateMessagesToResponsesInput: } ] result = _translate_messages(messages) - assert result[0]["content"] == [ - {"type": "input_image", "image_url": "https://example.com/img.jpg"} - ] + assert result[0]["content"] == [{"type": "input_image", "image_url": "https://example.com/img.jpg"}] def test_user_base64_image_empty_data_skipped(self): """Base64 image with empty data is skipped (no URL can be formed).""" @@ -341,9 +327,7 @@ class TestTranslateMessagesToResponsesInput: messages = [ { "role": "user", - "content": [ - {"type": "tool_result", "tool_use_id": "call_null", "content": None} - ], + "content": [{"type": "tool_result", "tool_use_id": "call_null", "content": None}], } ] result = _translate_messages(messages) @@ -370,9 +354,7 @@ class TestTranslateMessagesToResponsesInput: } ] result = _translate_messages(messages) - assert result[0]["content"] == [ - {"type": "output_text", "text": "Here is the answer."} - ] + assert result[0]["content"] == [{"type": "output_text", "text": "Here is the answer."}] def test_assistant_tool_use_becomes_function_call(self): """Assistant tool_use block becomes a top-level function_call item.""" @@ -404,15 +386,11 @@ class TestTranslateMessagesToResponsesInput: messages = [ { "role": "assistant", - "content": [ - {"type": "thinking", "thinking": "Let me reason step by step."} - ], + "content": [{"type": "thinking", "thinking": "Let me reason step by step."}], } ] result = _translate_messages(messages) - assert result[0]["content"] == [ - {"type": "output_text", "text": "Let me reason step by step."} - ] + assert result[0]["content"] == [{"type": "output_text", "text": "Let me reason step by step."}] def test_assistant_empty_thinking_block_skipped(self): """Assistant thinking block with empty thinking text is skipped.""" @@ -584,28 +562,27 @@ class TestTranslateToolsToResponsesAPI: class TestTranslateToolChoiceToResponsesAPI: - """Anthropic tool_choice -> Responses API tool_choice.""" + """Anthropic tool_choice -> Responses API tool_choice. - def test_auto_maps_to_auto(self): - assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "auto"}) == { - "type": "auto" - } + The Responses API's tool_choice schema (openai.types.responses.tool_choice_options) + is a bare Literal["none", "auto", "required"] for these simple cases - not an + object like {"type": "auto"}. Sending the object shape to an OpenAI-compatible + server gets rejected with a pydantic validation error. + """ - def test_any_maps_to_required(self): - assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "any"}) == { - "type": "required" - } + def test_auto_maps_to_bare_string_auto(self): + assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "auto"}) == "auto" + + def test_any_maps_to_bare_string_required(self): + assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "any"}) == "required" + + def test_none_maps_to_bare_string_none(self): + assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "none"}) == "none" def test_specific_tool_maps_to_function(self): - result = _ADAPTER.translate_tool_choice_to_responses_api( - {"type": "tool", "name": "get_weather"} - ) + result = _ADAPTER.translate_tool_choice_to_responses_api({"type": "tool", "name": "get_weather"}) assert result == {"type": "function", "name": "get_weather"} - def test_unknown_type_defaults_to_auto(self): - result = _ADAPTER.translate_tool_choice_to_responses_api({"type": "none"}) - assert result == {"type": "auto"} - # --------------------------------------------------------------------------- # translate_thinking_to_reasoning @@ -616,17 +593,13 @@ class TestTranslateThinkingToReasoning: """Anthropic thinking param -> Responses API reasoning param.""" def test_budget_high_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning( - {"type": "enabled", "budget_tokens": 10000} - ) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}) # Default (reasoning_auto_summary=False): only effort, no summary assert result == {"effort": "high"} assert result is not None and "summary" not in result def test_budget_above_threshold_high_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning( - {"type": "enabled", "budget_tokens": 50000} - ) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 50000}) assert result is not None assert result["effort"] == "high" assert "summary" not in result @@ -652,9 +625,7 @@ class TestTranslateThinkingToReasoning: assert result is not None and "summary" not in result def test_budget_minimal_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning( - {"type": "enabled", "budget_tokens": 500} - ) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 500}) assert result == {"effort": "minimal"} assert result is not None and "summary" not in result @@ -707,9 +678,7 @@ class TestTranslateThinkingToReasoning: original = litellm.reasoning_auto_summary try: litellm.reasoning_auto_summary = True - result = _ADAPTER.translate_thinking_to_reasoning( - {"type": "enabled", "budget_tokens": 10000} - ) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}) assert result == {"effort": "high", "summary": "detailed"} finally: litellm.reasoning_auto_summary = original @@ -789,11 +758,7 @@ class TestTranslateRequestBroaderCoverage: assert kwargs["top_p"] == 0.9 def test_tools_translated(self): - req = _make_request( - tools=[ - {"name": "calculator", "description": "Does math.", "input_schema": {}} - ] - ) + req = _make_request(tools=[{"name": "calculator", "description": "Does math.", "input_schema": {}}]) kwargs = _ADAPTER.translate_request(req) assert len(kwargs["tools"]) == 1 assert kwargs["tools"][0]["name"] == "calculator" @@ -929,9 +894,7 @@ class TestTranslateResponse: def test_multiple_text_parts(self): """Multiple output_text parts become multiple text content blocks.""" - response = _make_mock_response( - output=[_make_output_message(["Part 1", "Part 2"])] - ) + response = _make_mock_response(output=[_make_output_message(["Part 1", "Part 2"])]) result: Any = _ADAPTER.translate_response(response) assert len(result["content"]) == 2 assert result["content"][0]["text"] == "Part 1" diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 9f2a4168dec..fc12ead36a1 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -5767,3 +5767,73 @@ def test_message_level_cache_control_drops_ttl_for_unsupported_model(ttl_target) cache_points = _collect_cache_points(result) assert len(cache_points) == 1 assert "ttl" not in cache_points[0] + + +@pytest.mark.parametrize( + "model", + [ + "bedrock/converse/us.anthropic.claude-haiku-4-5", + "bedrock/converse/us.anthropic.claude-sonnet-4-5", + ], +) +def test_adaptive_thinking_translated_to_legacy_on_pre_46_converse(model): + """Raw thinking={type: adaptive} from callers like Claude Code must be + translated to legacy thinking={type: enabled, budget_tokens} for pre-4.6 + models on Bedrock Converse rather than forwarded as-is and rejected.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "adaptive"}, "max_tokens": 8192}, + optional_params={}, + model=model, + drop_params=False, + ) + + thinking = optional_params.get("thinking") + assert thinking is not None + assert thinking["type"] == "enabled" + assert isinstance(thinking.get("budget_tokens"), int) + assert thinking["budget_tokens"] < 8192 + + +@pytest.mark.parametrize( + "model", + [ + "bedrock/converse/us.anthropic.claude-opus-4-7", + "bedrock/converse/us.anthropic.claude-sonnet-4-6", + ], +) +def test_adaptive_thinking_passes_through_on_46_plus_converse(model): + """thinking={type: adaptive} must be forwarded unchanged for 4.6+ models + that natively support adaptive thinking.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "adaptive"}, "max_tokens": 8192}, + optional_params={}, + model=model, + drop_params=False, + ) + + assert optional_params.get("thinking") == {"type": "adaptive"} + + +def test_adaptive_thinking_dropped_when_max_tokens_too_small_converse(): + """When max_tokens can't fit even the minimum thinking budget, the raw + adaptive block must be dropped entirely rather than translated, so the + Bedrock Converse request still succeeds.""" + from litellm.constants import ANTHROPIC_MIN_THINKING_BUDGET_TOKENS + + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={ + "thinking": {"type": "adaptive"}, + "max_tokens": ANTHROPIC_MIN_THINKING_BUDGET_TOKENS, + }, + optional_params={}, + model="bedrock/converse/us.anthropic.claude-sonnet-4-5", + drop_params=False, + ) + + assert "thinking" not in optional_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() diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index d389b54b3f1..151c51f1ca0 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -44,6 +44,60 @@ class TestOpenAIResponsesAPIConfig: # The function should return the params unchanged assert result == test_params + @pytest.mark.parametrize("max_output_tokens", [1, 15]) + def test_map_openai_params_clamps_max_output_tokens_below_minimum(self, max_output_tokens): + """OpenAI's Responses API rejects max_output_tokens < 16. + + Claude Code (via the Anthropic Messages -> Responses adapter) sends a + max_tokens=1 warmup probe when running `/model`, which produced: + "Invalid 'max_output_tokens': integer below minimum value. + Expected a value >= 16, but got 1 instead." + Clamp anything below the minimum up to 16 instead of erroring. + """ + result = self.config.map_openai_params( + response_api_optional_params={"max_output_tokens": max_output_tokens}, + model=self.model, + drop_params=False, + ) + + assert result["max_output_tokens"] == 16 + + def test_map_openai_params_preserves_max_output_tokens_at_or_above_minimum(self): + """Values already >= 16 must pass through untouched.""" + result = self.config.map_openai_params( + response_api_optional_params={"max_output_tokens": 256}, + model=self.model, + drop_params=False, + ) + + assert result["max_output_tokens"] == 256 + + def test_map_openai_params_leaves_max_output_tokens_absent(self): + """A request without max_output_tokens must not gain the key.""" + result = self.config.map_openai_params( + response_api_optional_params={"input": "hi"}, + model=self.model, + drop_params=False, + ) + + assert "max_output_tokens" not in result + + @pytest.mark.parametrize( + "value, expected", + [ + (1, 16), + (15, 16), + (16, 16), + (17, 17), + (256, 256), + (None, None), + ], + ) + def test_enforce_min_max_output_tokens(self, value, expected): + """Below the minimum clamps to 16; the boundary, larger values, and None + are returned unchanged so no previously-valid request regresses.""" + assert self.config._enforce_min_max_output_tokens(value) == expected + def validate_responses_api_request_params(self, params, expected_fields): """ Validate that the params dict has the expected structure of ResponsesAPIRequestParams diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 51f13affa3f..40f9f4e7910 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -880,6 +880,12 @@ def test_vertex_ai_usage_metadata_with_image_tokens_in_prompt(): ) +def test_map_response_modalities_video(): + """The video modality maps to VIDEO instead of MODALITY_UNSPECIFIED, which Gemini rejects.""" + v = VertexGeminiConfig() + assert v.map_response_modalities(["text", "video"]) == ["TEXT", "VIDEO"] + + def test_vertex_ai_usage_metadata_accumulates_duplicate_modalities(): """Ensure _calculate_usage accumulates repeated modality entries.""" v = VertexGeminiConfig() 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..6f132aaae9c 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,38 @@ 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 _wrapped_user_lookup_error(original: BaseException) -> ValueError: + """Reproduce get_user_object's real exception contract (litellm/proxy/auth/auth_checks.py): it + catches every DB failure in a broad ``except`` and re-raises a bare ``ValueError``, so the + original error (a missing-user Exception or a real outage) survives only as ``__context__``. + Injecting a raw ConnectionError/Exception instead would exercise a shape production never + produces and let a chain-blind outage classifier pass. That wrapping fidelity is itself pinned by + test_get_user_object_wraps_db_outage_as_valueerror_preserving_context in test_auth_checks.""" + try: + raise original + except BaseException: + try: + raise ValueError(f"User doesn't exist in db. Got error - {original}") + except ValueError as wrapped: + return wrapped + @staticmethod def _mcp_request(path="/mcp/bridge_delegate_server"): """A minimal ``Request`` for direct ``_admit_dcr_bridge_delegate`` calls, mirroring how @@ -5060,6 +5099,155 @@ 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}, + 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() + (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_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 catches the missing row and re-raises a bare ValueError (it does not return None + on the production path), so the reload must fail closed rather than let it propagate as an opaque + 500, and must not mistake the wrapped ValueError for a DB outage. Regression for the missing-user + path surfacing as a 500.""" + 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(side_effect=self._wrapped_user_lookup_error(Exception())), + ): + 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_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. get_user_object wraps the outage in a bare ValueError, so this + exercises the chain-aware classifier; a raw ConnectionError would falsely pass even the old + chain-blind check because it is an OSError. 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=self._wrapped_user_lookup_error(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 + 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/faults/test_classify.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_classify.py new file mode 100644 index 00000000000..dc20d664a53 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_classify.py @@ -0,0 +1,147 @@ +"""Classification matrix for upstream OAuth/DCR rejections: who is blamed depends only on the §5.2 +code and whose credentials the gateway presented, never on the upstream's HTTP status.""" + +import httpx + +from litellm.proxy._experimental.mcp_server.faults.classify import ( + classify_upstream_dcr_rejection, + classify_upstream_token_rejection, +) +from litellm.proxy._experimental.mcp_server.faults.types import ( + CallerRejected, + GatewayRejected, + UpstreamProtocolFault, + UpstreamReportedFault, +) + + +def _response(status_code: int, *, json_body: object = None, text_body: str = "", headers: dict = None) -> httpx.Response: + request = httpx.Request("POST", "https://idp.example.com/token") + if json_body is not None: + return httpx.Response(status_code, json=json_body, request=request) + return httpx.Response(status_code, text=text_body, headers=headers or {}, request=request) + + +def test_caller_fault_code_classifies_as_caller_rejected_regardless_of_status(): + fault = classify_upstream_token_rejection( + _response(500, json_body={"error": "invalid_grant", "error_description": "Code expired."}), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, CallerRejected) + assert fault.code == "invalid_grant" + assert fault.description == "Code expired." + + +def test_credential_code_with_gateway_stored_credentials_indicts_gateway(): + fault = classify_upstream_token_rejection( + _response(401, json_body={"error": "invalid_client", "error_description": "not found"}), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, GatewayRejected) + assert fault.code == "invalid_client" + + +def test_credential_code_with_caller_supplied_credentials_stays_caller_fault(): + fault = classify_upstream_token_rejection( + _response(401, json_body={"error": "invalid_client"}), + credential_source="caller_supplied", + log_context="srv", + ) + assert isinstance(fault, CallerRejected) + assert fault.code == "invalid_client" + + +def test_unknown_code_relays_as_caller_rejected(): + fault = classify_upstream_token_rejection( + _response(400, json_body={"error": "slow_down", "error_description": "Polling too fast."}), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, CallerRejected) + assert fault.code == "slow_down" + + +def test_body_without_error_field_is_protocol_fault(): + fault = classify_upstream_token_rejection( + _response(404, text_body="not here"), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, UpstreamProtocolFault) + assert fault.note == "upstream token endpoint returned HTTP 404" + + +def test_unreadable_body_is_protocol_fault_not_exception(): + unreadable = httpx.Response( + 400, + stream=httpx.ByteStream(b"\x1f\x8bnot-gzip"), + headers={"content-encoding": "gzip"}, + request=httpx.Request("POST", "https://idp.example.com/token"), + ) + fault = classify_upstream_token_rejection(unreadable, credential_source="gateway_stored", log_context="srv") + assert isinstance(fault, UpstreamProtocolFault) + + +def test_wire_fields_are_bounded(): + fault = classify_upstream_token_rejection( + _response(400, json_body={"error": "invalid_request", "error_description": "x" * 5000}), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, CallerRejected) + assert len(fault.description) == 500 + + +def test_dcr_rejection_with_rfc7591_code_is_caller_rejected(): + fault = classify_upstream_dcr_rejection( + _response(400, json_body={"error": "invalid_redirect_uri", "error_description": "not allowed"}), + log_context="srv", + ) + assert isinstance(fault, CallerRejected) + assert fault.code == "invalid_redirect_uri" + + +def test_dcr_rejection_without_code_is_protocol_fault(): + fault = classify_upstream_dcr_rejection(_response(500, text_body="trace"), log_context="srv") + assert isinstance(fault, UpstreamProtocolFault) + assert fault.note == "upstream registration failed with HTTP 500" + + +def test_upstream_self_blame_codes_stay_upstream_faults(): + fault = classify_upstream_token_rejection( + _response(400, json_body={"error": "server_error", "error_description": "boom"}), + credential_source="caller_supplied", + log_context="srv", + ) + assert isinstance(fault, UpstreamReportedFault) + assert fault.code == "server_error" + + +def test_temporarily_unavailable_is_upstream_fault(): + fault = classify_upstream_token_rejection( + _response(503, json_body={"error": "temporarily_unavailable"}), + credential_source="gateway_stored", + log_context="srv", + ) + assert isinstance(fault, UpstreamReportedFault) + assert fault.code == "temporarily_unavailable" + + +def test_invalid_target_is_gateway_fault_even_with_caller_credentials(): + fault = classify_upstream_token_rejection( + _response(400, json_body={"error": "invalid_target"}), + credential_source="caller_supplied", + log_context="srv", + ) + assert isinstance(fault, GatewayRejected) + assert fault.code == "invalid_target" + + +def test_dcr_server_error_code_is_not_blamed_on_caller(): + fault = classify_upstream_dcr_rejection( + _response(500, json_body={"error": "server_error"}), + log_context="srv", + ) + assert isinstance(fault, UpstreamReportedFault) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_render_oauth.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_render_oauth.py new file mode 100644 index 00000000000..78513e315a7 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_render_oauth.py @@ -0,0 +1,96 @@ +"""Rendering contract: status, wire code, and prose all derive from the fault tag, so a caller-fault +code can never ship on a server-fault status and gateway-side faults never carry provider prose.""" + +import json + +from litellm.proxy._experimental.mcp_server.faults.render_oauth import ( + dcr_fault_detail, + render_token_fault, +) +from litellm.proxy._experimental.mcp_server.faults.types import ( + CallerRejected, + GatewayRejected, + UpstreamProtocolFault, + UpstreamReportedFault, +) + + +def test_caller_rejected_renders_code_derived_status(): + response = render_token_fault(CallerRejected(code="invalid_grant", description="Code expired.")) + assert response.status_code == 400 + assert json.loads(response.body) == {"error": "invalid_grant", "error_description": "Code expired."} + assert response.headers["cache-control"] == "no-store" + + +def test_caller_rejected_invalid_client_renders_401(): + response = render_token_fault(CallerRejected(code="invalid_client")) + assert response.status_code == 401 + assert json.loads(response.body) == {"error": "invalid_client"} + + +def test_caller_rejected_includes_error_uri_only_when_present(): + response = render_token_fault( + CallerRejected(code="invalid_scope", description="bad scope", error_uri="https://idp.example.com/e") + ) + assert json.loads(response.body) == { + "error": "invalid_scope", + "error_description": "bad scope", + "error_uri": "https://idp.example.com/e", + } + + +def test_gateway_rejected_renders_502_with_gateway_prose(): + response = render_token_fault(GatewayRejected(code="invalid_client")) + assert response.status_code == 502 + body = json.loads(response.body) + assert body["error"] == "server_error" + assert "invalid_client" in body["error_description"] + assert "client_id and client_secret" in body["error_description"] + + +def test_gateway_invalid_target_prose_names_resource_indicators(): + response = render_token_fault(GatewayRejected(code="invalid_target")) + body = json.loads(response.body) + assert response.status_code == 502 + assert "RFC 8707" in body["error_description"] + + +def test_protocol_fault_renders_502_note(): + response = render_token_fault(UpstreamProtocolFault(note="upstream token endpoint returned HTTP 503")) + assert response.status_code == 502 + assert json.loads(response.body) == { + "error": "server_error", + "error_description": "upstream token endpoint returned HTTP 503", + } + + +def test_dcr_caller_rejection_is_400_per_rfc7591_regardless_of_upstream_status(): + status_code, detail = dcr_fault_detail(CallerRejected(code="invalid_client_metadata", description="bad grant types")) + assert status_code == 400 + assert detail == "invalid_client_metadata: bad grant types" + + +def test_dcr_protocol_fault_is_502(): + status_code, detail = dcr_fault_detail(UpstreamProtocolFault(note="upstream registration failed with HTTP 500")) + assert status_code == 502 + assert detail == "upstream registration failed with HTTP 500" + + +def test_upstream_reported_server_error_renders_502_with_matching_code(): + response = render_token_fault(UpstreamReportedFault(code="server_error")) + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" + + +def test_upstream_reported_temporarily_unavailable_renders_503_with_matching_code(): + response = render_token_fault(UpstreamReportedFault(code="temporarily_unavailable")) + assert response.status_code == 503 + body = json.loads(response.body) + assert body["error"] == "temporarily_unavailable" + assert "retry" in body["error_description"] + + +def test_dcr_upstream_reported_fault_maps_to_5xx(): + status_code, detail = dcr_fault_detail(UpstreamReportedFault(code="server_error")) + assert status_code == 502 + assert "internal error" in detail 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..753a3d6a942 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 @@ -15,10 +15,14 @@ from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, BridgeEnvelopeInvalid, + BridgeRefreshInvalid, + BridgeRefreshOpened, NotBridgeEnvelope, + build_bridge_refresh_token_response, build_bridge_token_response, envelope_keys_from_master_key, is_bridge_envelope_shaped, + open_bridge_refresh_envelope, resolve_bridge_envelope, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( @@ -26,15 +30,17 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import EnvelopeIdentity, EnvelopeKeys, EnvelopeTooLarge, + RefreshCredential, 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 @@ -48,6 +54,76 @@ def _sealed_token(keys: EnvelopeKeys, now: datetime = _NOW, identity: EnvelopeId return sealed.token.get_secret_value() +_UPSTREAM_REFRESH = "upstream-refresh-do-not-leak-9b2c" + + +def _sealed_refresh(keys: EnvelopeKeys, now: datetime = _NOW, identity: EnvelopeIdentity = _IDENTITY) -> str: + sealed = build_bridge_refresh_token_response( + identity, RefreshCredential(refresh_token=SecretStr(_UPSTREAM_REFRESH)), keys, now + ) + assert isinstance(sealed, SealedEnvelope) + return sealed.token.get_secret_value() + + +def test_open_bridge_refresh_envelope_round_trips_identity_and_refresh(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = open_bridge_refresh_envelope(_sealed_refresh(keys), keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeRefreshOpened) + assert result.identity == _IDENTITY + assert result.refresh.refresh_token.get_secret_value() == _UPSTREAM_REFRESH + + +def test_open_bridge_refresh_envelope_strips_bearer_scheme(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = open_bridge_refresh_envelope(f"Bearer {_sealed_refresh(keys)}", keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeRefreshOpened) + + +def test_open_bridge_refresh_envelope_rejects_wrong_server(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = open_bridge_refresh_envelope(_sealed_refresh(keys), keys, _NOW, "a-different-server") + assert isinstance(result, BridgeRefreshInvalid) + + +def test_open_bridge_refresh_envelope_rejects_non_refresh_bearers(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + # an access envelope is not a refresh envelope; a raw upstream refresh token is not one either + assert isinstance(open_bridge_refresh_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID), BridgeRefreshInvalid) + assert isinstance(open_bridge_refresh_envelope("raw-refresh-token", keys, _NOW, _SERVER_ID), BridgeRefreshInvalid) + + +def test_open_bridge_refresh_envelope_rejects_under_wrong_master_key(): + minted = envelope_keys_from_master_key(_MASTER_KEY) + other = envelope_keys_from_master_key(_MASTER_KEY + "-rotated") + result = open_bridge_refresh_envelope(_sealed_refresh(minted), other, _NOW, _SERVER_ID) + assert isinstance(result, BridgeRefreshInvalid) + + +def test_refresh_envelope_is_never_admitted_at_the_tool_call_edge(): + """A refresh envelope must never authenticate a tool call. The admission edge engages the bridge arm + for it (is_bridge_envelope_shaped is true for either envelope kind), and the consumer rejects it as + BridgeEnvelopeInvalid, which admission fails closed (401): a refresh credential is only ever + presented back to the token endpoint.""" + keys = envelope_keys_from_master_key(_MASTER_KEY) + refresh = _sealed_refresh(keys) + assert is_bridge_envelope_shaped(refresh) is True + assert is_bridge_envelope_shaped(f"Bearer {refresh}") is True + result = resolve_bridge_envelope(refresh, keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeInvalid) + + +def test_refresh_jwt_wearing_the_access_prefix_is_rejected_at_the_edge(): + """Belt-and-suspenders against a swapped wire prefix: a refresh JWT re-prefixed as an access envelope + opens far enough to hit the signed kind claim, which rejects it, so admission fails closed rather + than forwarding a refresh credential's contents upstream.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import REFRESH_ENVELOPE_PREFIX + + keys = envelope_keys_from_master_key(_MASTER_KEY) + swapped = ENVELOPE_PREFIX + _sealed_refresh(keys).removeprefix(REFRESH_ENVELOPE_PREFIX) + result = resolve_bridge_envelope(swapped, keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeInvalid) + + def test_key_derivation_is_deterministic(): assert envelope_keys_from_master_key(_MASTER_KEY) == envelope_keys_from_master_key(_MASTER_KEY) @@ -138,7 +214,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 +231,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..ae196c9080b 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 @@ -24,6 +24,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ENVELOPE_PREFIX, MAX_ENVELOPE_BYTES, MAX_ENVELOPE_TTL_SECONDS, + MAX_REFRESH_ENVELOPE_TTL_SECONDS, + REFRESH_ENVELOPE_PREFIX, BadSignature, DecryptFailed, EnvelopeIdentity, @@ -33,11 +35,18 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import MalformedPayload, NotAnEnvelope, OpenedEnvelope, + OpenedRefreshEnvelope, + RefreshCredential, SealedEnvelope, UpstreamTokenGrant, is_envelope, + is_refresh_envelope, + key_hash_identity, mint_envelope, + mint_refresh_envelope, open_envelope, + open_refresh_envelope, + user_identity, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value @@ -51,7 +60,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,17 +146,99 @@ 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", "kind", "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["kind"] == "access" 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) +def _refresh_credential() -> RefreshCredential: + return RefreshCredential(refresh_token=SecretStr(_REFRESH_TOKEN), scope="read:tools", expires_in=None) + + +def _sealed_refresh_token(refresh: RefreshCredential | None = None, keys: EnvelopeKeys = _KEYS) -> str: + sealed = mint_refresh_envelope(_IDENTITY, refresh or _refresh_credential(), keys, _NOW) + assert isinstance(sealed, SealedEnvelope) + return sealed.token.get_secret_value() + + +def test_refresh_envelope_round_trips_identity_and_refresh_token(): + token = _sealed_refresh_token() + assert is_refresh_envelope(token) + assert not is_envelope(token) + opened = open_refresh_envelope(token, _KEYS, _NOW) + assert isinstance(opened, OpenedRefreshEnvelope) + assert opened.identity == _IDENTITY + assert opened.refresh.refresh_token.get_secret_value() == _REFRESH_TOKEN + assert opened.refresh.scope == "read:tools" + + +def test_refresh_envelope_ttl_is_min_of_upstream_refresh_lifetime_and_cap(): + short = mint_refresh_envelope( + _IDENTITY, RefreshCredential(refresh_token=SecretStr("r"), expires_in=120), _KEYS, _NOW + ) + assert isinstance(short, SealedEnvelope) + assert short.expires_at == _NOW + timedelta(seconds=120) + capped = mint_refresh_envelope( + _IDENTITY, + RefreshCredential(refresh_token=SecretStr("r"), expires_in=MAX_REFRESH_ENVELOPE_TTL_SECONDS + 86400), + _KEYS, + _NOW, + ) + assert isinstance(capped, SealedEnvelope) + assert capped.expires_at == _NOW + timedelta(seconds=MAX_REFRESH_ENVELOPE_TTL_SECONDS) + default = mint_refresh_envelope(_IDENTITY, RefreshCredential(refresh_token=SecretStr("r")), _KEYS, _NOW) + assert isinstance(default, SealedEnvelope) + assert default.expires_at == _NOW + timedelta(seconds=MAX_REFRESH_ENVELOPE_TTL_SECONDS) + + +def test_access_and_refresh_envelopes_do_not_cross_open(): + access = _sealed_token(_full_grant()) + refresh = _sealed_refresh_token() + # each opener rejects the other kind's prefix outright + assert isinstance(open_refresh_envelope(access, _KEYS, _NOW), NotAnEnvelope) + assert isinstance(open_envelope(refresh, _KEYS, _NOW), NotAnEnvelope) + + +def test_prefix_swap_is_rejected_by_the_signed_kind_claim(): + # the wire prefix is not signed, so swap it; the signed kind claim must still reject the cross-use + refresh = _sealed_refresh_token() + swapped_to_access = ENVELOPE_PREFIX + refresh.removeprefix(REFRESH_ENVELOPE_PREFIX) + assert isinstance(open_envelope(swapped_to_access, _KEYS, _NOW), MalformedPayload) + access = _sealed_token(_full_grant()) + swapped_to_refresh = REFRESH_ENVELOPE_PREFIX + access.removeprefix(ENVELOPE_PREFIX) + assert isinstance(open_refresh_envelope(swapped_to_refresh, _KEYS, _NOW), MalformedPayload) + + +def test_refresh_envelope_total_over_hostile_input(): + token = _sealed_refresh_token() + # expired against the injected clock + assert isinstance( + open_refresh_envelope(token, _KEYS, _NOW + timedelta(seconds=MAX_REFRESH_ENVELOPE_TTL_SECONDS)), Expired + ) + # wrong signing key + assert isinstance(open_refresh_envelope(token, _WRONG_SIGNING, _NOW), BadSignature) + # right signature, wrong encryption key + assert isinstance(open_refresh_envelope(token, _WRONG_ENCRYPTION, _NOW), DecryptFailed) + # not an envelope at all + assert isinstance(open_refresh_envelope("raw-upstream-refresh-token", _KEYS, _NOW), NotAnEnvelope) + + +def test_refresh_envelope_never_leaks_the_refresh_token_in_plaintext(): + token = _sealed_refresh_token() + assert _REFRESH_TOKEN not in token + claims = jwt.decode(token.removeprefix(REFRESH_ENVELOPE_PREFIX), options={"verify_signature": False}) + assert claims["kind"] == "refresh" + assert _REFRESH_TOKEN not in json.dumps(claims) + + @pytest.mark.parametrize( "expires_in, expected_ttl", [ @@ -226,11 +317,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 +554,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 +567,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 +591,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 c55a631c7b3..d1aceffa968 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 @@ -4274,6 +4274,9 @@ async def test_register_bridge_relay_surfaces_upstream_error_not_500(): error_response = MagicMock() error_response.status_code = 400 error_response.text = '{"error":"invalid_redirect_uri","error_description":"redirect_uri not allowed"}' + error_response.json = MagicMock( + return_value={"error": "invalid_redirect_uri", "error_description": "redirect_uri not allowed"} + ) error_response.raise_for_status = MagicMock( side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response) ) @@ -4307,9 +4310,10 @@ async def test_register_bridge_relay_surfaces_upstream_error_not_500(): @pytest.mark.asyncio -async def test_register_non_bridge_upstream_error_still_raises_500(): - """Non-bridge DCR keeps its pre-change behavior: raise_for_status propagates so the flag-off - contract is byte-identical; only the bridge relay arm relays the upstream status.""" +async def test_register_non_bridge_upstream_error_relays_status_not_500(): + """A non-bridge DCR rejection must relay the upstream status and RFC 7591 error body just like + the bridge relay arm; a raw HTTPStatusError would escape to the global handler and surface as an + opaque 500 that hides the real reason from the create-flow UI.""" import httpx from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -4319,6 +4323,7 @@ async def test_register_non_bridge_upstream_error_still_raises_500(): error_response = MagicMock() error_response.status_code = 400 error_response.text = '{"error":"invalid_client_metadata"}' + error_response.json = MagicMock(return_value={"error": "invalid_client_metadata"}) error_response.raise_for_status = MagicMock( side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response) ) @@ -4338,7 +4343,7 @@ async def test_register_non_bridge_upstream_error_still_raises_500(): return_value=False, ), ): - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(HTTPException) as exc: await register_client_with_server( request=_bridge_mock_request(), mcp_server=oauth2_server, @@ -4348,6 +4353,9 @@ async def test_register_non_bridge_upstream_error_still_raises_500(): token_endpoint_auth_method=None, ) + assert exc.value.status_code == 400 + assert "invalid_client_metadata" in str(exc.value.detail) + @pytest.mark.asyncio async def test_register_bridge_relay_never_persists(): @@ -4362,6 +4370,1247 @@ async def test_register_bridge_relay_never_persists(): mock_persist.assert_not_called() +_BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef" + + +async def _exchange_for_bridge_server(server, upstream_body, key_hash, code="auth-code", fake_client_out=None): + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _ResolvedKey + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + + fake_http_response = MagicMock() + fake_http_response.json.return_value = upstream_body + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + # The mint consumes _resolve_active_litellm_key's tagged result: an active key resolves to a + # _ResolvedKey carrying its hash; a request with no usable credential resolves to "no_active_key". + resolution = _ResolvedKey(key_hash=key_hash, key=MagicMock()) if key_hash is not None else "no_active_key" + key_resolver = AsyncMock(return_value=resolution) + if fake_client_out is not None: + fake_client_out["client"] = fake_http_client + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._resolve_active_litellm_key", + new=key_resolver, + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="authorization_code", + code=code, + redirect_uri="https://claude.ai/api/mcp/auth_callback", + client_id="dcr-client-123", + client_secret=None, + code_verifier="verifier", + ) + 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() + return response + + +@pytest.mark.asyncio +async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token(): + """A dcr_bridge oauth_delegate token exchange returns a gateway-bound envelope, not the raw + upstream token: the response access_token opens (under the same master-key-derived keys and the + server_id) to the caller's identity and the upstream Authorization, and the raw upstream token + never appears in the bearer the client receives.""" + from datetime import datetime, timezone + + 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) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + token = body["access_token"] + assert body["token_type"] == "Bearer" + assert body["expires_in"] > 0 + assert token.startswith("llm_env_") + assert "UPSTREAM-SECRET-TOKEN" not in token + assert "refresh_token" not in body + + 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 == "key_hash" + assert opened.identity.subject == "hashed-litellm-key-77" + 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 + identity-less envelope. It returns an RFC 6749 §5.2-shaped invalid_request (error at the top + level, not wrapped in detail) BEFORE exchanging the upstream code, so the single-use code is not + burned and the client can retry.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + 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, fake_client_out=captured) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + # identity resolution failed first, so the upstream single-use code was never exchanged (not burned) + captured["client"].post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_envelope_too_large_upstream_token_is_502(): + """An upstream token too large to seal into the envelope is an upstream-payload condition, so the + mint surfaces a 502 (as an RFC 6749 §5.2 error body, not a raised HTTPException) rather than a 500: + build_bridge_token_response returns EnvelopeTooLarge as a value, _finish_bridge_mint returns the + "too_large" failure, and _bridge_mint_error_response maps it to a truthful status.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "x" * 40000, "token_type": "Bearer", "expires_in": 3600} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" + + +@pytest.mark.asyncio +async def test_bridge_access_envelope_never_carries_upstream_refresh_token(): + """The upstream refresh token is never sealed into the ACCESS envelope, the bearer forwarded upstream + on every tool call: the opened access grant carries no refresh token even when the upstream returned + one, and the raw refresh token never appears in the access envelope. It rides only in the separate + refresh envelope returned as the response's refresh_token, encrypted, never in plaintext.""" + from datetime import datetime, timezone + + 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 ( + OpenedEnvelope, + open_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = { + "access_token": "UP", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "UPSTREAM-REFRESH", + } + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + assert "UPSTREAM-REFRESH" not in body["access_token"] + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = open_envelope(body["access_token"], keys, datetime.now(timezone.utc)) + assert isinstance(opened, OpenedEnvelope) + assert opened.grant.refresh_token is None + # the refresh token rides only in the separate, encrypted refresh envelope, never in plaintext + assert body["refresh_token"].startswith("llm_refresh_") + assert "UPSTREAM-REFRESH" not in body["refresh_token"] + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_with_non_envelope_is_invalid_grant_before_upstream(): + """A bridge oauth_delegate client only ever holds a refresh envelope, never a raw upstream refresh + token, so a refresh_token grant carrying a bare (non-envelope) value is invalid_grant, rejected in + _prepare_bridge_refresh BEFORE any upstream exchange. Rejecting before the exchange means a bad + refresh request can never consume or rotate an upstream refresh token.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock() + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="dcr-client-123", + client_secret=None, + code_verifier=None, + refresh_token="client-refresh-token", + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" + fake_http_client.post.assert_not_called() + + +def _mint_test_refresh_envelope( + server_id="bridge_srv", key_hash="hashed-litellm-key-77", upstream_refresh="UPSTREAM-REFRESH", identity=None, + scope=None, +): + """Mint a refresh envelope the way the producer does, for driving the refresh_token grant in tests. + Defaults to a key_hash subject; pass ``identity`` to seal a specific subject (e.g. a user_id), and + ``scope`` to seal the scope to re-request on refresh.""" + from datetime import datetime, timezone + + from pydantic import SecretStr + + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + build_bridge_refresh_token_response, + envelope_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + RefreshCredential, + SealedEnvelope, + key_hash_identity, + ) + + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + identity = identity if identity is not None else key_hash_identity(server_id=server_id, key_hash=key_hash) + sealed = build_bridge_refresh_token_response( + identity, RefreshCredential(refresh_token=SecretStr(upstream_refresh), scope=scope), keys, + datetime.now(timezone.utc), + ) + assert isinstance(sealed, SealedEnvelope) + return sealed.token.get_secret_value() + + +async def _refresh_for_bridge_server( + server, refresh_envelope_value, upstream_body, revalidate_result=None, fake_client_out=None +): + """Drive a refresh_token grant for a bridge server: the client presents ``refresh_envelope_value``, + the sealed subject re-validates to ``revalidate_result`` (``None`` when the key or user is still + active, or a failure literal like "no_active_key" when revoked/deactivated), and the upstream returns + ``upstream_body``. Patching the single subject-revalidation dispatch covers both a key_hash and a + user_id refresh envelope. Returns the response; the captured client exposes the POST call so a test + can assert what refresh token was actually sent upstream.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + + fake_http_response = MagicMock() + fake_http_response.json.return_value = upstream_body + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + if fake_client_out is not None: + fake_client_out["client"] = fake_http_client + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._revalidate_active_subject", + new=AsyncMock(return_value=revalidate_result), + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + return await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="dcr-client-123", + client_secret=None, + code_verifier=None, + refresh_token=refresh_envelope_value, + ) + + +@pytest.mark.asyncio +async def test_bridge_mint_returns_refresh_envelope_that_opens_to_upstream_refresh(): + """When the upstream returns a refresh token, the authorization_code mint returns a refresh envelope + alongside the access envelope. The refresh envelope is a distinct llm_refresh_ credential that opens + (under the same keys and server_id) to the upstream refresh token, so the client can renew later.""" + from datetime import datetime, timezone + + 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 ( + OpenedRefreshEnvelope, + open_refresh_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 3600, "refresh_token": "R-UP"} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + refresh_env = body["refresh_token"] + assert refresh_env.startswith("llm_refresh_") + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = open_refresh_envelope(refresh_env, keys, datetime.now(timezone.utc)) + assert isinstance(opened, OpenedRefreshEnvelope) + assert opened.identity.server_id == server.server_id + assert opened.refresh.refresh_token.get_secret_value() == "R-UP" + + +@pytest.mark.asyncio +async def test_bridge_mint_omits_refresh_envelope_when_upstream_has_no_refresh(): + """No refresh envelope is issued when the upstream returns no refresh token, so the response carries + only the access envelope; the client re-authenticates at access expiry (nothing to renew with).""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 3600} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + assert body["access_token"].startswith("llm_env_") + assert "refresh_token" not in body + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_sends_unwrapped_upstream_token_and_renews(): + """A refresh_token grant carrying a valid refresh envelope renews: the exchange unwraps the envelope + and sends the REAL upstream refresh token upstream (never the envelope), then returns a fresh access + envelope. This is the flow that lets the client renew without re-authenticating.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="UPSTREAM-REFRESH") + upstream = {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600} + captured: dict = {} + response = await _refresh_for_bridge_server(server, refresh_env, upstream, None, fake_client_out=captured) + + assert response.status_code == 200 + body = json.loads(response.body) + assert body["access_token"].startswith("llm_env_") + # the upstream exchange received the unwrapped upstream refresh token, never the client's envelope + sent = captured["client"].post.call_args.kwargs["data"] + assert sent["grant_type"] == "refresh_token" + assert sent["refresh_token"] == "UPSTREAM-REFRESH" + assert not sent["refresh_token"].startswith("llm_refresh_") + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_rotates_refresh_envelope_wrapping_new_upstream_token(): + """When the upstream rotates the refresh token on renewal, the client receives a new refresh envelope + that wraps the NEW upstream refresh token, so the rotation is carried through faithfully.""" + from datetime import datetime, timezone + + 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 ( + OpenedRefreshEnvelope, + open_refresh_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="OLD-UP-REFRESH") + upstream = { + "access_token": "NEW-ACCESS", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "NEW-UP-REFRESH", + } + response = await _refresh_for_bridge_server(server, refresh_env, upstream, None) + + body = json.loads(response.body) + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = open_refresh_envelope(body["refresh_token"], keys, datetime.now(timezone.utc)) + assert isinstance(opened, OpenedRefreshEnvelope) + assert opened.refresh.refresh_token.get_secret_value() == "NEW-UP-REFRESH" + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_with_revoked_key_is_invalid_grant_before_upstream(): + """A valid refresh envelope whose sealed litellm key has since been revoked cannot keep refreshing: + the reload gate reports no_active_key and the refresh is invalid_grant, returned BEFORE the upstream + exchange so the upstream refresh token is never consumed. Revocation kills renewal.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope(server_id=server.server_id) + captured: dict = {} + response = await _refresh_for_bridge_server( + server, refresh_env, {"access_token": "NEW"}, "no_active_key", fake_client_out=captured + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" + captured["client"].post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_refresh_envelope_for_another_server_is_invalid_grant(): + """A refresh envelope minted for one server cannot renew against another: the sealed server_id must + match the server the refresh targets, so a cross-server refresh envelope is invalid_grant and never + reaches the upstream exchange.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + foreign_env = _mint_test_refresh_envelope(server_id="some-other-server") + captured: dict = {} + response = await _refresh_for_bridge_server( + server, foreign_env, {"access_token": "NEW"}, None, fake_client_out=captured + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" + captured["client"].post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_renews_a_user_subject_envelope(): + """The interactive SSO client mints a user_id-subject envelope, so its refresh envelope carries a + user subject too. Renewing it re-validates the user (still active here), unwraps the upstream refresh + token, and returns a fresh access envelope that opens back to the same user_id subject; the upstream + exchange received the real upstream refresh token, not the client's envelope.""" + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + BridgeEnvelopeAdmitted, + envelope_keys_from_master_key, + resolve_bridge_envelope, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import user_identity + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + user_env = _mint_test_refresh_envelope( + identity=user_identity(server_id=server.server_id, user_id="sso-user-42"), upstream_refresh="UP-REFRESH-USER" + ) + upstream = {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600} + captured: dict = {} + response = await _refresh_for_bridge_server(server, user_env, upstream, None, fake_client_out=captured) + + assert response.status_code == 200 + body = json.loads(response.body) + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = resolve_bridge_envelope(body["access_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 captured["client"].post.call_args.kwargs["data"]["refresh_token"] == "UP-REFRESH-USER" + + +@pytest.mark.asyncio +async def test_bridge_refresh_re_requests_the_sealed_scope_when_client_omits_it(): + """A DCR/MCP client omits scope on the refresh request, so the gateway must re-request the scope sealed + at mint; dropping it lets a stricter upstream narrow the renewed token. The upstream POST must carry + the sealed scope even though the client sent none. Regression for the dropped sealed refresh scope.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope( + server_id=server.server_id, upstream_refresh="UP-REFRESH", scope="read:tools write:tools" + ) + captured: dict = {} + response = await _refresh_for_bridge_server( + server, refresh_env, {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600}, None, + fake_client_out=captured, + ) + + assert response.status_code == 200 + assert captured["client"].post.call_args.kwargs["data"]["scope"] == "read:tools write:tools" + + +@pytest.mark.asyncio +async def test_bridge_refresh_re_seals_scope_when_upstream_omits_it_so_the_chain_keeps_it(): + """RFC 6749 5.1 lets an upstream omit scope in a refresh response when it is unchanged. The re-minted + refresh envelope must still seal the scope that was requested, otherwise the NEXT refresh loses it and + a stricter upstream could narrow the token. The returned refresh envelope carries the scope even though + the upstream response had none, and a second refresh off it still re-requests the scope.""" + from datetime import datetime, timezone + + 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 ( + OpenedRefreshEnvelope, + open_refresh_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope( + server_id=server.server_id, upstream_refresh="UP-1", scope="mcp:read mcp:write" + ) + # the upstream rotates the refresh token but OMITS scope (valid when unchanged) + upstream_no_scope = {"access_token": "NEW", "token_type": "Bearer", "expires_in": 3600, "refresh_token": "UP-2"} + + captured: dict = {} + r1 = await _refresh_for_bridge_server(server, refresh_env, upstream_no_scope, None, fake_client_out=captured) + assert r1.status_code == 200 + assert captured["client"].post.call_args.kwargs["data"]["scope"] == "mcp:read mcp:write" + + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + new_env = json.loads(r1.body)["refresh_token"] + opened = open_refresh_envelope(new_env, keys, datetime.now(timezone.utc)) + assert isinstance(opened, OpenedRefreshEnvelope) + assert opened.refresh.scope == "mcp:read mcp:write" + + captured2: dict = {} + r2 = await _refresh_for_bridge_server(server, new_env, upstream_no_scope, None, fake_client_out=captured2) + assert r2.status_code == 200 + assert captured2["client"].post.call_args.kwargs["data"]["scope"] == "mcp:read mcp:write" + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_with_deactivated_user_is_invalid_grant_before_upstream(): + """A user_id-subject refresh envelope whose user has since been deactivated (SCIM offboarding, or + the user no longer exists) cannot keep refreshing: subject re-validation reports no_active_key and + the refresh is invalid_grant, returned BEFORE the upstream exchange. Revocation kills renewal for the + user subject exactly as it does for the key subject.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import user_identity + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + user_env = _mint_test_refresh_envelope(identity=user_identity(server_id=server.server_id, user_id="gone-user")) + captured: dict = {} + response = await _refresh_for_bridge_server( + server, user_env, {"access_token": "NEW"}, "no_active_key", fake_client_out=captured + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" + captured["client"].post.assert_not_called() + + +@pytest.mark.asyncio +async def test_revalidate_active_subject_dispatches_on_subject_type(): + """Subject re-validation routes a key_hash envelope to the key reload and a user_id envelope to the + user reload, so revocation gates renewal for either identity source through one dispatch point.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _ResolvedKey, + _revalidate_active_subject, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity, user_identity + + with ( + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_key_by_hash", + new=AsyncMock(return_value=_ResolvedKey(key_hash="kh", key=MagicMock())), + ) as key_reload, + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_user_by_id", + new=AsyncMock(return_value=None), + ) as user_reload, + ): + assert await _revalidate_active_subject(key_hash_identity(server_id="s", key_hash="kh")) is None + key_reload.assert_awaited_once_with("kh") + user_reload.assert_not_awaited() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_key_by_hash", + new=AsyncMock(), + ) as key_reload2, + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_user_by_id", + new=AsyncMock(return_value="no_active_key"), + ) as user_reload2, + ): + assert await _revalidate_active_subject(user_identity(server_id="s", user_id="u42")) == "no_active_key" + user_reload2.assert_awaited_once_with("u42") + key_reload2.assert_not_awaited() + + +def test_upstream_refresh_credential_expired_refresh_token_is_not_sealed(): + """An upstream that reports its refresh token already elapsed (refresh_expires_in non-positive) must + not be sealed: _upstream_refresh_credential returns None so the exchange degrades to an access-only + response, mirroring how the access grant refuses an already-elapsed access token rather than capping a + dead token to the full refresh TTL. A live or unspecified lifetime still yields a credential.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _upstream_refresh_credential + + assert _upstream_refresh_credential({"access_token": "A", "refresh_token": "R", "refresh_expires_in": 0}) is None + assert _upstream_refresh_credential({"refresh_token": "R", "refresh_expires_in": -5}) is None + live = _upstream_refresh_credential({"refresh_token": "R", "refresh_expires_in": 1800}) + assert live is not None and live.expires_in == 1800 + unspecified = _upstream_refresh_credential({"refresh_token": "R"}) + assert unspecified is not None and unspecified.expires_in is None + + +@pytest.mark.asyncio +async def test_bridge_refresh_upstream_invalid_grant_maps_to_invalid_grant(): + """When the sealed upstream refresh token has been revoked or expired at the IdP, the upstream returns + 400 invalid_grant. The bridge refresh path maps that to an RFC 6749 invalid_grant response so the OAuth + client re-runs authorization_code, rather than surfacing the opaque upstream error it cannot act on.""" + import httpx + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="LIVE-ENVELOPE-REFRESH") + + error_response = MagicMock() + error_response.status_code = 400 + error_response.text = '{"error": "invalid_grant", "error_description": "refresh token expired"}' + error_response.json = MagicMock(return_value={"error": "invalid_grant", "error_description": "refresh token expired"}) + error_response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response) + ) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=error_response) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._revalidate_active_subject", + new=AsyncMock(return_value=None), + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="dcr-client-123", + client_secret=None, + code_verifier=None, + refresh_token=refresh_env, + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_grant" + + +@pytest.mark.asyncio +async def test_bridge_refresh_upstream_error_detection_parses_json_not_substring(): + """The upstream invalid_grant detection reads the classified RFC 6749 5.2 error code, not a substring + of the body. An upstream error whose code is not invalid_grant (here invalid_client, with the string + invalid_grant only inside error_description) must NOT be mistaken for a dead refresh token: it renders + as the classified upstream rejection rather than triggering a spurious authorization_code re-run.""" + import httpx + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="UP") + + error_response = MagicMock() + error_response.status_code = 400 + error_response.text = '{"error": "invalid_client", "error_description": "this is not an invalid_grant problem"}' + error_response.json = MagicMock( + return_value={"error": "invalid_client", "error_description": "this is not an invalid_grant problem"} + ) + error_response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response) + ) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=error_response) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._revalidate_active_subject", + new=AsyncMock(return_value=None), + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="dcr-client-123", + client_secret=None, + code_verifier=None, + refresh_token=refresh_env, + ) + + assert response.status_code == 401 + body = json.loads(response.body) + assert body["error"] == "invalid_client" + + +@pytest.mark.asyncio +async def test_revalidate_key_subject_revoked_when_owner_scim_deactivated(proxy_globals): + """A key_hash refresh envelope whose key is still active but whose OWNING user was SCIM-deactivated must + fail closed to no_active_key, mirroring how admission's _reject_if_admitted_owner_scim_deactivated + revokes an offboarded owner's key. Without this, an offboarded user keeps renewing a live key.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _ResolvedKey, _revalidate_active_subject + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + resolved = _ResolvedKey(key_hash="kh", key=MagicMock(user_id="offboarded-owner")) + with ( + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_key_by_hash", + new=AsyncMock(return_value=resolved), + ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(return_value=MagicMock(metadata={"scim_active": False})), + ), + ): + result = await _revalidate_active_subject(key_hash_identity(server_id="s", key_hash="kh")) + + assert result == "no_active_key" + + +@pytest.mark.asyncio +async def test_revalidate_key_subject_active_owner_renews_and_missing_owner_fails_open(proxy_globals): + """The key-owner SCIM gate blocks only an explicit scim_active False: an active owner renews (None), and + a missing owner (get_user_object's wrapped ValueError) fails OPEN, since a key may outlive its owner + record and a transient blip must not revoke a live key.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _ResolvedKey, _revalidate_active_subject + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + resolved = _ResolvedKey(key_hash="kh", key=MagicMock(user_id="live-owner")) + identity = key_hash_identity(server_id="s", key_hash="kh") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_key_by_hash", + new=AsyncMock(return_value=resolved), + ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(return_value=MagicMock(metadata={"scim_active": True})), + ), + ): + assert await _revalidate_active_subject(identity) is None + + with ( + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_key_by_hash", + new=AsyncMock(return_value=resolved), + ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(side_effect=_wrapped_user_lookup_error(Exception())), + ), + ): + assert await _revalidate_active_subject(identity) is None + + +@pytest.mark.asyncio +async def test_bridge_mint_fails_closed_before_upstream_when_master_key_unset(): + """master_key is validated BEFORE the upstream exchange (in _prepare_bridge_mint), so a + misconfigured gateway returns a 500 server_error without consuming the single-use code, avoiding + the burn-then-fail the pre-exchange phase exists to prevent. The failure is returned as an RFC 6749 + error body, not raised.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock() + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._resolve_active_litellm_key", + new=AsyncMock(return_value="no_active_key"), + ), + patch("litellm.proxy.proxy_server.master_key", None), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://claude.ai/api/mcp/auth_callback", + client_id="dcr-client-123", + client_secret=None, + code_verifier="verifier", + ) + + assert response.status_code == 500 + assert json.loads(response.body)["error"] == "server_error" + fake_http_client.post.assert_not_called() + + +async def _prepare_only_bridge_exchange(resolver_result): + """Drive exchange_token_with_server for a bridge oauth_delegate authorization_code request with the + identity resolver stubbed to a given tagged result, returning (response, post_mock) so a test can + assert the mapped status and that the single-use code was never exchanged.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock() + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow._resolve_active_litellm_key", + new=AsyncMock(return_value=resolver_result), + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://claude.ai/api/mcp/auth_callback", + client_id="dcr-client-123", + client_secret=None, + code_verifier="verifier", + ) + return response, fake_http_client.post + + +@pytest.mark.asyncio +async def test_bridge_mint_db_outage_is_503_before_upstream(): + """A DB outage while resolving identity is a retryable gateway failure, so the mint returns 503 + temporarily_unavailable WITHOUT consuming the single-use code, matching how admission statuses the + same outage on the egress side. Collapsing every resolution failure to None used to blame the + client with 400 invalid_request for an infrastructure problem.""" + response, post = await _prepare_only_bridge_exchange("unavailable") + assert response.status_code == 503 + assert json.loads(response.body)["error"] == "temporarily_unavailable" + post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_mint_unresolvable_identity_is_500_before_upstream(): + """An unresolvable identity (no DB connection, or an unexpected resolution error) is a gateway + fault, so the mint returns 500 server_error before the exchange, a status distinct from both the + caller's 400 and the transient 503, matching admission's 500-vs-503 split for the same conditions.""" + response, post = await _prepare_only_bridge_exchange("unresolvable") + assert response.status_code == 500 + assert json.loads(response.body)["error"] == "server_error" + post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_mint_upstream_expired_lifetime_is_502(): + """An upstream token response reporting an already-elapsed lifetime (a parseable non-positive + expires_in) is rejected with 502 rather than sealed into an hour-long envelope around a dead + bearer. Regression for expires_in<=0 silently falling through to the 1h cap.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 0} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" + + +@pytest.mark.asyncio +async def test_bridge_mint_positive_sub_second_lifetime_mints_not_502(): + """A positive fractional expires_in in (0, 1) is a live token, not an elapsed one, so it mints a + (1s-floored) envelope rather than being truncated to 0 and rejected with 502 after the single-use + code was already consumed. Regression for classifying a sub-second remaining lifetime as expired.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 0.5} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 200 + body = json.loads(response.body) + assert body["access_token"].startswith("llm_env_") + assert body["expires_in"] >= 0 + + +@pytest.mark.asyncio +async def test_bridge_mint_unknown_lifetime_is_capped_not_rejected(): + """An absent or unparseable expires_in leaves the lifetime unknown, which the envelope caps (never + inventing a longer life than the upstream stated); it is NOT rejected. Only an explicitly-dead + lifetime fails, so a metadata glitch on an otherwise-valid token still mints a bounded envelope.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": "not-a-number"} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 200 + body = json.loads(response.body) + assert body["access_token"].startswith("llm_env_") + assert 0 < body["expires_in"] <= 3600 + + +@pytest.mark.asyncio +async def test_bridge_reported_expires_in_does_not_overstate_jwt_exp(): + """The reported expires_in is derived from the envelope JWT's second-truncated exp (rounding the + elapsed portion up), so the client is never told the bearer lives past the point admission expires + it. Regression for the sub-second overstatement of the raw (expires_at - now) delta.""" + import time + + import jwt as _jwt + + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 300} + before = int(time.time()) + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + body = json.loads(response.body) + claims = _jwt.decode(body["access_token"].removeprefix("llm_env_"), options={"verify_signature": False}) + # projecting the reported lifetime from a time no later than the mint must not exceed the JWT exp + assert before + body["expires_in"] <= claims["exp"] + + +def test_bridge_reported_expires_in_can_be_zero_at_jwt_exp_boundary(): + from datetime import datetime, timezone + + from fastapi.responses import JSONResponse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _BridgeMintReady, + _finish_bridge_mint, + ) + 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( + 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( + ready=ready, + mcp_server=_bridge_server(auth_type=MCPAuth.oauth_delegate), + token_response={"access_token": "UP", "expires_in": 1}, + now=datetime.fromtimestamp(100.25, tz=timezone.utc), + ) + + assert isinstance(response, JSONResponse) + assert json.loads(response.body)["expires_in"] == 0 + + +def test_classify_upstream_lifetime(): + """expires_in from an IdP may be an int, a float (3600.0), or a numeric string ("3600"); each + coerces to a positive number of seconds. Absent or unparseable input (bool, non-numeric, NaN/inf, + oversized) is "unspecified" so the envelope caps it, while a parseable non-positive value is + "expired": the upstream reporting an already-dead token, which the mint must reject rather than + silently give the 1h cap.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _classify_upstream_lifetime + + assert _classify_upstream_lifetime(300) == 300 + assert _classify_upstream_lifetime(300.0) == 300 + assert _classify_upstream_lifetime("300") == 300 + assert _classify_upstream_lifetime(" 300 ") == 300 + # explicit, parseable, non-positive -> the upstream says the token is already dead + assert _classify_upstream_lifetime(0) == "expired" + assert _classify_upstream_lifetime(-5) == "expired" + assert _classify_upstream_lifetime(-0.5) == "expired" + # a positive sub-second lifetime is alive, not elapsed; it clamps up to the envelope's 1s floor + # rather than truncating to 0 and being misread as expired + assert _classify_upstream_lifetime(0.5) == 1 + assert _classify_upstream_lifetime(0.001) == 1 + # a positive value >= 1 truncates toward zero (never overstating the stated lifetime) + assert _classify_upstream_lifetime(1.9) == 1 + # unknown lifetime -> cap (never invent a longer life than the upstream stated) + assert _classify_upstream_lifetime(None) == "unspecified" + assert _classify_upstream_lifetime(True) == "unspecified" + assert _classify_upstream_lifetime("nope") == "unspecified" + # hostile numerics must not raise (int(float(...)) can OverflowError) -> unspecified + assert _classify_upstream_lifetime("inf") == "unspecified" + assert _classify_upstream_lifetime("1e999") == "unspecified" + assert _classify_upstream_lifetime("-inf") == "unspecified" + assert _classify_upstream_lifetime("nan") == "unspecified" + assert _classify_upstream_lifetime(float("inf")) == "unspecified" + assert _classify_upstream_lifetime(10**400) == "unspecified" + + +def test_bridge_grant_honors_and_rejects_upstream_lifetime(): + """The grant validator honors a positive lifetime, leaves an unknown one None for the envelope to + cap, and rejects an explicitly-expired one with "expired_lifetime" so a dead upstream token is + never sealed into an hour-long envelope.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _bridge_grant_from_token_response + + def grant(v): + return _bridge_grant_from_token_response({"access_token": "x", "expires_in": v}) + + assert grant(300).expires_in == 300 + assert grant(120.0).expires_in == 120 + # unknown lifetime backs a grant whose expires_in the envelope caps; it is not a rejection + assert grant("nope").expires_in is None + assert _bridge_grant_from_token_response({"access_token": "x"}).expires_in is None + # an explicitly already-dead lifetime is rejected, not silently capped at 1h + assert grant(0) == "expired_lifetime" + assert grant(-5) == "expired_lifetime" + + +@pytest.mark.asyncio +async def test_bridge_token_exchange_honors_short_float_expires_in_ttl(): + """A short float expires_in from the upstream caps the envelope TTL, so the client-held envelope + does not outlive the upstream token. Before coercion a float was dropped and the envelope + defaulted to the 1h cap (3600), which would forward a stale bearer after the upstream token + expired.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 120.0} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert json.loads(response.body)["expires_in"] <= 120 + + +@pytest.mark.asyncio +async def test_oauth_delegate_bridge_token_exchange_missing_access_token_is_502_not_keyerror(): + """When the upstream token response has no access_token, a dcr_bridge oauth_delegate exchange + returns a clean 502 error body rather than raising a KeyError. _finish_bridge_mint asks + _bridge_grant_from_token_response for a typed grant, gets None, and returns the "no_upstream_token" + failure, which maps to 502; nothing indexes token_response["access_token"] on the bridge path.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"token_type": "Bearer", "expires_in": 3600} + + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" + + +@pytest.mark.asyncio +async def test_true_passthrough_bridge_token_exchange_returns_raw_upstream_token(): + """Only oauth_delegate mints. A true_passthrough dcr_bridge server relays the raw upstream token + to the client, since that mode has no litellm identity to bind and the caller owns the token.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.true_passthrough) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + assert body["access_token"] == "UPSTREAM-SECRET-TOKEN" + assert not body["access_token"].startswith("llm_env_") + + +@pytest.mark.asyncio +async def test_non_bridge_oauth_delegate_token_exchange_returns_raw_upstream_token(): + """An oauth_delegate server without dcr_bridge keeps the pre-change contract: the raw upstream + token is returned, so flag-off behavior is byte-identical.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate, dcr_bridge=None) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + assert body["access_token"] == "UPSTREAM-SECRET-TOKEN" + + async def _exchange_persistence_attempted_for_auth_type(auth_type) -> bool: """Run exchange_token_with_server for a server of ``auth_type`` and report whether it attempted to persist the exchanged token server-side. The client-forwarded token modes must not persist: @@ -4580,7 +5829,7 @@ async def test_extract_user_id_reads_x_litellm_api_key_header(proxy_globals): """The LiteLLM key arrives on x-litellm-api-key (what Claude Desktop/Code send), not Authorization. Reading only Authorization dropped the identity, so the per-user token was never stored and the egress 401'd forever. Resolution must honor x-litellm-api-key.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth, hash_token @@ -4605,7 +5854,7 @@ async def test_extract_user_id_rehydrates_cross_replica_dict_cache(proxy_globals """Cross-replica, async_get_cache hands back a serialized dict, not a UserAPIKeyAuth. Resolution must rehydrate it; the old getattr(cached, "user_id") returned None on a dict, which is exactly why a multi-replica gateway never found the stored token.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import hash_token @@ -4626,7 +5875,7 @@ async def test_extract_user_id_falls_back_to_db_on_cache_miss(proxy_globals): """A cache miss must read the key from the DB rather than returning None; the old code did a cache-only peek and skipped the DB, so any replica that hadn't just authenticated the key failed to store the token.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth @@ -4648,7 +5897,7 @@ async def test_extract_user_id_falls_back_to_db_on_cache_miss(proxy_globals): @pytest.mark.asyncio async def test_extract_user_id_none_without_litellm_key(proxy_globals): """No LiteLLM key on the request resolves to None without consulting the resolver.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -4665,7 +5914,7 @@ async def test_extract_user_id_rejects_blocked_key(proxy_globals): """A blocked LiteLLM key must not resolve an identity. get_key_object returns the DB row without checking blocked/expiry (the main auth pipeline does, and the public token endpoint bypasses it), so a revoked key could otherwise overwrite the stored per-user OAuth token for its user.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth @@ -4687,7 +5936,7 @@ async def test_extract_user_id_rejects_expired_key(proxy_globals): """An expired LiteLLM key must not resolve an identity, for the same reason as a blocked key.""" from datetime import datetime, timedelta, timezone - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth @@ -4706,6 +5955,219 @@ async def test_extract_user_id_rejects_expired_key(proxy_globals): assert await _extract_user_id_from_request(request) is None +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_returns_resolved_key_for_active_key(proxy_globals): + """The dcr_bridge mint seals the hash of the authorizing key so admission can reload the live + record. For an active key the resolver returns exactly hash_token(key), the same value + get_key_object and the whole cache/DB layer key the record by, so the sealed reference resolves + back to this key at admission.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + _ResolvedKey, + ) + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + key = "sk-alice-key" + cache = UserApiKeyCache() + await cache.async_set_cache( + hash_token(key), + UserAPIKeyAuth(token=hash_token(key), user_id="alice"), + model_type=UserAPIKeyAuth, + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = object() + + request = _token_request({"x-litellm-api-key": f"Bearer {key}"}) + resolved = await _resolve_active_litellm_key(request) + assert isinstance(resolved, _ResolvedKey) + assert resolved.key_hash == hash_token(key) + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_resolves_key_without_user_id(proxy_globals): + """A valid team-scoped or service-account key has no user_id but is a legitimate credential, so it + must still resolve to a hash and be able to mint a bridge envelope. Gating the resolver on user_id + presence wrongly rejected these keys with invalid_request; the active-state gate now checks only + blocked and expiry, and the key hash (not the user) is what the mint seals. The per-user token + store still gets no user for such a key, since there is none to key a stored credential by.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _extract_user_id_from_request, + _ResolvedKey, + _resolve_active_litellm_key, + ) + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + key = "sk-team-scoped-key" + cache = UserApiKeyCache() + await cache.async_set_cache( + hash_token(key), + UserAPIKeyAuth(token=hash_token(key), user_id=None, team_id="team-x"), + model_type=UserAPIKeyAuth, + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = object() + + request = _token_request({"x-litellm-api-key": f"Bearer {key}"}) + resolved = await _resolve_active_litellm_key(request) + assert isinstance(resolved, _ResolvedKey) + assert resolved.key_hash == hash_token(key) + assert await _extract_user_id_from_request(request) is None + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_rejects_blocked_key(proxy_globals): + """A blocked key must not yield a hash, so no gateway-bound envelope is minted for a revoked key; + the mint fails closed with invalid_request instead.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _FakePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + return UserAPIKeyAuth(token=token, user_id="blocked-user", blocked=True) + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _FakePrisma() + + request = _token_request({"x-litellm-api-key": "sk-blocked-key"}) + assert await _resolve_active_litellm_key(request) == "no_active_key" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_fails_closed_on_malformed_expiry(proxy_globals): + """A key whose stored expires string does not parse must fail closed to no-hash (the mint then + returns invalid_request), not surface an unhandled 500. The active-state check runs outside the + resolver's try, so it must be total over a bad expires rather than letting datetime.fromisoformat + raise. Before the fix this raised a ValueError instead of returning None.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _FakePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + return UserAPIKeyAuth(token=token, user_id="u", expires="not-a-parseable-date") + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _FakePrisma() + + request = _token_request({"x-litellm-api-key": "sk-bad-expiry-key"}) + assert await _resolve_active_litellm_key(request) == "no_active_key" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_no_active_key_without_litellm_key(proxy_globals): + """No LiteLLM key on the request yields no hash without consulting the resolver.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + request = _token_request({"content-type": "application/json"}) + assert await _resolve_active_litellm_key(request) == "no_active_key" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_db_outage_is_unavailable(proxy_globals): + """A database outage while resolving the presented key is a retryable infrastructure failure, not + the caller's fault, so the resolver reports "unavailable" (the mint statuses it 503) rather than + collapsing it to the same value as a missing credential. is_database_service_unavailable_error + classifies a connection error (an OSError) as an outage, matching admission's egress-side handling.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _OutagePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + raise ConnectionError("connection refused") + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _OutagePrisma() + + request = _token_request({"x-litellm-api-key": "sk-during-outage"}) + assert await _resolve_active_litellm_key(request) == "unavailable" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_no_database_is_unresolvable(proxy_globals): + """With no database connection configured the gateway cannot verify the presented key at all, so + the resolver reports "unresolvable" (the mint statuses it 500) instead of blaming the caller. + Mirrors admission, which 500s a missing prisma_client on the egress side.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = None + + request = _token_request({"x-litellm-api-key": "sk-no-db"}) + assert await _resolve_active_litellm_key(request) == "unresolvable" + + +def _wrapped_user_lookup_error(original: BaseException) -> ValueError: + """Reproduce get_user_object's real exception contract (litellm/proxy/auth/auth_checks.py): it + catches every DB failure in a broad ``except`` and re-raises a bare ``ValueError``, so the original + error (a missing-user Exception or a real outage) survives only as ``__context__``. Injecting a raw + ConnectionError/Exception instead would exercise a shape production never produces and let a + chain-blind outage classifier pass. The wrapping fidelity is pinned by + test_get_user_object_wraps_db_outage_as_valueerror_preserving_context in test_auth_checks.""" + try: + raise original + except BaseException: + try: + raise ValueError(f"User doesn't exist in db. Got error - {original}") + except ValueError as wrapped: + return wrapped + + +@pytest.mark.asyncio +async def test_reload_active_user_by_id_missing_user_is_no_active_key(proxy_globals): + """A user_id refresh envelope whose user has been deleted must fail closed to no_active_key (the + refresh path maps it to invalid_grant), not unresolvable/500. get_user_object catches the missing row + and re-raises a bare ValueError, so a missing user must not be misclassified as a DB outage or an + opaque gateway fault.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _reload_active_user_by_id + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + with patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(side_effect=_wrapped_user_lookup_error(Exception())), + ): + assert await _reload_active_user_by_id("gone-user") == "no_active_key" + + +@pytest.mark.asyncio +async def test_reload_active_user_by_id_db_outage_is_unavailable(proxy_globals): + """A transient DB outage while re-validating the user on refresh is a retryable outage, distinct from + a missing user, so the refresh path surfaces "unavailable" (a 503) rather than blaming the caller. + get_user_object wraps the outage in a bare ValueError, so this exercises the chain-aware classifier; a + raw ConnectionError would falsely pass even a chain-blind check because it is an OSError.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _reload_active_user_by_id + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + with patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(side_effect=_wrapped_user_lookup_error(ConnectionError("user database unreachable"))), + ): + assert await _reload_active_user_by_id("sso-user-7") == "unavailable" + + @pytest.mark.asyncio async def test_token_endpoint_uses_client_secret_basic_when_configured(): """LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the @@ -5306,3 +6768,364 @@ async def test_token_exchange_pairs_client_secret_with_server_client_id(): sent = mock_async_client.post.call_args.kwargs["data"] assert sent["client_id"] == "persisted-client" assert "client_secret" not in sent + + +def _upstream_token_response(status_code: int, *, json_body: object = None, text_body: str = "") -> "httpx.Response": + import httpx + + request = httpx.Request("POST", "https://oauth2.googleapis.com/token") + if json_body is not None: + return httpx.Response(status_code, json=json_body, request=request) + return httpx.Response(status_code, text=text_body, request=request) + + +async def _exchange_with_upstream_response(upstream_response, *, server_client_id="web-client.apps.googleusercontent.com"): + """Run the raw (non-bridge) authorization_code exchange against a canned upstream token-endpoint + response and return what the gateway would hand the client. ``server_client_id=None`` models the + caller-supplied-credentials flow (no stored client on the server).""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="gcal", + name="gcal", + server_name="gcal", + alias="gcal", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=server_client_id, + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + ) + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=upstream_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ): + return await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://litellm.example.com/ui/mcp/oauth/callback", + client_id="web-client.apps.googleusercontent.com", + client_secret=None, + code_verifier="verifier", + ) + + +@pytest.mark.asyncio +async def test_token_exchange_gateway_credential_rejection_is_502_with_gateway_prose(): + """When the gateway presented the server's stored client credentials and the IdP rejected them + (Google refusing a secret-less or unknown client), the fault is the operator's, not the caller's: + 502 server_error with gateway-authored prose naming the code, and the IdP's own prose stays in + server logs. Before the framework this either 500ed raw or relayed provider prose verbatim.""" + response = await _exchange_with_upstream_response( + _upstream_token_response( + 401, json_body={"error": "invalid_client", "error_description": "The OAuth client was not found."} + ) + ) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body["error"] == "server_error" + assert "invalid_client" in body["error_description"] + assert "client_id and client_secret" in body["error_description"] + assert "The OAuth client was not found." not in body["error_description"] + assert response.headers["cache-control"] == "no-store" + + +@pytest.mark.asyncio +async def test_token_exchange_caller_supplied_credential_rejection_relays_code(): + """When the caller supplied the client credentials themselves (no stored client on the server), + an invalid_client rejection is theirs to act on: the §5.2 code relays on the 401 that code + implies.""" + response = await _exchange_with_upstream_response( + _upstream_token_response( + 401, json_body={"error": "invalid_client", "error_description": "The OAuth client was not found."} + ), + server_client_id=None, + ) + + assert response.status_code == 401 + body = json.loads(response.body) + assert body == {"error": "invalid_client", "error_description": "The OAuth client was not found."} + + +@pytest.mark.asyncio +async def test_token_exchange_status_derives_from_error_code_not_upstream_status(): + """An upstream that pairs a caller-fault code with a server-fault status (invalid_grant on a 500) + must not produce a contradictory response: status derives from the classified fault, so the + caller sees 400 invalid_grant and knows to re-authorize rather than blaming the gateway.""" + response = await _exchange_with_upstream_response( + _upstream_token_response(500, json_body={"error": "invalid_grant", "error_description": "Code expired."}) + ) + + assert response.status_code == 400 + body = json.loads(response.body) + assert body == {"error": "invalid_grant", "error_description": "Code expired."} + + +@pytest.mark.asyncio +async def test_token_exchange_relays_only_rfc6749_error_fields(): + """Only error / error_description / error_uri cross the gateway; any other upstream body field is + dropped so an arbitrary rejection payload cannot ride the relay to the client.""" + response = await _exchange_with_upstream_response( + _upstream_token_response( + 400, + json_body={ + "error": "invalid_grant", + "error_description": "Code was already redeemed.", + "error_uri": "https://idp.example.com/errors/invalid_grant", + "internal_trace": "should never reach the client", + }, + ) + ) + + assert response.status_code == 400 + body = json.loads(response.body) + assert set(body.keys()) == {"error", "error_description", "error_uri"} + + +@pytest.mark.asyncio +async def test_token_exchange_maps_out_of_contract_rejection_to_502(): + """A rejection outside the §5.2 contract (no JSON error field, or a status the token-endpoint + contract does not define) is an upstream fault; 502 keeps it from being misread as a caller + mistake while the description still names the upstream status.""" + response = await _exchange_with_upstream_response( + _upstream_token_response(503, text_body="upstream maintenance") + ) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body["error"] == "server_error" + assert "HTTP 503" in body["error_description"] + + +@pytest.mark.asyncio +async def test_token_exchange_bounds_relayed_error_fields(): + """Relayed §5.2 fields are length-bounded so a hostile or broken upstream cannot bloat the + gateway response.""" + response = await _exchange_with_upstream_response( + _upstream_token_response(400, json_body={"error": "invalid_request", "error_description": "x" * 5000}) + ) + + assert response.status_code == 400 + body = json.loads(response.body) + assert len(body["error_description"]) == 500 + + +@pytest.mark.asyncio +async def test_token_exchange_200_without_access_token_is_502_not_keyerror(): + """A 200 whose body has no usable access_token used to KeyError into a 500; the raw arm now + answers 502 with the same wording as the bridge arm's no_upstream_token rejection.""" + response = await _exchange_with_upstream_response( + _upstream_token_response(200, json_body={"token_type": "Bearer"}) + ) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body["error"] == "server_error" + assert "access_token" in body["error_description"] + + +@pytest.mark.asyncio +async def test_token_exchange_relays_rejection_when_http_client_raises(): + """litellm's AsyncHTTPHandler.post raise_for_status()es internally and raises MaskedHTTPStatusError + at call time, so in production the rejection escapes from the post call itself rather than from the + explicit raise_for_status; the relay must catch it there too (proven live: a mock returning the + error response passed while the real proxy still 500ed).""" + import httpx + + rejection = _upstream_token_response( + 401, json_body={"error": "invalid_client", "error_description": "The OAuth client was not found."} + ) + raising_client = MagicMock() + raising_client.post = AsyncMock( + side_effect=httpx.HTTPStatusError("Client error '401 Unauthorized'", request=rejection.request, response=rejection) + ) + + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="gcal", + name="gcal", + server_name="gcal", + alias="gcal", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="web-client.apps.googleusercontent.com", + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + ) + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=raising_client, + ): + response = await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://litellm.example.com/ui/mcp/oauth/callback", + client_id="web-client.apps.googleusercontent.com", + client_secret=None, + code_verifier="verifier", + ) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body["error"] == "server_error" + assert "invalid_client" in body["error_description"] + + +@pytest.mark.asyncio +async def test_register_relays_rejection_when_http_client_raises(): + """Same live mechanism as the token exchange: the DCR rejection escapes from the post call itself, + so the register relay must catch it there, not only from the explicit raise_for_status.""" + import httpx + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client_with_server, + ) + + rejection = httpx.Response( + 400, + json={"error": "invalid_client_metadata"}, + request=httpx.Request("POST", "https://idp.example.com/register"), + ) + raising_client = MagicMock() + raising_client.post = AsyncMock( + side_effect=httpx.HTTPStatusError("Client error '400 Bad Request'", request=rejection.request, response=rejection) + ) + + oauth2_server = _bridge_server(auth_type=MCPAuth.oauth2, dcr_bridge=None) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=raising_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reuse_persisted_dcr_client_if_available", + new_callable=AsyncMock, + return_value=False, + ), + ): + with pytest.raises(HTTPException) as exc: + await register_client_with_server( + request=_bridge_mock_request(), + mcp_server=oauth2_server, + client_name="Claude", + grant_types=None, + response_types=None, + token_endpoint_auth_method=None, + ) + + assert exc.value.status_code == 400 + assert "invalid_client_metadata" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_token_exchange_never_relays_out_of_contract_body_to_client(): + """These endpoints serve unauthenticated OAuth clients, so a non-RFC6749 upstream body (HTML + error page, proxy banner, stack trace) must stay in server logs; the client sees only the + upstream status.""" + response = await _exchange_with_upstream_response( + _upstream_token_response(404, text_body="Error 404 stack trace: secret internals") + ) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 404"} + + +@pytest.mark.asyncio +async def test_register_never_relays_out_of_contract_body_to_client(): + """Same trust boundary for DCR: a non-RFC7591 rejection body is logged server-side and the + client detail names only the status.""" + import httpx + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client_with_server, + ) + + rejection = httpx.Response( + 500, + text="Tomcat stack trace with internals", + request=httpx.Request("POST", "https://idp.example.com/register"), + ) + raising_client = MagicMock() + raising_client.post = AsyncMock( + side_effect=httpx.HTTPStatusError("Server error '500'", request=rejection.request, response=rejection) + ) + + oauth2_server = _bridge_server(auth_type=MCPAuth.oauth2, dcr_bridge=None) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=raising_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reuse_persisted_dcr_client_if_available", + new_callable=AsyncMock, + return_value=False, + ), + ): + with pytest.raises(HTTPException) as exc: + await register_client_with_server( + request=_bridge_mock_request(), + mcp_server=oauth2_server, + client_name="Claude", + grant_types=None, + response_types=None, + token_endpoint_auth_method=None, + ) + + assert exc.value.status_code == 502 + assert str(exc.value.detail) == "upstream registration failed with HTTP 500" + assert "Tomcat" not in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): + """An upstream whose failure body cannot be read (unconsumed stream, lying content-encoding) + makes response.text/.json raise; the classifier must stay total so the caller still gets the + §5.2-shaped 502 instead of the opaque 500 this change set out to remove.""" + import httpx + + unreadable = httpx.Response( + 400, + stream=httpx.ByteStream(b"\x1f\x8bnot-actually-gzip"), + headers={"content-encoding": "gzip"}, + request=httpx.Request("POST", "https://oauth2.googleapis.com/token"), + ) + + response = await _exchange_with_upstream_response(unreadable) + + assert response.status_code == 502 + body = json.loads(response.body) + assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 32996905166..a983ac3ff48 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -5428,6 +5428,457 @@ def test_get_forwarded_auth_from_scope_skips_when_no_litellm_key_header(): assert _get_forwarded_auth_from_scope(scope) is None +def _delegate_auth_mcp_server(server_id: str = "delegate-1") -> MCPServer: + return MCPServer( + server_id=server_id, + name="delegate_test", + url="http://upstream:9401/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + oauth2_flow="authorization_code", + ) + + +def _delegate_scope(headers: list) -> dict: + return { + "type": "http", + "method": "POST", + "path": "/mcp/delegate_test", + "scheme": "http", + "server": ("localhost", 4000), + "headers": headers, + } + + +def _patch_delegate_resolver(server: MCPServer, *resolvable_names: str): + """Patch the admission-parity resolver the delegate probe gates on. Returns + ``server`` only for names admission's ``get_mcp_server_by_name`` would match + (alias / server_name / name); every other name (server_id, access group) yields + None, exactly as the real resolver does.""" + + def _resolve(name, client_ip=None): + return server if name in resolvable_names else None + + return patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + side_effect=_resolve, + ) + + +@pytest.mark.asyncio +async def test_delegate_bad_token_gets_connect_time_401(): + """Regression (LIT-4194): a rejected upstream token on a delegate-auth server + must fail the connect with 401 + ``error="invalid_token"``, not be absorbed + into HTTP 200 + an empty tool list by the tools/list handler. + + Delegate-mode clients send only ``Authorization`` (no ``x-litellm-api-key``), + so ``_get_forwarded_auth_from_scope`` returns None and, before the fix, the + preflight returned early without probing. + """ + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server() + scope = _delegate_scope([(b"authorization", b"Bearer bogus-token")]) + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, 'Bearer realm="upstream", error="invalid_token"')), + ) as probe: + with pytest.raises(HTTPException) as exc_info: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + assert exc_info.value.status_code == 401 + challenge = exc_info.value.headers["www-authenticate"] + assert 'error="invalid_token"' in challenge + assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge + probe.assert_awaited_once() + probe_url, probe_auth = probe.call_args.args + assert probe_url == "http://upstream:9401/mcp" + assert probe_auth == "Bearer bogus-token" + + +@pytest.mark.asyncio +async def test_delegate_valid_token_passes_preflight(): + """An upstream-accepted token must not be blocked by the delegate preflight.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server() + scope = _delegate_scope([(b"authorization", b"Bearer good-token")]) + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(200, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + probe.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_delegate_valid_token_forbidden_returns_403(): + """An upstream that accepts the token but forbids the caller (403) must surface + as a bare 403 with no ``WWW-Authenticate`` re-auth hint (a fresh token with the + same scopes would loop), not as an invalid_token challenge.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server() + scope = _delegate_scope([(b"authorization", b"Bearer scoped-out-token")]) + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(403, None)), + ): + with pytest.raises(HTTPException) as exc_info: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + assert exc_info.value.status_code == 403 + assert not (exc_info.value.headers or {}) + + +@pytest.mark.asyncio +async def test_delegate_tokenless_request_not_probed(): + """Tokenless delegate requests are the preemptive challenge's job; the + preflight must not probe upstream with an empty credential.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server() + scope = _delegate_scope([(b"content-type", b"application/json")]) + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delegate_preflight_skipped_on_multi_server_routes(): + """The delegate probe is gated to single-server routes so one rejected token + cannot 401 a multi-server aggregate connect (matching the OBO preflight).""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + servers = [_delegate_auth_mcp_server("delegate-1"), _delegate_auth_mcp_server("delegate-2")] + scope = _delegate_scope([(b"authorization", b"Bearer bogus-token")]) + + with _patch_delegate_resolver(servers[0], "delegate_test", "other_server"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=servers), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test", "other_server"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_bare_authorization_never_probes_passthrough_servers(): + """A bare ``Authorization`` header may be a LiteLLM key (backward-compat), so + only delegate servers (where admission classified it as an upstream token) + may be probed with it; ``is_oauth_passthrough`` servers still require the + unambiguous ``x-litellm-api-key`` + ``Authorization`` pair.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + passthrough_server = MCPServer( + server_id="pt-1", + name="pt_server", + url="http://upstream:9402/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + oauth_passthrough=True, + extra_headers=["Authorization"], + ) + scope = _delegate_scope([(b"authorization", b"Bearer ambiguous-token")]) + + with _patch_delegate_resolver(passthrough_server, "pt_server"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[passthrough_server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["pt_server"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delegate_not_probed_when_named_only_via_server_id(): + """Security regression (LIT-4194): a delegate server reachable by the requested + name only through its server_id (or an access group) is admitted as a real + LiteLLM key by ``process_mcp_request`` (its ``get_mcp_server_by_name`` misses), + so the bare ``Authorization`` header is that LiteLLM key. The probe must resolve + the target through the SAME resolver and therefore skip it, never forwarding the + key upstream, even though the widened allowed-server set still contains it.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server(server_id="delegate-secret-id") + # Admission's resolver matches alias/server_name/name only, never server_id: the + # requested server_id resolves to None here, mirroring the real divergence. + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/delegate-secret-id", + "scheme": "http", + "server": ("localhost", 4000), + "headers": [(b"authorization", b"Bearer sk-litellm-proxy-key")], + } + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(user_id="u1", api_key="hashed-sk"), + mcp_servers=["delegate-secret-id"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delegate_preflight_with_unpatched_probe(): + """Integration across the preflight and the unpatched ``_probe_upstream_auth``, + mocked only at the httpx-client boundary (tests/test_litellm is mocked-only; the + real-network proof lives in the PR's live-proxy evidence). The mock honors the + ``AsyncHTTPHandler.post`` contract by raising ``httpx.HTTPStatusError`` on the + upstream 401, so the production ``except httpx.HTTPStatusError`` branch is the one + exercised. A rejected token surfaces as the connect-time 401 challenge; an + accepted token passes untouched, and the caller's bearer reaches the delegate URL.""" + import httpx + + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + accepted = MagicMock() + accepted.status_code = 200 + accepted.headers = {} + rejected = MagicMock() + rejected.status_code = 401 + rejected.headers = {"www-authenticate": 'Bearer realm="stub-upstream", error="invalid_token"'} + + async def respond_by_token(url=None, headers=None, json=None, timeout=None, **kwargs): + if headers.get("Authorization") == "Bearer good-token": + return accepted + raise httpx.HTTPStatusError( + "401 Unauthorized", + request=httpx.Request("POST", url), + response=rejected, + ) + + mock_client = MagicMock() + mock_client.post = AsyncMock(side_effect=respond_by_token) + + server = _delegate_auth_mcp_server() + + with _patch_delegate_resolver(server, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server.get_async_httpx_client", + return_value=mock_client, + ): + with pytest.raises(HTTPException) as exc_info: + await _check_passthrough_upstream_auth( + scope=_delegate_scope([(b"authorization", b"Bearer bogus-token")]), + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + await _check_passthrough_upstream_auth( + scope=_delegate_scope([(b"authorization", b"Bearer good-token")]), + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + assert exc_info.value.status_code == 401 + challenge = exc_info.value.headers["www-authenticate"] + assert 'error="invalid_token"' in challenge + assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge + probed_urls = [call.kwargs["url"] for call in mock_client.post.await_args_list] + assert probed_urls == ["http://upstream:9401/mcp", "http://upstream:9401/mcp"] + + +@pytest.mark.asyncio +async def test_delegate_challenge_echoes_requested_alias(): + """An alias-routed delegate request must be probed, and the challenge must echo + the requested alias (not the canonical server name) so the resource_metadata + URL matches what the tokenless preemptive challenge emits for the same route.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = _delegate_auth_mcp_server().model_copy(update={"alias": "dt-alias"}) + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/dt-alias", + "scheme": "http", + "server": ("localhost", 4000), + "headers": [(b"authorization", b"Bearer bogus-token")], + } + + with _patch_delegate_resolver(server, "dt-alias"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, 'Bearer error="invalid_token"')), + ): + with pytest.raises(HTTPException) as exc_info: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["dt-alias"], + client_ip=None, + ) + + challenge = exc_info.value.headers["www-authenticate"] + assert 'error="invalid_token"' in challenge + assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/dt-alias"' in challenge + + +@pytest.mark.asyncio +async def test_delegate_probe_not_fanned_out_to_access_group_members(): + """A single access-group name passes the one-target route gate but must not fan + the delegate probe out to group-expanded member servers; the group name resolves + to no server under admission's resolver, so no probe fires.""" + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + group_member = _delegate_auth_mcp_server() + + with _patch_delegate_resolver(group_member, "delegate_test"), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[group_member]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe: + await _check_passthrough_upstream_auth( + scope=_delegate_scope([(b"authorization", b"Bearer bogus-token")]), + user_api_key_auth=UserAPIKeyAuth(), + mcp_servers=["prod_tools_group"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +def test_is_delegate_upstream_probe_target_fails_closed_on_m2m_shape(): + """An unstamped M2M-shape row (null ``oauth2_flow`` + client credentials) + resolves to ``client_credentials`` and must not be probed with the caller's + bearer; its stored client credentials drive egress instead.""" + from litellm.proxy._experimental.mcp_server.server import ( + _is_delegate_upstream_probe_target, + ) + + assert _is_delegate_upstream_probe_target(_delegate_auth_mcp_server()) is True + + m2m_shape = MCPServer( + server_id="delegate-m2m", + name="delegate_m2m", + url="http://upstream:9401/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + oauth2_flow=None, + token_url="http://idp:9000/token", + client_id="client", + client_secret="secret", + ) + assert _is_delegate_upstream_probe_target(m2m_shape) is False + + non_delegate = MCPServer( + server_id="oauth2-plain", + name="oauth2_plain", + url="http://upstream:9401/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + ) + assert _is_delegate_upstream_probe_target(non_delegate) is False + + @pytest.mark.asyncio async def test_create_mcp_client_sampling_disabled_by_default(): """Sampling callback must be None when allow_sampling is not set (default False).""" diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d12ff20ee5b..27f43c4948f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -701,6 +701,38 @@ async def test_default_internal_user_params_with_get_user_object(monkeypatch): assert creation_args["user_role"] == "internal_user" +@pytest.mark.asyncio +async def test_get_user_object_wraps_db_outage_as_valueerror_preserving_context(): + """Pin get_user_object's exception contract: it catches every DB failure in a broad except and + re-raises a bare ValueError, so a real outage survives only as __context__ rather than as the + exception type. The MCP dcr_bridge admission and refresh paths depend on this to tell a transient + outage (retry, 503) from a missing user (fail closed), which is why they classify across the cause + chain instead of the top exception's type. If this wrapping ever changes, that classification must + change with it, so this test guards the contract the callers rely on.""" + from unittest.mock import AsyncMock, MagicMock, patch + + mock_prisma_client = MagicMock() + mock_prisma_client.db = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + side_effect=ConnectionError("can't reach database server") + ) + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): + with pytest.raises(ValueError) as exc_info: + await get_user_object( + user_id="outage-contract-probe-user", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + user_id_upsert=False, + proxy_logging_obj=None, + ) + + assert isinstance(exc_info.value.__context__, ConnectionError) + + @pytest.mark.asyncio async def test_get_user_object_upsert_includes_user_email(): """Test that user_email is included when creating a new user via get_user_object upsert""" @@ -4413,3 +4445,85 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): request=MagicMock(spec=Request), ) assert "User=u1" in str(over.value) + + +@pytest.mark.asyncio +async def test_user_budget_enforced_on_team_key(): + """User budget must be enforced even when the key belongs to a team. + + Previously _user_max_budget_check skipped enforcement for team keys, + letting a user with a $100 personal budget spend unlimited through a + team key. This regression test ensures that is no longer the case. + """ + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0) + team = LiteLLM_TeamTable(team_id="t1", max_budget=2100.0) + token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1") + + async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): + return 999.0 if counter_key == "spend:user:u1" else 0.0 + + async def _no_membership(*a, **kw): + return None + + with patch("litellm.proxy.proxy_server.prisma_client", None), patch( + "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership): + with pytest.raises(litellm.BudgetExceededError) as over: + await common_checks( + request_body={"messages": [{"role": "user", "content": "hi"}]}, + team_object=team, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + assert "User=u1" in str(over.value) + + +@pytest.mark.asyncio +async def test_skip_user_budget_on_team_key_flag_restores_old_behavior(): + """Setting skip_user_budget_on_team_key=True skips user budget for team keys. + + This is the opt-in escape hatch that restores the legacy behavior where + user budgets were not enforced when the key belonged to a team. + """ + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0) + team = LiteLLM_TeamTable(team_id="t1", max_budget=2100.0) + token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1") + + async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): + return 999.0 if counter_key == "spend:user:u1" else 0.0 + + async def _no_membership(*a, **kw): + return None + + with patch("litellm.proxy.proxy_server.prisma_client", None), patch( + "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership): + result = await common_checks( + request_body={"messages": [{"role": "user", "content": "hi"}]}, + team_object=team, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={"skip_user_budget_on_team_key": True}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + assert result is True diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 90f46152837..a0248963cf1 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -3992,6 +3992,80 @@ async def test_non_admin_cli_session_token_reaches_production_auth_path(monkeypa assert result.is_session_token is True +@pytest.mark.asyncio +async def test_cli_session_token_authenticates_when_jwt_auth_enabled_without_license(monkeypatch): + """A lite login token is an encrypted (non-JWT) session blob. With + enable_jwt_auth on and no enterprise license (premium_user False), the JWT + premium gate used to fire for every request before the token was decoded, so + the CLI token 401'd with 'JWT Auth is an enterprise only feature' and was + never decrypted. The gate must apply only to actual JWTs; a non-JWT session + token has to keep authenticating on its own path regardless of license.""" + monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False) + cli_token = _mint_cli_session_token(monkeypatch) + + jwt_handler = MagicMock() + jwt_handler.is_jwt = JWTHandler.is_jwt + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + mock_request = MagicMock() + mock_request.url.path = "/v1/messages" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {cli_token}"} + mock_request.query_params = {} + + with ( + patch("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True}), + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + result = await user_api_key_auth( + request=mock_request, + api_key=f"Bearer {cli_token}", + ) + + assert result.user_id == "cli-admin" + assert result.team_id == "cli-team" + assert result.token is not None and result.token.startswith("cli-session-") + + +@pytest.mark.asyncio +async def test_real_jwt_still_requires_license_when_jwt_auth_enabled(monkeypatch): + """Guard for the reorder above: the enterprise gate must still reject an + actual JWT when there is no license. Moving the premium check inside the + is_jwt branch must not open JWT auth to non-premium deployments.""" + monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False) + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-cli-test") + + jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.sig" + jwt_handler = MagicMock() + jwt_handler.is_jwt = JWTHandler.is_jwt + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + mock_request = MagicMock() + mock_request.url.path = "/v1/messages" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + + with ( + patch("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True}), + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + with pytest.raises(Exception) as exc_info: + await user_api_key_auth( + request=mock_request, + api_key=f"Bearer {jwt_token}", + ) + + message = str(getattr(exc_info.value, "message", exc_info.value)) + assert "enterprise only feature" in message + + @pytest.mark.asyncio async def test_auth_path_caches_team_object_under_canonical_team_id_key(): """Regression for LIT-4000: the auth builder must cache the team object under diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 0634a01326c..23099177812 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -286,6 +286,46 @@ def test_is_database_service_unavailable_error_excludes_non_infra(error): ) +def _wrapped_like_get_user_object(original): + """Reproduce get_user_object's exception contract (litellm/proxy/auth/auth_checks.py): it catches + every DB failure in a broad ``except`` and re-raises a bare ``ValueError``, so the original error + survives only as ``__context__``. Building it by raising inside an ``except`` sets ``__context__`` + exactly as production does.""" + try: + raise original + except BaseException: + try: + raise ValueError("User doesn't exist in db. Got error - x") + except ValueError as wrapped: + return wrapped + + +def test_is_database_service_unavailable_error_in_chain_sees_through_wrapping(): + """The chain-aware classifier must see a real outage that a caller wrapped in a different type. + get_user_object turns a connection error into a bare ValueError whose type check reads as non-infra, + so the single-exception check returns False and only the chain walk recovers the outage. A missing + user (whose wrapped cause is a plain Exception) must stay non-infra on both.""" + outage = _wrapped_like_get_user_object(ConnectionError("can't reach database server")) + missing_user = _wrapped_like_get_user_object(Exception()) + + assert PrismaDBExceptionHandler.is_database_service_unavailable_error(outage) is False + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(outage) is True + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(missing_user) is False + # parity: a raw outage with no wrapper is still an outage, and a plain ValueError is not + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(ConnectionError("boom")) is True + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(ValueError("nope")) is False + + +def test_is_database_service_unavailable_error_in_chain_terminates_on_a_cause_cycle(): + """The walk must terminate on a pathological __cause__ cycle rather than hang. Neither link is an + outage, so the bounded walk returns False instead of looping forever.""" + first = ValueError("first") + second = ValueError("second") + first.__cause__ = second + second.__cause__ = first + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(first) is False + + def test_is_database_service_unavailable_error_asyncpg(monkeypatch): """asyncpg connection/interface errors map to service-unavailable. asyncpg is not a hard dependency, so inject a stand-in module to exercise the diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py index 5a84b6ebecd..16185cadbdf 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -693,6 +693,8 @@ class TestLassoGuardrail: assert prompt_payload["messages"] == messages assert prompt_payload["userId"] == "test-user" assert prompt_payload["sessionId"] == "test-conversation" + # Every call is attributed to the "litellm" integration for the "Used By" badge. + assert prompt_payload["source"] == {"type": "litellm"} # Test COMPLETION payload completion_messages = [{"role": "assistant", "content": "Test response"}] @@ -703,6 +705,7 @@ class TestLassoGuardrail: assert completion_payload["messages"] == completion_messages assert completion_payload["userId"] == "test-user" assert completion_payload["sessionId"] == "test-conversation" + assert completion_payload["source"] == {"type": "litellm"} def test_header_preparation(self): """Test header preparation.""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py index 60c595c64e5..6e601df897b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py @@ -1644,12 +1644,64 @@ class TestXecGuardLoggingHook: ) assert out_kwargs is kwargs assert out_result is result - info = kwargs["standard_logging_object"]["guardrail_information"] + info_list = kwargs["standard_logging_object"]["guardrail_information"] + assert isinstance(info_list, list), "guardrail_information must be a list" + assert len(info_list) == 1 + info = info_list[0] assert info["guardrail_mode"] == "logging_only" - assert info["guardrail_name"] == "xecguard" + assert info["guardrail_name"] == "test-xecguard" assert info["guardrail_status"] == "success" assert info["guardrail_response"]["trace_id"] == "lg-1" + @pytest.mark.asyncio + async def test_async_logging_hook_appends_to_existing_guardrail_info( + self, xecguard_guardrail, mock_request_data + ): + resp = _make_response({"decision": "SAFE", "trace_id": "lg-4"}) + prior_entry = {"guardrail_name": "other-guardrail"} + with patch.object(xecguard_guardrail.async_handler, "post", return_value=resp): + kwargs = { + **mock_request_data, + "standard_logging_object": {"guardrail_information": [prior_entry]}, + } + await xecguard_guardrail.async_logging_hook( + kwargs=kwargs, + result=_build_model_response("some answer"), + call_type="acompletion", + ) + info_list = kwargs["standard_logging_object"]["guardrail_information"] + assert len(info_list) == 2 + assert info_list[0] is prior_entry + assert info_list[1]["guardrail_name"] == "test-xecguard" + assert info_list[1]["guardrail_response"]["trace_id"] == "lg-4" + + @pytest.mark.asyncio + async def test_async_logging_hook_sanitizes_scan_result( + self, xecguard_guardrail, mock_request_data + ): + resp = _make_response( + { + "decision": "SAFE", + "trace_id": "lg-5", + "secret_fields": {"authorization": "Bearer xgs_raw"}, + "detections": [{"match": "raw matched span", "policy": "pii"}], + "api_key": "xgs_super_secret_value", + } + ) + with patch.object(xecguard_guardrail.async_handler, "post", return_value=resp): + kwargs = {**mock_request_data, "standard_logging_object": {}} + await xecguard_guardrail.async_logging_hook( + kwargs=kwargs, + result=_build_model_response("some answer"), + call_type="acompletion", + ) + info = kwargs["standard_logging_object"]["guardrail_information"][0] + guardrail_response = info["guardrail_response"] + assert "secret_fields" not in guardrail_response + assert guardrail_response["detections"][0]["match"] == "[REDACTED]" + assert guardrail_response["api_key"] != "xgs_super_secret_value" + assert guardrail_response["trace_id"] == "lg-5" + @pytest.mark.asyncio async def test_async_logging_hook_without_response_records_info( self, xecguard_guardrail, mock_request_data @@ -1680,7 +1732,9 @@ class TestXecGuardLoggingHook: result=_build_model_response("x"), call_type="acompletion", ) - info = kwargs["standard_logging_object"]["guardrail_information"] + info_list = kwargs["standard_logging_object"]["guardrail_information"] + assert isinstance(info_list, list), "guardrail_information must be a list" + info = info_list[0] assert info["guardrail_status"] == "guardrail_intervened" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index 34d92505359..3dfb98c12ea 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -524,3 +524,42 @@ def test_apply_redacted_messages_back_skips_input_when_not_string(): data = {"input": [{"type": "text", "text": "leak"}]} apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) assert data["input"] == [{"type": "text", "text": "leak"}] + + +# ------------------------------------------------------------------- +# LIT-4302: custom_tool_call_output walking +# ------------------------------------------------------------------- + +def test_iter_message_text_walks_custom_tool_call_output(): + """custom_tool_call_output items should yield their output text.""" + data = { + "input": [ + {"type": "custom_tool_call_output", "output": "tool-secret"}, + ] + } + from litellm.proxy.guardrails._content_utils import iter_message_text + texts = list(iter_message_text(data)) + assert "tool-secret" in texts + + +def test_walk_user_text_redacts_custom_tool_call_output(): + """walk_user_text should rewrite text inside custom_tool_call_output.""" + data = { + "input": [ + {"type": "custom_tool_call_output", "output": "PII-data"}, + ] + } + count = walk_user_text(data, lambda t: t.replace("PII-data", "[MASKED]")) + assert count >= 1 + assert data["input"][0]["output"] == "[MASKED]" + + +def test_build_inspection_messages_custom_tool_call_output(): + """build_inspection_messages should include custom_tool_call_output text.""" + data = { + "input": [ + {"type": "custom_tool_call_output", "output": "custom-tool-leak"}, + ] + } + msgs = build_inspection_messages(data) + assert any("custom-tool-leak" in m["content"] for m in msgs) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 813a0c5e38f..f289148101a 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -14,6 +14,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( _ProxyDBLogger, _get_budget_reservation_from_metadata, + _should_track_cost_callback, _update_database_and_spend_counters, ) @@ -1177,3 +1178,88 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com" ) + + +@pytest.mark.parametrize( + "call_type, expected", + [ + ("pass_through_endpoint", True), + ("llm_passthrough_route", True), + ("allm_passthrough_route", True), + ("acompletion", False), + ("call_mcp_tool", False), + (None, False), + ], +) +def test_should_track_cost_callback_pass_through_without_owner(call_type, expected): + """Regression for LIT-3782: unauthenticated pass-through requests (auth=false) + carry no key/user/team/end-user, yet must still be tracked so they land in + LiteLLM_SpendLogs. Other call types with no owner stay untracked.""" + assert ( + _should_track_cost_callback( + user_api_key=None, + user_id=None, + team_id=None, + end_user_id=None, + call_type=call_type, + ) + is expected + ) + + +@pytest.mark.parametrize( + "call_type, expect_spend_log", + [ + ("pass_through_endpoint", True), + ("acompletion", False), + (None, False), + ], +) +@pytest.mark.asyncio +async def test_track_cost_callback_logs_unauthenticated_pass_through_request( + call_type, expect_spend_log +): + """Regression for LIT-3782: a pass-through request with auth=false reaches the + cost callback with no key/user/team/end-user. Before the fix the spend-log + write was skipped and the request never appeared in request/usage logs. It + must now be written for pass-through call types while other unauthenticated + calls remain skipped.""" + logger = _ProxyDBLogger() + + kwargs = { + "call_type": call_type, + "model": "unknown", + "litellm_params": {"metadata": {}}, + "standard_logging_object": { + "response_cost": 0.0, + "request_tags": None, + }, + "stream": False, + } + + with ( + patch( + "litellm.proxy.proxy_server.increment_spend_counters", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.proxy_server.update_cache", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging, + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( + 1 if expect_spend_log else 0 + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index dd9cbcf5232..81840745d0e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -608,3 +608,320 @@ class TestValidateFiniteSpend: with pytest.raises(HTTPException) as exc_info: validate_finite_spend(bad) assert exc_info.value.status_code == 400 + + +class TestValidateFiniteSpendErrorDetail: + """The 400 for non-finite spend must carry the exact {"error": } body.""" + + def test_rejection_detail_is_exact(self): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.common_utils import ( + validate_finite_spend, + ) + + with pytest.raises(HTTPException) as exc_info: + validate_finite_spend(float("nan")) + + assert exc_info.value.detail == { + "error": "spend must be a finite number. Received: nan" + } + + +class TestRequireCallerUserIdErrorDetail: + """The 403 for a service-account key must carry the exact error body.""" + + def test_rejection_detail_is_exact(self): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.common_utils import ( + require_caller_user_id_for_non_admin, + ) + + service_account_key = UserAPIKeyAuth( + user_id=None, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + with pytest.raises(HTTPException) as exc_info: + require_caller_user_id_for_non_admin(service_account_key) + + assert exc_info.value.detail == { + "error": "Service-account keys cannot query user analytics. Use a user-bound key, or call as a proxy admin." + } + + +class TestCheckPassthroughRoutesCallerPermission: + """Only proxy admins may set allowed_passthrough_routes (top-level or under + metadata); non-admins get a 403 naming the entity.""" + + def _non_admin(self): + return UserAPIKeyAuth( + user_id="u1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + + def test_top_level_routes_rejected_with_default_entity(self): + from fastapi import HTTPException + from pydantic import BaseModel + + from litellm.proxy.management_endpoints.common_utils import ( + _check_passthrough_routes_caller_permission, + ) + + class _RouteData(BaseModel): + allowed_passthrough_routes: list | None = None + metadata: dict | None = None + + data = _RouteData(allowed_passthrough_routes=["/v1/foo"]) + with pytest.raises(HTTPException) as exc_info: + _check_passthrough_routes_caller_permission(data, self._non_admin()) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == { + "error": "Only proxy admins can set `allowed_passthrough_routes` on a key." + } + + def test_metadata_routes_rejected_with_default_entity(self): + from fastapi import HTTPException + from pydantic import BaseModel + + from litellm.proxy.management_endpoints.common_utils import ( + _check_passthrough_routes_caller_permission, + ) + + class _RouteData(BaseModel): + allowed_passthrough_routes: list | None = None + metadata: dict | None = None + + data = _RouteData(metadata={"allowed_passthrough_routes": ["/v1/foo"]}) + with pytest.raises(HTTPException) as exc_info: + _check_passthrough_routes_caller_permission(data, self._non_admin()) + + assert exc_info.value.detail == { + "error": "Only proxy admins can set `metadata.allowed_passthrough_routes` on a key." + } + + def test_tolerates_data_missing_passthrough_and_metadata_fields(self): + from pydantic import BaseModel + + from litellm.proxy.management_endpoints.common_utils import ( + _check_passthrough_routes_caller_permission, + ) + + class _Bare(BaseModel): + unrelated: str = "x" + + assert ( + _check_passthrough_routes_caller_permission(_Bare(), self._non_admin()) + is None + ) + + +class TestIsUserOrgAdminForTeam: + """The caller must be looked up with its exact identity; a nulled or omitted + lookup argument would silently mis-resolve org-admin status.""" + + @pytest.mark.asyncio + async def test_get_user_object_called_with_caller_identity(self): + from litellm.proxy.management_endpoints.common_utils import ( + _is_user_org_admin_for_team, + ) + + team = LiteLLM_TeamTable( + team_id="t1", organization_id="org1", members_with_roles=[] + ) + key = UserAPIKeyAuth( + user_id="u1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + fake_prisma, fake_cache, fake_logging = MagicMock(), MagicMock(), MagicMock() + mock_get_user = AsyncMock(return_value=None) + + with patch( + "litellm.proxy.proxy_server.prisma_client", fake_prisma + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", fake_cache + ), patch( + "litellm.proxy.proxy_server.proxy_logging_obj", fake_logging + ), patch( + "litellm.proxy.auth.auth_checks.get_user_object", mock_get_user + ): + result = await _is_user_org_admin_for_team(key, team) + + assert result is False + mock_get_user.assert_awaited_once_with( + user_id="u1", + prisma_client=fake_prisma, + user_api_key_cache=fake_cache, + user_id_upsert=False, + proxy_logging_obj=fake_logging, + ) + + +class TestTeamMemberHasPermission: + def test_requires_caller_to_be_a_team_member(self): + from litellm.proxy.management_endpoints.common_utils import ( + _team_member_has_permission, + ) + + team = LiteLLM_TeamTable( + team_id="t1", + team_member_permissions=["/key/generate"], + members_with_roles=[Member(user_id="someone-else", role="user")], + ) + key = UserAPIKeyAuth( + user_id="u1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + assert _team_member_has_permission(key, team, "/key/generate") is False + + +class TestUserHasAdminPrivilegesGuard: + @pytest.mark.asyncio + async def test_no_user_lookup_when_prisma_is_none(self): + """With no DB the guard short-circuits before any user lookup.""" + auth = UserAPIKeyAuth( + user_id="user1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + mock_get_user = AsyncMock(return_value=None) + with patch("litellm.proxy.auth.auth_checks.get_user_object", mock_get_user): + result = await _user_has_admin_privileges( + user_api_key_dict=auth, prisma_client=None + ) + assert result is False + mock_get_user.assert_not_called() + + @pytest.mark.asyncio + async def test_org_admin_membership_grants_privileges(self): + """With DB + user_id present, an ORG_ADMIN membership yields True.""" + auth = UserAPIKeyAuth( + user_id="user1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + now = datetime.now(timezone.utc) + user_obj = LiteLLM_UserTable( + user_id="user1", + organization_memberships=[ + LiteLLM_OrganizationMembershipTable( + user_id="user1", + organization_id="org1", + user_role=LitellmUserRoles.ORG_ADMIN.value, + created_at=now, + updated_at=now, + ) + ], + ) + mock_get_user = AsyncMock(return_value=user_obj) + with patch("litellm.proxy.auth.auth_checks.get_user_object", mock_get_user): + result = await _user_has_admin_privileges( + user_api_key_dict=auth, prisma_client=MagicMock() + ) + assert result is True + + +class TestAdminCanInviteUserGuard: + @pytest.mark.asyncio + async def test_no_user_lookup_when_prisma_is_none(self): + auth = UserAPIKeyAuth( + user_id="admin1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + mock_get_user = AsyncMock(return_value=None) + with patch("litellm.proxy.auth.auth_checks.get_user_object", mock_get_user): + result = await admin_can_invite_user( + target_user_id="target1", + user_api_key_dict=auth, + prisma_client=None, + ) + assert result is False + mock_get_user.assert_not_called() + + @pytest.mark.asyncio + async def test_org_admin_can_invite_user_in_shared_org(self): + now = datetime.now(timezone.utc) + auth = UserAPIKeyAuth( + user_id="admin1", api_key="sk-x", user_role=LitellmUserRoles.INTERNAL_USER + ) + + def membership(role): + return LiteLLM_OrganizationMembershipTable( + user_id="x", + organization_id="org1", + user_role=role, + created_at=now, + updated_at=now, + ) + + admin_obj = LiteLLM_UserTable( + user_id="admin1", + organization_memberships=[membership(LitellmUserRoles.ORG_ADMIN.value)], + ) + target_obj = LiteLLM_UserTable( + user_id="target1", + organization_memberships=[membership(LitellmUserRoles.INTERNAL_USER.value)], + ) + mock_get_user = AsyncMock(side_effect=[admin_obj, target_obj]) + with patch("litellm.proxy.auth.auth_checks.get_user_object", mock_get_user): + result = await admin_can_invite_user( + target_user_id="target1", + user_api_key_dict=auth, + prisma_client=MagicMock(), + ) + assert result is True + + +class TestTeamAdminCanInviteUserQuery: + @pytest.mark.asyncio + async def test_find_many_queries_admin_teams_with_exact_where(self): + mock_prisma = MagicMock() + mock_auth = MagicMock() + mock_auth.user_id = "admin" + admin_user = LiteLLM_UserTable(user_id="admin", teams=["t1", "t2"]) + target_user = LiteLLM_UserTable(user_id="target", teams=["t2"]) + + def make_team(tid): + obj = MagicMock() + obj.team_id = tid + obj.model_dump = lambda: { + "team_id": tid, + "members_with_roles": [{"user_id": "admin", "role": "admin"}], + } + return obj + + find_many = AsyncMock(return_value=[make_team("t1"), make_team("t2")]) + mock_prisma.db.litellm_teamtable.find_many = find_many + + await _team_admin_can_invite_user( + user_api_key_dict=mock_auth, + admin_user_obj=admin_user, + target_user_obj=target_user, + prisma_client=mock_prisma, + ) + + find_many.assert_awaited_once_with(where={"team_id": {"in": ["t1", "t2"]}}) + + +class TestSetObjectMetadataFieldPremiumArg: + def test_premium_check_receives_the_field_name(self): + team = LiteLLM_TeamTable(team_id="t1", metadata={}) + with patch( + "litellm.proxy.management_endpoints.common_utils._premium_user_check" + ) as mock_premium: + _set_object_metadata_field(team, "guardrails", ["g1"]) + mock_premium.assert_called_once_with("guardrails") + + +class TestUpdateMetadataFieldMove: + def test_none_valued_field_is_not_moved_into_metadata(self): + """A None value must leave the field untouched (guard requires non-None).""" + from litellm.proxy.management_endpoints.common_utils import ( + _update_metadata_field, + ) + + updated_kv = {"guardrails": None} + _update_metadata_field(updated_kv=updated_kv, field_name="guardrails") + assert updated_kv == {"guardrails": None} + + def test_set_premium_field_is_moved_into_metadata(self): + updated_kv = {"guardrails": ["g1"]} + with patch( + "litellm.proxy.management_endpoints.common_utils._premium_user_check" + ): + _update_metadata_fields(updated_kv) + assert "guardrails" not in updated_kv + assert updated_kv["metadata"]["guardrails"] == ["g1"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index a57e18df6ef..bc463d5e75d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -405,3 +405,98 @@ class TestResolveModelForCostLookup: assert resolved_model == "azure/openai/gpt-5.3-codex" assert provider is None + + def test_returns_custom_llm_provider_on_base_model_path(self): + """base_model path: the custom_llm_provider from litellm_params is + returned as the second tuple element, unchanged.""" + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "my-azure-model", + "litellm_params": { + "model": "azure/my-deployment", + "base_model": "azure/gpt-4o", + "custom_llm_provider": "azure", + }, + "model_info": {"id": "test-id"}, + } + ] + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + resolved_model, provider = _resolve_model_for_cost_lookup("my-azure-model") + + assert resolved_model == "azure/gpt-4o" + assert provider == "azure" + + def test_returns_custom_llm_provider_on_resolved_model_path(self): + """resolved-model path (no base_model): the custom_llm_provider from + litellm_params is returned alongside litellm_params.model.""" + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "openai/gpt-4", + "custom_llm_provider": "openai", + }, + "model_info": {"id": "test-id"}, + } + ] + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + resolved_model, provider = _resolve_model_for_cost_lookup("gpt-4") + + assert resolved_model == "openai/gpt-4" + assert provider == "openai" + + def test_resolves_base_model_when_deployment_has_no_litellm_params(self): + """A deployment can omit litellm_params entirely; base_model from + model_info must still resolve (the .get default must be {} not None, + else the later litellm_params.get(...) raises and resolution is lost).""" + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "my-azure-model", + "model_info": {"base_model": "azure/gpt-4o"}, + } + ] + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + resolved_model, provider = _resolve_model_for_cost_lookup("my-azure-model") + + assert resolved_model == "azure/gpt-4o" + assert provider is None + + def test_resolves_model_when_deployment_has_no_model_info(self): + """A deployment can omit model_info entirely; litellm_params.model must + still resolve (the .get default must be {} not None, else the earlier + model_info.get(...) raises and resolution is lost).""" + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + } + ] + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + resolved_model, provider = _resolve_model_for_cost_lookup("gpt-4") + + assert resolved_model == "openai/gpt-4" + assert provider is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index db6d3489830..aa24b0199ab 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -529,6 +529,56 @@ async def test_key_generation_with_object_permission(monkeypatch): assert key_insert_calls[0]["data"].get("object_permission_id") == "objperm123" +@pytest.mark.asyncio +async def test_generate_key_debug_log_never_contains_raw_token(monkeypatch, caplog): + """Regression for LIT-4356: /key/generate must never emit the raw virtual key + to a logger, even for short keys that bypass the regex-based + SecretRedactionFilter.""" + import hashlib + import logging + + mock_prisma_client = AsyncMock() + mock_prisma_client.jsonify_object = lambda data: data + mock_prisma_client.db = MagicMock() + + async def _insert_data_side_effect(*args, **kwargs): + if kwargs.get("table_name") == "user": + return MagicMock(models=[], spend=0) + return MagicMock( + token="hashed_token_456", + litellm_budget_table=None, + object_permission=None, + ) + + mock_prisma_client.insert_data = AsyncMock(side_effect=_insert_data_side_effect) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_fn, + ) + + raw_key = "sk-short-secret" + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + await generate_key_fn( + data=GenerateKeyRequest(key=raw_key), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="user-1", + ), + ) + + log_text = "\n".join(record.getMessage() for record in caplog.records) + assert raw_key not in log_text + assert hashlib.sha256(raw_key.encode()).hexdigest() in log_text + + @pytest.mark.asyncio @pytest.mark.parametrize( "field,request_kwargs,expected_in_error", @@ -12272,6 +12322,55 @@ async def test_bulk_update_team_keys_team_member_no_permission(monkeypatch): mock.update_data.assert_not_called() +def test_handle_key_type_persists_key_type_and_derives_routes(): + """`handle_key_type` keeps `key_type` in the payload (so it is persisted on + the token) while still deriving the `allowed_routes` preset. Regression for + the UI showing scoped keys as "All Proxy Models": the frontend now reads the + persisted `key_type` instead of reverse-mapping the preset string.""" + from litellm.proxy._types import GenerateKeyRequest, LiteLLMKeyType + from litellm.proxy.management_endpoints.key_management_endpoints import ( + handle_key_type, + ) + + cases = { + LiteLLMKeyType.MANAGEMENT: ("management", ["management_routes"]), + LiteLLMKeyType.READ_ONLY: ("read_only", ["info_routes"]), + LiteLLMKeyType.LLM_API: ("llm_api", ["llm_api_routes"]), + } + for key_type, (expected_type, expected_routes) in cases.items(): + data = GenerateKeyRequest(key_type=key_type) + out = handle_key_type(data, {"key_type": key_type}) + assert out["key_type"] == expected_type + assert out["allowed_routes"] == expected_routes + + +def test_handle_key_type_default_persists_type_without_forcing_routes(): + """`default` is persisted but must not overwrite an explicit `allowed_routes` + (e.g. a SCIM key created with `["/scim/*"]` and no explicit key_type).""" + from litellm.proxy._types import GenerateKeyRequest, LiteLLMKeyType + from litellm.proxy.management_endpoints.key_management_endpoints import ( + handle_key_type, + ) + + data = GenerateKeyRequest(key_type=LiteLLMKeyType.DEFAULT) + out = handle_key_type(data, {"allowed_routes": ["/scim/*"], "key_type": LiteLLMKeyType.DEFAULT}) + assert out["key_type"] == "default" + assert out["allowed_routes"] == ["/scim/*"] + + +def test_handle_key_type_none_drops_key_type(): + """When no `key_type` is supplied the payload must not carry a `key_type` + entry, so old keys stay `null` and the frontend keeps its route fallback.""" + from litellm.proxy._types import GenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + handle_key_type, + ) + + data = GenerateKeyRequest(key_type=None) + out = handle_key_type(data, {"key_type": None}) + assert "key_type" not in out + + # ---- pydantic-layer validation ------------------------------------------- diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 59b08a7ec4e..f7b2df45a85 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -18,6 +18,7 @@ sys.path.insert( from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth from litellm.proxy._types import ( LiteLLM_BudgetTableFull, + LiteLLM_ModelTable, LiteLLM_OrganizationMembershipTable, LiteLLM_OrganizationTable, LiteLLM_OrganizationTableWithMembers, @@ -30,6 +31,7 @@ from litellm.proxy._types import ( ProxyErrorTypes, ProxyException, TeamMemberAddRequest, + UpdateTeamRequest, ) from litellm.proxy.management_endpoints.team_endpoints import ( user_api_key_auth, # Assuming this dependency is needed @@ -40,6 +42,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _persist_deleted_team_records, _save_deleted_team_records, _transform_teams_to_deleted_records, + _update_model_table, _validate_and_populate_member_user_info, _verify_team_access, delete_team, @@ -9436,6 +9439,96 @@ async def test_team_info_forwards_key_limit_to_get_data(): assert mock_prisma.get_data.await_args.kwargs["limit"] == 7 +@pytest.mark.asyncio +async def test_team_info_returns_model_aliases(): + """/team/info must join LiteLLM_ModelTable so the response exposes the team's + current model aliases; without the ``litellm_model_table`` include the field + comes back null and the Admin UI can never display them. + """ + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + team_row = LiteLLM_TeamTable( + team_id="team-1", + litellm_model_table=LiteLLM_ModelTable( + id=1, + model_aliases={"gpt-4o": "gpt-4o-team-1"}, + created_by="admin", + updated_by="admin", + ), + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.get_data = AsyncMock(return_value=[]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object( + team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[]) + ), + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + include = mock_prisma.db.litellm_teamtable.find_unique.await_args.kwargs["include"] + assert include["litellm_model_table"] is True + + litellm_model_table = response["team_info"].litellm_model_table + assert litellm_model_table is not None + assert litellm_model_table.model_aliases == {"gpt-4o": "gpt-4o-team-1"} + + +@pytest.mark.asyncio +async def test_update_model_table_clears_aliases_with_empty_map(): + """``model_aliases={}`` on /team/update must persist an empty map (json.dumps({})) + so existing aliases are cleared, while ``model_aliases=None`` must be a no-op that + leaves the model table untouched. + """ + mock_prisma = MagicMock() + mock_prisma.db.litellm_modeltable.create = AsyncMock() + mock_prisma.db.litellm_modeltable.upsert = AsyncMock( + return_value=MagicMock(id="model-123") + ) + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin" + ) + + returned_model_id = await _update_model_table( + data=UpdateTeamRequest(team_id="team-1", model_aliases={}), + model_id="model-123", + prisma_client=mock_prisma, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="default_user_id", + ) + + mock_prisma.db.litellm_modeltable.upsert.assert_awaited_once() + upsert_kwargs = mock_prisma.db.litellm_modeltable.upsert.await_args.kwargs + assert upsert_kwargs["where"] == {"id": "model-123"} + assert upsert_kwargs["data"]["update"]["model_aliases"] == json.dumps({}) + assert upsert_kwargs["data"]["create"]["model_aliases"] == json.dumps({}) + assert returned_model_id == "model-123" + + mock_prisma.db.litellm_modeltable.create.reset_mock() + mock_prisma.db.litellm_modeltable.upsert.reset_mock() + + noop_model_id = await _update_model_table( + data=UpdateTeamRequest(team_id="team-1", model_aliases=None), + model_id="model-123", + prisma_client=mock_prisma, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="default_user_id", + ) + + mock_prisma.db.litellm_modeltable.create.assert_not_called() + mock_prisma.db.litellm_modeltable.upsert.assert_not_called() + assert noop_model_id == "model-123" + + class TestEmitTeamMembersMetric: """The _emit_team_members_metric seam between the team handlers and Prometheus.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 045e15f8b8b..2a4e2ed6b25 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -461,6 +461,53 @@ async def test_get_group_ids_from_service_principal_uses_configured_graph_endpoi ] +@pytest.mark.asyncio +async def test_get_group_ids_from_service_principal_paginates_through_all_pages(): + # Arrange + page_one = { + "@odata.nextLink": "https://graph.microsoft.com/v1.0/servicePrincipals/sp-123/appRoleAssignedTo?$skiptoken=page2", + "value": [ + { + "principalType": "Group", + "principalId": "group-on-page-1", + "principalDisplayName": "Group On Page 1", + } + ], + } + page_two = { + "value": [ + { + "principalType": "Group", + "principalId": "group-on-page-2", + "principalDisplayName": "Group On Page 2", + } + ], + } + responses = [page_one, page_two] + + async def mock_get(url, *args, **kwargs): + mock = MagicMock() + mock.json.return_value = responses.pop(0) + return mock + + async_client = MagicMock() + async_client.get = mock_get + + # Act + group_ids, teams = await MicrosoftSSOHandler.get_group_ids_from_service_principal( + service_principal_id="sp-123", + async_client=async_client, + access_token="mock_token", + ) + + # Assert + assert group_ids == ["group-on-page-1", "group-on-page-2"] + assert [team["principalId"] for team in teams] == [ + "group-on-page-1", + "group-on-page-2", + ] + + def test_get_group_ids_from_graph_api_response(): # Arrange mock_response = MicrosoftGraphAPIUserGroupResponse( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 3da103683ba..540f017ee88 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -576,6 +576,83 @@ async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_ await release_budget_reservation(reservation) +@pytest.mark.asyncio +async def test_should_reserve_user_budget_counter_for_team_key(spend_counter_state): + """A user's personal budget must be reserved even when the key belongs to a team. + + Regression for GitHub issue #12905: previously the reservation path skipped the + user spend counter whenever the key had a team, so a team key could overshoot the + user's personal max_budget under concurrency. + """ + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-user-on-team", + spend=0.0, + user_id="user-on-team", + team_id="team-no-budget", + ) + team_object = LiteLLM_TeamTable(team_id="team-no-budget", spend=0.0, max_budget=None) + user_object = LiteLLM_UserTable(user_id="user-on-team", spend=0.0, max_budget=5.0) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.3, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=user_object, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team") == pytest.approx(0.3) + + await release_budget_reservation(reservation) + + +@pytest.mark.asyncio +async def test_should_skip_user_budget_counter_for_team_key_when_flag_set(spend_counter_state): + """skip_user_budget_on_team_key=True restores the legacy behavior where a user's + personal budget is not reserved for a team key.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-user-on-team-skip", + spend=0.0, + user_id="user-on-team-skip", + team_id="team-no-budget-skip", + ) + team_object = LiteLLM_TeamTable(team_id="team-no-budget-skip", spend=0.0, max_budget=None) + user_object = LiteLLM_UserTable(user_id="user-on-team-skip", spend=0.0, max_budget=5.0) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.3, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=user_object, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + skip_user_budget_on_team_key=True, + ) + + assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team-skip") is None + + await release_budget_reservation(reservation) + + @pytest.mark.asyncio async def test_should_seed_org_counter_from_with_budget_cache(spend_counter_state): counter_cache, key_cache = spend_counter_state diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 603d5cc15b7..2f0924e9192 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8649,6 +8649,38 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): app.dependency_overrides.clear() +def test_get_config_list_includes_skip_user_budget_on_team_key(monkeypatch): + """Related to #12905: the opt-out flag must be discoverable via /config/list so + it renders as a Boolean toggle on the Admin UI General Settings table. This + requires both the ConfigGeneralSettings field and the allowed_args entry.""" + import types + from unittest.mock import AsyncMock, MagicMock + + from fastapi.testclient import TestClient + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.proxy_server import app + + mock_prisma = MagicMock() + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + client = TestClient(app) + resp = client.get("/config/list", params={"config_type": "general_settings"}) + assert resp.status_code == 200, resp.text + fields = {item["field_name"]: item for item in resp.json()} + assert "skip_user_budget_on_team_key" in fields + assert fields["skip_user_budget_on_team_key"]["field_type"] == "Boolean" + finally: + app.dependency_overrides.clear() + + def test_get_config_list_includes_budget_exceeded_throttle_percentage(monkeypatch): """The throttle fraction is a litellm_settings scalar surfaced on the General Settings table as a Float field so it sits with the other global limits; it diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py index 4e547b81acc..dd241397e87 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py @@ -10,6 +10,7 @@ from __future__ import annotations import hashlib import json +import logging from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -103,6 +104,35 @@ async def test_insert_data_user_organization_fk_raises_400( assert raised.status_code == 400 +@pytest.mark.asyncio +async def test_insert_data_debug_log_hashes_token( + prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture +) -> None: + """Regression for LIT-4356: the raw virtual key must never reach a logger, + even for short/nonstandard key formats that bypass the regex-based + SecretRedactionFilter.""" + token = "sk-short-secret" + expected_hash = hashlib.sha256(token.encode()).hexdigest() + prisma_client.db.litellm_verificationtoken.upsert = AsyncMock(return_value=SimpleNamespace(token=expected_hash)) + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + await prisma_client.insert_data(data={"token": token, "key_alias": "redaction-repro"}, table_name="key") + log_text = "\n".join(record.getMessage() for record in caplog.records) + assert token not in log_text + assert expected_hash in log_text + + +@pytest.mark.asyncio +async def test_insert_data_debug_log_tolerates_none_token( + prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture +) -> None: + """A None token must not crash the redacting debug log added for LIT-4356.""" + prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=SimpleNamespace(user_id="u1")) + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + result = await prisma_client.insert_data(data={"user_id": "u1", "token": None}, table_name="user") + assert result.user_id == "u1" + assert any("insert_data" in record.getMessage() for record in caplog.records) + + @pytest.mark.asyncio async def test_insert_data_logs_and_raises_generic_error( prisma_client: PrismaClient, diff --git a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py index cdace5f6327..24edf12fffe 100644 --- a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py @@ -39,8 +39,13 @@ def _output_item_added_chunk(): return SimpleNamespace(type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED) -def _completed_chunk(output): - response = ResponsesAPIResponse(id="resp-1", created_at=0, output=output) +def _created_chunk(response_id: str): + response = ResponsesAPIResponse(id=response_id, created_at=0, output=[]) + return SimpleNamespace(type=ResponsesAPIStreamEvents.RESPONSE_CREATED, response=response) + + +def _completed_chunk(output, response_id: str = "resp-1"): + response = ResponsesAPIResponse(id=response_id, created_at=0, output=output) return SimpleNamespace(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=response) @@ -52,12 +57,12 @@ def _text_message(text: str): return {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]} -def _tool_call_stream(call_id: str, tool_name: str) -> _FakeAsyncStream: - return _FakeAsyncStream([_completed_chunk([_function_call(call_id, tool_name)])]) +def _tool_call_stream(call_id: str, tool_name: str, response_id: str = "resp-1") -> _FakeAsyncStream: + return _FakeAsyncStream([_completed_chunk([_function_call(call_id, tool_name)], response_id=response_id)]) -def _text_only_stream(text: str) -> _FakeAsyncStream: - return _FakeAsyncStream([_completed_chunk([_text_message(text)])]) +def _text_only_stream(text: str, response_id: str = "resp-1") -> _FakeAsyncStream: + return _FakeAsyncStream([_completed_chunk([_text_message(text)], response_id=response_id)]) def _mock_mcp_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: @@ -170,3 +175,85 @@ async def test_tool_call_rounds_are_capped(monkeypatch): for call in aresponses_mock.call_args_list[:-1]: assert "tools" in call.kwargs assert "tools" not in aresponses_mock.call_args_list[-1].kwargs + + +@pytest.mark.asyncio +async def test_continuation_id_is_final_round_not_interim_tool_call(monkeypatch): + """ + Regression test for the broken `previous_response_id` continuation after a + gateway-executed MCP tool call. Each auto-execute round is a distinct + upstream response: the interim round holds only the model's function_call + (no tool output), the final round holds the answer. The client must be + handed the FINAL round's response id, because that is the one whose stored + chain includes the function_call_output. Pinning every event to the interim + round's id made the next turn continue from a response with a dangling + function_call, which the provider rejects with + "No tool output found for function call ...". + """ + _mock_mcp_environment(monkeypatch) + + aresponses_mock = AsyncMock(side_effect=[_text_only_stream("The first item is Alpha.", response_id="resp-final")]) + monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) + + iterator = _make_iterator( + [ + _created_chunk("resp-interim"), + _output_item_added_chunk(), + _completed_chunk([_function_call("call_1", "read_wiki_contents")], response_id="resp-interim"), + ] + ) + + chunks = [chunk async for chunk in iterator] + completed = [c for c in chunks if getattr(c, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED] + + assert completed[-1].response.output[0]["content"][0]["text"] == "The first item is Alpha." + assert completed[-1].response.id == "resp-final" + assert completed[-1].response.id != "resp-interim" + + +@pytest.mark.asyncio +async def test_follow_up_call_failure_emits_terminal_error_event(monkeypatch): + """ + Regression test for the silent-swallow path: when the follow-up LLM call + raises, the stream must emit a terminal `error` event instead of ending + silently after the tool events (which surfaced to clients as a successful + but empty completion). + """ + _mock_mcp_environment(monkeypatch) + + aresponses_mock = AsyncMock(side_effect=RuntimeError("boom from provider")) + monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) + + iterator = _make_iterator( + [ + _output_item_added_chunk(), + _completed_chunk([_function_call("call_1", "read_wiki_contents")]), + ] + ) + + chunks = [chunk async for chunk in iterator] + error_events = [c for c in chunks if getattr(c, "type", None) == ResponsesAPIStreamEvents.ERROR] + + assert len(error_events) == 1 + assert error_events[0].error.type == "mcp_gateway_error" + assert "boom from provider" in error_events[0].error.message + + +@pytest.mark.asyncio +async def test_initial_call_failure_is_stashed_for_eager_reraise(monkeypatch): + """ + Regression test: a failing initial LLM call must be stashed as + `_initial_creation_error` so `aresponses_api_with_mcp` can re-raise it as a + real 4xx before any SSE bytes are written, instead of returning HTTP 200 + with an empty stream. + """ + _mock_mcp_environment(monkeypatch) + + aresponses_mock = AsyncMock(side_effect=RuntimeError("initial boom")) + monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) + + iterator = _make_iterator([_output_item_added_chunk()]) + await iterator._create_initial_response_iterator() + + assert iterator._initial_creation_error is not None + assert "initial boom" in str(iterator._initial_creation_error) diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py index 93c4db90dad..cbf5635a5ae 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py @@ -7,9 +7,6 @@ from litellm.router_strategy.adaptive_router import adaptive_router as ar_module import pytest from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter -from litellm.router_strategy.adaptive_router.config import ( - OWNER_CACHE_TTL_SECONDS, -) from litellm.router_strategy.adaptive_router.signals import Turn from litellm.types.router import ( AdaptiveRouterConfig, @@ -22,9 +19,7 @@ def _make_router() -> AdaptiveRouter: cfg = AdaptiveRouterConfig(available_models=["fast", "smart"]) prefs = { "fast": AdaptiveRouterPreferences(quality_tier=1, strengths=[]), - "smart": AdaptiveRouterPreferences( - quality_tier=3, strengths=[RequestType.CODE_GENERATION] - ), + "smart": AdaptiveRouterPreferences(quality_tier=3, strengths=[RequestType.CODE_GENERATION]), } costs = {"fast": 0.0001, "smart": 0.001} return AdaptiveRouter( @@ -58,85 +53,6 @@ async def test_pick_model_min_quality_tier_filter_raises_when_no_eligible(): await r.pick_model(RequestType.GENERAL, min_quality_tier=4) -@pytest.mark.asyncio -async def test_pick_model_is_stateless_no_owner_cache_writes(): - """pick_model must not touch the owner cache — that's gated post-call.""" - r = _make_router() - for _ in range(5): - await r.pick_model(RequestType.GENERAL) - assert r._owner_cache == {} - - -# ---- claim_or_check_owner ----------------------------------------------- - - -def test_claim_or_check_owner_first_call_claims_and_returns_true(monkeypatch): - r = _make_router() - monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0) - - assert r.claim_or_check_owner("sess-A", "fast") is True - assert r._owner_cache["sess-A"] == ("fast", 1_000.0 + OWNER_CACHE_TTL_SECONDS) - assert r._skipped_updates_total == 0 - - -def test_claim_or_check_owner_same_model_returns_true_without_extending_ttl( - monkeypatch, -): - r = _make_router() - monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0) - r.claim_or_check_owner("sess-A", "fast") - original_expiry = r._owner_cache["sess-A"][1] - - monkeypatch.setattr(ar_module.time, "time", lambda: 1_500.0) - assert r.claim_or_check_owner("sess-A", "fast") is True - # No extension on hit — owner cache snapshots the first claim. - assert r._owner_cache["sess-A"][1] == original_expiry - - -def test_claim_or_check_owner_mismatch_skips_and_increments_counter(monkeypatch): - r = _make_router() - monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0) - r.claim_or_check_owner("sess-A", "fast") - - assert r.claim_or_check_owner("sess-A", "smart") is False - assert r._skipped_updates_total == 1 - # Owner unchanged. - assert r._owner_cache["sess-A"][0] == "fast" - - -def test_claim_or_check_owner_expired_owner_reclaims_for_new_model(monkeypatch): - r = _make_router() - monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0) - r.claim_or_check_owner("sess-A", "fast") - - monkeypatch.setattr( - ar_module.time, "time", lambda: 1_000.0 + OWNER_CACHE_TTL_SECONDS + 1 - ) - assert r.claim_or_check_owner("sess-A", "smart") is True - assert r._owner_cache["sess-A"][0] == "smart" - # Reclaim isn't a skip. - assert r._skipped_updates_total == 0 - - -def test_owner_cache_evicts_expired_entries_when_threshold_crossed(monkeypatch): - """Past _OWNER_CACHE_SWEEP_THRESHOLD live entries, new claims sweep stale.""" - r = _make_router() - monkeypatch.setattr(ar_module, "_OWNER_CACHE_SWEEP_THRESHOLD", 5) - monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0) - for i in range(5): - r.claim_or_check_owner(f"old-{i}", "fast") - assert len(r._owner_cache) == 5 - - # Jump past TTL so all "old-*" entries are now expired. - monkeypatch.setattr( - ar_module.time, "time", lambda: 1_000.0 + OWNER_CACHE_TTL_SECONDS + 1 - ) - r.claim_or_check_owner("new-1", "fast") - # Sweep ran -> only the new entry remains. - assert "new-1" in r._owner_cache - assert all(k.startswith("new-") for k in r._owner_cache) - - # ---- record_turn -------------------------------------------------------- @@ -185,9 +101,7 @@ async def test_record_turn_satisfaction_increments_alpha(): # Prime with 2 prior turns to clear the MIN_TURNS_FOR_CLEAN_CREDIT gate. # Use distinct content to avoid incidentally firing stagnation/misalignment. priming_turns = [ - Turn( - user_content="alpha bravo charlie", assistant_content="delta echo foxtrot" - ), + Turn(user_content="alpha bravo charlie", assistant_content="delta echo foxtrot"), Turn( user_content="golf hotel india juliet", assistant_content="kilo lima mike november", @@ -232,6 +146,128 @@ async def test_record_turn_failure_increments_beta(): assert cell_after.alpha == pytest.approx(cell_before.alpha) +@pytest.mark.asyncio +async def test_record_turn_detects_exhaustion_in_tool_results(): + r = _make_router() + + delta = await r.record_turn( + session_id="exhausted", + model_name="smart", + request_type=RequestType.GENERAL, + turn=Turn(tool_results=[{"content": "rate limit exceeded"}]), + ) + + assert delta.exhaustion == 1 + assert r._session_states[("exhausted", "smart")].exhaustion_count == 1 + + +@pytest.mark.asyncio +async def test_record_turn_attributes_user_feedback_to_previous_response_model(): + r = _make_router() + fast_before = r._cells[(RequestType.CODE_GENERATION, "fast")] + smart_before = r._cells[(RequestType.GENERAL, "smart")] + + await r.record_turn( + session_id="feedback-switch", + model_name="fast", + request_type=RequestType.CODE_GENERATION, + turn=Turn( + user_content="fix this python retry bug", + assistant_content="clear the cache on every retry", + ), + ) + await r.record_turn( + session_id="feedback-switch", + model_name="smart", + request_type=RequestType.GENERAL, + turn=Turn( + user_content="the python fix is still broken", + assistant_content="keep successful cache entries", + ), + ) + + fast_after = r._cells[(RequestType.CODE_GENERATION, "fast")] + smart_after = r._cells[(RequestType.GENERAL, "smart")] + assert fast_after.beta == pytest.approx(fast_before.beta + 1.0) + assert smart_after.beta == pytest.approx(smart_before.beta) + snapshot = await r.get_state_snapshot() + assert snapshot["feedback_attributed_total"] == 1 + assert snapshot["cross_model_feedback_total"] == 1 + assert snapshot["feedback_without_context_total"] == 0 + + +@pytest.mark.asyncio +async def test_record_turn_attributes_satisfaction_to_previous_response_model(): + r = _make_router() + await r.record_turn( + session_id="satisfaction-switch", + model_name="smart", + request_type=RequestType.CODE_GENERATION, + turn=Turn( + user_content="write a python retry helper", + assistant_content="first draft", + ), + ) + await r.record_turn( + session_id="satisfaction-switch", + model_name="fast", + request_type=RequestType.CODE_GENERATION, + turn=Turn( + user_content="add exponential backoff to the python helper", + assistant_content="updated draft", + ), + ) + fast_before = r._cells[(RequestType.CODE_GENERATION, "fast")] + smart_before = r._cells[(RequestType.GENERAL, "smart")] + + await r.record_turn( + session_id="satisfaction-switch", + model_name="smart", + request_type=RequestType.GENERAL, + turn=Turn( + user_content="thanks, that worked", + assistant_content="glad to help", + ), + ) + + fast_after = r._cells[(RequestType.CODE_GENERATION, "fast")] + smart_after = r._cells[(RequestType.GENERAL, "smart")] + assert fast_after.alpha == pytest.approx(fast_before.alpha + 1.0) + assert smart_after.alpha == pytest.approx(smart_before.alpha) + + +@pytest.mark.asyncio +async def test_record_turn_bounds_feedback_contexts_and_evicts_least_recent_session(): + r = _make_router() + context_limit = ar_module._FEEDBACK_CONTEXT_MAX_ENTRIES + + for index in range(context_limit): + await r.record_turn( + session_id=f"session-{index}", + model_name="fast", + request_type=RequestType.GENERAL, + turn=Turn(user_content="question", assistant_content="answer"), + ) + + await r.record_turn( + session_id="session-0", + model_name="fast", + request_type=RequestType.GENERAL, + turn=Turn(user_content="follow up", assistant_content="updated answer"), + ) + await r.record_turn( + session_id="overflow", + model_name="fast", + request_type=RequestType.GENERAL, + turn=Turn(user_content="question", assistant_content="answer"), + ) + + assert len(r._feedback_contexts) == context_limit + assert "session-0" in r._feedback_contexts + assert "session-1" not in r._feedback_contexts + assert "overflow" in r._feedback_contexts + + @pytest.mark.asyncio async def test_load_state_from_db_overrides_cold_start(): r = _make_router() @@ -270,9 +306,7 @@ async def test_load_state_from_db_handles_unknown_request_type(): good_row.beta = 3.0 prisma = MagicMock() - prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock( - return_value=[bad_row, good_row] - ) + prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[bad_row, good_row]) await r.load_state_from_db(prisma) # Unknown skipped; good applied. diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py index 9786832b4ae..3071f916ef1 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py @@ -110,23 +110,6 @@ async def test_pick_record_flush_full_cycle(): assert session_call.kwargs["data"]["create"]["model_name"] == chosen -@pytest.mark.asyncio -async def test_owner_cache_pins_attribution_to_first_picked_model(): - """First call claims ownership; matching model returns True, mismatch False.""" - router = _make_router() - chosen = await router.pick_model(RequestType.GENERAL) - assert router.claim_or_check_owner("sess-own", chosen) is True - - # Same model on later turns keeps attributing. - for _ in range(5): - assert router.claim_or_check_owner("sess-own", chosen) is True - - # A different model on a later turn is rejected. - other = "gpt-4o" if chosen == "gpt-4o-mini" else "gpt-4o-mini" - assert router.claim_or_check_owner("sess-own", other) is False - assert router._skipped_updates_total == 1 - - @pytest.mark.asyncio async def test_pick_model_returns_valid_models_without_error(): router = _make_router() diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_hooks.py b/tests/test_litellm/router_strategy/adaptive_router/test_hooks.py index a2b85f2ce53..ad61f43c5a0 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_hooks.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_hooks.py @@ -16,10 +16,9 @@ from litellm.router_strategy.adaptive_router.hooks import ( from litellm.router_strategy.adaptive_router.signals import Turn -def _make_hook(claim: bool = True) -> AdaptiveRouterPostCallHook: +def _make_hook() -> AdaptiveRouterPostCallHook: fake_router = MagicMock() fake_router.record_turn = AsyncMock() - fake_router.claim_or_check_owner = MagicMock(return_value=claim) return AdaptiveRouterPostCallHook(adaptive_router=fake_router) @@ -151,7 +150,24 @@ async def test_hook_skips_when_below_signal_gate(): kwargs = _kwargs(messages=short) await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0) hook.adaptive_router.record_turn.assert_not_awaited() - hook.adaptive_router.claim_or_check_owner.assert_not_called() + + +@pytest.mark.asyncio +async def test_hook_tracks_short_conversation_with_explicit_session_id(): + hook = _make_hook() + kwargs = _kwargs( + messages=[{"role": "user", "content": "hi"}], + extra_litellm_params={"litellm_session_id": "explicit-short"}, + ) + await hook.async_log_success_event( + kwargs, + _resp_with_content("hello"), + 0.0, + 1.0, + ) + assert hook.adaptive_router.record_turn.await_args.kwargs["session_id"] == ( + "explicit-short" + ) @pytest.mark.asyncio @@ -168,22 +184,19 @@ async def test_hook_skips_when_chosen_model_missing_from_metadata(): kwargs = _kwargs(chosen=None) await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0) hook.adaptive_router.record_turn.assert_not_awaited() - hook.adaptive_router.claim_or_check_owner.assert_not_called() @pytest.mark.asyncio -async def test_hook_skips_when_owner_cache_mismatch(): - """A different model owns this conversation -> no attribution.""" - hook = _make_hook(claim=False) +async def test_hook_records_when_model_changes(): + hook = _make_hook() kwargs = _kwargs(chosen="fast") await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0) - hook.adaptive_router.claim_or_check_owner.assert_called_once() - hook.adaptive_router.record_turn.assert_not_awaited() + hook.adaptive_router.record_turn.assert_awaited_once() @pytest.mark.asyncio -async def test_hook_records_turn_when_owner_claims(): - hook = _make_hook(claim=True) +async def test_hook_records_turn(): + hook = _make_hook() kwargs = _kwargs(chosen="smart", messages=_long_messages("ask")) await hook.async_log_success_event( kwargs, _resp_with_content("answer here"), 0.0, 1.0 @@ -205,8 +218,6 @@ async def test_hook_uses_explicit_session_id_when_provided(): extra_litellm_params={"litellm_session_id": "explicit-sess"}, ) await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0) - args, _ = hook.adaptive_router.claim_or_check_owner.call_args - assert args[0] == "explicit-sess" assert hook.adaptive_router.record_turn.await_args.kwargs["session_id"] == ( "explicit-sess" ) diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py b/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py index 753a449791b..d6d89c8e811 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py @@ -1,7 +1,6 @@ """Tests for the GET /adaptive_router/state introspection endpoint and the underlying `AdaptiveRouter.get_state_snapshot()` helper.""" -import time from unittest.mock import MagicMock import pytest @@ -9,7 +8,7 @@ from fastapi import HTTPException from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter -from litellm.router_strategy.adaptive_router.bandit import BanditCell, apply_delta +from litellm.router_strategy.adaptive_router.bandit import apply_delta from litellm.types.router import ( AdaptiveRouterConfig, AdaptiveRouterPreferences, @@ -47,8 +46,6 @@ async def test_get_state_snapshot_returns_cell_per_request_type_per_model(): assert snap["available_models"] == ["fast", "smart"] assert snap["weights"] == {"quality": 0.7, "cost": 0.3} assert snap["model_costs"] == {"fast": 0.0001, "smart": 0.001} - assert snap["owner_cache_live"] == 0 - assert snap["skipped_updates_total"] == 0 assert set(snap["queue"].keys()) == { "state_pending", "session_pending", @@ -95,26 +92,6 @@ async def test_get_state_snapshot_quality_mean_matches_alpha_over_total(): assert cell["quality_mean"] == pytest.approx(expected_mean) -@pytest.mark.asyncio -async def test_get_state_snapshot_counts_only_live_owner_cache_entries(): - r = _make_router() - now = time.time() - r._owner_cache["live-1"] = ("fast", now + 3600) - r._owner_cache["live-2"] = ("smart", now + 3600) - r._owner_cache["expired-1"] = ("fast", now - 1) - - snap = await r.get_state_snapshot() - assert snap["owner_cache_live"] == 2 - - -@pytest.mark.asyncio -async def test_get_state_snapshot_exposes_skipped_updates_total(): - r = _make_router() - r._skipped_updates_total = 7 - snap = await r.get_state_snapshot() - assert snap["skipped_updates_total"] == 7 - - # ---- endpoint -------------------------------------------------------- diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index f47c19b2baa..f7c9f343f80 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -14,13 +14,12 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from pydantic import ValidationError -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path import litellm from litellm import Router from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache from litellm.router_strategy.complexity_router.complexity_router import ( ComplexityRouter, DimensionScore, @@ -130,9 +129,7 @@ class TestTokenScoring: tier, score, signals = complexity_router.classify("What is Python?") # Should be classified as SIMPLE due to short length and simple indicator assert tier == ComplexityTier.SIMPLE - assert any("short" in s.lower() for s in signals) or any( - "simple" in s.lower() for s in signals - ) + assert any("short" in s.lower() for s in signals) or any("simple" in s.lower() for s in signals) def test_long_prompt_positive_score(self, complexity_router): """Long prompts should get positive scores (complex indicator).""" @@ -143,9 +140,7 @@ class TestTokenScoring: tier, score, signals = complexity_router.classify(long_prompt) # Should have positive score and detect long token count or technical terms assert score > 0, f"Expected positive score for long prompt, got {score}" - assert any("long" in s.lower() for s in signals) or any( - "technical" in s.lower() for s in signals - ) + assert any("long" in s.lower() for s in signals) or any("technical" in s.lower() for s in signals) class TestCodePresenceScoring: @@ -220,9 +215,7 @@ class TestMultiStepPatterns: def test_first_then_pattern(self, complexity_router): """'First...then' patterns should increase complexity.""" - prompt = ( - "First analyze the data, then create a visualization, then write a report" - ) + prompt = "First analyze the data, then create a visualization, then write a report" tier, score, signals = complexity_router.classify(prompt) assert any("multi-step" in s.lower() for s in signals) @@ -266,9 +259,7 @@ class TestTierAssignment: ) tier, score, signals = complexity_router.classify(prompt) # Should detect technical terms - assert any( - "technical" in s.lower() for s in signals - ), f"Expected technical signals, got {signals}" + assert any("technical" in s.lower() for s in signals), f"Expected technical signals, got {signals}" # Score should be positive due to technical content assert score > 0, f"Expected positive score, got {score}" @@ -315,6 +306,36 @@ class TestModelSelection: model = router.get_model_for_tier(ComplexityTier.SIMPLE) assert model == "fallback-model" + def test_get_model_for_tier_list_random_choice(self, mock_router_instance): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": ["cheap", "premium"], "MEDIUM": "mid"}, + "default_model": "mid", + }, + ) + pool = ["cheap", "premium"] + with patch( + "litellm.router_strategy.complexity_router.complexity_router.random.choice", + return_value="premium", + ) as choice: + assert router.get_model_for_tier(ComplexityTier.SIMPLE) == "premium" + choice.assert_called_once_with(pool) + assert router.get_model_for_tier(ComplexityTier.MEDIUM) == "mid" + + def test_get_model_for_tier_empty_pool_raises(self, mock_router_instance): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": []}, + "default_model": "mid", + }, + ) + with pytest.raises(ValueError, match="Empty model pool for tier SIMPLE"): + router.get_model_for_tier(ComplexityTier.SIMPLE) + class TestPreRoutingHook: """Test the async_pre_routing_hook method.""" @@ -438,13 +459,9 @@ class TestConfigOverrides: complexity_router_config=config, ) # With very low thresholds, even neutral prompts should be COMPLEX or higher - tier, score, signals = router.classify( - "Explain how HTTP works with REST APIs and distributed systems" - ) + tier, score, signals = router.classify("Explain how HTTP works with REST APIs and distributed systems") # With boundaries this low, should be at least MEDIUM (anything above -0.5) - assert ( - tier != ComplexityTier.SIMPLE - ), f"Expected non-SIMPLE tier, got {tier} with score {score}" + assert tier != ComplexityTier.SIMPLE, f"Expected non-SIMPLE tier, got {tier} with score {score}" def test_custom_token_thresholds(self, mock_router_instance): """Test custom token thresholds work correctly.""" @@ -469,9 +486,7 @@ class TestConfigOverrides: long_prompt = "This is a test prompt " * 30 # ~120 tokens tier, score, signals = router.classify(long_prompt) # Should get token length signal indicating "long" - assert any( - "long" in s.lower() if s else False for s in signals - ), f"Expected 'long' signal, got {signals}" + assert any("long" in s.lower() if s else False for s in signals), f"Expected 'long' signal, got {signals}" class TestCustomTechnicalKeywords: @@ -486,9 +501,7 @@ class TestCustomTechnicalKeywords: ) assert router.technical_keywords == DEFAULT_TECHNICAL_KEYWORDS + ["udp", "kafka"] - def test_custom_keywords_appended_to_technical_keywords_override( - self, mock_router_instance - ): + def test_custom_keywords_appended_to_technical_keywords_override(self, mock_router_instance): """Custom keywords should be appended to a technical_keywords override.""" router = ComplexityRouter( model_name="test-router", @@ -505,9 +518,7 @@ class TestCustomTechnicalKeywords: router = ComplexityRouter( model_name="test-router", litellm_router_instance=mock_router_instance, - complexity_router_config={ - "custom_technical_keywords": ["TCP", "udp", "UDP", "kafka"] - }, + complexity_router_config={"custom_technical_keywords": ["TCP", "udp", "UDP", "kafka"]}, ) lowered = [kw.lower() for kw in router.technical_keywords] assert lowered == [kw.lower() for kw in DEFAULT_TECHNICAL_KEYWORDS] + [ @@ -530,9 +541,7 @@ class TestCustomTechnicalKeywords: assert router_absent.technical_keywords == DEFAULT_TECHNICAL_KEYWORDS assert router_none.technical_keywords == DEFAULT_TECHNICAL_KEYWORDS - def test_prompt_with_only_custom_keywords_scores_technical( - self, mock_router_instance, basic_config - ): + def test_prompt_with_only_custom_keywords_scores_technical(self, mock_router_instance, basic_config): """A prompt matching only custom keywords should score higher on technicalTerms.""" prompt = "Configure udp multicast between kafka brokers" baseline_router = ComplexityRouter( @@ -551,9 +560,7 @@ class TestCustomTechnicalKeywords: _, baseline_score, baseline_signals = baseline_router.classify(prompt) _, custom_score, custom_signals = custom_router.classify(prompt) assert not any("technical" in s.lower() for s in baseline_signals) - assert any( - "technical" in s.lower() for s in custom_signals - ), f"Expected technical signal, got {custom_signals}" + assert any("technical" in s.lower() for s in custom_signals), f"Expected technical signal, got {custom_signals}" assert custom_score > baseline_score @@ -746,9 +753,7 @@ class TestKeywordFalsePositives: prompt = "What is the capital of France?" tier, score, signals = complexity_router.classify(prompt) # Should NOT detect code presence from 'api' in 'capital' - assert not any( - "code" in s.lower() for s in signals - ), "False positive: got code signal from 'capital'" + assert not any("code" in s.lower() for s in signals), "False positive: got code signal from 'capital'" # Should be SIMPLE (definition question) assert tier == ComplexityTier.SIMPLE @@ -757,9 +762,7 @@ class TestKeywordFalsePositives: prompt = "Explain digital marketing strategies" tier, score, signals = complexity_router.classify(prompt) # Should NOT detect code presence from 'git' in 'digital' - assert not any( - "code" in s.lower() for s in signals - ), "False positive: got code signal from 'digital'" + assert not any("code" in s.lower() for s in signals), "False positive: got code signal from 'digital'" def test_try_not_in_entry(self, complexity_router): """'try' should not match in 'entry'.""" @@ -773,43 +776,33 @@ class TestKeywordFalsePositives: """'error' should not match in 'terrorism'.""" prompt = "The country is dealing with terrorism" tier, score, signals = complexity_router.classify(prompt) - assert not any( - "code" in s.lower() for s in signals - ), "False positive: got code signal from 'terrorism'" + assert not any("code" in s.lower() for s in signals), "False positive: got code signal from 'terrorism'" def test_class_not_in_classical(self, complexity_router): """'class' should not match in 'classical'.""" prompt = "I enjoy listening to classical music" tier, score, signals = complexity_router.classify(prompt) - assert not any( - "code" in s.lower() for s in signals - ), "False positive: got code signal from 'classical'" + assert not any("code" in s.lower() for s in signals), "False positive: got code signal from 'classical'" def test_merge_not_in_emerged(self, complexity_router): """'merge' should not match in 'emerged'.""" prompt = "A new leader emerged from the crowd" tier, score, signals = complexity_router.classify(prompt) - assert not any( - "code" in s.lower() for s in signals - ), "False positive: got code signal from 'emerged'" + assert not any("code" in s.lower() for s in signals), "False positive: got code signal from 'emerged'" def test_actual_api_keyword_detected(self, complexity_router): """Actual 'api' usage should be detected.""" prompt = "How do I call the REST api endpoint?" tier, score, signals = complexity_router.classify(prompt) # Should detect code presence from actual 'api' usage - assert any( - "code" in s.lower() for s in signals - ), f"Expected code signal for 'api', got {signals}" + assert any("code" in s.lower() for s in signals), f"Expected code signal for 'api', got {signals}" def test_actual_git_keyword_detected(self, complexity_router): """Actual 'git' usage should be detected.""" prompt = "How do I use git to commit changes?" tier, score, signals = complexity_router.classify(prompt) # Should detect code presence from actual 'git' usage - assert any( - "code" in s.lower() for s in signals - ), f"Expected code signal for 'git', got {signals}" + assert any("code" in s.lower() for s in signals), f"Expected code signal for 'git', got {signals}" class TestEdgeCases: @@ -829,9 +822,7 @@ class TestEdgeCases: # Should have positive score due to length assert score > 0, f"Expected positive score for very long prompt, got {score}" # Should detect long token count - assert any( - "long" in s.lower() for s in signals - ), f"Expected 'long' signal, got {signals}" + assert any("long" in s.lower() for s in signals), f"Expected 'long' signal, got {signals}" def test_unicode_prompt(self, complexity_router): """Test handling of unicode characters.""" @@ -849,9 +840,7 @@ class TestEdgeCases: """ tier, score, signals = complexity_router.classify(prompt) # The "step N" pattern should be detected - assert any( - "multi-step" in s.lower() for s in signals - ), f"Expected multi-step signal, got {signals}" + assert any("multi-step" in s.lower() for s in signals), f"Expected multi-step signal, got {signals}" class TestRouterComplexityDeploymentMethods: @@ -918,6 +907,60 @@ class TestRouterComplexityDeploymentMethods: router.init_complexity_router_deployment(deployment) assert "auto_router/complexity_router/test-router" in router.complexity_routers + def test_hybrid_initialization_waits_for_later_pool_deployments(self): + router = Router( + model_list=[ + { + "model_name": "hybrid", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "cheap", + "complexity_router_config": { + "adaptive": True, + "tiers": { + "SIMPLE": ["cheap"], + "MEDIUM": ["cheap", "premium"], + }, + }, + }, + }, + { + "model_name": "cheap", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "input_cost_per_token": 0.00000015, + }, + "model_info": { + "adaptive_router_preferences": { + "quality_tier": 1, + "strengths": [], + } + }, + }, + { + "model_name": "premium", + "litellm_params": { + "model": "openai/gpt-4o", + "input_cost_per_token": 0.000005, + }, + "model_info": { + "adaptive_router_preferences": { + "quality_tier": 3, + "strengths": [], + } + }, + }, + ] + ) + + adaptive = router.adaptive_routers["hybrid"] + assert adaptive.model_to_cost == { + "cheap": pytest.approx(0.00000015), + "premium": pytest.approx(0.000005), + } + assert adaptive.model_to_prefs["cheap"].quality_tier == 1 + assert adaptive.model_to_prefs["premium"].quality_tier == 3 + class TestAsyncPreRoutingHookMultiFormat: """Test async_pre_routing_hook with multiple input formats.""" @@ -935,9 +978,7 @@ class TestAsyncPreRoutingHookMultiFormat: assert result.messages is not None @pytest.mark.asyncio - async def test_should_route_with_responses_api_string_input( - self, complexity_router - ): + async def test_should_route_with_responses_api_string_input(self, complexity_router): """Test routing with Responses API string input via handler dispatch.""" from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, @@ -1025,9 +1066,7 @@ class TestAsyncPreRoutingHookMultiFormat: assert result.model is not None @pytest.mark.asyncio - async def test_should_return_none_when_no_messages_or_input( - self, complexity_router - ): + async def test_should_return_none_when_no_messages_or_input(self, complexity_router): """Test that None is returned when neither messages nor input is available.""" result = await complexity_router.async_pre_routing_hook( model="test-model", @@ -1038,9 +1077,7 @@ class TestAsyncPreRoutingHookMultiFormat: assert result is None @pytest.mark.asyncio - async def test_should_prefer_original_messages_over_conversion( - self, complexity_router - ): + async def test_should_prefer_original_messages_over_conversion(self, complexity_router): """Test that original messages are used when both messages and input are available.""" messages = [{"role": "user", "content": "What is 2+2?"}] result = await complexity_router.async_pre_routing_hook( @@ -1052,9 +1089,7 @@ class TestAsyncPreRoutingHookMultiFormat: assert result.messages == messages @pytest.mark.asyncio - async def test_should_include_instructions_in_classification( - self, complexity_router - ): + async def test_should_include_instructions_in_classification(self, complexity_router): """Test that Responses API instructions influence classification via system message.""" from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, @@ -1091,9 +1126,7 @@ class TestExtractUserMessageAndSystemPrompt: {"role": "assistant", "content": "Hi!"}, {"role": "user", "content": "How are you?"}, ] - user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( - messages - ) + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt(messages) assert user_msg == "How are you?" assert sys_prompt == "You are helpful." @@ -1103,9 +1136,7 @@ class TestExtractUserMessageAndSystemPrompt: {"role": "system", "content": "You are helpful."}, {"role": "assistant", "content": "Hi!"}, ] - user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( - messages - ) + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt(messages) assert user_msg is None assert sys_prompt == "You are helpful." @@ -1123,17 +1154,13 @@ class TestExtractUserMessageAndSystemPrompt: ], } ] - user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( - messages - ) + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt(messages) assert user_msg == "Describe this image" assert sys_prompt is None def test_should_handle_empty_messages(self): """Test with empty messages list.""" - user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( - [] - ) + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt([]) assert user_msg is None assert sys_prompt is None @@ -1198,17 +1225,13 @@ class TestLLMClassifier: assert tier == ComplexityTier.SIMPLE @pytest.mark.asyncio - async def test_aclassify_llm_success_routes_by_llm_verdict( - self, llm_complexity_router, mock_router_instance - ): + async def test_aclassify_llm_success_routes_by_llm_verdict(self, llm_complexity_router, mock_router_instance): """A well-formed structured LLM response should decide the tier directly. Uses a prompt that heuristic scoring alone would classify as SIMPLE, to prove the LLM verdict -- not the heuristic scorer -- is what decided the tier. """ - mock_router_instance.acompletion = AsyncMock( - return_value=_llm_response('{"tier": "COMPLEX"}') - ) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')) tier, score, signals = await llm_complexity_router.aclassify("hi") assert tier == ComplexityTier.COMPLEX assert "llm-classifier:COMPLEX" in signals @@ -1227,13 +1250,9 @@ class TestLLMClassifier: sees no user_api_key/team_id/user_id and silently drops all spend logging and budget accounting for the classifier call. """ - mock_router_instance.acompletion = AsyncMock( - return_value=_llm_response('{"tier": "SIMPLE"}') - ) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) request_metadata = {"user_api_key": "sk-abc", "user_api_key_team_id": "team-1"} - await llm_complexity_router.aclassify( - "hi", request_kwargs={"litellm_metadata": request_metadata} - ) + await llm_complexity_router.aclassify("hi", request_kwargs={"litellm_metadata": request_metadata}) call_kwargs = mock_router_instance.acompletion.call_args.kwargs assert call_kwargs["metadata"] == request_metadata @@ -1249,18 +1268,14 @@ class TestLLMClassifier: business touching, so it must be stripped while the rest of the attribution metadata (key/team) is preserved. """ - mock_router_instance.acompletion = AsyncMock( - return_value=_llm_response('{"tier": "SIMPLE"}') - ) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) request_metadata = { "user_api_key": "sk-abc", "user_api_key_team_id": "team-1", "user_api_key_budget_reservation": {"reserved_cost": 1.0}, "user_api_key_auth": {"models": ["gpt-4o"], "budget_reservation": {"reserved_cost": 1.0}}, } - await llm_complexity_router.aclassify( - "hi", request_kwargs={"litellm_metadata": request_metadata} - ) + await llm_complexity_router.aclassify("hi", request_kwargs={"litellm_metadata": request_metadata}) call_kwargs = mock_router_instance.acompletion.call_args.kwargs # user_api_key_budget_reservation is stripped (budget enforcement) while # user_api_key_auth is kept so _filter_deployments_by_model_access_groups @@ -1307,13 +1322,9 @@ class TestLLMClassifier: assert tier == ComplexityTier.SIMPLE @pytest.mark.asyncio - async def test_pre_routing_hook_uses_llm_classifier_end_to_end( - self, llm_complexity_router, mock_router_instance - ): + async def test_pre_routing_hook_uses_llm_classifier_end_to_end(self, llm_complexity_router, mock_router_instance): """The full pre-routing hook should route using the LLM classifier's verdict.""" - mock_router_instance.acompletion = AsyncMock( - return_value=_llm_response('{"tier": "REASONING"}') - ) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}')) request_metadata = {"user_api_key": "sk-abc", "user_api_key_team_id": "team-1"} result = await llm_complexity_router.async_pre_routing_hook( model="test-model", @@ -1326,6 +1337,389 @@ class TestLLMClassifier: assert call_kwargs["metadata"] == request_metadata +class TestRouterPreRoutingAliasOverrides: + """ + Regression tests for: litellm_params configured on a complexity-router alias + entry (e.g. `cache_control_injection_points`, `drop_params`) were silently + dropped, because `async_pre_routing_hook` swaps `model` from the alias name + to the selected tier's model *before* the deployment lookup - so the actual + outbound call only ever merges in the tier deployment's own litellm_params, + never the alias's. + """ + + def _make_router(self) -> Router: + return Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "drop_params": True, + "cache_control_injection_points": [{"location": "message", "role": "system"}], + "complexity_router_config": { + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + } + }, + "complexity_router_default_model": "gpt-4o", + }, + }, + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + }, + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o"}, + }, + ] + ) + + @pytest.mark.asyncio + async def test_alias_litellm_params_applied_to_request_kwargs(self): + """cache_control_injection_points/drop_params set on the alias entry + reach the outbound request even though the tier deployment is what + actually gets called.""" + router = self._make_router() + request_kwargs: Dict = {} + + result = await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert result is not None + assert request_kwargs["drop_params"] is True + assert request_kwargs["cache_control_injection_points"] == [{"location": "message", "role": "system"}] + + @pytest.mark.asyncio + async def test_alias_overrides_exclude_only_model(self): + """`model` (the alias marker, e.g. auto_router/complexity_router) is + excluded since it's never a real provider model. Router-only fields + like complexity_router_config DO flow through into request_kwargs at + this layer - they're filtered from the actual outbound LLM call + downstream by litellm.types.utils.all_litellm_params instead, not by + the router's pre-routing hook. See test_router_init_only_params_are_ + never_sent_to_a_provider for the guard on that downstream filter.""" + router = self._make_router() + request_kwargs: Dict = {} + + await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert "model" not in request_kwargs + assert request_kwargs["complexity_router_config"] == { + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + } + } + assert request_kwargs["complexity_router_default_model"] == "gpt-4o" + + def test_router_init_only_params_are_never_sent_to_a_provider(self): + """The router's pre-routing hook only excludes `model` (see + test_alias_overrides_exclude_only_model above) - every other alias + litellm_param, including router-init-only fields like + complexity_router_config, flows into request_kwargs unfiltered. That's + only safe because litellm.completion()/acompletion() itself strips + anything listed in all_litellm_params before building the provider + request. If one of these keys is ever removed from that list, it + ships raw to the real provider as extra_body - verified live via + litellm.completion(..., complexity_router_config={...}) landing in + extra_body before this list included it.""" + from litellm.types.utils import all_litellm_params + + router_init_only_params = ( + "auto_router_config_path", + "auto_router_config", + "auto_router_default_model", + "auto_router_embedding_model", + "complexity_router_config", + "complexity_router_default_model", + "adaptive_router_config", + "adaptive_router_default_model", + "quality_router_config", + "quality_router_default_model", + ) + for param in router_init_only_params: + assert param in all_litellm_params, ( + f"{param} must stay in litellm.types.utils.all_litellm_params - " + "removing it means it ships raw to the real provider as extra_body" + ) + + @pytest.mark.asyncio + async def test_caller_supplied_kwargs_are_not_overwritten(self): + """A value the caller already passed for this request takes + precedence over the alias's configured default.""" + router = self._make_router() + request_kwargs: Dict = {"drop_params": False} + + await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert request_kwargs["drop_params"] is False + + @pytest.mark.asyncio + async def test_non_alias_model_is_untouched(self): + """A plain (non-router-alias) model name is not affected by the + alias-override merge at all.""" + router = self._make_router() + request_kwargs: Dict = {} + + result = await router.async_pre_routing_hook( + model="gpt-4o-mini", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert result is None + assert request_kwargs == {} + + @pytest.mark.asyncio + async def test_adaptive_router_alias_overrides_survive_reload(self): + """Alias litellm_params are read fresh from self.model_list at request + time (not cached at init), so a set_model_list() reload (e.g. + /config/reload) - which rebuilds self.model_list but leaves an + already-built AdaptiveRouter alone - can't leave them stale.""" + model_list = [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "drop_params": True, + "adaptive_router_config": {"available_models": ["gpt-4o-mini"]}, + }, + }, + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + }, + ] + router = Router(model_list=model_list) + router.set_model_list(model_list) + assert "smart-router" in router.adaptive_routers + + request_kwargs: Dict = {} + await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert request_kwargs["drop_params"] is True + + +class TestAdaptiveSoftFloors: + def test_adaptive_defaults_use_cost_weighted_cold_policy(self): + config = ComplexityRouterConfig( + adaptive=True, + tiers={"SIMPLE": ["cheap"]}, + ) + assert config.adaptive_weights.quality == pytest.approx(0.3) + assert config.adaptive_weights.cost == pytest.approx(0.7) + assert config.tier_distance_penalty == pytest.approx(0.5) + + @pytest.fixture + def adaptive_router_instance(self): + router = MagicMock() + router.model_list = [ + { + "model_name": "cheap", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "input_cost_per_token": 0.00000015, + }, + "model_info": {"adaptive_router_preferences": {"quality_tier": 1, "strengths": []}}, + }, + { + "model_name": "premium", + "litellm_params": { + "model": "openai/gpt-4o", + "input_cost_per_token": 0.000005, + }, + "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, + }, + ] + router.model_name_to_deployment_indices = {"cheap": [0], "premium": [1]} + return router + + @pytest.fixture + def hybrid_config(self) -> Dict: + return { + "adaptive": True, + "adaptive_weights": {"quality": 0.7, "cost": 0.3}, + "tier_distance_penalty": 0.15, + "tiers": { + "SIMPLE": ["cheap"], + "MEDIUM": ["cheap"], + "COMPLEX": ["premium"], + "REASONING": ["premium"], + }, + "default_model": "cheap", + } + + def test_adaptive_config_requires_non_empty_pools(self): + with pytest.raises(ValidationError): + ComplexityRouterConfig(adaptive=True, tiers={"SIMPLE": []}) + + def test_cold_start_randomly_samples_unobserved_classified_tier_models(self, adaptive_router_instance): + cr = ComplexityRouter( + model_name="hybrid", + litellm_router_instance=adaptive_router_instance, + complexity_router_config={ + "adaptive": True, + "tiers": { + "SIMPLE": ["cheap", "premium"], + "MEDIUM": ["premium"], + }, + }, + ) + request_kwargs: Dict = {"metadata": {}} + + with patch( + "litellm.router_strategy.complexity_router.complexity_router.random.choice", + return_value="premium", + ) as choice: + picked = cr._soft_floor_pick(ComplexityTier.SIMPLE, "hi", request_kwargs) + + assert picked == "premium" + choice.assert_called_once_with(("cheap", "premium")) + decision = request_kwargs["metadata"]["adaptive_router_decision"] + assert decision["phase"] == "cold_start" + assert {candidate["model"] for candidate in decision["candidates"]} == { + "cheap", + "premium", + } + + def test_get_model_for_tier_list_without_adaptive_random_choice(self, mock_router_instance): + router = ComplexityRouter( + model_name="test", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "adaptive": False, + "tiers": {"SIMPLE": ["cheap", "premium"], "MEDIUM": "mid"}, + "default_model": "mid", + }, + ) + pool = ["cheap", "premium"] + with patch( + "litellm.router_strategy.complexity_router.complexity_router.random.choice", + return_value="premium", + ) as choice: + assert router.get_model_for_tier(ComplexityTier.SIMPLE) == "premium" + choice.assert_called_once_with(pool) + assert router.get_model_for_tier(ComplexityTier.MEDIUM) == "mid" + + def test_soft_floor_prefers_home_tier_when_posteriors_equal(self, adaptive_router_instance, hybrid_config): + from litellm.router_strategy.adaptive_router.bandit import BanditCell + from litellm.types.router import RequestType + + cr = ComplexityRouter( + model_name="hybrid", + litellm_router_instance=adaptive_router_instance, + complexity_router_config=hybrid_config, + ) + adaptive = cr._ensure_adaptive_router() + assert adaptive is not None + for model in ("cheap", "premium"): + adaptive._cells[(RequestType.GENERAL, model)] = BanditCell(alpha=5.0, beta=5.0) + + # Equal quality samples; home-tier penalty should favor cheap for SIMPLE. + with patch( + "litellm.router_strategy.adaptive_router.bandit.thompson_sample", + return_value=0.5, + ): + picked = cr._soft_floor_pick(ComplexityTier.SIMPLE, "hi") + assert picked == "cheap" + + def test_soft_floor_allows_cross_tier_when_posterior_dominates(self, adaptive_router_instance, hybrid_config): + from litellm.router_strategy.adaptive_router.bandit import BanditCell + from litellm.types.router import RequestType + + cr = ComplexityRouter( + model_name="hybrid", + litellm_router_instance=adaptive_router_instance, + complexity_router_config=hybrid_config, + ) + adaptive = cr._ensure_adaptive_router() + assert adaptive is not None + adaptive._cells[(RequestType.GENERAL, "cheap")] = BanditCell(alpha=1.0, beta=20.0) + adaptive._cells[(RequestType.GENERAL, "premium")] = BanditCell(alpha=20.0, beta=1.0) + + with patch( + "litellm.router_strategy.adaptive_router.bandit.thompson_sample", + side_effect=lambda cell, rng=None: cell.alpha / (cell.alpha + cell.beta), + ): + picked = cr._soft_floor_pick(ComplexityTier.SIMPLE, "hi") + assert picked == "premium" + + def test_reused_model_has_zero_distance_in_each_configured_tier(self, adaptive_router_instance): + from litellm.router_strategy.adaptive_router.bandit import BanditCell + from litellm.types.router import RequestType + + cr = ComplexityRouter( + model_name="hybrid", + litellm_router_instance=adaptive_router_instance, + complexity_router_config={ + "adaptive": True, + "tiers": { + "SIMPLE": ["cheap"], + "MEDIUM": ["cheap", "premium"], + "COMPLEX": ["premium"], + }, + }, + ) + adaptive = cr._ensure_adaptive_router() + assert adaptive is not None + for model in ("cheap", "premium"): + adaptive._cells[(RequestType.GENERAL, model)] = BanditCell(alpha=6.0, beta=5.0) + request_kwargs: Dict = {"metadata": {}} + + with patch( + "litellm.router_strategy.adaptive_router.bandit.thompson_sample", + return_value=0.5, + ): + cr._soft_floor_pick(ComplexityTier.MEDIUM, "hi", request_kwargs) + + candidates = request_kwargs["metadata"]["adaptive_router_decision"]["candidates"] + assert {candidate["model"]: candidate["tier_distance"] for candidate in candidates} == { + "cheap": 0, + "premium": 0, + } + + @pytest.mark.asyncio + async def test_pre_routing_hook_adaptive_stashes_chosen_model(self, adaptive_router_instance, hybrid_config): + cr = ComplexityRouter( + model_name="hybrid", + litellm_router_instance=adaptive_router_instance, + complexity_router_config=hybrid_config, + ) + request_kwargs: Dict = {"metadata": {}} + result = await cr.async_pre_routing_hook( + model="hybrid", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + assert result is not None + assert result.model in {"cheap", "premium"} + assert request_kwargs["metadata"].get("adaptive_router_chosen_model") == result.model + decision = request_kwargs["metadata"]["adaptive_router_decision"] + assert decision["phase"] == "cold_start" + assert decision["classified_tier"] == "SIMPLE" + assert decision["request_type"] == "general" + assert decision["eligible_mode"] == "classified_tier" + assert decision["chosen_model"] == result.model + assert {candidate["model"] for candidate in decision["candidates"]} == {"cheap"} + + class TestLexicalKeywordTierRules: """Test deterministic (literal) keyword_tier_rules overrides.""" @@ -1339,9 +1733,7 @@ class TestLexicalKeywordTierRules: } @pytest.mark.asyncio - async def test_matching_rule_overrides_scoring( - self, mock_router_instance, rule_config - ): + async def test_matching_rule_overrides_scoring(self, mock_router_instance, rule_config): """A prompt hitting a rule keyword routes to that tier, not the scored tier.""" router = ComplexityRouter( model_name="test-router", @@ -1427,9 +1819,7 @@ class TestLexicalKeywordTierRules: assert router._lexical_tier_override("nothing relevant here") is None @pytest.mark.asyncio - async def test_no_rule_match_falls_back_to_scoring( - self, mock_router_instance, basic_config - ): + async def test_no_rule_match_falls_back_to_scoring(self, mock_router_instance, basic_config): """A prompt that matches no rule is classified by the scorer as usual.""" config = { **basic_config, @@ -1450,9 +1840,7 @@ class TestLexicalKeywordTierRules: assert result is not None assert result.model == "gpt-4o-mini" # SIMPLE via scoring, rule did not fire - def test_word_boundary_avoids_substring_false_positive( - self, mock_router_instance, basic_config - ): + def test_word_boundary_avoids_substring_false_positive(self, mock_router_instance, basic_config): """A single-word rule keyword must not match inside a larger word.""" config = { **basic_config, @@ -1470,10 +1858,7 @@ class TestLexicalKeywordTierRules: def _make_embedding_response(vectors: List[List[float]]) -> "litellm.EmbeddingResponse": return litellm.EmbeddingResponse( model="fake-embed", - data=[ - {"embedding": vec, "index": idx, "object": "embedding"} - for idx, vec in enumerate(vectors) - ], + data=[{"embedding": vec, "index": idx, "object": "embedding"} for idx, vec in enumerate(vectors)], object="list", ) @@ -1500,8 +1885,7 @@ class FakeEmbeddingRouter: def _vectors(self, docs: List[str]) -> List[List[float]]: return [ - [1.0, 0.0] if any(marker in doc.lower() for marker in self._CLUSTER_MARKERS) else [0.0, 1.0] - for doc in docs + [1.0, 0.0] if any(marker in doc.lower() for marker in self._CLUSTER_MARKERS) else [0.0, 1.0] for doc in docs ] @staticmethod @@ -2040,9 +2424,7 @@ class TestRoutingDecisionCauseLogging: verbose_router_logger.removeHandler(caplog.handler) @pytest.mark.asyncio - async def test_literal_keyword_match_logs_its_cause( - self, mock_router_instance, basic_config, router_log_capture - ): + async def test_literal_keyword_match_logs_its_cause(self, mock_router_instance, basic_config, router_log_capture): config = { **basic_config, "keyword_tier_rules": [{"keywords": ["deploy to k8s"], "tier": "REASONING"}], @@ -2088,9 +2470,7 @@ class TestRoutingDecisionCauseLogging: assert "cause=literal_keyword_match" not in router_log_capture.text @pytest.mark.asyncio - async def test_complexity_scorer_logs_its_cause( - self, mock_router_instance, basic_config, router_log_capture - ): + async def test_complexity_scorer_logs_its_cause(self, mock_router_instance, basic_config, router_log_capture): # No keyword rules -> the scorer decides, and its line must be tagged as such. router = ComplexityRouter( model_name="test-router", @@ -2106,3 +2486,217 @@ class TestRoutingDecisionCauseLogging: assert "score=" in router_log_capture.text assert "cause=literal_keyword_match" not in router_log_capture.text assert "cause=semantic_keyword_match" not in router_log_capture.text + + +class TestSessionAffinity: + """Test the opt-in session_affinity sticky-routing behavior.""" + + REASONING_MESSAGE = [ + { + "role": "user", + "content": "Let's think step by step and reason through this problem carefully.", + } + ] + SIMPLE_MESSAGE = [{"role": "user", "content": "Hello!"}] + + @pytest.fixture + def session_affinity_config(self, basic_config) -> Dict: + return {**basic_config, "session_affinity": True} + + @staticmethod + def _request_kwargs(session_id: str) -> Dict: + return {"metadata": {"session_id": session_id}} + + @pytest.mark.asyncio + async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): + """Regression: session_affinity defaults to False, so a shared session_id must + not pin the model -- each turn is still classified independently.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=basic_config, + ) + request_kwargs = self._request_kwargs("session-1") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.model == "o1-preview" + assert second.model == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_pins_model_after_first_turn(self, mock_router_instance, session_affinity_config): + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + request_kwargs = self._request_kwargs("session-1") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + assert first.model == "o1-preview" + + with patch.object(router, "aclassify", wraps=router.aclassify) as spy_aclassify: + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + spy_aclassify.assert_not_called() + # Pinned to the first turn's model, not re-classified down to SIMPLE. + assert second.model == "o1-preview" + + @pytest.mark.asyncio + async def test_different_sessions_classify_independently(self, mock_router_instance, session_affinity_config): + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + reasoning = await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("session-a"), messages=self.REASONING_MESSAGE + ) + simple = await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("session-b"), messages=self.SIMPLE_MESSAGE + ) + assert reasoning.model == "o1-preview" + assert simple.model == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_respects_ttl_seconds(self, mock_router_instance, basic_config): + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value=None) + mock_router_instance.cache = cache + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "session_affinity": True, + "session_affinity_ttl_seconds": 120, + }, + ) + await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("session-1"), messages=self.SIMPLE_MESSAGE + ) + cache.async_set_cache.assert_called_once() + call_kwargs = cache.async_set_cache.call_args.kwargs + assert call_kwargs["ttl"] == 120 + assert call_kwargs["value"] == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_ttl_refreshed_on_cache_hit(self, mock_router_instance, basic_config): + """Regression: a pinned turn must refresh the TTL, not just the first write -- + otherwise a session outliving session_affinity_ttl_seconds silently loses its pin.""" + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value="o1-preview") + mock_router_instance.cache = cache + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "session_affinity": True, + "session_affinity_ttl_seconds": 90, + }, + ) + result = await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("session-1"), messages=self.SIMPLE_MESSAGE + ) + assert result.model == "o1-preview" + cache.async_set_cache.assert_called_once() + call_kwargs = cache.async_set_cache.call_args.kwargs + assert call_kwargs["value"] == "o1-preview" + assert call_kwargs["ttl"] == 90 + + @pytest.mark.asyncio + async def test_different_api_keys_do_not_share_pin(self, mock_router_instance, session_affinity_config): + """A session_id is client-supplied and unauthenticated; two different callers + (API keys) reusing the same session_id must not poison each other's pin.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + caller_a_kwargs = {"metadata": {"session_id": "shared-session", "user_api_key_hash": "key-a"}} + caller_b_kwargs = {"metadata": {"session_id": "shared-session", "user_api_key_hash": "key-b"}} + + pinned_for_a = await router.async_pre_routing_hook( + model="test-model", request_kwargs=caller_a_kwargs, messages=self.REASONING_MESSAGE + ) + assert pinned_for_a.model == "o1-preview" + + # Caller B reuses the same session_id but has a different API key; its trivial + # message must classify fresh, not inherit caller A's REASONING-tier pin. + result_for_b = await router.async_pre_routing_hook( + model="test-model", request_kwargs=caller_b_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert result_for_b.model == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_no_session_id_falls_back_to_reclassify(self, mock_router_instance, session_affinity_config): + cache = AsyncMock() + mock_router_instance.cache = cache + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + result = await router.async_pre_routing_hook( + model="test-model", request_kwargs={}, messages=self.SIMPLE_MESSAGE + ) + assert result.model == "gpt-4o-mini" + cache.async_get_cache.assert_not_called() + cache.async_set_cache.assert_not_called() + + @pytest.mark.asyncio + async def test_adaptive_pinned_turn_still_stamps_chosen_model_metadata(self, mock_router_instance): + """Regression: skipping classification on a pinned turn must not break the + adaptive bandit's reward-feedback loop, which only records a turn's outcome + when ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY is present in the request metadata.""" + mock_router_instance.cache = DualCache() + mock_router_instance.model_list = [ + { + "model_name": "cheap", + "litellm_params": {"model": "openai/gpt-4o-mini", "input_cost_per_token": 0.0}, + "model_info": {}, + }, + ] + mock_router_instance.model_name_to_deployment_indices = {"cheap": [0]} + router = ComplexityRouter( + model_name="hybrid", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "adaptive": True, + "session_affinity": True, + "tiers": { + "SIMPLE": ["cheap"], + "MEDIUM": ["cheap"], + "COMPLEX": ["cheap"], + "REASONING": ["cheap"], + }, + "default_model": "cheap", + }, + ) + first = await router.async_pre_routing_hook( + model="hybrid", + request_kwargs=self._request_kwargs("session-1"), + messages=[{"role": "user", "content": "hi"}], + ) + assert first.model == "cheap" + + request_kwargs_2 = self._request_kwargs("session-1") + with patch.object(router, "aclassify", wraps=router.aclassify) as spy_aclassify: + second = await router.async_pre_routing_hook( + model="hybrid", + request_kwargs=request_kwargs_2, + messages=[{"role": "user", "content": "hi again"}], + ) + spy_aclassify.assert_not_called() + assert second.model == "cheap" + assert request_kwargs_2["metadata"]["adaptive_router_chosen_model"] == "cheap" diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/test_litellm/router_strategy/test_router_routing_plugins.py new file mode 100644 index 00000000000..e9c12d009e2 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_router_routing_plugins.py @@ -0,0 +1,220 @@ +""" +Tests for Router(plugins=[...]) -- a pipeline of routing plugins that run +before the routing decision is made, narrowing the candidate deployment pool. + +Discussion: https://github.com/BerriAI/litellm/discussions/32168 +""" + +import pytest + +from litellm import Router +from litellm.types.router import RoutingContext + + +class LanguageDetector: + async def run(self, context: RoutingContext) -> RoutingContext: + context.signals["language-detector"] = {"lang": "en"} + return context + + +class DomainClassifier: + async def run(self, context: RoutingContext) -> RoutingContext: + context.signals["domain-classifier"] = {"domain": "coding", "confidence": 0.93} + return context + + +class TenantPolicy: + ALLOWED_PROVIDERS = {"acme-corp": {"openai", "anthropic"}} + + async def run(self, context: RoutingContext) -> RoutingContext: + tenant = context.metadata.get("tenant", "default") + allowed = self.ALLOWED_PROVIDERS.get(tenant, {"openai", "anthropic", "self-hosted"}) + context.candidate_models = [m for m in context.candidate_models if m.split("/")[0] in allowed] + context.signals["tenant-policy"] = {"tenant": tenant, "allowed_providers": sorted(allowed)} + return context + + +class BudgetPolicy: + COST_CAP_PER_TOKEN = 0.000005 + COST_BY_MODEL = { + "openai/gpt-4o-mini": 0.00000015, + "anthropic/claude-haiku-4-5": 0.000001, + "openai/gpt-5.1": 0.00003, + } + + async def run(self, context: RoutingContext) -> RoutingContext: + context.candidate_models = [ + m for m in context.candidate_models if self.COST_BY_MODEL.get(m, 0) <= self.COST_CAP_PER_TOKEN + ] + context.signals["budget-policy"] = {"daily_limit": 100} + return context + + +class BlockEverything: + async def run(self, context: RoutingContext) -> RoutingContext: + context.candidate_models = [] + return context + + +def _smart_router_model_list(): + return [ + { + "model_name": "smart-router", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "cheap openai"}, + "model_info": {"tags": ["openai"]}, + }, + { + "model_name": "smart-router", + "litellm_params": {"model": "anthropic/claude-haiku-4-5", "mock_response": "anthropic"}, + "model_info": {"tags": ["anthropic"]}, + }, + { + "model_name": "smart-router", + "litellm_params": {"model": "openai/gpt-5.1", "mock_response": "expensive openai"}, + "model_info": {"tags": ["openai"]}, + }, + { + "model_name": "smart-router", + "litellm_params": {"model": "ollama/llama-3-70b", "mock_response": "self hosted"}, + "model_info": {"tags": ["self-hosted"]}, + }, + ] + + +@pytest.mark.asyncio +async def test_routing_plugin_pipeline_matches_jeann2013_e2e_scenario(): + """ + https://github.com/BerriAI/litellm/discussions/32168#discussioncomment-17608820 + + language plugin -> domain classifier -> tenant policy (openai+anthropic only) + -> budget policy (drops over-cap models) -> Router picks the best remaining + candidate. Must never land on the self-hosted or over-budget deployment. + """ + router = Router( + model_list=_smart_router_model_list(), + plugins=[LanguageDetector(), DomainClassifier(), TenantPolicy(), BudgetPolicy()], + ) + + response = await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "Write a function to reverse a linked list."}], + metadata={"tenant": "acme-corp"}, + ) + + # response.model is the bare model name (litellm strips the provider/ prefix + # on the response), so compare against bare names rather than litellm_params.model + routed_model = response.model + + assert routed_model in {"gpt-4o-mini", "claude-haiku-4-5"} + assert routed_model not in {"llama-3-70b", "gpt-5.1"} + + +@pytest.mark.asyncio +async def test_routing_plugin_narrowing_to_zero_candidates_raises(): + """A plugin narrowing to nothing is a policy decision -- must raise, not silently + fall back to the unfiltered pool (that would defeat the policy it enforces).""" + router = Router( + model_list=_smart_router_model_list(), + plugins=[BlockEverything()], + ) + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "hi"}], + ) + + +def test_sync_get_available_deployment_rejects_configured_plugins(): + """ + Router.completion() (and any other sync entry point) resolves deployments via + the synchronous get_available_deployment(), which never runs the routing-plugin + pipeline. Silently allowing that would let a deny-all policy plugin be bypassed + just by calling the sync API -- must fail closed instead. + """ + router = Router(model_list=_smart_router_model_list(), plugins=[TenantPolicy()]) + + with pytest.raises(ValueError, match="routing-plugin pipeline"): + router.get_available_deployment(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + +def test_sync_router_completion_rejects_configured_plugins(): + """End-to-end: Router.completion() (the sync API) must not silently skip plugins either.""" + router = Router(model_list=_smart_router_model_list(), plugins=[TenantPolicy()]) + + with pytest.raises(ValueError, match="routing-plugin pipeline"): + router.completion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + +@pytest.mark.asyncio +async def test_async_completion_with_unsupported_strategy_rejects_configured_plugins(): + """ + async_get_available_deployment() itself delegates to the synchronous selector + for routing strategies outside {simple-shuffle, usage-based-routing-v2, + cost-based-routing, latency-based-routing, least-busy} -- e.g. "usage-based-routing" + (v1, not v2) -- which would silently bypass the plugin pipeline on the async path too. + """ + router = Router( + model_list=_smart_router_model_list(), + plugins=[TenantPolicy()], + routing_strategy="usage-based-routing", + ) + + with pytest.raises(ValueError, match="routing-plugin pipeline"): + await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + +@pytest.mark.asyncio +async def test_router_without_plugins_is_unaffected(): + """Regression guard: a Router with no `plugins` configured behaves exactly as before.""" + router = Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "hi"}, + }, + ], + ) + response = await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "hi"}], + ) + assert response.choices[0].message.content == "hi" + + +@pytest.mark.asyncio +async def test_run_routing_plugins_narrows_candidates_and_records_signals(): + """Unit-level check of _run_routing_plugins in isolation, independent of acompletion.""" + router = Router( + model_list=_smart_router_model_list(), + plugins=[LanguageDetector(), DomainClassifier(), TenantPolicy(), BudgetPolicy()], + ) + request_kwargs = {"metadata": {"tenant": "acme-corp"}} + + context = await router._run_routing_plugins( + model="smart-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert context.candidate_models == ["openai/gpt-4o-mini", "anthropic/claude-haiku-4-5"] + assert context.signals["domain-classifier"]["domain"] == "coding" + assert request_kwargs["metadata"]["_routing_plugin_candidate_models"] == context.candidate_models + + +def test_filter_by_routing_plugin_candidates_narrows_and_raises_when_empty(): + """Unit-level check of _filter_by_routing_plugin_candidates in isolation.""" + router = Router(model_list=_smart_router_model_list(), plugins=[TenantPolicy()]) + healthy_deployments = router.model_list + + narrowed = router._filter_by_routing_plugin_candidates( + healthy_deployments=healthy_deployments, + request_kwargs={"metadata": {"_routing_plugin_candidate_models": ["openai/gpt-4o-mini"]}}, + ) + assert [d["litellm_params"]["model"] for d in narrowed] == ["openai/gpt-4o-mini"] + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + router._filter_by_routing_plugin_candidates( + healthy_deployments=healthy_deployments, + request_kwargs={"metadata": {"_routing_plugin_candidate_models": ["nonexistent/model"]}}, + ) diff --git a/tests/test_litellm/test_claude_opus_4_8_config.py b/tests/test_litellm/test_claude_opus_4_8_config.py index 32f7d249e05..8eead8a9c84 100644 --- a/tests/test_litellm/test_claude_opus_4_8_config.py +++ b/tests/test_litellm/test_claude_opus_4_8_config.py @@ -95,6 +95,8 @@ def test_opus_4_8_model_pricing_and_capabilities(): assert info["supports_tool_choice"] is True assert info["supports_vision"] is True + assert model_data["claude-opus-4-8"]["supports_native_structured_output"] is True + def test_opus_4_8_bedrock_regional_model_pricing(): model_data = _load_root_cost_map() @@ -165,6 +167,7 @@ def test_opus_4_8_present_in_bundled_backup(): "azure_ai/claude-opus-4-8", ): assert model_name in backup, f"Missing from backup cost map: {model_name}" + assert backup["claude-opus-4-8"]["supports_native_structured_output"] is True def test_opus_4_8_registered_for_bedrock_converse(): diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 28cf4fa0744..4611aafa3c1 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2081,3 +2081,79 @@ 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 +@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, aws_credential_kwargs: dict +): + 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", + 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_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 + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 26ae6ff7f4c..6d515ecdc73 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -615,6 +615,8 @@ def validate_model_cost_values(model_data, exceptions=None): "input_cost_per_audio_token", "output_cost_per_audio_token", "output_cost_per_image_token", + "input_cost_per_video_token", + "output_cost_per_video_token", "input_cost_per_audio_per_second", "input_cost_per_video_per_second", "input_cost_per_token_above_128k_tokens", @@ -732,6 +734,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_image": {"type": "number"}, "input_cost_per_image_above_128k_tokens": {"type": "number"}, "input_cost_per_image_token": {"type": "number"}, + "input_cost_per_video_token": {"type": "number"}, "input_cost_per_token_above_200k_tokens": {"type": "number"}, "input_cost_per_token_above_256k_tokens": {"type": "number"}, "input_cost_per_token_above_272k_tokens": {"type": "number"}, @@ -807,6 +810,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_character_above_128k_tokens": {"type": "number"}, "output_cost_per_image": {"type": "number"}, "output_cost_per_image_token": {"type": "number"}, + "output_cost_per_video_token": {"type": "number"}, "output_cost_per_pixel": {"type": "number"}, "output_cost_per_second": {"type": "number"}, "output_cost_per_second_1080p": {"type": "number"}, @@ -4710,3 +4714,30 @@ class TestValidateEnvironmentTencent: assert "TENCENT_API_KEY" in result["missing_keys"] + +@pytest.mark.parametrize( + "model", + [ + "vertex_ai/gemini-2.5-flash-image", + "vertex_ai/gemini-3-pro-image", + "vertex_ai/gemini-3-pro-image-preview", + "vertex_ai/gemini-3.1-flash-image", + "vertex_ai/gemini-3.1-flash-image-preview", + "gemini/gemini-2.5-flash-image", + "gemini/gemini-3-pro-image", + "gemini/gemini-3-pro-image-preview", + "gemini/gemini-3.1-flash-image", + "gemini/gemini-3.1-flash-image-preview", + ], +) +def test_gemini_image_models_do_not_support_reasoning( + model: str, local_model_cost_map: None +) -> None: + assert model in litellm.model_cost, ( + f"{model} is missing from the local model cost map. " + "Add its entry to litellm/model_prices_and_context_window_backup.json." + ) + assert litellm.supports_reasoning(model) is False, ( + f"{model} incorrectly classified as reasoning-capable. " + "Add 'supports_reasoning: false' to its model_cost entry." + ) diff --git a/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts b/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts index 6b5e7f7bb8d..a58ece16f9c 100644 --- a/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts +++ b/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts @@ -41,7 +41,7 @@ export async function dismissFeedbackPopup(page: PlaywrightPage): Promise /** * Click on a team ID in the table. Team IDs are rendered differently depending * on the component version — try button first (Tremor Button), fall back to - * clickable span (OldTeams Typography.Text). + * clickable span (Teams Typography.Text). */ export async function clickTeamId(page: PlaywrightPage, teamId: string): Promise { const cell = page.locator("td").filter({ hasText: teamId }).first(); diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts index b30bb8aca7b..e7f67d7367f 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts @@ -96,7 +96,9 @@ test.describe("Proxy Admin - Teams", () => { const teamRow = page.locator("tr", { hasText: E2E_TEAM_DELETE_ALIAS }).first(); await expect(teamRow).toBeVisible({ timeout: 10_000 }); - await teamRow.locator("svg, img").last().click(); + // Actions live in a kebab menu: open it, then click "Delete team". + await teamRow.locator('[data-testid^="team-actions-"]').click(); + await page.getByTestId("team-action-delete").click(); const modal = page.locator(".ant-modal:visible"); await expect(modal).toBeVisible({ timeout: 5_000 }); diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 48b400536bd..67139d2c932 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -301,17 +301,6 @@ "count": 1 } }, - "src/app/(dashboard)/guardrails/_components/edit_guardrail_form.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "no-restricted-syntax": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/app/(dashboard)/guardrails/_components/guardrail_info.tsx": { "max-params": { "count": 1 @@ -339,14 +328,6 @@ "count": 1 } }, - "src/app/(dashboard)/guardrails/_components/guardrail_table.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/guardrails/_components/tool_permission/ToolPermissionRulesEditor.tsx": { "no-restricted-imports": { "count": 1 @@ -515,6 +496,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 @@ -858,9 +985,6 @@ "src/app/(dashboard)/projects/_components/ProjectDetailsPage.tsx": { "no-nested-ternary": { "count": 3 - }, - "no-restricted-imports": { - "count": 1 } }, "src/app/(dashboard)/projects/_components/ProjectKeysSection.tsx": { @@ -1079,6 +1203,51 @@ "count": 1 } }, + "src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx": { + "no-nested-ternary": { + "count": 1 + } + }, + "src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "react-hooks/immutability": { + "count": 1 + } + }, + "src/app/(dashboard)/usage/_components/components/UsagePageView.tsx": { + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/purity": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 3 + } + }, + "src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts": { + "react-hooks/refs": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 1 + } + }, "src/app/(dashboard)/users/_components/DefaultUserSettings.tsx": { "no-restricted-imports": { "count": 1 @@ -1358,22 +1527,6 @@ "count": 1 } }, - "src/components/OldTeams.test.tsx": { - "max-nested-callbacks": { - "count": 4 - } - }, - "src/components/OldTeams.tsx": { - "no-nested-ternary": { - "count": 4 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 4 - } - }, "src/components/SCIM.tsx": { "no-restricted-imports": { "count": 1 @@ -1464,6 +1617,22 @@ "count": 1 } }, + "src/components/Teams.test.tsx": { + "max-nested-callbacks": { + "count": 4 + } + }, + "src/components/Teams.tsx": { + "no-nested-ternary": { + "count": 2 + }, + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 3 + } + }, "src/components/ToolDetail.tsx": { "unused-imports/no-unused-imports": { "count": 2 @@ -1502,69 +1671,6 @@ "count": 1 } }, - "src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageBarChart.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageLineChart.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx": { - "no-nested-ternary": { - "count": 1 - } - }, - "src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "react-hooks/immutability": { - "count": 1 - } - }, - "src/app/(dashboard)/usage/_components/components/UsagePageView.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/purity": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 3 - } - }, - "src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts": { - "react-hooks/refs": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "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 @@ -1809,11 +1915,6 @@ "count": 1 } }, - "src/components/common_components/user_search_modal.tsx": { - "react-hooks/use-memo": { - "count": 1 - } - }, "src/components/constants.tsx": { "react-hooks/set-state-in-effect": { "count": 1 @@ -1873,45 +1974,11 @@ "count": 1 } }, - "src/components/mcp_tools/ByokCredentialModal.tsx": { - "no-restricted-syntax": { - "count": 1 - } - }, - "src/components/mcp_tools/MCPLogoSelector.test.tsx": { - "unused-imports/no-unused-imports": { - "count": 1 - } - }, - "src/components/mcp_tools/MCPNetworkSettings.tsx": { - "react-hooks/immutability": { - "count": 2 - } - }, - "src/components/mcp_tools/MCPSubmissionsTab.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/mcp_tools/MCPToolArgumentsForm.tsx": { "no-nested-ternary": { "count": 5 } }, - "src/components/mcp_tools/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 @@ -1920,123 +1987,6 @@ "count": 1 } }, - "src/components/mcp_tools/OAuthFormFields.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/mcp_tools/OpenAPIQuickPicker.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, - "src/components/mcp_tools/ToolTestPanel.tsx": { - "no-nested-ternary": { - "count": 3 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, - "src/components/mcp_tools/UserEnvVarsModal.tsx": { - "no-nested-ternary": { - "count": 2 - } - }, - "src/components/mcp_tools/create_mcp_server.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 4 - } - }, - "src/components/mcp_tools/mcp_connect.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/static-components": { - "count": 4 - } - }, - "src/components/mcp_tools/mcp_connection_status.tsx": { - "no-nested-ternary": { - "count": 3 - }, - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/mcp_tools/mcp_discovery.tsx": { - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, - "src/components/mcp_tools/mcp_server_cost_config.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/mcp_tools/mcp_server_cost_display.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/mcp_tools/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/components/mcp_tools/mcp_server_view.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/mcp_tools/mcp_servers.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, - "src/components/mcp_tools/mcp_tool_configuration.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/mcp_tools/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 @@ -2130,9 +2080,6 @@ "src/components/molecules/filter.tsx": { "no-nested-ternary": { "count": 2 - }, - "react-hooks/use-memo": { - "count": 1 } }, "src/components/molecules/models/columns.test.tsx": { @@ -2194,9 +2141,6 @@ }, "react-hooks/set-state-in-effect": { "count": 4 - }, - "react-hooks/use-memo": { - "count": 1 } }, "src/components/organization/organization_view.tsx": { diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 648503abb27..91cf705060e 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -14,7 +14,7 @@ "@base-ui/react": "^1.6.0", "@headlessui/tailwindcss": "0.2.2", "@heroicons/react": "1.0.6", - "@tanstack/react-pacer": "0.2.0", + "@tanstack/react-pacer": "0.22.1", "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", "@tremor/react": "3.18.7", @@ -29,6 +29,7 @@ "next": "16.2.6", "openai": "4.104.0", "openapi-fetch": "^0.17.0", + "openapi-react-query": "^0.5.4", "papaparse": "5.5.3", "react": "18.3.1", "react-copy-to-clipboard": "5.1.1", @@ -50,7 +51,6 @@ "@testing-library/jest-dom": "6.9.1", "@testing-library/react": "16.3.2", "@testing-library/user-event": "14.6.1", - "@types/lodash": "4.17.23", "@types/node": "20.19.37", "@types/react": "18.2.48", "@types/react-copy-to-clipboard": "5.0.7", @@ -3618,11 +3618,31 @@ "tailwindcss": "4.3.2" } }, - "node_modules/@tanstack/pacer": { - "version": "0.2.0", - "resolved": "https://registry.npmjs.org/@tanstack/pacer/-/pacer-0.2.0.tgz", - "integrity": "sha512-fUJs3NpSwtAL/tfq8kuYdgvm9HbbJvHsOG6aHY2dFDfff0NBFNwjvyGreWZZRPs2zgoIbr4nOk+rRV7aQgmf+A==", + "node_modules/@tanstack/devtools-event-client": { + "version": "0.4.4", + "resolved": "https://registry.npmjs.org/@tanstack/devtools-event-client/-/devtools-event-client-0.4.4.tgz", + "integrity": "sha512-6T5Yop/793YI+H+5J8Hsyj4kCih9sl4t3ElLgKioW5hk3ocn+ZdSJ94tT7vL7uabxSugWYBZlOTMPzEw2puvQw==", "license": "MIT", + "bin": { + "intent": "bin/intent.js" + }, + "engines": { + "node": ">=18" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + } + }, + "node_modules/@tanstack/pacer": { + "version": "0.21.1", + "resolved": "https://registry.npmjs.org/@tanstack/pacer/-/pacer-0.21.1.tgz", + "integrity": "sha512-hB01dd4rlsYcTCNP7wK186jgAe6K5qimgM1Y5Jtvz+9PUaILvpmeLLjmQNUNSO1l23lIt+CeQR6mO1mjlPvRtQ==", + "license": "MIT", + "dependencies": { + "@tanstack/devtools-event-client": "^0.4.3", + "@tanstack/store": "^0.11.0" + }, "engines": { "node": ">=18" }, @@ -3642,12 +3662,13 @@ } }, "node_modules/@tanstack/react-pacer": { - "version": "0.2.0", - "resolved": "https://registry.npmjs.org/@tanstack/react-pacer/-/react-pacer-0.2.0.tgz", - "integrity": "sha512-KU5GtjkKSeNdYCilen5Dc+Pu/6BPQbsQshKrUUjrg7URyJIiGBCz6ZZFre1QjDz/aeUeqUJWMWSm+2Dsh64v+w==", + "version": "0.22.1", + "resolved": "https://registry.npmjs.org/@tanstack/react-pacer/-/react-pacer-0.22.1.tgz", + "integrity": "sha512-CenQqK0GluSPIrnsG1yuD7w5uMSQ/4lI9AcGEFxBrRd66r260boWcYRIsS5+eHtXb238FoZYhKmJPGlhRzmHRw==", "license": "MIT", "dependencies": { - "@tanstack/pacer": "0.2.0" + "@tanstack/pacer": "0.21.1", + "@tanstack/react-store": "^0.11.0" }, "engines": { "node": ">=18" @@ -3677,6 +3698,24 @@ "react": "^18 || ^19" } }, + "node_modules/@tanstack/react-store": { + "version": "0.11.0", + "resolved": "https://registry.npmjs.org/@tanstack/react-store/-/react-store-0.11.0.tgz", + "integrity": "sha512-tX4YXh3PDkmpvGQWkWqKpzs/MSqbtuwY9dWdWhtV9Q50PmO+jOkUKIWIX4G85dwt7lxdHLXsiaEKPdKmC8F41w==", + "license": "MIT", + "dependencies": { + "@tanstack/store": "0.11.0", + "use-sync-external-store": "^1.6.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + }, + "peerDependencies": { + "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0", + "react-dom": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" + } + }, "node_modules/@tanstack/react-table": { "version": "8.21.3", "resolved": "https://registry.npmjs.org/@tanstack/react-table/-/react-table-8.21.3.tgz", @@ -3714,6 +3753,16 @@ "react-dom": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, + "node_modules/@tanstack/store": { + "version": "0.11.0", + "resolved": "https://registry.npmjs.org/@tanstack/store/-/store-0.11.0.tgz", + "integrity": "sha512-WlzzCt3xi0G6pCAJu1U+2jiECwabETDpQDi3hfkFZvJii9AuZqEKbOiVarX1/bWhTNjU486yQtJCCasi/0q+Cw==", + "license": "MIT", + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + } + }, "node_modules/@tanstack/table-core": { "version": "8.21.3", "resolved": "https://registry.npmjs.org/@tanstack/table-core/-/table-core-8.21.3.tgz", @@ -4059,13 +4108,6 @@ "dev": true, "license": "MIT" }, - "node_modules/@types/lodash": { - "version": "4.17.23", - "resolved": "https://registry.npmjs.org/@types/lodash/-/lodash-4.17.23.tgz", - "integrity": "sha512-RDvF6wTulMPjrNdCoYRC8gNR880JNGT8uB+REUpC2Ns4pRqQJhGz90wh7rgdXDPpCczF3VGktDuFGVnz8zP7HA==", - "dev": true, - "license": "MIT" - }, "node_modules/@types/mdast": { "version": "4.0.4", "resolved": "https://registry.npmjs.org/@types/mdast/-/mdast-4.0.4.tgz", @@ -10544,6 +10586,19 @@ "openapi-typescript-helpers": "^0.1.0" } }, + "node_modules/openapi-react-query": { + "version": "0.5.4", + "resolved": "https://registry.npmjs.org/openapi-react-query/-/openapi-react-query-0.5.4.tgz", + "integrity": "sha512-V9lRiozjHot19/BYSgXYoyznDxDJQhEBSdi26+SJ0UqjMANLQhkni4XG+Z7e3Ag7X46ZLMrL9VxYkghU3QvbWg==", + "license": "MIT", + "dependencies": { + "openapi-typescript-helpers": "^0.1.0" + }, + "peerDependencies": { + "@tanstack/react-query": "^5.80.0", + "openapi-fetch": "^0.17.0" + } + }, "node_modules/openapi-typescript": { "version": "7.13.0", "resolved": "https://registry.npmjs.org/openapi-typescript/-/openapi-typescript-7.13.0.tgz", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 40c495cebf1..7aea571ea8b 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -19,6 +19,7 @@ "e2e:migration": "playwright test e2e_tests/tests/migration/migratedPages.spec.ts --config e2e_tests/playwright.config.ts", "e2e:migration:root": "playwright test --config e2e_tests/migration.serverRootPath.config.ts", "knip": "knip", + "knip:ci": "knip --exclude exports,nsExports,types,nsTypes,enumMembers,classMembers,duplicates", "knip:fix": "knip --fix", "gen:api": "node scripts/gen-api-types.mjs" }, @@ -29,7 +30,7 @@ "@base-ui/react": "^1.6.0", "@headlessui/tailwindcss": "0.2.2", "@heroicons/react": "1.0.6", - "@tanstack/react-pacer": "0.2.0", + "@tanstack/react-pacer": "0.22.1", "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", "@tremor/react": "3.18.7", @@ -44,6 +45,7 @@ "next": "16.2.6", "openai": "4.104.0", "openapi-fetch": "^0.17.0", + "openapi-react-query": "^0.5.4", "papaparse": "5.5.3", "react": "18.3.1", "react-copy-to-clipboard": "5.1.1", @@ -65,7 +67,6 @@ "@testing-library/jest-dom": "6.9.1", "@testing-library/react": "16.3.2", "@testing-library/user-event": "14.6.1", - "@types/lodash": "4.17.23", "@types/node": "20.19.37", "@types/react": "18.2.48", "@types/react-copy-to-clipboard": "5.0.7", @@ -93,7 +94,6 @@ "js-yaml": "4.2.0", "glob": "13.0.0", "minimatch": "10.2.4", - "lodash": "4.18.1", "ws": "8.21.0", "braces": "3.0.3", "axios": "1.13.6", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx index dc70aaa409b..4ee6332c54e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx @@ -1,6 +1,6 @@ import React from "react"; import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { screen, waitFor } from "@testing-library/react"; +import { act, fireEvent, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "@/../tests/test-utils"; import AgentCardDiscovery from "./agent_card_discovery"; @@ -243,6 +243,54 @@ describe("AgentCardDiscovery", () => { expect(selection.selected_card.capabilities.streaming).toBe(false); }); + it("does not fire discovery before the debounce wait and fires once with the last URL", async () => { + mockDiscover.mockResolvedValue({ + url: "https://last.example.com", + agent_card: sampleCard, + }); + renderWithProviders(); + const input = screen.getByPlaceholderText("https://upstream-agent.example.com"); + + act(() => { + fireEvent.change(input, { target: { value: "https://first.example.com" } }); + }); + act(() => { + vi.advanceTimersByTime(399); + }); + expect(mockDiscover).not.toHaveBeenCalled(); + + act(() => { + fireEvent.change(input, { target: { value: "https://last.example.com" } }); + }); + act(() => { + vi.advanceTimersByTime(399); + }); + expect(mockDiscover).not.toHaveBeenCalled(); + + await act(async () => { + await vi.advanceTimersByTimeAsync(1); + }); + expect(mockDiscover).toHaveBeenCalledTimes(1); + expect(mockDiscover).toHaveBeenCalledWith("tok", "https://last.example.com", undefined); + }); + + it("fires no discovery when unmounted mid-wait", () => { + const { unmount } = renderWithProviders(); + const input = screen.getByPlaceholderText("https://upstream-agent.example.com"); + + act(() => { + fireEvent.change(input, { target: { value: "https://first.example.com" } }); + }); + act(() => { + vi.advanceTimersByTime(200); + }); + unmount(); + act(() => { + vi.advanceTimersByTime(2000); + }); + expect(mockDiscover).not.toHaveBeenCalled(); + }); + it("blocks discover when no access token is provided", async () => { const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime }); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx index 5ea7458f643..e979b2dbe3f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.tsx @@ -11,6 +11,7 @@ import { SearchOutlined, } from "@ant-design/icons"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { DiscoveredAgentCard, discoverAgentCardCall } from "@/components/networking"; import { ALLOWED_CAPABILITY_KEYS, @@ -22,6 +23,8 @@ import { const { Text, Paragraph } = Typography; const { Panel } = Collapse; +const DISCOVERY_DEBOUNCE_WAIT_MS = 400; + export interface DiscoveredAgentCardSelection { /** Full upstream card the proxy fetched, unmodified. */ raw_card: DiscoveredAgentCard; @@ -171,6 +174,14 @@ const AgentCardDiscovery: React.FC = ({ // eslint-disable-next-line react-hooks/exhaustive-deps }, [accessToken, effectiveUrl, isParentDriven, discoveryMode, discoveryParamsKey]); + const debouncedDiscover = useDebouncedCallback( + () => { + if (!accessToken || !effectiveUrl.trim()) return; + void handleDiscover(); + }, + { wait: DISCOVERY_DEBOUNCE_WAIT_MS }, + ); + // Auto-discover when the URL (or parent plan) becomes available. Debounce // is applied uniformly so rapid changes from a watched parent form (e.g. // typing into a LangGraph api_base / assistant_id field) don't fire one @@ -186,11 +197,8 @@ const AgentCardDiscovery: React.FC = ({ return; } - const timer = window.setTimeout(() => { - void handleDiscover(); - }, 400); - return () => window.clearTimeout(timer); - }, [accessToken, effectiveUrl, handleDiscover]); + debouncedDiscover(); + }, [accessToken, effectiveUrl, handleDiscover, debouncedDiscover]); const toggleSkill = (id: string, checked: boolean) => { setSelectedSkillIds((prev) => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx index adaf7e4c63f..2ad332ae9ed 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx @@ -105,6 +105,6 @@ describe("GuardrailsPanel", () => { expect(screen.getByText("Guardrails")).toBeInTheDocument(); // Activate the Guardrails tab so its content (including the Add button) is rendered fireEvent.click(screen.getByText("Guardrails")); - expect(screen.getByText("+ Add New Guardrail")).toBeInTheDocument(); + expect(screen.getByText("Add New Guardrail")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index af4c85d0b9f..33c6933634c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -1,7 +1,15 @@ import React, { useState, useEffect } from "react"; -import { Button, Dropdown, Tabs } from "antd"; -import { DownOutlined, PlusOutlined, CodeOutlined } from "@ant-design/icons"; +import { Tabs } from "antd"; +import { ChevronDown, Code, Plus } from "lucide-react"; import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; import AddGuardrailForm from "./add_guardrail_form"; import GuardrailTable from "./guardrail_table"; import { isAdminRole } from "@/utils/roles"; @@ -133,30 +141,26 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole children: ( <>
- , - label: "Add Provider Guardrail", - onClick: handleAddGuardrail, - }, - { - key: "custom_code", - icon: , - label: "Create Custom Code Guardrail", - onClick: handleAddCustomCodeGuardrail, - }, - ], - }} - trigger={["click"]} - disabled={!accessToken} - > - - + + + + Add New Guardrail + + + + + + Add Provider Guardrail + + + + Create Custom Code Guardrail + + +
{selectedGuardrailId ? ( @@ -171,9 +175,6 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole guardrailsList={guardrailsList} isLoading={isLoading} onDeleteClick={handleDeleteClick} - accessToken={accessToken} - onGuardrailUpdated={fetchGuardrails} - isAdmin={isAdmin} onGuardrailClick={(id) => setSelectedGuardrailId(id)} /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx index 07cdcbc888c..4217a765732 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx @@ -1,6 +1,8 @@ "use client"; import React, { useState, useEffect, useCallback } from "react"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { SearchIcon, PlusIcon, @@ -728,6 +730,7 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { rejected: 0, }); const [search, setSearch] = useState(""); + const [searchDebounced] = useDebouncedValue(search, { wait: DEBOUNCE_WAIT_MS }); const [statusFilter, setStatusFilter] = useState<"all" | GuardrailStatus>("all"); const [selectedId, setSelectedId] = useState(null); const [expandedHeaders, setExpandedHeaders] = useState>(new Set()); @@ -737,16 +740,10 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { } | null>(null); const [isLoading, setIsLoading] = useState(true); const [error, setError] = useState(null); - const [searchDebounced, setSearchDebounced] = useState(""); const [isSubmitModalOpen, setIsSubmitModalOpen] = useState(false); const [submitForm] = Form.useForm(); const registerGuardrail = useRegisterGuardrail(); - useEffect(() => { - const t = setTimeout(() => setSearchDebounced(search), 300); - return () => clearTimeout(t); - }, [search]); - const fetchSubmissions = useCallback(async () => { if (!accessToken) { setIsLoading(false); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/edit_guardrail_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/edit_guardrail_form.tsx deleted file mode 100644 index a916984f8d7..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/edit_guardrail_form.tsx +++ /dev/null @@ -1,491 +0,0 @@ -import React, { useState, useEffect } from "react"; -import { Form, Typography, Select, Input, Switch, Modal } from "antd"; -import { Button, TextInput } from "@tremor/react"; -import { - guardrail_provider_map, - guardrailLogoMap, - getGuardrailProviders, - getSupportedModesForProvider, - toModeArray, - type SkipSystemMessageChoice, - type SkipToolMessageChoice, -} from "./guardrail_info_helpers"; -import { resolveLogoSrc } from "@/lib/assetPaths"; -import { getGuardrailUISettings, getGlobalLitellmHeaderName } from "@/components/networking"; -import PiiConfiguration from "./pii_configuration"; -import NotificationsManager from "@/components/molecules/notifications_manager"; - -const { Title, Text } = Typography; -const { Option } = Select; - -interface EditGuardrailFormProps { - visible: boolean; - onClose: () => void; - accessToken: string | null; - onSuccess: () => void; - guardrailId: string; - /** Full stored params merged into PUT so optional fields (e.g. content filter) are preserved. */ - fullLitellmParams?: Record | null; - initialValues: { - guardrail_name: string; - provider: string; - mode: string; - default_on: boolean; - pii_entities_config?: { [key: string]: string }; - skip_system_message_choice?: SkipSystemMessageChoice; - skip_tool_message_choice?: SkipToolMessageChoice; - [key: string]: unknown; - }; -} - -interface GuardrailSettings { - supported_entities: string[]; - supported_actions: string[]; - supported_modes: string[]; - supported_modes_by_provider?: Record; - pii_entity_categories: Array<{ - category: string; - entities: string[]; - }>; -} - -const EditGuardrailForm: React.FC = ({ - visible, - onClose, - accessToken, - onSuccess, - guardrailId, - fullLitellmParams, - initialValues, -}) => { - const [form] = Form.useForm(); - const [loading, setLoading] = useState(false); - const [selectedProvider, setSelectedProvider] = useState(initialValues?.provider || null); - const [guardrailSettings, setGuardrailSettings] = useState(null); - const [selectedEntities, setSelectedEntities] = useState([]); - const [selectedActions, setSelectedActions] = useState<{ [key: string]: string }>({}); - - // Fetch guardrail settings when the component mounts - useEffect(() => { - const fetchGuardrailSettings = async () => { - try { - if (!accessToken) return; - - const data = await getGuardrailUISettings(accessToken); - setGuardrailSettings(data); - } catch (error) { - console.error("Error fetching guardrail settings:", error); - NotificationsManager.fromBackend("Failed to load guardrail settings"); - } - }; - - fetchGuardrailSettings(); - }, [accessToken]); - - // Initialize selected entities and actions from initialValues - useEffect(() => { - if (initialValues?.pii_entities_config && Object.keys(initialValues.pii_entities_config).length > 0) { - const entities = Object.keys(initialValues.pii_entities_config); - setSelectedEntities(entities); - setSelectedActions(initialValues.pii_entities_config); - } - }, [initialValues]); - - const handleProviderChange = (value: string) => { - setSelectedProvider(value); - // Reset form fields that are provider-specific - form.setFieldsValue({ - config: undefined, - }); - - // Reset PII selections when changing provider - setSelectedEntities([]); - setSelectedActions({}); - }; - - const handleEntitySelect = (entity: string) => { - setSelectedEntities((prev) => { - if (prev.includes(entity)) { - return prev.filter((e) => e !== entity); - } else { - return [...prev, entity]; - } - }); - }; - - const handleActionSelect = (entity: string, action: string) => { - setSelectedActions((prev) => ({ - ...prev, - [entity]: action, - })); - }; - - const handleSubmit = async () => { - try { - setLoading(true); - const values = await form.validateFields(); - - // Get the guardrail provider value from the map - const guardrailProvider = guardrail_provider_map[values.provider]; - - const litellm_params: Record = - fullLitellmParams && typeof fullLitellmParams === "object" ? { ...fullLitellmParams } : {}; - - litellm_params.guardrail = guardrailProvider; - litellm_params.mode = values.mode; - litellm_params.default_on = values.default_on; - - const skipChoice = values.skip_system_message_choice as SkipSystemMessageChoice | undefined; - if (skipChoice === "yes") { - litellm_params.skip_system_message_in_guardrail = true; - } else if (skipChoice === "no") { - litellm_params.skip_system_message_in_guardrail = false; - } else { - delete litellm_params.skip_system_message_in_guardrail; - } - - const skipToolChoice = values.skip_tool_message_choice as SkipToolMessageChoice | undefined; - if (skipToolChoice === "yes") { - litellm_params.skip_tool_message_in_guardrail = true; - } else if (skipToolChoice === "no") { - litellm_params.skip_tool_message_in_guardrail = false; - } else { - delete litellm_params.skip_tool_message_in_guardrail; - } - - let guardrail_info: Record = {}; - - // For Presidio PII, add the entity and action configurations - if (values.provider === "PresidioPII" && selectedEntities.length > 0) { - const piiEntitiesConfig: { [key: string]: string } = {}; - selectedEntities.forEach((entity) => { - piiEntitiesConfig[entity] = selectedActions[entity] || "MASK"; // Default to MASK if no action selected - }); - - litellm_params.pii_entities_config = piiEntitiesConfig; - } - // Add config values to the guardrail_info if provided - else if (values.config) { - try { - const configObj = JSON.parse(values.config); - // For some guardrails, the config values need to be in litellm_params - // Especially for providers like Bedrock that need guardrailIdentifier and guardrailVersion - if (values.provider === "Bedrock" && configObj) { - if (configObj.guardrail_id) { - litellm_params.guardrailIdentifier = configObj.guardrail_id; - } - if (configObj.guardrail_version) { - litellm_params.guardrailVersion = configObj.guardrail_version; - } - } else { - // For other providers, add the config to guardrail_info - guardrail_info = configObj; - } - } catch (error) { - NotificationsManager.fromBackend("Invalid JSON in configuration"); - setLoading(false); - return; - } - } - - const guardrailData: { - guardrail_id: string; - guardrail: { - guardrail_name: string; - litellm_params: Record; - guardrail_info: Record; - }; - } = { - guardrail_id: guardrailId, - guardrail: { - guardrail_name: values.guardrail_name, - litellm_params, - guardrail_info, - }, - }; - - if (!accessToken) { - throw new Error("No access token available"); - } - - // Call the update endpoint - const url = `/guardrails/${guardrailId}`; - const response = await fetch(url, { - method: "PUT", - headers: { - [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, - "Content-Type": "application/json", - }, - body: JSON.stringify(guardrailData), - }); - - if (!response.ok) { - const errorData = await response.text(); - throw new Error(errorData || "Failed to update guardrail"); - } - - NotificationsManager.success("Guardrail updated successfully"); - - // Reset and close - onSuccess(); - onClose(); - } catch (error) { - console.error("Failed to update guardrail:", error); - NotificationsManager.fromBackend( - "Failed to update guardrail: " + (error instanceof Error ? error.message : String(error)), - ); - } finally { - setLoading(false); - } - }; - - const renderPiiConfiguration = () => { - if (!guardrailSettings || !selectedProvider || selectedProvider !== "PresidioPII") return null; - - return ( - - ); - }; - - const renderProviderSpecificFields = () => { - if (!selectedProvider) return null; - - // For Presidio, we use the new PII configuration UI - if (selectedProvider === "PresidioPII") { - return renderPiiConfiguration(); - } - - switch (selectedProvider) { - case "Aporia": - return ( - - - - ); - case "AimSecurity": - return ( - - - - ); - case "Bedrock": - return ( - - - - ); - case "CatoNetworks": - return ( - - - - ); - case "GuardrailsAI": - return ( - - - - ); - case "LakeraAI": - return ( - - - - ); - case "PromptInjection": - return ( - - - - ); - default: - return ( - - - - ); - } - }; - - return ( - -
- - - - - - - - - - - - - - - - - - - - - - - - - {renderProviderSpecificFields()} - -
- - -
-
- - ); -}; - -export default EditGuardrailForm; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx new file mode 100644 index 00000000000..9ceb6ba244b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx @@ -0,0 +1,176 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { MoreHorizontal, Trash2 } from "lucide-react"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdentityCell, StatusBadge } from "@/components/shared/table_cells"; +import { Guardrail, GuardrailDefinitionLocation } from "@/components/guardrails/types"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; + +import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; + +const CONFIG_DELETE_HINT = "Config guardrails are defined in the config file and cannot be deleted from the dashboard."; + +function GuardrailProviderCell({ provider }: { provider: string }) { + const { logo, displayName } = getGuardrailLogoAndName(provider); + return ( +
+ {logo ? ( + { + (event.currentTarget as HTMLImageElement).style.display = "none"; + }} + /> + ) : null} + {displayName} +
+ ); +} + +interface GuardrailRowActionsProps { + guardrail: Guardrail; + onDeleteClick: (guardrailId: string, guardrailName: string) => void; +} + +function GuardrailRowActions({ guardrail, onDeleteClick }: GuardrailRowActionsProps) { + const isConfigGuardrail = guardrail.guardrail_definition_location === GuardrailDefinitionLocation.CONFIG; + + return ( + + + + + + onDeleteClick(guardrail.guardrail_id, guardrail.guardrail_name || "Unnamed Guardrail")} + > + + Delete + + + + ); +} + +interface GuardrailTableColumnsDeps { + onGuardrailClick: (guardrailId: string) => void; + onDeleteClick: (guardrailId: string, guardrailName: string) => void; +} + +export const getGuardrailTableColumns = ({ + onGuardrailClick, + onDeleteClick, +}: GuardrailTableColumnsDeps): ColumnDef[] => [ + { + id: "guardrail_id", + accessorKey: "guardrail_id", + meta: { title: "Guardrail ID" }, + header: ({ column }) => , + size: 200, + enableSorting: true, + cell: ({ row }) => ( + onGuardrailClick(row.original.guardrail_id)} + /> + ), + }, + { + id: "guardrail_name", + accessorKey: "guardrail_name", + meta: { title: "Name" }, + header: ({ column }) => , + size: 200, + enableSorting: true, + cell: ({ row }) => { + const name = row.original.guardrail_name; + return ( + + {name || "-"} + + ); + }, + }, + { + id: "provider", + meta: { title: "Provider" }, + header: "Provider", + size: 180, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "mode", + meta: { title: "Mode" }, + header: "Mode", + size: 130, + enableSorting: false, + cell: ({ row }) => ( + {row.original.litellm_params.mode} + ), + }, + { + id: "default_on", + meta: { title: "Default On" }, + header: "Default On", + size: 120, + enableSorting: false, + cell: ({ row }) => { + const isDefaultOn = !!row.original.litellm_params?.default_on; + return ( + + ); + }, + }, + { + id: "created_at", + accessorKey: "created_at", + meta: { title: "Created At" }, + header: ({ column }) => , + size: 150, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "updated_at", + accessorKey: "updated_at", + meta: { title: "Updated At" }, + header: ({ column }) => , + size: 150, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + }, +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx index 70b5998ad95..4f556e74c16 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx @@ -1,56 +1,61 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, it, expect, vi } from "vitest"; + import GuardrailTable from "./guardrail_table"; -import { render } from "@testing-library/react"; -import { describe, it, expect } from "vitest"; -import { GuardrailDefinitionLocation } from "@/components/guardrails/types"; +import { Guardrail, GuardrailDefinitionLocation } from "@/components/guardrails/types"; + +const baseProps = { + isLoading: false, + onDeleteClick: vi.fn(), + onGuardrailClick: vi.fn(), +}; + +const makeGuardrail = (overrides: Partial = {}): Guardrail => ({ + guardrail_id: "gr-1", + guardrail_name: "PII Redaction", + litellm_params: { guardrail: "presidio", mode: "pre_call", default_on: true }, + guardrail_info: null, + created_at: "2021-01-01", + updated_at: "2021-01-02", + guardrail_definition_location: GuardrailDefinitionLocation.DB, + ...overrides, +}); + describe("GuardrailTable", () => { - it("should render", () => { - const { getByText } = render( - {}} - accessToken={null} - onGuardrailUpdated={() => {}} - onGuardrailClick={() => {}} - />, - ); - expect(getByText("Guardrail ID")).toBeInTheDocument(); - expect(getByText("Name")).toBeInTheDocument(); - expect(getByText("Provider")).toBeInTheDocument(); - expect(getByText("Mode")).toBeInTheDocument(); - expect(getByText("Default On")).toBeInTheDocument(); - expect(getByText("Created At")).toBeInTheDocument(); - expect(getByText("Updated At")).toBeInTheDocument(); + it("renders every column header", () => { + render(); + for (const header of ["Guardrail ID", "Name", "Provider", "Mode", "Default On", "Created At", "Updated At"]) { + expect(screen.getByText(header)).toBeInTheDocument(); + } }); - it("should not allow deletion of config guardrails", () => { - const { getByTestId } = render( - {}} - accessToken={null} - onGuardrailUpdated={() => {}} - onGuardrailClick={() => {}} - />, - ); + it("deletes a DB guardrail through the actions menu", async () => { + const user = userEvent.setup(); + const onDeleteClick = vi.fn(); + const guardrail = makeGuardrail({ guardrail_id: "gr-9", guardrail_name: "Toxicity Filter" }); + render(); - const deleteGuardrailButton = getByTestId("config-delete-icon"); - expect(deleteGuardrailButton).toBeInTheDocument(); - expect(deleteGuardrailButton).toHaveClass("cursor-not-allowed text-gray-400"); - expect(deleteGuardrailButton).toHaveAttribute( - "title", - "Config guardrail cannot be deleted on the dashboard. Please delete it from the config file.", - ); + await user.click(screen.getByTestId("guardrail-actions-gr-9")); + await user.click(await screen.findByTestId("guardrail-action-delete")); + + expect(onDeleteClick).toHaveBeenCalledWith("gr-9", "Toxicity Filter"); + }); + + it("disables deletion for config guardrails so they cannot be removed from the dashboard", async () => { + const user = userEvent.setup(); + const onDeleteClick = vi.fn(); + const guardrail = makeGuardrail({ + guardrail_id: "cfg-1", + guardrail_name: "Config Guardrail", + guardrail_definition_location: GuardrailDefinitionLocation.CONFIG, + }); + render(); + + await user.click(screen.getByTestId("guardrail-actions-cfg-1")); + const deleteItem = await screen.findByTestId("guardrail-action-delete"); + + expect(deleteItem).toHaveAttribute("data-disabled"); + expect(onDeleteClick).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx index 4ea4505f506..e6a14b2b2f4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx @@ -1,284 +1,61 @@ -import React, { useState } from "react"; -import { Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Icon } from "@tremor/react"; -import { TrashIcon, SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon } from "@heroicons/react/outline"; -import { Tooltip } from "antd"; -import { DateCell, IdCell, StatusBadge } from "@/components/shared/table_cells"; -import { - ColumnDef, - flexRender, - getCoreRowModel, - getSortedRowModel, - SortingState, - useReactTable, -} from "@tanstack/react-table"; -import { - getGuardrailLogoAndName, - guardrail_provider_map, - skipSystemMessageToChoice, - skipToolMessageToChoice, -} from "./guardrail_info_helpers"; -import EditGuardrailForm from "./edit_guardrail_form"; -import { Guardrail, GuardrailDefinitionLocation } from "@/components/guardrails/types"; +"use client"; + +import { SortingState } from "@tanstack/react-table"; +import { Inbox } from "lucide-react"; +import React, { useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; +import { Guardrail } from "@/components/guardrails/types"; + +import { getGuardrailTableColumns } from "./guardrailTableColumns"; interface GuardrailTableProps { guardrailsList: Guardrail[]; isLoading: boolean; onDeleteClick: (guardrailId: string, guardrailName: string) => void; - accessToken: string | null; - onGuardrailUpdated: () => void; - isAdmin?: boolean; onGuardrailClick: (id: string) => void; } +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +function EmptyState() { + return ( +
+
+ +
+
No guardrails yet
+
Add a guardrail to start filtering requests and responses.
+
+ ); +} + const GuardrailTable: React.FC = ({ guardrailsList, isLoading, onDeleteClick, - accessToken, - onGuardrailUpdated, - isAdmin = false, onGuardrailClick, }) => { - const [sorting, setSorting] = useState([{ id: "created_at", desc: true }]); - const [editModalVisible, setEditModalVisible] = useState(false); - const [selectedGuardrail, setSelectedGuardrail] = useState(null); + const [sorting, setSorting] = useState(DEFAULT_SORTING); - const handleEditClick = (guardrail: Guardrail) => { - setSelectedGuardrail(guardrail); - setEditModalVisible(true); - }; - - const handleEditSuccess = () => { - setEditModalVisible(false); - setSelectedGuardrail(null); - onGuardrailUpdated(); - }; - - const columns: ColumnDef[] = [ - { - header: "Guardrail ID", - accessorKey: "guardrail_id", - cell: (info: any) => , - }, - { - header: "Name", - accessorKey: "guardrail_name", - cell: ({ row }) => { - const guardrail = row.original; - return ( - - {guardrail.guardrail_name || "-"} - - ); - }, - }, - { - header: "Provider", - accessorKey: "litellm_params.guardrail", - cell: ({ row }) => { - const guardrail = row.original; - const { logo, displayName } = getGuardrailLogoAndName(guardrail.litellm_params.guardrail); - return ( -
- {logo && ( - {`${displayName} { - // Hide broken image - (e.target as HTMLImageElement).style.display = "none"; - }} - /> - )} - {displayName} -
- ); - }, - }, - { - header: "Mode", - accessorKey: "litellm_params.mode", - cell: ({ row }) => { - const guardrail = row.original; - return {guardrail.litellm_params.mode}; - }, - }, - { - header: "Default On", - accessorKey: "litellm_params.default_on", - cell: ({ row }) => { - const isDefaultOn = !!row.original.litellm_params?.default_on; - return ( - - ); - }, - }, - { - header: "Created At", - accessorKey: "created_at", - cell: ({ row }) => , - }, - { - header: "Updated At", - accessorKey: "updated_at", - cell: ({ row }) => , - }, - { - id: "actions", - header: "Actions", - cell: ({ row }) => { - const guardrail = row.original; - const isConfigGuardrail = guardrail.guardrail_definition_location === GuardrailDefinitionLocation.CONFIG; - return ( -
- {isConfigGuardrail ? ( - - - - ) : ( - - - guardrail.guardrail_id && - onDeleteClick(guardrail.guardrail_id, guardrail.guardrail_name || "Unnamed Guardrail") - } - className="cursor-pointer hover:text-red-500" - /> - - )} -
- ); - }, - }, - ]; - - const table = useReactTable({ - data: guardrailsList, - columns, - state: { - sorting, - }, - onSortingChange: setSorting, - getCoreRowModel: getCoreRowModel(), - getSortedRowModel: getSortedRowModel(), - enableSorting: true, - }); + const columns = useMemo( + () => getGuardrailTableColumns({ onGuardrailClick, onDeleteClick }), + [onGuardrailClick, onDeleteClick], + ); return ( -
-
- - - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - -
-
- {header.isPlaceholder ? null : flexRender(header.column.columnDef.header, header.getContext())} -
- {header.id !== "actions" && ( -
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
- )} -
- - ))} - - ))} - - - {isLoading ? ( - - -
-

Loading...

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

No guardrails found

-
- - - )} - -
-
- - {/* Edit Modal */} - {selectedGuardrail && ( - setEditModalVisible(false)} - accessToken={accessToken} - onSuccess={handleEditSuccess} - guardrailId={selectedGuardrail.guardrail_id || ""} - fullLitellmParams={selectedGuardrail.litellm_params} - initialValues={{ - guardrail_name: selectedGuardrail.guardrail_name || "", - provider: - Object.keys(guardrail_provider_map).find( - (key) => guardrail_provider_map[key] === selectedGuardrail?.litellm_params.guardrail, - ) || "", - mode: selectedGuardrail.litellm_params.mode, - default_on: selectedGuardrail.litellm_params.default_on, - pii_entities_config: selectedGuardrail.litellm_params.pii_entities_config, - skip_system_message_choice: skipSystemMessageToChoice( - selectedGuardrail.litellm_params?.skip_system_message_in_guardrail, - ), - skip_tool_message_choice: skipToolMessageToChoice( - selectedGuardrail.litellm_params?.skip_tool_message_in_guardrail, - ), - ...selectedGuardrail.guardrail_info, - }} - /> - )} -
+ guardrail.guardrail_id || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + isLoading={isLoading} + loadingMessage="Loading guardrails…" + noDataMessage={} + size="compact" + /> ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts index 1e614b709e2..b09ab4498f7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts @@ -1,12 +1,10 @@ -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { renderHook, waitFor } from "@testing-library/react"; -import React, { ReactNode } from "react"; +import { renderHook } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { useCustomers, type EndUser } from "./useCustomers"; -const mockGet = vi.fn(); +const useQueryMock = vi.fn(); vi.mock("@/lib/http/api", () => ({ - fetchClient: { GET: (...args: unknown[]) => mockGet(...args) }, + $api: { useQuery: (...args: unknown[]) => useQueryMock(...args) }, })); const mockUseAuthorized = vi.fn(); @@ -14,91 +12,55 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => mockUseAuthorized(), })); -const mockCustomers: EndUser[] = [ - { user_id: "customer-1", alias: "Test Customer 1", spend: 150.5, blocked: false }, - { user_id: "customer-2", alias: null, spend: 0, blocked: true }, -]; +const authorized = { accessToken: "test-access-token", userRole: "Admin" }; -const authorized = { - accessToken: "test-access-token", - userRole: "Admin", - userId: "test-user-id", - token: "test-token", - userEmail: "test@example.com", - premiumUser: false, - disabledPersonalKeyCreation: null, - showSSOBanner: false, +type QueryOptions = { enabled: boolean; select: (data: EndUser[] | undefined) => EndUser[] }; + +const lastCallOptions = (): QueryOptions => { + const calls = useQueryMock.mock.calls; + return calls[calls.length - 1][3] as QueryOptions; }; describe("useCustomers", () => { - let queryClient: QueryClient; - beforeEach(() => { - queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); vi.clearAllMocks(); + useQueryMock.mockReturnValue({ data: [] }); mockUseAuthorized.mockReturnValue(authorized); }); - const wrapper = ({ children }: { children: ReactNode }) => - React.createElement(QueryClientProvider, { client: queryClient }, children); - - it("fetches /customer/list and returns the typed list on success", async () => { - mockGet.mockResolvedValue({ data: mockCustomers }); - - const { result } = renderHook(() => useCustomers(), { wrapper }); - - expect(result.current.isLoading).toBe(true); - - await waitFor(() => { - expect(result.current.isSuccess).toBe(true); - }); - - expect(result.current.data).toEqual(mockCustomers); - expect(mockGet).toHaveBeenCalledWith("/customer/list"); - expect(mockGet).toHaveBeenCalledTimes(1); + it("queries GET /customer/list with a derived key (no hand-written queryKey)", () => { + renderHook(() => useCustomers()); + expect(useQueryMock).toHaveBeenCalledWith("get", "/customer/list", {}, expect.any(Object)); }); - it("surfaces an error when the request rejects", async () => { - const testError = new Error("Failed to fetch customers"); - mockGet.mockRejectedValue(testError); - - const { result } = renderHook(() => useCustomers(), { wrapper }); - - await waitFor(() => { - expect(result.current.isError).toBe(true); - }); - - expect(result.current.error).toEqual(testError); - expect(result.current.data).toBeUndefined(); + it("enables the query only for an admin holding an access token", () => { + renderHook(() => useCustomers()); + expect(lastCallOptions().enabled).toBe(true); }); - it("falls back to an empty list when the response has no body", async () => { - mockGet.mockResolvedValue({ data: undefined }); - - const { result } = renderHook(() => useCustomers(), { wrapper }); - - await waitFor(() => { - expect(result.current.isSuccess).toBe(true); - }); - - expect(result.current.data).toEqual([]); + it("disables the query when the access token is missing", () => { + mockUseAuthorized.mockReturnValue({ ...authorized, accessToken: null }); + renderHook(() => useCustomers()); + expect(lastCallOptions().enabled).toBe(false); }); - it("does not fetch when the access token is missing", () => { - mockUseAuthorized.mockReturnValue({ ...authorized, accessToken: null, token: null }); - - const { result } = renderHook(() => useCustomers(), { wrapper }); - - expect(result.current.isFetched).toBe(false); - expect(mockGet).not.toHaveBeenCalled(); - }); - - it("does not fetch when the user is not an admin", () => { + it("disables the query for a non-admin role", () => { mockUseAuthorized.mockReturnValue({ ...authorized, userRole: "member" }); + renderHook(() => useCustomers()); + expect(lastCallOptions().enabled).toBe(false); + }); - const { result } = renderHook(() => useCustomers(), { wrapper }); + it("selects an empty list when the response body is missing", () => { + renderHook(() => useCustomers()); + expect(lastCallOptions().select(undefined)).toEqual([]); + }); - expect(result.current.isFetched).toBe(false); - expect(mockGet).not.toHaveBeenCalled(); + it("selects the customer list through unchanged", () => { + const customers: EndUser[] = [ + { user_id: "customer-1", alias: "Test Customer 1", spend: 150.5, blocked: false }, + { user_id: "customer-2", alias: null, spend: 0, blocked: true }, + ]; + renderHook(() => useCustomers()); + expect(lastCallOptions().select(customers)).toEqual(customers); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.ts index 25e2e3f5e90..ebea4618b7f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.ts @@ -1,19 +1,19 @@ -import { useQuery } from "@tanstack/react-query"; -import { createQueryKeys } from "../common/queryKeysFactory"; -import { fetchClient } from "@/lib/http/api"; +import { $api } from "@/lib/http/api"; import { all_admin_roles } from "@/utils/roles"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import type { components } from "@/lib/http/schema"; export type EndUser = components["schemas"]["CustomerResponse"]; -const customersKeys = createQueryKeys("customers"); - export const useCustomers = () => { const { accessToken, userRole } = useAuthorized(); - return useQuery({ - queryKey: customersKeys.list({}), - queryFn: async () => (await fetchClient.GET("/customer/list")).data ?? [], - enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!), - }); + return $api.useQuery( + "get", + "/customer/list", + {}, + { + enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!), + select: (data) => data ?? [], + }, + ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts index 490e5e3dad1..f532c44ffd7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts @@ -84,6 +84,24 @@ export const teamListCall = async ( } }; +export const teamsTableKeys = createQueryKeys("teamsTable"); + +export const useTeamsTable = ( + page: number, + pageSize: number, + options: TeamListCallOptions = {}, +): UseQueryResult => { + const { accessToken } = useAuthorized(); + + return useQuery({ + queryKey: teamsTableKeys.list({ page, limit: pageSize, ...options }), + queryFn: async () => await teamListCall(accessToken!, page, pageSize, options), + enabled: Boolean(accessToken), + staleTime: 30000, + placeholderData: keepPreviousData, + }); +}; + const teamKeys = createQueryKeys("teams"); export const useTeams = (): UseQueryResult => { const { accessToken, userId, userRole } = useAuthorized(); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/DcrBridgeToggle.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/DcrBridgeToggle.tsx similarity index 95% rename from ui/litellm-dashboard/src/components/mcp_tools/DcrBridgeToggle.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/DcrBridgeToggle.tsx index f1b642293c0..49c182aa6be 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/DcrBridgeToggle.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/DcrBridgeToggle.tsx @@ -1,7 +1,7 @@ import React from "react"; import { Form, Switch, Tooltip } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; -import { isClientForwardedTokenMode } from "./types"; +import { isClientForwardedTokenMode } from "@/components/mcp_tools/types"; /** * DCR-bridge toggle for the client-forwarded token modes (true_passthrough / diff --git a/ui/litellm-dashboard/src/components/mcp_tools/EnvVarsSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/EnvVarsSection.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPLogoSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPLogoSelector.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.test.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPLogoSelector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPLogoSelector.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPLogoSelector.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPNetworkSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPNetworkSettings.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx index 00323465731..7ab240389f3 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPNetworkSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx @@ -1,13 +1,13 @@ import React, { useState, useEffect } from "react"; import { Select, Button, Card, Typography, Spin, Tag } from "antd"; import { SaveOutlined, PlusOutlined } from "@ant-design/icons"; -import { DeprecationBanner } from "../DeprecationBanner"; +import { DeprecationBanner } from "@/components/DeprecationBanner"; import { getGeneralSettingsCall, updateConfigFieldSetting, deleteConfigFieldSetting, fetchMCPClientIp, -} from "../networking"; +} from "@/components/networking"; const { Text } = Typography; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPPermissionManagement.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPPermissionManagement.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.test.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPPermissionManagement.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPPermissionManagement.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx index 27cbdf2ea34..aae13d4b467 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPPermissionManagement.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx @@ -1,7 +1,7 @@ import React, { useEffect } from "react"; import { Alert, Form, Select, Tooltip, Collapse, Input, Space, Button, Switch } from "antd"; import { InfoCircleOutlined, MinusCircleOutlined, PlusOutlined } from "@ant-design/icons"; -import { MCPServer, AUTH_TYPE } from "./types"; +import { MCPServer, AUTH_TYPE } from "@/components/mcp_tools/types"; const { Panel } = Collapse; interface MCPPermissionManagementProps { diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPServerCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.test.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPServerCard.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.test.tsx index 8463a23cc16..a0998b587fb 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPServerCard.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.test.tsx @@ -2,7 +2,7 @@ import React from "react"; import { render, screen } from "@testing-library/react"; import { describe, it, expect, vi } from "vitest"; import MCPServerCard from "./MCPServerCard"; -import type { MCPServer } from "./types"; +import type { MCPServer } from "@/components/mcp_tools/types"; const baseServer: MCPServer = { server_id: "srv-1", diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPServerCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPServerCard.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.tsx index c16ad87980e..4282cdba278 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPServerCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.tsx @@ -8,7 +8,7 @@ import { MoreOutlined, ThunderboltOutlined, } from "@ant-design/icons"; -import { AUTH_TYPE, type MCPServer } from "./types"; +import { AUTH_TYPE, type MCPServer } from "@/components/mcp_tools/types"; import { getMaskedAndFullUrl } from "./utils"; const { Text } = Typography; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPStandardsSettings.test.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPStandardsSettings.test.tsx index fd8d9f92c48..08f65a72537 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPStandardsSettings.test.tsx @@ -1,6 +1,6 @@ import { describe, it, expect } from "vitest"; import { FIELD_GROUPS, MCP_REQUIRED_FIELD_DEFS, SETTINGS_KEY } from "./MCPStandardsSettings"; -import { MCPServer } from "./types"; +import { MCPServer } from "@/components/mcp_tools/types"; const makeServer = (overrides: Partial = {}): MCPServer => ({ server_id: "s1", diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPStandardsSettings.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPStandardsSettings.tsx index fb38e392631..13a9b7c171e 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPStandardsSettings.tsx @@ -1,6 +1,6 @@ "use client"; -import { MCPServer } from "./types"; +import { MCPServer } from "@/components/mcp_tools/types"; export interface RequiredFieldDef { key: string; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPSubmissionsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPSubmissionsTab.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPSubmissionsTab.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPSubmissionsTab.tsx index 50c88a63582..de030420bc3 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPSubmissionsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPSubmissionsTab.tsx @@ -18,7 +18,7 @@ import { getGeneralSettingsCall, updateConfigFieldSetting, } from "@/components/networking"; -import { MCPServer, MCPSubmissionsSummary } from "./types"; +import { MCPServer, MCPSubmissionsSummary } from "@/components/mcp_tools/types"; import { FIELD_GROUPS, MCP_REQUIRED_FIELD_DEFS, SETTINGS_KEY } from "./MCPStandardsSettings"; import NotificationsManager from "@/components/molecules/notifications_manager"; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolsetsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPToolsetsTab.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx index 5e7e99ee8ec..fa44694887b 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolsetsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx @@ -7,9 +7,15 @@ import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolset import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useQueryClient } from "@tanstack/react-query"; import { DateCell, IdCell } from "@/components/shared/table_cells"; -import { DataTable } from "../view_logs/table"; -import { createMCPToolset, updateMCPToolset, deleteMCPToolset, listMCPTools, getProxyBaseUrl } from "../networking"; -import { MCPToolset, MCPToolsetTool } from "./types"; +import { DataTable } from "@/components/view_logs/table"; +import { + createMCPToolset, + updateMCPToolset, + deleteMCPToolset, + listMCPTools, + getProxyBaseUrl, +} from "@/components/networking"; +import { MCPToolset, MCPToolsetTool } from "@/components/mcp_tools/types"; const { Text: AntdText } = Typography; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.test.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx index 93afefe4358..f359bdee065 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx @@ -2,7 +2,7 @@ import React from "react"; import { Form, Input, InputNumber, Select, Tooltip } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TextInput } from "@tremor/react"; -import { OAUTH_FLOW } from "./types"; +import { OAUTH_FLOW } from "@/components/mcp_tools/types"; import TokenEndpointAuthMethodField from "./TokenEndpointAuthMethodField"; interface OAuthFlowStatus { diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OpenAPIFormSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OpenAPIFormSection.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/mcp_tools/OpenAPIFormSection.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OpenAPIFormSection.tsx index d606b8ba1fc..073780b359f 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/OpenAPIFormSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OpenAPIFormSection.tsx @@ -2,7 +2,7 @@ import React, { useState } from "react"; import { Form, Input, Tooltip } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { FormInstance } from "antd/es/form"; -import { AUTH_TYPE, OAUTH_FLOW } from "./types"; +import { AUTH_TYPE, OAUTH_FLOW } from "@/components/mcp_tools/types"; import OpenAPIQuickPicker, { OpenAPIRegistryEntry, OpenAPIKeyTool } from "./OpenAPIQuickPicker"; interface OpenAPIFormSectionProps { diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OpenAPIQuickPicker.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OpenAPIQuickPicker.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/mcp_tools/OpenAPIQuickPicker.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OpenAPIQuickPicker.tsx index a8208a09286..0aec81fdf4b 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/OpenAPIQuickPicker.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OpenAPIQuickPicker.tsx @@ -1,6 +1,6 @@ import React, { useEffect, useState } from "react"; import { Spin } from "antd"; -import { fetchOpenAPIRegistry } from "../networking"; +import { fetchOpenAPIRegistry } from "@/components/networking"; export interface OpenAPIKeyTool { name: string; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/PassthroughAuthorizeSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/PassthroughAuthorizeSection.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.test.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/PassthroughAuthorizeSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/PassthroughAuthorizeSection.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.tsx index 0ed4ee555d1..dc10f0f1392 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/PassthroughAuthorizeSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.tsx @@ -1,7 +1,7 @@ import React from "react"; import { Button, Checkbox, Form, Input } from "antd"; import DcrBridgeToggle from "./DcrBridgeToggle"; -import { credentialAuthClass, isClientForwardedTokenMode } from "./types"; +import { credentialAuthClass, isClientForwardedTokenMode } from "@/components/mcp_tools/types"; interface PassthroughOAuthFlow { startOAuthFlow: () => void | Promise; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/StdioConfiguration.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/StdioConfiguration.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/StdioConfiguration.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/StdioConfiguration.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/TokenEndpointAuthMethodField.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/TokenEndpointAuthMethodField.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/TokenEndpointAuthMethodField.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/TokenEndpointAuthMethodField.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/TokenExchangeFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/TokenExchangeFormFields.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/TokenExchangeFormFields.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/TokenExchangeFormFields.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.test.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.test.tsx index 4bd351216b8..0613e84feed 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.test.tsx @@ -3,9 +3,9 @@ import { render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import { ToolTestPanel } from "./ToolTestPanel"; -import { InputSchema, MCPTool } from "./types"; +import { InputSchema, MCPTool } from "@/components/mcp_tools/types"; -vi.mock("../molecules/notifications_manager", () => ({ +vi.mock("@/components/molecules/notifications_manager", () => ({ default: { success: vi.fn(), fromBackend: vi.fn(), diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.tsx index c15150eee58..8f042445c01 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.tsx @@ -1,10 +1,10 @@ import React from "react"; import { Button, TextInput } from "@tremor/react"; -import { MCPTool, InputSchema, InputSchemaProperty } from "./types"; +import { MCPTool, InputSchema, InputSchemaProperty } from "@/components/mcp_tools/types"; import { resolveLogoSrc } from "@/lib/assetPaths"; import { Form, Select, Tooltip } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; -import NotificationsManager from "../molecules/notifications_manager"; +import NotificationsManager from "@/components/molecules/notifications_manager"; const isPlainObject = (value: unknown): value is Record => typeof value === "object" && value !== null && !Array.isArray(value); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/TruePassthroughWarning.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/TruePassthroughWarning.tsx similarity index 94% rename from ui/litellm-dashboard/src/components/mcp_tools/TruePassthroughWarning.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/TruePassthroughWarning.tsx index b52d3f4c672..9c57cbd7d14 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/TruePassthroughWarning.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/TruePassthroughWarning.tsx @@ -1,6 +1,6 @@ import React from "react"; import { Alert } from "antd"; -import { AUTH_TYPE } from "./types"; +import { AUTH_TYPE } from "@/components/mcp_tools/types"; /** * Warning shown in the create/edit MCP server forms when auth_type diff --git a/ui/litellm-dashboard/src/components/mcp_tools/UserEnvVarsModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx similarity index 95% rename from ui/litellm-dashboard/src/components/mcp_tools/UserEnvVarsModal.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx index 08a285cd56b..f76aef365f7 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/UserEnvVarsModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx @@ -1,9 +1,9 @@ import React from "react"; import { Modal, Form, Input, Button, Alert, Spin, Tag, Typography } from "antd"; import { useMutation, useQuery } from "@tanstack/react-query"; -import { MCPServer, MCPUserEnvVarsStatus } from "./types"; -import { getMCPUserEnvVars, storeMCPUserEnvVars } from "../networking"; -import NotificationsManager from "../molecules/notifications_manager"; +import { MCPServer, MCPUserEnvVarsStatus } from "@/components/mcp_tools/types"; +import { getMCPUserEnvVars, storeMCPUserEnvVars } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; const { Text, Title } = Typography; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/create_mcp_server.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/create_mcp_server.test.tsx index 32fe439a316..6ce7f5c75ed 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/create_mcp_server.test.tsx @@ -1,12 +1,12 @@ import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import * as networking from "../networking"; +import * as networking from "@/components/networking"; import { setToken } from "@/utils/mcpTokenStore"; import CreateMCPServer from "./create_mcp_server"; import { selectAntOption } from "./testUtils"; -vi.mock("../networking", () => ({ +vi.mock("@/components/networking", () => ({ createMCPServer: vi.fn(), fetchOpenAPIRegistry: vi.fn().mockResolvedValue({ apis: [] }), registerMCPServer: vi.fn(), diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/create_mcp_server.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/create_mcp_server.tsx index b9a6f5d229d..70838c592c8 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/create_mcp_server.tsx @@ -2,7 +2,7 @@ import React, { useState } from "react"; import { Modal, Tooltip, Form, Select, Input, InputNumber, Switch, Collapse } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TextInput } from "@tremor/react"; -import { createMCPServer, registerMCPServer, storeMCPOAuthUserCredential } from "../networking"; +import { createMCPServer, registerMCPServer, storeMCPOAuthUserCredential } from "@/components/networking"; import { setToken } from "@/utils/mcpTokenStore"; import { AUTH_TYPE, @@ -20,7 +20,7 @@ import { isHeldOAuthTokenStale, preservedDeclaredAppCredentials, withoutMintedTokenCredentials, -} from "./types"; +} from "@/components/mcp_tools/types"; import OAuthFormFields from "./OAuthFormFields"; import TruePassthroughWarning from "./TruePassthroughWarning"; import PassthroughAuthorizeSection from "./PassthroughAuthorizeSection"; @@ -35,7 +35,7 @@ import MCPLogoSelector from "./MCPLogoSelector"; import EnvVarsSection from "./EnvVarsSection"; import { isAdminRole } from "@/utils/roles"; import { validateMCPServerUrl, validateMCPServerName, normalizeEnvVars, TOOL_DISPLAY_NAME_PATTERN } from "./utils"; -import NotificationsManager from "../molecules/notifications_manager"; +import NotificationsManager from "@/components/molecules/notifications_manager"; import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; import { useTestMCPConnection } from "@/hooks/useTestMCPConnection"; import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/index.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/index.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/index.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx index 42c557f142f..7bdfd9c6b8f 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx @@ -4,8 +4,8 @@ import React, { useState } from "react"; import { Card, Typography, Space, Alert, Button, Switch, Form, Collapse } from "antd"; import { TabPanel, TabPanels, TabGroup, TabList, Tab, Title as TremorTitle, Text as TremorText } from "@tremor/react"; import { CopyIcon, Code, Terminal, Globe, CheckIcon, ExternalLinkIcon, KeyIcon, ServerIcon, Zap } from "lucide-react"; -import { getProxyBaseUrl } from "../networking"; -import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; +import { getProxyBaseUrl } from "@/components/networking"; +import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils"; const { Title, Text } = Typography; const { Panel } = Collapse; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connection_status.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connection_status.test.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connection_status.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connection_status.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_discovery.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_discovery.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_discovery.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_discovery.tsx index 189b52dd6b2..6fcff011ba6 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_discovery.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_discovery.tsx @@ -1,7 +1,7 @@ import React, { useState, useMemo, useEffect } from "react"; import { Modal, Input, Typography } from "antd"; -import { fetchDiscoverableMCPServers } from "../networking"; -import { DiscoverableMCPServer, DiscoverMCPServersResponse } from "./types"; +import { fetchDiscoverableMCPServers } from "@/components/networking"; +import { DiscoverableMCPServer, DiscoverMCPServersResponse } from "@/components/mcp_tools/types"; import { mcpLogoImg } from "./create_mcp_server"; import { resolveLogoSrc } from "@/lib/assetPaths"; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_cost_config.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_config.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_server_cost_config.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_config.tsx index 3f3986d5a2e..89c41693a4c 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_cost_config.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_config.tsx @@ -2,7 +2,7 @@ import React from "react"; import { Tooltip, InputNumber, Collapse, Badge } from "antd"; import { InfoCircleOutlined, DollarOutlined, ToolOutlined } from "@ant-design/icons"; import { Card, Title, Text } from "@tremor/react"; -import { MCPServerCostInfo } from "./types"; +import { MCPServerCostInfo } from "@/components/mcp_tools/types"; interface MCPServerCostConfigProps { value?: MCPServerCostInfo; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_cost_display.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_display.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_server_cost_display.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_display.tsx index e41a06b879e..f26f7ba2320 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_cost_display.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_cost_display.tsx @@ -1,6 +1,6 @@ import React from "react"; import { Text } from "@tremor/react"; -import { MCPServerCostInfo } from "./types"; +import { MCPServerCostInfo } from "@/components/mcp_tools/types"; interface MCPServerCostDisplayProps { costConfig?: MCPServerCostInfo | null; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx index 5b8d0aac8c0..278bf6f7e13 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx @@ -4,18 +4,18 @@ import { render, screen, waitFor, fireEvent, act } from "@testing-library/react" import userEvent from "@testing-library/user-event"; import MCPServerEdit, { EDIT_OAUTH_UI_STATE_KEY } from "./mcp_server_edit"; import { setSecureItem } from "@/utils/secureStorage"; -import * as networking from "../networking"; -import NotificationsManager from "../molecules/notifications_manager"; +import * as networking from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; import { selectAntOption } from "./testUtils"; -vi.mock("../networking", () => ({ +vi.mock("@/components/networking", () => ({ updateMCPServer: vi.fn(), listMCPTools: vi.fn().mockResolvedValue({ tools: [], error: null }), storeMCPOAuthUserCredential: vi.fn().mockResolvedValue({}), testMCPToolsListRequest: vi.fn().mockResolvedValue({ tools: [], error: null }), })); -vi.mock("../molecules/notifications_manager", () => ({ +vi.mock("@/components/molecules/notifications_manager", () => ({ default: { success: vi.fn(), fromBackend: vi.fn(), diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index 6709eb02c65..3b184cc6f1e 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -18,8 +18,13 @@ import { TRANSPORT, getMcpOAuthMode, oauth2FlowToFormValue, -} from "./types"; -import { updateMCPServer, listMCPTools, storeMCPOAuthUserCredential, testMCPToolsListRequest } from "../networking"; +} from "@/components/mcp_tools/types"; +import { + updateMCPServer, + listMCPTools, + storeMCPOAuthUserCredential, + testMCPToolsListRequest, +} from "@/components/networking"; import { getToken, isTokenValid, removeToken, setToken } from "@/utils/mcpTokenStore"; import { buildMcpPassthroughAuthHeader } from "@/utils/mcpHeaderUtils"; import MCPServerCostConfig from "./mcp_server_cost_config"; @@ -39,7 +44,7 @@ import { normalizeToolOverrideMap, TOOL_DISPLAY_NAME_PATTERN, } from "./utils"; -import NotificationsManager from "../molecules/notifications_manager"; +import NotificationsManager from "@/components/molecules/notifications_manager"; import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx index 7ed00f74225..620f76a739a 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx @@ -2,7 +2,7 @@ import React, { useState } from "react"; import { ArrowLeftIcon, EyeIcon, EyeOffIcon } from "@heroicons/react/outline"; import { Title, Card, Button, Text, Grid, TabGroup, TabList, TabPanel, TabPanels, Tab, Icon } from "@tremor/react"; -import { MCPServer, handleTransport, handleAuth } from "./types"; +import { MCPServer, handleTransport, handleAuth } from "@/components/mcp_tools/types"; // TODO: Move Tools viewer from index file import { MCPToolsViewer } from "."; import MCPServerEdit, { EDIT_OAUTH_UI_STATE_KEY } from "./mcp_server_edit"; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx index 446fdd8c22d..d61bc23c757 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx @@ -3,10 +3,10 @@ import { render, waitFor, screen, fireEvent, act } from "@testing-library/react" import { describe, it, expect, vi, beforeEach } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import MCPServers from "./mcp_servers"; -import * as networking from "../networking"; +import * as networking from "@/components/networking"; // Mock the networking module -vi.mock("../networking", () => ({ +vi.mock("@/components/networking", () => ({ fetchMCPServers: vi.fn(), fetchMCPServerHealth: vi.fn(), deleteMCPServer: vi.fn(), @@ -19,7 +19,7 @@ vi.mock("../networking", () => ({ })); // Mock NotificationsManager -vi.mock("../molecules/notifications_manager", () => ({ +vi.mock("@/components/molecules/notifications_manager", () => ({ default: { success: vi.fn(), fromBackend: vi.fn(), diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx index f0d3e60b32e..f186fef22da 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx @@ -1,29 +1,35 @@ import { isAdminRole } from "@/utils/roles"; import { QuestionCircleOutlined, SearchOutlined } from "@ant-design/icons"; import { Button, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; -import NewBadge from "../common_components/NewBadge"; +import NewBadge from "@/components/common_components/NewBadge"; import { Descriptions, Empty, Input, Modal, Select, Spin, Tooltip, Typography } from "antd"; import React, { useEffect, useState, useMemo, useCallback } from "react"; import { useQuery } from "@tanstack/react-query"; -import { useMCPServers } from "../../app/(dashboard)/hooks/mcpServers/useMCPServers"; -import { useMCPServerHealth } from "../../app/(dashboard)/hooks/mcpServers/useMCPServerHealth"; -import NotificationsManager from "../molecules/notifications_manager"; -import { deleteMCPServer } from "../networking"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPServerHealth } from "@/app/(dashboard)/hooks/mcpServers/useMCPServerHealth"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { deleteMCPServer } from "@/components/networking"; import { MCPSubmissionsTab } from "./MCPSubmissionsTab"; import { MCPToolsetsTab } from "./MCPToolsetsTab"; import CreateMCPServer from "./create_mcp_server"; import MCPConnect from "./mcp_connect"; import MCPServerCard from "./MCPServerCard"; import { MCPServerView } from "./mcp_server_view"; -import type { DiscoverableMCPServer, MCPServer, MCPServerProps, MCPUserEnvVarsStatus, Team } from "./types"; -import MCPSemanticFilterSettings from "../Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings"; +import type { + DiscoverableMCPServer, + MCPServer, + MCPServerProps, + MCPUserEnvVarsStatus, + Team, +} from "@/components/mcp_tools/types"; +import MCPSemanticFilterSettings from "@/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings"; import MCPNetworkSettings from "./MCPNetworkSettings"; import MCPDiscovery from "./mcp_discovery"; -import { ByokCredentialModal } from "./ByokCredentialModal"; +import { ByokCredentialModal } from "@/components/mcp_tools/ByokCredentialModal"; import { getSecureItem } from "@/utils/secureStorage"; import { TOOLS_OAUTH_UI_STATE_KEY } from "@/hooks/mcpOAuthUtils"; import UserEnvVarsModal from "./UserEnvVarsModal"; -import { listMCPUserEnvVarStatus } from "../networking"; +import { listMCPUserEnvVarStatus } from "@/components/networking"; type SortKey = "created_desc" | "updated_desc" | "name_asc" | "health"; @@ -678,7 +684,6 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) refetch(); setByokModalServer(null); }} - accessToken={accessToken || ""} /> )} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.test.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx index 1ebc07eac86..60c4c264c3c 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx @@ -2,7 +2,7 @@ import React, { useEffect, useMemo, useRef, useState } from "react"; import { Card, Title, Text } from "@tremor/react"; import { ToolOutlined, CheckCircleOutlined, SearchOutlined, EditOutlined } from "@ant-design/icons"; import { Badge, Spin, Checkbox, Input, Radio } from "antd"; -import McpCrudPermissionPanel from "./McpCrudPermissionPanel"; +import McpCrudPermissionPanel from "@/components/mcp_tools/McpCrudPermissionPanel"; import { TOOL_DISPLAY_NAME_PATTERN } from "./utils"; interface KeyTool { diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tools.test.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tools.test.tsx index 2e8fa901f6d..8b0e6d62f66 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tools.test.tsx @@ -2,10 +2,10 @@ import { render, screen, waitFor } from "@testing-library/react"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { describe, expect, it, vi, beforeEach } from "vitest"; import MCPToolsViewer from "./mcp_tools"; -import { listMCPTools, getMCPOAuthUserCredentialStatus } from "../networking"; +import { listMCPTools, getMCPOAuthUserCredentialStatus } from "@/components/networking"; import { isTokenValid, getToken } from "@/utils/mcpTokenStore"; -vi.mock("../networking", () => ({ +vi.mock("@/components/networking", () => ({ listMCPTools: vi.fn(), callMCPTool: vi.fn(), getMCPOAuthUserCredentialStatus: vi.fn(), diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tools.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tools.tsx index 928a2e3c6bb..428c10da284 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tools.tsx @@ -9,8 +9,8 @@ import { MCPContent, CallMCPToolResponse, getMcpOAuthMode, -} from "./types"; -import { listMCPTools, callMCPTool, getMCPOAuthUserCredentialStatus } from "../networking"; +} from "@/components/mcp_tools/types"; +import { listMCPTools, callMCPTool, getMCPOAuthUserCredentialStatus } from "@/components/networking"; import { isTokenValid, getToken, removeToken } from "@/utils/mcpTokenStore"; import { sanitizeMcpAliasForHeader, buildMcpPassthroughAuthHeader } from "@/utils/mcpHeaderUtils"; import { useToolsOAuthFlow } from "@/hooks/useToolsOAuthFlow"; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/testUtils.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/testUtils.ts similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/testUtils.ts rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/testUtils.ts diff --git a/ui/litellm-dashboard/src/components/mcp_tools/utils.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/utils.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/mcp_tools/utils.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/utils.test.tsx diff --git a/ui/litellm-dashboard/src/components/mcp_tools/utils.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/utils.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/mcp_tools/utils.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/utils.tsx index 7d6e24fc480..4738e1e8fba 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/utils.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/utils.tsx @@ -1,4 +1,4 @@ -import { MCPEnvVar, MCPEnvVarScope } from "./types"; +import { MCPEnvVar, MCPEnvVarScope } from "@/components/mcp_tools/types"; export const extractMCPToken = (url: string): { token: string | null; baseUrl: string } => { try { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/page.tsx index dfc7ca15896..462c48360cd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/page.tsx @@ -1,6 +1,6 @@ "use client"; -import { MCPServers } from "@/components/mcp_tools"; +import { MCPServers } from "./_components"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; export default function McpServers() { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx index 560a4ec15d0..4bf4af0c7f6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx @@ -296,7 +296,7 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te } return ( -
+
{/* Model Management Header */} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index 576729d18d3..20c9e805a0d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -14,11 +14,13 @@ import { useQueryClient } from "@tanstack/react-query"; import { Grid, TabPanel } from "@tremor/react"; import { Badge, Button, Select, Skeleton, Space, Typography } from "antd"; import ModelSettingsModal from "@/components/model_dashboard/ModelSettingsModal/ModelSettingsModal"; -import debounce from "lodash/debounce"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { useEffect, useMemo, useState } from "react"; import { useModelsInfo } from "../../hooks/models/useModels"; import { transformModelData } from "../utils/modelDataTransformer"; type ModelViewMode = "all" | "current_team"; + +const SEARCH_DEBOUNCE_WAIT_MS = 200; const { Text } = Typography; interface AllModelsTabProps { @@ -59,23 +61,17 @@ const AllModelsTab = ({ const [sorting, setSorting] = useState([]); const [isModelSettingsModalVisible, setIsModelSettingsModalVisible] = useState(false); - // Debounce search input - const debouncedUpdateSearch = useMemo( - () => - debounce((value: string) => { - setDebouncedSearch(value); - // Reset to page 1 when search changes - setCurrentPage(1); - setPagination((prev: PaginationState) => ({ ...prev, pageIndex: 0 })); - }, 200), - [], + const debouncedUpdateSearch = useDebouncedCallback( + (value: string) => { + setDebouncedSearch(value); + setCurrentPage(1); + setPagination((prev: PaginationState) => ({ ...prev, pageIndex: 0 })); + }, + { wait: SEARCH_DEBOUNCE_WAIT_MS }, ); useEffect(() => { debouncedUpdateSearch(modelNameSearch); - return () => { - debouncedUpdateSearch.cancel(); - }; }, [modelNameSearch, debouncedUpdateSearch]); // Determine teamId to pass to the query - only pass if not "personal" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx index d3af5b62668..87d8010759d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx @@ -191,7 +191,7 @@ const OrganizationsTable: React.FC = ({ } return ( -
+
{(userRole === "Admin" || userRole === "Org Admin") && ( 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..d8f927b9a63 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 @@ -77,6 +77,7 @@ import { A2ATaskMetadata, MessageType } from "@/components/chat_ui/types"; import { useCodeInterpreter } from "../../hooks/useCodeInterpreter"; import { useChatHistory } from "../../hooks/useChatHistory"; import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; const { TextArea } = Input; const { Dragger } = Upload; @@ -99,6 +100,8 @@ interface ChatUIProps { const MCP_SUPPORTED_ENDPOINTS = new Set([EndpointType.CHAT, EndpointType.RESPONSES, EndpointType.MCP]); +const CUSTOM_MODEL_DEBOUNCE_WAIT_MS = 500; + const ChatUI: React.FC = ({ accessToken, token, @@ -185,7 +188,9 @@ const ChatUI: React.FC = ({ const [modelInfo, setModelInfo] = useState([]); const [agentInfo, setAgentInfo] = useState([]); const [selectedAgent, setSelectedAgent] = useState(undefined); - const customModelTimeout = useRef(null); + const debouncedSetSelectedModel = useDebouncedCallback((value: string) => setSelectedModel(value), { + wait: CUSTOM_MODEL_DEBOUNCE_WAIT_MS, + }); const [endpointType, setEndpointType] = useState( () => sessionStorage.getItem("endpointType") || EndpointType.CHAT, ); @@ -1255,16 +1260,7 @@ const ChatUI: React.FC = ({ { - // Using setTimeout to create a simple debounce effect - if (customModelTimeout.current) { - clearTimeout(customModelTimeout.current); - } - - customModelTimeout.current = setTimeout(() => { - setSelectedModel(value); - }, 500); // 500ms delay after typing stops - }} + onValueChange={debouncedSetSelectedModel} /> )}
@@ -2186,7 +2182,6 @@ const ChatUI: React.FC = ({ loadMCPServers(); setByokModalServer(null); }} - accessToken={accessToken || ""} /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx index 65157206f21..3b396924aba 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx @@ -1,7 +1,9 @@ "use client"; import NotificationsManager from "@/components/molecules/notifications_manager"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { ClearOutlined, DeleteOutlined, FilePdfOutlined, PlusOutlined } from "@ant-design/icons"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { Button, Input, Select, Tooltip } from "antd"; import { useEffect, useMemo, useState } from "react"; import { v4 as uuidv4 } from "uuid"; @@ -105,14 +107,8 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }: disabledPersonalKeyCreation ? "custom" : "session", ); const [customApiKey, setCustomApiKey] = useState(""); - const [debouncedCustomApiKey, setDebouncedCustomApiKey] = useState(""); + const [debouncedCustomApiKey] = useDebouncedValue(customApiKey, { wait: DEBOUNCE_WAIT_MS }); const [customProxyBaseUrl] = useState(() => sessionStorage.getItem("customProxyBaseUrl") || ""); - useEffect(() => { - const timer = setTimeout(() => { - setDebouncedCustomApiKey(customApiKey); - }, 300); - return () => clearTimeout(timer); - }, [customApiKey]); useEffect(() => { return () => { if (uploadedFilePreviewUrl) { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.test.ts index b61f27cf717..762e17a6f91 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.test.ts @@ -1,5 +1,5 @@ import { renderHook, act } from "@testing-library/react"; -import { describe, it, expect, beforeEach, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { useChatHistory } from "./useChatHistory"; describe("useChatHistory", () => { @@ -499,6 +499,80 @@ describe("useChatHistory", () => { }); }); + describe("debounced chatHistory persistence", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + }); + + it("should not write chatHistory to sessionStorage before the debounce wait elapses", () => { + const setItemSpy = vi.spyOn(Storage.prototype, "setItem"); + const { result } = renderHook(() => useChatHistory({ simplified: false })); + + act(() => { + result.current.updateTextUI("user", "hello"); + }); + act(() => { + vi.advanceTimersByTime(499); + }); + + expect(setItemSpy.mock.calls.filter(([key]) => key === "chatHistory")).toHaveLength(0); + + setItemSpy.mockRestore(); + }); + + it("should write chatHistory exactly once with the last value after the wait", () => { + const setItemSpy = vi.spyOn(Storage.prototype, "setItem"); + const { result } = renderHook(() => useChatHistory({ simplified: false })); + + act(() => { + result.current.updateTextUI("user", "h"); + }); + act(() => { + vi.advanceTimersByTime(300); + }); + act(() => { + result.current.updateTextUI("user", "i"); + }); + act(() => { + vi.advanceTimersByTime(499); + }); + + expect(setItemSpy.mock.calls.filter(([key]) => key === "chatHistory")).toHaveLength(0); + + act(() => { + vi.advanceTimersByTime(1); + }); + + const writes = setItemSpy.mock.calls.filter(([key]) => key === "chatHistory"); + expect(writes).toHaveLength(1); + expect(JSON.parse(writes[0][1])).toEqual([{ role: "user", content: "hi" }]); + + setItemSpy.mockRestore(); + }); + + it("should not write chatHistory when unmounted mid-wait", () => { + const setItemSpy = vi.spyOn(Storage.prototype, "setItem"); + const { result, unmount } = renderHook(() => useChatHistory({ simplified: false })); + + act(() => { + result.current.updateTextUI("user", "hello"); + }); + unmount(); + act(() => { + vi.advanceTimersByTime(1000); + }); + + expect(setItemSpy.mock.calls.filter(([key]) => key === "chatHistory")).toHaveLength(0); + + setItemSpy.mockRestore(); + }); + }); + describe("simplified mode session isolation", () => { it("should not hydrate messageTraceId from sessionStorage in simplified mode", () => { sessionStorage.setItem("messageTraceId", "trace-from-playground"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.ts index 6e2a263f599..7ee30757524 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.ts @@ -1,9 +1,12 @@ import React, { useState, useEffect } from "react"; +import { useDebouncer } from "@tanstack/react-pacer/debouncer"; import { MessageType, A2ATaskMetadata } from "@/components/chat_ui/types"; import { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; import { MCPEvent } from "@/components/mcp_tools/types"; import { truncateString } from "@/utils/textUtils"; +const CHAT_HISTORY_PERSIST_WAIT_MS = 500; + export interface UseChatHistoryReturn { // State chatHistory: MessageType[]; @@ -64,20 +67,20 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat return saved ? JSON.parse(saved) : true; // Default to API session management }); - // Debounced chatHistory persistence - useEffect(() => { - if (simplified) return; // Do not persist chat history in simplified (embedded) mode - // When chatHistory is empty (e.g. after clearChatHistory removed the key), - // don't re-write an empty array back into sessionStorage. - if (chatHistory.length === 0) return; - const handler = setTimeout(() => { - sessionStorage.setItem("chatHistory", JSON.stringify(chatHistory)); - }, 500); // Debounce by 500ms + const persistDebouncer = useDebouncer( + (history: MessageType[]) => { + sessionStorage.setItem("chatHistory", JSON.stringify(history)); + }, + { wait: CHAT_HISTORY_PERSIST_WAIT_MS }, + ); - return () => { - clearTimeout(handler); - }; - }, [chatHistory, simplified]); + useEffect(() => { + if (simplified || chatHistory.length === 0) { + persistDebouncer.cancel(); + return; + } + persistDebouncer.maybeExecute(chatHistory); + }, [chatHistory, simplified, persistDebouncer]); // messageTraceId/responsesSessionId/useApiSessionManagement persistence useEffect(() => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx index 043d229e3c1..9c5a5592d5e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx @@ -115,7 +115,7 @@ const TagManagement: React.FC = ({ accessToken, userID, userRole }) => }, [accessToken]); return ( -
+
{selectedTagId ? ( ; + return ; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/EndpointUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/EndpointUsage.tsx index 64fdb13a0b9..51e451ca770 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/EndpointUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/EndpointUsage.tsx @@ -59,7 +59,7 @@ const EndpointUsage: React.FC = ({ userSpendData }) => {
- +
); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageBarChart.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageBarChart.test.tsx index f200b8a6a18..a9e65b21f4b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageBarChart.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageBarChart.test.tsx @@ -1,39 +1,70 @@ -import { render, screen } from "@testing-library/react"; -import { describe, expect, it, vi } from "vitest"; +import { screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; +import { renderWithProviders } from "@/../tests/test-utils"; +import { MetricWithMetadata } from "@/components/UsagePage/types"; import EndpointUsageBarChart from "./EndpointUsageBarChart"; -vi.mock("@tremor/react", async () => { - const React = await import("react"); - - function Card({ children }: any) { - return React.createElement("div", { "data-testid": "tremor-card" }, children); - } - (Card as any).displayName = "Card"; - - function Title({ children }: any) { - return React.createElement("h2", { "data-testid": "tremor-title" }, children); - } - (Title as any).displayName = "Title"; - - function BarChart(_props: any) { - return React.createElement("div", { "data-testid": "tremor-bar-chart" }, "Bar Chart"); - } - (BarChart as any).displayName = "BarChart"; - - return { Card, Title, BarChart }; +const metric = (successful: number, failed: number): MetricWithMetadata => ({ + metrics: { + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: successful + failed, + successful_requests: successful, + failed_requests: failed, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + metadata: {}, + api_key_breakdown: {}, }); -vi.mock("@/components/common_components/chartUtils", () => ({ - CustomLegend: ({ categories }: any) =>
{categories.join(", ")}
, - CustomTooltip: () =>
Tooltip
, -})); +const endpointData = { + "/chat/completions": metric(120, 5), + "/embeddings": metric(40, 2), +}; describe("EndpointUsageBarChart", () => { - it("should render", () => { - render(); + it("renders the title and the header legend labels", () => { + renderWithProviders(); - expect(screen.getByTestId("tremor-card")).toBeInTheDocument(); expect(screen.getByText("Success vs Failed Requests by Endpoint")).toBeInTheDocument(); - expect(screen.getByTestId("tremor-bar-chart")).toBeInTheDocument(); + expect(screen.getByText("Successful Requests")).toBeInTheDocument(); + expect(screen.getByText("Failed Requests")).toBeInTheDocument(); + }); + + it("renders stacked green and red bars per endpoint", () => { + const { container } = renderWithProviders(); + + expect(container.querySelectorAll(".recharts-bar")).toHaveLength(2); + const rectangles = Array.from(container.querySelectorAll("path.recharts-rectangle")); + expect(rectangles).toHaveLength(4); + const fills = new Set(rectangles.map((rect) => rect.getAttribute("fill"))); + expect(fills).toEqual(new Set(["var(--color-green-500, #22c55e)", "var(--color-red-500, #ef4444)"])); + + const xPositions = rectangles.map((rect) => rect.getAttribute("d")?.split(",")[0]); + expect(new Set(xPositions).size).toBe(2); + }); + + it("labels the x axis with endpoint names", () => { + renderWithProviders(); + + expect(screen.getAllByText("/chat/completions").length).toBeGreaterThan(0); + expect(screen.getAllByText("/embeddings").length).toBeGreaterThan(0); + }); + + it("keeps the chart's own legend off; only the header legend is shown", () => { + const { container } = renderWithProviders(); + + expect(container.querySelector(".recharts-legend-wrapper")).toBeNull(); + expect(screen.queryByText("metrics.successful_requests")).not.toBeInTheDocument(); + }); + + it("renders an empty chart without bars when endpointData is absent", () => { + const { container } = renderWithProviders(); + + expect(screen.getByText("Success vs Failed Requests by Endpoint")).toBeInTheDocument(); + expect(container.querySelectorAll("path.recharts-rectangle")).toHaveLength(0); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageBarChart.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageBarChart.tsx index 2badbe30868..bf9868d77cf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageBarChart.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageBarChart.tsx @@ -1,6 +1,6 @@ import React from "react"; -import { BarChart, Card, Title } from "@tremor/react"; -import { CustomLegend, CustomTooltip } from "@/components/common_components/chartUtils"; +import { BarChart, CustomLegend, CustomTooltip } from "@/components/shared/charts"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { MetricWithMetadata } from "@/components/UsagePage/types"; interface EndpointUsageBarChartProps { @@ -8,11 +8,9 @@ interface EndpointUsageBarChartProps { } const EndpointUsageBarChart: React.FC = ({ endpointData }) => { - const dataToUse = endpointData || {}; - // Transform endpoint data into chart format const chartData = React.useMemo(() => { - return Object.entries(dataToUse).map(([endpoint, data]) => ({ + return Object.entries(endpointData || {}).map(([endpoint, data]) => ({ endpoint, "metrics.successful_requests": data.metrics.successful_requests, "metrics.failed_requests": data.metrics.failed_requests, @@ -21,31 +19,34 @@ const EndpointUsageBarChart: React.FC = ({ endpointD failed_requests: data.metrics.failed_requests, }, })); - }, [dataToUse]); + }, [endpointData]); const valueFormatter = (value: number) => value.toLocaleString(); return ( -
- Success vs Failed Requests by Endpoint - +
+ Success vs Failed Requests by Endpoint + +
+ + + -
- + ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageLineChart.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageLineChart.test.tsx index ec99825289c..33914e627dc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageLineChart.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageLineChart.test.tsx @@ -1,34 +1,103 @@ -import { render, screen } from "@testing-library/react"; -import { describe, expect, it, vi } from "vitest"; +import { screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; +import { renderWithProviders } from "@/../tests/test-utils"; +import { DailyData, MetricWithMetadata, SpendMetrics } from "@/components/UsagePage/types"; import EndpointUsageLineChart from "./EndpointUsageLineChart"; -vi.mock("@tremor/react", async () => { - const React = await import("react"); - - function Card({ children }: any) { - return React.createElement("div", { "data-testid": "tremor-card" }, children); - } - (Card as any).displayName = "Card"; - - function Title({ children }: any) { - return React.createElement("h2", { "data-testid": "tremor-title" }, children); - } - (Title as any).displayName = "Title"; - - function LineChart(_props: any) { - return React.createElement("div", { "data-testid": "tremor-line-chart" }, "Line Chart"); - } - (LineChart as any).displayName = "LineChart"; - - return { Card, Title, LineChart }; +const spendMetrics = (apiRequests: number): SpendMetrics => ({ + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: apiRequests, + successful_requests: apiRequests, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }); +const endpointMetric = (apiRequests: number): MetricWithMetadata => ({ + metrics: spendMetrics(apiRequests), + metadata: {}, + api_key_breakdown: {}, +}); + +const day = (date: string, endpoints: Record): DailyData => ({ + date, + metrics: spendMetrics(0), + breakdown: { + models: {}, + model_groups: {}, + mcp_servers: {}, + providers: {}, + api_keys: {}, + entities: {}, + endpoints: Object.fromEntries( + Object.entries(endpoints).map(([name, requests]) => [name, endpointMetric(requests)]), + ), + }, +}); + +const dailyData = { + results: [ + day("2026-06-03T12:00:00", { "/chat/completions": 4000, "/embeddings": 900 }), + day("2026-06-02T12:00:00", { "/chat/completions": 2500, "/embeddings": 700 }), + day("2026-06-01T12:00:00", { "/chat/completions": 1200 }), + ], +}; + describe("EndpointUsageLineChart", () => { - it("should render", () => { - render(); + it("renders the title", () => { + renderWithProviders(); - expect(screen.getByTestId("tremor-card")).toBeInTheDocument(); expect(screen.getByText("Endpoint Usage Trends")).toBeInTheDocument(); - expect(screen.getByTestId("tremor-line-chart")).toBeInTheDocument(); + }); + + it("renders one line per endpoint with the tremor palette strokes", () => { + const { container } = renderWithProviders(); + + const curves = Array.from(container.querySelectorAll("path.recharts-line-curve")); + expect(curves).toHaveLength(2); + expect(new Set(curves.map((curve) => curve.getAttribute("stroke")))).toEqual( + new Set(["var(--color-blue-500, #3b82f6)", "var(--color-cyan-500, #06b6d4)"]), + ); + }); + + it("shows a legend with the endpoint names", () => { + const { container } = renderWithProviders(); + + const legend = container.querySelector(".recharts-legend-wrapper"); + expect(legend).not.toBeNull(); + expect(legend!.textContent).toContain("/chat/completions"); + expect(legend!.textContent).toContain("/embeddings"); + }); + + it("orders formatted dates oldest to newest on the x axis", () => { + const { container } = renderWithProviders(); + + const tickLabels = Array.from(container.querySelectorAll(".recharts-xAxis-tick-labels text")).map( + (tick) => tick.textContent, + ); + expect(tickLabels).toEqual(["Jun 1", "Jun 2", "Jun 3"]); + }); + + it("formats y axis ticks with toLocaleString", () => { + renderWithProviders(); + + expect(screen.getAllByText(/^\d,\d{3}$/).length).toBeGreaterThan(0); + }); + + it("draws smooth natural curves", () => { + const { container } = renderWithProviders(); + + const path = container.querySelector("path.recharts-line-curve")?.getAttribute("d") ?? ""; + expect(path).toContain("C"); + }); + + it("renders an empty chart without lines when dailyData is absent", () => { + const { container } = renderWithProviders(); + + expect(screen.getByText("Endpoint Usage Trends")).toBeInTheDocument(); + expect(container.querySelectorAll("path.recharts-line-curve")).toHaveLength(0); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageLineChart.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageLineChart.tsx index 9f838d6156a..483a30b1639 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageLineChart.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageLineChart.tsx @@ -1,10 +1,10 @@ -import { Card, LineChart, Title } from "@tremor/react"; import { useMemo } from "react"; +import { LineChart, type ChartColor } from "@/components/shared/charts"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { DailyData } from "@/components/UsagePage/types"; interface EndpointUsageLineChartProps { dailyData?: { results: DailyData[] }; - endpointData?: Record; } // Transform daily data into chart format @@ -42,7 +42,7 @@ function transformDailyDataToChart(dailyData: DailyData[]): Array { if (!dailyData?.results || dailyData.results.length === 0) { return []; @@ -59,26 +59,39 @@ export function EndpointUsageLineChart({ dailyData, endpointData }: EndpointUsag }, [chartData]); // Tremor color palette for multiple lines - const colors = ["blue", "cyan", "indigo", "violet", "purple", "fuchsia", "pink", "rose", "red", "orange"]; + const colors: readonly ChartColor[] = [ + "blue", + "cyan", + "indigo", + "violet", + "purple", + "fuchsia", + "pink", + "rose", + "red", + "orange", + ]; return ( -
- Endpoint Usage Trends -
- value.toLocaleString()} - showLegend={true} - showGridLines={true} - yAxisWidth={60} - connectNulls={true} - curveType="natural" - /> + + Endpoint Usage Trends + + + value.toLocaleString()} + showLegend={true} + showGridLines={true} + yAxisWidth={60} + connectNulls={true} + curveType="natural" + /> + ); } 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/app/(dashboard)/vector-stores/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx index fd02effc972..565e645f7b5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx @@ -138,7 +138,7 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID />
) : ( -
+

Vector Store Management

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/app/globals.css b/ui/litellm-dashboard/src/app/globals.css index 0b555fcdd52..4589d0f528a 100644 --- a/ui/litellm-dashboard/src/app/globals.css +++ b/ui/litellm-dashboard/src/app/globals.css @@ -11,6 +11,25 @@ @custom-variant dark (&:where(.dark, .dark *)); +/* shadcn Base UI primitives reference these variants; upstream omits them (shadcn-ui/ui#9196) */ +@custom-variant data-open (&:where([data-state="open"], [data-open]:not([data-open="false"]))); +@custom-variant data-closed (&:where([data-state="closed"], [data-closed]:not([data-closed="false"]))); +@custom-variant data-checked (&:where([data-state="checked"], [data-checked]:not([data-checked="false"]))); +@custom-variant data-unchecked (&:where([data-state="unchecked"], [data-unchecked]:not([data-unchecked="false"]))); +@custom-variant data-selected (&:where([data-selected="true"])); +@custom-variant data-disabled (&:where([data-disabled="true"], [data-disabled]:not([data-disabled="false"]))); +@custom-variant data-active (&:where([data-state="active"], [data-active]:not([data-active="false"]))); +@custom-variant data-horizontal (&:where([data-orientation="horizontal"])); +@custom-variant data-vertical (&:where([data-orientation="vertical"])); + +@utility no-scrollbar { + -ms-overflow-style: none; + scrollbar-width: none; + &::-webkit-scrollbar { + display: none; + } +} + :root { --radius: 0.5rem; --background: oklch(1 0 0); diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index d3dd914f5c5..eb47f775d90 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -382,7 +382,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, } return ( -
+
{publicPage == false ? (
{/* Header with Title, Description and URL */} 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/OldTeams.test.tsx b/ui/litellm-dashboard/src/components/OldTeams.test.tsx deleted file mode 100644 index 0b6e5786aaf..00000000000 --- a/ui/litellm-dashboard/src/components/OldTeams.test.tsx +++ /dev/null @@ -1,1148 +0,0 @@ -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; -import React from "react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; -import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key"; -import { fetchMCPAccessGroups, getGuardrailsList, teamCreateCall } from "./networking"; -import OldTeams from "./OldTeams"; -import { teamListCall } from "@/app/(dashboard)/hooks/teams/useTeams"; - -const mockTeamInfoView = vi.fn(); -const mockUseOrganizations = vi.fn(); - -vi.mock("./networking", () => ({ - teamCreateCall: vi.fn(), - teamDeleteCall: vi.fn(), - fetchMCPAccessGroups: vi.fn(), - v2TeamListCall: vi.fn(), - getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), - getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), -})); - -vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ - teamListCall: vi.fn().mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 0 }), -})); - -vi.mock("./molecules/notifications_manager", () => ({ - default: { - info: vi.fn(), - success: vi.fn(), - error: vi.fn(), - fromBackend: vi.fn(), - }, -})); - -vi.mock("./key_team_helpers/fetch_available_models_team_key", () => ({ - fetchAvailableModelsForTeamOrKey: vi.fn(), - getModelDisplayName: vi.fn((model: string) => model), - unfurlWildcardModelsInList: vi.fn((teamModels: string[], allModels: string[]) => { - const wildcardDisplayNames: string[] = []; - const expandedModels: string[] = []; - - teamModels.forEach((teamModel) => { - if (teamModel.endsWith("/*")) { - const provider = teamModel.replace("/*", ""); - const matchingModels = allModels.filter((model) => model.startsWith(provider + "/")); - expandedModels.push(...matchingModels); - wildcardDisplayNames.push(teamModel); - } else { - expandedModels.push(teamModel); - } - }); - - return [...wildcardDisplayNames, ...expandedModels].filter((item, index, array) => array.indexOf(item) === index); - }), -})); - -vi.mock("@/components/team/TeamInfo", () => ({ - __esModule: true, - default: (props: any) => { - mockTeamInfoView(props); - return
; - }, -})); - -vi.mock("./ModelSelect/ModelSelect", () => { - const ModelSelect = React.forwardRef(({ value, onChange, dataTestId, id }: any, ref: any) => { - return ( - { - if (onChange) { - const newVal = e.target.value - ? e.target.value - .split(",") - .map((s: string) => s.trim()) - .filter(Boolean) - : []; - onChange(newVal); - } - }} - /> - ); - }); - ModelSelect.displayName = "ModelSelect"; - return { - ModelSelect, - }; -}); - -vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ - useOrganizations: () => mockUseOrganizations(), -})); - -vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ - useAccessGroups: vi.fn().mockReturnValue({ - data: [ - { access_group_id: "ag-1", access_group_name: "Group 1" }, - { access_group_id: "ag-2", access_group_name: "Group 2" }, - ], - isLoading: false, - isError: false, - }), -})); - -vi.mock("./common_components/AccessGroupSelector", () => ({ - default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( - onChange?.(e.target.value ? e.target.value.split(",").map((s) => s.trim()) : [])} - /> - ), -})); - -const createQueryClient = () => { - return new QueryClient({ - defaultOptions: { - queries: { - retry: false, - }, - }, - }); -}; - -const renderWithQueryClient = (component: React.ReactElement) => { - const queryClient = createQueryClient(); - return render({component}); -}; - -describe("OldTeams - handleCreate organization handling", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockTeamInfoView.mockClear(); - vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue([]); - vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); - vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); - mockUseOrganizations.mockReturnValue({ data: null }); - }); - - it("should not include organization_id when it's an empty string", async () => { - const mockAccessToken = "test-token"; - const formValues: Record = { - team_alias: "Test Team", - organization_id: "", // Empty string - models: [], - }; - - // Simulate the handleCreate logic - let organizationId = formValues?.organization_id || null; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - formValues.organization_id = organizationId.trim(); - } - - expect(formValues.organization_id).toBeNull(); - expect(formValues.organization_id).not.toBe(""); - }); - - it("should set organization_id to null when it's not a string type", async () => { - const formValues: Record = { - team_alias: "Test Team", - organization_id: undefined, - models: [], - }; - - // Simulate the handleCreate logic - let organizationId = formValues?.organization_id || null; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - formValues.organization_id = organizationId.trim(); - } - - expect(formValues.organization_id).toBeNull(); - }); - - it("should trim and keep valid organization_id string", async () => { - const formValues: Record = { - team_alias: "Test Team", - organization_id: " org-123 ", // String with whitespace - models: [], - }; - - // Simulate the handleCreate logic - let organizationId = formValues?.organization_id || null; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - formValues.organization_id = organizationId.trim(); - } - - expect(formValues.organization_id).toBe("org-123"); - }); - - it("should keep valid organization_id without modification", async () => { - const formValues: Record = { - team_alias: "Test Team", - organization_id: "f874bb43-b898-4813-beca-4054d224eafc", - models: [], - }; - - // Simulate the handleCreate logic - let organizationId = formValues?.organization_id || null; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - formValues.organization_id = organizationId.trim(); - } - - expect(formValues.organization_id).toBe("f874bb43-b898-4813-beca-4054d224eafc"); - }); - - it("should not send organization_id field when converting empty string to null", async () => { - const formValues: Record = { - team_alias: "Test Team", - organization_id: "", - models: ["gpt-4"], - max_budget: 100, - }; - - // Simulate the handleCreate logic - let organizationId = formValues?.organization_id || null; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - formValues.organization_id = organizationId.trim(); - } - - // Verify the structure - expect(formValues).toEqual({ - team_alias: "Test Team", - organization_id: null, - models: ["gpt-4"], - max_budget: 100, - }); - - // Verify we're not sending an empty string - expect(formValues.organization_id).not.toBe(""); - - // Verify it's explicitly null, not undefined - expect(formValues.organization_id).toBeNull(); - }); - - it("should handle when currentOrg is used as fallback", async () => { - const currentOrg = { - organization_id: "fallback-org-id", - organization_alias: "Fallback Org", - models: [], - members: [], - }; - - const formValues: Record = { - team_alias: "Test Team", - models: [], - }; - - // Simulate the handleCreate logic with currentOrg fallback - let organizationId = formValues?.organization_id || currentOrg?.organization_id; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - formValues.organization_id = organizationId.trim(); - } - - expect(formValues.organization_id).toBe("fallback-org-id"); - }); - - it("should not include organizations as an empty array in the request payload", async () => { - const mockTeamCreateCall = vi.mocked(teamCreateCall); - const mockAccessToken = "test-token"; - - const formValues = { - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - organizations: [], // This should never be sent - }; - - // Remove organizations key if it's empty - if (Array.isArray(formValues.organizations) && formValues.organizations.length === 0) { - delete (formValues as any).organizations; - } - - // Verify organizations key is removed - expect(formValues).not.toHaveProperty("organizations"); - expect(formValues).toEqual({ - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - }); - }); - - it("should handle organization_id validation for org admins", () => { - // This test simulates the validation that should happen for org admins - const isOrgAdmin = true; - const formValues: Record = { - team_alias: "Test Team", - // organization_id is missing/undefined - }; - - // For org admins, organization_id should be required - const hasOrganization = - formValues.organization_id !== undefined && - formValues.organization_id !== null && - formValues.organization_id !== ""; - - if (isOrgAdmin && !hasOrganization) { - // This should trigger validation error - expect(hasOrganization).toBe(false); - } - }); - - it("should allow null organization_id for global admins", () => { - const isAdmin = true; - const formValues: Record = { - team_alias: "Test Team", - organization_id: null, - models: [], - }; - - // Global admins can create teams without an organization - if (isAdmin) { - expect(formValues.organization_id).toBeNull(); - // This is valid for admins - } - }); - - it("should ensure organization_id is never an empty list", () => { - const invalidFormValues: Record = { - team_alias: "Test Team", - organization_id: [], // Wrong type - should be string or null - }; - - // Type check: organization_id should never be an array - expect(Array.isArray(invalidFormValues.organization_id)).toBe(true); - - // Correct it to null - if (Array.isArray(invalidFormValues.organization_id)) { - invalidFormValues.organization_id = null; - } - - expect(invalidFormValues.organization_id).toBeNull(); - expect(Array.isArray(invalidFormValues.organization_id)).toBe(false); - }); - - it("should clear the delete modal when the cancel button is clicked", async () => { - mockUseOrganizations.mockReturnValue({ data: [] }); - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - await waitFor(() => { - expect(screen.getByTestId("delete-team-button")).toBeInTheDocument(); - }); - const deleteTeamButton = screen.getByTestId("delete-team-button"); - act(() => { - fireEvent.click(deleteTeamButton); - }); - expect(screen.getByText("Delete Team?")).toBeInTheDocument(); - }); -}); - -describe("OldTeams - empty state", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - it("should display empty state message when teams array is empty", async () => { - vi.mocked(teamListCall).mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 1 }); - renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("No teams yet")).toBeInTheDocument(); - }); - expect( - screen.getByText("Create your first team to organize members and manage access to models."), - ).toBeInTheDocument(); - }); - - it("should display empty state message when teams is null", async () => { - vi.mocked(teamListCall).mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 1 }); - renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("No teams yet")).toBeInTheDocument(); - }); - expect( - screen.getByText("Create your first team to organize members and manage access to models."), - ).toBeInTheDocument(); - }); - - it("should not display empty state when teams array has items", async () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("Test Team")).toBeInTheDocument(); - }); - expect(screen.queryByText("No teams yet")).not.toBeInTheDocument(); - expect( - screen.queryByText("Create your first team to organize members and manage access to models."), - ).not.toBeInTheDocument(); - }); -}); - -describe("OldTeams - helper functions", () => { - describe("getAdminOrganizations", () => { - it("should return all organizations for Admin role", () => { - const organizations = [ - { - organization_id: "org-1", - organization_alias: "Org 1", - models: [], - members: [], - }, - { - organization_id: "org-2", - organization_alias: "Org 2", - models: [], - members: [], - }, - ]; - - // Simulate getAdminOrganizations logic for Admin - const userRole = "Admin"; - const result = userRole === "Admin" ? organizations : []; - - expect(result).toEqual(organizations); - expect(result.length).toBe(2); - }); - - it("should return only org_admin organizations for Org Admin role", () => { - const userID = "user-123"; - const userRole = "Org Admin"; - const organizations = [ - { - organization_id: "org-1", - organization_alias: "Org 1", - models: [], - members: [{ user_id: "user-123", user_role: "org_admin" }], - }, - { - organization_id: "org-2", - organization_alias: "Org 2", - models: [], - members: [{ user_id: "user-456", user_role: "org_admin" }], - }, - { - organization_id: "org-3", - organization_alias: "Org 3", - models: [], - members: [{ user_id: "user-123", user_role: "member" }], - }, - ]; - - // Simulate getAdminOrganizations logic - const result = organizations.filter((org) => - org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), - ); - - expect(result.length).toBe(1); - expect(result[0].organization_id).toBe("org-1"); - }); - - it("should return empty array when user is not admin of any organization", () => { - const userID = "user-999"; - const organizations = [ - { - organization_id: "org-1", - organization_alias: "Org 1", - models: [], - members: [{ user_id: "user-123", user_role: "org_admin" }], - }, - ]; - - // Simulate getAdminOrganizations logic - const result = organizations.filter((org) => - org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), - ); - - expect(result.length).toBe(0); - }); - }); - - describe("canCreateOrManageTeams", () => { - it("should return true for Admin role", () => { - const userRole = "Admin"; - const result = userRole === "Admin"; - expect(result).toBe(true); - }); - - it("should return true for org_admin in any organization", () => { - const userID = "user-123"; - const organizations = [ - { - organization_id: "org-1", - organization_alias: "Org 1", - models: [], - members: [{ user_id: "user-123", user_role: "org_admin" }], - }, - ]; - - const result = organizations.some((org) => - org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), - ); - - expect(result).toBe(true); - }); - - it("should return false when user has no admin permissions", () => { - const userID = "user-123"; - const userRole: string = "User"; - const organizations = [ - { - organization_id: "org-1", - organization_alias: "Org 1", - models: [], - members: [{ user_id: "user-123", user_role: "member" }], - }, - ]; - - const isAdmin = userRole === "Admin"; - const isOrgAdmin = organizations.some((org) => - org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), - ); - - expect(isAdmin || isOrgAdmin).toBe(false); - }); - }); -}); - -describe("OldTeams - premium props", () => { - beforeEach(() => { - mockTeamInfoView.mockClear(); - vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue([]); - vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); - vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - it("passes premiumUser flag to TeamInfoView", async () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "team-123456789", - team_alias: "Premium Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - const teamIdElement = await screen.findByText("team-123456789"); - act(() => { - fireEvent.click(teamIdElement); - }); - - await waitFor(() => expect(mockTeamInfoView).toHaveBeenCalled()); - - expect(mockTeamInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ premiumUser: true })); - }); -}); - -describe("OldTeams - Default Team Settings tab visibility", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - it("should show Default Team Settings tab for Admin role", () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - expect(screen.getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument(); - }); - - it("should show Default Team Settings tab for proxy_admin role", () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - expect(screen.getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument(); - }); - - it("should not show Default Team Settings tab for proxy_admin_viewer role", () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - expect(screen.queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument(); - }); - - it("should not show Default Team Settings tab for Admin Viewer role", () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - expect(screen.queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument(); - }); -}); - -describe("OldTeams - access_group_ids in team create", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockTeamInfoView.mockClear(); - vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); - vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); - vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); - vi.mocked(teamCreateCall).mockResolvedValue({ - team_id: "new-team-1", - team_alias: "Test Team", - models: ["gpt-4"], - organization_id: null, - keys: [], - members_with_roles: [], - spend: 0, - }); - mockUseOrganizations.mockReturnValue({ - data: [{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }], - }); - }); - - it("should pass access_group_ids to teamCreateCall when creating team", async () => { - vi.mocked(teamListCall).mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 1 }); - renderWithQueryClient(); - - const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; - act(() => { - fireEvent.click(createButton); - }); - - await waitFor(() => { - expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); - }); - - const teamNameInput = screen.getByLabelText(/team name/i); - fireEvent.change(teamNameInput, { target: { value: "Test Team" } }); - - const modelsInput = screen.getByTestId("create-team-models-select"); - fireEvent.change(modelsInput, { target: { value: "gpt-4" } }); - - const additionalSettingsAccordion = screen.getByText("Additional Settings"); - fireEvent.click(additionalSettingsAccordion); - - await waitFor(() => { - expect(screen.getByTestId("access-group-selector")).toBeInTheDocument(); - }); - - const accessGroupInput = screen.getByTestId("access-group-selector"); - fireEvent.change(accessGroupInput, { target: { value: "ag-1,ag-2" } }); - - const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); - const createTeamSubmitButton = createTeamSubmitButtons[createTeamSubmitButtons.length - 1]; - fireEvent.click(createTeamSubmitButton); - - await waitFor(() => { - expect(teamCreateCall).toHaveBeenCalledWith( - "test-token", - expect.objectContaining({ - team_alias: "Test Team", - models: ["gpt-4"], - access_group_ids: ["ag-1", "ag-2"], - }), - ); - }); - }); -}); - -describe("OldTeams - models dropdown options", () => { - beforeEach(() => { - vi.clearAllMocks(); - vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - it("should not render all-proxy-models option in models select", async () => { - vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); - - vi.mocked(teamListCall).mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 1 }); - renderWithQueryClient(); - - await waitFor(() => { - expect(fetchAvailableModelsForTeamOrKey).toHaveBeenCalled(); - }); - - const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; - act(() => { - fireEvent.click(createButton); - }); - - await waitFor(() => { - expect(screen.getByLabelText(/models/i)).toBeInTheDocument(); - }); - const allProxyModelsOption = screen.queryByText("All Proxy Models"); - expect(allProxyModelsOption).not.toBeInTheDocument(); - }); -}); - -describe("OldTeams - organization alias display", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - it("should display organization alias instead of organization id", async () => { - const mockOrganizations = [ - { - organization_id: "org-123", - organization_alias: "Test Organization", - budget_id: "budget-1", - metadata: {}, - models: [], - spend: 0, - model_spend: {}, - created_at: new Date().toISOString(), - created_by: "user-1", - updated_at: new Date().toISOString(), - updated_by: "user-1", - litellm_budget_table: null, - teams: null, - users: null, - members: null, - }, - ]; - - mockUseOrganizations.mockReturnValue({ data: mockOrganizations }); - - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("Test Organization")).toBeInTheDocument(); - }); - expect(screen.queryByText("org-123")).not.toBeInTheDocument(); - }); - - it("should display organization id when alias is not found", async () => { - mockUseOrganizations.mockReturnValue({ data: [] }); - - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-unknown", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - await waitFor(() => { - expect(screen.getByText("org-unknown")).toBeInTheDocument(); - }); - }); - - it("should display N/A when organization_id is null", async () => { - mockUseOrganizations.mockReturnValue({ data: [] }); - - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Test Team", - organization_id: null, - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - - await waitFor(() => { - // When organization_id is null, the table shows "—" in the Organization column - expect(screen.getAllByText("—").length).toBeGreaterThan(0); - }); - }); -}); - -describe("OldTeams - Resources column keys badge", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - it("renders keys_count from the v2 payload in the Resources badge", async () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "1", - team_alias: "Team With Keys", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [], - keys_count: 3, - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - const { container } = renderWithQueryClient( - , - ); - - await waitFor(() => { - expect(screen.getByText("Team With Keys")).toBeInTheDocument(); - }); - const cyanTag = container.querySelector(".ant-tag-cyan"); - expect(cyanTag).not.toBeNull(); - expect(cyanTag?.textContent).toContain("3"); - }); - - it("falls back to keys.length when keys_count is absent", async () => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [ - { - team_id: "2", - team_alias: "Legacy Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - keys: [{ token: "t1" }, { token: "t2" }], - members_with_roles: [], - spend: 0, - }, - ], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - const { container } = renderWithQueryClient( - , - ); - - await waitFor(() => { - expect(screen.getByText("Legacy Team")).toBeInTheDocument(); - }); - const cyanTag = container.querySelector(".ant-tag-cyan"); - expect(cyanTag).not.toBeNull(); - expect(cyanTag?.textContent).toContain("2"); - }); -}); - -describe("OldTeams - delete team warning copy", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockUseOrganizations.mockReturnValue({ data: [] }); - }); - - const openDeleteModal = async (team: any) => { - vi.mocked(teamListCall).mockResolvedValue({ - teams: [team], - total: 1, - page: 1, - page_size: 100, - total_pages: 1, - }); - renderWithQueryClient(); - await waitFor(() => { - expect(screen.getByTestId("delete-team-button")).toBeInTheDocument(); - }); - act(() => { - fireEvent.click(screen.getByTestId("delete-team-button")); - }); - expect(screen.getByText("Delete Team?")).toBeInTheDocument(); - }; - - const baseTeam = { - team_id: "1", - team_alias: "Test Team", - organization_id: "org-123", - models: ["gpt-4"], - max_budget: 100, - budget_duration: "1d", - tpm_limit: 1000, - rpm_limit: 1000, - created_at: new Date().toISOString(), - members_with_roles: [], - spend: 0, - }; - - it("warns that the team's models are deleted when the team has keys", async () => { - await openDeleteModal({ ...baseTeam, keys: [], keys_count: 5 }); - - expect(screen.getByText(/Warning: This team has 5 keys associated with it/i)).toHaveTextContent( - /along with any models created for this team/i, - ); - expect(screen.getByText(/Are you sure you want to delete this team/i)).toHaveTextContent( - /any models created for it/i, - ); - }); - - it("still warns about model deletion in the confirmation message when the team has no keys", async () => { - await openDeleteModal({ ...baseTeam, keys: [], keys_count: 0 }); - - expect(screen.queryByText(/Warning: This team has/i)).not.toBeInTheDocument(); - expect(screen.getByText(/Are you sure you want to delete this team/i)).toHaveTextContent( - /any models created for it/i, - ); - }); -}); - -describe("OldTeams - LIT-2530 organization stays optional for proxy admin with a single org", () => { - beforeEach(() => { - vi.clearAllMocks(); - mockTeamInfoView.mockClear(); - vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4"]); - vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); - vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); - vi.mocked(teamListCall).mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 1 }); - vi.mocked(teamCreateCall).mockResolvedValue({ - team_id: "new-team-1", - team_alias: "No Org Team", - models: ["gpt-4"], - organization_id: null, - keys: [], - members_with_roles: [], - spend: 0, - }); - mockUseOrganizations.mockReturnValue({ - data: [{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }], - }); - }); - - it("creates a team with no organization when exactly one organization exists", async () => { - renderWithQueryClient(); - - const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; - act(() => { - fireEvent.click(createButton); - }); - - await waitFor(() => { - expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); - }); - - fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "No Org Team" } }); - fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); - - const submitButtons = screen.getAllByRole("button", { name: /create team/i }); - fireEvent.click(submitButtons[submitButtons.length - 1]); - - await waitFor(() => { - expect(teamCreateCall).toHaveBeenCalledWith( - "test-token", - expect.objectContaining({ team_alias: "No Org Team", organization_id: null }), - ); - }); - }); -}); diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx new file mode 100644 index 00000000000..7065b1a5fb6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -0,0 +1,661 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import React from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key"; +import { fetchMCPAccessGroups, getGuardrailsList, teamCreateCall } from "./networking"; +import Teams from "./Teams"; + +const mockTeamInfoView = vi.fn(); +const mockUseOrganizations = vi.fn(); + +// The teams grid is unit-tested in TeamsPage/TeamsTable.test.tsx. Here we stub it and drive its callbacks +// directly so we can test the Teams shell wiring (delete modal, detail view) without the real DataTable. +let mockTeamsTableProps: any = null; +vi.mock("./TeamsPage/TeamsTable", () => ({ + TeamsTable: (props: any) => { + mockTeamsTableProps = props; + return
; + }, +})); + +vi.mock("./networking", () => ({ + teamCreateCall: vi.fn(), + teamDeleteCall: vi.fn(), + fetchMCPAccessGroups: vi.fn(), + v2TeamListCall: vi.fn(), + getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), + getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), +})); + +// Teams invalidates teamsTableKeys on mutations; the selected team is passed up from the table. +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + teamsTableKeys: { all: ["teamsTable"] }, +})); + +vi.mock("./molecules/notifications_manager", () => ({ + default: { + info: vi.fn(), + success: vi.fn(), + error: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +vi.mock("./key_team_helpers/fetch_available_models_team_key", () => ({ + fetchAvailableModelsForTeamOrKey: vi.fn(), + getModelDisplayName: vi.fn((model: string) => model), + unfurlWildcardModelsInList: vi.fn((teamModels: string[], allModels: string[]) => { + const wildcardDisplayNames: string[] = []; + const expandedModels: string[] = []; + + teamModels.forEach((teamModel) => { + if (teamModel.endsWith("/*")) { + const provider = teamModel.replace("/*", ""); + const matchingModels = allModels.filter((model) => model.startsWith(provider + "/")); + expandedModels.push(...matchingModels); + wildcardDisplayNames.push(teamModel); + } else { + expandedModels.push(teamModel); + } + }); + + return [...wildcardDisplayNames, ...expandedModels].filter((item, index, array) => array.indexOf(item) === index); + }), +})); + +vi.mock("@/components/team/TeamInfo", () => ({ + __esModule: true, + default: (props: any) => { + mockTeamInfoView(props); + return
; + }, +})); + +vi.mock("./ModelSelect/ModelSelect", () => { + const ModelSelect = React.forwardRef(({ value, onChange, dataTestId, id }: any, ref: any) => { + return ( + { + if (onChange) { + const newVal = e.target.value + ? e.target.value + .split(",") + .map((s: string) => s.trim()) + .filter(Boolean) + : []; + onChange(newVal); + } + }} + /> + ); + }); + ModelSelect.displayName = "ModelSelect"; + return { + ModelSelect, + }; +}); + +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + useOrganizations: () => mockUseOrganizations(), +})); + +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: vi.fn().mockReturnValue({ + data: [ + { access_group_id: "ag-1", access_group_name: "Group 1" }, + { access_group_id: "ag-2", access_group_name: "Group 2" }, + ], + isLoading: false, + isError: false, + }), +})); + +vi.mock("./common_components/AccessGroupSelector", () => ({ + default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( + onChange?.(e.target.value ? e.target.value.split(",").map((s) => s.trim()) : [])} + /> + ), +})); + +const baseTableTeam = { + team_id: "1", + team_alias: "Test Team", + organization_id: "org-123", + models: ["gpt-4"], + max_budget: 100, + budget_duration: "1d", + tpm_limit: 1000, + rpm_limit: 1000, + created_at: new Date().toISOString(), + keys: [], + members_with_roles: [], + spend: 0, +}; + +const createQueryClient = () => { + return new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); +}; + +const renderWithQueryClient = (component: React.ReactElement) => { + const queryClient = createQueryClient(); + return render({component}); +}; + +// Re-establish safe defaults before every test (clearAllMocks keeps return values, so restore them here). +beforeEach(() => { + mockTeamsTableProps = null; +}); + +describe("Teams - handleCreate organization handling", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + mockTeamsTableProps = null; + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue([]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + mockUseOrganizations.mockReturnValue({ data: null }); + }); + + it("should not include organization_id when it's an empty string", async () => { + const formValues: Record = { + team_alias: "Test Team", + organization_id: "", // Empty string + models: [], + }; + + // Simulate the handleCreate logic + let organizationId = formValues?.organization_id || null; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues.organization_id).toBeNull(); + expect(formValues.organization_id).not.toBe(""); + }); + + it("should set organization_id to null when it's not a string type", async () => { + const formValues: Record = { + team_alias: "Test Team", + organization_id: undefined, + models: [], + }; + + let organizationId = formValues?.organization_id || null; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues.organization_id).toBeNull(); + }); + + it("should trim and keep valid organization_id string", async () => { + const formValues: Record = { + team_alias: "Test Team", + organization_id: " org-123 ", + models: [], + }; + + let organizationId = formValues?.organization_id || null; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues.organization_id).toBe("org-123"); + }); + + it("should keep valid organization_id without modification", async () => { + const formValues: Record = { + team_alias: "Test Team", + organization_id: "f874bb43-b898-4813-beca-4054d224eafc", + models: [], + }; + + let organizationId = formValues?.organization_id || null; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues.organization_id).toBe("f874bb43-b898-4813-beca-4054d224eafc"); + }); + + it("should not send organization_id field when converting empty string to null", async () => { + const formValues: Record = { + team_alias: "Test Team", + organization_id: "", + models: ["gpt-4"], + max_budget: 100, + }; + + let organizationId = formValues?.organization_id || null; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues).toEqual({ + team_alias: "Test Team", + organization_id: null, + models: ["gpt-4"], + max_budget: 100, + }); + expect(formValues.organization_id).not.toBe(""); + expect(formValues.organization_id).toBeNull(); + }); + + it("should handle when currentOrg is used as fallback", async () => { + const currentOrg = { + organization_id: "fallback-org-id", + organization_alias: "Fallback Org", + models: [], + members: [], + }; + + const formValues: Record = { + team_alias: "Test Team", + models: [], + }; + + let organizationId = formValues?.organization_id || currentOrg?.organization_id; + if (organizationId === "" || typeof organizationId !== "string") { + formValues.organization_id = null; + } else { + formValues.organization_id = organizationId.trim(); + } + + expect(formValues.organization_id).toBe("fallback-org-id"); + }); + + it("opens the delete modal when the table's delete action fires", async () => { + mockUseOrganizations.mockReturnValue({ data: [] }); + renderWithQueryClient(); + + await waitFor(() => expect(mockTeamsTableProps).not.toBeNull()); + await act(async () => { + mockTeamsTableProps.onDeleteTeam(baseTableTeam); + }); + + expect(screen.getByText("Delete Team?")).toBeInTheDocument(); + }); +}); + +describe("Teams - helper functions", () => { + describe("getAdminOrganizations", () => { + it("should return all organizations for Admin role", () => { + const organizations = [ + { organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }, + { organization_id: "org-2", organization_alias: "Org 2", models: [], members: [] }, + ]; + + const userRole = "Admin"; + const result = userRole === "Admin" ? organizations : []; + + expect(result).toEqual(organizations); + expect(result.length).toBe(2); + }); + + it("should return only org_admin organizations for Org Admin role", () => { + const userID = "user-123"; + const organizations = [ + { + organization_id: "org-1", + organization_alias: "Org 1", + models: [], + members: [{ user_id: "user-123", user_role: "org_admin" }], + }, + { + organization_id: "org-2", + organization_alias: "Org 2", + models: [], + members: [{ user_id: "user-456", user_role: "org_admin" }], + }, + { + organization_id: "org-3", + organization_alias: "Org 3", + models: [], + members: [{ user_id: "user-123", user_role: "member" }], + }, + ]; + + const result = organizations.filter((org) => + org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), + ); + + expect(result.length).toBe(1); + expect(result[0].organization_id).toBe("org-1"); + }); + + it("should return empty array when user is not admin of any organization", () => { + const userID = "user-999"; + const organizations = [ + { + organization_id: "org-1", + organization_alias: "Org 1", + models: [], + members: [{ user_id: "user-123", user_role: "org_admin" }], + }, + ]; + + const result = organizations.filter((org) => + org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), + ); + + expect(result.length).toBe(0); + }); + }); + + describe("canCreateOrManageTeams", () => { + it("should return true for Admin role", () => { + const userRole = "Admin"; + expect(userRole === "Admin").toBe(true); + }); + + it("should return true for org_admin in any organization", () => { + const userID = "user-123"; + const organizations = [ + { + organization_id: "org-1", + organization_alias: "Org 1", + models: [], + members: [{ user_id: "user-123", user_role: "org_admin" }], + }, + ]; + + const result = organizations.some((org) => + org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), + ); + + expect(result).toBe(true); + }); + + it("should return false when user has no admin permissions", () => { + const userID = "user-123"; + const userRole: string = "User"; + const organizations = [ + { + organization_id: "org-1", + organization_alias: "Org 1", + models: [], + members: [{ user_id: "user-123", user_role: "member" }], + }, + ]; + + const isAdmin = userRole === "Admin"; + const isOrgAdmin = organizations.some((org) => + org.members?.some((member) => member.user_id === userID && member.user_role === "org_admin"), + ); + + expect(isAdmin || isOrgAdmin).toBe(false); + }); + }); +}); + +describe("Teams - premium props", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue([]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + mockUseOrganizations.mockReturnValue({ data: [] }); + }); + + it("passes premiumUser flag to TeamInfoView when a team is opened", async () => { + const premiumTeam = { ...baseTableTeam, team_id: "team-123456789", team_alias: "Premium Team" }; + renderWithQueryClient(); + + await waitFor(() => expect(mockTeamsTableProps).not.toBeNull()); + act(() => mockTeamsTableProps.onSelectTeam(premiumTeam)); + + await waitFor(() => expect(mockTeamInfoView).toHaveBeenCalled()); + expect(mockTeamInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ premiumUser: true })); + }); +}); + +describe("Teams - Create Team CTA is grouped with the tabs on the left", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseOrganizations.mockReturnValue({ data: [] }); + }); + + it("renders the Create Team button inside the tab bar, ahead of the tabs", () => { + const { container } = renderWithQueryClient(); + + const createButton = screen.getByTestId("create-team-button"); + const tabNav = container.querySelector(".ant-tabs-nav"); + + // The CTA lives in the tab bar's left slot, not the standalone page header. + expect(tabNav).not.toBeNull(); + expect(tabNav!.contains(createButton)).toBe(true); + + // It reads as the left end of the cluster: it precedes the first tab in DOM order. + const firstTab = screen.getByRole("tab", { name: "Your Teams" }); + expect(createButton.compareDocumentPosition(firstTab) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + }); + + it("omits the Create Team CTA for a role that cannot manage teams", () => { + renderWithQueryClient(); + expect(screen.queryByTestId("create-team-button")).not.toBeInTheDocument(); + }); +}); + +describe("Teams - Default Team Settings tab visibility", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseOrganizations.mockReturnValue({ data: [] }); + }); + + it("should show Default Team Settings tab for Admin role", () => { + renderWithQueryClient(); + expect(screen.getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument(); + }); + + it("should show Default Team Settings tab for proxy_admin role", () => { + renderWithQueryClient(); + expect(screen.getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument(); + }); + + it("should not show Default Team Settings tab for proxy_admin_viewer role", () => { + renderWithQueryClient(); + expect(screen.queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument(); + }); + + it("should not show Default Team Settings tab for Admin Viewer role", () => { + renderWithQueryClient(); + expect(screen.queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument(); + }); +}); + +describe("Teams - access_group_ids in team create", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(teamCreateCall).mockResolvedValue({ + team_id: "new-team-1", + team_alias: "Test Team", + models: ["gpt-4"], + organization_id: null, + keys: [], + members_with_roles: [], + spend: 0, + }); + mockUseOrganizations.mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }], + }); + }); + + it("should pass access_group_ids to teamCreateCall when creating team", async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Test Team" } }); + fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); + + fireEvent.click(screen.getByText("Additional Settings")); + + await waitFor(() => { + expect(screen.getByTestId("access-group-selector")).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByTestId("access-group-selector"), { target: { value: "ag-1,ag-2" } }); + + const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + team_alias: "Test Team", + models: ["gpt-4"], + access_group_ids: ["ag-1", "ag-2"], + }), + ); + }); + }); +}); + +describe("Teams - models dropdown options", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); + mockUseOrganizations.mockReturnValue({ data: [] }); + }); + + it("should not render all-proxy-models option in models select", async () => { + renderWithQueryClient(); + + await waitFor(() => { + expect(fetchAvailableModelsForTeamOrKey).toHaveBeenCalled(); + }); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/models/i)).toBeInTheDocument(); + }); + expect(screen.queryByText("All Proxy Models")).not.toBeInTheDocument(); + }); +}); + +describe("Teams - delete team warning copy", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseOrganizations.mockReturnValue({ data: [] }); + }); + + const openDeleteModal = async (team: any) => { + renderWithQueryClient(); + await waitFor(() => expect(mockTeamsTableProps).not.toBeNull()); + await act(async () => { + mockTeamsTableProps.onDeleteTeam(team); + }); + expect(screen.getByText("Delete Team?")).toBeInTheDocument(); + }; + + it("warns that the team's models are deleted when the team has keys", async () => { + await openDeleteModal({ ...baseTableTeam, keys: [], keys_count: 5 }); + + expect(screen.getByText(/Warning: This team has 5 keys associated with it/i)).toHaveTextContent( + /along with any models created for this team/i, + ); + expect(screen.getByText(/Are you sure you want to delete this team/i)).toHaveTextContent( + /any models created for it/i, + ); + }); + + it("still warns about model deletion in the confirmation message when the team has no keys", async () => { + await openDeleteModal({ ...baseTableTeam, keys: [], keys_count: 0 }); + + expect(screen.queryByText(/Warning: This team has/i)).not.toBeInTheDocument(); + expect(screen.getByText(/Are you sure you want to delete this team/i)).toHaveTextContent( + /any models created for it/i, + ); + }); +}); + +describe("Teams - LIT-2530 organization stays optional for proxy admin with a single org", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(teamCreateCall).mockResolvedValue({ + team_id: "new-team-1", + team_alias: "No Org Team", + models: ["gpt-4"], + organization_id: null, + keys: [], + members_with_roles: [], + spend: 0, + }); + mockUseOrganizations.mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }], + }); + }); + + it("creates a team with no organization when exactly one organization exists", async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "No Org Team" } }); + fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); + + const submitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(submitButtons[submitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ team_alias: "No Org Team", organization_id: null }), + ); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/OldTeams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx similarity index 71% rename from ui/litellm-dashboard/src/components/OldTeams.tsx rename to ui/litellm-dashboard/src/components/Teams.tsx index e83d1acf4e5..2627f2c1d7f 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -3,38 +3,16 @@ import AvailableTeamsPanel from "@/components/team/available_teams"; import TeamInfoView from "@/components/team/TeamInfo"; import TeamSSOSettings from "@/components/TeamSSOSettings"; import { isProxyAdminRole } from "@/utils/roles"; -import { InfoCircleOutlined, PlusOutlined, TeamOutlined, ReloadOutlined } from "@ant-design/icons"; +import { InfoCircleOutlined } from "@ant-design/icons"; import { Accordion, AccordionBody, AccordionHeader, TextInput } from "@tremor/react"; -import { - Button, - Card, - Flex, - Form, - Input, - Layout, - Modal, - Pagination, - Progress, - Select, - Space, - Switch, - Table, - Tabs, - Tag, - theme, - Tooltip, - Typography, - message, -} from "antd"; -import type { ColumnsType } from "antd/es/table"; -import type { SorterResult } from "antd/es/table/interface"; -import { KeyIcon, LayersIcon, SearchIcon, UsersIcon } from "lucide-react"; -import React, { useEffect, useMemo, useRef, useState } from "react"; -import { AntDLoadingSpinner } from "@/components/ui/AntDLoadingSpinner"; -import { DateCell, IdCell } from "@/components/shared/table_cells"; -import OrganizationDropdown from "./common_components/OrganizationDropdown"; -import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; -import { teamListCall as v2TeamListCall, type TeamsResponse } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { Button, Form, Input, Layout, Modal, Select, Switch, Tabs, theme, Tooltip, Typography } from "antd"; +import { Plus, Users } from "lucide-react"; +import React, { useEffect, useState } from "react"; +import { useQueryClient } from "@tanstack/react-query"; +import { PageHeader } from "@/components/shared/PageHeader"; +import { Button as UIButton } from "@/components/ui/button"; +import { teamsTableKeys } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { TeamsTable } from "./TeamsPage/TeamsTable"; import AccessGroupSelector from "./common_components/AccessGroupSelector"; import PassThroughRoutesSelector from "./common_components/PassThroughRoutesSelector"; import AgentSelector from "./agent_management/AgentSelector"; @@ -45,7 +23,7 @@ import { fetchAvailableModelsForTeamOrKey, unfurlWildcardModelsInList, } from "./key_team_helpers/fetch_available_models_team_key"; -import type { KeyResponse, Team } from "./key_team_helpers/key_list"; +import type { Team } from "./key_team_helpers/key_list"; import MCPServerSelector from "./mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "./mcp_server_management/MCPToolPermissions"; import NotificationsManager from "./molecules/notifications_manager"; @@ -61,13 +39,6 @@ interface TeamProps { premiumUser?: boolean; } -interface FilterState { - search: string; - organization_id: string; - sort_by: string; - sort_order: "asc" | "desc"; -} - interface EditTeamModalProps { visible: boolean; onCancel: () => void; @@ -75,21 +46,10 @@ interface EditTeamModalProps { onSubmit: (data: FormData) => void; // Assuming FormData is the type of data to be submitted } -import { updateExistingKeys } from "@/utils/dataUtils"; import DeleteResourceModal from "./common_components/DeleteResourceModal"; -import { Member, teamCreateCall } from "./networking"; +import { teamCreateCall } from "./networking"; import { ModelSelect } from "./ModelSelect/ModelSelect"; -interface TeamInfo { - members_with_roles: Member[]; -} - -interface PerTeamInfo { - keys: KeyResponse[]; - keys_count: number; - team_info: TeamInfo; -} - const getOrganizationModels = (organization: Organization | null, userModels: string[]) => { let tempModelsToPick = []; @@ -164,70 +124,17 @@ const getOrganizationAlias = ( const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser = false }) => { const { data: organizationsData } = useOrganizations(); const organizations = organizationsData ?? null; - const [teams, setTeams] = useState(null); - const [isLoading, setIsLoading] = useState(true); - const [fetchError, setFetchError] = useState(null); - const [currentPage, setCurrentPage] = useState(1); - const [pageSize, setPageSize] = useState(10); - const [totalTeams, setTotalTeams] = useState(0); - const [currentOrg, setCurrentOrg] = useState(null); + const queryClient = useQueryClient(); + const refreshTeams = () => queryClient.invalidateQueries({ queryKey: teamsTableKeys.all }); + const [currentOrg] = useState(null); const [currentOrgForCreateTeam, setCurrentOrgForCreateTeam] = useState(null); - const [filters, setFilters] = useState({ - search: "", - organization_id: "", - sort_by: "created_at", - sort_order: "desc", - }); - const searchDebounceRef = useRef | null>(null); - const [isSearching, setIsSearching] = useState(false); - - const fetchTeamsV2 = async ( - opts: { - page?: number; - size?: number; - sortBy?: string; - sortOrder?: string; - organizationID?: string; - search?: string; - } = {}, - ) => { - if (!accessToken) return; - const page = opts.page ?? currentPage; - const size = opts.size ?? pageSize; - const sortBy = opts.sortBy ?? filters.sort_by; - const sortOrder = opts.sortOrder ?? filters.sort_order; - const organizationID = opts.organizationID ?? filters.organization_id; - const search = opts.search ?? filters.search; - - setIsLoading(true); - setFetchError(null); - try { - const response: TeamsResponse = await v2TeamListCall(accessToken, page, size, { - organizationID: organizationID || null, - search: search || null, - userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null, - sortBy: sortBy || null, - sortOrder: sortOrder || null, - }); - setTeams(response.teams ?? []); - setTotalTeams(response.total ?? 0); - } catch (err: any) { - setFetchError(err?.message || "Failed to fetch teams"); - } finally { - setIsLoading(false); - } - }; - - useEffect(() => { - fetchTeamsV2(); - }, [accessToken]); const [form] = Form.useForm(); const [memberForm] = Form.useForm(); const [value, setValue] = useState(""); const [editModalVisible, setEditModalVisible] = useState(false); - const [selectedTeam, setSelectedTeam] = useState(null); + const [selectedTeam, setSelectedTeam] = useState(null); const [selectedTeamId, setSelectedTeamId] = useState(null); const [editTeam, setEditTeam] = useState(false); @@ -238,7 +145,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [teamToDelete, setTeamToDelete] = useState(null); const [modelsToPick, setModelsToPick] = useState([]); - const [perTeamInfo, setPerTeamInfo] = useState>({}); const [isTeamDeleting, setIsTeamDeleting] = useState(false); // Add this state near the other useState declarations const [guardrailsList, setGuardrailsList] = useState([]); @@ -325,30 +231,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser fetchMcpAccessGroups(); }, [accessToken]); - useEffect(() => { - const fetchTeamInfo = () => { - if (!teams) return; - - const newPerTeamInfo = teams.reduce( - (acc, team) => { - acc[team.team_id] = { - keys: team.keys || [], - keys_count: team.keys_count ?? team.keys?.length ?? 0, - team_info: { - members_with_roles: team.members_with_roles || [], - }, - }; - return acc; - }, - {} as Record, - ); - - setPerTeamInfo(newPerTeamInfo); - }; - - fetchTeamInfo(); - }, [teams]); - const handleOk = () => { setIsTeamModalVisible(false); form.resetFields(); @@ -386,14 +268,14 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser }; const confirmDelete = async () => { - if (teamToDelete == null || teams == null || accessToken == null) { + if (teamToDelete == null || accessToken == null) { return; } try { setIsTeamDeleting(true); await teamDeleteCall(accessToken, teamToDelete.team_id); - await fetchTeamsV2(); + await refreshTeams(); NotificationsManager.success("Team deleted successfully"); } catch (error) { NotificationsManager.fromBackend("Error deleting the team: " + error); @@ -425,13 +307,11 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser }; fetchUserModels(); - }, [accessToken, userID, userRole, teams]); + }, [accessToken, userID, userRole]); const handleCreate = async (formValues: Record) => { try { if (accessToken != null) { - const newTeamAlias = formValues?.team_alias; - const existingTeamAliases = teams?.map((t) => t.team_alias) ?? []; let organizationId = formValues?.organization_id || currentOrg?.organization_id; if (organizationId === "" || typeof organizationId !== "string") { formValues.organization_id = null; @@ -439,11 +319,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser formValues.organization_id = organizationId.trim(); } - // Remove guardrails from top level since it's now in metadata - if (existingTeamAliases.includes(newTeamAlias)) { - throw new Error(`Team alias ${newTeamAlias} already exists, please pick another alias`); - } - NotificationsManager.info("Creating Team"); // Handle logging settings in metadata @@ -565,10 +440,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser await teamCreateCall(accessToken, formValues); NotificationsManager.success("Team created"); - await fetchTeamsV2({ - page: currentPage, - size: pageSize, - }); + await refreshTeams(); form.resetFields(); setLoggingSettings([]); setModelAliases({}); @@ -595,352 +467,31 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser return false; }; - const handleSearchChange = (value: string) => { - if (searchDebounceRef.current) clearTimeout(searchDebounceRef.current); - setIsSearching(true); - searchDebounceRef.current = setTimeout(async () => { - try { - setFilters((prev) => ({ ...prev, search: value })); - setCurrentPage(1); - await fetchTeamsV2({ page: 1, search: value }); - } finally { - setIsSearching(false); - } - }, 300); - }; - - const handleFilterChange = async (key: keyof FilterState, value: string) => { - const newFilters = { ...filters, [key]: value }; - setFilters(newFilters); - setCurrentPage(1); - if (!accessToken) return; - try { - const response: TeamsResponse = await v2TeamListCall(accessToken, 1, pageSize, { - organizationID: newFilters.organization_id || null, - search: newFilters.search || null, - userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null, - sortBy: newFilters.sort_by || null, - sortOrder: newFilters.sort_order || null, - }); - setTeams(response.teams ?? []); - setTotalTeams(response.total ?? 0); - } catch (error) { - console.error("Error fetching teams:", error); - } - }; - - const handleFilterReset = () => { - if (searchDebounceRef.current) clearTimeout(searchDebounceRef.current); - setIsSearching(false); - const resetFilters: FilterState = { - search: "", - organization_id: "", - sort_by: "created_at", - sort_order: "desc", - }; - setFilters(resetFilters); - setCurrentPage(1); - fetchTeamsV2({ page: 1, organizationID: "", search: "", sortBy: "created_at", sortOrder: "desc" }); - }; - const { token } = theme.useToken(); - const { Title, Text } = Typography; + const { Text } = Typography; const { Content } = Layout; - const handleRetry = () => { - fetchTeamsV2(); - }; - - const handleTableSort = ( - _pagination: unknown, - _filters: unknown, - sorter: SorterResult | SorterResult[], - ) => { - const s = Array.isArray(sorter) ? sorter[0] : sorter; - const sortBy = s.order ? (s.columnKey as string) : "created_at"; - const sortOrder = s.order === "ascend" ? "asc" : s.order === "descend" ? "desc" : "desc"; - setFilters((prev) => ({ ...prev, sort_by: sortBy, sort_order: sortOrder })); - fetchTeamsV2({ sortBy, sortOrder }); - }; - - const teamColumns: ColumnsType = useMemo( - () => [ - { - title: "Team ID", - dataIndex: "team_id", - key: "team_id", - width: 170, - ellipsis: true, - render: (id: string) => ( - setSelectedTeamId(teamId)} dataTestId="team-id-cell" /> - ), - }, - { - title: "Team Alias", - dataIndex: "team_alias", - key: "team_alias", - ellipsis: true, - sorter: true, - render: (alias: string | undefined) => ( - - {alias || ( - - — - - )} - - ), - }, - { - title: "Organization", - key: "organization", - width: 160, - ellipsis: true, - render: (_: unknown, record: Team) => { - const orgAlias = getOrganizationAlias(record.organization_id, organizations); - return record.organization_id ? ( - - {orgAlias} - - ) : ( - — - ); - }, - }, - { - title: "Resources", - key: "resources", - width: 240, - render: (_: unknown, record: Team) => { - const memberCount = perTeamInfo?.[record.team_id]?.team_info?.members_with_roles?.length ?? 0; - const modelCount = record.models?.length ?? 0; - const keyCount = perTeamInfo?.[record.team_id]?.keys_count ?? 0; - return ( - - - - - - {memberCount} - - - - - - - - {modelCount} - - - - - - - - {keyCount} - - - - - ); - }, - }, - { - title: "Spend / Budget", - key: "spend", - width: 200, - sorter: true, - render: (_: unknown, record: Team) => { - const spendVal = record.spend ?? 0; - const budgetVal = record.max_budget; - const spendStr = `$${spendVal.toLocaleString(undefined, { minimumFractionDigits: 2, maximumFractionDigits: 2 })}`; - const budgetStr = - budgetVal != null - ? `$${budgetVal.toLocaleString(undefined, { minimumFractionDigits: 2, maximumFractionDigits: 2 })}` - : "Unlimited"; - const percent = budgetVal != null && budgetVal > 0 ? Math.min((spendVal / budgetVal) * 100, 100) : null; - return ( - - - {spendStr} - - {" / "} - {budgetStr} - - - {percent != null && ( - = 90 ? "#ff4d4f" : percent >= 70 ? "#faad14" : "#1677ff"} - style={{ marginBottom: 0 }} - /> - )} - - ); - }, - }, - { - title: "Created", - dataIndex: "created_at", - key: "created_at", - width: 130, - ellipsis: true, - sorter: true, - render: (date: string | undefined) => , - }, - { - title: "Actions", - key: "actions", - width: 120, - align: "right" as const, - render: (_: unknown, record: Team) => ( - - { - navigator.clipboard - .writeText(record.team_id) - .then(() => message.success("Team ID copied")) - .catch(() => message.error("Failed to copy")); - }} - /> - {userRole === "Admin" && ( - <> - { - setSelectedTeamId(record.team_id); - setEditTeam(true); - }} - /> - handleDelete(record)} - /> - - )} - - ), - }, - ], - [userRole, perTeamInfo, organizations], - ); - - const displayTeams = useMemo(() => teams ?? [], [teams]); - - const renderTeamsContent = () => { - if (isLoading) { - return ( - - - - ); - } - - if (fetchError) { - return ( - - - Failed to load teams - - - {fetchError} - - - - ); - } - - return ( - - columns={teamColumns} - dataSource={displayTeams} - rowKey="team_id" - pagination={false} - onChange={handleTableSort} - locale={{ - emptyText: ( -
- -
- No teams yet -
-
- - Create your first team to organize members and manage access to models. - -
- {canCreateOrManageTeams(userRole, userID, organizations) && ( - - )} -
- ), - }} - scroll={{ x: 1000 }} - size="middle" - /> - ); - }; - const tabItems = [ { key: "your-teams", label: "Your Teams", children: ( <> - - - - } - suffix={isSearching ? : null} - placeholder="Search teams by name or ID..." - onChange={(e) => handleSearchChange(e.target.value)} - allowClear - style={{ maxWidth: 400 }} - /> - handleFilterChange("organization_id", value || "")} - loading={isLoading} - /> - - { - setCurrentPage(page); - setPageSize(size); - fetchTeamsV2({ page, size }); - }} - size="small" - showTotal={(total) => `${total} teams`} - showSizeChanger - pageSizeOptions={["10", "20", "50"]} - /> - - - {renderTeamsContent()} - + { + setSelectedTeam(team); + setSelectedTeamId(team.team_id); + setEditTeam(false); + }} + onEditTeam={(team) => { + setSelectedTeam(team); + setSelectedTeamId(team.team_id); + setEditTeam(true); + }} + onDeleteTeam={handleDelete} + /> = ({ accessToken, userID, userRole, premiumUser {selectedTeamId ? ( { - setTeams((teams) => { - if (teams == null) { - return teams; - } - return teams.map((team) => { - if (data.team_id === team.team_id) { - return updateExistingKeys(team, data); - } - return team; - }); - }); - fetchTeamsV2(); + onUpdate={() => { + refreshTeams(); }} onClose={() => { + setSelectedTeam(null); setSelectedTeamId(null); setEditTeam(false); }} accessToken={accessToken} - is_team_admin={is_team_admin(teams?.find((team) => team.team_id === selectedTeamId))} + is_team_admin={is_team_admin(selectedTeam)} is_proxy_admin={userRole == "Admin"} userModels={userModels} editTeam={editTeam} @@ -1018,27 +559,28 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser /> ) : ( <> - - - - <TeamOutlined style={{ marginRight: 8 }} /> - Teams - - Manage teams, members, and their access to models and budgets - - {canCreateOrManageTeams(userRole, userID, organizations) && ( - - )} - +
+ } + title="Teams" + subtitle="Manage teams, members, and their access to models and budgets" + /> +
- + + setIsTeamModalVisible(true)} data-testid="create-team-button"> + + Create Team + +
+
+ ) : undefined, + }} + /> )} diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx new file mode 100644 index 00000000000..9469be12128 --- /dev/null +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx @@ -0,0 +1,330 @@ +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, MockedFunction, vi } from "vitest"; + +import { renderWithProviders } from "../../../tests/test-utils"; +import { Team } from "../key_team_helpers/key_list"; +import { TeamsResponse, useTeamsTable } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { TeamsTable } from "./TeamsTable"; + +// Resolve debounced values synchronously so an applied filter lands in the useTeamsTable query within the test tick. +vi.mock("@tanstack/react-pacer/debouncer", async () => { + const React = await vi.importActual("react"); + return { + useDebouncedValue: (value: unknown) => [value, { cancel: vi.fn(), flush: vi.fn() }], + useDebouncedState: (initial: unknown) => { + const [value, setValue] = React.useState(initial); + return [value, setValue, { cancel: vi.fn(), flush: vi.fn() }]; + }, + }; +}); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(() => ({ + accessToken: "test-token", + userId: "test-user", + userRole: "Admin", + premiumUser: true, + token: "test-token", + })), +})); + +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useTeamsTable: vi.fn(), + teamsTableKeys: { all: ["teamsTable"] }, +})); + +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + useOrganizations: vi.fn().mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "Test Organization" }], + }), +})); + +const mockTeam: Team = { + team_id: "team-1", + team_alias: "Acme Team", + models: ["gpt-4", "gpt-3.5-turbo", "claude-3", "claude-3-5-sonnet"], + max_budget: 100, + budget_duration: "1mo", + tpm_limit: 5000, + rpm_limit: 500, + organization_id: "org-1", + created_at: "2024-10-01T10:00:00Z", + updated_at: "2024-11-01T10:00:00Z", + keys: [], + keys_count: 3, + members_with_roles: [ + { user_id: "u1", user_email: "a@x.com", role: "admin" }, + { user_id: "u2", user_email: "b@x.com", role: "user" }, + ] as unknown as Team["members_with_roles"], + spend: 42.5, +}; + +const mockUseTeamsTable = useTeamsTable as MockedFunction; + +const teamsResult = (teams: Team[], data: Partial = {}, extra: Record = {}) => + ({ + data: { + teams, + total: teams.length, + page: 1, + page_size: 50, + total_pages: 1, + ...data, + } as TeamsResponse, + isPending: false, + isFetching: false, + isError: false, + refetch: vi.fn(), + ...extra, + }) as any; + +const noop = () => {}; + +const renderTable = (props: Partial> = {}) => + renderWithProviders( + , + ); + +const openFilters = () => fireEvent.click(screen.getByRole("button", { name: "Filters" })); +const lastOptions = () => mockUseTeamsTable.mock.calls[mockUseTeamsTable.mock.calls.length - 1][2] ?? {}; + +beforeEach(() => { + vi.clearAllMocks(); + mockUseTeamsTable.mockReturnValue(teamsResult([mockTeam])); +}); + +it("renders a team row with alias, organization, and spend/budget", async () => { + renderTable(); + + await waitFor(() => { + expect(screen.getByText("Acme Team")).toBeInTheDocument(); + expect(screen.getByText("Test Organization")).toBeInTheDocument(); + expect(screen.getByText("$42.5000")).toBeInTheDocument(); + expect(screen.getByText("of $100")).toBeInTheDocument(); + }); +}); + +it("renders the Resources cell with member, model, and key counts", () => { + renderTable(); + + expect(screen.getByTitle("2 members")).toBeInTheDocument(); + expect(screen.getByTitle("4 models")).toBeInTheDocument(); + expect(screen.getByTitle("3 keys")).toBeInTheDocument(); +}); + +it("shows 'No teams found' when the list is empty", () => { + mockUseTeamsTable.mockReturnValue(teamsResult([])); + renderTable(); + expect(screen.getByText("No teams found")).toBeInTheDocument(); +}); + +it("shows a loading state on initial load and hides the data", () => { + mockUseTeamsTable.mockReturnValue(teamsResult([], {}, { data: null, isPending: true, isFetching: true })); + renderTable(); + + expect(screen.getByText("Loading teams...")).toBeInTheDocument(); + expect(screen.queryByText("Acme Team")).not.toBeInTheDocument(); +}); + +describe("sort contract – only backend-sortable columns are sortable", () => { + it("requests the default created_at descending sort on first render", () => { + renderTable(); + expect(lastOptions()).toMatchObject({ sortBy: "created_at", sortOrder: "desc" }); + }); + + it("sorts by the backend team_alias field (not the label) when the Team header is clicked", async () => { + renderTable(); + fireEvent.click(screen.getByText("Team").closest("button") as HTMLElement); + + await waitFor(() => { + expect(mockUseTeamsTable).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: "team_alias" })); + }); + }); + + it("does not make Spend / Budget sortable (the backend rejects sort_by=spend)", () => { + renderTable(); + expect(screen.getByText("Spend / Budget").closest("button")).toBeNull(); + // Team and Created are the only sortable headers. + expect(screen.getByText("Team").closest("button")).not.toBeNull(); + expect(screen.getByText("Created").closest("button")).not.toBeNull(); + }); +}); + +describe("server-side filtering maps controls to the right query params", () => { + it("sends no filter params when nothing is applied", () => { + renderTable(); + expect(lastOptions()).toMatchObject({ organizationID: undefined, team_alias: undefined, teamID: undefined }); + }); + + it("threads an applied Team alias filter into the query", async () => { + renderTable(); + openFilters(); + + fireEvent.change(await screen.findByPlaceholderText(/Enter team alias/), { target: { value: "acme" } }); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => { + expect(mockUseTeamsTable).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ team_alias: "acme" })); + }); + }); + + it("threads an applied Team ID filter into the query", async () => { + renderTable(); + openFilters(); + + fireEvent.change(await screen.findByPlaceholderText(/Enter team ID/), { target: { value: "team-xyz" } }); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => { + expect(mockUseTeamsTable).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ teamID: "team-xyz" })); + }); + }); + + it("threads the toolbar search into the search param", async () => { + renderTable(); + fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "platform" } }); + + await waitFor(() => { + expect(mockUseTeamsTable).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: "platform" })); + }); + }); +}); + +describe("non-admin scoping", () => { + it("scopes the list to the current user when the role is not an admin role", () => { + renderTable({ userRole: "Internal User", userID: "user-42" }); + expect(lastOptions()).toMatchObject({ userID: "user-42" }); + }); + + it("does not scope by user for the Admin role", () => { + renderTable({ userRole: "Admin", userID: "admin-1" }); + expect(lastOptions()).toMatchObject({ userID: undefined }); + }); +}); + +describe("row actions", () => { + it("opens the team detail when the team cell is clicked", () => { + const onSelectTeam = vi.fn(); + renderTable({ onSelectTeam }); + + fireEvent.click(screen.getByText("Acme Team")); + expect(onSelectTeam).toHaveBeenCalledWith(expect.objectContaining({ team_id: "team-1" })); + }); + + it("offers Edit and Delete to an Admin and wires them to the callbacks", async () => { + const onEditTeam = vi.fn(); + const onDeleteTeam = vi.fn(); + const user = userEvent.setup(); + renderTable({ userRole: "Admin", onEditTeam, onDeleteTeam }); + + await user.click(screen.getByTestId("team-actions-team-1")); + + await user.click(await screen.findByText("Edit team")); + expect(onEditTeam).toHaveBeenCalledWith(expect.objectContaining({ team_id: "team-1" })); + + await user.click(screen.getByTestId("team-actions-team-1")); + await user.click(await screen.findByText("Delete team")); + expect(onDeleteTeam).toHaveBeenCalledWith(expect.objectContaining({ team_id: "team-1" })); + }); + + it("hides Edit and Delete from a non-admin, leaving only Copy team ID", async () => { + const user = userEvent.setup(); + renderTable({ userRole: "Internal User" }); + + await user.click(screen.getByTestId("team-actions-team-1")); + + expect(await screen.findByText("Copy team ID")).toBeInTheDocument(); + expect(screen.queryByText("Edit team")).not.toBeInTheDocument(); + expect(screen.queryByText("Delete team")).not.toBeInTheDocument(); + }); +}); + +describe("pagination total comes from the query response", () => { + it("shows the total count and page count from the response", async () => { + mockUseTeamsTable.mockReturnValue(teamsResult([mockTeam], { total: 137, total_pages: 3 })); + renderTable(); + + await waitFor(() => { + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-50 of 137"); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 3"); + }); + }); +}); + +describe("refresh control", () => { + it("calls refetch when clicked", () => { + const refetch = vi.fn(); + mockUseTeamsTable.mockReturnValue(teamsResult([mockTeam], {}, { refetch })); + renderTable(); + + fireEvent.click(screen.getByTestId("datatable-refresh")); + expect(refetch).toHaveBeenCalledTimes(1); + }); + + it("keeps rows visible but disables refresh while a background fetch is in flight", () => { + mockUseTeamsTable.mockReturnValue(teamsResult([mockTeam], {}, { isFetching: true })); + renderTable(); + + expect(screen.getByTestId("datatable-refresh")).toBeDisabled(); + expect(screen.getByText("Acme Team")).toBeInTheDocument(); + }); +}); + +describe("column rendering details", () => { + it("shows the organization alias when the id resolves, and the raw id when it does not", async () => { + mockUseTeamsTable.mockReturnValue( + teamsResult([ + { ...mockTeam, team_id: "a", organization_id: "org-1" }, + { ...mockTeam, team_id: "b", team_alias: "Orphan Team", organization_id: "org-unknown" }, + ]), + ); + renderTable(); + + await waitFor(() => { + expect(screen.getByText("Test Organization")).toBeInTheDocument(); + expect(screen.getByText("org-unknown")).toBeInTheDocument(); + }); + }); + + it("renders an em dash for a team with no organization", () => { + mockUseTeamsTable.mockReturnValue(teamsResult([{ ...mockTeam, organization_id: null as unknown as string }])); + renderTable(); + expect(screen.getByText("—")).toBeInTheDocument(); + }); + + it("falls back to keys.length when keys_count is absent", () => { + mockUseTeamsTable.mockReturnValue( + teamsResult([ + { + ...mockTeam, + keys_count: undefined, + keys: [{ token: "t1" }, { token: "t2" }] as unknown as Team["keys"], + }, + ]), + ); + renderTable(); + expect(screen.getByTitle("2 keys")).toBeInTheDocument(); + }); +}); + +describe("hidden-by-default columns", () => { + it("hides Members, Models, Rate Limits, and Updated until toggled on", async () => { + const user = userEvent.setup(); + renderTable(); + + expect(screen.queryByText("Rate Limits")).not.toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Columns" })); + const menu = await screen.findByRole("menu"); + expect(within(menu).getByText("Rate Limits")).toBeInTheDocument(); + expect(within(menu).getByText("Updated")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx new file mode 100644 index 00000000000..3b75db52b16 --- /dev/null +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx @@ -0,0 +1,201 @@ +"use client"; + +import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { useTeamsTable } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { + DataTable, + DataTableFilterDrawer, + DataTableFilterField, + DataTableToolbar, +} from "@/components/shared/DataTable"; +import { SearchSelect } from "@/components/shared/SearchSelect"; +import { Input } from "@/components/ui/input"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; +import { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; +import React, { useCallback, useMemo, useState } from "react"; + +import { Team } from "../key_team_helpers/key_list"; +import { getTeamTableColumns, TEAM_TABLE_HIDDEN_COLUMNS } from "./teamTableColumns"; + +interface TeamsTableProps { + userRole: string | null; + userID: string | null; + onSelectTeam: (team: Team) => void; + onEditTeam: (team: Team) => void; + onDeleteTeam: (team: Team) => void; +} + +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 FILTER_LABELS: Record = { + org_id: "Organization", + alias: "Team alias", + team_id: "Team ID", +}; + +export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDeleteTeam }: TeamsTableProps) { + const { data: fetchedOrganizations } = useOrganizations(); + const organizations = useMemo(() => fetchedOrganizations ?? [], [fetchedOrganizations]); + + 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 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 isAdminView = userRole === "Admin" || userRole === "Admin Viewer"; + + const teamListOptions = { + organizationID: getFilterValue("org_id"), + team_alias: getFilterValue("alias"), + teamID: getFilterValue("team_id"), + search: searchQuery.trim() || undefined, + userID: isAdminView ? undefined : userID ?? undefined, + sortBy: sorting[0]?.id, + sortOrder: toSortOrder(sorting), + }; + + const { + data: teamsResponse, + isPending: isLoading, + isFetching, + refetch, + } = useTeamsTable(tablePagination.pageIndex + 1, tablePagination.pageSize, teamListOptions); + + const teamList = useMemo(() => teamsResponse?.teams ?? [], [teamsResponse]); + const rowCount = teamsResponse?.total ?? 0; + + const handleSearchChange = useCallback((value: string) => { + setSearchInput(value); + setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); + + const handleSortingChange = useCallback>((updaterOrValue) => { + setSorting(updaterOrValue); + setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); + + const handleColumnFiltersChange = useCallback>((updaterOrValue) => { + setColumnFilters(updaterOrValue); + setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); + + const columns = useMemo(() => { + const columnDeps = { organizations, userRole, onSelectTeam, onEditTeam, onDeleteTeam }; + return getTeamTableColumns(columnDeps); + }, [organizations, userRole, onSelectTeam, onEditTeam, onDeleteTeam]); + + 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], + ); + + const formatFilterValue = useCallback( + (columnId: string, value: unknown): string => { + const raw = String(value); + if (columnId === "org_id") { + return organizations.find((org) => org.organization_id === raw)?.organization_alias || raw; + } + return raw; + }, + [organizations], + ); + + return ( + row.team_id} + defaultColumnVisibility={TEAM_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 teams..." + noDataMessage="No teams found" + maxBodyHeight="calc(75vh - 210px)" + size="compact" + toolbar={(table) => ( + <> + refetch?.()} + isRefreshing={isFetching} + onOpenFilters={() => setFiltersOpen(true)} + filterLabels={FILTER_LABELS} + formatFilterValue={formatFilterValue} + /> + + {({ get, set }) => ( + <> + + set("org_id", value)} + placeholder="Select an organization…" + emptyText="No organizations found" + /> + + + set("alias", event.target.value)} + placeholder="Enter team alias…" + /> + + + set("team_id", event.target.value)} + placeholder="Enter team ID…" + /> + + + )} + + + )} + /> + ); +} diff --git a/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx new file mode 100644 index 00000000000..ecf7387ee83 --- /dev/null +++ b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx @@ -0,0 +1,287 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { Copy, KeyRound, Layers, MoreHorizontal, Pencil, Trash2, Users } from "lucide-react"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdentityCell, SpendBudgetCell } from "@/components/shared/table_cells"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { Skeleton } from "@/components/ui/skeleton"; +import { cn } from "@/lib/cva.config"; +import { copyToClipboard, formatNumberWithCommas } from "@/utils/dataUtils"; + +import { Team } from "../key_team_helpers/key_list"; +import { Organization } from "../networking"; + +interface ResourceTone { + icon: typeof Users; + className: string; +} + +const RESOURCE_TONES: Record<"members" | "models" | "keys", ResourceTone> = { + members: { icon: Users, className: "bg-violet-50 text-violet-700 ring-violet-600/20" }, + models: { icon: Layers, className: "bg-sky-50 text-sky-700 ring-sky-600/20" }, + keys: { icon: KeyRound, className: "bg-emerald-50 text-emerald-700 ring-emerald-600/20" }, +}; + +const teamMemberCount = (team: Team): number => team.members_count ?? team.members_with_roles?.length ?? 0; +const teamModelCount = (team: Team): number => team.models?.length ?? 0; +const teamKeyCount = (team: Team): number => team.keys_count ?? team.keys?.length ?? 0; + +function ResourcesCell({ team }: { team: Team }) { + const items = [ + { key: "members" as const, label: "members", count: teamMemberCount(team) }, + { key: "models" as const, label: "models", count: teamModelCount(team) }, + { key: "keys" as const, label: "keys", count: teamKeyCount(team) }, + ]; + + return ( +
+ {items.map((item) => { + const tone = RESOURCE_TONES[item.key]; + const Icon = tone.icon; + return ( + + + {item.count} + + ); + })} +
+ ); +} + +function RateLimitLine({ label, value }: { label: string; value: number | null }) { + return ( +
+ {label} + {value != null ? formatNumberWithCommas(value) : "Unlimited"} +
+ ); +} + +interface TeamRowActionsProps { + team: Team; + canManage: boolean; + onEditTeam: (team: Team) => void; + onDeleteTeam: (team: Team) => void; +} + +function TeamRowActions({ team, canManage, onEditTeam, onDeleteTeam }: TeamRowActionsProps) { + const handleCopy = () => { + void copyToClipboard(team.team_id, "Team ID copied"); + }; + + return ( + + + + + + {canManage && ( + onEditTeam(team)} data-testid="team-action-edit"> + + Edit team + + )} + + + Copy team ID + + {canManage && ( + <> + + onDeleteTeam(team)} data-testid="team-action-delete"> + + Delete team + + + )} + + + ); +} + +interface TeamTableColumnsDeps { + organizations: Organization[]; + userRole: string | null; + onSelectTeam: (team: Team) => void; + onEditTeam: (team: Team) => void; + onDeleteTeam: (team: Team) => void; +} + +export const getTeamTableColumns = ({ + organizations, + userRole, + onSelectTeam, + onEditTeam, + onDeleteTeam, +}: TeamTableColumnsDeps): ColumnDef[] => { + const canManage = userRole === "Admin"; + + return [ + { + id: "team_alias", + accessorKey: "team_alias", + meta: { + title: "Team", + renderSkeleton: () => ( +
+ + +
+ ), + }, + header: ({ column }) => , + size: 260, + enableSorting: true, + cell: ({ row }) => { + const team = row.original; + const hasAlias = Boolean(team.team_alias); + return ( + onSelectTeam(team)} + /> + ); + }, + }, + { + id: "organization_alias", + accessorKey: "organization_id", + meta: { title: "Organization" }, + header: "Organization", + size: 160, + 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: "resources", + meta: { + title: "Resources", + renderSkeleton: () => ( +
+ + + +
+ ), + }, + header: "Resources", + size: 210, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "spend", + accessorKey: "spend", + meta: { title: "Spend / Budget", skeleton: "meter" }, + header: "Spend / Budget", + size: 200, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "created_at", + accessorKey: "created_at", + meta: { title: "Created" }, + header: ({ column }) => , + size: 130, + enableSorting: true, + cell: (info) => , + }, + { + id: "members", + meta: { title: "Members" }, + header: "Members", + size: 110, + enableSorting: false, + cell: ({ row }) => {teamMemberCount(row.original)}, + }, + { + id: "models", + meta: { title: "Models" }, + header: "Models", + size: 100, + enableSorting: false, + cell: ({ row }) => {teamModelCount(row.original)}, + }, + { + id: "rate_limits", + meta: { title: "Rate Limits", skeleton: "twoLine" }, + header: "Rate Limits", + size: 140, + enableSorting: false, + cell: ({ row }) => ( +
+ + +
+ ), + }, + { + id: "updated_at", + accessorKey: "updated_at", + meta: { title: "Updated" }, + header: "Updated", + size: 130, + enableSorting: false, + cell: (info) => , + }, + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 60, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + }, + ]; +}; + +export const TEAM_TABLE_HIDDEN_COLUMNS: Record = { + members: false, + models: false, + rate_limits: false, + updated_at: false, +}; diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 02f5d588149..513054aae7a 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -16,6 +16,8 @@ vi.mock("@tanstack/react-pacer/debouncer", async () => { const [value, setValue] = React.useState(initial); return [value, setValue, { cancel: vi.fn(), flush: vi.fn() }]; }, + useDebouncedCallback: (fn: (...args: unknown[]) => void) => fn, + useDebouncer: (fn: (...args: unknown[]) => void) => ({ maybeExecute: fn, cancel: vi.fn(), flush: vi.fn() }), }; }); @@ -63,7 +65,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 +156,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 +174,21 @@ it("should render VirtualKeysTable component", () => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); +it("left-anchors the create-key CTA below the title, between the header and the table toolbar", () => { + renderWithProviders(Create New Key} />); + + const heading = screen.getByRole("heading", { name: "Virtual Keys" }); + const ctas = screen.getAllByRole("button", { name: "Create New Key" }); + expect(ctas).toHaveLength(1); + const cta = ctas[0]; + const search = screen.getByPlaceholderText(/Search by key alias/); + + // The CTA follows the title row... + expect(heading.compareDocumentPosition(cta) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + // ...and precedes the table's search toolbar, so it sits in its own row above the table. + expect(cta.compareDocumentPosition(search) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); +}); + it("should display key information correctly", async () => { renderWithProviders(); @@ -177,6 +196,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 +208,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 +261,98 @@ 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", { selector: "[data-sort-field='spend']" })).toBeInTheDocument(); + expect(screen.getByText("Budget", { selector: "[data-sort-field='max_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("sorts by the backend max_budget field when 'Budget descending' is chosen from the Spend / Budget menu", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByTestId("sort-trigger-spend")); + await user.click(await screen.findByText("Budget descending")); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 1, + 50, + expect.objectContaining({ sortBy: "max_budget", sortOrder: "desc" }), + ); + }); +}); + +it("emphasizes the active field in the Spend / Budget header so the sorted column reads without opening the menu", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByTestId("sort-trigger-spend")); + await user.click(await screen.findByText("Budget descending")); + + await waitFor(() => { + expect(screen.getByText("Budget", { selector: "[data-sort-field='max_budget']" }).className).toContain( + "font-semibold", + ); + }); + expect(screen.getByText("Spend", { selector: "[data-sort-field='spend']" }).className).toContain( + "text-muted-foreground", + ); +}); + +it("sorts by spend ascending when 'Spend ascending' is chosen from the Spend / Budget menu", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByTestId("sort-trigger-spend")); + await user.click(await screen.findByText("Spend ascending")); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: "spend", sortOrder: "asc" })); + }); +}); + +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 +374,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 +381,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 +396,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 +417,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 +445,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 +456,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 +503,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 +526,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 cae6dc54df5..fe929dc0179 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -1,810 +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 { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; 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: 300 }); + 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." + /> + {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..133ff89a898 --- /dev/null +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -0,0 +1,364 @@ +"use client"; + +import { InfoCircleOutlined } from "@ant-design/icons"; +import { ColumnDef } from "@tanstack/react-table"; +import { Popover, Typography } from "antd"; + +import { DataTableMultiSortHeader, DataTableSortHeader, type DataTableSortField } 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 SPEND_BUDGET_SORT_FIELDS: DataTableSortField[] = [ + { id: "spend", label: "Spend" }, + { id: "max_budget", label: "Budget" }, +]; + +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: ({ table }) => , + 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/add_model/AdaptiveRoutingConfig.tsx b/ui/litellm-dashboard/src/components/add_model/AdaptiveRoutingConfig.tsx new file mode 100644 index 00000000000..720b6f88db3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/AdaptiveRoutingConfig.tsx @@ -0,0 +1,134 @@ +import { Card, InputNumber, Radio, Slider, Space, Switch, Typography } from "antd"; +import React from "react"; +import { + AdaptiveEligible, + ComplexityRouterConfigValue, + DEFAULT_ADAPTIVE_WEIGHTS, + DEFAULT_TIER_DISTANCE_PENALTY, +} from "./ComplexityRouterConfig"; + +const { Text } = Typography; + +interface AdaptiveRoutingConfigProps { + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; +} + +const AdaptiveRoutingConfig: React.FC = ({ value, onChange }) => { + const adaptiveWeights = value.adaptive_weights ?? DEFAULT_ADAPTIVE_WEIGHTS; + const adaptiveEligible = value.adaptive_eligible ?? "all"; + const tierDistancePenalty = value.tier_distance_penalty ?? DEFAULT_TIER_DISTANCE_PENALTY; + + const handleAdaptiveToggle = (adaptive: boolean) => { + const nextValue: ComplexityRouterConfigValue = { + ...value, + adaptive, + adaptive_weights: adaptiveWeights, + adaptive_eligible: adaptiveEligible, + tier_distance_penalty: tierDistancePenalty, + }; + onChange(nextValue); + }; + + const handleQualityWeightChange = (qualityPercent: number) => { + const quality = qualityPercent / 100; + onChange({ ...value, adaptive_weights: { quality, cost: Math.round((1 - quality) * 100) / 100 } }); + }; + + const handleAdaptiveEligibleChange = (eligible: AdaptiveEligible) => { + onChange({ ...value, adaptive_eligible: eligible }); + }; + + const handleTierDistancePenaltyChange = (penalty: number | null) => { + onChange({ ...value, tier_distance_penalty: penalty ?? DEFAULT_TIER_DISTANCE_PENALTY }); + }; + + return ( + <> +
+ + Enable adaptive bandit selection +
+ + When disabled, each request always uses the model assigned to its classified tier. + + + + + How Adaptive Routing Works + + + It learns from how each conversation actually goes: does the user have to rephrase or correct the model, does + it get stuck repeating itself, does it run out of tool calls, does the user seem satisfied. Combined with + cost, this live feedback shifts future routing toward the models that are actually working well, and improves + as more conversations come in. Until there's enough feedback, it defaults to the classified tier's + model. + + + + {value.adaptive && ( +
+
+ + Quality vs. Cost ({Math.round(adaptiveWeights.quality * 100)}% quality /{" "} + {Math.round(adaptiveWeights.cost * 100)}% cost) + + `${v}% quality / ${100 - (v ?? 0)}% cost` }} + /> + + Higher quality weight favors more capable (pricier) models; higher cost weight favors cheaper models when + the bandit has feedback to act on. Recommended: 30% quality / 70% cost split. + +
+ +
+ + Eligible Model Pool + + handleAdaptiveEligibleChange(e.target.value)} + className="w-full" + > + + + All tiers (soft floor){" "} + — router can pick across tiers, depending on the best fit for the prompt + + + Classified tier only{" "} + — router can only pick models within tier + + + +
+ + {adaptiveEligible === "all" && ( +
+ + Tier Distance Penalty + + + + Score penalty applied per tier-step away from the classified tier. + +
+ )} +
+ )} + + ); +}; + +export default AdaptiveRoutingConfig; diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx new file mode 100644 index 00000000000..92df8029edc --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -0,0 +1,171 @@ +import { InfoCircleOutlined } from "@ant-design/icons"; +import { Select as AntdSelect, Card, InputNumber, Radio, Space, Tooltip, Typography } from "antd"; +import React from "react"; +import { ClassifierType, ComplexityRouterConfigValue, DEFAULT_CLASSIFIER_TIMEOUT_MS } from "./ComplexityRouterConfig"; + +const { Text } = Typography; + +interface ClassificationMethodConfigProps { + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; + modelOptions: { value: string; label: string }[]; + customTechnicalKeywords?: string[]; + onCustomTechnicalKeywordsChange?: (keywords: string[]) => void; + showValidationErrors?: boolean; +} + +const ClassificationMethodConfig: React.FC = ({ + value, + onChange, + modelOptions, + customTechnicalKeywords, + onCustomTechnicalKeywordsChange, + showValidationErrors = false, +}) => { + const classifierModelMissing = + showValidationErrors && value.classifier_type === "llm" && !value.classifier_llm_config?.model; + + const handleClassifierTypeChange = (classifierType: ClassifierType) => { + onChange({ + ...value, + classifier_type: classifierType, + classifier_llm_config: + classifierType === "llm" + ? value.classifier_llm_config ?? { model: "", timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS } + : undefined, + }); + }; + + const handleClassifierModelChange = (model: string) => { + onChange({ + ...value, + classifier_llm_config: { + model, + timeout_ms: value.classifier_llm_config?.timeout_ms ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, + }, + }); + }; + + const handleClassifierTimeoutChange = (timeoutMs: number | null) => { + onChange({ + ...value, + classifier_llm_config: { + model: value.classifier_llm_config?.model ?? "", + timeout_ms: timeoutMs ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, + }, + }); + }; + + return ( + <> + handleClassifierTypeChange(e.target.value)} + className="w-full" + > + + + Heuristic{" "} + (default) — rule-based scoring, no API calls, <1ms latency + + + LLM Classifier{" "} + — use a model to decide the tier (e.g. a small/fast model) + + + + + {value.classifier_type === "llm" && ( +
+
+ + Classifier Model + + + {classifierModelMissing && ( + + A classifier model is required + + )} +
+
+ + Timeout (ms) + + + + Falls back to the heuristic scorer if the classifier call errors, times out, or returns an unparseable + response. + +
+
+ )} + + {value.classifier_type === "heuristic" && ( +
+
+ Custom Technical Keywords + + + +
+ + Optional: Add terms to the built-in list to improve classification accuracy on the technical dimension. + (e.g., udp, kafka, terraform). + + onCustomTechnicalKeywordsChange?.(keywords)} + placeholder="Type a keyword and press Enter, or paste a comma-separated list" + tokenSeparators={[","]} + open={false} + suffixIcon={null} + style={{ width: "100%" }} + allowClear + /> +
+ )} + + + + How Classification Works + + + The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical + terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the + tier: + +
    +
  • + SIMPLE: Score < 0.15 +
    • +
  • + MEDIUM: Score 0.15 - 0.35 +
    • +
  • + COMPLEX: Score 0.35 - 0.60 +
    • +
  • + REASONING: Score > 0.60 (or 2+ reasoning markers) +
    • +
+ + + ); +}; + +export default ClassificationMethodConfig; diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index a433808085b..e1f90296770 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -4,17 +4,18 @@ import { vi } from "vitest"; import ComplexityRouterConfig, { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; const mockModelInfo = [ - { model_group: "gpt-4" }, - { model_group: "gpt-3.5-turbo" }, - { model_group: "claude-3-opus" }, + { model_group: "gpt-4", mode: "chat" }, + { model_group: "gpt-3.5-turbo", mode: "chat" }, + { model_group: "claude-3-opus", mode: "chat" }, + { model_group: "text-embedding-3-small", mode: "embedding" }, ] as any[]; const defaultValue: ComplexityRouterConfigValue = { tiers: { - SIMPLE: "gpt-3.5-turbo", - MEDIUM: "gpt-3.5-turbo", - COMPLEX: "gpt-4", - REASONING: "claude-3-opus", + SIMPLE: ["gpt-3.5-turbo"], + MEDIUM: ["gpt-3.5-turbo"], + COMPLEX: ["gpt-4"], + REASONING: ["claude-3-opus"], }, classifier_type: "heuristic", }; @@ -57,11 +58,13 @@ describe("ComplexityRouterConfig", () => { it("should display the how classification works section", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText("How Classification Works")).toBeInTheDocument(); }); it("should show score thresholds in the classification section", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText(/Score < 0.15/)).toBeInTheDocument(); expect(screen.getByText(/Score 0.15 - 0.35/)).toBeInTheDocument(); expect(screen.getByText(/Score 0.35 - 0.60/)).toBeInTheDocument(); @@ -106,6 +109,7 @@ describe("ComplexityRouterConfig", () => { it("should render the custom technical keywords field", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText("Custom Technical Keywords")).toBeInTheDocument(); }); @@ -117,6 +121,7 @@ describe("ComplexityRouterConfig", () => { onCustomTechnicalKeywordsChange={vi.fn()} />, ); + fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText("udp")).toBeInTheDocument(); expect(screen.getByText("kafka")).toBeInTheDocument(); }); @@ -131,14 +136,16 @@ describe("ComplexityRouterConfig", () => { onCustomTechnicalKeywordsChange={onCustomTechnicalKeywordsChange} />, ); - const keywordsCard = screen.getByText("Custom Technical Keywords").closest(".ant-card") as HTMLElement; - const input = within(keywordsCard).getByRole("combobox"); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + const keywordsSection = screen.getByText("Custom Technical Keywords").closest("div")?.parentElement as HTMLElement; + const input = within(keywordsSection).getByRole("combobox"); await user.type(input, "udp,"); expect(onCustomTechnicalKeywordsChange).toHaveBeenCalledWith(["udp"]); }); it("should render an empty state when no keyword tier rules exist", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); expect(screen.getByText("Keyword Tier Overrides")).toBeInTheDocument(); expect(screen.getByText("No keyword tier overrides configured")).toBeInTheDocument(); }); @@ -157,6 +164,7 @@ describe("ComplexityRouterConfig", () => { const user = userEvent.setup(); const onKeywordTierRulesChange = vi.fn(); renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); await user.click(screen.getByRole("button", { name: /add keyword rule/i })); expect(onKeywordTierRulesChange).toHaveBeenCalledTimes(1); const newRules = onKeywordTierRulesChange.mock.calls[0][0]; @@ -174,6 +182,7 @@ describe("ComplexityRouterConfig", () => { onKeywordTierRulesChange={onKeywordTierRulesChange} />, ); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); expect(screen.getByText("invoice")).toBeInTheDocument(); expect(screen.getByText("refund")).toBeInTheDocument(); @@ -183,6 +192,7 @@ describe("ComplexityRouterConfig", () => { it("should not show embedding model or match score fields when semantic matching is disabled", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); expect(screen.getByText("Semantic keyword matching")).toBeInTheDocument(); expect(screen.queryByText("Embedding model")).not.toBeInTheDocument(); expect(screen.queryByText("Minimum match score")).not.toBeInTheDocument(); @@ -190,6 +200,7 @@ describe("ComplexityRouterConfig", () => { it("should show embedding model and match score fields when semantic matching is enabled", () => { renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); expect(screen.getByText("Embedding model")).toBeInTheDocument(); expect(screen.getByText("Minimum match score")).toBeInTheDocument(); }); @@ -204,7 +215,58 @@ describe("ComplexityRouterConfig", () => { onSemanticMatchingEnabledChange={onSemanticMatchingEnabledChange} />, ); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); await user.click(screen.getByRole("switch")); expect(onSemanticMatchingEnabledChange).toHaveBeenCalledWith(true, expect.anything()); }); + + it("excludes embedding-mode models from the tier and classifier dropdowns", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const simpleTierSection = screen.getByText("Simple Tier").closest(".mb-4") as HTMLElement; + const combobox = within(simpleTierSection).getByRole("combobox"); + await user.click(combobox); + + expect((await screen.findAllByText("gpt-3.5-turbo")).length).toBeGreaterThan(0); + expect(screen.queryAllByText("text-embedding-3-small")).toHaveLength(0); + }); + + it("does not show tier validation errors by default", () => { + renderWithProviders(); + expect(screen.queryByText("This tier is required")).not.toBeInTheDocument(); + }); + + it("shows an inline error on the classifier model select when llm is selected without a model", () => { + const llmValue: ComplexityRouterConfigValue = { + ...defaultValue, + classifier_type: "llm", + classifier_llm_config: { model: "", timeout_ms: 3000 }, + }; + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.getByText("A classifier model is required")).toBeInTheDocument(); + }); + + it("does not show the classifier model error once a classifier model is set", () => { + const llmValue: ComplexityRouterConfigValue = { + ...defaultValue, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 }, + }; + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.queryByText("A classifier model is required")).not.toBeInTheDocument(); + }); + + it("shows a validation error only under unfilled tiers when showValidationErrors is true", () => { + renderWithProviders( + , + ); + expect(screen.getAllByText("This tier is required")).toHaveLength(1); + }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 555648db8ad..855a1b27df9 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -1,19 +1,22 @@ import { InfoCircleOutlined } from "@ant-design/icons"; -import { Select as AntdSelect, Card, Collapse, Divider, InputNumber, Radio, Space, Tooltip, Typography } from "antd"; +import { Select as AntdSelect, Card, Collapse, Divider, Space, Tooltip, Typography } from "antd"; import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; +import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; +import ClassificationMethodConfig from "./ClassificationMethodConfig"; import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules"; import SemanticKeywordMatching from "./SemanticKeywordMatching"; const { Text } = Typography; export const DEFAULT_CLASSIFIER_TIMEOUT_MS = 3000; +export const DEFAULT_TIER_DISTANCE_PENALTY = 0.5; export interface ComplexityTiers { - SIMPLE: string; - MEDIUM: string; - COMPLEX: string; - REASONING: string; + SIMPLE: string[]; + MEDIUM: string[]; + COMPLEX: string[]; + REASONING: string[]; } export interface ClassifierLLMConfig { @@ -23,10 +26,23 @@ export interface ClassifierLLMConfig { export type ClassifierType = "heuristic" | "llm"; +export interface AdaptiveRouterWeights { + quality: number; + cost: number; +} + +export const DEFAULT_ADAPTIVE_WEIGHTS: AdaptiveRouterWeights = { quality: 0.3, cost: 0.7 }; + +export type AdaptiveEligible = "all" | "classified_tier"; + export interface ComplexityRouterConfigValue { tiers: ComplexityTiers; classifier_type: ClassifierType; classifier_llm_config?: ClassifierLLMConfig; + adaptive?: boolean; + adaptive_weights?: AdaptiveRouterWeights; + tier_distance_penalty?: number; + adaptive_eligible?: AdaptiveEligible; } interface ComplexityRouterConfigProps { @@ -45,6 +61,7 @@ interface ComplexityRouterConfigProps { onEmbeddingModelChange?: (model: string) => void; matchThreshold?: number; onMatchThresholdChange?: (threshold: number) => void; + showValidationErrors?: boolean; } const TIER_DESCRIPTIONS: Record = { @@ -84,48 +101,20 @@ const ComplexityRouterConfig: React.FC = ({ onEmbeddingModelChange = () => {}, matchThreshold = 0.5, onMatchThresholdChange = () => {}, + showValidationErrors = false, }) => { - // Prepare model options for dropdowns - const modelOptions = modelInfo.map((model) => ({ - value: model.model_group, - label: model.model_group, - })); + // Embedding models can't serve a chat-completion role, so they're excluded here. + const modelOptions = modelInfo + .filter((model) => model.mode !== "embedding") + .map((model) => ({ + value: model.model_group, + label: model.model_group, + })); - const handleTierChange = (tier: keyof ComplexityTiers, model: string) => { + const handleTierChange = (tier: keyof ComplexityTiers, models: string[]) => { onChange({ ...value, - tiers: { ...value.tiers, [tier]: model }, - }); - }; - - const handleClassifierTypeChange = (classifierType: ClassifierType) => { - onChange({ - ...value, - classifier_type: classifierType, - classifier_llm_config: - classifierType === "llm" - ? value.classifier_llm_config ?? { model: "", timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS } - : undefined, - }); - }; - - const handleClassifierModelChange = (model: string) => { - onChange({ - ...value, - classifier_llm_config: { - model, - timeout_ms: value.classifier_llm_config?.timeout_ms ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, - }, - }); - }; - - const handleClassifierTimeoutChange = (timeoutMs: number | null) => { - onChange({ - ...value, - classifier_llm_config: { - model: value.classifier_llm_config?.model ?? "", - timeout_ms: timeoutMs ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, - }, + tiers: { ...value.tiers, [tier]: models }, }); }; @@ -135,19 +124,20 @@ const ComplexityRouterConfig: React.FC = ({ Complexity Tier Configuration - + The complexity router automatically classifies requests by complexity using rule-based scoring (no API calls, - <1ms latency). Configure which model handles each tier. + <1ms latency). Configure which model(s) handle each tier. {(Object.keys(TIER_DESCRIPTIONS) as Array).map((tier, index) => { const tierInfo = TIER_DESCRIPTIONS[tier]; + const tierMissing = showValidationErrors && value.tiers[tier].length === 0; return (
{index > 0 && } @@ -164,13 +154,26 @@ const ComplexityRouterConfig: React.FC = ({ Examples: {tierInfo.examples} handleTierChange(tier, model)} - placeholder={`Select model for ${tierInfo.label.toLowerCase()} queries`} + onChange={(models) => handleTierChange(tier, models)} + placeholder={`Select model(s) for ${tierInfo.label.toLowerCase()} queries`} showSearch style={{ width: "100%" }} options={modelOptions} + status={tierMissing ? "error" : undefined} /> + {value.tiers[tier].length > 1 && ( + + Multiple models selected — the router randomly picks among them per request (or Thompson-samples + within the pool when adaptive routing is on). + + )} + {tierMissing && ( + + This tier is required + + )}
); @@ -191,141 +194,61 @@ const ComplexityRouterConfig: React.FC = ({ ), children: ( - <> - handleClassifierTypeChange(e.target.value)} - className="w-full" - > - - - Heuristic{" "} - (default) — rule-based scoring, no API calls, <1ms latency - - - LLM Classifier{" "} - — use a model to decide the tier (e.g. a small/fast model) - - - - - {value.classifier_type === "llm" && ( -
-
- - Classifier Model - - -
-
- - Timeout (ms) - - - - Falls back to the heuristic scorer if the classifier call errors, times out, or returns an - unparseable response. - -
-
- )} - + ), }, + { + key: "adaptive", + label: ( + + Advanced: Adaptive Routing + + ), + children: , + }, + ...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange + ? [ + { + key: "keyword-semantic", + label: ( + + Advanced: Keyword/Semantic Matching + + ), + children: ( + <> + {onKeywordTierRulesChange && ( + + )} + {onKeywordTierRulesChange && onSemanticMatchingEnabledChange && ( + + )} + {onSemanticMatchingEnabledChange && ( + + )} + + ), + }, + ] + : []), ]} /> - - - - -
- - Custom Technical Keywords - - - - -
- - Optional: Add terms to the built-in list to improve classification accuracy on the technical dimension. (e.g., - udp, kafka, terraform). - - onCustomTechnicalKeywordsChange?.(keywords)} - placeholder="Type a keyword and press Enter, or paste a comma-separated list" - tokenSeparators={[","]} - open={false} - suffixIcon={null} - style={{ width: "100%" }} - allowClear - /> - - - - - - - How Classification Works - - - The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical - terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the - tier: - -
    -
  • - SIMPLE: Score < 0.15 -
    • -
  • - MEDIUM: Score 0.15 - 0.35 -
    • -
  • - COMPLEX: Score 0.35 - 0.60 -
    • -
  • - REASONING: Score > 0.60 (or 2+ reasoning markers) -
    • -
- - - {/* Keyword-tier and semantic sections only render when their change handlers are - wired (the add-router flow). The edit-auto-router modal doesn't pass them yet, so - they stay hidden there rather than rendering interactive-but-dead controls. */} - {onKeywordTierRulesChange && ( - <> - - - - )} - - {onSemanticMatchingEnabledChange && ( - <> - - - - )}
); }; diff --git a/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.test.tsx b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.test.tsx new file mode 100644 index 00000000000..2336e6faf43 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.test.tsx @@ -0,0 +1,53 @@ +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import userEvent from "@testing-library/user-event"; +import { vi } from "vitest"; +import SemanticKeywordMatching from "./SemanticKeywordMatching"; + +const mockModelInfo = [ + { model_group: "gpt-4", mode: "chat" }, + { model_group: "text-embedding-3-small", mode: "embedding" }, + { model_group: "voyage-3-5", mode: "embedding" }, + { model_group: "legacy-model" }, +] as any[]; + +const baseProps = { + enabled: true, + onEnabledChange: vi.fn(), + embeddingModel: undefined, + onEmbeddingModelChange: vi.fn(), + matchThreshold: 0.5, + onMatchThresholdChange: vi.fn(), + modelInfo: mockModelInfo, +}; + +describe("SemanticKeywordMatching", () => { + it("only lists embedding-mode models in the embedding model dropdown", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const combobox = screen.getByRole("combobox"); + await user.click(combobox); + + expect((await screen.findAllByText("text-embedding-3-small")).length).toBeGreaterThan(0); + expect(screen.getAllByText("voyage-3-5").length).toBeGreaterThan(0); + expect(screen.queryAllByText("gpt-4")).toHaveLength(0); + expect(screen.queryAllByText("legacy-model")).toHaveLength(0); + }); + + it("does not show a validation error by default", () => { + renderWithProviders(); + expect(screen.queryByText("An embedding model is required")).not.toBeInTheDocument(); + }); + + it("shows a validation error when showValidationErrors is true and no embedding model is set", () => { + renderWithProviders(); + expect(screen.getByText("An embedding model is required")).toBeInTheDocument(); + }); + + it("hides the validation error once an embedding model is set", () => { + renderWithProviders( + , + ); + expect(screen.queryByText("An embedding model is required")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx index 0f9907ac6c9..c2252843b29 100644 --- a/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx +++ b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx @@ -1,5 +1,5 @@ import { InfoCircleOutlined } from "@ant-design/icons"; -import { Card, InputNumber, Select as AntdSelect, Switch, Tooltip, Typography } from "antd"; +import { InputNumber, Select as AntdSelect, Switch, Tooltip, Typography } from "antd"; import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; @@ -15,6 +15,7 @@ interface SemanticKeywordMatchingProps { matchThreshold: number; onMatchThresholdChange: (threshold: number) => void; modelInfo: ModelGroup[]; + showValidationErrors?: boolean; } const SemanticKeywordMatching: React.FC = ({ @@ -25,14 +26,17 @@ const SemanticKeywordMatching: React.FC = ({ matchThreshold, onMatchThresholdChange, modelInfo, + showValidationErrors = false, }) => { - const modelOptions = Array.from(new Set(modelInfo.map((model) => model.model_group))).map((model_group) => ({ + const embeddingModels = modelInfo.filter((model) => model.mode === "embedding"); + const modelOptions = Array.from(new Set(embeddingModels.map((model) => model.model_group))).map((model_group) => ({ value: model_group, label: model_group, })); + const embeddingModelMissing = showValidationErrors && !embeddingModel; return ( - +
@@ -60,7 +64,13 @@ const SemanticKeywordMatching: React.FC = ({ showSearch style={{ width: "100%" }} options={modelOptions} + status={embeddingModelMissing ? "error" : undefined} /> + {embeddingModelMissing && ( + + An embedding model is required + + )}
Minimum match score @@ -76,7 +86,7 @@ const SemanticKeywordMatching: React.FC = ({
)} - +
); }; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx new file mode 100644 index 00000000000..4713f8c6869 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -0,0 +1,40 @@ +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import userEvent from "@testing-library/user-event"; +import { vi } from "vitest"; +import { Form } from "antd"; +import AddAutoRouterTab from "./add_auto_router_tab"; +import NotificationManager from "../molecules/notifications_manager"; + +vi.mock("../networking", () => ({ + modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), +})); + +vi.mock("@/components/llm_calls/fetch_models", () => ({ + fetchAvailableModels: vi.fn().mockResolvedValue([]), +})); + +vi.mock("./handle_add_auto_router_submit", () => ({ + handleAddAutoRouterSubmit: vi.fn(), +})); + +vi.mock("../molecules/notifications_manager", () => ({ + default: { fromBackend: vi.fn() }, +})); + +const Harness = () => { + const [form] = Form.useForm(); + return ; +}; + +describe("AddAutoRouterTab", () => { + it("flags every mandatory field when Add Auto Router is clicked with nothing filled", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + expect(await screen.findByText("Auto router name is required")).toBeInTheDocument(); + expect(screen.getAllByText("This tier is required")).toHaveLength(4); + expect(NotificationManager.fromBackend).toHaveBeenCalledWith("Please enter an Auto Router Name"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 79c07040210..8724c27b41a 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -1,5 +1,5 @@ import React, { useEffect, useState } from "react"; -import { Card, Form, Button, Tooltip, Typography, Select as AntdSelect, Radio, Badge, Space } from "antd"; +import { Card, Form, Button, Tooltip, Typography, Select as AntdSelect, Radio, Badge, Space, Modal } from "antd"; import type { FormInstance } from "antd"; import { ThunderboltOutlined, BranchesOutlined } from "@ant-design/icons"; import { Text, TextInput } from "@tremor/react"; @@ -8,10 +8,20 @@ import { all_admin_roles } from "@/utils/roles"; import { handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import RouterConfigBuilder from "./RouterConfigBuilder"; -import ComplexityRouterConfig, { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import ComplexityRouterConfig, { + ComplexityRouterConfigValue, + DEFAULT_ADAPTIVE_WEIGHTS, + DEFAULT_TIER_DISTANCE_PENALTY, +} from "./ComplexityRouterConfig"; import { KeywordTierRule } from "./KeywordTierRules"; import { DEFAULT_MATCH_THRESHOLD } from "./SemanticKeywordMatching"; -import { buildComplexityRouterConfig, getSemanticConfigError } from "./build_complexity_router_config"; +import { + buildComplexityRouterConfig, + getMissingTiersError, + getSemanticConfigError, +} from "./build_complexity_router_config"; +import { buildAutoRouterTestTargets, AutoRouterTestTarget } from "./build_auto_router_test_targets"; +import AutoRouterConnectionTest from "./auto_router_connection_test"; import NotificationManager from "../molecules/notifications_manager"; interface AddAutoRouterTabProps { @@ -32,7 +42,7 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc const [routerType, setRouterType] = useState("recommended"); const [complexityRouterConfig, setComplexityRouterConfig] = useState({ - tiers: { SIMPLE: "", MEDIUM: "", COMPLEX: "", REASONING: "" }, + tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic", }); @@ -41,10 +51,16 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc const [semanticMatchingEnabled, setSemanticMatchingEnabled] = useState(false); const [embeddingModel, setEmbeddingModel] = useState(undefined); const [matchThreshold, setMatchThreshold] = useState(DEFAULT_MATCH_THRESHOLD); + const [showValidationErrors, setShowValidationErrors] = useState(false); // Semantic router config (existing) const [routerConfig, setRouterConfig] = useState(null); + const [isTestModalVisible, setIsTestModalVisible] = useState(false); + const [isTestingConnection, setIsTestingConnection] = useState(false); + const [connectionTestId, setConnectionTestId] = useState(0); + const [testTargets, setTestTargets] = useState([]); + useEffect(() => { const fetchModelAccessGroups = async () => { const response = await modelAvailableCall(accessToken, "", "", false, null, true, true); @@ -77,26 +93,33 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc tiers, classifier_type: classifierType, classifier_llm_config: classifierLlmConfig, + adaptive = false, + adaptive_weights: adaptiveWeights = DEFAULT_ADAPTIVE_WEIGHTS, + tier_distance_penalty: tierDistancePenalty = DEFAULT_TIER_DISTANCE_PENALTY, + adaptive_eligible: adaptiveEligible = "all", } = complexityRouterConfig; - const filledTiers = Object.values(tiers).filter(Boolean); - if (filledTiers.length === 0) { - NotificationManager.fromBackend("Please select at least one model for a complexity tier"); + const missingTiersError = getMissingTiersError(tiers); + if (missingTiersError) { + setShowValidationErrors(true); + NotificationManager.fromBackend(missingTiersError); return; } if (classifierType === "llm" && !classifierLlmConfig?.model) { + setShowValidationErrors(true); NotificationManager.fromBackend("Please select a classifier model, or switch back to Heuristic"); return; } const semanticError = getSemanticConfigError({ semanticMatchingEnabled, embeddingModel, keywordTierRules }); if (semanticError) { + setShowValidationErrors(true); NotificationManager.fromBackend(semanticError); return; } - const defaultModel = tiers.MEDIUM || tiers.SIMPLE || tiers.COMPLEX || tiers.REASONING; + const defaultModel = tiers.MEDIUM[0] || tiers.SIMPLE[0] || tiers.COMPLEX[0] || tiers.REASONING[0]; form.setFieldsValue({ custom_llm_provider: "auto_router", @@ -117,6 +140,10 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc semanticMatchingEnabled, embeddingModel, matchThreshold, + adaptive, + adaptiveWeights, + tierDistancePenalty, + adaptiveEligible, }; const submitValues = { @@ -183,6 +210,8 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc const handleAutoRouterSubmit = () => { const name = form.getFieldValue("auto_router_name"); if (!name) { + setShowValidationErrors(true); + form.validateFields(["auto_router_name"]).catch(() => undefined); NotificationManager.fromBackend("Please enter an Auto Router Name"); return; } @@ -194,6 +223,24 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc } }; + const handleTestConnection = () => { + const targets = buildAutoRouterTestTargets({ + tiers: complexityRouterConfig.tiers, + semanticMatchingEnabled, + embeddingModel, + }); + + if (targets.length === 0) { + NotificationManager.fromBackend("Please select at least one model for a complexity tier"); + return; + } + + setTestTargets(targets); + setConnectionTestId((id) => id + 1); + setIsTestingConnection(true); + setIsTestModalVisible(true); + }; + return ( <> Add Auto Router @@ -205,7 +252,14 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc
Router Type - setRouterType(e.target.value)} className="w-full"> + { + setRouterType(e.target.value); + setShowValidationErrors(false); + }} + className="w-full" + >
@@ -271,6 +325,7 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc onEmbeddingModelChange={setEmbeddingModel} matchThreshold={matchThreshold} onMatchThresholdChange={setMatchThreshold} + showValidationErrors={showValidationErrors} />
) : ( @@ -355,10 +410,15 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc Need Help?
- {/* TODO: add back a Test Connection or JSON preview action here. Test Connection was removed - because prepareModelAddRequest can't build a valid pre-save payload for an auto router - (tiers are model-group references, not litellm_params); a JSON preview of the - complexity_router_config would be a good alternative. */} + {routerType === "recommended" && ( + + )}
+ + { + setIsTestModalVisible(false); + setIsTestingConnection(false); + }} + footer={[ + , + ]} + width={700} + > + {isTestModalVisible && ( + setIsTestingConnection(false)} + /> + )} + ); }; diff --git a/ui/litellm-dashboard/src/components/add_model/auto_router_connection_test.test.tsx b/ui/litellm-dashboard/src/components/add_model/auto_router_connection_test.test.tsx new file mode 100644 index 00000000000..b07270b5ced --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/auto_router_connection_test.test.tsx @@ -0,0 +1,84 @@ +import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; +import { vi } from "vitest"; +import AutoRouterConnectionTest from "./auto_router_connection_test"; +import { AutoRouterTestTarget } from "./build_auto_router_test_targets"; + +vi.mock("../networking", async () => { + const actual = await vi.importActual("../networking"); + return { + ...actual, + testModelGroupConnection: vi.fn(), + }; +}); + +const getMock = async () => vi.mocked((await import("../networking")).testModelGroupConnection); + +const targets: AutoRouterTestTarget[] = [ + { labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }, + { labels: ["MEDIUM", "COMPLEX"], modelGroup: "claude-sonnet-4", mode: "chat" }, + { labels: ["Embedding"], modelGroup: "voyage-3-5", mode: "embedding" }, +]; + +describe("AutoRouterConnectionTest", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("probes each target once with the right model and mode (chat for tiers, embedding for the embedding model)", async () => { + const mock = await getMock(); + mock.mockResolvedValue({ status: "success" }); + + renderWithProviders(); + + await waitFor(() => expect(mock).toHaveBeenCalledTimes(3)); + + expect(mock).toHaveBeenCalledWith("sk-test", "gpt-4o-mini", "chat"); + expect(mock).toHaveBeenCalledWith("sk-test", "claude-sonnet-4", "chat"); + expect(mock).toHaveBeenCalledWith("sk-test", "voyage-3-5", "embedding"); + }); + + it("shows a success indicator per target when the routing probe passes", async () => { + const mock = await getMock(); + mock.mockResolvedValue({ status: "success" }); + + renderWithProviders(); + + await waitFor(() => expect(screen.getAllByTestId("test-status-success")).toHaveLength(3)); + expect(screen.queryByTestId("test-status-error")).toBeNull(); + expect(screen.getByText("MEDIUM, COMPLEX")).toBeInTheDocument(); + }); + + it("renders the provider error message (litellm prefix stripped) for a failing target while others pass", async () => { + const mock = await getMock(); + mock.mockImplementation((_token, modelGroup) => + Promise.resolve( + modelGroup === "claude-sonnet-4" + ? { status: "error", error: "litellm.AuthenticationError: invalid api key" } + : { status: "success" }, + ), + ); + + renderWithProviders(); + + await waitFor(() => expect(screen.getByTestId("test-error-message")).toBeInTheDocument()); + expect(screen.getByTestId("test-error-message")).toHaveTextContent("invalid api key"); + expect(screen.getByTestId("test-error-message")).not.toHaveTextContent("litellm.AuthenticationError"); + expect(screen.getAllByTestId("test-status-success")).toHaveLength(2); + }); + + it("renders a non-litellm error string verbatim", async () => { + const mock = await getMock(); + mock.mockResolvedValue({ status: "error", error: "Connection test failed: 404 Not Found" }); + + renderWithProviders( + , + ); + + await waitFor(() => + expect(screen.getByTestId("test-error-message")).toHaveTextContent("Connection test failed: 404 Not Found"), + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/auto_router_connection_test.tsx b/ui/litellm-dashboard/src/components/add_model/auto_router_connection_test.tsx new file mode 100644 index 00000000000..5badd155da8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/auto_router_connection_test.tsx @@ -0,0 +1,107 @@ +import React from "react"; +import { Typography } from "antd"; +import { CheckCircleTwoTone, CloseCircleTwoTone, LoadingOutlined } from "@ant-design/icons"; +import { testModelGroupConnection, ModelGroupConnectionResult } from "../networking"; +import { AutoRouterTestTarget } from "./build_auto_router_test_targets"; + +const { Text } = Typography; + +interface AutoRouterConnectionTestProps { + accessToken: string; + targets: AutoRouterTestTarget[]; + onTestComplete?: () => void; +} + +type TargetResult = { status: "pending" } | ModelGroupConnectionResult; + +const cleanErrorMessage = (error: string): string => { + const mainError = error.split("stack trace:")[0].trim(); + return mainError.replace(/^litellm\.(.*?)Error: /, ""); +}; + +const AutoRouterConnectionTest: React.FC = ({ + accessToken, + targets, + onTestComplete, +}) => { + const [results, setResults] = React.useState(() => targets.map(() => ({ status: "pending" }))); + + React.useEffect(() => { + let cancelled = false; + const run = async () => { + await Promise.all( + targets.map(async (target, index) => { + const result = await testModelGroupConnection(accessToken, target.modelGroup, target.mode); + if (cancelled) return; + const cleaned: TargetResult = + result.status === "error" ? { status: "error", error: cleanErrorMessage(result.error) } : result; + setResults((prev) => prev.map((r, i) => (i === index ? cleaned : r))); + }), + ); + if (!cancelled && onTestComplete) onTestComplete(); + }; + run(); + return () => { + cancelled = true; + }; + // eslint-disable-next-line react-hooks/exhaustive-deps -- probes run once per mount; the parent remounts via `key` to start a fresh test, and re-running on prop identity changes would refire paid requests + }, []); + + if (targets.length === 0) { + return No complexity tiers are configured yet, so there is nothing to test.; + } + + return ( +
+ + Each configured tier routes to a saved model group. Test Connection sends a minimal request through the proxy to + each one, exactly as the auto router would. + + {targets.map((target, index) => { + const result = results[index] ?? { status: "pending" }; + return ( +
+
+ {result.status === "pending" && } + {result.status === "success" && ( + + )} + {result.status === "error" && ( + + )} +
+
+ {target.labels.join(", ")}{" "} + + {"->"} {target.modelGroup} + {target.mode === "embedding" ? " (embedding)" : ""} + + {result.status === "error" && ( + + {result.error} + + )} +
+
+ ); + })} +
+ ); +}; + +export default AutoRouterConnectionTest; diff --git a/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.test.ts b/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.test.ts new file mode 100644 index 00000000000..f8ab4bab903 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.test.ts @@ -0,0 +1,80 @@ +import { buildAutoRouterTestTargets } from "./build_auto_router_test_targets"; + +const tiers = { + SIMPLE: ["gpt-4o-mini"], + MEDIUM: ["claude-sonnet-4"], + COMPLEX: ["claude-sonnet-4"], + REASONING: ["o3"], +}; + +describe("buildAutoRouterTestTargets", () => { + it("dedups tiers that share a model group into one chat target carrying both labels", () => { + const targets = buildAutoRouterTestTargets({ tiers, semanticMatchingEnabled: false, embeddingModel: undefined }); + expect(targets).toEqual([ + { labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }, + { labels: ["MEDIUM", "COMPLEX"], modelGroup: "claude-sonnet-4", mode: "chat" }, + { labels: ["REASONING"], modelGroup: "o3", mode: "chat" }, + ]); + }); + + it("emits a target per model when a tier has more than one, and dedups across tiers", () => { + const targets = buildAutoRouterTestTargets({ + tiers: { SIMPLE: ["gpt-4o-mini", "claude-sonnet-4"], MEDIUM: ["claude-sonnet-4"], COMPLEX: [], REASONING: [] }, + semanticMatchingEnabled: false, + embeddingModel: undefined, + }); + expect(targets).toEqual([ + { labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }, + { labels: ["SIMPLE", "MEDIUM"], modelGroup: "claude-sonnet-4", mode: "chat" }, + ]); + }); + + it("drops empty/whitespace tiers", () => { + const targets = buildAutoRouterTestTargets({ + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [" "], REASONING: [] }, + semanticMatchingEnabled: false, + embeddingModel: undefined, + }); + expect(targets).toEqual([{ labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }]); + }); + + it("returns [] when no tier is configured", () => { + expect( + buildAutoRouterTestTargets({ + tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, + semanticMatchingEnabled: false, + embeddingModel: undefined, + }), + ).toEqual([]); + }); + + it("appends an embedding target only when semantic matching is on and a model is set", () => { + const targets = buildAutoRouterTestTargets({ + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] }, + semanticMatchingEnabled: true, + embeddingModel: "voyage-3-5", + }); + expect(targets).toEqual([ + { labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }, + { labels: ["Embedding"], modelGroup: "voyage-3-5", mode: "embedding" }, + ]); + }); + + it("omits the embedding target when semantic matching is on but no model is chosen", () => { + const targets = buildAutoRouterTestTargets({ + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] }, + semanticMatchingEnabled: true, + embeddingModel: undefined, + }); + expect(targets).toEqual([{ labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }]); + }); + + it("omits the embedding target when a model is set but semantic matching is off", () => { + const targets = buildAutoRouterTestTargets({ + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] }, + semanticMatchingEnabled: false, + embeddingModel: "voyage-3-5", + }); + expect(targets).toEqual([{ labels: ["SIMPLE"], modelGroup: "gpt-4o-mini", mode: "chat" }]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.ts b/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.ts new file mode 100644 index 00000000000..708a25c16f0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/build_auto_router_test_targets.ts @@ -0,0 +1,51 @@ +import { ComplexityTiers } from "./ComplexityRouterConfig"; + +export type AutoRouterTestMode = "chat" | "embedding"; + +export interface AutoRouterTestTarget { + labels: string[]; + modelGroup: string; + mode: AutoRouterTestMode; +} + +export interface BuildAutoRouterTestTargetsParams { + tiers: ComplexityTiers; + semanticMatchingEnabled: boolean; + embeddingModel: string | undefined; +} + +// Keys drive iteration order; `satisfies Record` makes it a +// compile error to add a tier to ComplexityTiers without listing it here (and vice versa). +const TIER_ORDER = Object.keys({ + SIMPLE: null, + MEDIUM: null, + COMPLEX: null, + REASONING: null, +} satisfies Record) as (keyof ComplexityTiers)[]; + +export const buildAutoRouterTestTargets = ({ + tiers, + semanticMatchingEnabled, + embeddingModel, +}: BuildAutoRouterTestTargetsParams): AutoRouterTestTarget[] => { + const groupedByModel = TIER_ORDER.reduce>((acc, tier) => { + return (tiers[tier] ?? []).reduce((tierAcc, rawModel) => { + const modelGroup = rawModel?.trim(); + if (!modelGroup) return tierAcc; + return { ...tierAcc, [modelGroup]: [...(tierAcc[modelGroup] ?? []), tier] }; + }, acc); + }, {}); + + const tierTargets: AutoRouterTestTarget[] = Object.entries(groupedByModel).map(([modelGroup, labels]) => ({ + labels, + modelGroup, + mode: "chat" as const, + })); + + const embeddingTarget: AutoRouterTestTarget[] = + semanticMatchingEnabled && embeddingModel?.trim() + ? [{ labels: ["Embedding"], modelGroup: embeddingModel.trim(), mode: "embedding" as const }] + : []; + + return [...tierTargets, ...embeddingTarget]; +}; diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 3c252646b57..85a15ffad45 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -1,14 +1,15 @@ import { buildComplexityRouterConfig, + getMissingTiersError, getSemanticConfigError, BuildComplexityRouterConfigParams, } from "./build_complexity_router_config"; const tiers = { - SIMPLE: "gpt-4o-mini", - MEDIUM: "gpt-4o", - COMPLEX: "claude-sonnet-4", - REASONING: "o1-preview", + SIMPLE: ["gpt-4o-mini"], + MEDIUM: ["gpt-4o"], + COMPLEX: ["claude-sonnet-4"], + REASONING: ["o1-preview"], }; const baseParams: BuildComplexityRouterConfigParams = { @@ -20,6 +21,10 @@ const baseParams: BuildComplexityRouterConfigParams = { semanticMatchingEnabled: false, embeddingModel: undefined, matchThreshold: 0.5, + adaptive: false, + adaptiveWeights: { quality: 0.3, cost: 0.7 }, + tierDistancePenalty: 0.5, + adaptiveEligible: "all", }; describe("buildComplexityRouterConfig", () => { @@ -28,6 +33,14 @@ describe("buildComplexityRouterConfig", () => { expect(config).toEqual({ tiers, classifier_type: "heuristic" }); }); + it("passes through a tier configured with more than one model as a pool", () => { + const config = buildComplexityRouterConfig({ + ...baseParams, + tiers: { ...tiers, SIMPLE: ["gpt-4o-mini", "gpt-4o", "claude-haiku-4-5"] }, + }); + expect(config.tiers.SIMPLE).toEqual(["gpt-4o-mini", "gpt-4o", "claude-haiku-4-5"]); + }); + it("includes classifier_llm_config only when classifier_type is llm", () => { const config = buildComplexityRouterConfig({ ...baseParams, @@ -122,6 +135,76 @@ describe("buildComplexityRouterConfig", () => { const config = buildComplexityRouterConfig(params); expect(config.keyword_tier_rules).toBeUndefined(); }); + + it("omits adaptive fields when adaptive is disabled even if weights linger in state", () => { + const config = buildComplexityRouterConfig({ + ...baseParams, + adaptive: false, + adaptiveWeights: { quality: 0.9, cost: 0.1 }, + tierDistancePenalty: 2, + adaptiveEligible: "classified_tier", + }); + expect(config.adaptive).toBeUndefined(); + expect(config.adaptive_weights).toBeUndefined(); + expect(config.tier_distance_penalty).toBeUndefined(); + expect(config.adaptive_eligible).toBeUndefined(); + }); + + it("includes tier_distance_penalty when adaptive is enabled with eligible='all'", () => { + const config = buildComplexityRouterConfig({ + ...baseParams, + adaptive: true, + adaptiveWeights: { quality: 0.6, cost: 0.4 }, + tierDistancePenalty: 0.75, + adaptiveEligible: "all", + }); + expect(config.adaptive).toBe(true); + expect(config.adaptive_weights).toEqual({ quality: 0.6, cost: 0.4 }); + expect(config.tier_distance_penalty).toBe(0.75); + expect(config.adaptive_eligible).toBe("all"); + }); + + it("omits tier_distance_penalty when eligible='classified_tier', since the penalty doesn't apply there", () => { + const config = buildComplexityRouterConfig({ + ...baseParams, + adaptive: true, + adaptiveWeights: { quality: 0.6, cost: 0.4 }, + tierDistancePenalty: 0.75, + adaptiveEligible: "classified_tier", + }); + expect(config.adaptive).toBe(true); + expect(config.adaptive_eligible).toBe("classified_tier"); + expect(config.tier_distance_penalty).toBeUndefined(); + }); +}); + +describe("getMissingTiersError", () => { + it("returns null when all four tiers have a model", () => { + expect(getMissingTiersError(tiers)).toBeNull(); + }); + + it("names the specific missing tier when only one is blank", () => { + expect(getMissingTiersError({ ...tiers, REASONING: [] })).toBe( + "Select a model for the following tier(s): REASONING", + ); + }); + + it("names multiple missing tiers in SIMPLE/MEDIUM/COMPLEX/REASONING order", () => { + expect(getMissingTiersError({ ...tiers, SIMPLE: [], REASONING: [] })).toBe( + "Select a model for the following tier(s): SIMPLE, REASONING", + ); + }); + + it("names all four tiers when none are filled", () => { + const noTiers = { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }; + expect(getMissingTiersError(noTiers)).toBe( + "Select a model for the following tier(s): SIMPLE, MEDIUM, COMPLEX, REASONING", + ); + }); + + it("treats a tier with more than one model as filled", () => { + expect(getMissingTiersError({ ...tiers, SIMPLE: ["gpt-4o-mini", "gpt-4o"] })).toBeNull(); + }); }); describe("getSemanticConfigError", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 82ea4f8c12f..3c3f21163b3 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -1,12 +1,11 @@ import { KeywordTierRule } from "./KeywordTierRules"; -import { ClassifierLLMConfig, ClassifierType } from "./ComplexityRouterConfig"; - -export interface ComplexityTiers { - SIMPLE: string; - MEDIUM: string; - COMPLEX: string; - REASONING: string; -} +import { + AdaptiveEligible, + AdaptiveRouterWeights, + ClassifierLLMConfig, + ClassifierType, + ComplexityTiers, +} from "./ComplexityRouterConfig"; export interface BuildComplexityRouterConfigParams { tiers: ComplexityTiers; @@ -17,6 +16,10 @@ export interface BuildComplexityRouterConfigParams { semanticMatchingEnabled: boolean; embeddingModel: string | undefined; matchThreshold: number; + adaptive: boolean; + adaptiveWeights: AdaptiveRouterWeights; + tierDistancePenalty: number; + adaptiveEligible: AdaptiveEligible; } export interface ComplexityRouterConfigPayload { @@ -28,8 +31,20 @@ export interface ComplexityRouterConfigPayload { semantic_keyword_matching?: boolean; embedding_model?: string; match_threshold?: number; + adaptive?: boolean; + adaptive_weights?: AdaptiveRouterWeights; + tier_distance_penalty?: number; + adaptive_eligible?: AdaptiveEligible; } +const TIER_KEYS: Array = ["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"]; + +export const getMissingTiersError = (tiers: ComplexityTiers): string | null => { + const missing = TIER_KEYS.filter((tier) => tiers[tier].length === 0); + if (missing.length === 0) return null; + return `Select a model for the following tier(s): ${missing.join(", ")}`; +}; + export const getSemanticConfigError = ({ semanticMatchingEnabled, embeddingModel, @@ -54,6 +69,10 @@ export const buildComplexityRouterConfig = ({ semanticMatchingEnabled, embeddingModel, matchThreshold, + adaptive, + adaptiveWeights, + tierDistancePenalty, + adaptiveEligible, }: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => { // Trim keywords and drop empty ones; drop any rule left with no keywords. Clicking // "Add keyword rule" seeds a rule with an empty keywords list, so without this an @@ -74,5 +93,11 @@ export const buildComplexityRouterConfig = ({ embedding_model: embeddingModel, match_threshold: matchThreshold, }), + ...(adaptive && { + adaptive: true, + adaptive_weights: adaptiveWeights, + ...(adaptiveEligible === "all" && { tier_distance_penalty: tierDistancePenalty }), + adaptive_eligible: adaptiveEligible, + }), }; }; diff --git a/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.test.tsx b/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.test.tsx index e96936e179a..0fc9d9fd6fb 100644 --- a/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.test.tsx @@ -56,4 +56,23 @@ describe("FilterInput", () => { expect(input.value).toBe("a"); }); + + it("should not call onChange when unmounted mid-debounce", () => { + const onChange = vi.fn(); + const { unmount } = render(); + + const input = screen.getByPlaceholderText("Search..."); + + act(() => { + fireEvent.change(input, { target: { value: "test" } }); + }); + + unmount(); + + act(() => { + vi.advanceTimersByTime(300); + }); + + expect(onChange).not.toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.tsx b/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.tsx index 3590619ba25..abb6d591386 100644 --- a/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.tsx +++ b/ui/litellm-dashboard/src/components/common_components/Filters/FilterInput.tsx @@ -1,8 +1,9 @@ import { cx } from "@/lib/cva.config"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { Input } from "antd"; -import debounce from "lodash/debounce"; import { LucideIcon } from "lucide-react"; -import React, { useCallback, useEffect, useMemo, useState } from "react"; +import React, { useEffect, useState } from "react"; interface FilterInputProps { placeholder?: string; @@ -13,8 +14,6 @@ interface FilterInputProps { style?: React.CSSProperties; } -const DEBOUNCE_DELAY = 300; - export const FilterInput: React.FC = ({ placeholder, value, onChange, icon: Icon, className }) => { const [localValue, setLocalValue] = useState(value); @@ -22,22 +21,13 @@ export const FilterInput: React.FC = ({ placeholder, value, on setLocalValue(value); }, [value]); - const debouncedOnChange = useMemo(() => debounce((val: string) => onChange(val), DEBOUNCE_DELAY), [onChange]); + const debouncedOnChange = useDebouncedCallback((val: string) => onChange(val), { wait: DEBOUNCE_WAIT_MS }); - useEffect(() => { - return () => { - debouncedOnChange.cancel(); - }; - }, [debouncedOnChange]); - - const handleChange = useCallback( - (e: React.ChangeEvent) => { - const newValue = e.target.value; - setLocalValue(newValue); - debouncedOnChange(newValue); - }, - [debouncedOnChange], - ); + const handleChange = (e: React.ChangeEvent) => { + const newValue = e.target.value; + setLocalValue(newValue); + debouncedOnChange(newValue); + }; return ( ({ + fetchAvailableModels: vi.fn().mockResolvedValue([]), +})); + +const openCustomModelInput = () => { + const selector = document.querySelector(".ant-select-selector"); + expect(selector).toBeTruthy(); + act(() => { + fireEvent.mouseDown(selector!); + }); + act(() => { + fireEvent.click(screen.getByText("Enter custom model")); + }); + return screen.getByPlaceholderText("Enter custom model name"); +}; + +describe("ModelSelector custom model debounce", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + act(() => { + vi.runOnlyPendingTimers(); + }); + vi.useRealTimers(); + }); + + it("does not call onChange before the debounce wait elapses", () => { + const onChange = vi.fn(); + render(); + + const input = openCustomModelInput(); + + act(() => { + fireEvent.change(input, { target: { value: "gpt-4o" } }); + }); + + expect(onChange).not.toHaveBeenCalled(); + + act(() => { + vi.advanceTimersByTime(499); + }); + + expect(onChange).not.toHaveBeenCalled(); + }); + + it("calls onChange exactly once with the last typed value after the wait", () => { + const onChange = vi.fn(); + render(); + + const input = openCustomModelInput(); + + act(() => { + fireEvent.change(input, { target: { value: "g" } }); + fireEvent.change(input, { target: { value: "gp" } }); + fireEvent.change(input, { target: { value: "gpt-5.2" } }); + }); + + expect(onChange).not.toHaveBeenCalled(); + + act(() => { + vi.advanceTimersByTime(500); + }); + + expect(onChange).toHaveBeenCalledTimes(1); + expect(onChange).toHaveBeenCalledWith("gpt-5.2"); + }); + + it("does not call onChange when unmounted mid-wait", () => { + const onChange = vi.fn(); + const { unmount } = render(); + + const input = openCustomModelInput(); + + act(() => { + fireEvent.change(input, { target: { value: "gpt-4o" } }); + }); + + unmount(); + + act(() => { + vi.advanceTimersByTime(500); + }); + + expect(onChange).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx index 1121647c041..f2621cd1acb 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx @@ -1,9 +1,12 @@ -import React, { useState, useEffect, useRef } from "react"; +import React, { useState, useEffect } from "react"; import { TextInput, Text } from "@tremor/react"; import { Select } from "antd"; import { RobotOutlined } from "@ant-design/icons"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; +const MODEL_SELECT_DEBOUNCE_MS = 500; + interface ModelSelectorProps { accessToken: string; value?: string; @@ -30,7 +33,6 @@ const ModelSelector: React.FC = ({ const [selectedModel, setSelectedModel] = useState(value); const [showCustomModelInput, setShowCustomModelInput] = useState(false); const [modelInfo, setModelInfo] = useState([]); - const customModelTimeout = useRef(null); useEffect(() => { setSelectedModel(value); @@ -67,19 +69,13 @@ const ModelSelector: React.FC = ({ } }; - const handleCustomModelChange = (value: string) => { - // Using setTimeout to create a simple debounce effect - if (customModelTimeout.current) { - clearTimeout(customModelTimeout.current); - } - - customModelTimeout.current = setTimeout(() => { + const debouncedSelect = useDebouncedCallback( + (value: string) => { setSelectedModel(value); - if (onChange) { - onChange(value); - } - }, 500); // 500ms delay after typing stops - }; + onChange?.(value); + }, + { wait: MODEL_SELECT_DEBOUNCE_MS }, + ); return (
@@ -109,7 +105,7 @@ const ModelSelector: React.FC = ({ )} diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx new file mode 100644 index 00000000000..a70b7602e5b --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx @@ -0,0 +1,97 @@ +import { act, fireEvent, render, screen } from "@testing-library/react"; +import type { ReactNode } from "react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import type { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm"; +import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "./RouterSettingsAccordion"; + +vi.mock("../networking", () => ({ + getRouterSettingsCall: vi.fn().mockResolvedValue({}), +})); + +vi.mock("@/components/llm_calls/fetch_models", () => ({ + fetchAvailableModels: vi.fn().mockResolvedValue([]), +})); + +vi.mock("../Settings/RouterSettings/Fallbacks/FallbackSelectionForm", () => ({ + FallbackSelectionForm: () => null, +})); + +vi.mock("@tremor/react", () => ({ + TabGroup: ({ children }: { children: ReactNode }) =>
{children}
, + TabList: ({ children }: { children: ReactNode }) =>
{children}
, + Tab: ({ children }: { children: ReactNode }) =>
{children}
, + TabPanels: ({ children }: { children: ReactNode }) =>
{children}
, + TabPanel: ({ children }: { children: ReactNode }) =>
{children}
, +})); + +vi.mock("../router_settings/RouterSettingsForm", () => ({ + default: ({ + value, + onChange, + }: { + value: RouterSettingsFormValue; + onChange: (value: RouterSettingsFormValue) => void; + }) => ( +
+ + +
+ ), +})); + +describe("RouterSettingsAccordion", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + }); + + const flushInitialPropagation = async (onChange: ReturnType) => { + await act(async () => { + vi.advanceTimersByTime(100); + }); + onChange.mockClear(); + }; + + it("debounces propagation and calls onChange once with the last value", async () => { + const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>(); + render(); + await flushInitialPropagation(onChange); + + fireEvent.click(screen.getByText("set-least-busy")); + act(() => { + vi.advanceTimersByTime(50); + }); + fireEvent.click(screen.getByText("set-usage-based")); + + expect(onChange).not.toHaveBeenCalled(); + + act(() => { + vi.advanceTimersByTime(99); + }); + expect(onChange).not.toHaveBeenCalled(); + + act(() => { + vi.advanceTimersByTime(1); + }); + expect(onChange).toHaveBeenCalledTimes(1); + expect(onChange.mock.calls[0][0].router_settings.routing_strategy).toBe("usage-based-routing"); + }); + + it("does not call onChange when unmounted mid-wait", async () => { + const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>(); + const { unmount } = render(); + await flushInitialPropagation(onChange); + + fireEvent.click(screen.getByText("set-least-busy")); + unmount(); + + act(() => { + vi.advanceTimersByTime(500); + }); + expect(onChange).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx index 0aa274b5749..08b917e302f 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx @@ -1,5 +1,6 @@ import React, { useEffect, useState, useImperativeHandle, forwardRef, useRef } from "react"; import { TabPanel, TabPanels, TabGroup, TabList, Tab } from "@tremor/react"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { getRouterSettingsCall } from "../networking"; import RouterSettingsForm, { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm"; import { Fallbacks } from "../Settings/RouterSettings/Fallbacks/AddFallbacks"; @@ -35,6 +36,8 @@ export interface RouterSettingsAccordionRef { getValue: () => RouterSettingsAccordionValue; } +const PROPAGATE_WAIT_MS = 100; + const RouterSettingsAccordion = forwardRef( ({ accessToken, value, onChange, modelData }, ref) => { const [formValue, setFormValue] = useState({ @@ -304,21 +307,26 @@ const RouterSettingsAccordion = forwardRef { - if (!onChange) { - return; - } - - const timeoutId = setTimeout(() => { + const debouncedPropagate = useDebouncedCallback( + () => { + if (!onChange) { + return; + } isInternalUpdateRef.current = true; const finalRouterSettings = buildRouterSettings(); onChange({ router_settings: finalRouterSettings, }); - }, 100); + }, + { wait: PROPAGATE_WAIT_MS }, + ); - return () => clearTimeout(timeoutId); + // Update parent when form values change (with debounce to avoid infinite loops) + useEffect(() => { + if (!onChange) { + return; + } + debouncedPropagate(); // eslint-disable-next-line react-hooks/exhaustive-deps }, [formValue, fallbacks]); 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/common_components/user_search_modal.test.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx new file mode 100644 index 00000000000..72b0e10e5d6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx @@ -0,0 +1,68 @@ +import { act, fireEvent, render, screen, within } from "@testing-library/react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import UserSearchModal from "./user_search_modal"; +import { userFilterUICall } from "@/components/networking"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; + +vi.mock("@/components/networking", () => ({ + userFilterUICall: vi.fn().mockResolvedValue([]), +})); + +const renderModal = () => + render(); + +const getEmailSearchInput = () => within(screen.getByTestId("member-email-search")).getByRole("combobox"); + +describe("UserSearchModal", () => { + beforeEach(() => { + vi.useFakeTimers(); + vi.mocked(userFilterUICall).mockClear(); + }); + + afterEach(() => { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + }); + + it("debounces the user search and fires exactly once with the last typed value", async () => { + renderModal(); + const input = getEmailSearchInput(); + + act(() => { + fireEvent.change(input, { target: { value: "a" } }); + fireEvent.change(input, { target: { value: "ab" } }); + fireEvent.change(input, { target: { value: "abc" } }); + }); + + act(() => { + vi.advanceTimersByTime(DEBOUNCE_WAIT_MS - 1); + }); + expect(userFilterUICall).not.toHaveBeenCalled(); + + await act(async () => { + vi.advanceTimersByTime(1); + await Promise.resolve(); + }); + + expect(userFilterUICall).toHaveBeenCalledTimes(1); + const params = vi.mocked(userFilterUICall).mock.calls[0][1]; + expect(params.get("user_email")).toBe("abc"); + }); + + it("does not fire the search when unmounted mid-wait", () => { + const { unmount } = renderModal(); + const input = getEmailSearchInput(); + + act(() => { + fireEvent.change(input, { target: { value: "abc" } }); + }); + + unmount(); + + act(() => { + vi.advanceTimersByTime(DEBOUNCE_WAIT_MS * 2); + }); + + expect(userFilterUICall).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx index 866d7cbec7f..fafafd8e5d5 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx @@ -1,8 +1,9 @@ -import { useState, useCallback } from "react"; +import { useState } from "react"; import { Modal, Form, Button, Select, Tooltip } from "antd"; import { UserAddOutlined } from "@ant-design/icons"; -import debounce from "lodash/debounce"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { userFilterUICall } from "@/components/networking"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; interface User { user_id: string; user_email: string; @@ -93,9 +94,9 @@ const UserSearchModal: React.FC = ({ } }; - const debouncedSearch = useCallback( - debounce((text: string, fieldName: "user_email" | "user_id") => fetchUsers(text, fieldName), 300), - [], + const debouncedSearch = useDebouncedCallback( + (text: string, fieldName: "user_email" | "user_id") => fetchUsers(text, fieldName), + { wait: DEBOUNCE_WAIT_MS }, ); const handleSearch = (value: string, fieldName: "user_email" | "user_id"): void => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts new file mode 100644 index 00000000000..cd8093928d5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts @@ -0,0 +1,104 @@ +import { buildUpdatedComplexityRouterConfig } from "./edit_auto_router_modal"; + +const storedConfigValue = { + tiers: { + SIMPLE: "old-simple", + MEDIUM: "old-medium", + COMPLEX: "old-complex", + REASONING: "old-reasoning", + }, + classifier_type: "llm", + classifier_llm_config: { model: "old-classifier", timeout_ms: 1200 }, + custom_technical_keywords: ["kafka", "terraform"], + keyword_tier_rules: [{ keywords: ["invoice", "refund"], tier: "MEDIUM" }], + semantic_keyword_matching: true, + embedding_model: "voyage-4-large", + match_threshold: 0.65, + adaptive: true, + adaptive_weights: { quality: 0.3, cost: 0.7 }, + tier_distance_penalty: 0.8, + adaptive_eligible: "all", +}; + +const storedConfig = JSON.stringify(storedConfigValue); + +const tiers = { + SIMPLE: ["gpt-4o-mini"], + MEDIUM: ["gpt-4o-mini"], + COMPLEX: ["anthropic-sonnet-4-5"], + REASONING: ["anthropic-sonnet-4-5"], +}; + +const classifiedTierValue = { + tiers, + classifier_type: "heuristic" as const, + adaptive: true, + adaptive_weights: { quality: 0.4, cost: 0.6 }, + tier_distance_penalty: 0.8, + adaptive_eligible: "classified_tier" as const, +}; + +const expectedClassifiedTierConfig = { + tiers, + classifier_type: "heuristic", + custom_technical_keywords: ["kafka", "terraform"], + keyword_tier_rules: [{ keywords: ["invoice", "refund"], tier: "MEDIUM" }], + semantic_keyword_matching: true, + embedding_model: "voyage-4-large", + match_threshold: 0.65, + adaptive: true, + adaptive_weights: { quality: 0.4, cost: 0.6 }, + adaptive_eligible: "classified_tier", +}; + +const adaptiveDisabledValue = { + tiers, + classifier_type: "heuristic" as const, + adaptive: false, +}; + +const expectedAdaptiveDisabledConfig = { + tiers, + classifier_type: "heuristic", + custom_technical_keywords: ["kafka", "terraform"], + keyword_tier_rules: [{ keywords: ["invoice", "refund"], tier: "MEDIUM" }], + semantic_keyword_matching: true, + embedding_model: "voyage-4-large", + match_threshold: 0.65, +}; + +describe("buildUpdatedComplexityRouterConfig", () => { + it("preserves unrelated options and omits the penalty for classified-tier routing", () => { + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, classifiedTierValue); + + expect(updatedConfig).toEqual(expectedClassifiedTierConfig); + }); + + it("removes managed adaptive and classifier fields when they are disabled", () => { + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, adaptiveDisabledValue); + + expect(updatedConfig).toEqual(expectedAdaptiveDisabledConfig); + }); + + it("updates custom technical keywords when they are edited", () => { + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, classifiedTierValue, ["postgres"]); + + expect(updatedConfig.custom_technical_keywords).toEqual(["postgres"]); + }); + + it("removes custom technical keywords when they are cleared", () => { + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, classifiedTierValue, []); + + expect(updatedConfig.custom_technical_keywords).toBeUndefined(); + }); + + it("preserves a tier configured with more than one model as a pool", () => { + const multiModelValue = { + ...classifiedTierValue, + tiers: { ...tiers, SIMPLE: ["gpt-4o-mini", "claude-haiku-4-5"] }, + }; + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, multiModelValue); + + expect(updatedConfig.tiers).toMatchObject({ SIMPLE: ["gpt-4o-mini", "claude-haiku-4-5"] }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 9c97809bfde..ec54c9b7bad 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -4,13 +4,23 @@ import { Text, TextInput } from "@tremor/react"; import { modelAvailableCall, modelPatchUpdateCall } from "../networking"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import RouterConfigBuilder from "../add_model/RouterConfigBuilder"; -import ComplexityRouterConfig, { ComplexityRouterConfigValue } from "../add_model/ComplexityRouterConfig"; +import ComplexityRouterConfig, { + ComplexityRouterConfigValue, + DEFAULT_ADAPTIVE_WEIGHTS, + DEFAULT_TIER_DISTANCE_PENALTY, +} from "../add_model/ComplexityRouterConfig"; import NotificationsManager from "../molecules/notifications_manager"; const isComplexityRouterModel = (modelData: any): boolean => modelData?.litellm_params?.model?.startsWith("auto_router/complexity_router") || modelData?.litellm_params?.complexity_router_config != null; +const normalizeTierModels = (value: unknown): string[] => { + if (Array.isArray(value)) return value; + if (typeof value === "string" && value) return [value]; + return []; +}; + interface EditAutoRouterModalProps { isVisible: boolean; onCancel: () => void; @@ -20,6 +30,57 @@ interface EditAutoRouterModalProps { userRole: string; } +const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ + "tiers", + "classifier_type", + "classifier_llm_config", + "adaptive", + "adaptive_weights", + "tier_distance_penalty", + "adaptive_eligible", +]); + +const toRecord = (value: unknown): Record => { + const parsed: unknown = typeof value === "string" ? JSON.parse(value) : value; + return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed) + ? (parsed as Record) + : {}; +}; + +export const buildUpdatedComplexityRouterConfig = ( + storedConfig: unknown, + value: ComplexityRouterConfigValue, + customTechnicalKeywords?: string[], +): Record => { + const preservedConfig = Object.fromEntries( + Object.entries(toRecord(storedConfig)).filter( + ([key]) => + !MANAGED_COMPLEXITY_ROUTER_KEYS.has(key) && + (customTechnicalKeywords === undefined || key !== "custom_technical_keywords"), + ), + ); + const adaptiveEligible = value.adaptive_eligible ?? "all"; + + return { + ...preservedConfig, + tiers: value.tiers, + classifier_type: value.classifier_type, + ...(value.classifier_type === "llm" ? { classifier_llm_config: value.classifier_llm_config } : {}), + ...(customTechnicalKeywords && + customTechnicalKeywords.length > 0 && { + custom_technical_keywords: customTechnicalKeywords, + }), + ...(value.adaptive && { + adaptive: true, + adaptive_weights: value.adaptive_weights ?? DEFAULT_ADAPTIVE_WEIGHTS, + ...(adaptiveEligible === "all" && { + tier_distance_penalty: value.tier_distance_penalty ?? DEFAULT_TIER_DISTANCE_PENALTY, + }), + adaptive_eligible: adaptiveEligible, + }), + }; +}; + const EditAutoRouterModal: React.FC = ({ isVisible, onCancel, @@ -35,8 +96,9 @@ const EditAutoRouterModal: React.FC = ({ const [showCustomDefaultModel, setShowCustomDefaultModel] = useState(false); const [showCustomEmbeddingModel, setShowCustomEmbeddingModel] = useState(false); const [routerConfig, setRouterConfig] = useState(null); + const [customTechnicalKeywords, setCustomTechnicalKeywords] = useState([]); const [complexityRouterConfig, setComplexityRouterConfig] = useState({ - tiers: { SIMPLE: "", MEDIUM: "", COMPLEX: "", REASONING: "" }, + tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic", }); const isComplexityRouter = isComplexityRouterModel(modelData); @@ -85,14 +147,21 @@ const EditAutoRouterModal: React.FC = ({ setComplexityRouterConfig({ tiers: { - SIMPLE: parsedConfig.tiers?.SIMPLE || "", - MEDIUM: parsedConfig.tiers?.MEDIUM || "", - COMPLEX: parsedConfig.tiers?.COMPLEX || "", - REASONING: parsedConfig.tiers?.REASONING || "", + SIMPLE: normalizeTierModels(parsedConfig.tiers?.SIMPLE), + MEDIUM: normalizeTierModels(parsedConfig.tiers?.MEDIUM), + COMPLEX: normalizeTierModels(parsedConfig.tiers?.COMPLEX), + REASONING: normalizeTierModels(parsedConfig.tiers?.REASONING), }, classifier_type: parsedConfig.classifier_type || "heuristic", classifier_llm_config: parsedConfig.classifier_llm_config, + adaptive: parsedConfig.adaptive || false, + adaptive_weights: parsedConfig.adaptive_weights, + tier_distance_penalty: parsedConfig.tier_distance_penalty, + adaptive_eligible: parsedConfig.adaptive_eligible || "all", }); + setCustomTechnicalKeywords( + Array.isArray(parsedConfig.custom_technical_keywords) ? parsedConfig.custom_technical_keywords : [], + ); form.setFieldsValue({ auto_router_name: modelData.model_name, @@ -138,7 +207,7 @@ const EditAutoRouterModal: React.FC = ({ if (isComplexityRouter) { const { tiers, classifier_type, classifier_llm_config } = complexityRouterConfig; - if (Object.values(tiers).filter(Boolean).length === 0) { + if (Object.values(tiers).every((models) => models.length === 0)) { NotificationsManager.fromBackend("Please select at least one model for a complexity tier"); return; } @@ -147,14 +216,14 @@ const EditAutoRouterModal: React.FC = ({ return; } - const defaultModel = tiers.MEDIUM || tiers.SIMPLE || tiers.COMPLEX || tiers.REASONING; + const defaultModel = tiers.MEDIUM[0] || tiers.SIMPLE[0] || tiers.COMPLEX[0] || tiers.REASONING[0]; const updatedLitellmParams = { ...modelData.litellm_params, - complexity_router_config: { - tiers, - classifier_type, - ...(classifier_type === "llm" ? { classifier_llm_config } : {}), - }, + complexity_router_config: buildUpdatedComplexityRouterConfig( + modelData.litellm_params?.complexity_router_config, + complexityRouterConfig, + customTechnicalKeywords, + ), complexity_router_default_model: defaultModel, }; const updatedModelInfo = { @@ -264,6 +333,8 @@ const EditAutoRouterModal: React.FC = ({ onChange={(config) => { setComplexityRouterConfig(config); }} + customTechnicalKeywords={customTechnicalKeywords} + onCustomTechnicalKeywordsChange={setCustomTechnicalKeywords} />
) : ( diff --git a/ui/litellm-dashboard/src/components/key_scope.test.ts b/ui/litellm-dashboard/src/components/key_scope.test.ts new file mode 100644 index 00000000000..7c20a1baece --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_scope.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it } from "vitest"; + +import { deriveKeyModelScope } from "./key_scope"; + +describe("deriveKeyModelScope", () => { + it("treats unrestricted keys (null/empty allowed_routes) as full model access", () => { + expect(deriveKeyModelScope(null)).toEqual({ hasModelAccess: true, label: null }); + expect(deriveKeyModelScope(undefined)).toEqual({ hasModelAccess: true, label: null }); + expect(deriveKeyModelScope([])).toEqual({ hasModelAccess: true, label: null }); + }); + + it("classifies SCIM keys as no model access", () => { + expect(deriveKeyModelScope(["/scim/*"])).toEqual({ hasModelAccess: false, label: "SCIM" }); + expect(deriveKeyModelScope(["/scim/v2/Users", "/scim/v2/Groups"])).toEqual({ + hasModelAccess: false, + label: "SCIM", + }); + }); + + it("classifies management-only keys as no model access", () => { + expect(deriveKeyModelScope(["management_routes"])).toEqual({ hasModelAccess: false, label: "Management" }); + }); + + it("classifies read-only keys as no model access", () => { + expect(deriveKeyModelScope(["info_routes"])).toEqual({ hasModelAccess: false, label: "Read-only" }); + }); + + it("leaves LLM-API and custom scopes with model access (default rendering)", () => { + expect(deriveKeyModelScope(["llm_api_routes"])).toEqual({ hasModelAccess: true, label: null }); + expect(deriveKeyModelScope(["/chat/completions"])).toEqual({ hasModelAccess: true, label: null }); + expect(deriveKeyModelScope(["management_routes", "llm_api_routes"])).toEqual({ + hasModelAccess: true, + label: null, + }); + }); + + it("prefers a persisted key_type over allowed_routes for the no-inference buckets", () => { + expect(deriveKeyModelScope([], "management")).toEqual({ hasModelAccess: false, label: "Management" }); + expect(deriveKeyModelScope([], "read_only")).toEqual({ hasModelAccess: false, label: "Read-only" }); + expect(deriveKeyModelScope(["some_future_mgmt_preset"], "management")).toEqual({ + hasModelAccess: false, + label: "Management", + }); + }); + + it("falls back to allowed_routes for null/default/llm_api key_type", () => { + expect(deriveKeyModelScope(["/scim/*"], null)).toEqual({ hasModelAccess: false, label: "SCIM" }); + expect(deriveKeyModelScope(["/scim/*"], "default")).toEqual({ hasModelAccess: false, label: "SCIM" }); + expect(deriveKeyModelScope([], "default")).toEqual({ hasModelAccess: true, label: null }); + expect(deriveKeyModelScope([], "llm_api")).toEqual({ hasModelAccess: true, label: null }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/key_scope.ts b/ui/litellm-dashboard/src/components/key_scope.ts new file mode 100644 index 00000000000..01dc595b4f1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_scope.ts @@ -0,0 +1,49 @@ +export interface KeyModelScope { + hasModelAccess: boolean; + label: string | null; +} + +const MANAGEMENT_ROUTES_PRESET = "management_routes"; +const INFO_ROUTES_PRESET = "info_routes"; +const SCIM_ROUTE_PREFIX = "/scim"; + +const MANAGEMENT_SCOPE: KeyModelScope = { hasModelAccess: false, label: "Management" }; +const READ_ONLY_SCOPE: KeyModelScope = { hasModelAccess: false, label: "Read-only" }; +const SCIM_SCOPE: KeyModelScope = { hasModelAccess: false, label: "SCIM" }; +const FULL_MODEL_ACCESS: KeyModelScope = { hasModelAccess: true, label: null }; + +const isScimRoute = (route: string): boolean => route.startsWith(SCIM_ROUTE_PREFIX); + +const isOnlyPreset = (allowedRoutes: string[], preset: string): boolean => + allowedRoutes.length === 1 && allowedRoutes[0] === preset; + +export const deriveKeyModelScope = ( + allowedRoutes: string[] | null | undefined, + keyType?: string | null, +): KeyModelScope => { + if (keyType === "management") { + return MANAGEMENT_SCOPE; + } + + if (keyType === "read_only") { + return READ_ONLY_SCOPE; + } + + if (!Array.isArray(allowedRoutes) || allowedRoutes.length === 0) { + return FULL_MODEL_ACCESS; + } + + if (allowedRoutes.every(isScimRoute)) { + return SCIM_SCOPE; + } + + if (isOnlyPreset(allowedRoutes, MANAGEMENT_ROUTES_PRESET)) { + return MANAGEMENT_SCOPE; + } + + if (isOnlyPreset(allowedRoutes, INFO_ROUTES_PRESET)) { + return READ_ONLY_SCOPE; + } + + return FULL_MODEL_ACCESS; +}; diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index 285daa8f156..e1c1fcb232c 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -12,8 +12,10 @@ export interface Team { rpm_limit: number | null; organization_id: string; created_at: string; + updated_at?: string | null; keys: KeyResponse[]; keys_count?: number; + members_count?: number; members_with_roles: Member[]; spend: number; access_group_ids?: string[]; @@ -45,6 +47,7 @@ export interface KeyResponse { budget_reset_at: string; allowed_cache_controls: string[]; allowed_routes: string[]; + key_type: string | null; permissions: Record; model_spend: Record; model_max_budget: Record; diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index c24cebe7e7d..92dd3c849ef 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -7,9 +7,9 @@ import { getProxyBaseUrl } from "@/components/networking"; import { useTheme } from "@/contexts/ThemeContext"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; +import { ScrollArea } from "@/components/ui/scroll-area"; import { Sidebar, - SidebarContent, SidebarFooter, SidebarGroup, SidebarGroupLabel, @@ -608,15 +608,17 @@ const Sidebar_: React.FC = ({
- - {visibleGroups.map((group, gi) => ( - - {gi > 0 && } - {group.groupLabel} - {group.items.map((item) => renderItem(item))} - - ))} - + + + {isAdminRole(userRole) && ( 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/molecules/filter.test.tsx b/ui/litellm-dashboard/src/components/molecules/filter.test.tsx index d956cd93168..66c671b12d4 100644 --- a/ui/litellm-dashboard/src/components/molecules/filter.test.tsx +++ b/ui/litellm-dashboard/src/components/molecules/filter.test.tsx @@ -547,6 +547,44 @@ describe("FilterComponent", () => { }); }); + it("cancels a pending debounced search when the component unmounts mid-type", async () => { + const user = userEvent.setup({ delay: null }); + const mockSearchFn = vi.fn().mockResolvedValue([{ label: "Result", value: "result" }]); + + const options: FilterOption[] = [ + { + name: "model", + label: "Model", + isSearchable: true, + searchFn: mockSearchFn, + }, + ]; + + const { unmount } = renderWithProviders( + , + ); + + await user.click(screen.getByRole("button", { name: "Filters" })); + + await waitFor(() => { + expect(mockSearchFn).toHaveBeenCalledWith(""); + }); + + vi.clearAllMocks(); + + const modelLabel = screen.getByText("Model"); + const modelSelect = within(modelLabel.closest("div")!).getByRole("combobox"); + await user.click(modelSelect); + await user.type(modelSelect, "test"); + + expect(mockSearchFn).not.toHaveBeenCalled(); + + unmount(); + + await new Promise((resolve) => setTimeout(resolve, 400)); + expect(mockSearchFn).not.toHaveBeenCalled(); + }); + it("should reset all filter values when reset button is clicked", async () => { const user = userEvent.setup({ delay: null }); renderWithProviders( diff --git a/ui/litellm-dashboard/src/components/molecules/filter.tsx b/ui/litellm-dashboard/src/components/molecules/filter.tsx index 45ad3a6b9ca..8218de41a19 100644 --- a/ui/litellm-dashboard/src/components/molecules/filter.tsx +++ b/ui/litellm-dashboard/src/components/molecules/filter.tsx @@ -1,6 +1,7 @@ +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { FilterIcon } from "@heroicons/react/outline"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { Button, Input, Select } from "antd"; -import debounce from "lodash/debounce"; import React, { useCallback, useEffect, useState } from "react"; export interface FilterOptionCustomComponentProps { @@ -54,8 +55,8 @@ const FilterComponent: React.FC = ({ [key: string]: boolean; }>({}); - const debouncedSearch = useCallback( - debounce(async (value: string, option: FilterOption) => { + const debouncedSearch = useDebouncedCallback( + async (value: string, option: FilterOption) => { if (!option.isSearchable || !option.searchFn) return; setSearchLoadingMap((prev) => ({ ...prev, [option.name]: true })); @@ -68,8 +69,8 @@ const FilterComponent: React.FC = ({ } finally { setSearchLoadingMap((prev) => ({ ...prev, [option.name]: false })); } - }, 300), - [], + }, + { wait: DEBOUNCE_WAIT_MS }, ); // Load initial options for searchable filters diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index 43f85cc8674..e6ee4de2735 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -514,3 +514,83 @@ describe("sessionSpendLogsCall", () => { expect(parsed.searchParams.get("page_size")).toBe("100"); }); }); + +describe("buildModelGroupTestRequest", () => { + it("builds a chat completion request with NO max_tokens (reasoning models 400 on a tiny cap)", () => { + const { path, body } = Networking.buildModelGroupTestRequest("o3", "chat"); + expect(path).toBe("/v1/chat/completions"); + expect(body).toEqual({ model: "o3", messages: [{ role: "user", content: "test from litellm" }] }); + expect(body).not.toHaveProperty("max_tokens"); + expect(body).not.toHaveProperty("max_completion_tokens"); + }); + + it("builds an embeddings request for embedding mode", () => { + const { path, body } = Networking.buildModelGroupTestRequest("text-embedding-3-small", "embedding"); + expect(path).toBe("/v1/embeddings"); + 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 da6bb079876..7e3af46b931 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2315,6 +2315,46 @@ export const testConnectionRequest = async ( } }; +export type ModelGroupConnectionResult = { status: "success" } | { status: "error"; error: string }; + +/** + * Test an existing model group by routing a minimal request through the proxy + * exactly as production would (by public model_group name). Unlike + * /health/test_connection, this needs no litellm_params resolution: the router + * resolves the group, credentials, and provider. Used by the auto-router Test + * Connection to probe each tier's model group and the embedding model. + */ +/** + * Build the minimal request that probes a model group by public name. No + * max_tokens: reasoning models (o1/o3/...) reject a tiny cap with "max_tokens + * reached" because reasoning tokens count against it, which would show a false + * failure for a reachable tier. + */ +export const buildModelGroupTestRequest = ( + modelGroup: string, + mode: "chat" | "embedding", +): { path: string; body: Record } => + mode === "embedding" + ? { path: "/v1/embeddings", body: { model: modelGroup, input: "test from litellm" } } + : { + path: "/v1/chat/completions", + body: { model: modelGroup, messages: [{ role: "user", content: "test from litellm" }] }, + }; + +export const testModelGroupConnection = async ( + accessToken: string, + modelGroup: string, + mode: "chat" | "embedding", +): Promise => { + const { path, body } = buildModelGroupTestRequest(modelGroup, mode); + try { + await apiClient.post(path, { accessToken, body }); + return { status: "success" }; + } catch (error) { + return { status: "error", error: error instanceof Error ? error.message : String(error) }; + } +}; + // ... existing code ... export const keyInfoV1Call = async (accessToken: string, key: string) => { try { @@ -6610,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}`; @@ -6804,7 +6847,11 @@ export const exchangeMcpOAuthToken = async ({ const data = await response.json(); if (!response.ok) { - const errorMessage = deriveErrorMessage(data) || data?.detail || "OAuth token exchange failed"; + const oauthErrorMessage = + typeof data?.error === "string" && typeof data?.error_description === "string" + ? `${data.error}: ${data.error_description}` + : undefined; + const errorMessage = oauthErrorMessage || deriveErrorMessage(data) || data?.detail || "OAuth token exchange failed"; throw new Error(errorMessage); } return data; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx index 0fe8adb70e1..f84b95b8d4d 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx @@ -1,7 +1,8 @@ import { act, fireEvent, within } from "@testing-library/react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; import { Team } from "../key_team_helpers/key_list"; +import { userFilterUICall } from "../networking"; import CreateKey from "./create_key_button"; const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyCreateCall, teamDropdownTeamsRef } = @@ -134,14 +135,16 @@ vi.mock("antd", () => { const Select = ({ children, onChange, + onSearch, options, ...props }: { children?: any; onChange?: (value: string) => void; + onSearch?: (value: string) => void; options?: Array<{ value: string; label: string }>; - }) => - React.createElement( + }) => { + const select = React.createElement( "select", { ...props, @@ -151,6 +154,21 @@ vi.mock("antd", () => { options?.map((opt: any) => React.createElement("option", { key: opt.value, value: opt.value }, opt.label)), ); + if (!onSearch) { + return select; + } + + return React.createElement( + React.Fragment, + null, + React.createElement("input", { + "data-testid": "select-search-input", + onChange: (event: React.ChangeEvent) => onSearch(event.target.value), + }), + select, + ); + }; + Select.Option = ({ children, ...props }: { children?: any }) => React.createElement("option", props, children); const Input = (props: any) => React.createElement("input", props); @@ -641,6 +659,80 @@ describe("CreateKey", () => { }); }); + describe("user search debounce", () => { + const mockUserFilterUICall = vi.mocked(userFilterUICall); + + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + }); + + const renderUserSearch = () => { + const view = renderWithProviders( + , + ); + return { input: screen.getByTestId("select-search-input"), unmount: view.unmount }; + }; + + it("should not fire the search before the wait elapses", () => { + const { input } = renderUserSearch(); + + act(() => { + fireEvent.change(input, { target: { value: "alice" } }); + }); + + expect(mockUserFilterUICall).not.toHaveBeenCalled(); + + act(() => { + vi.advanceTimersByTime(299); + }); + + expect(mockUserFilterUICall).not.toHaveBeenCalled(); + }); + + it("should fire exactly one search carrying the last value after the wait", async () => { + const { input } = renderUserSearch(); + + act(() => { + fireEvent.change(input, { target: { value: "a" } }); + vi.advanceTimersByTime(100); + fireEvent.change(input, { target: { value: "al" } }); + vi.advanceTimersByTime(100); + fireEvent.change(input, { target: { value: "alice" } }); + }); + + expect(mockUserFilterUICall).not.toHaveBeenCalled(); + + await act(async () => { + vi.advanceTimersByTime(300); + }); + + expect(mockUserFilterUICall).toHaveBeenCalledTimes(1); + const params = mockUserFilterUICall.mock.calls[0][1] as URLSearchParams; + expect(params.get("user_email")).toBe("alice"); + }); + + it("should fire nothing when unmounted mid-wait", () => { + const { input, unmount } = renderUserSearch(); + + act(() => { + fireEvent.change(input, { target: { value: "alice" } }); + }); + + unmount(); + + act(() => { + vi.advanceTimersByTime(1000); + }); + + expect(mockUserFilterUICall).not.toHaveBeenCalled(); + }); + }); + describe("tags dropdown", () => { it("should populate tags dropdown with options from useTags hook", async () => { renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index ef2ddab70ed..4bfd869f17a 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -10,8 +10,9 @@ import { InfoCircleOutlined } from "@ant-design/icons"; import { useQueryClient } from "@tanstack/react-query"; import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react"; import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip, Typography } from "antd"; -import debounce from "lodash/debounce"; -import React, { useCallback, useEffect, useState } from "react"; +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; +import React, { useEffect, useState } from "react"; import { rolesWithWriteAccess } from "../../utils/roles"; import AgentSelector from "../agent_management/AgentSelector"; import { mapDisplayToInternalNames } from "../callback_info_helpers"; @@ -688,14 +689,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp } }; - const debouncedSearch = useCallback( - debounce((text: string) => fetchUsers(text), 300), - [accessToken], - ); - - const handleUserSearch = (value: string): void => { - debouncedSearch(value); - }; + const handleUserSearch = useDebouncedCallback((text: string) => fetchUsers(text), { wait: DEBOUNCE_WAIT_MS }); const handleUserSelect = (_value: string, option: UserOption): void => { const selectedUser = option.user; diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index 6f4ad17313b..46811e0106b 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -566,7 +566,7 @@ const Settings: React.FC = ({ accessToken, userRole, userID, } return ( -
+
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..60547415908 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -5,7 +5,7 @@ import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; import { DataTable } from "./DataTable"; -import { DataTableSortHeader } from "./DataTableSortHeader"; +import { DataTableMultiSortHeader, DataTableSortHeader } from "./DataTableSortHeader"; import { DataTableViewOptions } from "./DataTableViewOptions"; interface Person { @@ -55,6 +55,23 @@ const dropdownSortColumns: ColumnDef[] = [ }, ]; +const multiSortColumns: ColumnDef[] = [ + { + id: "spend", + accessorKey: "name", + header: ({ table }) => ( + + ), + cell: ({ row }) => {row.original.name}, + }, +]; + const nameEmailColumns: ColumnDef[] = [ { accessorKey: "name", @@ -165,6 +182,60 @@ describe("DataTable sorting", () => { await user.click(await screen.findByText("Reset")); expect(names()).toEqual(["Charlie", "Alice", "Bob"]); }); + + it("multi-sort header emits the chosen field id (not the column id) as the sort key", async () => { + const user = userEvent.setup(); + const onSortingChange = vi.fn(); + render( + , + ); + + await user.click(screen.getByTestId("sort-trigger-spend")); + await user.click(await screen.findByText("Budget descending")); + expect(onSortingChange).toHaveBeenLastCalledWith([{ id: "max_budget", desc: true }]); + + await user.click(screen.getByTestId("sort-trigger-spend")); + await user.click(await screen.findByText("Spend ascending")); + expect(onSortingChange).toHaveBeenLastCalledWith([{ id: "spend", desc: false }]); + }); + + it("multi-sort header reflects the active field and direction, and Reset clears it", async () => { + const user = userEvent.setup(); + const onSortingChange = vi.fn(); + render( + , + ); + + await user.click(screen.getByTestId("sort-trigger-spend")); + // The header trigger shows the active (descending) indicator while sorted by a field it owns. + expect(screen.getByTestId("sort-trigger-spend").querySelector("[data-sort-indicator='desc']")).not.toBeNull(); + + await user.click(await screen.findByText("Reset")); + expect(onSortingChange).toHaveBeenLastCalledWith([]); + }); +}); + +describe("DataTable layout", () => { + it("stretches the table to fill the container when resizing is on, so hidden columns leave no right-side gap", () => { + const { container } = render(); + + const table = container.querySelector("table"); + expect(table).not.toBeNull(); + // width pins the natural column total (horizontal scroll on overflow); minWidth:100% fills the gap on underflow. + expect(table?.style.minWidth).toBe("100%"); + }); }); describe("DataTable pagination", () => { @@ -287,6 +358,23 @@ describe("DataTable loading", () => { expect(names()).toEqual(["Charlie", "Alice", "Bob"]); }); + it("gives compact skeleton rows the same height as loaded rows so loading does not shrink the table", () => { + const { rerender } = render( + , + ); + const skeletonRow = screen.getAllByTestId("skeleton-row").at(0); + const loadedRowHeight = "h-8"; + expect(skeletonRow?.className).toContain(loadedRowHeight); + + rerender(); + expect(document.querySelector("[data-row-id]")?.className).toContain(loadedRowHeight); + }); + + it("does not force the compact height on default-size skeleton rows", () => { + render(); + expect(screen.getAllByTestId("skeleton-row").at(0)?.className).not.toContain("h-8"); + }); + it("varies skeleton shape and width per column instead of one fixed bar", () => { const columns: ColumnDef[] = [ { accessorKey: "name", header: "Name", meta: { skeleton: "twoLine" }, cell: () => null }, @@ -303,6 +391,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..fa0e672026e 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 ; } @@ -365,7 +389,11 @@ function SkeletonRows({ return ( {rowKeys.map((rowKey) => ( - + {cells.map((column, columnKey) => ( @@ -508,7 +536,7 @@ export function DataTable(props: DataTableProps { if (paginationSlot !== undefined) { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.tsx index 1cf09ce4f47..8988bd6c2d0 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.tsx @@ -1,8 +1,8 @@ "use client"; import { Menu } from "@base-ui/react/menu"; -import type { Column, SortDirection } from "@tanstack/react-table"; -import { ChevronDown, ChevronsUpDown, ChevronUp, X } from "lucide-react"; +import type { Column, SortDirection, Table } from "@tanstack/react-table"; +import { Check, ChevronDown, ChevronsUpDown, ChevronUp, X } from "lucide-react"; import type * as React from "react"; import { cn } from "@/lib/cva.config"; @@ -94,3 +94,101 @@ export function DataTableSortHeader({ ); } + +export interface DataTableSortField { + /** Backend sort column, sent verbatim as the sorting state id (e.g. "spend", "max_budget"). */ + id: string; + label: string; +} + +interface DataTableMultiSortHeaderProps { + table: Table; + fields: DataTableSortField[]; + className?: string; +} + +/** + * Sort header for a column that merges several backend-sortable fields into one cell + * (e.g. a combined Spend / Budget cell). The header label is the field labels joined by " / ", + * with the field currently driving the sort emphasized so the active column reads at a glance + * without opening the menu. The chevron opens a menu offering each field in both directions. + */ +export function DataTableMultiSortHeader({ table, fields, className }: DataTableMultiSortHeaderProps) { + const active = table.getState().sorting[0]; + const activeField = active !== undefined && fields.some((field) => field.id === active.id) ? active : undefined; + const activeDirection: SortDirection = activeField?.desc === true ? "desc" : "asc"; + const sorted: false | SortDirection = activeField === undefined ? false : activeDirection; + + const options = fields.flatMap((field) => [ + { key: `${field.id}-asc`, id: field.id, desc: false, label: `${field.label} ascending`, Icon: ChevronUp }, + { key: `${field.id}-desc`, id: field.id, desc: true, label: `${field.label} descending`, Icon: ChevronDown }, + ]); + + const segmentClass = (isActive: boolean): string => { + if (isActive) return "font-semibold text-foreground"; + if (activeField) return "text-muted-foreground"; + return ""; + }; + + const labelSegments = fields.flatMap((field, index) => { + const isActive = activeField?.id === field.id; + const segment = ( + + {field.label} + + ); + if (index === 0) return [segment]; + return [ + + {" / "} + , + segment, + ]; + }); + + return ( +
+ {labelSegments} + + field.label).join(" or ")}`} + onClick={(event) => event.stopPropagation()} + className={cn( + "inline-flex size-6 items-center justify-center rounded-md hover:bg-muted", + sorted ? "text-primary" : "text-muted-foreground", + )} + > + + + } + /> + + + + {options.map((option) => { + const isActive = activeField?.id === option.id && activeField.desc === option.desc; + return ( + table.setSorting([{ id: option.id, desc: option.desc }])} + > + {option.label} + {isActive && } + + ); + })} + table.setSorting([])}> + Reset + + + + + +
+ ); +} 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/index.ts b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts index c4218f6051a..1ee1eed1258 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/index.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts @@ -5,7 +5,12 @@ export { DataTableFilterDrawer, DataTableFilterField, type FilterDraft } from ". export { DataTablePagination, DEFAULT_PAGE_SIZE_OPTIONS } from "./DataTablePagination"; export { DataTableToolbar } from "./DataTableToolbar"; export { DataTableViewOptions } from "./DataTableViewOptions"; -export { DataTableSortHeader, type DataTableSortVariant } from "./DataTableSortHeader"; +export { + DataTableSortHeader, + DataTableMultiSortHeader, + type DataTableSortVariant, + type DataTableSortField, +} from "./DataTableSortHeader"; export type { DataTablePaginationProps } from "./DataTablePagination"; export type { ColumnPinnedSide, 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..e314e8e8bc2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/PageHeader.tsx @@ -0,0 +1,25 @@ +"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..eff5cfc804d --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx @@ -0,0 +1,78 @@ +"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; +} + +const matchesQuery = (option: SearchSelectOption, query: string): boolean => { + const q = query.trim().toLowerCase(); + if (!q) return true; + return option.label.toLowerCase().includes(q) || (option.sublabel?.toLowerCase().includes(q) ?? false); +}; + +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={matchesQuery} + disabled={disabled} + > + + + {emptyText} + + {(item: SearchSelectOption) => ( + + + {item.label} + {item.sublabel != null && item.sublabel !== "" && ( + {item.sublabel} + )} + + + )} + + + + ); +} diff --git a/ui/litellm-dashboard/src/components/shared/charts/index.ts b/ui/litellm-dashboard/src/components/shared/charts/index.ts index ba0a7544ddb..8383c767064 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/index.ts +++ b/ui/litellm-dashboard/src/components/shared/charts/index.ts @@ -10,3 +10,4 @@ export { } from "./chart_tooltip"; export { CHART_COLOR_HEX, DEFAULT_COLOR_CYCLE, categoryFills, chartColorValue, type ChartColor } from "./colors"; export { DonutChart, type DonutChartProps } from "./donut_chart"; +export { LineChart, type LineChartCurveType, type LineChartProps } from "./line_chart"; diff --git a/ui/litellm-dashboard/src/components/shared/charts/line_chart.test.tsx b/ui/litellm-dashboard/src/components/shared/charts/line_chart.test.tsx new file mode 100644 index 00000000000..9385dc49494 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/charts/line_chart.test.tsx @@ -0,0 +1,117 @@ +import { render, screen } from "@testing-library/react"; +import React from "react"; +import { describe, expect, it } from "vitest"; +import { LineChart } from "./line_chart"; + +const data = [ + { date: "Jun 1", "/chat/completions": 10, "/embeddings": 4 }, + { date: "Jun 2", "/chat/completions": 15, "/embeddings": 6 }, + { date: "Jun 3", "/chat/completions": 12, "/embeddings": 9 }, +]; + +describe("LineChart", () => { + it("renders one line per category with the mapped tremor stroke colors", () => { + const { container } = render( + , + ); + + const curves = Array.from(container.querySelectorAll("path.recharts-line-curve")); + expect(curves).toHaveLength(2); + expect(curves.map((curve) => curve.getAttribute("stroke"))).toEqual([ + "var(--color-blue-500, #3b82f6)", + "var(--color-cyan-500, #06b6d4)", + ]); + }); + + it("falls back to the tremor default color cycle when no colors are passed", () => { + const { container } = render( + , + ); + + const strokes = Array.from(container.querySelectorAll("path.recharts-line-curve")).map((curve) => + curve.getAttribute("stroke"), + ); + expect(strokes).toEqual(["var(--color-blue-500, #3b82f6)", "var(--color-cyan-500, #06b6d4)"]); + }); + + it("applies valueFormatter to the value axis ticks", () => { + render( + `${v} req`} + />, + ); + + expect(screen.getAllByText(/ req$/).length).toBeGreaterThan(0); + }); + + it("renders a legend by default, matching tremor, and hides it when showLegend is false", () => { + const { container, rerender } = render( + , + ); + expect(screen.getByText("/chat/completions")).toBeInTheDocument(); + expect(container.querySelector(".recharts-legend-wrapper")).not.toBeNull(); + + rerender( + , + ); + expect(screen.queryByText("/chat/completions")).not.toBeInTheDocument(); + }); + + it("draws straight segments by default and curved segments for curveType natural", () => { + const { container: linear } = render( + , + ); + const { container: natural } = render( + , + ); + + const linearPath = linear.querySelector("path.recharts-line-curve")?.getAttribute("d") ?? ""; + const naturalPath = natural.querySelector("path.recharts-line-curve")?.getAttribute("d") ?? ""; + expect(linearPath).not.toContain("C"); + expect(naturalPath).toContain("C"); + }); + + it("bridges gaps over null values only when connectNulls is set", () => { + const gappedData = [ + { date: "Jun 1", "/chat/completions": 10 }, + { date: "Jun 2", "/chat/completions": null }, + { date: "Jun 3", "/chat/completions": 12 }, + { date: "Jun 4", "/chat/completions": 15 }, + ]; + + const { container: broken } = render( + , + ); + const { container: bridged } = render( + , + ); + + const brokenPath = broken.querySelector("path.recharts-line-curve")?.getAttribute("d") ?? ""; + const bridgedPath = bridged.querySelector("path.recharts-line-curve")?.getAttribute("d") ?? ""; + expect((brokenPath.match(/M/g) ?? []).length).toBeGreaterThan(1); + expect((bridgedPath.match(/M/g) ?? []).length).toBe(1); + }); + + it("renders an empty chart without lines when there are no categories", () => { + const { container } = render(); + + expect(container.querySelector("[data-slot='chart']")).not.toBeNull(); + expect(container.querySelectorAll("path.recharts-line-curve")).toHaveLength(0); + }); + + it("emits no per-chart style tag; colors flow through strokes, not CSS vars", () => { + const { container } = render( + , + ); + expect(container.querySelector("style")).toBeNull(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/charts/line_chart.tsx b/ui/litellm-dashboard/src/components/shared/charts/line_chart.tsx new file mode 100644 index 00000000000..2dc8747118d --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/charts/line_chart.tsx @@ -0,0 +1,99 @@ +"use client"; + +import * as React from "react"; +import { CartesianGrid, Line, LineChart as RechartsLineChart, XAxis, YAxis } from "recharts"; +import { ChartContainer, ChartLegend, ChartLegendContent, ChartTooltip, type ChartConfig } from "@/components/ui/chart"; +import { cn } from "@/lib/cva.config"; +import { ValueTooltip, type ChartTooltipComponent } from "./chart_tooltip"; +import { categoryFills, type ChartColor } from "./colors"; + +export type LineChartCurveType = "linear" | "natural" | "monotone" | "step"; + +export type LineChartProps> = { + data: readonly TDatum[]; + index: string; + categories: readonly string[]; + colors?: readonly ChartColor[]; + valueFormatter?: (value: number) => string; + yAxisWidth?: number; + tickGap?: number; + showLegend?: boolean; + showXAxis?: boolean; + showGridLines?: boolean; + showTooltip?: boolean; + customTooltip?: ChartTooltipComponent; + connectNulls?: boolean; + curveType?: LineChartCurveType; + className?: string; + style?: React.CSSProperties; +}; + +export function LineChart>({ + data, + index, + categories, + colors, + valueFormatter, + yAxisWidth = 56, + tickGap = 5, + showLegend = true, + showXAxis = true, + showGridLines = true, + showTooltip = true, + customTooltip, + connectNulls = false, + curveType = "linear", + className, + style, +}: LineChartProps) { + const fills = categoryFills(categories.length, colors); + const config: ChartConfig = Object.fromEntries(categories.map((category) => [category, { label: category }])); + const TooltipContent = customTooltip ?? ValueTooltip; + + return ( + + + {showGridLines && } + + + {showTooltip && ( + ( + + )} + /> + )} + {showLegend && ( + } + /> + )} + {categories.map((category, i) => ( + + ))} + + + ); +} 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..f3ec203adbb --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.test.tsx @@ -0,0 +1,41 @@ +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 that signals interactivity 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(); + // The clickable area must read as clickable: a hover background and a pointer cursor. + expect(button.className).toContain("hover:bg-muted"); + expect(button.className).toContain("cursor-pointer"); + 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..97263436454 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.tsx @@ -0,0 +1,51 @@ +"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..3a618114613 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx @@ -0,0 +1,69 @@ +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("shows 'No model access' for scope-restricted keys with an empty model list", () => { + const { rerender } = render(); + expect(screen.getByText("No model access")).toBeInTheDocument(); + expect(screen.queryByText("All Proxy Models")).not.toBeInTheDocument(); + rerender(); + expect(screen.getByText("No model access")).toBeInTheDocument(); + rerender(); + expect(screen.getByText("No model access")).toBeInTheDocument(); + }); + + it("still shows 'All Proxy Models' for empty models when the key is not scope-restricted", () => { + render(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + expect(screen.queryByText("No model access")).not.toBeInTheDocument(); + }); + + it("uses a persisted key_type to render 'No model access' regardless of allowed_routes", () => { + const { rerender } = render(); + expect(screen.getByText("No model access")).toBeInTheDocument(); + expect(screen.queryByText("All Proxy Models")).not.toBeInTheDocument(); + rerender(); + expect(screen.getByText("No model access")).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..81928d6fe87 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx @@ -0,0 +1,72 @@ +"use client"; + +import { deriveKeyModelScope } from "@/components/key_scope"; +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; + allowedRoutes?: string[] | null; + keyType?: string | null; +} + +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, allowedRoutes, keyType }: ModelsCellProps) { + if (!Array.isArray(models) || models.length === 0) { + const scope = deriveKeyModelScope(allowedRoutes, keyType); + if (!scope.hasModelAccess) { + return ( + + No model access + + } + /> + ); + } + 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/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index f0e27607ffb..a57676d07b8 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -91,6 +91,16 @@ vi.mock("@/components/team/member_permissions", () => ({ default: vi.fn(() =>
Member Permissions
), })); +vi.mock("@/components/common_components/ModelAliasManager", () => ({ + default: vi.fn(({ initialModelAliases, onAliasUpdate }) => ( +
+
{JSON.stringify(initialModelAliases)}
+ + +
+ )), +})); + vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ useAccessGroups: vi.fn().mockReturnValue({ data: [ @@ -858,4 +868,160 @@ describe("TeamInfoView", () => { }); }); }); + + describe("model aliases", () => { + const openSettingsEditor = async (user: ReturnType) => { + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /edit settings/i })); + + await waitFor(() => { + expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); + }); + }; + + it("should render existing model aliases in the read-only settings view", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + litellm_model_table: { model_aliases: { "my-smart-model": "gpt-4", "my-fast-model": "gpt-3.5-turbo" } }, + }), + ); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => { + expect(screen.getByText("Team Settings")).toBeInTheDocument(); + }); + + expect(screen.getByText("Model Aliases")).toBeInTheDocument(); + expect(screen.getByText("my-smart-model")).toBeInTheDocument(); + expect(screen.getByText("gpt-4")).toBeInTheDocument(); + expect(screen.getByText("my-fast-model")).toBeInTheDocument(); + expect(screen.getByText("gpt-3.5-turbo")).toBeInTheDocument(); + }); + + it("should show an empty state when the team has no model aliases", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ litellm_model_table: null })); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => { + expect(screen.getByText("Team Settings")).toBeInTheDocument(); + }); + + expect(screen.getByText("No model aliases configured")).toBeInTheDocument(); + }); + + it("should seed the alias editor from existing team aliases", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + models: ["gpt-4"], + litellm_model_table: { model_aliases: { "my-smart-model": "gpt-4" } }, + }), + ); + + renderWithProviders(); + + await openSettingsEditor(user); + + expect(screen.getByTestId("alias-editor-initial")).toHaveTextContent( + JSON.stringify({ "my-smart-model": "gpt-4" }), + ); + }); + + it("should pass model_aliases to teamUpdateCall when aliases are added", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["gpt-4"] })); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + + await user.click(screen.getByRole("button", { name: "Set Alias" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + team_id: "123", + model_aliases: { "gpt-4o": "gpt-4" }, + }), + ); + }); + }); + + it("should send an empty model_aliases map to clear existing aliases", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + models: ["gpt-4"], + litellm_model_table: { model_aliases: { "my-smart-model": "gpt-4" } }, + }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + + await user.click(screen.getByRole("button", { name: "Clear Aliases" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + const payload = vi.mocked(networking.teamUpdateCall).mock.calls[0][1] as Record; + expect(payload.model_aliases).toEqual({}); + }); + + it("should not include model_aliases when the team has none and the editor is untouched", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ models: ["gpt-4"], litellm_model_table: null }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + const payload = vi.mocked(networking.teamUpdateCall).mock.calls[0][1] as Record; + expect(payload).not.toHaveProperty("model_aliases"); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 3b80c598d34..74f201b40c6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -34,6 +34,7 @@ import { CheckIcon, CopyIcon } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; +import ModelAliasManager from "../common_components/ModelAliasManager"; import AgentSelector from "../agent_management/AgentSelector"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; import DurationSelect from "../common_components/DurationSelect"; @@ -106,7 +107,7 @@ export interface TeamData { budget_reset_at: string | null; model_id: string | null; litellm_model_table: { - model_aliases: Record; + model_aliases: Record | null; } | null; created_at: string; access_group_ids?: string[]; @@ -206,6 +207,7 @@ const TeamInfoView: React.FC = ({ const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [isDeleting, setIsDeleting] = useState(false); const [isTeamSaving, setIsTeamSaving] = useState(false); + const [teamModelAliases, setTeamModelAliases] = useState>({}); const routerSettingsRef = React.useRef(null); const [organization, setOrganization] = useState(null); const { userRole, userId } = useAuthorized(); @@ -619,6 +621,11 @@ const TeamInfoView: React.FC = ({ updateData.default_team_member_models = values.default_team_member_models; } + const previousModelAliases = info.litellm_model_table?.model_aliases ?? {}; + if (Object.keys(teamModelAliases).length > 0 || Object.keys(previousModelAliases).length > 0) { + updateData.model_aliases = teamModelAliases; + } + // Handle router_settings - read fresh values from DOM at save time. const currentRouterSettings = routerSettingsRef.current?.getValue(); if (currentRouterSettings?.router_settings) { @@ -905,7 +912,13 @@ const TeamInfoView: React.FC = ({
Team Settings {canEditTeam && !isEditing && ( - )} @@ -1026,6 +1039,24 @@ const TeamInfoView: React.FC = ({ /> + + Model Aliases{" "} + + + + + } + > + + + @@ -1520,6 +1551,26 @@ const TeamInfoView: React.FC = ({
)} +
+ Model Aliases + {(() => { + const aliasEntries = Object.entries(info.litellm_model_table?.model_aliases ?? {}); + if (aliasEntries.length === 0) { + return
No model aliases configured
; + } + return ( +
+ {aliasEntries.map(([alias, target]) => ( +
+ {alias} + {" -> "} + {target} +
+ ))} +
+ ); + })()} +
Rate Limits
TPM: {info.tpm_limit || "Unlimited"}
diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index 207e6f2ccfe..66690e5478f 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"; @@ -17,6 +18,7 @@ import { Popover, Tooltip, Typography } from "antd"; import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; import React, { useCallback, useEffect, useMemo, useState } from "react"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; +import { deriveKeyModelScope } from "../key_scope"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; import KeyInfoView from "../templates/key_info_view"; @@ -43,7 +45,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); @@ -338,14 +340,24 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi enableSorting: false, cell: (info) => { const models = info.getValue() as string[]; + const scope = deriveKeyModelScope(info.row.original.allowed_routes, info.row.original.key_type); + const emptyModelsBadge = !scope.hasModelAccess ? ( + + + No model access + + + ) : ( + + All Proxy Models + + ); return (
{Array.isArray(models) ? (
{models.length === 0 ? ( - - All Proxy Models - + emptyModelsBadge ) : ( <>
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/dropdown-menu.tsx b/ui/litellm-dashboard/src/components/ui/dropdown-menu.tsx new file mode 100644 index 00000000000..03fefba7d19 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ui/dropdown-menu.tsx @@ -0,0 +1,254 @@ +"use client"; + +import * as React from "react"; +import { Menu as MenuPrimitive } from "@base-ui/react/menu"; + +import { cn } from "@/lib/cva.config"; +import { ChevronRightIcon, CheckIcon } from "lucide-react"; + +function DropdownMenu({ ...props }: MenuPrimitive.Root.Props) { + return ; +} + +function DropdownMenuPortal({ ...props }: MenuPrimitive.Portal.Props) { + return ; +} + +function DropdownMenuTrigger({ ...props }: MenuPrimitive.Trigger.Props) { + return ; +} + +function DropdownMenuContent({ + align = "start", + alignOffset = 0, + side = "bottom", + sideOffset = 4, + className, + ...props +}: MenuPrimitive.Popup.Props & Pick) { + return ( + + + + + + ); +} + +function DropdownMenuGroup({ ...props }: MenuPrimitive.Group.Props) { + return ; +} + +function DropdownMenuLabel({ + className, + inset, + ...props +}: MenuPrimitive.GroupLabel.Props & { + inset?: boolean; +}) { + return ( + + ); +} + +function DropdownMenuItem({ + className, + inset, + variant = "default", + ...props +}: MenuPrimitive.Item.Props & { + inset?: boolean; + variant?: "default" | "destructive"; +}) { + return ( + + ); +} + +function DropdownMenuSub({ ...props }: MenuPrimitive.SubmenuRoot.Props) { + return ; +} + +function DropdownMenuSubTrigger({ + className, + inset, + children, + ...props +}: MenuPrimitive.SubmenuTrigger.Props & { + inset?: boolean; +}) { + return ( + + {children} + + + ); +} + +function DropdownMenuSubContent({ + align = "start", + alignOffset = -3, + side = "right", + sideOffset = 0, + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function DropdownMenuCheckboxItem({ + className, + children, + checked, + inset, + ...props +}: MenuPrimitive.CheckboxItem.Props & { + inset?: boolean; +}) { + return ( + + + + + + + {children} + + ); +} + +function DropdownMenuRadioGroup({ ...props }: MenuPrimitive.RadioGroup.Props) { + return ; +} + +function DropdownMenuRadioItem({ + className, + children, + inset, + ...props +}: MenuPrimitive.RadioItem.Props & { + inset?: boolean; +}) { + return ( + + + + + + + {children} + + ); +} + +function DropdownMenuSeparator({ className, ...props }: MenuPrimitive.Separator.Props) { + return ( + + ); +} + +function DropdownMenuShortcut({ className, ...props }: React.ComponentProps<"span">) { + return ( + + ); +} + +export { + DropdownMenu, + DropdownMenuPortal, + DropdownMenuTrigger, + DropdownMenuContent, + DropdownMenuGroup, + DropdownMenuLabel, + DropdownMenuItem, + DropdownMenuCheckboxItem, + DropdownMenuRadioGroup, + DropdownMenuRadioItem, + DropdownMenuSeparator, + DropdownMenuShortcut, + DropdownMenuSub, + DropdownMenuSubTrigger, + DropdownMenuSubContent, +}; 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 ( +