From 737e052a7ae07963474285b062ad64875bb119c6 Mon Sep 17 00:00:00 2001 From: Sundar Raghavan Date: Thu, 3 Sep 2026 12:31:38 -0700 Subject: [PATCH 1/3] feat(tools): add WebSearchClient for invoking Amazon Web Search The SDK can create a web search connector target on a gateway, but there is no way to call the resulting tool from Python. Anything that is not already an MCP client has to hand-roll the SigV4 signed MCP handshake to use it. WebSearchClient closes that gap: client = WebSearchClient(region="us-east-1", gateway_id="my-gateway-abc123") for result in client.search("what is agentcore", max_results=5): print(result.title, result.url) Details: - Transport sits behind a WebSearchBackend seam. GatewayMcpBackend speaks the slice of MCP streamable HTTP that one tool call needs (initialize, the initialized notification, optionally tools/list, then tools/call) using urllib3 and botocore.auth.SigV4Auth, so no new dependency is added. If the direct web search API arrives later, it is a second backend behind the same search() signature. - Both response framings are handled, since a gateway may answer a POST with application/json or with text/event-stream. - Tool name resolution accounts for Gateway prefixing every tool with its target name: target_name derives "___WebSearch" directly, otherwise tools/list is walked (following nextCursor) and the WebSearch tool is picked, erroring when the choice is ambiguous. - Inputs are validated against the documented limits before the call: query is required and capped at 200 characters, max_results at 1 to 25. Domain and published-date filters need connector version 1.2.0 or later on the target. - Results are returned as WebSearchResult with url, title and published_date retained, because citations have to be displayable in anything shown to an end user. - get_gateway_mcp_endpoint validates the gateway identifier as a DNS label before interpolating it into the hostname, matching the existing region validation, so a crafted identifier cannot redirect the request off AWS. - A region outside the connector's availability warns instead of failing, so a stale constant never blocks a call to a newly added region. --- src/bedrock_agentcore/_utils/endpoints.py | 36 + src/bedrock_agentcore/tools/__init__.py | 12 + .../tools/web_search_client.py | 648 ++++++++++++++++ .../test_region_validation.py | 51 ++ .../tools/test_web_search_client.py | 693 ++++++++++++++++++ tests_integ/tools/test_web_search_client.py | 101 +++ 6 files changed, 1541 insertions(+) create mode 100644 src/bedrock_agentcore/tools/web_search_client.py create mode 100644 tests/bedrock_agentcore/tools/test_web_search_client.py create mode 100644 tests_integ/tools/test_web_search_client.py diff --git a/src/bedrock_agentcore/_utils/endpoints.py b/src/bedrock_agentcore/_utils/endpoints.py index 9f37fba9..946b760a 100644 --- a/src/bedrock_agentcore/_utils/endpoints.py +++ b/src/bedrock_agentcore/_utils/endpoints.py @@ -13,6 +13,19 @@ # Uses \A and \Z anchors to prevent newline injection bypass that $ allows. _VALID_REGION_PATTERN = re.compile(r"\A[a-z]{2}(-[a-z]+)+-\d+\Z") +# A gateway identifier becomes a DNS label in the gateway's MCP endpoint, so it is +# constrained to the characters a label allows. Anchored with \A and \Z for the same +# reason as the region pattern. +_VALID_GATEWAY_ID_PATTERN = re.compile(r"\A[a-zA-Z0-9][a-zA-Z0-9-]{0,62}\Z") + + +class InvalidGatewayIdentifierError(ValueError): + """Raised when a gateway identifier is not a valid DNS label. + + The identifier is interpolated into the endpoint hostname, so an + unvalidated value could redirect requests to a non-AWS host. + """ + class InvalidRegionError(ValueError): """Raised when an invalid AWS region string is provided. @@ -79,3 +92,26 @@ def get_control_plane_endpoint(region: str = DEFAULT_REGION) -> str: validate_region(region) url = f"https://bedrock-agentcore-control.{region}.amazonaws.com" return _validate_endpoint_url(url) + + +def get_gateway_mcp_endpoint(gateway_id: str, region: str = DEFAULT_REGION) -> str: + """Build the MCP endpoint URL for a gateway. + + Args: + gateway_id: The gateway identifier (not an ARN). + region: The region the gateway lives in. + + Returns: + The gateway's streamable HTTP MCP endpoint URL. + + Raises: + InvalidGatewayIdentifierError: If the identifier is not a valid DNS label. + InvalidRegionError: If the region is malformed or the URL resolves off-AWS. + """ + if not isinstance(gateway_id, str) or not _VALID_GATEWAY_ID_PATTERN.match(gateway_id): + raise InvalidGatewayIdentifierError( + f"Invalid gateway identifier: {gateway_id!r}. Expected a gateway ID such as 'my-gateway-abc123'." + ) + validate_region(region) + url = f"https://{gateway_id}.gateway.bedrock-agentcore.{region}.amazonaws.com/mcp" + return _validate_endpoint_url(url) diff --git a/src/bedrock_agentcore/tools/__init__.py b/src/bedrock_agentcore/tools/__init__.py index 8ff3af90..f33bf3f9 100644 --- a/src/bedrock_agentcore/tools/__init__.py +++ b/src/bedrock_agentcore/tools/__init__.py @@ -26,6 +26,13 @@ VpcConfig, create_browser_config, ) +from .web_search_client import ( + WebSearchBackend, + WebSearchClient, + WebSearchError, + WebSearchResponse, + WebSearchResult, +) __all__ = [ "BasicAuth", @@ -53,5 +60,10 @@ "SessionConfiguration", "ViewportConfiguration", "VpcConfig", + "WebSearchBackend", + "WebSearchClient", + "WebSearchError", + "WebSearchResponse", + "WebSearchResult", "create_browser_config", ] diff --git a/src/bedrock_agentcore/tools/web_search_client.py b/src/bedrock_agentcore/tools/web_search_client.py new file mode 100644 index 00000000..5c1a1c24 --- /dev/null +++ b/src/bedrock_agentcore/tools/web_search_client.py @@ -0,0 +1,648 @@ +"""Client for AgentCore Web Search. + +Web Search is reachable today as an AgentCore Gateway connector target, which the +agent calls as an MCP tool. This module wraps that so callers get a plain +``search()`` method and a normalized result type instead of MCP content blocks: + + >>> from bedrock_agentcore.tools import WebSearchClient + >>> + >>> client = WebSearchClient(region="us-east-1", gateway_id="my-gateway-abc123") + >>> response = client.search("what shipped in python 3.13", max_results=5) + >>> for result in response.results: + ... print(result.title, result.url) + +The transport lives behind :class:`WebSearchBackend` so the same ``search()`` +signature and the same :class:`WebSearchResult` can be served by a different +backend later without changing callers. +""" + +import json +import logging +import threading +from dataclasses import dataclass, field +from typing import Any, Dict, Iterator, List, Optional, Sequence + +import urllib3 +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest + +from bedrock_agentcore._utils.endpoints import get_gateway_mcp_endpoint +from bedrock_agentcore._utils.user_agent import SDK_VERSION, build_user_agent_suffix + +logger = logging.getLogger(__name__) + +#: Name of the MCP tool the web search connector exposes. Fixed by the service. +WEB_SEARCH_TOOL_NAME = "WebSearch" + +#: Gateway prefixes every tool it exposes with the name of the target it came from. +GATEWAY_TOOL_NAME_DELIMITER = "___" + +#: Default target name used by ``GatewayClient.create_web_search_target``. +DEFAULT_TARGET_NAME = "amazon-web-search" + +#: Service name the gateway data plane signs as. +GATEWAY_SIGNING_SERVICE = "bedrock-agentcore" + +#: Documented input limits for the WebSearch tool. +MAX_QUERY_LENGTH = 200 +MIN_MAX_RESULTS = 1 +MAX_MAX_RESULTS = 25 + +#: Regions where the web search connector is offered. Used for a warning only, never +#: to block a call, so that a newly added region does not require an SDK release. +KNOWN_REGIONS = ("us-east-1", "eu-west-1", "ap-northeast-1") + +_MCP_PROTOCOL_VERSION = "2025-06-18" +_DEFAULT_TIMEOUT = 30 + + +class WebSearchError(RuntimeError): + """Raised when a web search call fails.""" + + +@dataclass(frozen=True) +class WebSearchResult: + """A single web search result. + + Attributes: + text: The extracted snippet relevant to the query. Always present. + url: URL of the source page. + title: Title of the source page. + published_date: Publication date of the page, as reported by the index. + """ + + text: str + url: Optional[str] = None + title: Optional[str] = None + published_date: Optional[str] = None + + @classmethod + def from_payload(cls, payload: Dict[str, Any]) -> "WebSearchResult": + """Build a result from one entry of a search response.""" + return cls( + text=payload.get("text") or "", + url=payload.get("url"), + title=payload.get("title"), + published_date=payload.get("publishedDate"), + ) + + +@dataclass(frozen=True) +class WebSearchResponse: + """The results of one web search. + + Attributes: + results: The results, in the order the service returned them. + search_id: Service-assigned identifier for the search, when present. + """ + + results: List[WebSearchResult] = field(default_factory=list) + search_id: Optional[str] = None + + def __len__(self) -> int: + """Number of results.""" + return len(self.results) + + def __iter__(self) -> Iterator[WebSearchResult]: + """Iterate over the results.""" + return iter(self.results) + + @classmethod + def from_payload(cls, payload: Dict[str, Any]) -> "WebSearchResponse": + """Build a response from the decoded search payload.""" + raw_results = payload.get("results") or [] + return cls( + results=[WebSearchResult.from_payload(item) for item in raw_results if isinstance(item, dict)], + search_id=payload.get("id"), + ) + + +def _build_arguments( + query: str, + max_results: Optional[int] = None, + include_domains: Optional[Sequence[str]] = None, + exclude_domains: Optional[Sequence[str]] = None, + published_after: Optional[str] = None, + published_before: Optional[str] = None, +) -> Dict[str, Any]: + """Validate search inputs and shape them into the tool's argument object. + + Raises: + ValueError: If the query is empty or over the documented length limit, or + if max_results falls outside the documented range. + """ + if not query or not query.strip(): + raise ValueError("query must be a non-empty string") + if len(query) > MAX_QUERY_LENGTH: + raise ValueError(f"query must be {MAX_QUERY_LENGTH} characters or fewer, got {len(query)}") + + arguments: Dict[str, Any] = {"query": query} + + if max_results is not None: + if not isinstance(max_results, int) or isinstance(max_results, bool): + raise ValueError(f"max_results must be an integer, got {type(max_results).__name__}") + if not MIN_MAX_RESULTS <= max_results <= MAX_MAX_RESULTS: + raise ValueError(f"max_results must be between {MIN_MAX_RESULTS} and {MAX_MAX_RESULTS}, got {max_results}") + arguments["maxResults"] = max_results + + filters: Dict[str, Any] = {} + domain_filter: Dict[str, List[str]] = {} + if include_domains: + domain_filter["include"] = list(include_domains) + if exclude_domains: + domain_filter["exclude"] = list(exclude_domains) + if domain_filter: + filters["domainFilter"] = domain_filter + + published_filter: Dict[str, str] = {} + if published_after: + published_filter["from"] = published_after + if published_before: + published_filter["to"] = published_before + if published_filter: + filters["publishedDateFilter"] = published_filter + + if filters: + arguments["filters"] = filters + + return arguments + + +def _extract_search_payload(result: Dict[str, Any]) -> Dict[str, Any]: + """Pull the search payload out of an MCP ``tools/call`` result. + + The connector returns the results as a JSON document inside a text content + block, so the text has to be decoded rather than read directly. + + Raises: + WebSearchError: If the tool reported an error or returned no decodable + text content. + """ + if result.get("isError"): + raise WebSearchError(f"Web search tool reported an error: {_first_text(result) or result}") + + text = _first_text(result) + if text is None: + raise WebSearchError(f"Web search response contained no text content: {result}") + + try: + payload = json.loads(text) + except json.JSONDecodeError as exc: + raise WebSearchError(f"Could not decode web search response as JSON: {text[:200]!r}") from exc + + if not isinstance(payload, dict): + raise WebSearchError(f"Expected a JSON object in the web search response, got {type(payload).__name__}") + return payload + + +def _first_text(result: Dict[str, Any]) -> Optional[str]: + """Return the first text content block of an MCP result, if any.""" + for block in result.get("content") or []: + if isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str): + return block["text"] + return None + + +class WebSearchBackend: + """How a :class:`WebSearchClient` reaches web search. + + A backend takes the tool's argument object and returns the decoded search + payload, meaning a dict shaped ``{"id": ..., "results": [...]}``. Everything + above this line is transport independent. + """ + + def search(self, arguments: Dict[str, Any]) -> Dict[str, Any]: + """Run one search and return the decoded payload.""" + raise NotImplementedError + + def close(self) -> None: + """Release any resources held by the backend.""" + + +class GatewayMcpBackend(WebSearchBackend): + """Reaches web search through an AgentCore Gateway target over MCP. + + Speaks the subset of MCP streamable HTTP that one tool call needs -- initialize, + the initialized notification, optionally ``tools/list``, then ``tools/call`` -- + signing each request with SigV4. It is deliberately narrow: it is not a general + MCP client, and it holds no dependency beyond what the SDK already requires. + + Both response framings the transport allows are handled, since a gateway may + answer a POST with either ``application/json`` or ``text/event-stream``. + """ + + def __init__( + self, + endpoint: str, + region: str, + *, + boto3_session: Optional[Any] = None, + tool_name: Optional[str] = None, + target_name: Optional[str] = None, + timeout: float = _DEFAULT_TIMEOUT, + integration_source: Optional[str] = None, + signing_service: str = GATEWAY_SIGNING_SERVICE, + ): + """Initialize the backend. + + Args: + endpoint: The gateway's MCP endpoint URL. + region: Region to sign for. + boto3_session: Session to take credentials from. Defaults to a new session. + tool_name: Fully qualified tool name. Skips discovery when given. + target_name: Target the connector was added under. Used to derive the tool + name without a ``tools/list`` round trip. + timeout: Per-request timeout in seconds. + integration_source: Optional framework identifier for the User-Agent. + signing_service: SigV4 service name. + """ + import boto3 + + self._endpoint = endpoint + self._region = region + self._session = boto3_session or boto3.Session() + self._signing_service = signing_service + self._timeout = timeout + self._user_agent = f"python-urllib3/{urllib3.__version__} {build_user_agent_suffix(integration_source)}" + + self._tool_name = tool_name + self._target_name = target_name + + # A single signed POST per call, so retries are left to the caller: replaying a + # tools/call is not always safe and the signature is only valid for a few minutes. + self._http = urllib3.PoolManager(retries=urllib3.Retry(total=0, redirect=0)) + + self._lock = threading.Lock() + self._mcp_session_id: Optional[str] = None + self._protocol_version = _MCP_PROTOCOL_VERSION + self._initialized = False + self._request_id = 0 + + # Transport + # ------------------------------------------------------------------------- + def _next_id(self) -> int: + self._request_id += 1 + return self._request_id + + def _signed_headers(self, body: bytes, extra: Dict[str, str]) -> Dict[str, str]: + """Sign a request body with SigV4 and return the headers to send.""" + credentials = self._session.get_credentials() + if credentials is None: + raise WebSearchError("No AWS credentials available. Configure credentials before calling web search.") + + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "Content-Length": str(len(body)), + "User-Agent": self._user_agent, + **extra, + } + request = AWSRequest(method="POST", url=self._endpoint, data=body, headers=headers) + SigV4Auth(credentials.get_frozen_credentials(), self._signing_service, self._region).add_auth(request) + return dict(request.headers) + + def _post(self, message: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """Send one JSON-RPC message and return the decoded reply, if there is one.""" + body = json.dumps(message).encode("utf-8") + + extra: Dict[str, str] = {} + if self._mcp_session_id: + extra["Mcp-Session-Id"] = self._mcp_session_id + if self._initialized: + extra["MCP-Protocol-Version"] = self._protocol_version + + response = self._http.request( + "POST", + self._endpoint, + body=body, + headers=self._signed_headers(body, extra), + timeout=urllib3.Timeout(total=self._timeout), + preload_content=True, + ) + + session_id = response.headers.get("Mcp-Session-Id") + if session_id: + self._mcp_session_id = session_id + + if response.status >= 400: + body_text = response.data.decode("utf-8", "replace")[:500] + raise WebSearchError(f"Web search request failed with HTTP {response.status}: {body_text}") + + if not response.data: + return None + + reply = _decode_jsonrpc(response.headers.get("Content-Type", ""), response.data) + if reply is None: + return None + if "error" in reply: + error = reply["error"] or {} + raise WebSearchError(f"Gateway returned a JSON-RPC error {error.get('code')}: {error.get('message')}") + return reply + + # MCP session + # ------------------------------------------------------------------------- + def _ensure_initialized(self) -> None: + if self._initialized: + return + + reply = self._post( + { + "jsonrpc": "2.0", + "id": self._next_id(), + "method": "initialize", + "params": { + "protocolVersion": _MCP_PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": {"name": "bedrock-agentcore-python", "version": SDK_VERSION}, + }, + } + ) + if reply is None: + raise WebSearchError("Gateway did not answer the MCP initialize request") + + negotiated = (reply.get("result") or {}).get("protocolVersion") + if isinstance(negotiated, str) and negotiated: + self._protocol_version = negotiated + + self._initialized = True + self._post({"jsonrpc": "2.0", "method": "notifications/initialized"}) + + def _ensure_tool_name(self) -> str: + """Resolve the fully qualified tool name, discovering it if necessary.""" + if self._tool_name: + return self._tool_name + + if self._target_name: + self._tool_name = f"{self._target_name}{GATEWAY_TOOL_NAME_DELIMITER}{WEB_SEARCH_TOOL_NAME}" + return self._tool_name + + candidates = [ + name + for name in self._list_tool_names() + if name == WEB_SEARCH_TOOL_NAME or name.endswith(f"{GATEWAY_TOOL_NAME_DELIMITER}{WEB_SEARCH_TOOL_NAME}") + ] + if not candidates: + raise WebSearchError( + f"No {WEB_SEARCH_TOOL_NAME} tool found on {self._endpoint}. " + "Add a web search connector target to the gateway, or pass target_name." + ) + if len(candidates) > 1: + raise WebSearchError( + f"Gateway exposes more than one {WEB_SEARCH_TOOL_NAME} tool ({', '.join(sorted(candidates))}). " + "Pass target_name to choose one." + ) + + self._tool_name = candidates[0] + logger.debug("Resolved web search tool name to %s", self._tool_name) + return self._tool_name + + def _list_tool_names(self) -> List[str]: + """List every tool the gateway exposes, following pagination.""" + names: List[str] = [] + cursor: Optional[str] = None + while True: + params: Dict[str, Any] = {"cursor": cursor} if cursor else {} + reply = self._post({"jsonrpc": "2.0", "id": self._next_id(), "method": "tools/list", "params": params}) + result = (reply or {}).get("result") or {} + for tool in result.get("tools") or []: + if isinstance(tool, dict) and isinstance(tool.get("name"), str): + names.append(tool["name"]) + cursor = result.get("nextCursor") + if not cursor: + return names + + # WebSearchBackend + # ------------------------------------------------------------------------- + def search(self, arguments: Dict[str, Any]) -> Dict[str, Any]: + """Call the WebSearch tool and return the decoded search payload.""" + with self._lock: + self._ensure_initialized() + tool_name = self._ensure_tool_name() + reply = self._post( + { + "jsonrpc": "2.0", + "id": self._next_id(), + "method": "tools/call", + "params": {"name": tool_name, "arguments": arguments}, + } + ) + if reply is None: + raise WebSearchError("Gateway did not answer the web search tool call") + return _extract_search_payload(reply.get("result") or {}) + + def close(self) -> None: + """Close the connection pool.""" + self._http.clear() + self._initialized = False + self._mcp_session_id = None + + +def _decode_jsonrpc(content_type: str, data: bytes) -> Optional[Dict[str, Any]]: + """Decode a JSON-RPC reply from either a JSON body or an SSE stream. + + Returns None when the body carries no JSON-RPC message, which is what a + notification acknowledgement looks like. + """ + text = data.decode("utf-8", "replace") + + if "text/event-stream" in content_type.lower(): + for line in text.splitlines(): + if not line.startswith("data:"): + continue + chunk = line[len("data:") :].strip() + if not chunk: + continue + try: + message = json.loads(chunk) + except json.JSONDecodeError: + continue + if isinstance(message, dict) and "jsonrpc" in message: + return message + return None + + try: + message = json.loads(text) + except json.JSONDecodeError as exc: + raise WebSearchError(f"Could not decode gateway response as JSON: {text[:200]!r}") from exc + return message if isinstance(message, dict) else None + + +class WebSearchClient: + """Client for AgentCore Web Search. + + Attributes: + region (str): The region being used. + backend (WebSearchBackend): The transport in use. + + Basic Usage: + >>> from bedrock_agentcore.tools import WebSearchClient + >>> + >>> client = WebSearchClient(region="us-east-1", gateway_id="my-gateway-abc123") + >>> response = client.search("latest boto3 release notes") + >>> response.results[0].url + + Context Manager: + >>> with WebSearchClient(region="us-east-1", gateway_id="my-gateway-abc123") as client: + ... response = client.search("who maintains urllib3") + """ + + def __init__( + self, + region: Optional[str] = None, + *, + gateway_id: Optional[str] = None, + gateway_arn: Optional[str] = None, + gateway_endpoint: Optional[str] = None, + target_name: Optional[str] = None, + tool_name: Optional[str] = None, + backend: Optional[WebSearchBackend] = None, + boto3_session: Optional[Any] = None, + timeout: float = _DEFAULT_TIMEOUT, + integration_source: Optional[str] = None, + ): + """Initialize the client. + + Exactly one of ``gateway_id``, ``gateway_arn``, ``gateway_endpoint`` or + ``backend`` identifies where the search goes. The gateway arguments are + keyword only and optional so that a future transport needing none of them + is an addition rather than a breaking change. + + Args: + region: Region to call. Defaults to the session's region. + gateway_id: ID of a gateway with a web search connector target. + gateway_arn: ARN of that gateway. The ID and region are read from it. + gateway_endpoint: A gateway MCP endpoint URL, if you already have one. + target_name: Name of the connector target. Supplying it avoids a + ``tools/list`` round trip on the first search. + tool_name: Fully qualified tool name, if you already know it. + backend: A backend to use as is. Overrides every gateway argument. + boto3_session: Session to take credentials and region from. + timeout: Per-request timeout in seconds. + integration_source: Optional framework identifier for the User-Agent. + + Raises: + ValueError: If no gateway is identified, or more than one is. + """ + import boto3 + + self._session = boto3_session or boto3.Session() + self._owns_backend = backend is None + + if backend is not None: + if any(value is not None for value in (gateway_id, gateway_arn, gateway_endpoint)): + raise ValueError("Pass either backend or one of gateway_id/gateway_arn/gateway_endpoint, not both") + self.region = region or self._session.region_name + self.backend: WebSearchBackend = backend + return + + given = [ + name + for name, value in ( + ("gateway_id", gateway_id), + ("gateway_arn", gateway_arn), + ("gateway_endpoint", gateway_endpoint), + ) + if value + ] + if len(given) > 1: + raise ValueError(f"Pass only one of gateway_id, gateway_arn or gateway_endpoint, got {', '.join(given)}") + + if gateway_arn: + gateway_id, arn_region = _parse_gateway_arn(gateway_arn) + region = region or arn_region + + self.region = region or self._session.region_name + if not self.region: + raise ValueError("region could not be determined. Pass region= or configure a default region.") + if self.region not in KNOWN_REGIONS: + logger.warning( + "Web search is offered in %s. Calling %s may fail if the connector is not available there.", + ", ".join(KNOWN_REGIONS), + self.region, + ) + + if gateway_id: + gateway_endpoint = get_gateway_mcp_endpoint(gateway_id, self.region) + if not gateway_endpoint: + raise ValueError("One of gateway_id, gateway_arn, gateway_endpoint or backend is required") + + self.backend = GatewayMcpBackend( + endpoint=gateway_endpoint, + region=self.region, + boto3_session=self._session, + tool_name=tool_name, + target_name=target_name, + timeout=timeout, + integration_source=integration_source, + ) + + def search( + self, + query: str, + *, + max_results: Optional[int] = None, + include_domains: Optional[Sequence[str]] = None, + exclude_domains: Optional[Sequence[str]] = None, + published_after: Optional[str] = None, + published_before: Optional[str] = None, + ) -> WebSearchResponse: + """Search the web. + + The filter arguments need connector version 1.2.0 or later on the target. + On an earlier version the tool accepts only ``query`` and ``max_results``. + Target level domain rules always apply on top and cannot be relaxed here. + + Args: + query: What to search for. 200 characters or fewer. + max_results: How many results to return, 1 to 25. Service default is 10. + include_domains: Restrict results to these domains. + exclude_domains: Drop results from these domains. + published_after: Earliest publication date, ISO-8601 UTC, inclusive. + published_before: Latest publication date, ISO-8601 UTC, inclusive. + + Returns: + The search results. + + Raises: + ValueError: If the query or max_results is outside the documented limits. + WebSearchError: If the call fails or the response cannot be decoded. + """ + arguments = _build_arguments( + query=query, + max_results=max_results, + include_domains=include_domains, + exclude_domains=exclude_domains, + published_after=published_after, + published_before=published_before, + ) + return WebSearchResponse.from_payload(self.backend.search(arguments)) + + def close(self) -> None: + """Release the backend, if this client created it.""" + if self._owns_backend: + self.backend.close() + + def __enter__(self) -> "WebSearchClient": + """Enter the context manager.""" + return self + + def __exit__(self, *exc_info: Any) -> None: + """Close the client on exit.""" + self.close() + + +def _parse_gateway_arn(arn: str) -> tuple: + """Pull the gateway ID and region out of a gateway ARN. + + Raises: + ValueError: If the ARN is not a gateway ARN. + """ + parts = arn.split(":") + if len(parts) < 6 or parts[0] != "arn" or not parts[5].startswith("gateway/"): + raise ValueError( + f"Not a gateway ARN: {arn!r}. Expected 'arn:aws:bedrock-agentcore:::gateway/'." + ) + gateway_id = parts[5].split("/", 1)[1] + if not gateway_id: + raise ValueError(f"Gateway ARN carries no gateway ID: {arn!r}") + return gateway_id, parts[3] diff --git a/tests/bedrock_agentcore/test_region_validation.py b/tests/bedrock_agentcore/test_region_validation.py index cee328f7..793e5194 100644 --- a/tests/bedrock_agentcore/test_region_validation.py +++ b/tests/bedrock_agentcore/test_region_validation.py @@ -7,10 +7,12 @@ import pytest from bedrock_agentcore._utils.endpoints import ( + InvalidGatewayIdentifierError, InvalidRegionError, _validate_endpoint_url, get_control_plane_endpoint, get_data_plane_endpoint, + get_gateway_mcp_endpoint, validate_region, ) @@ -173,6 +175,55 @@ def test_govcloud_regions(self): assert "us-gov-west-1" in url +class TestGatewayMcpEndpoint: + """Tests for get_gateway_mcp_endpoint. + + The gateway identifier becomes a DNS label in the hostname, so it needs the + same treatment as the region. + """ + + def test_valid_endpoint(self): + url = get_gateway_mcp_endpoint("my-gateway-abc123", "us-east-1") + assert url == "https://my-gateway-abc123.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp" + + def test_malicious_region_rejected(self): + with pytest.raises(InvalidRegionError): + get_gateway_mcp_endpoint("my-gateway", "x@attacker.com:443/#") + + @pytest.mark.parametrize( + "gateway_id", + [ + "evil.example.com", + "gw@attacker.com", + "gw/../../", + "gw:443", + "gw#fragment", + "gw?a=b", + "gw abc", + "gw\n", + "-gw", + "", + "a" * 64, + ], + ) + def test_malicious_gateway_id_rejected(self, gateway_id): + with pytest.raises(InvalidGatewayIdentifierError): + get_gateway_mcp_endpoint(gateway_id, "us-east-1") + + def test_non_string_gateway_id_rejected(self): + with pytest.raises(InvalidGatewayIdentifierError): + get_gateway_mcp_endpoint(None, "us-east-1") # type: ignore[arg-type] + + def test_error_is_valueerror_subclass(self): + with pytest.raises(ValueError): + get_gateway_mcp_endpoint("evil.example.com", "us-east-1") + + def test_arn_is_not_accepted_as_an_identifier(self): + arn = "arn:aws:bedrock-agentcore:us-east-1:123456789012:gateway/gw-abc123" + with pytest.raises(InvalidGatewayIdentifierError): + get_gateway_mcp_endpoint(arn, "us-east-1") + + # --------------------------------------------------------------------------- # build_runtime_url (ARN extraction path) # --------------------------------------------------------------------------- diff --git a/tests/bedrock_agentcore/tools/test_web_search_client.py b/tests/bedrock_agentcore/tools/test_web_search_client.py new file mode 100644 index 00000000..d23abe60 --- /dev/null +++ b/tests/bedrock_agentcore/tools/test_web_search_client.py @@ -0,0 +1,693 @@ +"""Tests for WebSearchClient.""" + +import json +from unittest.mock import MagicMock + +import pytest + +from bedrock_agentcore._utils.endpoints import InvalidGatewayIdentifierError, InvalidRegionError +from bedrock_agentcore.tools.web_search_client import ( + GatewayMcpBackend, + WebSearchBackend, + WebSearchClient, + WebSearchError, + WebSearchResponse, + WebSearchResult, + _build_arguments, + _decode_jsonrpc, + _parse_gateway_arn, +) + +ENDPOINT = "https://gw-abc123.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp" + +SEARCH_PAYLOAD = { + "id": "search-1", + "results": [ + { + "text": "urllib3 is a HTTP client for Python.", + "url": "https://urllib3.readthedocs.io/", + "title": "urllib3 docs", + "publishedDate": "2026-01-05T00:00:00Z", + }, + {"text": "Only text is required."}, + ], +} + + +def _http_response(status=200, body=b"", headers=None, content_type="application/json"): + response = MagicMock() + response.status = status + response.data = body + merged = {"Content-Type": content_type} + merged.update(headers or {}) + response.headers = merged + return response + + +def _json_rpc_response(result, request_id=1, **kwargs): + body = json.dumps({"jsonrpc": "2.0", "id": request_id, "result": result}).encode() + return _http_response(body=body, **kwargs) + + +def _initialize_response(): + return _json_rpc_response( + {"protocolVersion": "2025-06-18", "capabilities": {"tools": {}}, "serverInfo": {"name": "gateway"}}, + headers={"Mcp-Session-Id": "sess-1"}, + ) + + +def _tools_call_response(payload=None): + return _json_rpc_response( + {"content": [{"type": "text", "text": json.dumps(payload if payload is not None else SEARCH_PAYLOAD)}]}, + request_id=2, + ) + + +def _make_backend(responses, **kwargs): + """Build a backend whose HTTP layer replays the given responses in order.""" + session = MagicMock() + session.get_credentials.return_value.get_frozen_credentials.return_value = _frozen_credentials() + + kwargs.setdefault("tool_name", "amazon-web-search___WebSearch") + backend = GatewayMcpBackend(endpoint=ENDPOINT, region="us-east-1", boto3_session=session, **kwargs) + backend._http = MagicMock() + backend._http.request.side_effect = list(responses) + return backend + + +def _frozen_credentials(): + from botocore.credentials import ReadOnlyCredentials + + return ReadOnlyCredentials("AKIAEXAMPLE", "secret", None) + + +class TestBuildArguments: + """Tests for input validation and argument shaping.""" + + def test_minimal(self): + assert _build_arguments("hello") == {"query": "hello"} + + def test_all_options(self): + arguments = _build_arguments( + "hello", + max_results=5, + include_domains=["a.example"], + exclude_domains=["b.example"], + published_after="2026-01-01T00:00:00Z", + published_before="2026-06-01T00:00:00Z", + ) + assert arguments == { + "query": "hello", + "maxResults": 5, + "filters": { + "domainFilter": {"include": ["a.example"], "exclude": ["b.example"]}, + "publishedDateFilter": {"from": "2026-01-01T00:00:00Z", "to": "2026-06-01T00:00:00Z"}, + }, + } + + @pytest.mark.parametrize("query", ["", " "]) + def test_empty_query_rejected(self, query): + with pytest.raises(ValueError, match="non-empty"): + _build_arguments(query) + + def test_query_length_limit(self): + _build_arguments("x" * 200) + with pytest.raises(ValueError, match="200 characters or fewer"): + _build_arguments("x" * 201) + + @pytest.mark.parametrize("max_results", [0, 26, -1]) + def test_max_results_range(self, max_results): + with pytest.raises(ValueError, match="between 1 and 25"): + _build_arguments("hello", max_results=max_results) + + @pytest.mark.parametrize("max_results", [1, 25]) + def test_max_results_boundaries_allowed(self, max_results): + assert _build_arguments("hello", max_results=max_results)["maxResults"] == max_results + + @pytest.mark.parametrize("max_results", ["5", 5.0, True]) + def test_max_results_must_be_int(self, max_results): + with pytest.raises(ValueError, match="must be an integer"): + _build_arguments("hello", max_results=max_results) + + def test_empty_filter_lists_omitted(self): + assert "filters" not in _build_arguments("hello", include_domains=[], exclude_domains=[]) + + def test_only_one_date_bound(self): + arguments = _build_arguments("hello", published_after="2026-01-01T00:00:00Z") + assert arguments["filters"] == {"publishedDateFilter": {"from": "2026-01-01T00:00:00Z"}} + + +class TestResponseParsing: + """Tests for turning the tool payload into result objects.""" + + def test_from_payload(self): + response = WebSearchResponse.from_payload(SEARCH_PAYLOAD) + assert response.search_id == "search-1" + assert len(response) == 2 + first = response.results[0] + assert first.title == "urllib3 docs" + assert first.url == "https://urllib3.readthedocs.io/" + assert first.published_date == "2026-01-05T00:00:00Z" + + def test_optional_fields_default_to_none(self): + result = WebSearchResponse.from_payload(SEARCH_PAYLOAD).results[1] + assert result.text == "Only text is required." + assert result.url is None + assert result.title is None + assert result.published_date is None + + def test_empty_payload(self): + response = WebSearchResponse.from_payload({}) + assert len(response) == 0 + assert response.search_id is None + + def test_non_dict_entries_skipped(self): + response = WebSearchResponse.from_payload({"results": ["nope", {"text": "yes"}]}) + assert [r.text for r in response] == ["yes"] + + def test_missing_text_becomes_empty_string(self): + assert WebSearchResult.from_payload({"url": "https://example.com"}).text == "" + + def test_iterable(self): + assert [r.text for r in WebSearchResponse.from_payload(SEARCH_PAYLOAD)][0].startswith("urllib3") + + +class TestDecodeJsonRpc: + """Tests for both response framings the transport allows.""" + + def test_json_body(self): + message = _decode_jsonrpc("application/json", b'{"jsonrpc":"2.0","id":1,"result":{}}') + assert message["id"] == 1 + + def test_event_stream_body(self): + body = b'event: message\ndata: {"jsonrpc":"2.0","id":1,"result":{"ok":true}}\n\n' + message = _decode_jsonrpc("text/event-stream", body) + assert message["result"] == {"ok": True} + + def test_event_stream_skips_non_data_and_undecodable_lines(self): + body = b': ping\nid: 7\ndata: not json\ndata: {"jsonrpc":"2.0","id":1,"result":{}}\n' + assert _decode_jsonrpc("text/event-stream", body)["id"] == 1 + + def test_event_stream_without_message_returns_none(self): + assert _decode_jsonrpc("text/event-stream", b"event: ping\ndata: \n\n") is None + + def test_undecodable_json_raises(self): + with pytest.raises(WebSearchError, match="Could not decode gateway response"): + _decode_jsonrpc("application/json", b"gateway error") + + def test_non_object_json_returns_none(self): + assert _decode_jsonrpc("application/json", b"[1, 2]") is None + + +class TestGatewayMcpBackendHandshake: + """Tests for the MCP request sequence and its signed headers.""" + + def test_initialize_then_notify_then_call(self): + backend = _make_backend([_initialize_response(), _http_response(status=202, body=b""), _tools_call_response()]) + + payload = backend.search({"query": "hello"}) + + assert payload == SEARCH_PAYLOAD + methods = [json.loads(call.kwargs["body"])["method"] for call in backend._http.request.call_args_list] + assert methods == ["initialize", "notifications/initialized", "tools/call"] + + def test_session_reused_across_searches(self): + backend = _make_backend( + [ + _initialize_response(), + _http_response(status=202, body=b""), + _tools_call_response(), + _tools_call_response(), + ] + ) + + backend.search({"query": "one"}) + backend.search({"query": "two"}) + + methods = [json.loads(call.kwargs["body"])["method"] for call in backend._http.request.call_args_list] + assert methods == ["initialize", "notifications/initialized", "tools/call", "tools/call"] + + def test_notification_carries_no_id(self): + backend = _make_backend([_initialize_response(), _http_response(status=202, body=b""), _tools_call_response()]) + backend.search({"query": "hello"}) + + notification = json.loads(backend._http.request.call_args_list[1].kwargs["body"]) + assert "id" not in notification + + def test_session_id_and_protocol_version_sent_after_initialize(self): + backend = _make_backend([_initialize_response(), _http_response(status=202, body=b""), _tools_call_response()]) + backend.search({"query": "hello"}) + + initialize_headers = backend._http.request.call_args_list[0].kwargs["headers"] + assert "Mcp-Session-Id" not in initialize_headers + assert "MCP-Protocol-Version" not in initialize_headers + + call_headers = backend._http.request.call_args_list[2].kwargs["headers"] + assert call_headers["Mcp-Session-Id"] == "sess-1" + assert call_headers["MCP-Protocol-Version"] == "2025-06-18" + + def test_negotiated_protocol_version_is_echoed_back(self): + negotiated = _json_rpc_response({"protocolVersion": "2025-03-26"}, headers={"Mcp-Session-Id": "sess-1"}) + backend = _make_backend([negotiated, _http_response(status=202, body=b""), _tools_call_response()]) + + backend.search({"query": "hello"}) + + call_headers = backend._http.request.call_args_list[2].kwargs["headers"] + assert call_headers["MCP-Protocol-Version"] == "2025-03-26" + + def test_requests_are_sigv4_signed(self): + backend = _make_backend([_initialize_response(), _http_response(status=202, body=b""), _tools_call_response()]) + backend.search({"query": "hello"}) + + headers = backend._http.request.call_args_list[2].kwargs["headers"] + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIAEXAMPLE/") + assert "/us-east-1/bedrock-agentcore/aws4_request" in headers["Authorization"] + assert "X-Amz-Date" in headers + assert headers["Content-Type"] == "application/json" + assert headers["Accept"] == "application/json, text/event-stream" + assert headers["Content-Length"] == str(len(backend._http.request.call_args_list[2].kwargs["body"])) + + def test_connection_header_is_never_signed(self): + backend = _make_backend([_initialize_response(), _http_response(status=202, body=b""), _tools_call_response()]) + backend.search({"query": "hello"}) + + for call in backend._http.request.call_args_list: + signed = call.kwargs["headers"]["Authorization"].split("SignedHeaders=")[1].split(",")[0] + assert "connection" not in signed + + def test_user_agent_reports_the_sdk(self): + backend = _make_backend( + [_initialize_response(), _http_response(status=202, body=b""), _tools_call_response()], + integration_source="langchain", + ) + backend.search({"query": "hello"}) + + user_agent = backend._http.request.call_args_list[0].kwargs["headers"]["User-Agent"] + assert "bedrock-agentcore/" in user_agent + assert "integration_source=langchain" in user_agent + + def test_no_credentials_raises(self): + session = MagicMock() + session.get_credentials.return_value = None + backend = GatewayMcpBackend(endpoint=ENDPOINT, region="us-east-1", boto3_session=session, tool_name="t") + backend._http = MagicMock() + + with pytest.raises(WebSearchError, match="No AWS credentials"): + backend.search({"query": "hello"}) + + def test_keepalive_only_notification_reply_is_tolerated(self): + """A body carrying no JSON-RPC message is not an error, it is an ack.""" + keepalive = _http_response(status=202, body=b"event: ping\ndata: \n\n", content_type="text/event-stream") + backend = _make_backend([_initialize_response(), keepalive, _tools_call_response()]) + + assert backend.search({"query": "hello"}) == SEARCH_PAYLOAD + + def test_event_stream_tools_call_is_parsed(self): + sse_body = ( + b"event: message\ndata: " + + json.dumps( + { + "jsonrpc": "2.0", + "id": 2, + "result": {"content": [{"type": "text", "text": json.dumps(SEARCH_PAYLOAD)}]}, + } + ).encode() + + b"\n\n" + ) + backend = _make_backend( + [ + _initialize_response(), + _http_response(status=202, body=b""), + _http_response(body=sse_body, content_type="text/event-stream"), + ] + ) + + assert backend.search({"query": "hello"}) == SEARCH_PAYLOAD + + def test_close_resets_the_session(self): + backend = _make_backend([_initialize_response(), _http_response(status=202, body=b""), _tools_call_response()]) + backend.search({"query": "hello"}) + + backend.close() + + assert backend._initialized is False + assert backend._mcp_session_id is None + backend._http.clear.assert_called_once() + + +class TestGatewayMcpBackendErrors: + """Tests for the failure paths.""" + + def test_http_error_is_surfaced(self): + backend = _make_backend([_http_response(status=403, body=b"not authorized")]) + + with pytest.raises(WebSearchError, match="HTTP 403"): + backend.search({"query": "hello"}) + + def test_json_rpc_error_is_surfaced(self): + error_body = json.dumps( + {"jsonrpc": "2.0", "id": 1, "error": {"code": -32602, "message": "Unknown tool"}} + ).encode() + backend = _make_backend([_initialize_response(), _http_response(status=202), _http_response(body=error_body)]) + + with pytest.raises(WebSearchError, match="Unknown tool"): + backend.search({"query": "hello"}) + + def test_tool_error_flag_is_surfaced(self): + error_result = _json_rpc_response( + {"isError": True, "content": [{"type": "text", "text": "query too long"}]}, request_id=2 + ) + backend = _make_backend([_initialize_response(), _http_response(status=202), error_result]) + + with pytest.raises(WebSearchError, match="query too long"): + backend.search({"query": "hello"}) + + def test_missing_text_content_is_surfaced(self): + backend = _make_backend( + [_initialize_response(), _http_response(status=202), _json_rpc_response({"content": []}, request_id=2)] + ) + + with pytest.raises(WebSearchError, match="no text content"): + backend.search({"query": "hello"}) + + def test_undecodable_tool_payload_is_surfaced(self): + bad = _json_rpc_response({"content": [{"type": "text", "text": "not json"}]}, request_id=2) + backend = _make_backend([_initialize_response(), _http_response(status=202), bad]) + + with pytest.raises(WebSearchError, match="Could not decode web search response"): + backend.search({"query": "hello"}) + + def test_non_object_tool_payload_is_surfaced(self): + bad = _json_rpc_response({"content": [{"type": "text", "text": "[1,2]"}]}, request_id=2) + backend = _make_backend([_initialize_response(), _http_response(status=202), bad]) + + with pytest.raises(WebSearchError, match="Expected a JSON object"): + backend.search({"query": "hello"}) + + def test_empty_initialize_reply_is_surfaced(self): + backend = _make_backend([_http_response(status=202, body=b"")]) + + with pytest.raises(WebSearchError, match="did not answer the MCP initialize"): + backend.search({"query": "hello"}) + + def test_empty_tools_call_reply_is_surfaced(self): + backend = _make_backend( + [_initialize_response(), _http_response(status=202), _http_response(status=202, body=b"")] + ) + + with pytest.raises(WebSearchError, match="did not answer the web search tool call"): + backend.search({"query": "hello"}) + + +class TestToolNameResolution: + """Tests for finding the fully qualified tool name.""" + + def test_target_name_derives_the_prefixed_name(self): + backend = _make_backend( + [_initialize_response(), _http_response(status=202), _tools_call_response()], + tool_name=None, + target_name="amazon-web-search", + ) + backend.search({"query": "hello"}) + + params = json.loads(backend._http.request.call_args_list[2].kwargs["body"])["params"] + assert params["name"] == "amazon-web-search___WebSearch" + methods = [json.loads(c.kwargs["body"])["method"] for c in backend._http.request.call_args_list] + assert "tools/list" not in methods + + def test_explicit_tool_name_skips_discovery(self): + backend = _make_backend( + [_initialize_response(), _http_response(status=202), _tools_call_response()], + tool_name="custom___WebSearch", + ) + backend.search({"query": "hello"}) + + params = json.loads(backend._http.request.call_args_list[2].kwargs["body"])["params"] + assert params["name"] == "custom___WebSearch" + + def test_discovery_picks_the_prefixed_tool(self): + tools_list = _json_rpc_response( + {"tools": [{"name": "other___Lookup"}, {"name": "amazon-web-search___WebSearch"}]} + ) + backend = _make_backend( + [_initialize_response(), _http_response(status=202), tools_list, _tools_call_response()], + tool_name=None, + ) + backend.search({"query": "hello"}) + + params = json.loads(backend._http.request.call_args_list[3].kwargs["body"])["params"] + assert params["name"] == "amazon-web-search___WebSearch" + + def test_discovery_accepts_an_unprefixed_tool(self): + tools_list = _json_rpc_response({"tools": [{"name": "WebSearch"}]}) + backend = _make_backend( + [_initialize_response(), _http_response(status=202), tools_list, _tools_call_response()], + tool_name=None, + ) + backend.search({"query": "hello"}) + + params = json.loads(backend._http.request.call_args_list[3].kwargs["body"])["params"] + assert params["name"] == "WebSearch" + + def test_discovery_follows_pagination(self): + page_one = _json_rpc_response({"tools": [{"name": "other___Lookup"}], "nextCursor": "c1"}) + page_two = _json_rpc_response({"tools": [{"name": "amazon-web-search___WebSearch"}]}) + backend = _make_backend( + [_initialize_response(), _http_response(status=202), page_one, page_two, _tools_call_response()], + tool_name=None, + ) + backend.search({"query": "hello"}) + + second_page = json.loads(backend._http.request.call_args_list[3].kwargs["body"]) + assert second_page["params"] == {"cursor": "c1"} + assert json.loads(backend._http.request.call_args_list[4].kwargs["body"])["params"]["name"] == ( + "amazon-web-search___WebSearch" + ) + + def test_discovery_resolves_once_and_is_cached(self): + tools_list = _json_rpc_response({"tools": [{"name": "amazon-web-search___WebSearch"}]}) + backend = _make_backend( + [ + _initialize_response(), + _http_response(status=202), + tools_list, + _tools_call_response(), + _tools_call_response(), + ], + tool_name=None, + ) + + backend.search({"query": "one"}) + backend.search({"query": "two"}) + + methods = [json.loads(c.kwargs["body"])["method"] for c in backend._http.request.call_args_list] + assert methods.count("tools/list") == 1 + + def test_no_web_search_tool_raises(self): + tools_list = _json_rpc_response({"tools": [{"name": "other___Lookup"}]}) + backend = _make_backend([_initialize_response(), _http_response(status=202), tools_list], tool_name=None) + + with pytest.raises(WebSearchError, match="No WebSearch tool found"): + backend.search({"query": "hello"}) + + def test_ambiguous_web_search_tools_raise(self): + tools_list = _json_rpc_response({"tools": [{"name": "a___WebSearch"}, {"name": "b___WebSearch"}]}) + backend = _make_backend([_initialize_response(), _http_response(status=202), tools_list], tool_name=None) + + with pytest.raises(WebSearchError, match="more than one WebSearch tool"): + backend.search({"query": "hello"}) + + +class TestParseGatewayArn: + """Tests for reading a gateway ID and region out of an ARN.""" + + def test_valid_arn(self): + gateway_id, region = _parse_gateway_arn("arn:aws:bedrock-agentcore:eu-west-1:123456789012:gateway/gw-abc123") + assert (gateway_id, region) == ("gw-abc123", "eu-west-1") + + @pytest.mark.parametrize( + "arn", + [ + "gw-abc123", + "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/r-1", + "not:an:arn:at:all:gateway/gw-1", + ], + ) + def test_invalid_arn(self, arn): + with pytest.raises(ValueError, match="gateway ARN"): + _parse_gateway_arn(arn) + + def test_arn_without_id(self): + with pytest.raises(ValueError, match="no gateway ID"): + _parse_gateway_arn("arn:aws:bedrock-agentcore:us-east-1:123456789012:gateway/") + + +class _RecordingBackend(WebSearchBackend): + """A backend that records the arguments it was asked to search with.""" + + def __init__(self, payload=None): + self.payload = payload if payload is not None else SEARCH_PAYLOAD + self.arguments = None + self.closed = False + + def search(self, arguments): + self.arguments = arguments + return self.payload + + def close(self): + self.closed = True + + +class TestWebSearchClient: + """Tests for the client surface.""" + + def test_search_returns_results(self): + backend = _RecordingBackend() + client = WebSearchClient(region="us-east-1", backend=backend) + + response = client.search("who maintains urllib3", max_results=2) + + assert backend.arguments == {"query": "who maintains urllib3", "maxResults": 2} + assert len(response) == 2 + assert response.results[0].title == "urllib3 docs" + + def test_search_passes_filters_through(self): + backend = _RecordingBackend() + client = WebSearchClient(region="us-east-1", backend=backend) + + client.search( + "agentcore", + include_domains=["docs.aws.amazon.com"], + exclude_domains=["spam.example"], + published_after="2026-01-01T00:00:00Z", + ) + + assert backend.arguments["filters"] == { + "domainFilter": {"include": ["docs.aws.amazon.com"], "exclude": ["spam.example"]}, + "publishedDateFilter": {"from": "2026-01-01T00:00:00Z"}, + } + + def test_validation_happens_before_the_call(self): + backend = _RecordingBackend() + client = WebSearchClient(region="us-east-1", backend=backend) + + with pytest.raises(ValueError): + client.search("x" * 201) + + assert backend.arguments is None + + def test_gateway_id_builds_the_endpoint(self): + client = WebSearchClient(region="us-east-1", gateway_id="gw-abc123", boto3_session=MagicMock()) + + assert client.backend._endpoint == ENDPOINT + assert client.region == "us-east-1" + + def test_gateway_arn_supplies_the_region(self): + client = WebSearchClient( + gateway_arn="arn:aws:bedrock-agentcore:eu-west-1:123456789012:gateway/gw-abc123", + boto3_session=MagicMock(), + ) + + assert client.region == "eu-west-1" + assert client.backend._endpoint.startswith("https://gw-abc123.gateway.bedrock-agentcore.eu-west-1.") + + def test_explicit_region_wins_over_the_arn(self): + client = WebSearchClient( + region="us-east-1", + gateway_arn="arn:aws:bedrock-agentcore:eu-west-1:123456789012:gateway/gw-abc123", + boto3_session=MagicMock(), + ) + + assert client.region == "us-east-1" + + def test_gateway_endpoint_used_as_given(self): + endpoint = "https://gw-abc123.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp" + client = WebSearchClient(region="us-east-1", gateway_endpoint=endpoint, boto3_session=MagicMock()) + + assert client.backend._endpoint == endpoint + + def test_region_from_the_session(self): + session = MagicMock() + session.region_name = "us-east-1" + client = WebSearchClient(gateway_id="gw-abc123", boto3_session=session) + + assert client.region == "us-east-1" + + def test_missing_region_raises(self): + session = MagicMock() + session.region_name = None + + with pytest.raises(ValueError, match="region could not be determined"): + WebSearchClient(gateway_id="gw-abc123", boto3_session=session) + + def test_unknown_region_warns_but_proceeds(self, caplog): + with caplog.at_level("WARNING"): + client = WebSearchClient(region="us-west-2", gateway_id="gw-abc123", boto3_session=MagicMock()) + + assert client.region == "us-west-2" + assert "us-east-1, eu-west-1, ap-northeast-1" in caplog.text + + def test_no_gateway_raises(self): + with pytest.raises(ValueError, match="is required"): + WebSearchClient(region="us-east-1", boto3_session=MagicMock()) + + def test_two_gateways_raise(self): + with pytest.raises(ValueError, match="only one of"): + WebSearchClient( + region="us-east-1", + gateway_id="gw-abc123", + gateway_endpoint=ENDPOINT, + boto3_session=MagicMock(), + ) + + def test_backend_and_gateway_together_raise(self): + with pytest.raises(ValueError, match="not both"): + WebSearchClient(region="us-east-1", gateway_id="gw-abc123", backend=_RecordingBackend()) + + def test_invalid_gateway_id_raises(self): + with pytest.raises(InvalidGatewayIdentifierError): + WebSearchClient(region="us-east-1", gateway_id="evil.example.com/", boto3_session=MagicMock()) + + def test_invalid_region_raises(self): + with pytest.raises(InvalidRegionError): + WebSearchClient(region="not a region", gateway_id="gw-abc123", boto3_session=MagicMock()) + + def test_close_only_closes_an_owned_backend(self): + backend = _RecordingBackend() + WebSearchClient(region="us-east-1", backend=backend).close() + + assert backend.closed is False + + def test_context_manager_closes_an_owned_backend(self): + with WebSearchClient(region="us-east-1", gateway_id="gw-abc123", boto3_session=MagicMock()) as client: + client.backend._http = MagicMock() + http = client.backend._http + + http.clear.assert_called_once() + + def test_target_name_is_handed_to_the_backend(self): + client = WebSearchClient( + region="us-east-1", + gateway_id="gw-abc123", + target_name="amazon-web-search", + boto3_session=MagicMock(), + ) + + assert client.backend._target_name == "amazon-web-search" + + +class TestBackendProtocol: + """Tests for the extension point.""" + + def test_base_search_is_not_implemented(self): + with pytest.raises(NotImplementedError): + WebSearchBackend().search({"query": "hello"}) + + def test_base_close_is_a_no_op(self): + assert WebSearchBackend().close() is None + + +def test_exported_from_the_tools_package(): + from bedrock_agentcore.tools import WebSearchClient as exported + + assert exported is WebSearchClient diff --git a/tests_integ/tools/test_web_search_client.py b/tests_integ/tools/test_web_search_client.py new file mode 100644 index 00000000..cb067928 --- /dev/null +++ b/tests_integ/tools/test_web_search_client.py @@ -0,0 +1,101 @@ +"""Integration tests for WebSearchClient. + +These tests call the real gateway, so they need a gateway that already has a web +search connector target on it. The web search connector is enabled per account, so +they skip rather than fail when the account is not entitled. + +Run with: + uv run pytest tests_integ/tools/test_web_search_client.py -xvs + +Requires environment variables: + WEB_SEARCH_GATEWAY_ID: ID of a gateway with a web search connector target + BEDROCK_TEST_REGION: AWS region (default: us-east-1). The connector is only + offered in us-east-1, eu-west-1 and ap-northeast-1. + WEB_SEARCH_TARGET_NAME: Optional. The target name, if it is not the SDK default. +""" + +import os + +import pytest + +from bedrock_agentcore.tools.web_search_client import WebSearchClient, WebSearchError + + +@pytest.mark.integration +class TestWebSearchClient: + """Integration tests for WebSearchClient over a gateway target.""" + + @classmethod + def setup_class(cls): + cls.gateway_id = os.environ.get("WEB_SEARCH_GATEWAY_ID") + if not cls.gateway_id: + pytest.skip("WEB_SEARCH_GATEWAY_ID must be set") + cls.region = os.environ.get("BEDROCK_TEST_REGION", "us-east-1") + cls.target_name = os.environ.get("WEB_SEARCH_TARGET_NAME") + + def _client(self): + return WebSearchClient( + region=self.region, + gateway_id=self.gateway_id, + target_name=self.target_name, + ) + + def _search(self, client, query, **kwargs): + """Search, skipping the test when the account is not entitled to the connector.""" + try: + return client.search(query, **kwargs) + except WebSearchError as e: + if "not available for this account" in str(e): + pytest.skip(f"web-search connector not enabled for this account: {e}") + raise + + def test_search_returns_results(self): + with self._client() as client: + response = self._search(client, "what is amazon bedrock agentcore", max_results=3) + + assert len(response) > 0 + assert len(response) <= 3 + first = response.results[0] + assert first.text + # Citations must be retained for any output shown to an end user, so the + # client has to surface the source URL. + assert first.url + + def test_search_respects_max_results(self): + with self._client() as client: + response = self._search(client, "python urllib3 release notes", max_results=1) + + assert len(response) == 1 + + def test_search_with_domain_filter(self): + """Needs connector version 1.2.0 or later on the target.""" + with self._client() as client: + response = self._search( + client, + "agentcore gateway connector targets", + max_results=5, + include_domains=["docs.aws.amazon.com"], + ) + + assert len(response) > 0 + for result in response: + assert result.url is None or "aws.amazon.com" in result.url + + def test_tool_name_discovery(self): + """Without target_name the client finds the tool through tools/list.""" + with WebSearchClient(region=self.region, gateway_id=self.gateway_id) as client: + self._search(client, "bedrock agentcore gateway", max_results=1) + + assert client.backend._tool_name.endswith("WebSearch") + + def test_session_is_reused_across_searches(self): + with self._client() as client: + self._search(client, "first query", max_results=1) + self._search(client, "second query", max_results=1) + + assert client.backend._mcp_session_id + + def test_oversized_query_is_rejected_locally(self): + with self._client() as client: + with pytest.raises(ValueError, match="200 characters or fewer"): + client.search("x" * 201) From aed1f972d94022a761b535f83738b53954cce9f4 Mon Sep 17 00:00:00 2001 From: Sundar Raghavan Date: Thu, 3 Sep 2026 12:50:22 -0700 Subject: [PATCH 2/3] docs(tools): document filter composition and the IAM actions a search needs Web search takes no API key of its own: the caller's credentials need bedrock-agentcore:InvokeGateway on the gateway ARN and the gateway service role needs bedrock-agentcore:InvokeWebSearch on the connector. AccessDenied is almost always the first of those, so say so where a caller will look. Also spell out how request filters compose with the target's own domain rules. A result is dropped if its domain is on either exclude list, and returned only if it is on every include list that is set, so passing include_domains against a target that already has an include list narrows to the intersection and disjoint lists return nothing. That empty result is silent rather than an error, which is worth knowing before debugging it. --- .../tools/web_search_client.py | 24 ++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/src/bedrock_agentcore/tools/web_search_client.py b/src/bedrock_agentcore/tools/web_search_client.py index 5c1a1c24..9d36c9f1 100644 --- a/src/bedrock_agentcore/tools/web_search_client.py +++ b/src/bedrock_agentcore/tools/web_search_client.py @@ -470,6 +470,13 @@ def _decode_jsonrpc(content_type: str, data: bytes) -> Optional[Dict[str, Any]]: class WebSearchClient: """Client for AgentCore Web Search. + Calls are authenticated with SigV4 using ordinary AWS credentials. There is no + web search API key: the credentials this client resolves need + ``bedrock-agentcore:InvokeGateway`` on the gateway ARN, and the gateway's own + service role needs ``bedrock-agentcore:InvokeWebSearch`` on the connector. + An ``AccessDeniedException`` from a search is almost always the first of those + two missing. + Attributes: region (str): The region being used. backend (WebSearchBackend): The transport in use. @@ -590,15 +597,26 @@ def search( The filter arguments need connector version 1.2.0 or later on the target. On an earlier version the tool accepts only ``query`` and ``max_results``. - Target level domain rules always apply on top and cannot be relaxed here. + + Request filters compose with the target's own domain rules and can never + widen them. A domain is dropped if it appears on either exclude list. A + domain is returned only if it appears on every include list that is set, + so when the target already has an include list, passing ``include_domains`` + narrows to the intersection of the two. If the two share no domains the + search returns nothing, which is a silent empty result rather than an + error, so check the target's configuration when a filtered search comes + back empty. Args: query: What to search for. 200 characters or fewer. max_results: How many results to return, 1 to 25. Service default is 10. - include_domains: Restrict results to these domains. - exclude_domains: Drop results from these domains. + include_domains: Restrict results to these domains. Up to 100. A root + domain matches its subdomains. + exclude_domains: Drop results from these domains. Up to 100. published_after: Earliest publication date, ISO-8601 UTC, inclusive. + Applies to web results only. published_before: Latest publication date, ISO-8601 UTC, inclusive. + Applies to web results only. Returns: The search results. From 0ed088ea59175313ac07270ec7112961a1ade199 Mon Sep 17 00:00:00 2001 From: Sundar Raghavan Date: Sun, 6 Sep 2026 15:25:52 -0700 Subject: [PATCH 3/3] fix(tools): address review on the web search client Match every JSON-RPC reply to the id of the request it answers, so a server notification arriving ahead of the reply is not read as the answer, and require a reply to carry a result or an error. A JSON-RPC error is still surfaced whatever id it carries, since the spec allows a null id there. Parse SSE the way the specification defines it: consecutive data lines of one event join with a newline and a blank line ends the event, so a message split across lines decodes instead of being dropped. Wrap urllib3 transport failures in WebSearchError, so a connection drop or a read timeout cannot escape as a urllib3 exception from a method documented to raise WebSearchError. Treat HTTP 404 against a session we hold as the spec's signal that the session is gone, and drop it so the next call hands shake again. Roll the initialized flag back if the notifications/initialized POST fails, rather than leaving the client claiming a session the gateway never acknowledged. Apply the SDK's endpoint host check to a caller-supplied gateway_endpoint, since signed requests carry the caller's credentials either way. Reject a domain list longer than the documented maximum of 100 before calling, create the default boto3 session only when a region has to come from one, export GatewayMcpBackend, type _parse_gateway_arn as Tuple[str, str] and drop DEFAULT_TARGET_NAME, which duplicated a default that belongs to the gateway helper. Tests: replies are now numbered from the request they answer, so the fixtures no longer hard-code ids that did not match. Adds coverage for each fix above. Drops the signing test that only asserted botocore's own behavior. --- src/bedrock_agentcore/tools/__init__.py | 2 + .../tools/web_search_client.py | 196 +++++++++++---- .../tools/test_web_search_client.py | 228 +++++++++++++++--- tests_integ/tools/test_web_search_client.py | 29 ++- 4 files changed, 370 insertions(+), 85 deletions(-) diff --git a/src/bedrock_agentcore/tools/__init__.py b/src/bedrock_agentcore/tools/__init__.py index f33bf3f9..a548b52b 100644 --- a/src/bedrock_agentcore/tools/__init__.py +++ b/src/bedrock_agentcore/tools/__init__.py @@ -27,6 +27,7 @@ create_browser_config, ) from .web_search_client import ( + GatewayMcpBackend, WebSearchBackend, WebSearchClient, WebSearchError, @@ -50,6 +51,7 @@ "EnterprisePolicyS3Location", "ExtensionS3Location", "ExternalProxy", + "GatewayMcpBackend", "NetworkConfiguration", "ProfileConfiguration", "ProxyConfiguration", diff --git a/src/bedrock_agentcore/tools/web_search_client.py b/src/bedrock_agentcore/tools/web_search_client.py index 9d36c9f1..23b69f72 100644 --- a/src/bedrock_agentcore/tools/web_search_client.py +++ b/src/bedrock_agentcore/tools/web_search_client.py @@ -20,13 +20,13 @@ import logging import threading from dataclasses import dataclass, field -from typing import Any, Dict, Iterator, List, Optional, Sequence +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple import urllib3 from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest -from bedrock_agentcore._utils.endpoints import get_gateway_mcp_endpoint +from bedrock_agentcore._utils.endpoints import _validate_endpoint_url, get_gateway_mcp_endpoint from bedrock_agentcore._utils.user_agent import SDK_VERSION, build_user_agent_suffix logger = logging.getLogger(__name__) @@ -37,9 +37,6 @@ #: Gateway prefixes every tool it exposes with the name of the target it came from. GATEWAY_TOOL_NAME_DELIMITER = "___" -#: Default target name used by ``GatewayClient.create_web_search_target``. -DEFAULT_TARGET_NAME = "amazon-web-search" - #: Service name the gateway data plane signs as. GATEWAY_SIGNING_SERVICE = "bedrock-agentcore" @@ -47,6 +44,7 @@ MAX_QUERY_LENGTH = 200 MIN_MAX_RESULTS = 1 MAX_MAX_RESULTS = 25 +MAX_DOMAIN_FILTER_ENTRIES = 100 #: Regions where the web search connector is offered. Used for a warning only, never #: to block a call, so that a newly added region does not require an SDK release. @@ -128,8 +126,9 @@ def _build_arguments( """Validate search inputs and shape them into the tool's argument object. Raises: - ValueError: If the query is empty or over the documented length limit, or - if max_results falls outside the documented range. + ValueError: If the query is empty or over the documented length limit, if + max_results falls outside the documented range, or if either domain + list is longer than the documented maximum. """ if not query or not query.strip(): raise ValueError("query must be a non-empty string") @@ -148,9 +147,9 @@ def _build_arguments( filters: Dict[str, Any] = {} domain_filter: Dict[str, List[str]] = {} if include_domains: - domain_filter["include"] = list(include_domains) + domain_filter["include"] = _checked_domains("include_domains", include_domains) if exclude_domains: - domain_filter["exclude"] = list(exclude_domains) + domain_filter["exclude"] = _checked_domains("exclude_domains", exclude_domains) if domain_filter: filters["domainFilter"] = domain_filter @@ -168,6 +167,14 @@ def _build_arguments( return arguments +def _checked_domains(argument: str, domains: Sequence[str]) -> List[str]: + """Return the domain list, rejecting one longer than the documented maximum.""" + values = list(domains) + if len(values) > MAX_DOMAIN_FILTER_ENTRIES: + raise ValueError(f"{argument} accepts at most {MAX_DOMAIN_FILTER_ENTRIES} domains, got {len(values)}") + return values + + def _extract_search_payload(result: Dict[str, Any]) -> Dict[str, Any]: """Pull the search payload out of an MCP ``tools/call`` result. @@ -302,8 +309,15 @@ def _signed_headers(self, body: bytes, extra: Dict[str, str]) -> Dict[str, str]: return dict(request.headers) def _post(self, message: Dict[str, Any]) -> Optional[Dict[str, Any]]: - """Send one JSON-RPC message and return the decoded reply, if there is one.""" + """Send one JSON-RPC message and return the reply to it, if there is one. + + Raises: + WebSearchError: If the transport fails, the gateway answers with an HTTP + error or a JSON-RPC error, or the reply carries neither a result nor + an error. + """ body = json.dumps(message).encode("utf-8") + expected_id = message.get("id") extra: Dict[str, str] = {} if self._mcp_session_id: @@ -311,14 +325,20 @@ def _post(self, message: Dict[str, Any]) -> Optional[Dict[str, Any]]: if self._initialized: extra["MCP-Protocol-Version"] = self._protocol_version - response = self._http.request( - "POST", - self._endpoint, - body=body, - headers=self._signed_headers(body, extra), - timeout=urllib3.Timeout(total=self._timeout), - preload_content=True, - ) + try: + response = self._http.request( + "POST", + self._endpoint, + body=body, + headers=self._signed_headers(body, extra), + timeout=urllib3.Timeout(total=self._timeout), + preload_content=True, + ) + except urllib3.exceptions.HTTPError as exc: + # Connection refused, DNS failure, TLS errors and read timeouts all land + # here. search() promises WebSearchError, so they cannot escape as urllib3 + # exceptions. + raise WebSearchError(f"Web search request to {self._endpoint} failed: {exc}") from exc session_id = response.headers.get("Mcp-Session-Id") if session_id: @@ -326,19 +346,31 @@ def _post(self, message: Dict[str, Any]) -> Optional[Dict[str, Any]]: if response.status >= 400: body_text = response.data.decode("utf-8", "replace")[:500] + # 404 against a session we hold is the spec's signal that the session is + # gone. Forget it, so the next call re-initializes rather than failing + # this way until someone calls close(). + if response.status == 404 and self._mcp_session_id: + self._reset_session() raise WebSearchError(f"Web search request failed with HTTP {response.status}: {body_text}") if not response.data: return None - reply = _decode_jsonrpc(response.headers.get("Content-Type", ""), response.data) + reply = _decode_jsonrpc(response.headers.get("Content-Type", ""), response.data, expected_id) if reply is None: return None if "error" in reply: error = reply["error"] or {} raise WebSearchError(f"Gateway returned a JSON-RPC error {error.get('code')}: {error.get('message')}") + if expected_id is not None and "result" not in reply: + raise WebSearchError(f"Gateway reply carried neither a result nor an error: {reply}") return reply + def _reset_session(self) -> None: + """Forget the MCP session, so the next call performs the handshake again.""" + self._initialized = False + self._mcp_session_id = None + # MCP session # ------------------------------------------------------------------------- def _ensure_initialized(self) -> None: @@ -364,8 +396,16 @@ def _ensure_initialized(self) -> None: if isinstance(negotiated, str) and negotiated: self._protocol_version = negotiated + # Set before the notification is sent, because from here on every request + # carries MCP-Protocol-Version. If the notification fails, the handshake did + # not complete, so the flag is rolled back rather than left claiming a session + # the gateway never acknowledged. self._initialized = True - self._post({"jsonrpc": "2.0", "method": "notifications/initialized"}) + try: + self._post({"jsonrpc": "2.0", "method": "notifications/initialized"}) + except Exception: + self._reset_session() + raise def _ensure_tool_name(self) -> str: """Resolve the fully qualified tool name, discovering it if necessary.""" @@ -433,30 +473,22 @@ def search(self, arguments: Dict[str, Any]) -> Dict[str, Any]: def close(self) -> None: """Close the connection pool.""" self._http.clear() - self._initialized = False - self._mcp_session_id = None + self._reset_session() -def _decode_jsonrpc(content_type: str, data: bytes) -> Optional[Dict[str, Any]]: - """Decode a JSON-RPC reply from either a JSON body or an SSE stream. +def _decode_jsonrpc(content_type: str, data: bytes, expected_id: Optional[int] = None) -> Optional[Dict[str, Any]]: + """Decode the JSON-RPC reply to ``expected_id`` from a JSON body or an SSE stream. - Returns None when the body carries no JSON-RPC message, which is what a - notification acknowledgement looks like. + Returns None when the body carries no reply to that id, which is what a + notification acknowledgement looks like. A message addressed to another id, such + as a server notification arriving ahead of the reply, is skipped rather than + mistaken for the answer. """ text = data.decode("utf-8", "replace") if "text/event-stream" in content_type.lower(): - for line in text.splitlines(): - if not line.startswith("data:"): - continue - chunk = line[len("data:") :].strip() - if not chunk: - continue - try: - message = json.loads(chunk) - except json.JSONDecodeError: - continue - if isinstance(message, dict) and "jsonrpc" in message: + for message in _iter_sse_messages(text): + if _answers(message, expected_id): return message return None @@ -464,7 +496,59 @@ def _decode_jsonrpc(content_type: str, data: bytes) -> Optional[Dict[str, Any]]: message = json.loads(text) except json.JSONDecodeError as exc: raise WebSearchError(f"Could not decode gateway response as JSON: {text[:200]!r}") from exc - return message if isinstance(message, dict) else None + if not isinstance(message, dict): + return None + return message if _answers(message, expected_id) else None + + +def _answers(message: Any, expected_id: Optional[int]) -> bool: + """Whether a decoded JSON-RPC message is the reply to ``expected_id``.""" + if not isinstance(message, dict) or "jsonrpc" not in message: + return False + if "error" in message: + # A JSON-RPC error is allowed to carry a null id, so it is always surfaced + # rather than filtered out for not matching. + return True + if expected_id is None: + return True + return message.get("id") == expected_id + + +def _iter_sse_messages(text: str) -> Iterator[Dict[str, Any]]: + """Yield the JSON objects carried by the ``data`` field of each SSE event. + + Consecutive ``data:`` lines belonging to one event are joined with a newline, as + the SSE specification requires, so a message split across lines is decoded rather + than dropped. A blank line ends an event. + """ + buffer: List[str] = [] + + def flush() -> Optional[Dict[str, Any]]: + if not buffer: + return None + joined = "\n".join(buffer) + buffer.clear() + try: + message = json.loads(joined) + except json.JSONDecodeError: + return None + return message if isinstance(message, dict) else None + + for line in text.splitlines(): + if line.startswith("data:"): + value = line[len("data:") :] + buffer.append(value[1:] if value.startswith(" ") else value) + continue + if line.strip(): + # Any other SSE field (event:, id:, retry:) or a comment line. + continue + message = flush() + if message is not None: + yield message + + message = flush() + if message is not None: + yield message class WebSearchClient: @@ -518,7 +602,9 @@ def __init__( region: Region to call. Defaults to the session's region. gateway_id: ID of a gateway with a web search connector target. gateway_arn: ARN of that gateway. The ID and region are read from it. - gateway_endpoint: A gateway MCP endpoint URL, if you already have one. + gateway_endpoint: A gateway MCP endpoint URL, if you already have one. The + host has to be an AWS one, the same check applied to endpoints this SDK + builds itself. target_name: Name of the connector target. Supplying it avoids a ``tools/list`` round trip on the first search. tool_name: Fully qualified tool name, if you already know it. @@ -529,16 +615,19 @@ def __init__( Raises: ValueError: If no gateway is identified, or more than one is. + InvalidGatewayIdentifierError: If ``gateway_id`` is not a gateway ID. + InvalidRegionError: If the region is malformed, or a supplied + ``gateway_endpoint`` does not point at an AWS host. """ - import boto3 - - self._session = boto3_session or boto3.Session() + self._session = boto3_session self._owns_backend = backend is None if backend is not None: if any(value is not None for value in (gateway_id, gateway_arn, gateway_endpoint)): raise ValueError("Pass either backend or one of gateway_id/gateway_arn/gateway_endpoint, not both") - self.region = region or self._session.region_name + # Only resolve a session if the region has to come from one, so supplying a + # backend does not pay for the credential provider chain it will not use. + self.region = region or self._resolve_session().region_name self.backend: WebSearchBackend = backend return @@ -558,7 +647,7 @@ def __init__( gateway_id, arn_region = _parse_gateway_arn(gateway_arn) region = region or arn_region - self.region = region or self._session.region_name + self.region = region or self._resolve_session().region_name if not self.region: raise ValueError("region could not be determined. Pass region= or configure a default region.") if self.region not in KNOWN_REGIONS: @@ -570,19 +659,31 @@ def __init__( if gateway_id: gateway_endpoint = get_gateway_mcp_endpoint(gateway_id, self.region) + elif gateway_endpoint: + # Signed requests carry the caller's credentials, including a session token, + # so a supplied endpoint gets the same host check as one this SDK builds. + gateway_endpoint = _validate_endpoint_url(gateway_endpoint) if not gateway_endpoint: raise ValueError("One of gateway_id, gateway_arn, gateway_endpoint or backend is required") self.backend = GatewayMcpBackend( endpoint=gateway_endpoint, region=self.region, - boto3_session=self._session, + boto3_session=self._resolve_session(), tool_name=tool_name, target_name=target_name, timeout=timeout, integration_source=integration_source, ) + def _resolve_session(self) -> Any: + """Return the boto3 session, creating a default one only when first needed.""" + if self._session is None: + import boto3 + + self._session = boto3.Session() + return self._session + def search( self, query: str, @@ -622,7 +723,8 @@ def search( The search results. Raises: - ValueError: If the query or max_results is outside the documented limits. + ValueError: If the query, max_results or either domain list is outside the + documented limits. WebSearchError: If the call fails or the response cannot be decoded. """ arguments = _build_arguments( @@ -649,7 +751,7 @@ def __exit__(self, *exc_info: Any) -> None: self.close() -def _parse_gateway_arn(arn: str) -> tuple: +def _parse_gateway_arn(arn: str) -> Tuple[str, str]: """Pull the gateway ID and region out of a gateway ARN. Raises: diff --git a/tests/bedrock_agentcore/tools/test_web_search_client.py b/tests/bedrock_agentcore/tools/test_web_search_client.py index d23abe60..f76953d3 100644 --- a/tests/bedrock_agentcore/tools/test_web_search_client.py +++ b/tests/bedrock_agentcore/tools/test_web_search_client.py @@ -1,9 +1,10 @@ """Tests for WebSearchClient.""" import json -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest +import urllib3 from bedrock_agentcore._utils.endpoints import InvalidGatewayIdentifierError, InvalidRegionError from bedrock_agentcore.tools.web_search_client import ( @@ -34,19 +35,29 @@ } -def _http_response(status=200, body=b"", headers=None, content_type="application/json"): - response = MagicMock() - response.status = status - response.data = body - merged = {"Content-Type": content_type} - merged.update(headers or {}) - response.headers = merged - return response +class _FakeResponse: + """A urllib3-shaped response, with a flag the fake transport uses to number ids.""" + def __init__(self, status=200, body=b"", headers=None, content_type="application/json", echo_id=False): + self.status = status + self.data = body + self.headers = {"Content-Type": content_type, **(headers or {})} + self.echo_id = echo_id -def _json_rpc_response(result, request_id=1, **kwargs): + +def _http_response(status=200, body=b"", headers=None, content_type="application/json", echo_id=False): + return _FakeResponse(status=status, body=body, headers=headers, content_type=content_type, echo_id=echo_id) + + +def _json_rpc_response(result, request_id=None, **kwargs): + """A JSON-RPC reply. + + By default the id is filled in from the request it answers, which is what a real + gateway does. Pass ``request_id`` to pin it, which is how the id-matching tests + build a reply addressed to something else. + """ body = json.dumps({"jsonrpc": "2.0", "id": request_id, "result": result}).encode() - return _http_response(body=body, **kwargs) + return _http_response(body=body, echo_id=request_id is None, **kwargs) def _initialize_response(): @@ -59,22 +70,42 @@ def _initialize_response(): def _tools_call_response(payload=None): return _json_rpc_response( {"content": [{"type": "text", "text": json.dumps(payload if payload is not None else SEARCH_PAYLOAD)}]}, - request_id=2, ) def _make_backend(responses, **kwargs): - """Build a backend whose HTTP layer replays the given responses in order.""" + """Build a backend whose HTTP layer replays the given responses in order. + + A reply built without an explicit ``request_id`` is numbered with the id of the + request it answers, the way a gateway does, so the fixtures do not have to track + the client's request counter. A reply with a pinned id is replayed untouched. + """ session = MagicMock() session.get_credentials.return_value.get_frozen_credentials.return_value = _frozen_credentials() kwargs.setdefault("tool_name", "amazon-web-search___WebSearch") backend = GatewayMcpBackend(endpoint=ENDPOINT, region="us-east-1", boto3_session=session, **kwargs) backend._http = MagicMock() - backend._http.request.side_effect = list(responses) + backend._http.request.side_effect = _replay(responses) return backend +def _replay(responses): + """Return a side effect that serves the responses in order, numbering ids.""" + queue = list(responses) + + def request(*args, **kwargs): + response = queue.pop(0) + if getattr(response, "echo_id", False): + sent = json.loads(kwargs["body"]) + message = json.loads(response.data) + message["id"] = sent.get("id") + response.data = json.dumps(message).encode() + return response + + return request + + def _frozen_credentials(): from botocore.credentials import ReadOnlyCredentials @@ -136,6 +167,19 @@ def test_only_one_date_bound(self): arguments = _build_arguments("hello", published_after="2026-01-01T00:00:00Z") assert arguments["filters"] == {"publishedDateFilter": {"from": "2026-01-01T00:00:00Z"}} + @pytest.mark.parametrize("argument", ["include_domains", "exclude_domains"]) + def test_domain_list_maximum(self, argument): + domains = [f"d{index}.example" for index in range(100)] + assert _build_arguments("hello", **{argument: domains})["filters"]["domainFilter"] + + with pytest.raises(ValueError, match="at most 100 domains, got 101"): + _build_arguments("hello", **{argument: domains + ["one.too.many"]}) + + def test_a_domain_tuple_is_accepted(self): + """Any sequence works, and the shaped argument is always a list for JSON.""" + arguments = _build_arguments("hello", include_domains=("a.example", "b.example")) + assert arguments["filters"]["domainFilter"]["include"] == ["a.example", "b.example"] + class TestResponseParsing: """Tests for turning the tool payload into result objects.""" @@ -184,13 +228,34 @@ def test_event_stream_body(self): message = _decode_jsonrpc("text/event-stream", body) assert message["result"] == {"ok": True} - def test_event_stream_skips_non_data_and_undecodable_lines(self): - body = b': ping\nid: 7\ndata: not json\ndata: {"jsonrpc":"2.0","id":1,"result":{}}\n' + def test_event_stream_skips_comments_and_other_fields(self): + body = b': ping\nid: 7\nretry: 500\ndata: {"jsonrpc":"2.0","id":1,"result":{}}\n\n' + assert _decode_jsonrpc("text/event-stream", body)["id"] == 1 + + def test_event_stream_joins_data_lines_of_one_event(self): + """The spec joins consecutive data lines with a newline, so a split message decodes.""" + body = b'event: message\ndata: {"jsonrpc":"2.0","id":1,\ndata: "result":{"ok":true}}\n\n' + assert _decode_jsonrpc("text/event-stream", body)["result"] == {"ok": True} + + def test_event_stream_skips_an_undecodable_event(self): + body = b'data: not json\n\ndata: {"jsonrpc":"2.0","id":1,"result":{}}\n\n' + assert _decode_jsonrpc("text/event-stream", body)["id"] == 1 + + def test_event_stream_reads_a_final_event_without_a_trailing_blank_line(self): + body = b'data: {"jsonrpc":"2.0","id":1,"result":{}}' assert _decode_jsonrpc("text/event-stream", body)["id"] == 1 def test_event_stream_without_message_returns_none(self): assert _decode_jsonrpc("text/event-stream", b"event: ping\ndata: \n\n") is None + def test_event_stream_skips_a_non_json_rpc_object(self): + """An event carrying JSON that is not a JSON-RPC message is not an answer.""" + body = b'data: {"type":"ping"}\n\ndata: {"jsonrpc":"2.0","id":1,"result":{}}\n\n' + assert _decode_jsonrpc("text/event-stream", body)["id"] == 1 + + def test_non_json_rpc_json_body_returns_none(self): + assert _decode_jsonrpc("application/json", b'{"type":"ping"}') is None + def test_undecodable_json_raises(self): with pytest.raises(WebSearchError, match="Could not decode gateway response"): _decode_jsonrpc("application/json", b"gateway error") @@ -198,6 +263,22 @@ def test_undecodable_json_raises(self): def test_non_object_json_returns_none(self): assert _decode_jsonrpc("application/json", b"[1, 2]") is None + def test_a_reply_to_another_id_is_not_the_answer(self): + body = b'{"jsonrpc":"2.0","id":9,"result":{}}' + assert _decode_jsonrpc("application/json", body, expected_id=2) is None + + def test_an_error_is_returned_whatever_id_it_carries(self): + """JSON-RPC allows an error to carry a null id, so it is never filtered out.""" + body = b'{"jsonrpc":"2.0","id":null,"error":{"code":-32600,"message":"bad request"}}' + assert _decode_jsonrpc("application/json", body, expected_id=2)["error"]["code"] == -32600 + + def test_event_stream_skips_a_notification_ahead_of_the_reply(self): + body = ( + b'data: {"jsonrpc":"2.0","method":"notifications/message","params":{}}\n\n' + b'data: {"jsonrpc":"2.0","id":2,"result":{"ok":true}}\n\n' + ) + assert _decode_jsonrpc("text/event-stream", body, expected_id=2)["result"] == {"ok": True} + class TestGatewayMcpBackendHandshake: """Tests for the MCP request sequence and its signed headers.""" @@ -267,14 +348,6 @@ def test_requests_are_sigv4_signed(self): assert headers["Accept"] == "application/json, text/event-stream" assert headers["Content-Length"] == str(len(backend._http.request.call_args_list[2].kwargs["body"])) - def test_connection_header_is_never_signed(self): - backend = _make_backend([_initialize_response(), _http_response(status=202, body=b""), _tools_call_response()]) - backend.search({"query": "hello"}) - - for call in backend._http.request.call_args_list: - signed = call.kwargs["headers"]["Authorization"].split("SignedHeaders=")[1].split(",")[0] - assert "connection" not in signed - def test_user_agent_reports_the_sdk(self): backend = _make_backend( [_initialize_response(), _http_response(status=202, body=b""), _tools_call_response()], @@ -354,9 +427,7 @@ def test_json_rpc_error_is_surfaced(self): backend.search({"query": "hello"}) def test_tool_error_flag_is_surfaced(self): - error_result = _json_rpc_response( - {"isError": True, "content": [{"type": "text", "text": "query too long"}]}, request_id=2 - ) + error_result = _json_rpc_response({"isError": True, "content": [{"type": "text", "text": "query too long"}]}) backend = _make_backend([_initialize_response(), _http_response(status=202), error_result]) with pytest.raises(WebSearchError, match="query too long"): @@ -364,21 +435,21 @@ def test_tool_error_flag_is_surfaced(self): def test_missing_text_content_is_surfaced(self): backend = _make_backend( - [_initialize_response(), _http_response(status=202), _json_rpc_response({"content": []}, request_id=2)] + [_initialize_response(), _http_response(status=202), _json_rpc_response({"content": []})] ) with pytest.raises(WebSearchError, match="no text content"): backend.search({"query": "hello"}) def test_undecodable_tool_payload_is_surfaced(self): - bad = _json_rpc_response({"content": [{"type": "text", "text": "not json"}]}, request_id=2) + bad = _json_rpc_response({"content": [{"type": "text", "text": "not json"}]}) backend = _make_backend([_initialize_response(), _http_response(status=202), bad]) with pytest.raises(WebSearchError, match="Could not decode web search response"): backend.search({"query": "hello"}) def test_non_object_tool_payload_is_surfaced(self): - bad = _json_rpc_response({"content": [{"type": "text", "text": "[1,2]"}]}, request_id=2) + bad = _json_rpc_response({"content": [{"type": "text", "text": "[1,2]"}]}) backend = _make_backend([_initialize_response(), _http_response(status=202), bad]) with pytest.raises(WebSearchError, match="Expected a JSON object"): @@ -398,6 +469,72 @@ def test_empty_tools_call_reply_is_surfaced(self): with pytest.raises(WebSearchError, match="did not answer the web search tool call"): backend.search({"query": "hello"}) + def test_reply_without_result_or_error_is_surfaced(self): + """A reply to the outgoing id has to carry one of the two, or it means nothing.""" + empty = _http_response(body=b'{"jsonrpc":"2.0","id":2}') + backend = _make_backend([_initialize_response(), _http_response(status=202), empty]) + + with pytest.raises(WebSearchError, match="neither a result nor an error"): + backend.search({"query": "hello"}) + + def test_transport_failure_becomes_a_web_search_error(self): + """A connection drop must not escape as a urllib3 exception.""" + backend = _make_backend([]) + backend._http.request.side_effect = urllib3.exceptions.ProtocolError("connection aborted") + + with pytest.raises(WebSearchError, match="failed: connection aborted"): + backend.search({"query": "hello"}) + + def test_notification_ahead_of_the_reply_is_skipped(self): + """A server notification arriving first must not be read as the answer.""" + stream = ( + b'data: {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info"}}\n\n' + b'data: {"jsonrpc":"2.0","id":2,"result":{"content":[{"type":"text","text":' + + json.dumps(json.dumps(SEARCH_PAYLOAD)).encode() + + b"}]}}\n\n" + ) + backend = _make_backend( + [ + _initialize_response(), + _http_response(status=202), + _http_response(body=stream, content_type="text/event-stream"), + ] + ) + + assert backend.search({"query": "hello"}) == SEARCH_PAYLOAD + + def test_failed_notification_leaves_the_handshake_unset(self): + """The protocol version header is only claimed once the handshake completed.""" + backend = _make_backend([_initialize_response(), _http_response(status=500, body=b"boom")]) + + with pytest.raises(WebSearchError, match="HTTP 500"): + backend.search({"query": "hello"}) + + assert backend._initialized is False + assert backend._mcp_session_id is None + + def test_expired_session_is_dropped_and_the_next_search_reinitializes(self): + """404 against a session we hold is the spec's signal to hand shake again.""" + backend = _make_backend( + [ + _initialize_response(), + _http_response(status=202), + _http_response(status=404, body=b"session not found"), + _initialize_response(), + _http_response(status=202), + _tools_call_response(), + ] + ) + + with pytest.raises(WebSearchError, match="HTTP 404"): + backend.search({"query": "hello"}) + assert backend._initialized is False + assert backend._mcp_session_id is None + + assert backend.search({"query": "hello"}) == SEARCH_PAYLOAD + methods = [json.loads(c.kwargs["body"])["method"] for c in backend._http.request.call_args_list] + assert methods.count("initialize") == 2 + class TestToolNameResolution: """Tests for finding the fully qualified tool name.""" @@ -675,6 +812,37 @@ def test_target_name_is_handed_to_the_backend(self): assert client.backend._target_name == "amazon-web-search" + def test_non_aws_gateway_endpoint_is_rejected(self): + """A supplied endpoint gets the same host check as one the SDK builds itself.""" + with pytest.raises(InvalidRegionError, match="non-AWS host"): + WebSearchClient( + region="us-east-1", + gateway_endpoint="https://evil.example.com/mcp", + boto3_session=MagicMock(), + ) + + def test_a_backend_and_a_region_need_no_session(self): + """Supplying both must not pay for the credential provider chain it will not use.""" + with patch("boto3.Session") as session_factory: + WebSearchClient(region="us-east-1", backend=_RecordingBackend()) + + session_factory.assert_not_called() + + def test_a_default_session_is_created_when_none_is_given(self): + with patch("boto3.Session") as session_factory: + session_factory.return_value.region_name = "us-east-1" + client = WebSearchClient(gateway_id="gw-abc123") + + assert client.region == "us-east-1" + session_factory.assert_called_once_with() + + def test_a_backend_without_a_region_falls_back_to_a_session(self): + session = MagicMock() + session.region_name = "eu-west-1" + client = WebSearchClient(backend=_RecordingBackend(), boto3_session=session) + + assert client.region == "eu-west-1" + class TestBackendProtocol: """Tests for the extension point.""" diff --git a/tests_integ/tools/test_web_search_client.py b/tests_integ/tools/test_web_search_client.py index cb067928..a814dd42 100644 --- a/tests_integ/tools/test_web_search_client.py +++ b/tests_integ/tools/test_web_search_client.py @@ -11,7 +11,10 @@ WEB_SEARCH_GATEWAY_ID: ID of a gateway with a web search connector target BEDROCK_TEST_REGION: AWS region (default: us-east-1). The connector is only offered in us-east-1, eu-west-1 and ap-northeast-1. - WEB_SEARCH_TARGET_NAME: Optional. The target name, if it is not the SDK default. + WEB_SEARCH_TARGET_NAME: Optional. The name of the web search target on that + gateway. Supplying it saves a tools/list call, since Gateway prefixes every + tool with the name of the target it came from. Without it the client + discovers the tool instead. """ import os @@ -68,14 +71,24 @@ def test_search_respects_max_results(self): assert len(response) == 1 def test_search_with_domain_filter(self): - """Needs connector version 1.2.0 or later on the target.""" + """Needs connector version 1.2.0 or later on the target. + + An include filter set per request is only accepted from 1.2.0 on, and the + version of an existing gateway's target is not this test's to choose, so an + older target skips rather than fails. + """ with self._client() as client: - response = self._search( - client, - "agentcore gateway connector targets", - max_results=5, - include_domains=["docs.aws.amazon.com"], - ) + try: + response = self._search( + client, + "agentcore gateway connector targets", + max_results=5, + include_domains=["docs.aws.amazon.com"], + ) + except WebSearchError as e: + if "domainFilter" in str(e) or "include" in str(e): + pytest.skip(f"target's connector version does not accept an include filter: {e}") + raise assert len(response) > 0 for result in response: