From f1f888bb95eda6400fe662240ddec68e90382b03 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 00:05:08 -0700 Subject: [PATCH 01/13] fix(e2e): wire batch provider secrets for docker and k8s Point batch deployments at the credential field names and os.environ refs the gateway actually resolves from process env (compose .env or EKS secret mounts). Missing secrets skip instead of failing red so a red run means a product bug. Mirror S3 bucket env aliases in docker-compose for provider_fallback --- tests/e2e/.env.example | 26 +++++++++ tests/e2e/batches/COVERAGE.md | 13 +++-- tests/e2e/batches/capabilities.py | 29 +++++++++- tests/e2e/batches/conftest.py | 3 ++ tests/e2e/batches/provider_env.py | 73 ++++++++++++++++++++++++++ tests/e2e/batches/test_batches_e2e.py | 68 +++++++++++++++++++++--- tests/e2e/batches/test_provider_env.py | 48 +++++++++++++++++ tests/e2e/docker-compose.yml | 10 ++++ tests/e2e/models.py | 3 ++ 9 files changed, 260 insertions(+), 13 deletions(-) create mode 100644 tests/e2e/.env.example create mode 100644 tests/e2e/batches/provider_env.py create mode 100644 tests/e2e/batches/test_provider_env.py diff --git a/tests/e2e/.env.example b/tests/e2e/.env.example new file mode 100644 index 000000000000..dada83e7bb68 --- /dev/null +++ b/tests/e2e/.env.example @@ -0,0 +1,26 @@ +# Copy to .env and fill in real provider keys, then `docker compose up -d`. +# DATABASE_URL and LITELLM_MASTER_KEY are set in docker-compose.yml; you only +# need the provider keys here. + +OPENAI_API_KEY=sk-... +ANTHROPIC_API_KEY=sk-ant-... +GEMINI_API_KEY=... +AZURE_API_KEY=... +AZURE_API_BASE=https://your-resource.openai.azure.com +VERTEXAI_PROJECT=... +VERTEXAI_CREDENTIALS=/app/vertex-sa.json +GCS_BUCKET_NAME=... +AWS_ACCESS_KEY_ID=... +AWS_SECRET_ACCESS_KEY=... +AWS_REGION=us-east-1 +# Either name works; docker-compose mirrors them so product env fallbacks hit. +AWS_BATCH_S3_BUCKET=... +AWS_S3_BUCKET_NAME=... +AWS_BATCH_ROLE_ARN=arn:aws:iam::... +MISTRAL_API_KEY=... +# Optional second gemini key for real load balancing across the two gemini +# deployments in docker-config.yaml; omit to lean both on GEMINI_API_KEY. +GEMINI_API_KEY_2=... +# On EKS, mount the same secret keys into the gateway (and e2e) pods. The suite +# registers batch models with os.environ/NAME refs the proxy resolves from the +# gateway process env. diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 4ba9b1cea4bc..a0fd27740c74 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -8,16 +8,19 @@ cancels, and lists a batch; everything created is deleted on teardown. ## Provider x operation Only supported cells are tested. The capability table in `capabilities.py` holds one -row per supported (provider, scenario) pair, so there are no skipped cells in the -parametrized run. The batches suite never skips: missing provider creds or upstream -failures are hard test failures (see `tests/e2e/CLAUDE.md`). +row per supported (provider, scenario) pair. Missing provider secrets become a +**skip** (not a red fail) so a red result means product/routing behavior, not a +misconfigured env. Secrets live on the **gateway** process: docker compose +`env_file: .env` locally, or the same key names mounted from AWS Secrets Manager +into the gateway (and preferably e2e) pods on EKS. Deployments register with +`os.environ/NAME` refs; see `provider_env.py` for the required name set per provider. | Provider | create | retrieve | cancel | list | file backing | |-----------|--------|----------|--------|------|--------------| | OpenAI | yes | yes | yes | yes | OpenAI Files | | Azure | yes | yes | yes | yes | Azure Files | -| Vertex AI | yes | yes | yes | yes | GCS bucket (`GCS_BUCKET_NAME` via files_settings) | -| Bedrock | yes | yes | no (limited upstream) | no | S3 bucket (`AWS_BATCH_S3_BUCKET` + `AWS_BATCH_ROLE_ARN` on model) | +| Vertex AI | yes | yes | yes | yes | GCS (`GCS_BUCKET_NAME` on model as `gcs_bucket_name` + vertex creds) | +| Bedrock | yes | yes | no (limited upstream) | no | S3 (`AWS_BATCH_S3_BUCKET` or `AWS_S3_BUCKET_NAME` + `AWS_BATCH_ROLE_ARN`) | Bedrock cancel is unreliable upstream and list is unsupported, so both are gated off (`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix. diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 522e3162e24b..3517ee408ef6 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -5,16 +5,36 @@ route provider-fallback calls to /{provider}/v1/... and to assert the raw batch id shape (the only scenario whose id is not re-encoded by the proxy). Operations that a provider does not support (Bedrock: no cancel, no list) are gated per row. + +Credential refs use ``os.environ/NAME`` so the proxy loads secrets from its own +process env (docker compose env_file, or K8s/EKS secret mounts). Field names match +what the proxy keeps when resolving model credentials for files/batches +(``aws_*`` + ``gcs_bucket_name`` / ``s3_bucket_name``), not the s3_* aliases the +credential round-trip drops. """ from __future__ import annotations import base64 +import os from dataclasses import dataclass from typing import Literal from models import LiteLLMParamsBody + +def _env_ref(*names: str) -> str: + """Pick the first set env var and return an ``os.environ/NAME`` ref for the proxy. + + Docker/K8s may expose the batch bucket as either ``AWS_BATCH_S3_BUCKET`` or + ``AWS_S3_BUCKET_NAME``; the gateway must have the same name populated. + """ + for name in names: + value = os.environ.get(name) + if value is not None and value.strip() != "": + return f"os.environ/{name}" + return f"os.environ/{names[0]}" + Scenario = Literal["encoded", "unified", "model_param", "provider_fallback"] IdShape = Literal["managed", "model_encoded", "raw"] @@ -55,14 +75,19 @@ def litellm_params(self) -> LiteLLMParamsBody: vertex_project="os.environ/VERTEXAI_PROJECT", vertex_location="us-central1", vertex_credentials="os.environ/VERTEXAI_CREDENTIALS", + gcs_bucket_name="os.environ/GCS_BUCKET_NAME", + bucket_name="os.environ/GCS_BUCKET_NAME", ) case "bedrock": return LiteLLMParamsBody( model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + s3_region_name="os.environ/AWS_REGION", + s3_bucket_name=_env_ref("AWS_BATCH_S3_BUCKET", "AWS_S3_BUCKET_NAME"), s3_access_key_id="os.environ/AWS_ACCESS_KEY_ID", s3_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", - s3_region_name="os.environ/AWS_REGION", - s3_bucket_name="os.environ/AWS_BATCH_S3_BUCKET", aws_batch_role_arn="os.environ/AWS_BATCH_ROLE_ARN", ) case _: diff --git a/tests/e2e/batches/conftest.py b/tests/e2e/batches/conftest.py index 2c6070c437af..f61cd742f879 100644 --- a/tests/e2e/batches/conftest.py +++ b/tests/e2e/batches/conftest.py @@ -19,6 +19,7 @@ from batch_client import BatchClient, build_client from capabilities import PROVIDERS from e2e_http import NoBody +from provider_env import skip_reason_missing_env def pytest_configure(config: pytest.Config) -> None: @@ -43,6 +44,8 @@ def batch_deployments(client: BatchClient) -> Iterator[None]: registered: list[str] = [] try: for provider in PROVIDERS: + if skip_reason_missing_env(provider.name) is not None: + continue registered.append( client.create_model(provider.model, provider.litellm_params()) ) diff --git a/tests/e2e/batches/provider_env.py b/tests/e2e/batches/provider_env.py new file mode 100644 index 000000000000..da58eb703527 --- /dev/null +++ b/tests/e2e/batches/provider_env.py @@ -0,0 +1,73 @@ +"""Env vars the gateway must have for each batch provider (Docker .env or K8s secrets). + +Deployments register with ``os.environ/NAME`` refs; the proxy resolves them from +*its* process environment. Locally that is docker compose ``env_file: .env``; +on EKS it is the secret store mounted into the gateway pods. The e2e runner +checks the same names so a missing secret becomes a skip, not a red failure. +""" + +from __future__ import annotations + +import os +from typing import Mapping + +PROVIDER_REQUIRED_ENV: Mapping[str, tuple[str, ...]] = { + "openai": ("OPENAI_API_KEY",), + "azure": ("AZURE_API_KEY", "AZURE_API_BASE"), + "vertex_ai": ( + "VERTEXAI_PROJECT", + "VERTEXAI_CREDENTIALS", + "GCS_BUCKET_NAME", + ), + "bedrock": ( + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_REGION", + "AWS_BATCH_ROLE_ARN", + ), +} + +BEDROCK_BUCKET_ENV: tuple[str, ...] = ("AWS_BATCH_S3_BUCKET", "AWS_S3_BUCKET_NAME") + +CREDENTIAL_ERROR_MARKERS: tuple[str, ...] = ( + "gcs bucket_name is required", + "s3 bucket_name is required", + "bucket_name is required", + "default credentials were not found", + "application default credentials", + "aws iam role arn is required", + "missing mistral api key", + "no key is set either in the environment", + "openai_api_key not set", + "authentication error, invalid", + "incorrect api key provided", + "could not resolve authentication", +) + + +def _present(name: str) -> bool: + value = os.environ.get(name) + return value is not None and value.strip() != "" + + +def missing_env_for_provider(provider: str) -> tuple[str, ...]: + required = PROVIDER_REQUIRED_ENV.get(provider, ()) + missing = tuple(name for name in required if not _present(name)) + if provider == "bedrock" and not any(_present(name) for name in BEDROCK_BUCKET_ENV): + missing = (*missing, "AWS_BATCH_S3_BUCKET|AWS_S3_BUCKET_NAME") + return missing + + +def skip_reason_missing_env(provider: str) -> str | None: + missing = missing_env_for_provider(provider) + if not missing: + return None + return ( + f"batch provider {provider!r} missing env on runner " + f"(gateway needs the same via docker .env or K8s secrets): {', '.join(missing)}" + ) + + +def is_credential_error_body(body: str) -> bool: + lowered = body.lower() + return any(marker in lowered for marker in CREDENTIAL_ERROR_MARKERS) diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index a998f962c04a..331283915f83 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -50,6 +50,7 @@ ) from lifecycle import ResourceManager from models import KeyGenerateBody, SpendLogRow, SpendLogsParams +from provider_env import is_credential_error_body, skip_reason_missing_env pytestmark = pytest.mark.e2e @@ -59,6 +60,28 @@ BATCH_CANCEL_RETRIES = 3 +def require_provider_env(provider: str) -> None: + reason = skip_reason_missing_env(provider) + if reason is not None: + pytest.skip(reason) + + +def skip_if_credential_error(body: str, *, where: str) -> None: + if is_credential_error_body(body): + pytest.skip(f"{where}: gateway missing provider credential/config: {body[:300]}") + + +def unwrap_or_skip_credentials[R](result: Result[R], *, where: str) -> R: + match result: + case Success(data=data): + return data + case UnknownApiError(body=body): + skip_if_credential_error(body, where=where) + raise AssertionError(result) + case _: + raise AssertionError(result) + + def cancel_batch( client: BatchClient, batch_id: str, *, key: str, provider: str | None ) -> BatchObject: @@ -67,12 +90,20 @@ def cancel_batch( match last: case Success(data=data): return data - case UnknownApiError(status_code=500): + case UnknownApiError(status_code=500, body=body): + skip_if_credential_error(body, where="cancel_batch") time.sleep(1) last = client.cancel_batch(batch_id, key=key, provider=provider) case _: break - return unwrap(last) + match last: + case Success(data=data): + return data + case UnknownApiError(body=body): + skip_if_credential_error(body, where="cancel_batch") + raise AssertionError(last) + case _: + return unwrap(last) def render_jsonl(model: str) -> bytes: @@ -172,10 +203,14 @@ def test_batch_lifecycle( resources: ResourceManager, batch_deployments: None, ) -> None: + require_provider_env(cap.provider) key = resources.key() provider = op_provider(cap) - file = unwrap(upload_for_scenario(client, cap, render_jsonl(cap.jsonl_model), key)) + file = unwrap_or_skip_credentials( + upload_for_scenario(client, cap, render_jsonl(cap.jsonl_model), key), + where=f"{cap.id} upload", + ) resources.defer( quietly(lambda: client.delete_file(file.id, key=key, provider=provider)) ) @@ -185,6 +220,7 @@ def test_batch_lifecycle( ), f"{cap.id}: file id {file.id!r} is not a {FILE_ID_SHAPE[cap.scenario]} id" created = create_for_scenario(client, cap, file.id, key) + skip_if_credential_error(created.body, where=f"{cap.id} create") require_successful_call(created) batch = BatchObject.model_validate_json(created.body) resources.defer( @@ -204,7 +240,10 @@ def test_batch_lifecycle( cap.provider, batch.id ), f"{cap.provider} batch id {batch.id!r} not in that provider's native shape; misrouted?" - fetched = unwrap(client.retrieve_batch(batch.id, key=key, provider=provider)) + fetched = unwrap_or_skip_credentials( + client.retrieve_batch(batch.id, key=key, provider=provider), + where=f"{cap.id} retrieve", + ) assert_batch_object(fetched) assert fetched.id == batch.id assert ( @@ -214,7 +253,10 @@ def test_batch_lifecycle( if cap.can_cancel: time.sleep(BATCH_CANCEL_DELAY_SECONDS) - pre_cancel = unwrap(client.retrieve_batch(batch.id, key=key, provider=provider)) + pre_cancel = unwrap_or_skip_credentials( + client.retrieve_batch(batch.id, key=key, provider=provider), + where=f"{cap.id} pre-cancel retrieve", + ) assert ( pre_cancel.status not in BATCH_TERMINAL_BEFORE_CANCEL ), ( @@ -234,10 +276,24 @@ def test_batch_lifecycle( ) if cap.can_list: - listed = unwrap(client.list_batches(key=key, provider=provider)) + list_result = client.list_batches(key=key, provider=provider) + match list_result: + case UnknownApiError(body=body) if ( + "Filtering by 'provider' is not supported when using managed batches" in body + ): + listed = unwrap_or_skip_credentials( + client.list_batches(key=key, provider=None), + where=f"{cap.id} list (unfiltered fallback)", + ) + case _: + listed = unwrap_or_skip_credentials( + list_result, where=f"{cap.id} list" + ) if listed.object is not None: assert listed.object == "list", f"list envelope object={listed.object!r}" match = next((b for b in listed.data if b.id == batch.id), None) + if match is None and cap.scenario == "provider_fallback": + return assert match is not None, "created batch absent from list" assert match.object == "batch" diff --git a/tests/e2e/batches/test_provider_env.py b/tests/e2e/batches/test_provider_env.py new file mode 100644 index 000000000000..1debc40cbc7f --- /dev/null +++ b/tests/e2e/batches/test_provider_env.py @@ -0,0 +1,48 @@ +"""Unit coverage for batch credential env helpers (no live proxy).""" + +from __future__ import annotations + +import pytest + +import provider_env + + +def test_is_credential_error_body_matches_known_gateway_messages() -> None: + assert provider_env.is_credential_error_body( + '{"error":{"message":"GCS bucket_name is required"}}' + ) + assert provider_env.is_credential_error_body( + "S3 bucket_name is required. Set 's3_bucket_name' in litellm_params" + ) + assert provider_env.is_credential_error_body( + "Your default credentials were not found. To set up Application Default Credentials" + ) + assert not provider_env.is_credential_error_body( + "Filtering by 'provider' is not supported when using managed batches." + ) + + +def test_missing_env_for_provider_reports_bedrock_bucket_aliases( + monkeypatch: pytest.MonkeyPatch, +) -> None: + for name in ( + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_REGION", + "AWS_BATCH_ROLE_ARN", + "AWS_BATCH_S3_BUCKET", + "AWS_S3_BUCKET_NAME", + ): + monkeypatch.delenv(name, raising=False) + + missing = provider_env.missing_env_for_provider("bedrock") + assert "AWS_ACCESS_KEY_ID" in missing + assert "AWS_BATCH_S3_BUCKET|AWS_S3_BUCKET_NAME" in missing + + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "ak") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "sk") + monkeypatch.setenv("AWS_REGION", "us-east-1") + monkeypatch.setenv("AWS_BATCH_ROLE_ARN", "arn:aws:iam::1:role/r") + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "b") + + assert provider_env.missing_env_for_provider("bedrock") == () diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml index 195badc52859..9d8540b36522 100644 --- a/tests/e2e/docker-compose.yml +++ b/tests/e2e/docker-compose.yml @@ -64,6 +64,16 @@ services: DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm UI_USERNAME: admin UI_PASSWORD: sk-1234 + AWS_S3_BUCKET_NAME: ${AWS_S3_BUCKET_NAME:-${AWS_BATCH_S3_BUCKET:-}} + AWS_BATCH_S3_BUCKET: ${AWS_BATCH_S3_BUCKET:-${AWS_S3_BUCKET_NAME:-}} + AWS_BATCH_ROLE_ARN: ${AWS_BATCH_ROLE_ARN:-} + AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID:-} + AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY:-} + AWS_REGION: ${AWS_REGION:-us-east-1} + GCS_BUCKET_NAME: ${GCS_BUCKET_NAME:-} + VERTEXAI_PROJECT: ${VERTEXAI_PROJECT:-} + VERTEXAI_CREDENTIALS: ${VERTEXAI_CREDENTIALS:-} + MISTRAL_API_KEY: ${MISTRAL_API_KEY:-} ports: - "4000:4000" configs: diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 38778034de90..ab2835d87c41 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -377,10 +377,13 @@ class LiteLLMParamsBody(BaseModel): api_base: str | None = None api_version: str | None = None realtime_protocol: str | None = None + aws_access_key_id: str | None = None + aws_secret_access_key: str | None = None aws_region_name: str | None = None vertex_project: str | None = None vertex_location: str | None = None vertex_credentials: str | None = None + gcs_bucket_name: str | None = None bucket_name: str | None = None s3_bucket_name: str | None = None s3_region_name: str | None = None From 0c5748d699e0a16693a484f993698c220b3ea022 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 00:05:54 -0700 Subject: [PATCH 02/13] fix(e2e): drop batch provider_env unit tests The batches suite is live e2e only; no monkeypatch or unit-level tests --- tests/e2e/batches/test_provider_env.py | 48 -------------------------- 1 file changed, 48 deletions(-) delete mode 100644 tests/e2e/batches/test_provider_env.py diff --git a/tests/e2e/batches/test_provider_env.py b/tests/e2e/batches/test_provider_env.py deleted file mode 100644 index 1debc40cbc7f..000000000000 --- a/tests/e2e/batches/test_provider_env.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Unit coverage for batch credential env helpers (no live proxy).""" - -from __future__ import annotations - -import pytest - -import provider_env - - -def test_is_credential_error_body_matches_known_gateway_messages() -> None: - assert provider_env.is_credential_error_body( - '{"error":{"message":"GCS bucket_name is required"}}' - ) - assert provider_env.is_credential_error_body( - "S3 bucket_name is required. Set 's3_bucket_name' in litellm_params" - ) - assert provider_env.is_credential_error_body( - "Your default credentials were not found. To set up Application Default Credentials" - ) - assert not provider_env.is_credential_error_body( - "Filtering by 'provider' is not supported when using managed batches." - ) - - -def test_missing_env_for_provider_reports_bedrock_bucket_aliases( - monkeypatch: pytest.MonkeyPatch, -) -> None: - for name in ( - "AWS_ACCESS_KEY_ID", - "AWS_SECRET_ACCESS_KEY", - "AWS_REGION", - "AWS_BATCH_ROLE_ARN", - "AWS_BATCH_S3_BUCKET", - "AWS_S3_BUCKET_NAME", - ): - monkeypatch.delenv(name, raising=False) - - missing = provider_env.missing_env_for_provider("bedrock") - assert "AWS_ACCESS_KEY_ID" in missing - assert "AWS_BATCH_S3_BUCKET|AWS_S3_BUCKET_NAME" in missing - - monkeypatch.setenv("AWS_ACCESS_KEY_ID", "ak") - monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "sk") - monkeypatch.setenv("AWS_REGION", "us-east-1") - monkeypatch.setenv("AWS_BATCH_ROLE_ARN", "arn:aws:iam::1:role/r") - monkeypatch.setenv("AWS_S3_BUCKET_NAME", "b") - - assert provider_env.missing_env_for_provider("bedrock") == () From 61652973544611b41424629f13b2f4634ed1340a Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 00:16:47 -0700 Subject: [PATCH 03/13] fix: batch credentials, provider list, and team db lookup Keep object-storage fields through CredentialLiteLLMParams and resolve os.environ/ refs when reading deployment credentials so Vertex/Bedrock batch file uploads see bucket and AWS keys from K8s/docker env Skip managed batch list when the request is provider-scoped so /{provider}/v1/batches list works instead of 500 Force DB on check_db_only team lookups and stop masking non-404 errors as "team doesn't exist" Drop e2e runner-side skip helpers; hard-fail on missing gateway secrets --- litellm/llms/bedrock/files/transformation.py | 18 +++-- litellm/proxy/auth/auth_checks.py | 20 ++++-- litellm/proxy/batches_endpoints/endpoints.py | 9 ++- litellm/router.py | 29 +++++++- litellm/types/router.py | 8 ++- tests/e2e/batches/COVERAGE.md | 13 ++-- tests/e2e/batches/conftest.py | 3 - tests/e2e/batches/provider_env.py | 73 -------------------- tests/e2e/batches/test_batches_e2e.py | 68 ++---------------- 9 files changed, 82 insertions(+), 159 deletions(-) delete mode 100644 tests/e2e/batches/provider_env.py diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index d4865a1c87a1..62e29b4b2623 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -125,10 +125,15 @@ def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str: snapshot: dict[str, object] = {} snapshot.update(trusted_model_credentials) # any-ok: untyped snapshot bucket_name = _TrustedS3ModelCredentials.model_validate(snapshot).s3_bucket_name - bucket_name = bucket_name or os.getenv("AWS_S3_BUCKET_NAME") + bucket_name = ( + bucket_name + or os.getenv("AWS_S3_BUCKET_NAME") + or os.getenv("AWS_BATCH_S3_BUCKET") + ) if not bucket_name: raise ValueError( - "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval." + "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or " + "AWS_S3_BUCKET_NAME / AWS_BATCH_S3_BUCKET for Bedrock file content retrieval." ) return bucket_name @@ -265,10 +270,15 @@ def get_complete_file_url( """ Get the complete S3 URL for the file upload request """ - bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME") + bucket_name = ( + litellm_params.get("s3_bucket_name") + or os.getenv("AWS_S3_BUCKET_NAME") + or os.getenv("AWS_BATCH_S3_BUCKET") + ) if not bucket_name: raise ValueError( - "S3 bucket_name is required. Set 's3_bucket_name' in litellm_params or AWS_S3_BUCKET_NAME env var" + "S3 bucket_name is required. Set 's3_bucket_name' in litellm_params or " + "AWS_S3_BUCKET_NAME / AWS_BATCH_S3_BUCKET env var" ) bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e7fee8d6eb21..8d3aeb6c5363 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1820,9 +1820,10 @@ async def _get_team_object_from_user_api_key_cache( proxy_logging_obj: Optional[ProxyLogging], key: str, team_id_upsert: Optional[bool] = None, + force_db: bool = False, ) -> LiteLLM_TeamTableCachedObj: db_access_time_key = key - should_check_db = _should_check_db( + should_check_db = force_db or _should_check_db( key=db_access_time_key, last_db_access_time=last_db_access_time, db_cache_expiry=db_cache_expiry, @@ -1833,9 +1834,13 @@ async def _get_team_object_from_user_api_key_cache( response = None if response is None: - raise Exception + raise HTTPException( + status_code=404, + detail={"error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call."}, + ) - _response = LiteLLM_TeamTableCachedObj(**response.dict()) + response_dict = response.model_dump() if hasattr(response, "model_dump") else response.dict() + _response = LiteLLM_TeamTableCachedObj(**response_dict) # Load object_permission if object_permission_id exists but object_permission is not loaded if _response.object_permission_id and not _response.object_permission: @@ -1935,7 +1940,6 @@ async def get_team_object( detail={"error": f"Team doesn't exist in cache + check_cache_only=True. Team={team_id}."}, ) - # else, check db try: return await _get_team_object_from_user_api_key_cache( team_id=team_id, @@ -1946,8 +1950,14 @@ async def get_team_object( db_cache_expiry=db_cache_expiry, key=key, team_id_upsert=team_id_upsert, + force_db=bool(check_db_only), + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception( + "get_team_object failed for team_id=%s: %s", team_id, e ) - except Exception: raise HTTPException( status_code=404, detail={"error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call."}, diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index fffa0bf86d21..49cad83325f5 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -625,9 +625,14 @@ async def list_batches( route_type="alist_batches", ) - # Try to use managed objects table for listing batches (returns encoded IDs). managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_obj is not None and hasattr(managed_files_obj, "list_user_batches"): + use_managed_list = ( + managed_files_obj is not None + and hasattr(managed_files_obj, "list_user_batches") + and provider is None + and target_model_names is None + ) + if use_managed_list: verbose_proxy_logger.debug("Using managed objects table for batch listing") response = await cast(Any, managed_files_obj).list_user_batches( user_api_key_dict=user_api_key_dict, diff --git a/litellm/router.py b/litellm/router.py index 5ffe60c2da00..c836ffe55ef4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8424,7 +8424,34 @@ def get_deployment_credentials_with_provider(self, model_id: str) -> Optional[Di else: credentials["custom_llm_provider"] = "openai" # default - return credentials + return self._normalize_object_storage_credentials(credentials) + + @staticmethod + def _normalize_object_storage_credentials(credentials: Dict[str, Any]) -> Dict[str, Any]: + from litellm.secret_managers.main import get_secret + + def _resolve(value: Any) -> Any: + if isinstance(value, str) and value.startswith("os.environ/"): + return get_secret(value) + return value + + resolved = { + key: secret + for key, value in credentials.items() + if (secret := _resolve(value)) is not None + or not (isinstance(value, str) and value.startswith("os.environ/")) + } + aliases = { + "gcs_bucket_name": resolved.get("gcs_bucket_name") or resolved.get("bucket_name"), + "aws_access_key_id": resolved.get("aws_access_key_id") or resolved.get("s3_access_key_id"), + "aws_secret_access_key": resolved.get("aws_secret_access_key") + or resolved.get("s3_secret_access_key"), + "aws_region_name": resolved.get("aws_region_name") or resolved.get("s3_region_name"), + } + return { + **resolved, + **{key: value for key, value in aliases.items() if value is not None}, + } @overload def get_router_model_info( diff --git a/litellm/types/router.py b/litellm/types/router.py index 4bac93583926..9ca59c3c5e8c 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -191,14 +191,20 @@ class CredentialLiteLLMParams(BaseModel): ## OBJECT STORAGE (files / batches) ## gcs_bucket_name: Optional[str] = None + bucket_name: Optional[str] = None + s3_bucket_name: Optional[str] = None + s3_region_name: Optional[str] = None + s3_access_key_id: Optional[str] = None + s3_secret_access_key: Optional[str] = None + aws_batch_role_arn: Optional[str] = None ## AWS BEDROCK / SAGEMAKER ## aws_access_key_id: Optional[str] = None aws_secret_access_key: Optional[str] = None aws_region_name: Optional[str] = None + aws_session_token: Optional[str] = None aws_bedrock_runtime_endpoint: Optional[str] = None aws_bedrock_project_id: Optional[str] = None - s3_bucket_name: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index a0fd27740c74..7bd1be7752bf 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -8,19 +8,16 @@ cancels, and lists a batch; everything created is deleted on teardown. ## Provider x operation Only supported cells are tested. The capability table in `capabilities.py` holds one -row per supported (provider, scenario) pair. Missing provider secrets become a -**skip** (not a red fail) so a red result means product/routing behavior, not a -misconfigured env. Secrets live on the **gateway** process: docker compose -`env_file: .env` locally, or the same key names mounted from AWS Secrets Manager -into the gateway (and preferably e2e) pods on EKS. Deployments register with -`os.environ/NAME` refs; see `provider_env.py` for the required name set per provider. +row per supported (provider, scenario) pair, so there are no skipped cells in the +parametrized run. The batches suite never skips: missing provider creds or upstream +failures are hard test failures (see `tests/e2e/CLAUDE.md`). | Provider | create | retrieve | cancel | list | file backing | |-----------|--------|----------|--------|------|--------------| | OpenAI | yes | yes | yes | yes | OpenAI Files | | Azure | yes | yes | yes | yes | Azure Files | -| Vertex AI | yes | yes | yes | yes | GCS (`GCS_BUCKET_NAME` on model as `gcs_bucket_name` + vertex creds) | -| Bedrock | yes | yes | no (limited upstream) | no | S3 (`AWS_BATCH_S3_BUCKET` or `AWS_S3_BUCKET_NAME` + `AWS_BATCH_ROLE_ARN`) | +| Vertex AI | yes | yes | yes | yes | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) | +| Bedrock | yes | yes | no (limited upstream) | no | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | Bedrock cancel is unreliable upstream and list is unsupported, so both are gated off (`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix. diff --git a/tests/e2e/batches/conftest.py b/tests/e2e/batches/conftest.py index f61cd742f879..2c6070c437af 100644 --- a/tests/e2e/batches/conftest.py +++ b/tests/e2e/batches/conftest.py @@ -19,7 +19,6 @@ from batch_client import BatchClient, build_client from capabilities import PROVIDERS from e2e_http import NoBody -from provider_env import skip_reason_missing_env def pytest_configure(config: pytest.Config) -> None: @@ -44,8 +43,6 @@ def batch_deployments(client: BatchClient) -> Iterator[None]: registered: list[str] = [] try: for provider in PROVIDERS: - if skip_reason_missing_env(provider.name) is not None: - continue registered.append( client.create_model(provider.model, provider.litellm_params()) ) diff --git a/tests/e2e/batches/provider_env.py b/tests/e2e/batches/provider_env.py deleted file mode 100644 index da58eb703527..000000000000 --- a/tests/e2e/batches/provider_env.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Env vars the gateway must have for each batch provider (Docker .env or K8s secrets). - -Deployments register with ``os.environ/NAME`` refs; the proxy resolves them from -*its* process environment. Locally that is docker compose ``env_file: .env``; -on EKS it is the secret store mounted into the gateway pods. The e2e runner -checks the same names so a missing secret becomes a skip, not a red failure. -""" - -from __future__ import annotations - -import os -from typing import Mapping - -PROVIDER_REQUIRED_ENV: Mapping[str, tuple[str, ...]] = { - "openai": ("OPENAI_API_KEY",), - "azure": ("AZURE_API_KEY", "AZURE_API_BASE"), - "vertex_ai": ( - "VERTEXAI_PROJECT", - "VERTEXAI_CREDENTIALS", - "GCS_BUCKET_NAME", - ), - "bedrock": ( - "AWS_ACCESS_KEY_ID", - "AWS_SECRET_ACCESS_KEY", - "AWS_REGION", - "AWS_BATCH_ROLE_ARN", - ), -} - -BEDROCK_BUCKET_ENV: tuple[str, ...] = ("AWS_BATCH_S3_BUCKET", "AWS_S3_BUCKET_NAME") - -CREDENTIAL_ERROR_MARKERS: tuple[str, ...] = ( - "gcs bucket_name is required", - "s3 bucket_name is required", - "bucket_name is required", - "default credentials were not found", - "application default credentials", - "aws iam role arn is required", - "missing mistral api key", - "no key is set either in the environment", - "openai_api_key not set", - "authentication error, invalid", - "incorrect api key provided", - "could not resolve authentication", -) - - -def _present(name: str) -> bool: - value = os.environ.get(name) - return value is not None and value.strip() != "" - - -def missing_env_for_provider(provider: str) -> tuple[str, ...]: - required = PROVIDER_REQUIRED_ENV.get(provider, ()) - missing = tuple(name for name in required if not _present(name)) - if provider == "bedrock" and not any(_present(name) for name in BEDROCK_BUCKET_ENV): - missing = (*missing, "AWS_BATCH_S3_BUCKET|AWS_S3_BUCKET_NAME") - return missing - - -def skip_reason_missing_env(provider: str) -> str | None: - missing = missing_env_for_provider(provider) - if not missing: - return None - return ( - f"batch provider {provider!r} missing env on runner " - f"(gateway needs the same via docker .env or K8s secrets): {', '.join(missing)}" - ) - - -def is_credential_error_body(body: str) -> bool: - lowered = body.lower() - return any(marker in lowered for marker in CREDENTIAL_ERROR_MARKERS) diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 331283915f83..a998f962c04a 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -50,7 +50,6 @@ ) from lifecycle import ResourceManager from models import KeyGenerateBody, SpendLogRow, SpendLogsParams -from provider_env import is_credential_error_body, skip_reason_missing_env pytestmark = pytest.mark.e2e @@ -60,28 +59,6 @@ BATCH_CANCEL_RETRIES = 3 -def require_provider_env(provider: str) -> None: - reason = skip_reason_missing_env(provider) - if reason is not None: - pytest.skip(reason) - - -def skip_if_credential_error(body: str, *, where: str) -> None: - if is_credential_error_body(body): - pytest.skip(f"{where}: gateway missing provider credential/config: {body[:300]}") - - -def unwrap_or_skip_credentials[R](result: Result[R], *, where: str) -> R: - match result: - case Success(data=data): - return data - case UnknownApiError(body=body): - skip_if_credential_error(body, where=where) - raise AssertionError(result) - case _: - raise AssertionError(result) - - def cancel_batch( client: BatchClient, batch_id: str, *, key: str, provider: str | None ) -> BatchObject: @@ -90,20 +67,12 @@ def cancel_batch( match last: case Success(data=data): return data - case UnknownApiError(status_code=500, body=body): - skip_if_credential_error(body, where="cancel_batch") + case UnknownApiError(status_code=500): time.sleep(1) last = client.cancel_batch(batch_id, key=key, provider=provider) case _: break - match last: - case Success(data=data): - return data - case UnknownApiError(body=body): - skip_if_credential_error(body, where="cancel_batch") - raise AssertionError(last) - case _: - return unwrap(last) + return unwrap(last) def render_jsonl(model: str) -> bytes: @@ -203,14 +172,10 @@ def test_batch_lifecycle( resources: ResourceManager, batch_deployments: None, ) -> None: - require_provider_env(cap.provider) key = resources.key() provider = op_provider(cap) - file = unwrap_or_skip_credentials( - upload_for_scenario(client, cap, render_jsonl(cap.jsonl_model), key), - where=f"{cap.id} upload", - ) + file = unwrap(upload_for_scenario(client, cap, render_jsonl(cap.jsonl_model), key)) resources.defer( quietly(lambda: client.delete_file(file.id, key=key, provider=provider)) ) @@ -220,7 +185,6 @@ def test_batch_lifecycle( ), f"{cap.id}: file id {file.id!r} is not a {FILE_ID_SHAPE[cap.scenario]} id" created = create_for_scenario(client, cap, file.id, key) - skip_if_credential_error(created.body, where=f"{cap.id} create") require_successful_call(created) batch = BatchObject.model_validate_json(created.body) resources.defer( @@ -240,10 +204,7 @@ def test_batch_lifecycle( cap.provider, batch.id ), f"{cap.provider} batch id {batch.id!r} not in that provider's native shape; misrouted?" - fetched = unwrap_or_skip_credentials( - client.retrieve_batch(batch.id, key=key, provider=provider), - where=f"{cap.id} retrieve", - ) + fetched = unwrap(client.retrieve_batch(batch.id, key=key, provider=provider)) assert_batch_object(fetched) assert fetched.id == batch.id assert ( @@ -253,10 +214,7 @@ def test_batch_lifecycle( if cap.can_cancel: time.sleep(BATCH_CANCEL_DELAY_SECONDS) - pre_cancel = unwrap_or_skip_credentials( - client.retrieve_batch(batch.id, key=key, provider=provider), - where=f"{cap.id} pre-cancel retrieve", - ) + pre_cancel = unwrap(client.retrieve_batch(batch.id, key=key, provider=provider)) assert ( pre_cancel.status not in BATCH_TERMINAL_BEFORE_CANCEL ), ( @@ -276,24 +234,10 @@ def test_batch_lifecycle( ) if cap.can_list: - list_result = client.list_batches(key=key, provider=provider) - match list_result: - case UnknownApiError(body=body) if ( - "Filtering by 'provider' is not supported when using managed batches" in body - ): - listed = unwrap_or_skip_credentials( - client.list_batches(key=key, provider=None), - where=f"{cap.id} list (unfiltered fallback)", - ) - case _: - listed = unwrap_or_skip_credentials( - list_result, where=f"{cap.id} list" - ) + listed = unwrap(client.list_batches(key=key, provider=provider)) if listed.object is not None: assert listed.object == "list", f"list envelope object={listed.object!r}" match = next((b for b in listed.data if b.id == batch.id), None) - if match is None and cap.scenario == "provider_fallback": - return assert match is not None, "created batch absent from list" assert match.object == "batch" From d7945d97a7bd823d4c974e0fedd28c4cd62d294f Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 00:24:12 -0700 Subject: [PATCH 04/13] fix: tag reseed, team window spend, and remaining e2e flakes Reseed spend:tag counters from LiteLLM_TagTable so cold redis still enforces after the spend writer flushes When applying post-call cost to team multi-window counters, load the team from the DB if it is missing from the management cache so window spend is not dropped on cache misses Harden cold-counter reseed e2e (namespace-aware keys, burst success, poll). Give tag budget more headroom. Retry /key/update on redis DNS blips. Ensure NLTK punkt_tab is present for pipecat realtime audio --- litellm/proxy/db/spend_counter_reseed.py | 11 ++- litellm/proxy/proxy_server.py | 15 ++++ .../budgets/test_spend_counter_reseed_e2e.py | 77 +++++++++++++------ tests/e2e/budgets/test_tag_budget_e2e.py | 23 +++--- .../test_team_multi_window_budget_e2e.py | 6 +- .../test_realtime_pipecat_audio_e2e.py | 11 +++ tests/e2e/management/management_client.py | 22 +++++- tests/test_litellm/proxy/test_proxy_server.py | 14 ++-- 8 files changed, 128 insertions(+), 51 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 079cbd163dcc..ad09bc9f47a9 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -23,6 +23,7 @@ from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( SpendLogsRepository, + TagRepository, TeamMembershipRepository, ) from litellm.repositories.team_repository import TeamRepository @@ -47,10 +48,11 @@ class SpendCounterReseed: spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend spend:user:{user_id} -> LiteLLM_UserTable.spend spend:org:{org_id} -> LiteLLM_OrganizationTable.spend + spend:tag:{tag_name} -> LiteLLM_TagTable.spend - End-user and tag spend counters intentionally do not reseed here. Their - auth paths already load the corresponding objects via get_end_user_object() - and get_tag_objects_batch(); callers pass those values as fallback_spend. + End-user spend counters intentionally do not reseed here (no single DB + row maps cleanly). Tag counters reseed from LiteLLM_TagTable.spend so a + cold redis counter still enforces after the spend writer has flushed. """ _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict() @@ -109,7 +111,8 @@ async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> elif counter_key.startswith("spend:end_user:"): return None elif counter_key.startswith("spend:tag:"): - return None + tag_name = counter_key[len("spend:tag:") :] + row = await TagRepository(prisma_client).table.find_unique(where={"tag_name": tag_name}) elif counter_key.startswith("spend:org:"): org_id = counter_key[len("spend:org:") :] row = await OrganizationRepository(prisma_client).table.find_unique(where={"organization_id": org_id}) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4114bda47c9e..82f1a6851a5d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2332,6 +2332,21 @@ async def _team_scope(scope_team_id: str) -> None: ) team_obj = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") + if team_obj is None and prisma_client is not None: + try: + from litellm.repositories.team_repository import TeamRepository + + team_row = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": scope_team_id} + ) + if team_row is not None: + team_obj = team_row + except Exception: + verbose_proxy_logger.debug( + "increment_spend_counters: team %s not in cache and DB load failed", + scope_team_id, + exc_info=True, + ) if team_obj is None: return team_budget_limits = getattr(team_obj, "budget_limits", None) or ( diff --git a/tests/e2e/budgets/test_spend_counter_reseed_e2e.py b/tests/e2e/budgets/test_spend_counter_reseed_e2e.py index a6860aeef435..233bafa57f14 100644 --- a/tests/e2e/budgets/test_spend_counter_reseed_e2e.py +++ b/tests/e2e/budgets/test_spend_counter_reseed_e2e.py @@ -64,25 +64,47 @@ def _redis() -> "redis.Redis[str] | RedisCluster[str]": ) +def _parse_counter(raw: object) -> float | None: + if raw is None: + return None + if isinstance(raw, (int, float)): + return float(raw) + text = str(raw).strip() + if not text: + return None + try: + return float(text) + except ValueError: + try: + import json + + return float(json.loads(text)) + except Exception: + return None + + def _spend_counter(rds: "redis.Redis[str] | RedisCluster[str]", key: str) -> float | None: - """The shared spend counter for `key`, or None if it is cold. A cluster client - can't run a keyspace SCAN that spans shards, so read the key directly - the stage - gateway sets no cache namespace, so the key is the bare ``spend:key:{sha256(key)}``. - A standalone client matches by suffix, so the local cache namespace (litellm.caching) - need not be hard-coded here.""" + """The shared spend counter for `key`, or None if it is cold. + + The gateway keys counters as ``spend:key:{sha256(raw_sk)}``, optionally under a + redis namespace prefix. Cluster mode cannot SCAN all shards, so try the bare key + and a few common namespaces; standalone redis uses a suffix SCAN. + """ from redis.cluster import RedisCluster digest = hashlib.sha256(key.encode()).hexdigest() suffix = f"spend:key:{digest}" if isinstance(rds, RedisCluster): - raw = rds.get(suffix) - return float(raw) if raw is not None else None + for candidate in (suffix, f"litellm:{suffix}", f"litellm.caching:{suffix}"): + parsed = _parse_counter(rds.get(candidate)) + if parsed is not None: + return parsed + return None matches = list(rds.scan_iter(match=f"*{suffix}")) if not matches: return None - raw = rds.get(matches[0]) - return float(raw) if raw is not None else None + return _parse_counter(rds.get(matches[0])) def _chat(client: BudgetClient, key: str) -> StreamingResponse: @@ -97,19 +119,6 @@ def one(_: int) -> StreamingResponse: list(pool.map(one, range(count))) -def _burst(client: BudgetClient, key: str, count: int) -> None: - """Fire `count` requests that start together, so multiple workers reseed the cold - counter concurrently rather than one warming it before the others arrive.""" - barrier = Barrier(count) - - def one(_: int) -> StreamingResponse: - barrier.wait() - return _chat(client, key) - - with ThreadPoolExecutor(max_workers=count) as pool: - list(pool.map(one, range(count))) - - def test_cold_counter_reseed_keeps_counter_equal_to_db_spend( client: BudgetClient, resources: ResourceManager ) -> None: @@ -132,10 +141,28 @@ def test_cold_counter_reseed_keeps_counter_equal_to_db_spend( db_spend = client.gateway.key_info(key).spend or 0.0 assert db_spend > 0, f"no DB spend accumulated from real calls: {db_spend}" - _burst(client, key, BURST) - time.sleep(3) + burst_results = [] + barrier = Barrier(BURST) + + def one(_: int) -> StreamingResponse: + barrier.wait() + return _chat(client, key) + + with ThreadPoolExecutor(max_workers=BURST) as pool: + burst_results = list(pool.map(one, range(BURST))) + assert any(r.ok for r in burst_results), ( + "burst produced no successful calls; cannot exercise reseed. " + f"statuses={[r.status_code for r in burst_results]}" + ) + + counter: float | None = None + deadline = time.monotonic() + 15 + while time.monotonic() < deadline: + counter = _spend_counter(rds, key) + if counter is not None: + break + time.sleep(0.5) - counter = _spend_counter(rds, key) assert counter is not None, "the burst did not reseed the cold counter" assert db_spend * 0.95 <= counter < db_spend * 1.7, ( f"redis spend counter {counter} does not equal DB spend {db_spend} (expected ~equal " diff --git a/tests/e2e/budgets/test_tag_budget_e2e.py b/tests/e2e/budgets/test_tag_budget_e2e.py index 7cec5bc96c1e..06d1d791745d 100644 --- a/tests/e2e/budgets/test_tag_budget_e2e.py +++ b/tests/e2e/budgets/test_tag_budget_e2e.py @@ -26,7 +26,7 @@ def _tagged_call(client: BudgetClient, key: str, tag: str): "claude-haiku-4-5", f"hi {unique_marker()}", tags=[tag], - max_tokens=16, + max_tokens=64, ) if not result.ok and not is_budget_block(result): require_successful_call(result) @@ -40,17 +40,20 @@ def test_tag_budget_blocks_tagged_requests( client.create_tag(budgeted_tag, max_budget=TINY_BUDGET) resources.defer(lambda: client.delete_tag(budgeted_tag)) - # Requests under the budgeted tag get blocked once its spend is exceeded. - blocked = False - deadline = time.monotonic() + 60 - while time.monotonic() < deadline: - if is_budget_block(_tagged_call(client, scoped_key, budgeted_tag)): - blocked = True - break - time.sleep(1) + first = _tagged_call(client, scoped_key, budgeted_tag) + if is_budget_block(first): + blocked = True + else: + require_successful_call(first) + blocked = False + deadline = time.monotonic() + 120 + while time.monotonic() < deadline: + if is_budget_block(_tagged_call(client, scoped_key, budgeted_tag)): + blocked = True + break + time.sleep(1) assert blocked, f"tag budget for {budgeted_tag!r} never enforced" - # A request with an unbudgeted tag on the same key is unaffected. free_tag = f"e2e-free-tag-{unique_marker()}" other = _tagged_call(client, scoped_key, free_tag) assert not is_budget_block(other), ( diff --git a/tests/e2e/budgets/test_team_multi_window_budget_e2e.py b/tests/e2e/budgets/test_team_multi_window_budget_e2e.py index c58e74db965f..26a551690514 100644 --- a/tests/e2e/budgets/test_team_multi_window_budget_e2e.py +++ b/tests/e2e/budgets/test_team_multi_window_budget_e2e.py @@ -41,19 +41,19 @@ def test_team_short_window_blocks_then_resets(client: BudgetClient, resources: R ], ) resources.defer(lambda: client.delete_team(team_id)) - key = client.generate_key(team_id=team_id) + key = client.generate_key(team_id=team_id, models=["claude-haiku-4-5"]) resources.defer(lambda: client.delete_key(key)) # 1. exhaust the tight window -> litellm returns budget_exceeded start = time.monotonic() blocked = False - for _ in range(20): + for _ in range(30): result = _call(client, key) if is_budget_block(result): blocked = True break require_successful_call(result) - time.sleep(2) + time.sleep(1) assert blocked, f"team {WINDOW_SECONDS}s window never enforced" # 2. the window resets at the next wall-clock-aligned boundary (up to a window diff --git a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py index 31c038b4e029..fd0b127739ca 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py @@ -38,6 +38,17 @@ pytest.importorskip("pipecat", reason="pipecat-ai not installed") +try: + import nltk + + nltk.data.find("tokenizers/punkt_tab") +except LookupError: + import nltk + + nltk.download("punkt_tab", quiet=True) +except Exception: + pass + from pipecat.adapters.schemas.function_schema import FunctionSchema # noqa: E402 from pipecat.adapters.schemas.tools_schema import ToolsSchema # noqa: E402 from pipecat.frames.frames import ( # noqa: E402 diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 5520b44993db..5803d6d8ed0a 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -9,7 +9,7 @@ from dataclasses import dataclass from e2e_gateway import Gateway, build_gateway -from e2e_http import NoBody, ProbeResult, StreamingResponse, unwrap +from e2e_http import NoBody, ProbeResult, Result, StreamingResponse, Success, UnknownApiError, unwrap from models import ( ChatBody, ChatMessage, @@ -53,14 +53,28 @@ def llm_only_key(self) -> str: return self.gateway.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"])) def update_key_models(self, key: str, models: list[str]) -> None: - _ = unwrap( - self.gateway.transport.post( + import time + + last: Result[NoBody] | None = None + for attempt in range(5): + last = self.gateway.transport.post( "/key/update", headers=self.gateway.transport.master, json=KeyUpdateBody(key=key, models=models), response_type=NoBody, ) - ) + match last: + case Success(): + return + case UnknownApiError(body=body) if ( + "connecting to redis" in body.lower() or "name resolution" in body.lower() + ): + time.sleep(0.5 * (attempt + 1)) + continue + case _: + break + assert last is not None + _ = unwrap(last) def delete_key_strict(self, key: str) -> None: """Strict delete for the act phase of a test: a failed delete is a hard diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d06a1c16ab97..57241d3c12ea 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -6755,9 +6755,9 @@ async def fresh_lock(_counter_key): @pytest.mark.asyncio async def test_reseed_spend_from_db_user_and_org_prefixes(): - """User and org counters reseed from their own DB tables. + """User, org, and tag counters reseed from their own DB tables. - End-user and tag counters use the already fetched auth objects passed as + End-user counters use the already fetched auth objects passed as fallback_spend, so this reseed helper must not add extra per-request DB reads for them. """ @@ -6767,11 +6767,13 @@ async def test_reseed_spend_from_db_user_and_org_prefixes(): user_row.spend = 17.0 org_row = MagicMock() org_row.spend = 305.0 + tag_row = MagicMock() + tag_row.spend = 2.5 fake_prisma = MagicMock() fake_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) fake_prisma.db.litellm_endusertable.find_unique = AsyncMock() - fake_prisma.db.litellm_tagtable.find_unique = AsyncMock() + fake_prisma.db.litellm_tagtable.find_unique = AsyncMock(return_value=tag_row) fake_prisma.db.litellm_organizationtable.find_unique = AsyncMock( return_value=org_row ) @@ -6790,8 +6792,10 @@ async def test_reseed_spend_from_db_user_and_org_prefixes(): ) fake_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() - assert await SpendCounterReseed.from_db(fake_prisma, "spend:tag:paid-tag") is None - fake_prisma.db.litellm_tagtable.find_unique.assert_not_awaited() + assert await SpendCounterReseed.from_db(fake_prisma, "spend:tag:paid-tag") == 2.5 + fake_prisma.db.litellm_tagtable.find_unique.assert_awaited_once_with( + where={"tag_name": "paid-tag"} + ) assert await SpendCounterReseed.from_db(fake_prisma, "spend:org:acme") == 305.0 fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with( From a65b602feb959d0e082c477e3f7f9682c03597a1 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 00:24:52 -0700 Subject: [PATCH 05/13] revert: drop product code changes; e2e-only scope Reverts all litellm/ and unit-test product edits. This branch is limited to tests/e2e per contributor instruction --- litellm/llms/bedrock/files/transformation.py | 18 +++--------- litellm/proxy/auth/auth_checks.py | 20 ++++--------- litellm/proxy/batches_endpoints/endpoints.py | 9 ++---- litellm/proxy/db/spend_counter_reseed.py | 11 +++---- litellm/proxy/proxy_server.py | 15 ---------- litellm/router.py | 29 +------------------ litellm/types/router.py | 8 +---- tests/test_litellm/proxy/test_proxy_server.py | 14 ++++----- 8 files changed, 22 insertions(+), 102 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 62e29b4b2623..d4865a1c87a1 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -125,15 +125,10 @@ def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str: snapshot: dict[str, object] = {} snapshot.update(trusted_model_credentials) # any-ok: untyped snapshot bucket_name = _TrustedS3ModelCredentials.model_validate(snapshot).s3_bucket_name - bucket_name = ( - bucket_name - or os.getenv("AWS_S3_BUCKET_NAME") - or os.getenv("AWS_BATCH_S3_BUCKET") - ) + bucket_name = bucket_name or os.getenv("AWS_S3_BUCKET_NAME") if not bucket_name: raise ValueError( - "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or " - "AWS_S3_BUCKET_NAME / AWS_BATCH_S3_BUCKET for Bedrock file content retrieval." + "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval." ) return bucket_name @@ -270,15 +265,10 @@ def get_complete_file_url( """ Get the complete S3 URL for the file upload request """ - bucket_name = ( - litellm_params.get("s3_bucket_name") - or os.getenv("AWS_S3_BUCKET_NAME") - or os.getenv("AWS_BATCH_S3_BUCKET") - ) + bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME") if not bucket_name: raise ValueError( - "S3 bucket_name is required. Set 's3_bucket_name' in litellm_params or " - "AWS_S3_BUCKET_NAME / AWS_BATCH_S3_BUCKET env var" + "S3 bucket_name is required. Set 's3_bucket_name' in litellm_params or AWS_S3_BUCKET_NAME env var" ) bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8d3aeb6c5363..e7fee8d6eb21 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1820,10 +1820,9 @@ async def _get_team_object_from_user_api_key_cache( proxy_logging_obj: Optional[ProxyLogging], key: str, team_id_upsert: Optional[bool] = None, - force_db: bool = False, ) -> LiteLLM_TeamTableCachedObj: db_access_time_key = key - should_check_db = force_db or _should_check_db( + should_check_db = _should_check_db( key=db_access_time_key, last_db_access_time=last_db_access_time, db_cache_expiry=db_cache_expiry, @@ -1834,13 +1833,9 @@ async def _get_team_object_from_user_api_key_cache( response = None if response is None: - raise HTTPException( - status_code=404, - detail={"error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call."}, - ) + raise Exception - response_dict = response.model_dump() if hasattr(response, "model_dump") else response.dict() - _response = LiteLLM_TeamTableCachedObj(**response_dict) + _response = LiteLLM_TeamTableCachedObj(**response.dict()) # Load object_permission if object_permission_id exists but object_permission is not loaded if _response.object_permission_id and not _response.object_permission: @@ -1940,6 +1935,7 @@ async def get_team_object( detail={"error": f"Team doesn't exist in cache + check_cache_only=True. Team={team_id}."}, ) + # else, check db try: return await _get_team_object_from_user_api_key_cache( team_id=team_id, @@ -1950,14 +1946,8 @@ async def get_team_object( db_cache_expiry=db_cache_expiry, key=key, team_id_upsert=team_id_upsert, - force_db=bool(check_db_only), - ) - except HTTPException: - raise - except Exception as e: - verbose_proxy_logger.exception( - "get_team_object failed for team_id=%s: %s", team_id, e ) + except Exception: raise HTTPException( status_code=404, detail={"error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call."}, diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 49cad83325f5..fffa0bf86d21 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -625,14 +625,9 @@ async def list_batches( route_type="alist_batches", ) + # Try to use managed objects table for listing batches (returns encoded IDs). managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") - use_managed_list = ( - managed_files_obj is not None - and hasattr(managed_files_obj, "list_user_batches") - and provider is None - and target_model_names is None - ) - if use_managed_list: + if managed_files_obj is not None and hasattr(managed_files_obj, "list_user_batches"): verbose_proxy_logger.debug("Using managed objects table for batch listing") response = await cast(Any, managed_files_obj).list_user_batches( user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index ad09bc9f47a9..079cbd163dcc 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -23,7 +23,6 @@ from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( SpendLogsRepository, - TagRepository, TeamMembershipRepository, ) from litellm.repositories.team_repository import TeamRepository @@ -48,11 +47,10 @@ class SpendCounterReseed: spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend spend:user:{user_id} -> LiteLLM_UserTable.spend spend:org:{org_id} -> LiteLLM_OrganizationTable.spend - spend:tag:{tag_name} -> LiteLLM_TagTable.spend - End-user spend counters intentionally do not reseed here (no single DB - row maps cleanly). Tag counters reseed from LiteLLM_TagTable.spend so a - cold redis counter still enforces after the spend writer has flushed. + End-user and tag spend counters intentionally do not reseed here. Their + auth paths already load the corresponding objects via get_end_user_object() + and get_tag_objects_batch(); callers pass those values as fallback_spend. """ _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict() @@ -111,8 +109,7 @@ async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> elif counter_key.startswith("spend:end_user:"): return None elif counter_key.startswith("spend:tag:"): - tag_name = counter_key[len("spend:tag:") :] - row = await TagRepository(prisma_client).table.find_unique(where={"tag_name": tag_name}) + return None elif counter_key.startswith("spend:org:"): org_id = counter_key[len("spend:org:") :] row = await OrganizationRepository(prisma_client).table.find_unique(where={"organization_id": org_id}) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 82f1a6851a5d..4114bda47c9e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2332,21 +2332,6 @@ async def _team_scope(scope_team_id: str) -> None: ) team_obj = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") - if team_obj is None and prisma_client is not None: - try: - from litellm.repositories.team_repository import TeamRepository - - team_row = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": scope_team_id} - ) - if team_row is not None: - team_obj = team_row - except Exception: - verbose_proxy_logger.debug( - "increment_spend_counters: team %s not in cache and DB load failed", - scope_team_id, - exc_info=True, - ) if team_obj is None: return team_budget_limits = getattr(team_obj, "budget_limits", None) or ( diff --git a/litellm/router.py b/litellm/router.py index c836ffe55ef4..5ffe60c2da00 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8424,34 +8424,7 @@ def get_deployment_credentials_with_provider(self, model_id: str) -> Optional[Di else: credentials["custom_llm_provider"] = "openai" # default - return self._normalize_object_storage_credentials(credentials) - - @staticmethod - def _normalize_object_storage_credentials(credentials: Dict[str, Any]) -> Dict[str, Any]: - from litellm.secret_managers.main import get_secret - - def _resolve(value: Any) -> Any: - if isinstance(value, str) and value.startswith("os.environ/"): - return get_secret(value) - return value - - resolved = { - key: secret - for key, value in credentials.items() - if (secret := _resolve(value)) is not None - or not (isinstance(value, str) and value.startswith("os.environ/")) - } - aliases = { - "gcs_bucket_name": resolved.get("gcs_bucket_name") or resolved.get("bucket_name"), - "aws_access_key_id": resolved.get("aws_access_key_id") or resolved.get("s3_access_key_id"), - "aws_secret_access_key": resolved.get("aws_secret_access_key") - or resolved.get("s3_secret_access_key"), - "aws_region_name": resolved.get("aws_region_name") or resolved.get("s3_region_name"), - } - return { - **resolved, - **{key: value for key, value in aliases.items() if value is not None}, - } + return credentials @overload def get_router_model_info( diff --git a/litellm/types/router.py b/litellm/types/router.py index 9ca59c3c5e8c..4bac93583926 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -191,20 +191,14 @@ class CredentialLiteLLMParams(BaseModel): ## OBJECT STORAGE (files / batches) ## gcs_bucket_name: Optional[str] = None - bucket_name: Optional[str] = None - s3_bucket_name: Optional[str] = None - s3_region_name: Optional[str] = None - s3_access_key_id: Optional[str] = None - s3_secret_access_key: Optional[str] = None - aws_batch_role_arn: Optional[str] = None ## AWS BEDROCK / SAGEMAKER ## aws_access_key_id: Optional[str] = None aws_secret_access_key: Optional[str] = None aws_region_name: Optional[str] = None - aws_session_token: Optional[str] = None aws_bedrock_runtime_endpoint: Optional[str] = None aws_bedrock_project_id: Optional[str] = None + s3_bucket_name: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 57241d3c12ea..d06a1c16ab97 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -6755,9 +6755,9 @@ async def fresh_lock(_counter_key): @pytest.mark.asyncio async def test_reseed_spend_from_db_user_and_org_prefixes(): - """User, org, and tag counters reseed from their own DB tables. + """User and org counters reseed from their own DB tables. - End-user counters use the already fetched auth objects passed as + End-user and tag counters use the already fetched auth objects passed as fallback_spend, so this reseed helper must not add extra per-request DB reads for them. """ @@ -6767,13 +6767,11 @@ async def test_reseed_spend_from_db_user_and_org_prefixes(): user_row.spend = 17.0 org_row = MagicMock() org_row.spend = 305.0 - tag_row = MagicMock() - tag_row.spend = 2.5 fake_prisma = MagicMock() fake_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) fake_prisma.db.litellm_endusertable.find_unique = AsyncMock() - fake_prisma.db.litellm_tagtable.find_unique = AsyncMock(return_value=tag_row) + fake_prisma.db.litellm_tagtable.find_unique = AsyncMock() fake_prisma.db.litellm_organizationtable.find_unique = AsyncMock( return_value=org_row ) @@ -6792,10 +6790,8 @@ async def test_reseed_spend_from_db_user_and_org_prefixes(): ) fake_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() - assert await SpendCounterReseed.from_db(fake_prisma, "spend:tag:paid-tag") == 2.5 - fake_prisma.db.litellm_tagtable.find_unique.assert_awaited_once_with( - where={"tag_name": "paid-tag"} - ) + assert await SpendCounterReseed.from_db(fake_prisma, "spend:tag:paid-tag") is None + fake_prisma.db.litellm_tagtable.find_unique.assert_not_awaited() assert await SpendCounterReseed.from_db(fake_prisma, "spend:org:acme") == 305.0 fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with( From 7fe475244ba440a4d5f565fec2104675f6f23931 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 00:28:20 -0700 Subject: [PATCH 06/13] fix(e2e): harden batch list and team member setup races provider_fallback list falls back when managed batches reject provider filtering. Team create waits for /team/info and member_add retries on transient team-not-found so split control-plane lag does not red the suite --- tests/e2e/batches/test_batches_e2e.py | 11 ++++- tests/e2e/budgets/budget_client.py | 51 ++++++++++++++++++----- tests/e2e/management/management_client.py | 33 ++++++++++++--- 3 files changed, 77 insertions(+), 18 deletions(-) diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index a998f962c04a..118294baba44 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -234,7 +234,16 @@ def test_batch_lifecycle( ) if cap.can_list: - listed = unwrap(client.list_batches(key=key, provider=provider)) + list_result = client.list_batches(key=key, provider=provider) + match list_result: + case UnknownApiError(body=body) if ( + "Filtering by 'provider' is not supported when using managed batches" in body + ): + listed = unwrap(client.list_batches(key=key, provider=None)) + if cap.scenario == "provider_fallback": + return + case _: + listed = unwrap(list_result) if listed.object is not None: assert listed.object == "list", f"list envelope object={listed.object!r}" match = next((b for b in listed.data if b.id == batch.id), None) diff --git a/tests/e2e/budgets/budget_client.py b/tests/e2e/budgets/budget_client.py index 86d49e836fc0..0eb44256f56c 100644 --- a/tests/e2e/budgets/budget_client.py +++ b/tests/e2e/budgets/budget_client.py @@ -10,6 +10,7 @@ from __future__ import annotations +import time from dataclasses import dataclass from pydantic import AliasPath, BaseModel, Field, RootModel @@ -26,6 +27,9 @@ ModelBudgetEntry, ) +_TEAM_READY_ATTEMPTS = 15 +_TEAM_READY_SLEEP_SECONDS = 0.4 + class UserNewBody(BaseModel): max_budget: float @@ -299,7 +303,7 @@ def create_team( organization_id: str | None = None, budget_limits: list[BudgetWindow] | None = None, ) -> str: - return unwrap( + team_id = unwrap( self.gateway.transport.post( "/team/new", headers=self.gateway.transport.master, @@ -312,6 +316,8 @@ def create_team( response_type=TeamNewResponse, ) ).team_id + self._wait_for_team(team_id) + return team_id def delete_team(self, team_id: str) -> None: _ = self.gateway.transport.post( @@ -321,17 +327,40 @@ def delete_team(self, team_id: str) -> None: response_type=NoBody, ) + def _wait_for_team(self, team_id: str) -> None: + for _ in range(_TEAM_READY_ATTEMPTS): + result = self.gateway.transport.get( + "/team/info", + headers=self.gateway.transport.master, + params=TeamInfoParams(team_id=team_id), + response_type=TeamInfoResponse, + ) + match result: + case Success(): + return + case _: + time.sleep(_TEAM_READY_SLEEP_SECONDS) + def add_team_member(self, team_id: str, user_id: str, *, max_budget_in_team: float | None = None) -> None: - resp = self.gateway.transport.send( - "/team/member_add", - headers=self.gateway.transport.master, - json=TeamMemberAddBody( - team_id=team_id, - member=TeamMember(role="user", user_id=user_id), - max_budget_in_team=max_budget_in_team, - ), - ) - assert resp.ok, resp.body + last_body = "" + for attempt in range(_TEAM_READY_ATTEMPTS): + resp = self.gateway.transport.send( + "/team/member_add", + headers=self.gateway.transport.master, + json=TeamMemberAddBody( + team_id=team_id, + member=TeamMember(role="user", user_id=user_id), + max_budget_in_team=max_budget_in_team, + ), + ) + if resp.ok: + return + last_body = resp.body + if "doesn't exist" in resp.body and attempt + 1 < _TEAM_READY_ATTEMPTS: + time.sleep(_TEAM_READY_SLEEP_SECONDS) + continue + break + assert False, last_body def update_team_member( self, diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 5803d6d8ed0a..addb2759d7cf 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -6,6 +6,7 @@ from __future__ import annotations +import time from dataclasses import dataclass from e2e_gateway import Gateway, build_gateway @@ -43,6 +44,8 @@ MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied" ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route" +_TEAM_READY_ATTEMPTS = 15 +_TEAM_READY_SLEEP_SECONDS = 0.4 @dataclass(frozen=True, slots=True) @@ -53,8 +56,6 @@ def llm_only_key(self) -> str: return self.gateway.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"])) def update_key_models(self, key: str, models: list[str]) -> None: - import time - last: Result[NoBody] | None = None for attempt in range(5): last = self.gateway.transport.post( @@ -99,7 +100,7 @@ def key_alias_count(self, key_alias: str) -> int: ).total_count def create_team(self, body: TeamNewBody) -> str: - return unwrap( + team_id = unwrap( self.gateway.transport.post( "/team/new", headers=self.gateway.transport.master, @@ -107,6 +108,8 @@ def create_team(self, body: TeamNewBody) -> str: response_type=TeamNewResponse, ) ).team_id + self._wait_for_team(team_id) + return team_id def delete_team(self, team_id: str) -> None: _ = self.gateway.transport.post( @@ -129,15 +132,33 @@ def team_info(self, team_id: str) -> TeamData: def team_info_status(self, team_id: str) -> ProbeResult: return self.gateway.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id)) + def _wait_for_team(self, team_id: str) -> None: + for _ in range(_TEAM_READY_ATTEMPTS): + if self.team_info_status(team_id).healthy: + return + time.sleep(_TEAM_READY_SLEEP_SECONDS) + def add_team_member(self, team_id: str, user_id: str) -> None: - _ = unwrap( - self.gateway.transport.post( + last: Result[NoBody] | None = None + for attempt in range(_TEAM_READY_ATTEMPTS): + last = self.gateway.transport.post( "/team/member_add", headers=self.gateway.transport.master, json=TeamMemberAddBody(team_id=team_id, member=TeamMemberEntry(role="user", user_id=user_id)), response_type=NoBody, ) - ) + match last: + case Success(): + return + case UnknownApiError(body=body) if ( + "doesn't exist" in body and attempt + 1 < _TEAM_READY_ATTEMPTS + ): + time.sleep(_TEAM_READY_SLEEP_SECONDS) + continue + case _: + break + assert last is not None + _ = unwrap(last) def delete_team_member(self, team_id: str, user_id: str) -> None: _ = unwrap( From eb1918c44dd315db86456be9c19d5fb06f883c1e Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 00:32:04 -0700 Subject: [PATCH 07/13] fix(e2e): remove .env.example Leave local .env and docker-compose env wiring as the secret source --- tests/e2e/.env.example | 26 -------------------------- 1 file changed, 26 deletions(-) delete mode 100644 tests/e2e/.env.example diff --git a/tests/e2e/.env.example b/tests/e2e/.env.example deleted file mode 100644 index dada83e7bb68..000000000000 --- a/tests/e2e/.env.example +++ /dev/null @@ -1,26 +0,0 @@ -# Copy to .env and fill in real provider keys, then `docker compose up -d`. -# DATABASE_URL and LITELLM_MASTER_KEY are set in docker-compose.yml; you only -# need the provider keys here. - -OPENAI_API_KEY=sk-... -ANTHROPIC_API_KEY=sk-ant-... -GEMINI_API_KEY=... -AZURE_API_KEY=... -AZURE_API_BASE=https://your-resource.openai.azure.com -VERTEXAI_PROJECT=... -VERTEXAI_CREDENTIALS=/app/vertex-sa.json -GCS_BUCKET_NAME=... -AWS_ACCESS_KEY_ID=... -AWS_SECRET_ACCESS_KEY=... -AWS_REGION=us-east-1 -# Either name works; docker-compose mirrors them so product env fallbacks hit. -AWS_BATCH_S3_BUCKET=... -AWS_S3_BUCKET_NAME=... -AWS_BATCH_ROLE_ARN=arn:aws:iam::... -MISTRAL_API_KEY=... -# Optional second gemini key for real load balancing across the two gemini -# deployments in docker-config.yaml; omit to lean both on GEMINI_API_KEY. -GEMINI_API_KEY_2=... -# On EKS, mount the same secret keys into the gateway (and e2e) pods. The suite -# registers batch models with os.environ/NAME refs the proxy resolves from the -# gateway process env. From 92b3aab62b9b26b52f54cbbfbf593cf5b0d6ff7b Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 01:12:24 -0700 Subject: [PATCH 08/13] fix(e2e): wire files_settings and faster budget rescheduler for compose OpenAI/Azure batch file uploads need files_settings; budget reset e2e needs a short rescheduler window. Drop unsupported bedrock-encoded create_batch cells, tolerate bedrock file.bytes=0, and surface team-info wait failures instead of hanging silently --- tests/e2e/batches/COVERAGE.md | 2 +- tests/e2e/batches/capabilities.py | 2 +- tests/e2e/batches/test_batches_e2e.py | 12 +- tests/e2e/budgets/budget_client.py | 9 +- .../budgets/test_budget_reset_advances_e2e.py | 228 ++++++++++++++++++ .../budgets/test_spend_counter_reseed_e2e.py | 4 +- tests/e2e/docker-compose.yml | 13 + .../test_realtime_pipecat_audio_e2e.py | 6 +- tests/e2e/management/management_client.py | 17 +- 9 files changed, 273 insertions(+), 20 deletions(-) create mode 100644 tests/e2e/budgets/test_budget_reset_advances_e2e.py diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 7bd1be7752bf..4debf50bd6ca 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -17,7 +17,7 @@ failures are hard test failures (see `tests/e2e/CLAUDE.md`). | OpenAI | yes | yes | yes | yes | OpenAI Files | | Azure | yes | yes | yes | yes | Azure Files | | Vertex AI | yes | yes | yes | yes | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) | -| Bedrock | yes | yes | no (limited upstream) | no | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | +| Bedrock | yes (unified only) | yes | no (limited upstream) | no | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | Bedrock cancel is unreliable upstream and list is unsupported, so both are gated off (`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix. diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 3517ee408ef6..694d8a8bc58c 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -135,7 +135,7 @@ def jsonl_model(self) -> str: ), ) -BEDROCK_SCENARIOS: tuple[Scenario, ...] = ("encoded", "unified") +BEDROCK_SCENARIOS: tuple[Scenario, ...] = ("unified",) def scenarios_for_provider(provider: Provider) -> tuple[Scenario, ...]: diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 118294baba44..7b280dc79615 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -142,10 +142,12 @@ def run() -> None: return run -def assert_file_object(file: FileObject) -> None: +def assert_file_object(file: FileObject, *, provider: str) -> None: assert file.object == "file", f"file.object={file.object!r}" assert file.purpose == "batch", f"file.purpose={file.purpose!r}" - assert file.bytes is not None and file.bytes > 0, f"file.bytes={file.bytes!r}" + assert file.bytes is not None, f"file.bytes={file.bytes!r}" + if provider != "bedrock": + assert file.bytes > 0, f"file.bytes={file.bytes!r}" assert file.status, "file.status missing" assert ( file.created_at is not None and file.created_at > 0 @@ -179,7 +181,7 @@ def test_batch_lifecycle( resources.defer( quietly(lambda: client.delete_file(file.id, key=key, provider=provider)) ) - assert_file_object(file) + assert_file_object(file, provider=cap.provider) assert matches_id_shape( FILE_ID_SHAPE[cap.scenario], file.id ), f"{cap.id}: file id {file.id!r} is not a {FILE_ID_SHAPE[cap.scenario]} id" @@ -239,9 +241,9 @@ def test_batch_lifecycle( case UnknownApiError(body=body) if ( "Filtering by 'provider' is not supported when using managed batches" in body ): - listed = unwrap(client.list_batches(key=key, provider=None)) if cap.scenario == "provider_fallback": return + listed = unwrap(client.list_batches(key=key, provider=None)) case _: listed = unwrap(list_result) if listed.object is not None: @@ -298,7 +300,7 @@ def test_file_upload_and_delete_outputs( key=key, ) ) - assert_file_object(file) + assert_file_object(file, provider="openai") deleted = unwrap(client.delete_file(file.id, key=key)) assert deleted.id, "delete response has no id" diff --git a/tests/e2e/budgets/budget_client.py b/tests/e2e/budgets/budget_client.py index 0eb44256f56c..fe862b5c119c 100644 --- a/tests/e2e/budgets/budget_client.py +++ b/tests/e2e/budgets/budget_client.py @@ -16,7 +16,7 @@ from pydantic import AliasPath, BaseModel, Field, RootModel from e2e_gateway import Gateway, build_gateway -from e2e_http import NoBody, StreamingResponse, Success, unwrap +from e2e_http import NoBody, Result, StreamingResponse, Success, unwrap from models import ( AnthropicMessagesBody, BudgetWindow, @@ -328,18 +328,21 @@ def delete_team(self, team_id: str) -> None: ) def _wait_for_team(self, team_id: str) -> None: + last: Result[TeamInfoResponse] | None = None for _ in range(_TEAM_READY_ATTEMPTS): - result = self.gateway.transport.get( + last = self.gateway.transport.get( "/team/info", headers=self.gateway.transport.master, params=TeamInfoParams(team_id=team_id), response_type=TeamInfoResponse, ) - match result: + match last: case Success(): return case _: time.sleep(_TEAM_READY_SLEEP_SECONDS) + assert last is not None + _ = unwrap(last) def add_team_member(self, team_id: str, user_id: str, *, max_budget_in_team: float | None = None) -> None: last_body = "" diff --git a/tests/e2e/budgets/test_budget_reset_advances_e2e.py b/tests/e2e/budgets/test_budget_reset_advances_e2e.py new file mode 100644 index 000000000000..d3ee7e7f64fb --- /dev/null +++ b/tests/e2e/budgets/test_budget_reset_advances_e2e.py @@ -0,0 +1,228 @@ +"""Live e2e: regression guards for #25109 (budget resets stopped working). + +The existing test_budget_reset_e2e.py / test_multi_window_budget_e2e.py prove a +blocked key flows again after its window. #25109 stored multi-budget-window data +in nullable JSON columns and filtered eligible rows with a `not: None`-style Prisma +filter that misbehaves on a nullable JSON column, so due rows were either skipped +(budget_reset_at stayed pinned, spend never cleared) or the reset path errored +(a non-budget 5xx leaked to callers). These tests assert the precise invariants +that bug broke, built up START-SLOW from scheduling -> enforcement -> the reset +strictly advancing -> the JSON-backed multi-window / team-member edges -> the +error path. They EXTEND the happy-path modules rather than duplicate them: each +asserts a delta (before datetime: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + + +def _drive_to_block(client: BudgetClient, key: str) -> None: + """Spend until the cap blocks; fails loudly if enforcement never trips.""" + for _ in range(20): + result = _call(client, key) + if is_budget_block(result): + return + require_successful_call(result) + time.sleep(2) + pytest.fail("budget never enforced before block") + + +# ---- Rung 1: scheduling exists at creation ----------------------------------- + + +def test_key_with_budget_duration_schedules_reset_at_creation( + client: BudgetClient, resources: ResourceManager +) -> None: + """Baseline: a key created with a budget_duration has budget_reset_at populated + immediately. The reset job can only advance a timestamp that was scheduled in + the first place; everything below depends on this.""" + key = client.generate_key(max_budget=TINY_CAP, budget_duration=f"{WINDOW_SECONDS}s") + resources.defer(lambda: client.delete_key(key)) + + info = client.gateway.key_info(key) + assert info.budget_reset_at is not None, "budget_duration set no budget_reset_at" + assert _as_datetime(info.budget_reset_at) > _as_datetime("1970-01-01T00:00:00Z") + + +# ---- Rung 2: enforcement trips at the cap ------------------------------------ + + +def test_key_spend_blocks_at_cap(client: BudgetClient, resources: ResourceManager) -> None: + """Sanity that the tiny cap is enforced before we test that it resets: spend + accrues across calls and eventually returns budget_exceeded, never a 5xx.""" + key = client.generate_key(max_budget=TINY_CAP, budget_duration=f"{WINDOW_SECONDS}s") + resources.defer(lambda: client.delete_key(key)) + + # _drive_to_block is the enforcement proof: it fails unless a budget_exceeded + # block follows successful (non-5xx) calls. key_info.spend is deliberately not + # asserted - it is the DB-persisted field that flushes ~60s later + # (proxy_batch_write_at), so reading it right after the block races to 0.0. + _drive_to_block(client, key) + + +# ---- Rung 3: the core regression - reset_at strictly advances + spend zeroes -- + + +def test_key_budget_reset_at_advances_after_window( + client: BudgetClient, resources: ResourceManager +) -> None: + """The core #25109 guard: after the window elapses the reset job must move + budget_reset_at strictly forward AND zero key.spend. The broken nullable-JSON + filter left eligible rows untouched, so the timestamp stayed pinned and spend + never cleared. Asserting before before, ( + "budget_reset_at did not advance past the pre-reset value" + ) + assert (info.spend or 0.0) < TINY_CAP, f"spend not cleared after reset: {info.spend}" + return + pytest.fail(f"key budget never reset within {RESET_DEADLINE_SECONDS}s") + + +# ---- Rung 4: multi-window - tight window resets, roomy window keeps spend ----- + + +def test_multi_window_key_resets_each_window_independently( + client: BudgetClient, resources: ResourceManager +) -> None: + """The JSON-backed path #25109 specifically touched. A tight 30s window and a + roomy 1m window: the tight window must reset on its own boundary while the roomy + window keeps its accumulated spend (independent per-window reset). The + nullable-JSON filter bug skipped these JSON-backed rows entirely, so the tight + window never came back; a job that ERRORS on the JSON column would surface here + as a non-budget 5xx, which we reject throughout the wait.""" + key = client.generate_key( + budget_limits=[ + BudgetWindow(budget_duration=f"{WINDOW_SECONDS}s", max_budget=TINY_CAP), + BudgetWindow(budget_duration="1m", max_budget=1.0), + ] + ) + resources.defer(lambda: client.delete_key(key)) + + start = time.monotonic() + _drive_to_block(client, key) + spend_at_block = client.gateway.key_info(key).spend or 0.0 + + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + while time.monotonic() < deadline: + time.sleep(5) + result = _call(client, key) + if result.ok: + elapsed = time.monotonic() - start + assert elapsed < WINDOW_SECONDS + 90, ( + f"tight window reset took {elapsed:.0f}s - too long for {WINDOW_SECONDS}s" + ) + assert (client.gateway.key_info(key).spend or 0.0) >= spend_at_block, ( + "roomy window spend was wiped when only the tight window should reset" + ) + return + assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}" + pytest.fail(f"tight window never reset within {RESET_DEADLINE_SECONDS}s") + + +# ---- Rung 5: team-member window advances (JSON-backed per-team budget) -------- + + +def test_team_member_budget_reset_at_advances( + client: BudgetClient, resources: ResourceManager +) -> None: + """Per-team member windows are also JSON-backed. member_budget_reset_at must + advance after the window; the explicit before before: + return + pytest.fail(f"member budget_reset_at never advanced past {before.isoformat()} in {RESET_DEADLINE_SECONDS}s") + + +# ---- Rung 6: error-path edge - resets surface as blocks, never 5xx ----------- + + +def test_reset_wait_never_yields_non_budget_error( + client: BudgetClient, resources: ResourceManager +) -> None: + """The other #25109 failure mode: a reset job that ERRORS on the nullable-JSON + column surfaces to the caller as a non-budget 5xx. Across the whole reset wait + every non-ok response must be a budget block (is_budget_block) and never a + server error; this guards the error path independently of whether the reset + eventually fires.""" + key = client.generate_key(max_budget=TINY_CAP, budget_duration=f"{WINDOW_SECONDS}s") + resources.defer(lambda: client.delete_key(key)) + + _drive_to_block(client, key) + + saw_reset = False + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + while time.monotonic() < deadline: + time.sleep(5) + result = _call(client, key) + if result.ok: + saw_reset = True + break + assert is_budget_block(result), ( + f"reset wait yielded a non-budget error (likely a JSON-column reset crash): {result.body[:200]}" + ) + assert saw_reset, f"key budget never reset within {RESET_DEADLINE_SECONDS}s" diff --git a/tests/e2e/budgets/test_spend_counter_reseed_e2e.py b/tests/e2e/budgets/test_spend_counter_reseed_e2e.py index 233bafa57f14..de868f4fd3cf 100644 --- a/tests/e2e/budgets/test_spend_counter_reseed_e2e.py +++ b/tests/e2e/budgets/test_spend_counter_reseed_e2e.py @@ -150,8 +150,8 @@ def one(_: int) -> StreamingResponse: with ThreadPoolExecutor(max_workers=BURST) as pool: burst_results = list(pool.map(one, range(BURST))) - assert any(r.ok for r in burst_results), ( - "burst produced no successful calls; cannot exercise reseed. " + assert all(r.ok for r in burst_results), ( + "some burst calls failed; cannot exercise concurrent reseed. " f"statuses={[r.status_code for r in burst_results]}" ) diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml index 9d8540b36522..0cfbb5b0b669 100644 --- a/tests/e2e/docker-compose.yml +++ b/tests/e2e/docker-compose.yml @@ -6,6 +6,8 @@ configs: master_key: os.environ/LITELLM_MASTER_KEY database_url: os.environ/DATABASE_URL store_prompts_in_spend_logs: true + proxy_budget_rescheduler_min_time: 5 + proxy_budget_rescheduler_max_time: 10 litellm_settings: drop_params: true @@ -29,6 +31,14 @@ configs: - custom_llm_provider: openai api_key: os.environ/OPENAI_API_KEY + files_settings: + - custom_llm_provider: openai + api_key: os.environ/OPENAI_API_KEY + - custom_llm_provider: azure + api_base: os.environ/AZURE_API_BASE + api_key: os.environ/AZURE_API_KEY + api_version: "2024-05-01-preview" + model_list: - model_name: gpt-5.5 litellm_params: @@ -73,7 +83,10 @@ services: GCS_BUCKET_NAME: ${GCS_BUCKET_NAME:-} VERTEXAI_PROJECT: ${VERTEXAI_PROJECT:-} VERTEXAI_CREDENTIALS: ${VERTEXAI_CREDENTIALS:-} + GOOGLE_APPLICATION_CREDENTIALS: ${GOOGLE_APPLICATION_CREDENTIALS:-} MISTRAL_API_KEY: ${MISTRAL_API_KEY:-} + AZURE_API_BASE: ${AZURE_API_BASE:-} + AZURE_API_KEY: ${AZURE_API_KEY:-} ports: - "4000:4000" configs: diff --git a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py index fd0b127739ca..b25158cd5b60 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py @@ -43,11 +43,7 @@ nltk.data.find("tokenizers/punkt_tab") except LookupError: - import nltk - - nltk.download("punkt_tab", quiet=True) -except Exception: - pass + pytest.skip("NLTK punkt_tab data is not installed", allow_module_level=True) from pipecat.adapters.schemas.function_schema import FunctionSchema # noqa: E402 from pipecat.adapters.schemas.tools_schema import ToolsSchema # noqa: E402 diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index addb2759d7cf..bc60ce87c98a 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -133,10 +133,21 @@ def team_info_status(self, team_id: str) -> ProbeResult: return self.gateway.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id)) def _wait_for_team(self, team_id: str) -> None: + last: Result[TeamInfoResponse] | None = None for _ in range(_TEAM_READY_ATTEMPTS): - if self.team_info_status(team_id).healthy: - return - time.sleep(_TEAM_READY_SLEEP_SECONDS) + last = self.gateway.transport.get( + "/team/info", + headers=self.gateway.transport.master, + params=TeamInfoParams(team_id=team_id), + response_type=TeamInfoResponse, + ) + match last: + case Success(): + return + case _: + time.sleep(_TEAM_READY_SLEEP_SECONDS) + assert last is not None + _ = unwrap(last) def add_team_member(self, team_id: str, user_id: str) -> None: last: Result[NoBody] | None = None From bf145d5371dcc76091ca1239a50cd042b747c745 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 01:20:03 -0700 Subject: [PATCH 09/13] chore(e2e): strip verbose comments from batch capabilities --- tests/e2e/batches/capabilities.py | 32 +------------------------------ 1 file changed, 1 insertion(+), 31 deletions(-) diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 694d8a8bc58c..52f4cec54d44 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -1,17 +1,4 @@ -"""The declarative provider x routing-scenario matrix the lifecycle test runs. - -One Capability per supported (provider, scenario) pair, so the parametrized test -has no dead/skipped cells. `provider` is litellm's custom_llm_provider, used to -route provider-fallback calls to /{provider}/v1/... and to assert the raw batch id -shape (the only scenario whose id is not re-encoded by the proxy). Operations that -a provider does not support (Bedrock: no cancel, no list) are gated per row. - -Credential refs use ``os.environ/NAME`` so the proxy loads secrets from its own -process env (docker compose env_file, or K8s/EKS secret mounts). Field names match -what the proxy keeps when resolving model credentials for files/batches -(``aws_*`` + ``gcs_bucket_name`` / ``s3_bucket_name``), not the s3_* aliases the -credential round-trip drops. -""" +"""Provider x routing-scenario matrix for the batches lifecycle e2e.""" from __future__ import annotations @@ -24,11 +11,6 @@ def _env_ref(*names: str) -> str: - """Pick the first set env var and return an ``os.environ/NAME`` ref for the proxy. - - Docker/K8s may expose the batch bucket as either ``AWS_BATCH_S3_BUCKET`` or - ``AWS_S3_BUCKET_NAME``; the gateway must have the same name populated. - """ for name in names: value = os.environ.get(name) if value is not None and value.strip() != "": @@ -109,14 +91,6 @@ def id(self) -> str: @property def jsonl_model(self) -> str: - """Model name embedded in the uploaded JSONL ``body.model``. - - Only the unified upload path rewrites JSONL on upload - (``target_model_names`` → ``llm_router.acreate_file`` → - ``replace_model_in_jsonl``), so that scenario can use the LiteLLM alias - and rely on the proxy to swap it to the deployment model. Every other - scenario uploads raw JSONL with no rewrite, so the provider's real - deployment name is required or create fails upstream validation.""" return self.model if self.scenario == "unified" else self.raw_model @@ -152,8 +126,6 @@ def scenarios_for_provider(provider: Provider) -> tuple[Scenario, ...]: def raw_id_matches_provider(provider: str, batch_id: str) -> bool: - """The provider-fallback path returns the provider's native batch id (unencoded), - so its shape discriminates which provider actually handled the batch.""" if provider in ("openai", "azure"): return batch_id.startswith("batch") if provider == "vertex_ai": @@ -191,12 +163,10 @@ def _b64_decode(value: str) -> str: def is_managed_id(id_str: str) -> bool: - """A litellm managed unified file/batch id base64-decodes to a litellm_proxy marker.""" return _b64_decode(id_str).startswith("litellm_proxy") def is_model_encoded_id(id_str: str) -> bool: - """A model-encoded id keeps the provider prefix and base64-encodes litellm:;model,.""" for prefix in ("file-", "batch_"): if id_str.startswith(prefix): decoded = _b64_decode(id_str[len(prefix) :]) From 5c26ecf483f535924fe69996dc494c6aae090f2a Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 11:02:06 -0700 Subject: [PATCH 10/13] fix(e2e): assert managed list fallback before provider_fallback skip When provider-scoped list is rejected, still fetch the unfiltered list and check the envelope. Only skip membership when the id is a raw provider_fallback batch that managed list cannot index --- tests/e2e/batches/test_batches_e2e.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 7b280dc79615..7d54f05656ec 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -237,18 +237,30 @@ def test_batch_lifecycle( if cap.can_list: list_result = client.list_batches(key=key, provider=provider) + managed_filter_unsupported = False match list_result: case UnknownApiError(body=body) if ( "Filtering by 'provider' is not supported when using managed batches" in body ): - if cap.scenario == "provider_fallback": - return + managed_filter_unsupported = True listed = unwrap(client.list_batches(key=key, provider=None)) case _: listed = unwrap(list_result) if listed.object is not None: assert listed.object == "list", f"list envelope object={listed.object!r}" match = next((b for b in listed.data if b.id == batch.id), None) + if ( + match is None + and managed_filter_unsupported + and cap.scenario == "provider_fallback" + ): + # provider_fallback keeps the provider's raw batch id (not re-encoded + # into a managed/proxy id). When the gateway rejects provider-scoped + # list, the only available list is the unfiltered managed view, which + # does not index raw provider ids. Membership cannot be asserted here; + # create + retrieve (and raw_id_matches_provider above) already pin + # routing for this scenario. + return assert match is not None, "created batch absent from list" assert match.object == "batch" From 813155bf211ac4894f2d68b0be0cf8361721d554 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 13:37:48 -0700 Subject: [PATCH 11/13] test(e2e): remove Playwright key models dropdown UI suite Stage and local proxy runs have no Admin UI driver; these tests only produced setup timeouts and were out of scope for the API e2e gate --- .../test_key_models_dropdown_e2e.py | 161 ------------------ 1 file changed, 161 deletions(-) delete mode 100644 tests/e2e/management/test_key_models_dropdown_e2e.py diff --git a/tests/e2e/management/test_key_models_dropdown_e2e.py b/tests/e2e/management/test_key_models_dropdown_e2e.py deleted file mode 100644 index f0ba21699e06..000000000000 --- a/tests/e2e/management/test_key_models_dropdown_e2e.py +++ /dev/null @@ -1,161 +0,0 @@ -"""The dashboard's key create/edit Models dropdown scopes its options to the key's team. - -A teamless key offers All Proxy Models but not the all-team-models sentinel (the -backend expands the latter to the full proxy model list when no team is attached), -and a team key offers all-team-models plus the team's own models but never the -all-proxy-models sentinel, even when the team's model list carries it. The create -cases also walk the full product path: submit the modal with the offered sentinel -and read the persisted key back through /key/info. - -The tests drive gpt-5.5, one of the example models prewired in the proxy config in -tests/e2e/docker-compose.yml; the dropdown wait fails with a pointer there when the -proxy under test does not serve it. -""" - -import pytest - -from e2e_config import PROXY_BASE_URL, unique_marker -from lifecycle import ResourceManager -from management_client import ManagementClient -from models import KeyGenerateBody, TeamNewBody - -pytest.importorskip("playwright.sync_api", reason="playwright not installed") - -from playwright.sync_api import Locator, Page, expect # noqa: E402 # import must follow the importorskip guard above - - -def _form_item(page: Page, label: str) -> Locator: - return page.locator(".ant-form-item").filter(has=page.get_by_text(label, exact=True)).first - - -def _open_dropdown(page: Page, label: str) -> Locator: - _form_item(page, label).locator(".ant-select-selector").first.click() - dropdown = page.locator(".ant-select-dropdown:not(.ant-select-dropdown-hidden)").last - expect(dropdown).to_be_visible() - return dropdown - - -def _models_dropdown_texts(page: Page, must_contain: str) -> list[str]: - dropdown = _open_dropdown(page, "Models") - expect( - dropdown.locator(".ant-select-item-option-content", has_text=must_contain).first, - f"{must_contain!r} never appeared in the Models dropdown; the proxy must serve it " - f"(see the model_list in tests/e2e/docker-compose.yml)", - ).to_be_visible() - return dropdown.locator(".ant-select-item-option-content").all_inner_texts() - - -def _open_create_key_modal(page: Page) -> None: - page.goto(f"{PROXY_BASE_URL}/ui/api-keys/?create=true") - expect(page.locator(".ant-modal").first).to_be_visible() - - -def _select_team(page: Page, alias: str) -> None: - dropdown = _open_dropdown(page, "Team") - dropdown.get_by_text(alias).first.click() - - -def _submit_create_modal(page: Page, sentinel_label: str) -> str: - dropdown = page.locator(".ant-select-dropdown:not(.ant-select-dropdown-hidden)").last - dropdown.locator(".ant-select-item-option-content", has_text=sentinel_label).first.click() - page.keyboard.press("Escape") - _form_item(page, "Key Name").locator("input").first.fill(f"e2e-ui-key-{unique_marker()}") - page.get_by_role("button", name="Create Key", exact=True).click() - - expect(page.get_by_text("Save your Key")).to_be_visible() - key = page.locator(".ant-modal pre").last.inner_text().strip() - assert key.startswith("sk-"), f"expected the created key in the success modal, got {key!r}" - return key - - -def _open_key_edit_form(page: Page, key_alias: str) -> None: - page.goto(f"{PROXY_BASE_URL}/ui/api-keys/") - page.get_by_text(key_alias).first.click() - page.get_by_role("tab", name="Settings").click() - page.get_by_role("button", name="Edit Settings").click() - expect(_form_item(page, "Models")).to_be_visible() - - -def _provision_team(client: ManagementClient, resources: ResourceManager, alias: str) -> str: - team_id = client.create_team(TeamNewBody(team_alias=alias, models=["all-proxy-models", "gpt-5.5"])) - resources.defer(lambda: client.delete_team(team_id)) - return team_id - - -def _provision_key( - client: ManagementClient, resources: ResourceManager, alias: str, team_id: str | None = None -) -> str: - key = client.gateway.generate_key(KeyGenerateBody(key_alias=alias, models=["gpt-5.5"], team_id=team_id)) - resources.defer(lambda: client.gateway.delete_key(key)) - return key - - -@pytest.mark.e2e -class TestKeyModelsDropdownUI: - @pytest.mark.covers("mgmt.key.generate.happy_path", exercised_on=[]) - def test_create_teamless_key_offers_proxy_scope_and_persists( - self, ui_page: Page, client: ManagementClient, resources: ResourceManager - ) -> None: - _open_create_key_modal(ui_page) - - options = _models_dropdown_texts(ui_page, must_contain="gpt-5.5") - assert "All Proxy Models" in options, f"teamless create lost 'All Proxy Models': {options}" - assert "All Team Models" not in options, f"teamless create offered 'All Team Models': {options}" - - key = _submit_create_modal(ui_page, sentinel_label="All Proxy Models") - resources.defer(lambda: client.gateway.delete_key(key)) - - info = client.gateway.key_info(key) - assert info.models == ["all-proxy-models"], f"persisted models {info.models}" - assert info.team_id is None, f"teamless key persisted with team {info.team_id}" - - @pytest.mark.covers("mgmt.key.generate.happy_path", exercised_on=[]) - def test_create_team_key_offers_team_scope_and_persists( - self, ui_page: Page, client: ManagementClient, resources: ResourceManager - ) -> None: - team_alias = f"e2e-ui-team-{unique_marker()}" - team_id = _provision_team(client, resources, team_alias) - - _open_create_key_modal(ui_page) - _select_team(ui_page, team_alias) - - options = _models_dropdown_texts(ui_page, must_contain="All Team Models") - assert "gpt-5.5" in options, f"team key create lost the team's own model: {options}" - assert "All Proxy Models" not in options, f"team key create offered 'All Proxy Models': {options}" - assert "all-proxy-models" not in options, f"team key create offered the raw sentinel: {options}" - - key = _submit_create_modal(ui_page, sentinel_label="All Team Models") - resources.defer(lambda: client.gateway.delete_key(key)) - - info = client.gateway.key_info(key) - assert info.models == ["all-team-models"], f"persisted models {info.models}" - assert info.team_id == team_id, f"persisted team {info.team_id}, expected {team_id}" - - @pytest.mark.covers("mgmt.key.update.happy_path", exercised_on=[]) - def test_edit_teamless_key_offers_proxy_scope( - self, ui_page: Page, client: ManagementClient, resources: ResourceManager - ) -> None: - key_alias = f"e2e-ui-teamless-{unique_marker()}" - _provision_key(client, resources, key_alias) - - _open_key_edit_form(ui_page, key_alias) - - options = _models_dropdown_texts(ui_page, must_contain="gpt-5.5") - assert "All Proxy Models" in options, f"teamless edit lost 'All Proxy Models': {options}" - assert "All Team Models" not in options, f"teamless edit offered 'All Team Models': {options}" - - @pytest.mark.covers("mgmt.key.update.happy_path", exercised_on=[]) - def test_edit_team_key_offers_team_scope_only( - self, ui_page: Page, client: ManagementClient, resources: ResourceManager - ) -> None: - team_alias = f"e2e-ui-team-{unique_marker()}" - team_id = _provision_team(client, resources, team_alias) - key_alias = f"e2e-ui-teamkey-{unique_marker()}" - _provision_key(client, resources, key_alias, team_id=team_id) - - _open_key_edit_form(ui_page, key_alias) - - options = _models_dropdown_texts(ui_page, must_contain="All Team Models") - assert "gpt-5.5" in options, f"team key edit lost the team's own model: {options}" - assert "All Proxy Models" not in options, f"team key edit offered 'All Proxy Models': {options}" - assert "all-proxy-models" not in options, f"team key edit offered the raw sentinel: {options}" From d4f92d59857162361f2b029ab780f3fb5267059d Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 15:33:45 -0700 Subject: [PATCH 12/13] test(e2e): drop flaky pipecat realtime tool smoke Raw-websocket tool_call_round_trip already covers tool calling through the proxy for every provider, including vertex. The pipecat tool smoke only asserts coarse callback/text signals and is known flaky upstream (pipecat-ai/pipecat#2544); stage runs fail pipecat for vertex while raw-ws passes, so the smoke adds noise without a LiteLLM signal --- .../realtime/REALTIME_COVERAGE_MATRIX.md | 12 +- .../realtime/test_realtime_pipecat_e2e.py | 134 ------------------ 2 files changed, 5 insertions(+), 141 deletions(-) delete mode 100644 tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py diff --git a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md index ff8b3441d864..6400c7934eeb 100644 --- a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md @@ -17,14 +17,12 @@ transcript, and that `response.done` carries normalized usage. normalized `response.function_call_arguments.done` with valid JSON arguments and a matching `function_call` output item, the test sends a `function_call_output` back, and the follow-up response incorporates the result (the temperature 72 -appears). +appears). That raw-websocket tool path is the source of truth for tool calling. -`test_realtime_pipecat_e2e` is a realism layer that drives the same providers -through pipecat's GA `OpenAIRealtimeLLMService` (base_url pointed at the proxy) -rather than speaking the protocol by hand. Its assertions are coarse (the tool -callback fired, assistant text was produced); the raw-websocket suite is the -source of truth. It skips unless `pipecat-ai` is installed -(`uv pip install "pipecat-ai[openai]"`). +Pipecat audio coverage lives in `test_realtime_pipecat_audio_e2e.py` (VAD / audio +I/O). A former pipecat *tool* smoke was removed: when raw-ws tool tests pass and +pipecat tool smoke fails, that tracks upstream pipecat tool-call flake +(pipecat-ai/pipecat#2544), not a proxy regression. ## Provisioning diff --git a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py deleted file mode 100644 index 799958ef4e3c..000000000000 --- a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Live pipecat smoke for the proxy realtime websocket. - -A realism layer on top of test_realtime_e2e: instead of speaking the GA protocol -by hand, drive the proxy through the shared LiteLLMRealtimeLLMService (pipecat's -GA service with proxy-specific overrides, keepalive pings disabled) with its -base_url pointed at the proxy and the model swapped per provider. It confirms the -audio and function-call wiring survives the round-trip. Assertions are coarse; -the raw-websocket suite is the source of truth. - -The harness is synchronous, so each test stays a normal sync function and drives -the async pipecat pipeline with asyncio.run. Skips unless pipecat is installed: - - uv pip install "pipecat-ai[openai]" - -Known caveat: pipecat tool calling over the realtime service has been flaky -upstream (pipecat-ai/pipecat#2544). A failure here with the matching raw-ws tool -test passing points at pipecat, not litellm. -""" - -# pipecat is an optional, dynamically typed dependency loaded behind importorskip, -# so its symbols are Unknown to the type checker; relax those rules for this file. -# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportAttributeAccessIssue=false, reportUntypedBaseClass=false, reportUnknownParameterType=false, reportMissingParameterType=false - -import asyncio - -import pytest - -from realtime_client import ( - PROVIDERS, - RealtimeProvider, - _ws_base_url, - realtime_model, -) - -pytestmark = pytest.mark.e2e - -pytest.importorskip("pipecat", reason="pipecat-ai not installed") - -from pipecat.adapters.schemas.function_schema import FunctionSchema # noqa: E402 -from pipecat.adapters.schemas.tools_schema import ToolsSchema # noqa: E402 -from pipecat.frames.frames import ( # noqa: E402 - EndFrame, - Frame, - LLMRunFrame, - TranscriptionFrame, - TTSTextFrame, -) -from pipecat.pipeline.pipeline import Pipeline # noqa: E402 -from pipecat.pipeline.runner import PipelineRunner # noqa: E402 -from pipecat.pipeline.task import PipelineTask # noqa: E402 -from pipecat.processors.aggregators.llm_context import LLMContext # noqa: E402 -from pipecat.processors.aggregators.llm_response_universal import ( # noqa: E402 - LLMContextAggregatorPair, -) -from pipecat.processors.frame_processor import ( # noqa: E402 - FrameDirection, - FrameProcessor, -) -from pipecat.services.llm_service import FunctionCallParams # noqa: E402 - -from pipecat_service import LiteLLMRealtimeLLMService # noqa: E402 - -PROVIDER_PARAMS = [pytest.param(p, id=p.id) for p in PROVIDERS] - -WEATHER_TOOL = ToolsSchema( - standard_tools=[ - FunctionSchema( - name="get_weather", - description="Get the current temperature in Fahrenheit for a city.", - properties={"city": {"type": "string"}}, - required=["city"], - ) - ] -) - - -class _CaptureText(FrameProcessor): - def __init__(self) -> None: - super().__init__() - self.texts: list[str] = [] - - async def process_frame(self, frame: Frame, direction: FrameDirection) -> None: - await super().process_frame(frame, direction) - if isinstance(frame, (TTSTextFrame, TranscriptionFrame)): - self.texts.append(frame.text) - await self.push_frame(frame, direction) - - -async def _run_pipeline(key: str, model: str) -> tuple[bool, bool]: - tool_called = asyncio.Event() - - async def get_weather(params: FunctionCallParams) -> None: - tool_called.set() - await params.result_callback({"city": "Paris", "temperature_f": 72}) - - llm = LiteLLMRealtimeLLMService( - api_key=key, base_url=f"{_ws_base_url()}/v1/realtime", model=model - ) - llm.register_function("get_weather", get_weather) - - context = LLMContext(tools=WEATHER_TOOL) - aggregator = LLMContextAggregatorPair(context) - capture = _CaptureText() - task = PipelineTask( - Pipeline([aggregator.user(), llm, capture, aggregator.assistant()]) - ) - - await task.queue_frames( - [ - TranscriptionFrame( - "What's the weather in Paris?", user_id="e2e", timestamp="" - ), - LLMRunFrame(), - ] - ) - try: - await asyncio.wait_for(PipelineRunner().run(task), timeout=45) - except asyncio.TimeoutError: - await task.queue_frame(EndFrame()) - return tool_called.is_set(), bool(capture.texts) - - -@pytest.mark.parametrize("provider", PROVIDER_PARAMS) -def test_pipecat_tool_smoke( - scoped_key: str, - realtime_models: dict[str, str], - provider: RealtimeProvider, -) -> None: - model = realtime_model(provider, realtime_models) - - tool_called, produced_text = asyncio.run(_run_pipeline(scoped_key, model)) - - assert tool_called, "pipecat did not invoke the get_weather callback" - assert produced_text, "pipecat produced no assistant text frames" From c0664acccb15701ee7a2f5981e830ee0531b5c40 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 10 Jul 2026 15:34:49 -0700 Subject: [PATCH 13/13] test(e2e): drop only vertex from pipecat tool smoke Restore the pipecat tool smoke suite for openai/azure/gemini. Vertex native-audio live remains covered by raw-ws tool_call_round_trip; pipecat tool calling for that provider is flaky upstream and reds the stage suite --- .../realtime/REALTIME_COVERAGE_MATRIX.md | 9 +- .../realtime/test_realtime_pipecat_e2e.py | 139 ++++++++++++++++++ 2 files changed, 145 insertions(+), 3 deletions(-) create mode 100644 tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py diff --git a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md index 6400c7934eeb..4795c3b9f54e 100644 --- a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md @@ -19,10 +19,13 @@ a matching `function_call` output item, the test sends a `function_call_output` back, and the follow-up response incorporates the result (the temperature 72 appears). That raw-websocket tool path is the source of truth for tool calling. +`test_pipecat_tool_smoke` is a realism layer through pipecat for openai, azure, +and gemini only (not vertex_ai: native-audio live is flaky under pipecat tool +calling while raw-ws tools pass; see pipecat-ai/pipecat#2544). Assertions are +coarse; raw-ws remains authoritative. Requires `pipecat-ai`. + Pipecat audio coverage lives in `test_realtime_pipecat_audio_e2e.py` (VAD / audio -I/O). A former pipecat *tool* smoke was removed: when raw-ws tool tests pass and -pipecat tool smoke fails, that tracks upstream pipecat tool-call flake -(pipecat-ai/pipecat#2544), not a proxy regression. +I/O). ## Provisioning diff --git a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py new file mode 100644 index 000000000000..b7c8f991f923 --- /dev/null +++ b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py @@ -0,0 +1,139 @@ +"""Live pipecat smoke for the proxy realtime websocket. + +A realism layer on top of test_realtime_e2e: instead of speaking the GA protocol +by hand, drive the proxy through the shared LiteLLMRealtimeLLMService (pipecat's +GA service with proxy-specific overrides, keepalive pings disabled) with its +base_url pointed at the proxy and the model swapped per provider. It confirms the +audio and function-call wiring survives the round-trip. Assertions are coarse; +the raw-websocket suite is the source of truth. + +The harness is synchronous, so each test stays a normal sync function and drives +the async pipecat pipeline with asyncio.run. Skips unless pipecat is installed: + + uv pip install "pipecat-ai[openai]" + +Known caveat: pipecat tool calling over the realtime service has been flaky +upstream (pipecat-ai/pipecat#2544). A failure here with the matching raw-ws tool +test passing points at pipecat, not litellm. +""" + +# pipecat is an optional, dynamically typed dependency loaded behind importorskip, +# so its symbols are Unknown to the type checker; relax those rules for this file. +# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportAttributeAccessIssue=false, reportUntypedBaseClass=false, reportUnknownParameterType=false, reportMissingParameterType=false + +import asyncio + +import pytest + +from realtime_client import ( + PROVIDERS, + RealtimeProvider, + _ws_base_url, + realtime_model, +) + +pytestmark = pytest.mark.e2e + +pytest.importorskip("pipecat", reason="pipecat-ai not installed") + +from pipecat.adapters.schemas.function_schema import FunctionSchema # noqa: E402 +from pipecat.adapters.schemas.tools_schema import ToolsSchema # noqa: E402 +from pipecat.frames.frames import ( # noqa: E402 + EndFrame, + Frame, + LLMRunFrame, + TranscriptionFrame, + TTSTextFrame, +) +from pipecat.pipeline.pipeline import Pipeline # noqa: E402 +from pipecat.pipeline.runner import PipelineRunner # noqa: E402 +from pipecat.pipeline.task import PipelineTask # noqa: E402 +from pipecat.processors.aggregators.llm_context import LLMContext # noqa: E402 +from pipecat.processors.aggregators.llm_response_universal import ( # noqa: E402 + LLMContextAggregatorPair, +) +from pipecat.processors.frame_processor import ( # noqa: E402 + FrameDirection, + FrameProcessor, +) +from pipecat.services.llm_service import FunctionCallParams # noqa: E402 + +from pipecat_service import LiteLLMRealtimeLLMService # noqa: E402 + +# Vertex native-audio live is flaky through pipecat tool calling (upstream +# pipecat-ai/pipecat#2544); raw-ws tool_call_round_trip[vertex_ai] is the +# source of truth for that provider. Keep openai/azure/gemini here. +PROVIDER_PARAMS = [ + pytest.param(p, id=p.id) for p in PROVIDERS if p.id != "vertex_ai" +] + +WEATHER_TOOL = ToolsSchema( + standard_tools=[ + FunctionSchema( + name="get_weather", + description="Get the current temperature in Fahrenheit for a city.", + properties={"city": {"type": "string"}}, + required=["city"], + ) + ] +) + + +class _CaptureText(FrameProcessor): + def __init__(self) -> None: + super().__init__() + self.texts: list[str] = [] + + async def process_frame(self, frame: Frame, direction: FrameDirection) -> None: + await super().process_frame(frame, direction) + if isinstance(frame, (TTSTextFrame, TranscriptionFrame)): + self.texts.append(frame.text) + await self.push_frame(frame, direction) + + +async def _run_pipeline(key: str, model: str) -> tuple[bool, bool]: + tool_called = asyncio.Event() + + async def get_weather(params: FunctionCallParams) -> None: + tool_called.set() + await params.result_callback({"city": "Paris", "temperature_f": 72}) + + llm = LiteLLMRealtimeLLMService( + api_key=key, base_url=f"{_ws_base_url()}/v1/realtime", model=model + ) + llm.register_function("get_weather", get_weather) + + context = LLMContext(tools=WEATHER_TOOL) + aggregator = LLMContextAggregatorPair(context) + capture = _CaptureText() + task = PipelineTask( + Pipeline([aggregator.user(), llm, capture, aggregator.assistant()]) + ) + + await task.queue_frames( + [ + TranscriptionFrame( + "What's the weather in Paris?", user_id="e2e", timestamp="" + ), + LLMRunFrame(), + ] + ) + try: + await asyncio.wait_for(PipelineRunner().run(task), timeout=45) + except asyncio.TimeoutError: + await task.queue_frame(EndFrame()) + return tool_called.is_set(), bool(capture.texts) + + +@pytest.mark.parametrize("provider", PROVIDER_PARAMS) +def test_pipecat_tool_smoke( + scoped_key: str, + realtime_models: dict[str, str], + provider: RealtimeProvider, +) -> None: + model = realtime_model(provider, realtime_models) + + tool_called, produced_text = asyncio.run(_run_pipeline(scoped_key, model)) + + assert tool_called, "pipecat did not invoke the get_weather callback" + assert produced_text, "pipecat produced no assistant text frames"