From 28e9968c8193763ddc1915cbb744d2448ccacee2 Mon Sep 17 00:00:00 2001 From: Brendan Smith-Elion Date: Thu, 5 Nov 2026 08:14:22 -0500 Subject: [PATCH 1/2] feat(bedrock/agentcore): forward deployment-listed request params into the InvokeAgentRuntime payload An AgentCore deployment can now set forward_payload_fields, a list of request param names. Each listed param the caller sends is copied into the InvokeAgentRuntime payload under its own name, so an agent can take structured input such as a context object next to the prompt. Without the setting the payload stays exactly {"prompt": ...} The list is deployment-only. The proxy rejects forward_payload_fields in a request body like the other deployment opt-ins, because a caller-chosen list could copy deployment-set params into the payload. prompt and the AWS credential params are never forwarded, and a value that is not JSON fails the request with a 400 that names the field but not its value --- .../bedrock/chat/agentcore/transformation.py | 37 +++++++++-- litellm/proxy/auth/auth_utils.py | 2 + .../test_agentcore_transformation.py | 66 +++++++++++++++++++ tests/unit/proxy/auth/test_auth_utils.py | 30 +++++++++ 4 files changed, 131 insertions(+), 4 deletions(-) diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 2982ae6cd08..c3f1746b841 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -6,15 +6,17 @@ https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgen import json from collections.abc import AsyncGenerator, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Optional, Union from urllib.parse import quote import httpx -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) @@ -53,6 +55,17 @@ else: AsyncHTTPHandler = Any _DECODED_JSON: Final = TypeAdapter(object) +_PARAMS: Final = TypeAdapter(dict[str, object]) +_PAYLOAD_FIELD_NAMES: Final = TypeAdapter(tuple[str, ...]) +_PAYLOAD_FIELD_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_UNFORWARDABLE_PAYLOAD_FIELDS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({"prompt"}) + + +def _payload_field_value(name: str, value: object) -> JsonValue: + try: + return _PAYLOAD_FIELD_VALUE.validate_python(value) + except ValidationError: + raise BedrockError(status_code=400, message=f"forward_payload_fields: {name} is not a JSON value") from None class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): @@ -245,10 +258,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): OpenAI-shaped multimodal blocks. This is opt-in because an AgentCore agent must be explicitly written to read ``payload["content"]``; by default the payload stays byte-identical to the legacy ``{"prompt": "..."}`` shape. + - Each request param named in the deployment's ``forward_payload_fields`` + list is copied in under its own name, never ``prompt`` or AWS credentials. Returns: - dict: Payload dict containing the prompt and (optionally) the OpenAI - content list. + dict: Payload dict containing the prompt, the forwarded params and + (optionally) the OpenAI content list. """ verbose_logger.debug("AgentCore transform_request - optional_params keys: %s", list(optional_params.keys())) @@ -256,7 +271,10 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): prompt: Final = convert_content_list_to_str(messages[-1]) # Create the payload - this is what goes in the body (raw JSON) - payload: Final[dict] = {"prompt": prompt} + payload: Final[dict] = { + "prompt": prompt, + **self._forwarded_payload_fields(_PARAMS.validate_python(optional_params)), + } # Opt-in: when forward_multimodal_content is set, forward the OpenAI content # list verbatim under "content" so an attachment-aware agent can read the raw @@ -307,6 +325,17 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return bool(value) return False + @staticmethod + def _forwarded_payload_fields(optional_params: Mapping[str, object]) -> Mapping[str, JsonValue]: + names: Final = _PAYLOAD_FIELD_NAMES.validate_python(optional_params.get("forward_payload_fields") or ()) + return MappingProxyType( + { + name: _payload_field_value(name, optional_params[name]) + for name in names + if name in optional_params and name not in _UNFORWARDABLE_PAYLOAD_FIELDS + } + ) + def _extract_sse_json(self, line: str) -> dict | None: """Extract and parse JSON from an SSE data line.""" if not line.startswith("data:"): diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 161ff2fdfcd..4b4134d652a 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -460,6 +460,8 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( # Deployment opt-in: a caller-supplied false would switch off identity # forwarding and let the caller choose the `user` Fireworks sees. "fireworks_forward_user_id", + # Deployment opt-in: a caller-supplied list would copy deployment-set params into the AgentCore payload. + "forward_payload_fields", # SDK-only field; also rejected outright in is_request_body_safe. "model_list", "vertex_ai_credentials", diff --git a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py index acfd954215d..03bc4a7e407 100644 --- a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py +++ b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py @@ -13,6 +13,7 @@ import json import httpx import pytest import respx +from pydantic import TypeAdapter from typing import Final @@ -1365,3 +1366,68 @@ def test_get_complete_url_sends_a_non_string_runtime_endpoint_to_that_endpoint() ) assert url.startswith(f"{vpc_endpoint}/runtimes/") + + +def _payload_sent_to_agentcore(respx_mock: respx.Router, **params: object) -> object: + route: Final = respx_mock.post(_AGENTCORE_RUNTIME_URL).mock( + return_value=httpx.Response(200, json={"result": "done"}, headers={"content-type": "application/json"}) + ) + litellm.completion( + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/" + "hosted_agent_13sf6-cALnp38iZD", + messages=[{"role": "user", "content": "Summarize the document"}], + aws_access_key_id="AKIAAGENTCOREKEY", + aws_secret_access_key="agentcore-secret", + client=HTTPHandler(), + **params, + ) + request: Final[httpx.Request] = route.calls.last.request + return TypeAdapter(object).validate_json(request.content) + + +DOCUMENT_CONTEXT: Final = {"document": "Q3 revenue grew 4 percent.", "sections": [{"id": 1, "flags": [True, None]}]} + + +@pytest.mark.respx(assert_all_called=True) +def test_agentcore_payload_holds_only_the_prompt_when_no_fields_are_forwarded(respx_mock: respx.Router) -> None: + assert _payload_sent_to_agentcore(respx_mock, context=DOCUMENT_CONTEXT) == {"prompt": "Summarize the document"} + + +@pytest.mark.respx(assert_all_called=True) +def test_agentcore_payload_carries_each_forwarded_request_param(respx_mock: respx.Router) -> None: + payload: Final = _payload_sent_to_agentcore( + respx_mock, + forward_payload_fields=["context", "trace_label", "not_sent"], + context=DOCUMENT_CONTEXT, + trace_label="run-7", + unlisted_param="stays out", + ) + + assert payload == { + "prompt": "Summarize the document", + "context": DOCUMENT_CONTEXT, + "trace_label": "run-7", + } + + +@pytest.mark.respx(assert_all_called=True) +def test_agentcore_never_forwards_the_prompt_or_aws_credentials(respx_mock: respx.Router) -> None: + payload: Final = _payload_sent_to_agentcore( + respx_mock, + forward_payload_fields=["prompt", "aws_secret_access_key", "aws_access_key_id", "context"], + prompt="replaced prompt", + context=DOCUMENT_CONTEXT, + ) + + assert payload == {"prompt": "Summarize the document", "context": DOCUMENT_CONTEXT} + + +@pytest.mark.respx(assert_all_called=False) +def test_agentcore_rejects_a_forwarded_param_that_is_not_json(respx_mock: respx.Router) -> None: + secret_marker: Final = "do-not-echo-this-value" + + with pytest.raises(litellm.BadRequestError, match="forward_payload_fields: context") as raised: + _payload_sent_to_agentcore(respx_mock, forward_payload_fields=["context"], context={secret_marker: {1, 2}}) + + assert secret_marker not in str(raised.value) + assert not respx_mock.calls diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 928c73d4b04..0918da9193c 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -3780,6 +3780,36 @@ class TestIsRequestBodySafeBlocksFireworksForwardUserId: ) +class TestIsRequestBodySafeBlocksForwardPayloadFields: + @pytest.mark.parametrize( + "body", + [ + pytest.param({"forward_payload_fields": ["qualifier"]}, id="root"), + pytest.param({"extra_body": {"forward_payload_fields": ["qualifier"]}}, id="extra_body"), + pytest.param({"metadata": {"forward_payload_fields": ["qualifier"]}}, id="metadata"), + ], + ) + def test_forward_payload_fields_in_request_body_is_rejected(self, body: dict[str, object]) -> None: + with pytest.raises(ValueError, match="forward_payload_fields"): + is_request_body_safe( + request_body={"model": "agentcore-model", "context": {"document": "note"}, **body}, + general_settings={}, + llm_router=None, + model="agentcore-model", + ) + + def test_forwarded_value_without_the_list_is_allowed(self) -> None: + assert ( + is_request_body_safe( + request_body={"model": "agentcore-model", "context": {"document": "note"}}, + general_settings={}, + llm_router=None, + model="agentcore-model", + ) + is True + ) + + class TestIsRequestBodySafeBlocksRustOptIn: """``rust`` hands the whole call to the Rust core, which signs and sends with its own HTTP client rather than the one the deployment configured, and From 209602939a41fcef0d67f1b7f73fe5202345e193 Mon Sep 17 00:00:00 2001 From: Brendan Smith-Elion Date: Thu, 5 Nov 2026 08:31:11 -0500 Subject: [PATCH 2/2] fix(caching): key the response cache on params a deployment forwards into the payload With forward_payload_fields set, a forwarded param such as context changes the agent's answer, but the cache key skipped provider params unless enable_caching_on_provider_specific_optional_params was on. Two requests with the same messages and different context shared one entry, so the second caller got the first caller's answer Provider params now join the cache key whenever the request carries a non-empty forward_payload_fields --- litellm/caching/caching.py | 6 +++++- tests/unit/caching/test_caching.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 3b7ea66ad6b..f8dbc038f58 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -417,6 +417,10 @@ class Cache: combined_kwargs: Final = ModelParamHelper.get_all_llm_api_params() is_semantic_cache: Final = self._is_semantic_cache() scope_excluded_params: Final = self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS if is_semantic_cache else frozenset() + # A deployment that forwards request params into the provider payload makes their values part of the answer + keys_provider_params: Final = litellm.enable_caching_on_provider_specific_optional_params is True or bool( + kwargs.get("forward_payload_fields") + ) for param in kwargs: if param in scope_excluded_params: continue @@ -427,7 +431,7 @@ class Cache: if param_value is not None: cache_key += f"{param}: {param_value}" elif not is_litellm_owned_kwarg(param): - if litellm.enable_caching_on_provider_specific_optional_params is True: # feature flagged for now + if keys_provider_params: if kwargs[param] is None: continue # ignore None params cache_key += f"{param}: {kwargs[param]}" diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index 0911d5ef888..48543933a62 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -472,6 +472,20 @@ def test_cache_key_tells_apart_requests_that_differ_only_in_continuation_token() assert cache.get_cache_key(**request, continuation_token="") == first_segment_key +def test_cache_key_tells_apart_requests_that_forward_different_payload_fields() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + request: Final = { + "model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:111111111111:runtime/docs_agent", + "messages": [{"role": "user", "content": "Summarize the document"}], + "forward_payload_fields": ["context"], + } + + first_document_key: Final = cache.get_cache_key(**request, context={"document": "Q3 revenue grew 4 percent."}) + + assert cache.get_cache_key(**request, context={"document": "Q4 revenue fell 2 percent."}) != first_document_key + assert cache.get_cache_key(**request, context={"document": "Q3 revenue grew 4 percent."}) == first_document_key + + def test_cache_key_tells_apart_requests_that_ask_for_token_ids() -> None: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) request: Final = {"model": "hosted_vllm/qwen3", "messages": [{"role": "user", "content": "count"}]}