From c6fc5b01bb0033cf189734b415e1f4dc85c79513 Mon Sep 17 00:00:00 2001 From: Priyansh Nandwana <136072927+msdianprince-7@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:00:16 +0530 Subject: [PATCH] fix(bedrock): send batch jobs to a configured control-plane endpoint (#47260) * fix(bedrock): send batch jobs to a configured control-plane endpoint Bedrock batch create, retrieve and cancel always built the regional bedrock.{region} host, so a VPC endpoint or compatible gateway could not be used. Add aws_bedrock_control_plane_endpoint as a deployment setting read by the signed HTTP path and the boto3 path, carry it through the proxy credential filter, and refuse it in request bodies like the other endpoint overrides. api_base stays out of it because on a Bedrock deployment it points at the runtime host. Fixes #45056 * fix(bedrock): read the control-plane endpoint from typed kwargs in job status The named parameter added one unknown-argument error in litellm/batches/main.py, which unpacks untyped kwargs into this function * docs(bedrock): explain batch control-plane endpoint configuration * fix(bedrock): keep the batch control-plane endpoint out of chat request bodies A deployment that sets aws_bedrock_control_plane_endpoint also serves chat, and the unregistered key was swept into the Bedrock request body, so every chat call on it failed with "Extra inputs are not permitted". Register it with the other Bedrock batch connection fields so it stays a LiteLLM-owned kwarg on every route while batches still read it from litellm_params. * test(bedrock): verify endpoint signatures without internal mocks --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../batch_api/bedrock/README.md | 41 +++++++ .../litellm_core_utils/get_litellm_params.py | 1 + litellm/llms/bedrock/batches/handler.py | 26 +++- .../llms/bedrock/batches/transformation.py | 22 ++-- litellm/llms/bedrock/common_utils.py | 19 +++ litellm/proxy/auth/auth_utils.py | 1 + litellm/types/litellm_params.py | 1 + litellm/types/router.py | 1 + .../unit/llms/bedrock/batches/test_handler.py | 42 +++++++ .../bedrock/batches/test_transformation.py | 115 +++++++++++++++++- .../chat/test_converse_transformation.py | 50 ++++++++ tests/unit/proxy/auth/test_auth_utils.py | 1 + tests/unit/types/test_litellm_params.py | 3 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 14 files changed, 310 insertions(+), 17 deletions(-) create mode 100644 cookbook/litellm_proxy_server/batch_api/bedrock/README.md diff --git a/cookbook/litellm_proxy_server/batch_api/bedrock/README.md b/cookbook/litellm_proxy_server/batch_api/bedrock/README.md new file mode 100644 index 00000000000..1410d76658e --- /dev/null +++ b/cookbook/litellm_proxy_server/batch_api/bedrock/README.md @@ -0,0 +1,41 @@ +# Bedrock batch endpoints + +Set `aws_bedrock_control_plane_endpoint` on a deployment when batch create, retrieve, and cancel requests must use a custom Bedrock control-plane endpoint. The proxy signs these requests for the configured URL. Without this setting, batch requests use the standard regional Bedrock endpoint + +```yaml +model_list: + - model_name: bedrock-batch + litellm_params: + model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0 + aws_region_name: us-east-1 + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_batch_role_arn: os.environ/AWS_BATCH_ROLE_ARN + s3_output_bucket_name: os.environ/AWS_S3_OUTPUT_BUCKET_NAME + aws_bedrock_control_plane_endpoint: os.environ/BEDROCK_CONTROL_PLANE_ENDPOINT +``` + +Set `BEDROCK_CONTROL_PLANE_ENDPOINT` to your Bedrock control-plane VPC endpoint or compatible gateway base URL, without `/model-invocation-job`. The proxy adds the operation path. The endpoint must be reachable from the proxy, and the configured AWS region must match the endpoint's signing region + +`api_base` and `aws_bedrock_runtime_endpoint` configure runtime inference separately. They do not override the batch control-plane endpoint. `s3_endpoint_url` configures S3 file access separately + +The control-plane setting belongs in deployment configuration. Sending `aws_bedrock_control_plane_endpoint` in an API request body is rejected. There is no automatic `AWS_BEDROCK_CONTROL_PLANE_ENDPOINT` environment fallback, so reference the desired environment variable explicitly as shown above + +Start the proxy with this configuration and set `LITELLM_API_BASE` to its URL. Set `BEDROCK_BATCH_INPUT_URI` to a valid input JSONL file in S3 that the batch service role can read + +```bash +curl -sS "$LITELLM_API_BASE/v1/batches" \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" \ + -H 'Content-Type: application/json' \ + -d "{\"model\":\"bedrock-batch\",\"input_file_id\":\"$BEDROCK_BATCH_INPUT_URI\",\"endpoint\":\"/v1/chat/completions\",\"completion_window\":\"24h\"}" +``` + +Use the returned `id` as `BATCH_ID` to retrieve or cancel the same job through the deployment's configured endpoint + +```bash +curl -sS "$LITELLM_API_BASE/v1/batches/$BATCH_ID?model=bedrock-batch" \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" + +curl -sS -X POST "$LITELLM_API_BASE/v1/batches/$BATCH_ID/cancel?model=bedrock-batch" \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" +``` diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 6ca02cbbe64..657fd3abfbc 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -27,6 +27,7 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( "aws_external_id", "aws_session_tags", "aws_bedrock_runtime_endpoint", + "aws_bedrock_control_plane_endpoint", "aws_bedrock_project_id", } ) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 81057a1ee23..48e79d39b8e 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -5,7 +5,10 @@ from typing import TYPE_CHECKING, Final, Literal from openai.types.batch import BatchRequestCounts from openai.types.batch import Metadata as OpenAIBatchMetadata -from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.llms.bedrock.common_utils import ( + bedrock_control_plane_base_url, + configured_bedrock_control_plane_endpoint, +) from litellm.types.llms.bedrock import AwsAuthParams, AwsSessionTag from litellm.types.utils import LiteLLMBatch @@ -156,10 +159,14 @@ class BedrockBatchesHandler: aws_session_tags=aws_session_tags, ) creds: Final = BedrockBatchesConfig().resolve_credentials(auth_params, region) + control_plane_endpoint: Final = configured_bedrock_control_plane_endpoint( + kwargs.get("aws_bedrock_control_plane_endpoint") + ) client: Final = boto3.client( "bedrock", region_name=region, + endpoint_url=control_plane_endpoint, aws_access_key_id=creds.access_key, aws_secret_access_key=creds.secret_key, aws_session_token=creds.token, @@ -171,6 +178,7 @@ class BedrockBatchesHandler: batch_id=batch_id, aws_region_name=region, logging_obj=logging_obj, + aws_bedrock_control_plane_endpoint=control_plane_endpoint, **auth_params.model_dump(), ) @@ -270,7 +278,7 @@ class BedrockBatchesHandler: batch_id: str, aws_region_name: str | None = None, logging_obj=None, - **kwargs, + **kwargs: object, ) -> "LiteLLMBatch": """ Handle ``GetModelInvocationJob`` status check for AWS Bedrock bulk batch @@ -294,7 +302,10 @@ class BedrockBatchesHandler: ``aws_session_token``, ``aws_profile_name``, ``aws_role_name``, ``aws_session_name``, ``aws_web_identity_token``, ``aws_sts_endpoint``, - ``aws_external_id``, ``aws_session_tags``). Unknown keys are ignored. + ``aws_external_id``, ``aws_session_tags``), plus + ``aws_bedrock_control_plane_endpoint`` to send the call to a VPC + endpoint or gateway instead of the regional ``bedrock`` host. + Unknown keys are ignored. Returns: ``LiteLLMBatch`` shaped like an OpenAI Batch resource. @@ -318,10 +329,14 @@ class BedrockBatchesHandler: from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig creds: Final = BedrockBatchesConfig().resolve_credentials(AwsAuthParams.model_validate(kwargs), region) + control_plane_endpoint: Final = configured_bedrock_control_plane_endpoint( + kwargs.get("aws_bedrock_control_plane_endpoint") + ) client: Final = boto3.client( "bedrock", region_name=region, + endpoint_url=control_plane_endpoint, aws_access_key_id=creds.access_key, aws_secret_access_key=creds.secret_key, aws_session_token=creds.token, @@ -334,14 +349,13 @@ class BedrockBatchesHandler: # `GetModelInvocationJob` accepts either form, but only the bare id # produces a sensible-looking URL in logs. url_path_id: Final = _extract_job_id_from_arn(batch_id) or batch_id + control_plane_base: Final = bedrock_control_plane_base_url(control_plane_endpoint, region) logging_obj.pre_call( input=batch_id, api_key="", additional_args={ "complete_input_dict": {"jobIdentifier": batch_id}, - "api_base": ( - f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{url_path_id}" - ), + "api_base": f"{control_plane_base}/model-invocation-job/{url_path_id}", }, ) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 50fb911709a..84f40bec9ca 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Final, Literal, cast from httpx import Headers, Response from pydantic import TypeAdapter, ValidationError -from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix, is_bedrock_arn +from litellm.litellm_core_utils.aws_partition import is_bedrock_arn from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, ) @@ -32,6 +32,7 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders, Usage from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( CommonBatchFilesUtils, + bedrock_control_plane_base_url, merge_bedrock_aws_request_params, resolve_s3_bucket_owner, resolve_s3_encryption_key_id, @@ -177,11 +178,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) aws_region_name: Final = self._get_aws_region_name(request_params, model) - # Bedrock model invocation job endpoint - # Format: https://bedrock.{region}.{partition dns suffix}/model-invocation-job - bedrock_endpoint: Final = ( - f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job" + control_plane_base: Final = bedrock_control_plane_base_url( + request_params.get("aws_bedrock_control_plane_endpoint"), aws_region_name ) + bedrock_endpoint: Final = f"{control_plane_base}/model-invocation-job" return bedrock_endpoint @@ -279,9 +279,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # Use common utility for AWS signing request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) aws_region_name: Final = self._get_aws_region_name(request_params, model) - endpoint_url: Final = ( - f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job" + control_plane_base: Final = bedrock_control_plane_base_url( + request_params.get("aws_bedrock_control_plane_endpoint"), aws_region_name ) + endpoint_url: Final = f"{control_plane_base}/model-invocation-job" signed_headers, signed_data = self.common_utils.sign_aws_request( service_name="bedrock", data=bedrock_request, @@ -431,12 +432,13 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): import urllib.parse as _ul encoded_arn: Final = _ul.quote(batch_id, safe="") - endpoint_url: Final = ( - f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{encoded_arn}" + request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) + control_plane_base: Final = bedrock_control_plane_base_url( + request_params.get("aws_bedrock_control_plane_endpoint"), region ) + endpoint_url: Final = f"{control_plane_base}/model-invocation-job/{encoded_arn}" # Use common utility for AWS signing - request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) signed_headers, _ = self.common_utils.sign_aws_request( service_name="bedrock", data={}, # GET request has no body diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index bbbc929fad3..0cf3e895ad2 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -138,6 +138,25 @@ def merge_bedrock_aws_request_params( return request_params +def configured_bedrock_control_plane_endpoint(endpoint: object) -> str | None: + """An ``aws_bedrock_control_plane_endpoint`` value without a trailing slash, or None when unset.""" + if isinstance(endpoint, str) and endpoint.strip(): + return endpoint.strip().rstrip("/") + return None + + +def bedrock_control_plane_base_url(endpoint: object, aws_region_name: str) -> str: + """Base URL for Bedrock control-plane calls such as model invocation jobs. + + A configured ``aws_bedrock_control_plane_endpoint`` (a VPC endpoint or a compatible gateway) + wins over the regional ``bedrock`` host. ``api_base`` is deliberately not used here, since on a + Bedrock deployment it points at the runtime host, which does not serve these APIs. + """ + return configured_bedrock_control_plane_endpoint(endpoint) or ( + f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" + ) + + def s3_static_key_pair(params: Mapping[str, object]) -> tuple[str, str] | None: """The s3_access_key_id / s3_secret_access_key pair when both are set, otherwise None.""" s3_access_key_id: Final = params.get("s3_access_key_id") diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 73d20ae632c..996c28bebaa 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -419,6 +419,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( # tokens) to the attacker's host, or coerces the proxy into # authenticating against the attacker's host with admin secrets. "aws_bedrock_runtime_endpoint", + "aws_bedrock_control_plane_endpoint", # Bedrock project/workspace association. Deployments pin this to # enforce a data-retention policy, so a caller-supplied value would # re-route the request's retention and accounting to any project diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 978eb9c4557..6b08920e042 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -103,6 +103,7 @@ class ProviderConnection: class BedrockBatchConnection: # Bedrock rejects these names in request bodies, so register them as LiteLLM-owned aws_batch_role_arn: str | None = None + aws_bedrock_control_plane_endpoint: str | None = None s3_bucket_name: str | None = None s3_region_name: str | None = None s3_endpoint_url: str | None = None diff --git a/litellm/types/router.py b/litellm/types/router.py index 1f159cb8b16..2618bf698ee 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -413,6 +413,7 @@ class CredentialLiteLLMParams(LiteLLMBaseModel): aws_external_id: str | None = None aws_session_tags: Sequence[AwsSessionTag] | None = None aws_bedrock_runtime_endpoint: str | None = None + aws_bedrock_control_plane_endpoint: str | None = None aws_bedrock_project_id: str | None = None s3_bucket_name: str | None = None s3_endpoint_url: str | None = None diff --git a/tests/unit/llms/bedrock/batches/test_handler.py b/tests/unit/llms/bedrock/batches/test_handler.py index d920eabcdf6..602504f0c49 100644 --- a/tests/unit/llms/bedrock/batches/test_handler.py +++ b/tests/unit/llms/bedrock/batches/test_handler.py @@ -257,6 +257,25 @@ def test_explicit_region_overrides_arn(patched_boto3): assert kwargs["region_name"] == "eu-central-1" +def test_status_client_uses_regional_host_when_no_control_plane_endpoint( + patched_boto3: tuple[MagicMock, MagicMock], +) -> None: + _, boto_client_factory = patched_boto3 + BedrockBatchesHandler.handle_model_invocation_job_status(batch_id=JOB_ARN) + _, kwargs = boto_client_factory.call_args + assert kwargs["endpoint_url"] is None + + +def test_status_client_uses_configured_control_plane_endpoint(patched_boto3: tuple[MagicMock, MagicMock]) -> None: + _, boto_client_factory = patched_boto3 + BedrockBatchesHandler.handle_model_invocation_job_status( + batch_id=JOB_ARN, + aws_bedrock_control_plane_endpoint="https://bedrock-gateway.internal/", + ) + _, kwargs = boto_client_factory.call_args + assert kwargs["endpoint_url"] == "https://bedrock-gateway.internal" + + def test_failure_message_propagates(patched_boto3): fake_client, _ = patched_boto3 failed_response = _fake_boto3_response(status="Failed") @@ -479,6 +498,29 @@ def test_litellm_cancel_batch_dispatches_to_bedrock(patched_boto3): assert batch.status == "cancelled" +def test_litellm_batch_calls_send_every_bedrock_client_to_the_configured_endpoint( + patched_boto3: tuple[MagicMock, MagicMock], +) -> None: + import litellm + + fake_client, boto_client_factory = patched_boto3 + fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped") + + litellm.cancel_batch( + batch_id=JOB_ARN, + custom_llm_provider="bedrock", + aws_bedrock_control_plane_endpoint="https://bedrock-gateway.internal", + ) + litellm.retrieve_batch( + batch_id=JOB_ARN, + custom_llm_provider="bedrock", + aws_bedrock_control_plane_endpoint="https://bedrock-gateway.internal", + ) + + endpoints: Final = tuple(call.kwargs["endpoint_url"] for call in boto_client_factory.call_args_list) + assert endpoints == ("https://bedrock-gateway.internal",) * 3 + + class _TagGatedSTSClient: """Stands in for STS behind a trust policy that only admits sessions carrying ``tags``.""" diff --git a/tests/unit/llms/bedrock/batches/test_transformation.py b/tests/unit/llms/bedrock/batches/test_transformation.py index 5e987239c54..5360ac00018 100644 --- a/tests/unit/llms/bedrock/batches/test_transformation.py +++ b/tests/unit/llms/bedrock/batches/test_transformation.py @@ -14,10 +14,15 @@ URL/ARN handling, and the error class. AWS auth/sigv4 is the only external seam we mock; everything else runs for real. """ +from datetime import datetime, timezone +from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.credentials import Credentials from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig from litellm.types.utils import LlmProviders @@ -41,7 +46,7 @@ ARN = "arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/abc1234567" @pytest.fixture -def config(): +def config() -> BedrockBatchesConfig: return BedrockBatchesConfig() @@ -66,6 +71,90 @@ def test_get_complete_batch_url_uses_region(config): assert url == "https://bedrock.eu-central-1.amazonaws.com/model-invocation-job" +CONTROL_PLANE_ENDPOINT: Final = "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/" + + +def test_get_complete_batch_url_uses_configured_control_plane_endpoint(config: BedrockBatchesConfig) -> None: + url: Final = config.get_complete_batch_url( + api_base=None, + api_key=None, + model="anthropic.claude-3", + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={"aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT}, + data={"input_file_id": "s3://b/k"}, + ) + assert url == "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job" + + +def test_proxy_deployment_credentials_carry_the_control_plane_endpoint(config: BedrockBatchesConfig) -> None: + from litellm import Router + + router: Final = Router( + model_list=[ + { + "model_name": "claude-batch", + "litellm_params": { + "model": "bedrock/anthropic.claude-3", + "aws_region_name": "us-west-2", + "aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT, + }, + } + ] + ) + credentials: Final = router.get_deployment_credentials_with_provider("claude-batch") + + url: Final = config.get_complete_batch_url( + api_base=None, + api_key=None, + model="anthropic.claude-3", + optional_params={}, + litellm_params=credentials, + data={"input_file_id": "s3://b/k"}, + ) + assert url == "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job" + + +def test_get_complete_batch_url_ignores_runtime_api_base(config: BedrockBatchesConfig) -> None: + url: Final = config.get_complete_batch_url( + api_base="https://bedrock-runtime.us-west-2.amazonaws.com", + api_key=None, + model="anthropic.claude-3", + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={"api_base": "https://bedrock-runtime.us-west-2.amazonaws.com"}, + data={"input_file_id": "s3://b/k"}, + ) + assert url == "https://bedrock.us-west-2.amazonaws.com/model-invocation-job" + + +def test_create_request_sends_and_signs_for_configured_control_plane_endpoint(config: BedrockBatchesConfig) -> None: + expected_url: Final = "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job" + with patch("botocore.auth.get_current_datetime", return_value=datetime(2026, 1, 1, tzinfo=timezone.utc)): + result: Final = config.transform_create_batch_request( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + create_batch_data={ + "input_file_id": "s3://in-bucket/input.jsonl", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::123:role/my-batch-role", + "aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT, + "aws_access_key_id": "endpoint-test-key", + "aws_secret_access_key": "endpoint-test-secret", + }, + ) + expected_request: Final = AWSRequest( + method="POST", url=expected_url, data=result["data"], headers={"Content-Type": "application/json"} + ) + SigV4Auth(Credentials("endpoint-test-key", "endpoint-test-secret"), "bedrock", "us-west-2").add_auth( + expected_request + ) + + assert result["url"] == expected_url + assert result["headers"]["Authorization"] == expected_request.headers["Authorization"] + + # --------------------------------------------------------------------------- # # transform_create_batch_request - request construction (sign_aws_request mocked) # --------------------------------------------------------------------------- # @@ -608,6 +697,30 @@ def test_retrieve_request_builds_encoded_arn_url(config): assert mock_sign.call_args.kwargs["data"] == {} +def test_retrieve_request_sends_and_signs_for_configured_control_plane_endpoint(config: BedrockBatchesConfig) -> None: + expected_url: Final = ( + "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job/" + "arn%3Aaws%3Abedrock%3Aus-west-2%3A123456789012%3Amodel-invocation-job%2Fabc1234567" + ) + with patch("botocore.auth.get_current_datetime", return_value=datetime(2026, 1, 1, tzinfo=timezone.utc)): + result: Final = config.transform_retrieve_batch_request( + batch_id=ARN, + optional_params={}, + litellm_params={ + "aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT, + "aws_access_key_id": "endpoint-test-key", + "aws_secret_access_key": "endpoint-test-secret", + }, + ) + expected_request: Final = AWSRequest(method="GET", url=expected_url, data="", headers={}) + SigV4Auth(Credentials("endpoint-test-key", "endpoint-test-secret"), "bedrock", "us-west-2").add_auth( + expected_request + ) + + assert result["url"] == expected_url + assert result["headers"]["Authorization"] == expected_request.headers["Authorization"] + + def test_retrieve_request_rejects_non_arn(config): with pytest.raises(ValueError, match="Expected ARN"): config.transform_retrieve_batch_request( diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 1fb5542c964..154797bab5a 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -10739,3 +10739,53 @@ def test_a_refused_native_schema_is_sent_again_only_when_the_retry_drops_native_ assert route.call_count == expected_calls assert all("outputConfig" in json.loads(call.request.content) for call in route.calls[:1]) assert all("outputConfig" not in json.loads(call.request.content) for call in route.calls[1:]) + + +CONTROL_PLANE_ENDPOINT: Final = "https://bedrock-fips.us-east-1.amazonaws.com" + + +def _bedrock_chat_answer(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/invoke"): + return httpx.Response( + 200, + json={ + "id": "msg_control_plane", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + return httpx.Response( + 200, + json={ + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + }, + ) + + +@pytest.mark.parametrize("route_prefix", ["bedrock/", "bedrock/converse/", "bedrock/invoke/"]) +@pytest.mark.respx(assert_all_called=True) +def test_a_deployment_with_a_batch_control_plane_endpoint_keeps_it_out_of_chat_requests( + respx_mock: respx.Router, route_prefix: str +) -> None: + route: Final = respx_mock.post( + url__regex=r"^https://bedrock-runtime\.us-east-1\.amazonaws\.com/model/.+/(converse|invoke)$" + ).mock(side_effect=_bedrock_chat_answer) + + response: Final = litellm.completion( + model=f"{route_prefix}us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + aws_access_key_id="AKIACONTROLPLANE", + aws_secret_access_key="control-plane-secret", + aws_region_name="us-east-1", + aws_bedrock_control_plane_endpoint=CONTROL_PLANE_ENDPOINT, + client=HTTPHandler(), + ) + + assert response.choices[0].message.content == "ok" + assert "aws_bedrock_control_plane_endpoint" not in route.calls[0].request.content.decode() diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index f81ba155097..67ff7dd9344 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -3566,6 +3566,7 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: "field", [ "aws_bedrock_runtime_endpoint", + "aws_bedrock_control_plane_endpoint", "langsmith_base_url", "langfuse_host", "posthog_host", diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index e88d757519c..84ba85f4a68 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -95,6 +95,7 @@ CONNECTION_NAMES: Final = ( "use_xai_oauth", "fireworks_forward_user_id", "aws_batch_role_arn", + "aws_bedrock_control_plane_endpoint", "s3_bucket_name", "s3_region_name", "s3_endpoint_url", @@ -274,6 +275,7 @@ INTERNAL_STATE_NAMES: Final = ( BEDROCK_BATCH_NAMES: Final = ( "aws_batch_role_arn", + "aws_bedrock_control_plane_endpoint", "s3_bucket_name", "s3_region_name", "s3_endpoint_url", @@ -686,6 +688,7 @@ NAMES_SHARED_WITH_TYPED_MODELS: Final[Mapping[str, tuple[str, ...]]] = MappingPr "api_key", "api_version", "aws_batch_role_arn", + "aws_bedrock_control_plane_endpoint", "azure_password", "azure_scope", "azure_username", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b93f7dce2ca..e432c947116 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35048,6 +35048,8 @@ export interface components { aws_access_key_id?: string | null; /** Aws Batch Role Arn */ aws_batch_role_arn?: string | null; + /** Aws Bedrock Control Plane Endpoint */ + aws_bedrock_control_plane_endpoint?: string | null; /** Aws Bedrock Project Id */ aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */ @@ -50336,6 +50338,8 @@ export interface components { aws_access_key_id?: string | null; /** Aws Batch Role Arn */ aws_batch_role_arn?: string | null; + /** Aws Bedrock Control Plane Endpoint */ + aws_bedrock_control_plane_endpoint?: string | null; /** Aws Bedrock Project Id */ aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */