From 737e052a7ae07963474285b062ad64875bb119c6 Mon Sep 17 00:00:00 2001 From: Sundar Raghavan Date: Thu, 3 Sep 2026 12:31:38 -0700 Subject: [PATCH 1/2] 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/2] 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.