From ce735f586c65689b8539d6136822e3923169c70b Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 22:58:15 +0000 Subject: [PATCH] fix(proxy): scope Transcribe jobs to the key that started them and charge rewritten media the maximum Standard jobs are tagged litellm-owner on StartTranscriptionJob so GetTranscriptionJob and DeleteTranscriptionJob only work for the owner or a proxy admin, and account-wide operations need a proxy admin. Media rewritten after job creation is charged the eight hour maximum, and the success handler takes an injected log dispatch instead of tests patching its private method Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + .../llm_passthrough_endpoints.py | 37 +++- .../transcribe_passthrough_logging_handler.py | 113 ++++++++++- .../pass_through_endpoints/success_handler.py | 12 +- ..._transcribe_passthrough_logging_handler.py | 185 ++++++++++++++++-- .../test_llm_pass_through_endpoints.py | 107 ++++++++-- 6 files changed, 410 insertions(+), 45 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 0d681f4788e..dee3ae63a16 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1578,6 +1578,7 @@ TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS: Final = 28800 # Amazon Transcribe quota: TRANSCRIBE_MAX_MEDIA_BYTES: Final = 2 * 1024**3 # Amazon Transcribe quota: maximum audio file size TRANSCRIBE_MEDIA_DOWNLOAD_CONCURRENCY: Final = 1 TRANSCRIBE_MEDIA_FETCH_ATTEMPTS: Final = 3 +TRANSCRIBE_MEDIA_LAST_MODIFIED_TOLERANCE_SECONDS: Final = 1.0 # S3 Last-Modified carries whole seconds only TRANSCRIBE_MEASURABLE_MEDIA_FORMATS: Final = frozenset({"flac", "mp3", "ogg", "wav"}) # what libsndfile can read BATCH_STATUS_POLL_INTERVAL_SECONDS: Final = int(os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600)) # 1 hour diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 3115faca30f..46ff0266d0c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1329,16 +1329,26 @@ async def transcribe_proxy_route( """ Pass-through for the Amazon Transcribe API, e.g. `POST /transcribe/StartTranscriptionJob`. - The request body is forwarded as-is to the AWS JSON 1.1 API and signed with SigV4 - using the proxy's AWS credentials. Streaming transcription (`transcribestreaming`) - uses a separate HTTP/2 event-stream protocol and is not served by this route. + The request body is forwarded to the AWS JSON 1.1 API and signed with SigV4 using the + proxy's AWS credentials. Standard jobs are tagged with the calling key's owner so that + only that owner (or a proxy admin) can read or delete them; account-wide operations + such as ListTranscriptionJobs are limited to proxy admins. Streaming transcription + (`transcribestreaming`) uses a separate HTTP/2 event-stream protocol and is not served + by this route. [Docs](https://docs.litellm.ai/docs/pass_through/transcribe) """ from .llm_provider_handlers.transcribe_passthrough_logging_handler import ( TRANSCRIBE_CUSTOM_LLM_PROVIDER, + TRANSCRIBE_OWNED_JOB_OPERATIONS, + TRANSCRIBE_PRICED_OPERATION, TRANSCRIBE_TARGET_PREFIX, + TranscribeRefusal, + transcribe_admin_only_refusal, transcribe_cost_per_second, + transcribe_job_access_refusal, + transcribe_job_lookup, + transcribe_owned_start_request, transcribe_supported_operations, transcribe_unpriceable_request_reason, ) @@ -1371,6 +1381,23 @@ async def transcribe_proxy_route( unpriceable_reason: Final = transcribe_unpriceable_request_reason(operation, data, transcribe_cost_per_second()) if unpriceable_reason is not None: raise HTTPException(status_code=400, detail=unpriceable_reason) + admin_only_refusal: Final = transcribe_admin_only_refusal(operation, user_api_key_dict) + if admin_only_refusal is not None: + raise HTTPException(status_code=admin_only_refusal.status_code, detail=admin_only_refusal.detail) + request_body: Final = ( + transcribe_owned_start_request(data, user_api_key_dict) if operation == TRANSCRIBE_PRICED_OPERATION else data + ) + if isinstance(request_body, TranscribeRefusal): + raise HTTPException(status_code=request_body.status_code, detail=request_body.detail) + access_refusal: Final = ( + await transcribe_job_access_refusal( + data.get("TranscriptionJobName"), user_api_key_dict, transcribe_job_lookup(aws_region_name) + ) + if operation in TRANSCRIBE_OWNED_JOB_OPERATIONS + else None + ) + if access_refusal is not None: + raise HTTPException(status_code=access_refusal.status_code, detail=access_refusal.detail) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post @@ -1381,7 +1408,7 @@ async def transcribe_proxy_route( service_name="transcribe", aws_region_name=aws_region_name, url=target_url, - body=json.dumps(data), + body=json.dumps(request_body), headers=MappingProxyType( { "Content-Type": "application/x-amz-json-1.1", @@ -1396,7 +1423,7 @@ async def transcribe_proxy_route( custom_headers=prepped.headers, custom_llm_provider=TRANSCRIBE_CUSTOM_LLM_PROVIDER, ) - setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, data) + setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, request_body) setattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, prepped.body) return await endpoint_func(request, fastapi_response, user_api_key_dict) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py index 763b2437523..d76cd5117f5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py @@ -3,7 +3,9 @@ import json import math import tempfile from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass from datetime import datetime +from email.utils import parsedate_to_datetime from functools import lru_cache, partial from pathlib import Path from types import MappingProxyType @@ -24,6 +26,7 @@ from litellm.constants import ( TRANSCRIBE_MEASURABLE_MEDIA_FORMATS, TRANSCRIBE_MEDIA_DOWNLOAD_CONCURRENCY, TRANSCRIBE_MEDIA_FETCH_ATTEMPTS, + TRANSCRIBE_MEDIA_LAST_MODIFIED_TOLERANCE_SECONDS, ) from litellm.litellm_core_utils.audio_utils.utils import calculate_request_duration from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix @@ -32,7 +35,16 @@ from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client -from litellm.proxy._types import PassThroughEndpointLoggingResultValues, PassThroughEndpointLoggingTypedDict +from litellm.proxy._types import ( + PassThroughEndpointLoggingResultValues, + PassThroughEndpointLoggingTypedDict, + UserAPIKeyAuth, +) +from litellm.proxy.common_utils.resource_ownership import ( + get_primary_resource_owner_scope, + is_proxy_admin, + user_can_access_resource_owner, +) from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.utils import StandardPassThroughResponseObject @@ -45,9 +57,11 @@ TRANSCRIBE_UNPRICED_OPERATIONS: Final = frozenset( ) TRANSCRIBE_SURCHARGE_MEMBERS: Final = ("ContentRedaction", "ToxicityDetection") TRANSCRIBE_TERMINAL_JOB_STATUSES: Final = frozenset({"COMPLETED", "FAILED"}) +TRANSCRIBE_OWNER_TAG: Final = "litellm-owner" +TRANSCRIBE_OWNED_JOB_OPERATIONS: Final = frozenset({"GetTranscriptionJob", "DeleteTranscriptionJob"}) JobLookup: TypeAlias = Callable[[str], Awaitable[Mapping[str, object]]] # mutable-ok: Callable parameter syntax -MediaDurationProbe: TypeAlias = Callable[[str], Awaitable[float | None]] # mutable-ok: Callable parameter syntax +MediaDurationProbe: TypeAlias = Callable[[str, float], Awaitable[float | None]] # mutable-ok: Callable parameter syntax JobPricer: TypeAlias = Callable[[str, str, float], Awaitable[float]] # mutable-ok: Callable parameter syntax @@ -60,10 +74,18 @@ class _MediaRef(BaseModel): MediaFileUri: str | None = None +class _JobTag(BaseModel): + model_config = ConfigDict(frozen=True) + Key: str | None = None + Value: str | None = None + + class _TranscriptionJob(BaseModel): model_config = ConfigDict(frozen=True) TranscriptionJobStatus: str | None = None + CreationTime: float | None = None Media: _MediaRef | None = None + Tags: tuple[_JobTag, ...] = () class _GetTranscriptionJobResponse(BaseModel): @@ -77,6 +99,13 @@ class _PricedCostMapEntry(BaseModel): _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_JSON_OBJECTS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +@dataclass(frozen=True, slots=True) +class TranscribeRefusal: + status_code: int + detail: str class PassThroughLogDispatch(Protocol): @@ -178,6 +207,60 @@ def requested_media_format(request_body: Mapping[str, object]) -> str | None: return suffix.lower() if dot else None +def transcribe_admin_only_refusal(operation: str, user_api_key_dict: UserAPIKeyAuth) -> TranscribeRefusal | None: + if ( + operation == TRANSCRIBE_PRICED_OPERATION + or operation in TRANSCRIBE_OWNED_JOB_OPERATIONS + or is_proxy_admin(user_api_key_dict) + ): + return None + return TranscribeRefusal( + 403, + f"{operation} reaches every Amazon Transcribe resource in the AWS account, so only a proxy admin may call it;" + f" other keys may {TRANSCRIBE_PRICED_OPERATION} and {' or '.join(sorted(TRANSCRIBE_OWNED_JOB_OPERATIONS))}" + " for the jobs they started", + ) + + +def transcribe_owned_start_request( + request_body: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> dict[str, object] | TranscribeRefusal: + owner: Final = get_primary_resource_owner_scope(user_api_key_dict) + if owner is None: + return TranscribeRefusal(400, "The calling key has no identity to record as the owner of the transcription job") + try: + tags: Final = _JSON_OBJECTS.validate_python(request_body.get("Tags", ())) + except ValidationError: + return TranscribeRefusal(400, "Tags must be a list of objects with Key and Value members") + if any(tag.get("Key") == TRANSCRIBE_OWNER_TAG for tag in tags): + return TranscribeRefusal( + 400, f"The {TRANSCRIBE_OWNER_TAG} tag is assigned by LiteLLM and cannot be supplied by the caller" + ) + owner_tag: Final = _JobTag(Key=TRANSCRIBE_OWNER_TAG, Value=owner).model_dump() + return {**request_body, "Tags": (*tags, owner_tag)} # mutable-ok: json.dumps and the body state key take a dict + + +async def transcribe_job_access_refusal( + job_name: object, user_api_key_dict: UserAPIKeyAuth, get_job: JobLookup +) -> TranscribeRefusal | None: + if is_proxy_admin(user_api_key_dict): + return None + if not isinstance(job_name, str): + return TranscribeRefusal(400, "TranscriptionJobName must be a string") + not_found: Final = TranscribeRefusal( + 404, f"No transcription job named {job_name} was started through this proxy by the calling key" + ) + try: + job: Final = _GetTranscriptionJobResponse.model_validate(await get_job(job_name)).TranscriptionJob + except Exception as e: # noqa: BLE001 # a job that cannot be read cannot be shown to belong to the caller + verbose_proxy_logger.warning("Looking up Transcribe job %s for an ownership check failed: %s", job_name, e) + return not_found + owner: Final = ( + next((tag.Value for tag in job.Tags if tag.Key == TRANSCRIBE_OWNER_TAG), None) if job is not None else None + ) + return None if user_can_access_resource_owner(owner, user_api_key_dict) else not_found + + def transcription_job_cost(audio_seconds: float, cost_per_second: float) -> float: return math.ceil(audio_seconds) * cost_per_second @@ -211,13 +294,14 @@ async def await_transcription_job( async def measure_media_seconds( media_uri: str, + job_created_at: float, media_seconds: MediaDurationProbe, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, attempts: int = TRANSCRIBE_MEDIA_FETCH_ATTEMPTS, ) -> float | None: for attempt in range(1, attempts + 1): try: - return await media_seconds(media_uri) + return await media_seconds(media_uri, job_created_at) except Exception as e: # noqa: BLE001 # the media is retried, then charged at the maximum if still unreadable verbose_proxy_logger.warning("Measuring Transcribe media %s failed (attempt %d): %s", media_uri, attempt, e) if attempt < attempts: @@ -236,7 +320,9 @@ async def price_transcription_job( """ Amazon Transcribe bills every second of the media file, silence included, and reports no duration itself, so the job is polled to completion and the media it transcribed is measured. - Anything that stops the duration from being read is charged as the longest media AWS accepts. + The measurement only counts when the object has not been rewritten since the job was created, + which is what ties it to the bytes Transcribe read. Anything that stops the duration from + being read is charged as the longest media AWS accepts. """ job: Final = await await_transcription_job(job_name, get_job, sleep=sleep, max_attempts=max_attempts) if job is None: @@ -245,9 +331,9 @@ async def price_transcription_job( if job.TranscriptionJobStatus == "FAILED": return 0.0 media_uri: Final = job.Media.MediaFileUri if job.Media is not None else None - if media_uri is None: + if media_uri is None or job.CreationTime is None: return transcribe_max_job_cost(cost_per_second) - audio_seconds: Final = await measure_media_seconds(media_uri, media_seconds, sleep=sleep) + audio_seconds: Final = await measure_media_seconds(media_uri, job.CreationTime, media_seconds, sleep=sleep) if audio_seconds is None: return transcribe_max_job_cost(cost_per_second) return transcription_job_cost(audio_seconds, cost_per_second) @@ -303,6 +389,14 @@ def s3_media_url(media_uri: str, aws_region_name: str) -> str | None: return f"https://{bucket}.s3.{aws_region_name}.{dns_suffix}/{quote(key)}" +def media_predates_job(headers: Mapping[str, str], job_created_at: float) -> bool: + try: + modified_at: Final = parsedate_to_datetime(headers["last-modified"]).timestamp() + except (KeyError, TypeError, ValueError): + return False + return modified_at <= job_created_at + TRANSCRIBE_MEDIA_LAST_MODIFIED_TOLERANCE_SECONDS + + async def write_media_within_limit(response: httpx.Response, media_file: IO[bytes], max_bytes: int) -> bool: if int(response.headers.get("content-length", "0")) > max_bytes: return False @@ -325,7 +419,7 @@ def transcribe_media_duration_probe(aws_region_name: str, download_slots: asynci S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) return dict(aws_request.prepare().headers.items()) # mutable-ok: httpx request headers take a dict - async def media_seconds(media_uri: str) -> float | None: + async def media_seconds(media_uri: str, job_created_at: float) -> float | None: url: Final = s3_media_url(media_uri, aws_region_name) if url is None: return None @@ -335,6 +429,11 @@ def transcribe_media_duration_probe(aws_region_name: str, download_slots: asynci with tempfile.NamedTemporaryFile() as media_file: async with client.stream("GET", url, headers=headers) as response: _ = response.raise_for_status() + if not media_predates_job(response.headers, job_created_at): + verbose_proxy_logger.warning( + "Transcribe media %s was rewritten after the job was created, charging maximum", media_uri + ) + return None if not await write_media_within_limit(response, media_file, TRANSCRIBE_MAX_MEDIA_BYTES): verbose_proxy_logger.warning( "Transcribe media %s exceeds the size cap, charging maximum", media_uri diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 43b9355e5b5..59d853a3df8 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -29,6 +29,7 @@ from .llm_provider_handlers.gemini_passthrough_logging_handler import ( ) from .llm_provider_handlers.transcribe_passthrough_logging_handler import ( TRANSCRIBE_CUSTOM_LLM_PROVIDER, + PassThroughLogDispatch, TranscribePassthroughLoggingHandler, ) from .llm_provider_handlers.vertex_passthrough_logging_handler import ( @@ -52,10 +53,15 @@ def _safe_response_text(httpx_response: httpx.Response) -> str: class PassThroughEndpointLogging: - def __init__(self, transcribe_handler: TranscribePassthroughLoggingHandler | None = None): + def __init__( + self, + transcribe_handler: TranscribePassthroughLoggingHandler | None = None, + log_dispatch: PassThroughLogDispatch | None = None, + ): self.transcribe_passthrough_logging_handler: Final = ( transcribe_handler if transcribe_handler is not None else TranscribePassthroughLoggingHandler() ) + self._log_dispatch: Final = log_dispatch if log_dispatch is not None else self._handle_logging self.TRACKED_VERTEX_METHOD_ROUTES = ( "generateContent", "streamGenerateContent", @@ -351,7 +357,7 @@ class PassThroughEndpointLogging: end_time=end_time, cache_hit=cache_hit, request_body=request_body, - log=self._handle_logging, + log=self._log_dispatch, standard_pass_through_logging_payload=passthrough_logging_payload, **kwargs, ) @@ -385,7 +391,7 @@ class PassThroughEndpointLogging: kwargs=kwargs, ) - await self._handle_logging( + await self._log_dispatch( logging_obj=logging_obj, standard_logging_response_object=standard_logging_response_object, result=result, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_transcribe_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_transcribe_passthrough_logging_handler.py index 5765900f444..4f28feafe6e 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_transcribe_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_transcribe_passthrough_logging_handler.py @@ -7,13 +7,20 @@ import httpx import pytest import litellm +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_provider_handlers.transcribe_passthrough_logging_handler import ( TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS, + TRANSCRIBE_OWNER_TAG, TranscribePassthroughLoggingHandler, + TranscribeRefusal, + media_predates_job, price_transcription_job, requested_media_format, s3_media_url, + transcribe_admin_only_refusal, transcribe_cost_per_second, + transcribe_job_access_refusal, + transcribe_owned_start_request, transcribe_supported_operations, transcribe_unpriceable_request_reason, write_media_within_limit, @@ -46,23 +53,27 @@ async def _no_sleep(_: float) -> None: MEDIA_URI = "s3://b/a.wav" +CREATED_AT = 1_789_682_363.696 -def _job(status: str, media_uri: str | None = MEDIA_URI) -> dict[str, object]: +def _job( + status: str, media_uri: str | None = MEDIA_URI, created_at: float | None = CREATED_AT, **members: object +) -> dict[str, object]: media = {"Media": {"MediaFileUri": media_uri}} if media_uri else {} - return {"TranscriptionJob": {"TranscriptionJobStatus": status, **media}} + created = {"CreationTime": created_at} if created_at is not None else {} + return {"TranscriptionJob": {"TranscriptionJobStatus": status, **media, **created, **members}} -async def _no_media(uri: str) -> float | None: +async def _no_media(uri: str, created_at: float) -> float | None: raise AssertionError("the media must not be measured on this path") def _media_probe(*durations: float | None | Exception): remaining = list(durations) - measured: list[str] = [] + measured: list[tuple[str, float]] = [] - async def media_seconds(uri: str) -> float | None: - measured.append(uri) + async def media_seconds(uri: str, created_at: float) -> float | None: + measured.append((uri, created_at)) outcome = remaining.pop(0) if len(remaining) > 1 else remaining[0] if isinstance(outcome, Exception): raise outcome @@ -266,7 +277,7 @@ class TestPriceTranscriptionJob: assert cost == pytest.approx(18 * COST_PER_SECOND) assert seen == ["job-1", "job-1", "job-1"] - assert measured == [MEDIA_URI] + assert measured == [(MEDIA_URI, CREATED_AT)] @pytest.mark.asyncio async def test_a_failed_poll_is_retried_instead_of_ending_pricing(self): @@ -310,7 +321,7 @@ class TestPriceTranscriptionJob: cost = await price_transcription_job("job-1", COST_PER_SECOND, get_job, media_seconds, sleep=_no_sleep) assert cost == pytest.approx(TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS * COST_PER_SECOND) - assert measured == [MEDIA_URI] + assert measured == [(MEDIA_URI, CREATED_AT)] @pytest.mark.asyncio async def test_media_fetch_is_retried_then_charged_the_maximum(self): @@ -340,6 +351,157 @@ class TestPriceTranscriptionJob: assert cost == pytest.approx(TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS * COST_PER_SECOND) + @pytest.mark.asyncio + async def test_completed_job_without_creation_time_is_charged_the_maximum_unmeasured(self): + get_job, _ = _sequence(_job("COMPLETED", created_at=None)) + media_seconds, measured = _media_probe(60.0) + + cost = await price_transcription_job("job-1", COST_PER_SECOND, get_job, media_seconds, sleep=_no_sleep) + + assert cost == pytest.approx(TRANSCRIBE_MAX_MEDIA_DURATION_SECONDS * COST_PER_SECOND) + assert measured == [] + + +class TestMediaPredatesJob: + LAST_MODIFIED = "Thu, 17 Sep 2026 17:45:00 GMT" + LAST_MODIFIED_EPOCH = 1_789_667_100.0 + + def test_object_written_before_the_job_counts(self): + assert media_predates_job(httpx.Headers({"Last-Modified": self.LAST_MODIFIED}), self.LAST_MODIFIED_EPOCH + 30) + + def test_object_written_in_the_same_second_as_the_job_counts(self): + assert media_predates_job(httpx.Headers({"Last-Modified": self.LAST_MODIFIED}), self.LAST_MODIFIED_EPOCH - 0.4) + + def test_object_rewritten_after_the_job_does_not_count(self): + assert not media_predates_job( + httpx.Headers({"Last-Modified": self.LAST_MODIFIED}), self.LAST_MODIFIED_EPOCH - 30 + ) + + @pytest.mark.parametrize("headers", [{}, {"Last-Modified": "yesterday"}]) + def test_unknown_modification_time_does_not_count(self, headers: dict[str, str]): + assert not media_predates_job(httpx.Headers(headers), self.LAST_MODIFIED_EPOCH + 30) + + +VIRTUAL_KEY = UserAPIKeyAuth(api_key="hashed-key-a", user_id="user-a", team_id="team-a") +OTHER_VIRTUAL_KEY = UserAPIKeyAuth(api_key="hashed-key-b", user_id="user-b", team_id="team-b") +ADMIN_KEY = UserAPIKeyAuth(api_key="hashed-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + +class TestTranscribeAdminOnlyRefusal: + @pytest.mark.parametrize("operation", ["StartTranscriptionJob", "GetTranscriptionJob", "DeleteTranscriptionJob"]) + def test_job_scoped_operations_are_open_to_virtual_keys(self, operation: str): + assert transcribe_admin_only_refusal(operation, VIRTUAL_KEY) is None + + @pytest.mark.parametrize("operation", ["ListTranscriptionJobs", "ListVocabularies", "DeleteVocabulary"]) + def test_account_wide_operations_are_refused_for_virtual_keys(self, operation: str): + refusal = transcribe_admin_only_refusal(operation, VIRTUAL_KEY) + + assert refusal is not None + assert refusal.status_code == 403 + assert operation in refusal.detail + + @pytest.mark.parametrize("operation", ["ListTranscriptionJobs", "DeleteVocabulary"]) + def test_account_wide_operations_are_open_to_proxy_admins(self, operation: str): + assert transcribe_admin_only_refusal(operation, ADMIN_KEY) is None + + +class TestTranscribeOwnedStartRequest: + def test_the_caller_identity_is_appended_to_the_job_tags(self): + body = {"TranscriptionJobName": "j", "Tags": [{"Key": "env", "Value": "qa"}]} + + owned = transcribe_owned_start_request(body, VIRTUAL_KEY) + + assert owned == { + "TranscriptionJobName": "j", + "Tags": ({"Key": "env", "Value": "qa"}, {"Key": TRANSCRIBE_OWNER_TAG, "Value": "user-a"}), + } + assert body == {"TranscriptionJobName": "j", "Tags": [{"Key": "env", "Value": "qa"}]} + + def test_a_request_without_tags_gets_the_owner_tag(self): + owned = transcribe_owned_start_request({"TranscriptionJobName": "j"}, VIRTUAL_KEY) + + assert owned == {"TranscriptionJobName": "j", "Tags": ({"Key": TRANSCRIBE_OWNER_TAG, "Value": "user-a"},)} + + def test_the_caller_cannot_supply_the_owner_tag(self): + owned = transcribe_owned_start_request( + {"TranscriptionJobName": "j", "Tags": [{"Key": TRANSCRIBE_OWNER_TAG, "Value": "user-b"}]}, VIRTUAL_KEY + ) + + assert isinstance(owned, TranscribeRefusal) + assert owned.status_code == 400 + + @pytest.mark.parametrize("tags", ["env=qa", ["env"], {"Key": "env"}]) + def test_malformed_tags_are_refused(self, tags: object): + owned = transcribe_owned_start_request({"TranscriptionJobName": "j", "Tags": tags}, VIRTUAL_KEY) + + assert isinstance(owned, TranscribeRefusal) + assert owned.status_code == 400 + + def test_a_key_without_any_identity_is_refused(self): + owned = transcribe_owned_start_request({"TranscriptionJobName": "j"}, UserAPIKeyAuth()) + + assert isinstance(owned, TranscribeRefusal) + assert owned.status_code == 400 + + +def _tagged(owner: str | None) -> dict[str, object]: + tags = {"Tags": [{"Key": TRANSCRIBE_OWNER_TAG, "Value": owner}]} if owner is not None else {} + return _job("COMPLETED", **tags) + + +class TestTranscribeJobAccessRefusal: + @pytest.mark.asyncio + async def test_the_key_that_started_the_job_may_read_it(self): + get_job, seen = _sequence(_tagged("user-a")) + + assert await transcribe_job_access_refusal("job-1", VIRTUAL_KEY, get_job) is None + assert seen == ["job-1"] + + @pytest.mark.asyncio + async def test_a_job_started_by_another_key_is_reported_missing(self): + get_job, _ = _sequence(_tagged("user-b")) + + refusal = await transcribe_job_access_refusal("job-1", VIRTUAL_KEY, get_job) + + assert refusal is not None + assert refusal.status_code == 404 + + @pytest.mark.asyncio + async def test_a_job_started_outside_the_proxy_is_reported_missing(self): + get_job, _ = _sequence(_tagged(None)) + + refusal = await transcribe_job_access_refusal("job-1", OTHER_VIRTUAL_KEY, get_job) + + assert refusal is not None + assert refusal.status_code == 404 + + @pytest.mark.asyncio + async def test_a_job_that_cannot_be_looked_up_is_reported_missing(self): + async def get_job(job_name: str) -> dict[str, object]: + raise httpx.HTTPStatusError("boom", request=MagicMock(), response=MagicMock()) + + refusal = await transcribe_job_access_refusal("job-1", VIRTUAL_KEY, get_job) + + assert refusal is not None + assert refusal.status_code == 404 + + @pytest.mark.asyncio + async def test_a_non_string_job_name_is_refused_before_any_lookup(self): + get_job, seen = _sequence(_tagged("user-a")) + + refusal = await transcribe_job_access_refusal(["job-1"], VIRTUAL_KEY, get_job) + + assert refusal is not None + assert refusal.status_code == 400 + assert seen == [] + + @pytest.mark.asyncio + async def test_a_proxy_admin_reads_any_job_without_a_lookup(self): + get_job, seen = _sequence(_tagged("user-b")) + + assert await transcribe_job_access_refusal("job-1", ADMIN_KEY, get_job) is None + assert seen == [] + class TestTranscribePassthroughHandler: def test_records_model_provider_and_the_given_cost(self): @@ -453,13 +615,14 @@ class TestStartTranscriptionJobIsLoggedAtJobCost: scheduled.append(job_name) return 0.0 - logging = PassThroughEndpointLogging(TranscribePassthroughLoggingHandler(job_pricer=job_pricer)) immediate: list[dict[str, object]] = [] - async def handle_logging(**kwargs: object) -> None: + async def log_dispatch(**kwargs: object) -> None: immediate.append(kwargs) - logging._handle_logging = handle_logging # rebind-ok: the shared dispatch is the observable under test + logging = PassThroughEndpointLogging( + TranscribePassthroughLoggingHandler(job_pricer=job_pricer), log_dispatch=log_dispatch + ) await logging.pass_through_async_success_handler( httpx_response=_make_response("StartTranscriptionJob"), diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 85c21e7ee60..535a5fc826c 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -5287,10 +5287,17 @@ def transcribe_client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]: monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() - monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual")) + monkeypatch.setitem( + app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual", user_id="user-a") + ) yield TestClient(app) +def _owned_job(owner: str | None, status: str = "COMPLETED") -> dict[str, object]: + tags = {"Tags": [{"Key": "litellm-owner", "Value": owner}]} if owner is not None else {} + return {"TranscriptionJob": {"TranscriptionJobName": "litellm-job-1", "TranscriptionJobStatus": status, **tags}} + + class TestTranscribeProxyRoute: START_JOB_BODY: Final = MappingProxyType( { @@ -5299,6 +5306,7 @@ class TestTranscribeProxyRoute: "Media": {"MediaFileUri": "s3://bucket/audio.wav"}, } ) + OWNER_TAG: Final = MappingProxyType({"Key": "litellm-owner", "Value": "user-a"}) def test_signs_and_forwards_start_transcription_job(self, transcribe_client: TestClient) -> None: upstream_body = { @@ -5317,17 +5325,27 @@ class TestTranscribeProxyRoute: assert targets[0] == "Transcribe.StartTranscriptionJob" assert set(targets[1:]) <= {"Transcribe.GetTranscriptionJob"} sent = route.calls[0].request - assert json.loads(sent.content) == dict(self.START_JOB_BODY) + assert json.loads(sent.content) == {**dict(self.START_JOB_BODY), "Tags": [dict(self.OWNER_TAG)]} assert sent.headers["content-type"] == "application/x-amz-json-1.1" assert sent.headers["authorization"].startswith("AWS4-HMAC-SHA256 Credential=test-access-key/") assert "/us-west-2/transcribe/aws4_request" in sent.headers["authorization"] assert "x-amz-date" in sent.headers + def test_the_caller_cannot_forge_the_owner_tag(self, transcribe_client: TestClient) -> None: + with respx.mock(assert_all_called=False) as upstream: + route = upstream.post(TRANSCRIBE_UPSTREAM) + response = transcribe_client.post( + "/transcribe/StartTranscriptionJob", + json={**dict(self.START_JOB_BODY), "Tags": [{"Key": "litellm-owner", "Value": "user-b"}]}, + ) + + assert response.status_code == 400 + assert "litellm-owner" in response.json()["detail"] + assert not route.called + def test_sdk_route_reads_operation_from_x_amz_target_and_resigns(self, transcribe_client: TestClient) -> None: with respx.mock(assert_all_called=True) as upstream: - route = upstream.post(TRANSCRIBE_UPSTREAM).mock( - return_value=httpx.Response(200, json={"TranscriptionJob": {"TranscriptionJobStatus": "COMPLETED"}}) - ) + route = upstream.post(TRANSCRIBE_UPSTREAM).mock(return_value=httpx.Response(200, json=_owned_job("user-a"))) response = transcribe_client.post( "/transcribe", json={"TranscriptionJobName": "litellm-job-1"}, @@ -5338,21 +5356,76 @@ class TestTranscribeProxyRoute: }, ) - assert (response.status_code, response.json()) == ( - 200, - {"TranscriptionJob": {"TranscriptionJobStatus": "COMPLETED"}}, - ) + assert (response.status_code, response.json()) == (200, _owned_job("user-a")) + assert [call.request.headers["x-amz-target"] for call in route.calls] == ["Transcribe.GetTranscriptionJob"] * 2 sent = route.calls.last.request - assert sent.headers["x-amz-target"] == "Transcribe.GetTranscriptionJob" assert "Credential=test-access-key/" in sent.headers["authorization"] assert "sk-virtual" not in sent.headers["authorization"] + @pytest.mark.parametrize("operation", ["GetTranscriptionJob", "DeleteTranscriptionJob"]) + @pytest.mark.parametrize("owner", ["user-b", None]) + def test_jobs_started_by_others_are_not_reachable( + self, transcribe_client: TestClient, operation: str, owner: str | None + ) -> None: + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(TRANSCRIBE_UPSTREAM).mock(return_value=httpx.Response(200, json=_owned_job(owner))) + response = transcribe_client.post( + f"/transcribe/{operation}", json={"TranscriptionJobName": "litellm-job-1"} + ) + + assert response.status_code == 404 + assert [call.request.headers["x-amz-target"] for call in route.calls] == ["Transcribe.GetTranscriptionJob"] + + def test_the_owner_may_delete_the_job(self, transcribe_client: TestClient) -> None: + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(TRANSCRIBE_UPSTREAM) + route.side_effect = [httpx.Response(200, json=_owned_job("user-a")), httpx.Response(200, json={})] + response = transcribe_client.post( + "/transcribe/DeleteTranscriptionJob", json={"TranscriptionJobName": "litellm-job-1"} + ) + + assert (response.status_code, response.json()) == (200, {}) + assert [call.request.headers["x-amz-target"] for call in route.calls] == [ + "Transcribe.GetTranscriptionJob", + "Transcribe.DeleteTranscriptionJob", + ] + + @pytest.mark.parametrize("operation", ["ListTranscriptionJobs", "ListVocabularies", "DeleteVocabulary"]) + def test_account_wide_operations_need_a_proxy_admin(self, transcribe_client: TestClient, operation: str) -> None: + with respx.mock(assert_all_called=False) as upstream: + route = upstream.post(TRANSCRIBE_UPSTREAM) + response = transcribe_client.post(f"/transcribe/{operation}", json={}) + + assert response.status_code == 403 + assert operation in response.json()["detail"] + assert not route.called + + def test_a_proxy_admin_reaches_account_wide_operations( + self, transcribe_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.proxy_server import app + + monkeypatch.setitem( + app.dependency_overrides, + user_api_key_auth, + lambda: UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + with respx.mock(assert_all_called=True) as upstream: + upstream.post(TRANSCRIBE_UPSTREAM).mock( + return_value=httpx.Response(200, json={"TranscriptionJobSummaries": []}) + ) + response = transcribe_client.post("/transcribe/ListTranscriptionJobs", json={}) + + assert (response.status_code, response.json()) == (200, {"TranscriptionJobSummaries": []}) + def test_upstream_error_status_and_body_are_returned(self, transcribe_client: TestClient) -> None: aws_error = {"__type": "BadRequestException", "Message": "The requested job couldn't be found."} with respx.mock(assert_all_called=True) as upstream: upstream.post(TRANSCRIBE_UPSTREAM).mock(return_value=httpx.Response(400, json=aws_error)) response = transcribe_client.post( - "/transcribe/GetTranscriptionJob", json={"TranscriptionJobName": "missing"} + "/transcribe/StartTranscriptionJob", + json={**dict(self.START_JOB_BODY), "TranscriptionJobName": "missing"}, ) assert (response.status_code, response.json()) == (400, aws_error) @@ -5386,7 +5459,7 @@ class TestTranscribeProxyRoute: with respx.mock(assert_all_called=False) as upstream: route = upstream.post(TRANSCRIBE_UPSTREAM) response = transcribe_client.post( - "/transcribe/ListTranscriptionJobs", content=raw_body, headers={"Content-Type": "application/json"} + "/transcribe/GetTranscriptionJob", content=raw_body, headers={"Content-Type": "application/json"} ) assert response.status_code == 400 @@ -5399,7 +5472,7 @@ class TestTranscribeProxyRoute: monkeypatch.delenv(name, raising=False) with respx.mock(assert_all_called=False) as upstream: route = upstream.post(TRANSCRIBE_UPSTREAM) - response = transcribe_client.post("/transcribe/ListTranscriptionJobs", json={}) + response = transcribe_client.post("/transcribe/GetTranscriptionJob", json={}) assert response.status_code == 400 assert "AWS region" in response.json()["detail"] @@ -5495,9 +5568,7 @@ class TestVertexAILiveWebsocketPassthrough: ] ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) - monkeypatch.setattr( - passthrough_module.passthrough_endpoint_router, "default_vertex_config", None - ) + monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None) self._clear_vertex_env(monkeypatch) websocket = self._websocket() ensure_token = AsyncMock(return_value=("token-abc", "proj-db")) @@ -5639,9 +5710,7 @@ class TestVertexAILiveWebsocketPassthrough: ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) - monkeypatch.setattr( - passthrough_module.passthrough_endpoint_router, "default_vertex_config", None - ) + monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None) self._clear_vertex_env(monkeypatch) websocket = self._websocket() ensure_token = AsyncMock(side_effect=Exception("Unable to find your credentials"))