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..9d36c9f1 --- /dev/null +++ b/src/bedrock_agentcore/tools/web_search_client.py @@ -0,0 +1,666 @@ +"""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. + + 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. + + 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``. + + 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. 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. + + 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)