feat(tools): add WebSearchClient for invoking Amazon Web Search - #658
feat(tools): add WebSearchClient for invoking Amazon Web Search#658sundargthb wants to merge 2 commits into
Conversation
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 "<target>___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.
… 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.
✅ No Breaking Changes DetectedNo public API breaking changes found in this PR. |
|
Claude Security Review: no high-confidence findings. (run) |
tejaskash
left a comment
There was a problem hiding this comment.
Reviewed the transport layer against the MCP streamable HTTP spec. Tests pass locally. Inline comments below, ordered roughly by severity. The first three matter once this runs against a real gateway.
| message = json.loads(chunk) | ||
| except json.JSONDecodeError: | ||
| continue | ||
| if isinstance(message, dict) and "jsonrpc" in message: |
There was a problem hiding this comment.
This returns the first JSON-RPC message on the stream, not the response. If the gateway sends a notification (e.g. notifications/progress) before the tools/call result, the notification is returned and search() fails with "contained no text content".
_post knows the outgoing id, so matching on it fixes this. At minimum, require a result or error key.
| if self._initialized: | ||
| extra["MCP-Protocol-Version"] = self._protocol_version | ||
|
|
||
| response = self._http.request( |
There was a problem hiding this comment.
Connection refused, DNS failure, TLS errors and read timeouts escape here as urllib3 exceptions (MaxRetryError, ReadTimeoutError, ProtocolError). The search() docstring promises WebSearchError, so callers who follow it will miss the most common failure for a 30s HTTP call.
Suggest wrapping urllib3.exceptions.HTTPError into WebSearchError with from exc.
| if session_id: | ||
| self._mcp_session_id = session_id | ||
|
|
||
| if response.status >= 400: |
There was a problem hiding this comment.
If the gateway expires the session and returns 404 (the spec's re-initialize signal), this raises but leaves _initialized=True and _mcp_session_id set. Every later search() on this instance fails the same way. Only close() recovers, and that is not documented.
On 404 with a session id in flight, reset _initialized and _mcp_session_id before raising so the next call re-initializes.
| if isinstance(negotiated, str) and negotiated: | ||
| self._protocol_version = negotiated | ||
|
|
||
| self._initialized = True |
There was a problem hiding this comment.
Set after the notification POST on the next line succeeds. If that POST raises, the flag stays True and the next search() skips the handshake.
| text = data.decode("utf-8", "replace") | ||
|
|
||
| if "text/event-stream" in content_type.lower(): | ||
| for line in text.splitlines(): |
There was a problem hiding this comment.
Per the SSE spec, consecutive data: lines in one event are joined with \n. Parsing each line alone means a message split across two lines is dropped and reported as "Gateway did not answer". Unlikely in practice but the parser claims to handle this framing.
| """ | ||
| import boto3 | ||
|
|
||
| self._session = boto3_session or boto3.Session() |
There was a problem hiding this comment.
When backend is supplied this session is only used to read region_name, but it still does the full credential provider chain setup. Harmless, just wasted work.
| 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 |
There was a problem hiding this comment.
"Up to 100" is documented for both domain lists but not enforced in _build_arguments, while query length and max_results are. Not wrong, just inconsistent.
| self.close() | ||
|
|
||
|
|
||
| def _parse_gateway_arn(arn: str) -> tuple: |
There was a problem hiding this comment.
Tuple[str, str] to match browser_client.py.
| create_browser_config, | ||
| ) | ||
| from .web_search_client import ( | ||
| WebSearchBackend, |
There was a problem hiding this comment.
GatewayMcpBackend is not exported though WebSearchBackend is. Anyone building the concrete backend directly (e.g. to share one across clients) has to import from the module path.
| 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): |
There was a problem hiding this comment.
This tests botocore, not this PR. connection is in botocore's SIGNED_HEADERS_BLACKLIST and no connection header is ever built here, so the test cannot fail. Suggest removing it.
Description
The SDK can already create a web search connector target on a gateway (
GatewayClient.create_web_search_target, #656), but there is no way to call the resulting tool from Python. Today anything that is not already an MCP client has to hand-roll the SigV4 signed MCP handshake against the gateway endpoint to use web search.This adds
WebSearchClient:Design notes
No new dependency. The transport sits behind a
WebSearchBackendseam.GatewayMcpBackendspeaks only the slice of MCP streamable HTTP that one tool call needs (initialize, thenotifications/initializedack, optionallytools/list, thentools/call) usingurllib3andbotocore.auth.SigV4Auth, both already core dependencies.mcp-proxy-for-awswould do the signing for us, but it is currently dev/test only here and it is async, so it would change the SDK's dependency surface and the calling style for one capability. If reviewers would rather take that dependency as a published extra, the backend is one file to swap and the publicsearch()signature does not move.Two response framings. A gateway may answer a POST with
application/jsonor withtext/event-stream. Both are parsed, since which one comes back for a singletools/callis not something the client can assume.Tool name resolution. Gateway prefixes every tool with the name of the target it came from, delimited by three underscores, so the tool the agent sees is
<targetName>___WebSearch, notWebSearch. Passingtarget_namederives the name directly and skips a round trip. Otherwisetools/listis walked (followingnextCursor) and the WebSearch tool is picked, with a clear error when more than one target exposes one.Validation before the call.
queryis required and capped at 200 characters,max_resultsat 1 to 25, per the documented tool schema. The domain and published-date filters require connector version 1.2.0 or later on the target, which the docstring says.Citations are preserved.
WebSearchResultkeepsurl,titleandpublished_datealongsidetext, because the acceptable use terms require source citations and links to be retained and displayed in anything surfaced to an end user. Dropping them in the response type would make compliant use harder than it needs to be.Endpoint safety.
get_gateway_mcp_endpointvalidates the gateway identifier as a DNS label before interpolating it into the hostname, matching the region validation already in_utils/endpoints.py, so a crafted identifier cannot redirect the request off AWS. Theconnectionheader is excluded from signing, which otherwise produces a signature mismatch server-side.Region handling. The connector is offered in us-east-1, eu-west-1 and ap-northeast-1. Calling from another region logs a warning rather than raising, so a stale constant in the SDK never blocks a call to a region that was added after the release.
Testing
uv run pytest tests -qpasses: 3501 passed, 10 skipped, 4 xpassed. 97 new unit tests (81 for the client, 16 for the gateway identifier validation) cover the SigV4 header set, the request sequence, session reuse, both response framings, tool-name resolution including pagination and ambiguity, theisErrorand JSON-RPC error paths, ARN parsing, and the gateway identifier validation.ruff checkandruff format --checkare clean. Statement coverage of the new module is 100%.Not verified against a live service. The web search connector is enabled per account and this repo's integration test account is not entitled to it, so the invoke path has been exercised only against mocked HTTP responses. The request shapes come from the public tool schema and response format documentation, and the MCP sequence from the streamable HTTP transport spec, but nothing here has completed a real search.
tests_integ/tools/test_web_search_client.pyis included and will exercise the real path once run against an entitled account withWEB_SEARCH_GATEWAY_IDset; it skips rather than fails without one. Reviewers with an entitled account running that file would be the useful confirmation.Confirmed with the service team
The API surface this is built on has since been confirmed by the Web Search Tool team, and matches what is implemented here: SigV4 (AWS_IAM) with no search-specific API key or credential; availability in us-east-1, eu-west-1 and ap-northeast-1; access scoped through IAM policies on the gateway ARN rather than any key; and both
domainFilter(include and exclude, 100 domains per list) andpublishedDateFilter(from/to, ISO-8601 UTC, inclusive) at target and request level from connector version 1.2.0.Two clarifications from that exchange are now in the docstrings: the caller's credentials need
bedrock-agentcore:InvokeGatewayon the gateway ARN while the gateway service role needsbedrock-agentcore:InvokeWebSearchon the connector, and request-level include lists intersect with the target-level include list rather than replacing it, so disjoint lists return no results with no error raised.Checklist