scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

fix(search): harden AgentCore gateway trust, error and SSE handling

Browse source 
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>
This commit is contained in:
yassin 2026-08-19 19:07:17 +00:00
parent 17b72d5089
commit ae18f055ee
2 changed files with 161 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
litellm/llms/bedrock/search/transformation.py
View file

@ -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")

141
tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py
View file

@ -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)