diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index f80cb41dc3..6e655b03fe 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -32,6 +32,7 @@ _OPTIONAL_KWARGS_KEYS = frozenset( "aws_sts_endpoint", "aws_external_id", "aws_bedrock_runtime_endpoint", + "aws_bedrock_project_id", "tpm", "rpm", "use_xai_oauth", diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index ef0199031a..cbed2232be 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -48,6 +48,30 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): region = self._get_aws_region_name(optional_params=optional_params, model=model) return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + headers = super().validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + project_id = litellm_params.get("aws_bedrock_project_id") + if project_id: + headers["anthropic-workspace"] = project_id + return headers + def transform_request( self, model: str, diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index a78f696a05..900d9aa97d 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -6,7 +6,7 @@ AmazonAnthropicClaudeMessagesConfig. Overrides only the URL and model-prefix stripping that are specific to the bedrock-mantle endpoint. """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeMessagesConfig, @@ -45,6 +45,30 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): region = self._get_aws_region_name(optional_params=optional_params, model=model) return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + def validate_anthropic_messages_environment( + self, + headers: dict, + model: str, + messages: List[Any], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> Tuple[dict, Optional[str]]: + headers, api_base = super().validate_anthropic_messages_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + project_id = litellm_params.get("aws_bedrock_project_id") + if project_id: + headers["anthropic-workspace"] = project_id + return headers, api_base + def transform_anthropic_messages_request( self, model: str, diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 81a56030a5..ad37a1990d 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -8,11 +8,12 @@ Auth: AWS Bedrock API key as Bearer token (set via BEDROCK_MANTLE_API_KEY env va or region-aware key via BEDROCK_MANTLE_{REGION}_API_KEY. """ -from typing import Iterator, AsyncIterator, Any, Optional, Tuple, Union +from typing import Iterator, AsyncIterator, Any, List, Optional, Tuple, Union import litellm from litellm._logging import verbose_logger from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues from ...openai_like.chat.transformation import OpenAILikeChatConfig @@ -48,6 +49,30 @@ class BedrockMantleChatConfig(OpenAILikeChatConfig): dynamic_api_key = api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") return api_base, dynamic_api_key + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + headers = super().validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + project_id = litellm_params.get("aws_bedrock_project_id") + if project_id: + headers["OpenAI-Project"] = project_id + return headers + def get_supported_openai_params(self, model: str) -> list: base_params = super().get_supported_openai_params(model) try: diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index dfa108833a..29248e1ca5 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -115,6 +115,8 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): ) if api_key: headers["Authorization"] = f"Bearer {api_key}" + if litellm_params.aws_bedrock_project_id: + headers["OpenAI-Project"] = litellm_params.aws_bedrock_project_id return headers def supports_native_file_search(self) -> bool: diff --git a/litellm/main.py b/litellm/main.py index 2c416a595c..02609217dd 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1639,6 +1639,7 @@ def completion( # type: ignore # noqa: PLR0915 tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), use_xai_oauth=kwargs.get("use_xai_oauth", False), + aws_bedrock_project_id=kwargs.get("aws_bedrock_project_id"), ) cast(LiteLLMLoggingObj, logging).update_environment_variables( model=model, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 71cf5197de..c868d3d22b 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -271,6 +271,11 @@ _BANNED_REQUEST_BODY_PARAMS: 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", + # 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 + # reachable with the deployment's shared AWS credentials. + "aws_bedrock_project_id", # Provider-specific endpoint overrides that flow into the outbound # request via ``optional_params``. Same threat as ``api_base``: # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker diff --git a/litellm/types/router.py b/litellm/types/router.py index ed858557a6..5047cee424 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -178,6 +178,7 @@ class CredentialLiteLLMParams(BaseModel): aws_secret_access_key: Optional[str] = None aws_region_name: Optional[str] = None aws_bedrock_runtime_endpoint: Optional[str] = None + aws_bedrock_project_id: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None @@ -364,6 +365,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): aws_access_key_id: Optional[str] aws_secret_access_key: Optional[str] aws_region_name: Optional[str] + aws_bedrock_project_id: Optional[str] ## AWS S3 VECTORS ## vector_bucket_name: Optional[str] index_name: Optional[str] diff --git a/litellm/utils.py b/litellm/utils.py index 03c628b195..a0b66234a7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3837,6 +3837,10 @@ class PreProcessNonDefaultParams: additional_endpoint_specific_params: List[str], ) -> dict: for k, v in special_params.items(): + if k == "aws_bedrock_project_id": + # sent as a request header (read from litellm_params by the + # bedrock-mantle configs), never as a request body field + continue if k.startswith("aws_") and ( custom_llm_provider != "bedrock" and not custom_llm_provider.startswith("sagemaker") diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index a00057eaa6..bbefdd621f 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -1,10 +1,17 @@ """ Unit tests for the Bedrock Mantle (Claude Mythos Preview) integration. -Tests cover route detection, URL construction, and config dispatch for both -the /chat/completions and /messages endpoints. +Tests cover route detection, URL construction, config dispatch for both +the /chat/completions and /messages endpoints, and project (workspace) +association via `aws_bedrock_project_id`. """ +import json +from unittest.mock import patch + +import httpx +import pytest + from litellm.llms.bedrock.common_utils import BedrockModelInfo, get_bedrock_chat_config from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig from litellm.llms.bedrock.messages.mantle_transformation import ( @@ -12,6 +19,32 @@ from litellm.llms.bedrock.messages.mantle_transformation import ( ) +def _anthropic_response(url: str) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-mythos-preview", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + request=httpx.Request("POST", url), + ) + + +def _capture_request(url: str, headers: dict, data) -> dict: + raw_body = data.decode("utf-8") if isinstance(data, bytes) else data or "{}" + return { + "path": httpx.URL(url).path, + "headers": headers, + "body": json.loads(raw_body), + } + + def test_get_bedrock_route_mantle(): assert ( BedrockModelInfo.get_bedrock_route("mantle/anthropic.claude-mythos-preview") @@ -103,3 +136,114 @@ def test_mantle_transform_request_strips_prefix_and_adds_model(): ) assert request["model"] == "anthropic.claude-mythos-preview" assert "mantle/" not in request["model"] + + +def test_mantle_validate_environment_sets_workspace_header(): + config = AmazonMantleConfig() + headers = config.validate_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": "proj_abc123def456"}, + ) + assert headers["anthropic-workspace"] == "proj_abc123def456" + + +def test_mantle_validate_environment_without_project_id(): + config = AmazonMantleConfig() + headers = config.validate_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": None}, + ) + assert "anthropic-workspace" not in headers + + +def test_mantle_messages_validate_environment_sets_workspace_header(): + config = AmazonMantleMessagesConfig() + headers, api_base = config.validate_anthropic_messages_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": "proj_abc123def456"}, + api_base="https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages", + ) + assert headers["anthropic-workspace"] == "proj_abc123def456" + assert api_base == "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages" + + +def test_mantle_messages_validate_environment_without_project_id(): + config = AmazonMantleMessagesConfig() + headers, _ = config.validate_anthropic_messages_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + ) + assert "anthropic-workspace" not in headers + + +def test_mantle_completion_sends_workspace_header_and_clean_body(): + import litellm + + requests = [] + + def mock_post(self, url, data=None, headers=None, **kwargs): + requests.append(_capture_request(url=url, headers=headers or {}, data=data)) + return _anthropic_response(url) + + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post): + response = litellm.completion( + model="bedrock/mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + aws_bedrock_project_id="proj_abc123def456", + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + ) + + assert response.choices[0].message.content == "ok" + assert len(requests) == 1 + assert requests[0]["path"] == "/anthropic/v1/messages" + assert requests[0]["headers"]["anthropic-workspace"] == "proj_abc123def456" + assert "aws_bedrock_project_id" not in requests[0]["body"] + + +@pytest.mark.asyncio +async def test_mantle_anthropic_messages_sends_workspace_header_and_clean_body(): + import litellm + + requests = [] + + async def mock_post(self, url, data=None, headers=None, **kwargs): + requests.append(_capture_request(url=url, headers=headers or {}, data=data)) + return _anthropic_response(url) + + try: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=mock_post, + ): + response = await litellm.anthropic_messages( + model="bedrock/mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + aws_bedrock_project_id="proj_abc123def456", + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + ) + finally: + await litellm.close_litellm_async_clients() + + assert response["content"][0]["text"] == "ok" + assert len(requests) == 1 + assert requests[0]["path"] == "/anthropic/v1/messages" + assert requests[0]["headers"]["anthropic-workspace"] == "proj_abc123def456" + assert "aws_bedrock_project_id" not in requests[0]["body"] diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index e83992b6bd..c3de29bd9d 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -167,6 +167,26 @@ class TestBedrockMantleResponsesAuth: ) assert "Authorization" not in headers + def test_project_id_sets_openai_project_header(self): + cfg = BedrockMantleResponsesAPIConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-5.5", + litellm_params=GenericLiteLLMParams( + api_key="fake-key", aws_bedrock_project_id="proj_abc123def456" + ), + ) + assert headers["OpenAI-Project"] == "proj_abc123def456" + + def test_no_project_id_no_openai_project_header(self): + cfg = BedrockMantleResponsesAPIConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-5.5", + litellm_params=GenericLiteLLMParams(api_key="fake-key"), + ) + assert "OpenAI-Project" not in headers + def test_custom_llm_provider(self): cfg = BedrockMantleResponsesAPIConfig() assert cfg.custom_llm_provider == LlmProviders.BEDROCK_MANTLE diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 1725aa85d1..deaa053793 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -5,11 +5,14 @@ Bedrock Mantle is Amazon Bedrock's OpenAI-compatible inference engine (Project M API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.html """ +import json import os import sys +from unittest.mock import patch sys.path.insert(0, os.path.abspath("../../../../..")) +import httpx import pytest import litellm @@ -96,6 +99,79 @@ class TestBedrockMantleConfig: assert "max_tokens" in params +class TestBedrockMantleProjectHeader: + def test_validate_environment_sets_openai_project_header(self): + cfg = BedrockMantleChatConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-oss-120b", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": "proj_abc123def456"}, + api_key="fake-key", + ) + assert headers["OpenAI-Project"] == "proj_abc123def456" + assert headers["Authorization"] == "Bearer fake-key" + + def test_validate_environment_without_project_id(self): + cfg = BedrockMantleChatConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-oss-120b", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + api_key="fake-key", + ) + assert "OpenAI-Project" not in headers + + def test_completion_sends_openai_project_header_and_clean_body(self): + requests = [] + + def mock_post(self, url, data=None, headers=None, **kwargs): + raw_body = data.decode("utf-8") if isinstance(data, bytes) else data + requests.append( + {"headers": headers or {}, "body": json.loads(raw_body or "{}")} + ) + return httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1733529600, + "model": "openai.gpt-oss-120b", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + request=httpx.Request("POST", url), + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post + ): + response = litellm.completion( + model="bedrock_mantle/openai.gpt-oss-120b", + messages=[{"role": "user", "content": "hello"}], + api_key="fake-key", + aws_bedrock_project_id="proj_abc123def456", + ) + + assert response.choices[0].message.content == "ok" + assert len(requests) == 1 + assert requests[0]["headers"]["OpenAI-Project"] == "proj_abc123def456" + assert "aws_bedrock_project_id" not in requests[0]["body"] + + class TestBedrockMantleProviderResolution: def test_get_llm_provider_resolves_correctly(self): model, provider, _, _ = litellm.get_llm_provider( diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index d4ca55ca16..32b597376b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -1551,6 +1551,40 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: ) +class TestIsRequestBodySafeBlocksBedrockProjectOverride: + """``aws_bedrock_project_id`` pins a deployment to a Bedrock project so + that project's data-retention policy applies to its requests. A + caller-supplied value would run the request under any project reachable + with the deployment's shared AWS credentials, bypassing the configured + retention/accounting association.""" + + def test_project_id_in_request_body_is_rejected(self): + with pytest.raises(ValueError, match="aws_bedrock_project_id"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + "aws_bedrock_project_id": "proj_attacker000000", + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_admin_opt_in_proxy_wide_allows_project_id(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "aws_bedrock_project_id": "proj_byok000000", + }, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + # ── is_request_body_safe nested-config recursion (VERIA-6) ──────────────────── diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 62cb8154b6..b6c9e9c865 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4198,3 +4198,21 @@ class TestBedrockBaseModelLabelKeepsTools: ) assert "tools" not in result + + +def test_aws_bedrock_project_id_excluded_from_bedrock_optional_params(): + """`aws_bedrock_project_id` is sent as a bedrock-mantle request header, so it + must never reach optional_params (and from there the request body), while + other aws_* params keep flowing for boto3 auth.""" + from litellm.utils import get_optional_params + + result = get_optional_params( + model="mantle/anthropic.claude-mythos-preview", + custom_llm_provider="bedrock", + max_tokens=10, + aws_bedrock_project_id="proj_abc123def456", + aws_region_name="us-east-1", + ) + + assert "aws_bedrock_project_id" not in result + assert result["aws_region_name"] == "us-east-1" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8e47081955..15123bcdbf 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24891,6 +24891,8 @@ export interface components { auto_router_embedding_model?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; + /** Aws Bedrock Project Id */ + aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */ aws_bedrock_runtime_endpoint?: string | null; /** Aws Region Name */ @@ -32495,6 +32497,8 @@ export interface components { auto_router_embedding_model?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; + /** Aws Bedrock Project Id */ + aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */ aws_bedrock_runtime_endpoint?: string | null; /** Aws Region Name */