From 534ab1628dc462912179660774af395fdce50839 Mon Sep 17 00:00:00 2001 From: Riddhi04 Date: Wed, 22 Jul 2026 15:58:37 +0400 Subject: [PATCH 1/5] fix(proxy): enforce model access checks on Bedrock passthrough routes get_model_from_request could not resolve a model for /bedrock/... routes since it only checked the JSON body's model field and a small set of URL regexes, none matching Bedrock's passthrough path. This let common_checks skip the key/project model allowlist entirely for any Bedrock passthrough action (invoke, converse, and their streaming variants), while the same model was correctly blocked on /v1/chat/completions Extract the model from the Bedrock endpoint path using the same helper the passthrough handler itself relies on, so the existing allowlist check applies uniformly across auth methods and call paths --- litellm/proxy/auth/auth_utils.py | 11 ++++ .../proxy/auth/test_auth_utils.py | 50 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index d6007a2d56e..43e60e49e3b 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1981,6 +1981,17 @@ def get_model_from_request( if vertex_match: model = vertex_match.group(1) + if model is None and route.lower().startswith("/bedrock"): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _extract_model_from_bedrock_endpoint, + ) + + bedrock_endpoint = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE) + try: + model = _extract_model_from_bedrock_endpoint(bedrock_endpoint) + except ValueError: + model = None + return model diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index f513f397b64..f3426534c28 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -428,6 +428,56 @@ def test_get_model_from_request_openai_deployment_route_still_works(): ) +def test_get_model_from_request_bedrock_converse_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/converse", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_invoke_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_v2_converse_stream_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/v2/model/us.anthropic.claude-sonnet-4-6/converse-stream", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_model_id_with_slashes(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/aws/anthropic/model-name/invoke", + ) + == "aws/anthropic/model-name" + ) + + +def test_get_model_from_request_bedrock_unparseable_endpoint_returns_none(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/agents/some-agent-route", + ) + is None + ) + + def test_get_model_from_request_includes_file_endpoint_header_model(): assert ( get_model_from_request( From 6981ccf0ca1c3e1937f71c3fce287c898a15feba Mon Sep 17 00:00:00 2001 From: Riddhi04 Date: Fri, 24 Jul 2026 16:38:34 +0400 Subject: [PATCH 2/5] fix(proxy): make URL model authoritative for Bedrock path-routed passthrough actions The allowlist check read model from the request body first, while bedrock_llm_proxy_route dispatches purely on the path model for invoke, converse, and their streaming variants. A caller could put an allowed model in the JSON body while targeting a disallowed model in the URL and slip past the check. count_tokens keeps reading from the body since its route has no model segment in the path. --- litellm/proxy/auth/auth_utils.py | 22 +++++----- .../proxy/auth/test_auth_utils.py | 40 +++++++++++++++++++ 2 files changed, 53 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 43e60e49e3b..fd746d62e76 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1981,16 +1981,20 @@ def get_model_from_request( if vertex_match: model = vertex_match.group(1) - if model is None and route.lower().startswith("/bedrock"): - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - _extract_model_from_bedrock_endpoint, - ) - + if route.lower().startswith("/bedrock"): bedrock_endpoint = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE) - try: - model = _extract_model_from_bedrock_endpoint(bedrock_endpoint) - except ValueError: - model = None + is_bedrock_count_tokens_route = ( + "count_tokens" in bedrock_endpoint.lower() or "count-tokens" in bedrock_endpoint.lower() + ) + if not is_bedrock_count_tokens_route: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _extract_model_from_bedrock_endpoint, + ) + + try: + model = _extract_model_from_bedrock_endpoint(bedrock_endpoint) + except ValueError: + pass return model diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index f3426534c28..61d1a308b9d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -478,6 +478,46 @@ def test_get_model_from_request_bedrock_unparseable_endpoint_returns_none(): ) +def test_get_model_from_request_bedrock_url_model_overrides_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/converse", + ) + == "us.anthropic.claude-opus-4-6-v1" + ) + + +def test_get_model_from_request_bedrock_invoke_url_model_overrides_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/invoke", + ) + == "us.anthropic.claude-opus-4-6-v1" + ) + + +def test_get_model_from_request_bedrock_count_tokens_uses_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/v1/messages/count_tokens", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/agents/some-agent-route", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + def test_get_model_from_request_includes_file_endpoint_header_model(): assert ( get_model_from_request( From 32293295f8fe9811a18179b556e12cdb6c4a11bc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:40:35 -0700 Subject: [PATCH 3/5] refactor(proxy): resolve the Bedrock route model through an early return The Bedrock branch of get_model_from_request reassigned the already resolved model binding. Move the route parsing into a helper that returns the URL model or None so the caller picks between it and the body model without rebinding. --- litellm/proxy/auth/auth_utils.py | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index fd746d62e76..12e6f75889e 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1982,23 +1982,26 @@ def get_model_from_request( model = vertex_match.group(1) if route.lower().startswith("/bedrock"): - bedrock_endpoint = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE) - is_bedrock_count_tokens_route = ( - "count_tokens" in bedrock_endpoint.lower() or "count-tokens" in bedrock_endpoint.lower() - ) - if not is_bedrock_count_tokens_route: - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - _extract_model_from_bedrock_endpoint, - ) - - try: - model = _extract_model_from_bedrock_endpoint(bedrock_endpoint) - except ValueError: - pass + bedrock_model: Final = _model_from_bedrock_route(route) + return model if bedrock_model is None else bedrock_model return model +def _model_from_bedrock_route(route: str) -> str | None: + bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE) + if "count_tokens" in bedrock_endpoint.lower() or "count-tokens" in bedrock_endpoint.lower(): + return None + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _extract_model_from_bedrock_endpoint, + ) + + try: + return _extract_model_from_bedrock_endpoint(bedrock_endpoint) + except ValueError: + return None + + def abbreviate_api_key(api_key: str) -> str: if len(api_key) < MINIMUM_CUSTOM_KEY_LENGTH: return "sk-..." From bd9593b74d4a05e265225a43afa0fd6c8d837484 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:41:31 -0700 Subject: [PATCH 4/5] fix(proxy): share the Bedrock count-tokens predicate between auth and the passthrough handler --- litellm/proxy/auth/auth_utils.py | 7 ++++--- .../llm_passthrough_endpoints.py | 7 +++++-- tests/test_litellm/proxy/auth/test_auth_utils.py | 10 ++++++++++ 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 12e6f75889e..1e4836654a1 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1989,13 +1989,14 @@ def get_model_from_request( def _model_from_bedrock_route(route: str) -> str | None: - bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE) - if "count_tokens" in bedrock_endpoint.lower() or "count-tokens" in bedrock_endpoint.lower(): - return None from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _extract_model_from_bedrock_endpoint, + is_bedrock_count_tokens_endpoint, ) + bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE) + if is_bedrock_count_tokens_endpoint(bedrock_endpoint): + return None try: return _extract_model_from_bedrock_endpoint(bedrock_endpoint) except ValueError: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 29f216fd450..820312ac0fd 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -703,6 +703,10 @@ BEDROCK_ENDPOINT_ACTIONS: Final = { BEDROCK_STREAMING_ACTIONS: Final = {"invoke-with-response-stream", "converse-stream"} +def is_bedrock_count_tokens_endpoint(endpoint: str) -> bool: + return "count_tokens" in endpoint or "count-tokens" in endpoint + + def _extract_model_from_bedrock_endpoint(endpoint: str) -> str: """ Extract model name from Bedrock endpoint path. @@ -977,8 +981,7 @@ async def bedrock_llm_proxy_route( request_body: Final = await _read_request_body(request=request) - # Special handling for count_tokens endpoints - if "count_tokens" in endpoint or "count-tokens" in endpoint: + if is_bedrock_count_tokens_endpoint(endpoint): return await handle_bedrock_count_tokens( endpoint=endpoint, request=request, diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 61d1a308b9d..a996de4d40c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -508,6 +508,16 @@ def test_get_model_from_request_bedrock_count_tokens_uses_body_model(): ) +def test_get_model_from_request_bedrock_uppercase_count_tokens_segment_is_not_count_tokens(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-haiku-4-5-20251001-v1:0"}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke/COUNT_TOKENS", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model(): assert ( get_model_from_request( From 533a1c959c21f659a78fb130b17c700a3d6d1a06 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:42:46 -0700 Subject: [PATCH 5/5] fix(proxy): reserve budget for the Bedrock Converse prompt, not the context window Resolving the model from the /bedrock path sends passthrough calls through optimistic budget reservation, whose tokenizer cannot walk Converse content blocks and so fell back to the model's max_input_tokens. Count those messages as text and read inferenceConfig.maxTokens so a budgeted key reserves the request's cost. --- .../spend_tracking/budget_reservation.py | 37 ++++++++++++------- .../spend_tracking/test_budget_reservation.py | 31 +++++++++++++++- 2 files changed, 53 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 91d2ece7a51..5ed52a327d4 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1362,12 +1362,15 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int: def _count_input_tokens(request_body: dict, model: str) -> int | None: try: if "messages" in request_body: - return litellm.token_counter( - model=model, - messages=request_body.get("messages") or [], - tools=request_body.get("tools"), - tool_choice=request_body.get("tool_choice"), - ) + try: + return litellm.token_counter( + model=model, + messages=request_body.get("messages") or (), + tools=request_body.get("tools"), + tool_choice=request_body.get("tool_choice"), + ) + except ValueError: + return _count_text_tokens(model=model, text=request_body.get("messages")) if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: @@ -1415,11 +1418,7 @@ def _estimate_output_tokens( if _is_input_only_route(route=route): return 0 - requested: int | None = None - for key in ("max_completion_tokens", "max_tokens", "max_output_tokens"): - requested = _to_int(request_body.get(key)) - if requested is not None: - break + requested: Final = _requested_output_tokens(request_body) # Clamp at min(requested-or-default, model_max-or-default). Two purposes: # (1) Without an explicit cap we still need a finite reservation so the @@ -1430,9 +1429,19 @@ def _estimate_output_tokens( # at the cap — the model can only physically emit max_output_tokens # anyway, so reserving more is both wasteful and a DoS surface. model_ceiling: Final = _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK - if requested is None: - requested = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK - return min(requested, model_ceiling) + return min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK if requested is None else requested, model_ceiling) + + +_OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens") + + +def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: + inference_config: Final = request_body.get("inferenceConfig") + candidates: Final = ( + *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), + inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None, + ) + return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) def _count_text_tokens(model: str, text: object) -> int: diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index f65f68812a2..c7de8c943ad 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -5,7 +5,7 @@ import pytest from litellm.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.spend_tracking.budget_reservation import reserve_budget_for_request +from litellm.proxy.spend_tracking.budget_reservation import estimate_request_max_cost, reserve_budget_for_request from litellm.proxy.utils import ProxyLogging TOKEN_COUNTING_ROUTES: Final = ( @@ -46,3 +46,32 @@ async def test_non_exempt_llm_route_still_reserves_budget(): assert reservation is not None assert reservation["reserved_cost"] > 0 + + +BEDROCK_SONNET: Final = "us.anthropic.claude-sonnet-4-6" +CONVERSE_BODY: Final = { + "messages": [{"role": "user", "content": [{"text": "Reply with one word: pong"}]}], + "inferenceConfig": {"maxTokens": 5}, +} +INVOKE_BODY: Final = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 5, + "messages": [{"role": "user", "content": "Reply with one word: pong"}], +} + + +def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window(): + converse_cost: Final = estimate_request_max_cost( + request_body=CONVERSE_BODY, + route=f"/bedrock/model/{BEDROCK_SONNET}/converse", + llm_router=None, + input_token_counts={}, + ) + invoke_cost: Final = estimate_request_max_cost( + request_body=INVOKE_BODY, + route=f"/bedrock/model/{BEDROCK_SONNET}/invoke", + llm_router=None, + input_token_counts={}, + ) + assert converse_cost is not None and invoke_cost is not None + assert invoke_cost < converse_cost < 2 * invoke_cost