From ae18f055ee3ce5983b3023d533da7b7cc9676d22 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 19 Aug 2026 19:07:17 +0000 Subject: [PATCH] fix(search): harden AgentCore gateway trust, error and SSE handling Refuse to SigV4-sign requests to hosts that are neither an AgentCore gateway hostname nor AGENTCORE_GATEWAY_URL's host, match gateway hostnames on the URL host instead of anywhere in the URL, accept the env token when api_base is a real gateway, raise on tools/call responses with result.isError, and split CRLF-framed SSE events. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/search/transformation.py | 52 ++++++- .../test_agentcore_search_transformation.py | 141 +++++++++++++++--- 2 files changed, 161 insertions(+), 32 deletions(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 5567dc8403e..85c5818031b 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -71,13 +71,19 @@ AGENTCORE_TOOL_NAME_SUFFIX: Final = "___WebSearch" # servers that predate the header ignore it. AGENTCORE_MCP_PROTOCOL_VERSION: Final = "2025-06-18" -_GATEWAY_REGION_PATTERN: Final = re.compile(r"\.gateway\.bedrock-agentcore\.([a-z0-9-]+)\.amazonaws\.com") +# Matched against the URL host so a crafted path or query string can't pass for +# a gateway hostname. +_GATEWAY_HOST_PATTERN: Final = re.compile(r"[a-z0-9-]+\.gateway\.bedrock-agentcore\.([a-z0-9-]+)\.amazonaws\.com") -_SSE_EVENT_SEPARATOR: Final = re.compile(r"\n[ \t]*\n") +_SSE_EVENT_SEPARATOR: Final = re.compile(r"\r?\n[ \t]*\r?\n") _SSE_LINE_PREFIXES: Final = ("event:", "data:", ":", "id:", "retry:") +def _gateway_host_match(api_base: str) -> re.Match[str] | None: + return _GATEWAY_HOST_PATTERN.fullmatch(httpx.URL(api_base).host) + + def _string_field(item: Mapping[str, object], *keys: str) -> str | None: return next( (value for key in keys if isinstance(value := item.get(key), str) and value), @@ -245,16 +251,17 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): if not isinstance(request_data, dict): raise TypeError("AgentCore search expects a single dict request body") - # Server-managed token fallback is gated on the request targeting the - # operator-configured gateway host, otherwise an authenticated caller - # could point api_base at their own server (e.g. via - # /search_tools/test_connection) and receive AGENTCORE_GATEWAY_TOKEN. + # Server-managed credentials only go to a trusted host, otherwise an + # authenticated caller could point api_base at their own server (e.g. via + # /search_tools/test_connection) and collect AGENTCORE_GATEWAY_TOKEN or a + # SigV4 signature with the proxy's credential scope and session token. + gateway_host_match: Final = _gateway_host_match(api_base) bearer_token: Final = self.resolve_server_api_key( caller_api_key=api_key, caller_api_base=api_base, key_env_vars=("AGENTCORE_GATEWAY_TOKEN",), base_env_var="AGENTCORE_GATEWAY_URL", - default_api_base=None, + default_api_base=api_base if gateway_host_match else None, ) if bearer_token: bearer_headers: Final = { # mutable-ok: httpx request headers are a dict @@ -263,6 +270,13 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): } return bearer_headers, json.dumps(request_data).encode() + if gateway_host_match is None and not self._is_configured_gateway(api_base): + raise ValueError( + f"Refusing to send SigV4-signed AgentCore requests to '{api_base}': it is neither an " + "AgentCore gateway hostname nor the host in AGENTCORE_GATEWAY_URL. Set " + "AGENTCORE_GATEWAY_URL to authorize a custom gateway hostname." + ) + signing_params: Final = ( optional_params if optional_params.get("aws_region_name") is not None @@ -284,6 +298,13 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): api_key="", ) + @staticmethod + def _is_configured_gateway(api_base: str) -> bool: + configured: Final = get_secret_str("AGENTCORE_GATEWAY_URL") + if not configured: + return False + return httpx.URL(configured).host == httpx.URL(api_base).host + @staticmethod def _signing_region(api_base: str) -> str: """ @@ -296,7 +317,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): nothing rather than silently signing for a guessed region the gateway would reject with a confusing auth error. """ - match: Final = _GATEWAY_REGION_PATTERN.search(api_base) + match: Final = _gateway_host_match(api_base) if match: return match.group(1) @@ -336,6 +357,15 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): message=f"AgentCore gateway MCP error: {error}", ) + # A failed tools/call is reported in-band, as HTTP 200 with result.isError + # and the failure text where the results would be. + result: Final = response_json.get("result") + if isinstance(result, dict) and result.get("isError"): + raise BedrockError( + status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, + message=f"AgentCore web search tool error: {self._tool_error_message(response_json)}", + ) + return SearchResponse( results=[ # mutable-ok: SearchResponse.results is a pydantic list field _to_search_result(item) @@ -345,6 +375,12 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): object="search", ) + def _tool_error_message(self, response_json: Mapping[str, object]) -> str: + texts: Final = tuple( + text for block in self._text_blocks(response_json) if isinstance(text := block.get("text"), str) + ) + return " ".join(texts) if texts else json.dumps(response_json.get("result"))[:500] + @staticmethod def _text_blocks(response_json: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: result: Final = response_json.get("result") diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py index 63ec2286c3d..889f78c58b9 100644 --- a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py @@ -231,6 +231,37 @@ class TestAgentCoreSearch: with pytest.raises(Exception, match="tool not found"): config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + def test_transform_search_response_raises_on_tool_error(self): + """A failed tools/call comes back as HTTP 200 with result.isError; it must not be + reported to the caller as a successful search with zero results.""" + config = AgentCoreSearchConfig() + mock_response = _make_mock_response( + { + "jsonrpc": "2.0", + "id": 1, + "result": { + "isError": True, + "content": [{"type": "text", "text": "AccessDeniedException: not authorized"}], + }, + } + ) + with pytest.raises(Exception, match="AccessDeniedException"): + config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + + def test_transform_search_response_parses_crlf_framed_sse(self): + """SSE streams may be CRLF framed; events must still split into separate events.""" + config = AgentCoreSearchConfig() + progress = {"jsonrpc": "2.0", "method": "notifications/progress", "params": {"progress": 1}} + sse_text = ( + f"event: message\r\ndata: {json.dumps(progress)}\r\n\r\n" + f"event: message\r\ndata: {json.dumps(_mcp_response_body())}\r\n\r\n" + ) + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + assert response.results[0].title == "Test Result 1" + def test_sign_request_uses_bearer_token_when_api_key_set(self): """CUSTOM_JWT gateways: api_key is sent as a bearer token, no SigV4.""" config = AgentCoreSearchConfig() @@ -280,6 +311,54 @@ class TestAgentCoreSearch: os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) os.environ.pop("AGENTCORE_GATEWAY_URL", None) + def test_sign_request_uses_env_token_for_gateway_api_base_without_gateway_url(self): + """api_base pointing at a real gateway is a trusted destination for the env token, + so operators configuring api_base in yaml don't also need AGENTCORE_GATEWAY_URL.""" + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + try: + headers, _ = config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + assert headers["Authorization"] == "Bearer env-jwt-token" + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + + @pytest.mark.parametrize( + "untrusted_api_base", + [ + "https://attacker.example.com/mcp", + # gateway hostname in the path/query must not pass for the host + "https://attacker.example.com/gw.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp", + ], + ) + def test_sign_request_refuses_sigv4_to_untrusted_host(self, untrusted_api_base): + """A SigV4 signature carries the proxy's credential scope and session token, so it + must never be sent to a host that is not the operator's gateway.""" + config = AgentCoreSearchConfig() + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + try: + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + with pytest.raises(ValueError, match="Refusing to send"): + config.sign_request( + headers={}, + optional_params={"aws_region_name": "us-east-1"}, + request_data={"jsonrpc": "2.0"}, + api_base=untrusted_api_base, + ) + mock_base_sign.assert_not_called() + finally: + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + def test_sign_request_does_not_leak_bedrock_bearer_token(self): """AWS_BEARER_TOKEN_BEDROCK is a Bedrock Runtime credential — it must not replace SigV4 on requests to an AgentCore gateway.""" @@ -303,39 +382,49 @@ class TestAgentCoreSearch: def test_sign_request_custom_hostname_requires_region(self): """Custom hostname + empty AWS config chain → clear error, no guessed region.""" config = AgentCoreSearchConfig() + custom_url = "https://gateway.internal.example.com/mcp" + os.environ["AGENTCORE_GATEWAY_URL"] = custom_url mock_session = MagicMock() mock_session.region_name = None # nothing configured anywhere - with patch("boto3.Session", return_value=mock_session): - with pytest.raises(ValueError, match="signing region"): - config.sign_request( - headers={}, - optional_params={}, - request_data={"jsonrpc": "2.0"}, - api_base="https://gateway.internal.example.com/mcp", - ) + try: + with patch("boto3.Session", return_value=mock_session): + with pytest.raises(ValueError, match="signing region"): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=custom_url, + ) + finally: + os.environ.pop("AGENTCORE_GATEWAY_URL", None) def test_sign_request_custom_hostname_uses_shared_config_region(self): """Custom hostname + region from AWS shared config (profile) must be honored.""" config = AgentCoreSearchConfig() + custom_url = "https://gateway.internal.example.com/mcp" + os.environ["AGENTCORE_GATEWAY_URL"] = custom_url mock_session = MagicMock() mock_session.region_name = "eu-west-1" # e.g. from ~/.aws/config profile - with ( - patch("boto3.Session", return_value=mock_session), - patch.object( - AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM - "_sign_request", - return_value=({}, b"{}"), - ) as mock_base_sign, - ): - config.sign_request( - headers={}, - optional_params={}, - request_data={"jsonrpc": "2.0"}, - api_base="https://gateway.internal.example.com/mcp", - ) - assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-west-1" + try: + with ( + patch("boto3.Session", return_value=mock_session), + patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign, + ): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=custom_url, + ) + assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-west-1" + finally: + os.environ.pop("AGENTCORE_GATEWAY_URL", None) def test_sign_request_passes_explicit_aws_credentials(self): """Explicit aws_* params (e.g. from a proxy search_tools entry) reach the signer.""" @@ -434,7 +523,11 @@ class TestAgentCoreSearchEdgeCases: assert getattr(err, "status_code", None) == 503 assert "boom" in str(err) - def test_search_cost_lookup_is_mapped(self): + def test_search_cost_lookup_is_mapped(self, monkeypatch): + """Assert against the map in this checkout: the remote cost map litellm loads by + default only carries providers already released.""" + from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.search.cost_calculator import search_provider_cost_per_query + monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map()) assert search_provider_cost_per_query(model="agentcore/search", custom_llm_provider="agentcore") == (0.0, 0.0)