diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index b41b7f4effc..0ae216086fa 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/src/google/adk/agents/llm_agent.py @@ -891,6 +891,14 @@ def validate_generate_content_config( raise ValueError( 'Response schema must be set via LlmAgent.output_schema.' ) + if ( + generate_content_config.http_options + and generate_content_config.http_options.base_url + ): + raise ValueError( + 'Base URL is a transport setting and must be set on the model or' + ' its client, not via LlmAgent.generate_content_config.' + ) return generate_content_config @override diff --git a/src/google/adk/integrations/api_registry/api_registry.py b/src/google/adk/integrations/api_registry/api_registry.py index 966ad68b7d8..a020421d9e0 100644 --- a/src/google/adk/integrations/api_registry/api_registry.py +++ b/src/google/adk/integrations/api_registry/api_registry.py @@ -16,6 +16,7 @@ from typing import Any from typing import Callable +from urllib.parse import urlparse from google.adk.agents.readonly_context import ReadonlyContext from google.adk.tools.base_toolset import ToolPredicate @@ -28,6 +29,17 @@ API_REGISTRY_URL = "https://cloudapiregistry.googleapis.com" +def _is_google_api(url: str) -> bool: + """Checks if the given URL points to a Google API endpoint over https.""" + parsed_url = urlparse(url) + if parsed_url.scheme != "https" or not parsed_url.hostname: + return False + return ( + parsed_url.hostname == "googleapis.com" + or parsed_url.hostname.endswith(".googleapis.com") + ) + + class ApiRegistry: """Registry that provides McpToolsets for MCP servers registered in API Registry.""" @@ -110,12 +122,18 @@ def get_toolset( raise ValueError(f"MCP server {mcp_server_name} has no URLs.") mcp_server_url = server["urls"][0] - headers = self._get_auth_headers() # Only prepend "https://" if the URL doesn't already have a scheme if not mcp_server_url.startswith(("http://", "https://")): mcp_server_url = "https://" + mcp_server_url + # A registry entry can name any host, so the caller's own credentials are + # only attached to Google API endpoints. Other servers get their headers + # from the header_provider. + headers = ( + self._get_auth_headers() if _is_google_api(mcp_server_url) else None + ) + return McpToolset( connection_params=StreamableHTTPConnectionParams( url=mcp_server_url, diff --git a/src/google/adk/models/apigee_llm.py b/src/google/adk/models/apigee_llm.py index a1575bdce67..ac8f92210dc 100644 --- a/src/google/adk/models/apigee_llm.py +++ b/src/google/adk/models/apigee_llm.py @@ -63,6 +63,28 @@ _REFUSAL_PREFIX = '[[REFUSAL]]: ' +# Timeouts, in seconds, for the completions HTTP client. httpx applies no +# timeout at all unless one is given, so a stalled proxy would otherwise hold +# the connection and the streaming loop open indefinitely. +_CONNECT_TIMEOUT_SECONDS = 30.0 +_REQUEST_TIMEOUT_SECONDS = 600.0 + + +def _httpx_timeout(timeout_seconds: Optional[float] = None) -> httpx.Timeout: + """Returns the httpx timeout budget for a completions request. + + A bare float would spend the caller's whole budget on the connect phase too, + so the connect budget is always kept short enough to fail fast on an + unreachable proxy. + + Args: + timeout_seconds: The total budget for the request, or None for the default. + """ + return httpx.Timeout( + _REQUEST_TIMEOUT_SECONDS if timeout_seconds is None else timeout_seconds, + connect=_CONNECT_TIMEOUT_SECONDS, + ) + class ApigeeLlm(Gemini): """A BaseLlm implementation for calling Apigee proxy. @@ -427,8 +449,8 @@ def _client(self) -> httpx.AsyncClient: client = httpx.AsyncClient( base_url=self._base_url, headers=self._headers, - timeout=None, - follow_redirects=True, + timeout=_httpx_timeout(), + follow_redirects=False, ) atexit.register(self._cleanup_client, client) return client @@ -519,6 +541,7 @@ async def generate_content_async( ) -> AsyncGenerator[LlmResponse, None]: """Generates content using the OpenAI-compatible HTTP API.""" payload = self._construct_payload(llm_request, stream) + timeout = self._get_request_timeout_seconds(llm_request) headers = self._headers.copy() headers['Content-Type'] = 'application/json' @@ -530,26 +553,50 @@ async def generate_content_async( url = f"{url.rstrip('/')}/chat/completions" if stream: - async for stream_res in self._handle_streaming(url, payload, headers): + async for stream_res in self._handle_streaming( + url, payload, headers, timeout=timeout + ): yield stream_res else: - response = await self._httpx_post_with_retry(url, payload, headers) + response = await self._httpx_post_with_retry( + url, payload, headers, timeout=timeout + ) data = response.json() yield self._parse_response(data) + @staticmethod + def _get_request_timeout_seconds(llm_request: LlmRequest) -> float | None: + """Returns the request timeout converted from milliseconds to seconds.""" + if not llm_request.config or not llm_request.config.http_options: + return None + timeout_ms = llm_request.config.http_options.timeout + return timeout_ms / 1000 if timeout_ms is not None else None + async def _httpx_post_with_retry( - self, url: str, payload: dict[str, Any], headers: dict[str, str] + self, + url: str, + payload: dict[str, Any], + headers: dict[str, str], + *, + timeout: float | None, ) -> httpx.Response: """Sends a POST request and handles retries.""" retry_kwargs = self._get_retry_kwargs() async for attempt in tenacity.AsyncRetrying(**retry_kwargs): with attempt: - response = await self._client.post(url, json=payload, headers=headers) + response = await self._client.post( + url, json=payload, headers=headers, timeout=_httpx_timeout(timeout) + ) response.raise_for_status() return response async def _handle_streaming( - self, url: str, payload: dict[str, Any], headers: dict[str, str] + self, + url: str, + payload: dict[str, Any], + headers: dict[str, str], + *, + timeout: float | None, ) -> AsyncGenerator[LlmResponse, None]: """Handles streaming response from OpenAI-compatible API.""" accumulator = ChatCompletionsResponseHandler() @@ -558,6 +605,7 @@ async def _handle_streaming( url, json=payload, headers=headers, + timeout=_httpx_timeout(timeout), ) as resp: resp.raise_for_status() async for line in resp.aiter_lines(): diff --git a/src/google/adk/tools/computer_use/computer_use_toolset.py b/src/google/adk/tools/computer_use/computer_use_toolset.py index 9707112d835..43504f8d42d 100644 --- a/src/google/adk/tools/computer_use/computer_use_toolset.py +++ b/src/google/adk/tools/computer_use/computer_use_toolset.py @@ -33,11 +33,17 @@ from ..base_toolset import BaseToolset from ..tool_context import ToolContext from .base_computer import BaseComputer +from .base_computer import ComputerState from .computer_use_tool import ComputerUseTool # Methods that should be excluded when creating tools from BaseComputer methods EXCLUDED_METHODS = {"screen_size", "environment", "close", "prepare"} +_URL_REFUSED_ERROR = ( + "navigate refused: url must be http(s) and must not target a private or" + " link-local address." +) + logger = logging.getLogger("google_adk." + __name__) @@ -49,10 +55,22 @@ def __init__( *, computer: BaseComputer, excluded_predefined_functions: Optional[list[str]] = None, + allow_private_network_access: bool = False, ): + """Initializes the ComputerUseToolset. + + Args: + computer: The computer environment to expose as tools. + excluded_predefined_functions: Names of BaseComputer methods that should + not be exposed as tools. + allow_private_network_access: By default `navigate` refuses urls whose + host is not publicly routable. Set this to True when the agent is + meant to drive the browser against localhost or an internal host. + """ super().__init__() self._computer = computer self._excluded_predefined_functions = excluded_predefined_functions + self._allow_private_network_access = allow_private_network_access self._initialized = False self._tools = None @@ -107,6 +125,42 @@ async def wrapper( return wrapper + def _wrap_navigate_with_url_validation( + self, navigate_method: Callable[..., Any] + ) -> Callable[..., Any]: + """Checks a model-supplied url before `navigate` hands it to the browser.""" + + @functools.wraps(navigate_method) + async def wrapper(url: str) -> Any: + # Deferred to keep `requests` off the computer-use import path. + from ..load_web_page import _is_blocked_hostname + from ..load_web_page import _parse_request_target + from ..load_web_page import _resolve_direct_addresses + + try: + if not isinstance(url, str): + raise ValueError("url is not a string") + target = _parse_request_target(url) + # A browser ends the authority at "\" but urlparse does not: in + # `http://169.254.169.254\@example.com/` the host is example.com here + # and 169.254.169.254 in Chrome, so refuse instead of checking it. + if "\\" in target.parsed_url.netloc: + raise ValueError("backslash in hostname") + if not self._allow_private_network_access: + if _is_blocked_hostname(target.hostname): + raise ValueError("hostname is blocked") + # getaddrinfo blocks, so keep it off the event loop. + await asyncio.to_thread(_resolve_direct_addresses, target.hostname) + except ValueError: + logger.warning("Refusing navigate(): url failed safety validation.") + # The computer-use model rejects a function response with no url, + # so report the page the browser is currently on. + state: ComputerState = await self._computer.current_state() + return {"error": _URL_REFUSED_ERROR, "url": state.url} + return await navigate_method(url) + + return wrapper + @staticmethod async def adapt_computer_use_tool( method_name: str, @@ -217,6 +271,11 @@ async def get_tools( if attr is not None and callable(attr): # Get the corresponding method from the concrete instance instance_method = getattr(self._computer, method_name) + if method_name == "navigate": + # Check the url the model supplied before it reaches the browser. + instance_method = self._wrap_navigate_with_url_validation( + instance_method + ) # Wrap with state binding so session_state is set before each call wrapped_method = self._wrap_method_with_state_binding(instance_method) computer_methods.append(wrapped_method) diff --git a/src/google/adk/tools/load_web_page.py b/src/google/adk/tools/load_web_page.py index eb86c823321..43dc8982d95 100644 --- a/src/google/adk/tools/load_web_page.py +++ b/src/google/adk/tools/load_web_page.py @@ -156,8 +156,42 @@ def _is_blocked_hostname(hostname: str) -> bool: ) +_NAT64_WELL_KNOWN_PREFIX = ipaddress.ip_network('64:ff9b::/96') + + +def _embedded_ipv4(address: _ResolvedAddress) -> ipaddress.IPv4Address | None: + """Returns the IPv4 address embedded in an IPv6 address, if any. + + ``is_global`` on the outer IPv6 address does not reflect the reachability of + the embedded IPv4 target for IPv4-mapped (``::ffff:a.b.c.d``), IPv4-compatible + (``::a.b.c.d``), 6to4 (``2002::/16``) and NAT64 (``64:ff9b::/96``) addresses. + For example ``64:ff9b::169.254.169.254`` is reported as global but, on a + network with NAT64, routes to the internal ``169.254.169.254`` metadata + endpoint. Returning the embedded IPv4 lets the caller vet it directly. + """ + if not isinstance(address, ipaddress.IPv6Address): + return None + if address.ipv4_mapped is not None: + return address.ipv4_mapped + if address.sixtofour is not None: + return address.sixtofour + if address in _NAT64_WELL_KNOWN_PREFIX: + return ipaddress.IPv4Address(int(address) & 0xFFFFFFFF) + # IPv4-compatible ``::a.b.c.d`` (deprecated): top 96 bits zero, low 32 bits a + # non-trivial IPv4 (excluding ``::`` and ``::1``). + packed = int(address) + if packed >> 32 == 0 and (packed & 0xFFFFFFFF) not in (0, 1): + return ipaddress.IPv4Address(packed & 0xFFFFFFFF) + return None + + def _is_blocked_address(address: _ResolvedAddress) -> bool: - return not address.is_global + if not address.is_global: + return True + # Reject IPv6 addresses that embed a non-global IPv4 target (NAT64, + # IPv4-compatible, etc.), which `is_global` alone does not catch. + embedded = _embedded_ipv4(address) + return embedded is not None and not embedded.is_global def _resolve_host_addresses(hostname: str) -> tuple[_ResolvedAddress, ...]: diff --git a/tests/unittests/agents/test_llm_agent_fields.py b/tests/unittests/agents/test_llm_agent_fields.py index 1158138fbce..8bec6c37cb6 100644 --- a/tests/unittests/agents/test_llm_agent_fields.py +++ b/tests/unittests/agents/test_llm_agent_fields.py @@ -300,6 +300,31 @@ class Schema(BaseModel): ) +def test_validate_generate_content_config_http_options_base_url_throw(): + """Tests that a transport base URL cannot be set directly in config.""" + with pytest.raises(ValueError): + _ = LlmAgent( + name='test_agent', + generate_content_config=types.GenerateContentConfig( + http_options=types.HttpOptions(base_url='http://example.invalid') + ), + ) + + +def test_validate_generate_content_config_http_options_allowed(): + """Tests that request-time http options remain settable in config.""" + extra_body = {'tool_config': {'function_calling_config': {'mode': 'AUTO'}}} + agent = LlmAgent( + name='test_agent', + generate_content_config=types.GenerateContentConfig( + http_options=types.HttpOptions(timeout=1000, extra_body=extra_body) + ), + ) + + assert agent.generate_content_config.http_options.timeout == 1000 + assert agent.generate_content_config.http_options.extra_body == extra_body + + def test_allow_transfer_by_default(): sub_agent = LlmAgent(name='sub_agent') agent = LlmAgent(name='test_agent', sub_agents=[sub_agent]) diff --git a/tests/unittests/integrations/api_registry/test_api_registry.py b/tests/unittests/integrations/api_registry/test_api_registry.py index 7edaee9fecf..09137e71291 100644 --- a/tests/unittests/integrations/api_registry/test_api_registry.py +++ b/tests/unittests/integrations/api_registry/test_api_registry.py @@ -44,6 +44,14 @@ "name": "test-mcp-server-https", "urls": ["https://mcp.server_https.com"], }, + { + "name": "test-mcp-server-google", + "urls": ["mcp.us-central1.googleapis.com"], + }, + { + "name": "test-mcp-server-google-http", + "urls": ["http://mcp.us-central1.googleapis.com"], + }, ] } @@ -79,7 +87,7 @@ def test_init_success(self, MockHttpClient): api_registry_project_id=self.project_id, location=self.location ) - self.assertEqual(len(api_registry._mcp_servers), 5) + self.assertEqual(len(api_registry._mcp_servers), 7) self.assertIn("test-mcp-server-1", api_registry._mcp_servers) self.assertIn("test-mcp-server-2", api_registry._mcp_servers) self.assertIn("test-mcp-server-no-url", api_registry._mcp_servers) @@ -107,7 +115,7 @@ def test_init_with_quota_project_id_success(self, MockHttpClient): api_registry_project_id=self.project_id, location=self.location ) - self.assertEqual(len(api_registry._mcp_servers), 5) + self.assertEqual(len(api_registry._mcp_servers), 7) self.assertIn("test-mcp-server-1", api_registry._mcp_servers) self.assertIn("test-mcp-server-2", api_registry._mcp_servers) self.assertIn("test-mcp-server-no-url", api_registry._mcp_servers) @@ -240,7 +248,7 @@ async def test_get_toolset_success(self, MockHttpClient, MockMcpToolset): MockMcpToolset.assert_called_once_with( connection_params=StreamableHTTPConnectionParams( url="https://mcp.server1.com", - headers={"Authorization": "Bearer mock_token"}, + headers=None, ), tool_filter=None, tool_name_prefix=None, @@ -267,11 +275,11 @@ async def test_get_toolset_with_quota_project_id_success( api_registry_project_id=self.project_id, location=self.location ) - toolset = api_registry.get_toolset("test-mcp-server-1") + toolset = api_registry.get_toolset("test-mcp-server-google") MockMcpToolset.assert_called_once_with( connection_params=StreamableHTTPConnectionParams( - url="https://mcp.server1.com", + url="https://mcp.us-central1.googleapis.com", headers={ "Authorization": "Bearer mock_token", "x-goog-user-project": "quota-project", @@ -312,7 +320,7 @@ async def test_get_toolset_with_filter_and_prefix( MockMcpToolset.assert_called_once_with( connection_params=StreamableHTTPConnectionParams( url="https://mcp.server1.com", - headers={"Authorization": "Bearer mock_token"}, + headers=None, ), tool_filter=tool_filter, tool_name_prefix=tool_name_prefix, @@ -348,13 +356,46 @@ def test_get_toolset_url_scheme(self): MockMcpToolset.assert_called_once_with( connection_params=StreamableHTTPConnectionParams( url=mock_url, - headers={"Authorization": "Bearer mock_token"}, + headers=None, ), tool_filter=None, tool_name_prefix=None, header_provider=None, ) + def test_get_toolset_credentials_only_for_google_api_url(self): + params = [ + ("test-mcp-server-1", None), + ("test-mcp-server-http", None), + ("test-mcp-server-https", None), + ("test-mcp-server-google-http", None), + ("test-mcp-server-google", {"Authorization": "Bearer mock_token"}), + ] + for mock_server_name, expected_headers in params: + with self.subTest(server_name=mock_server_name): + with ( + patch.object(httpx, "Client", autospec=True) as MockHttpClient, + patch.object( + api_registry.api_registry, "McpToolset", autospec=True + ) as MockMcpToolset, + ): + mock_response = create_autospec(httpx.Response, instance=True) + mock_response.json.return_value = MOCK_MCP_SERVERS_LIST + mock_client_instance = MockHttpClient.return_value + mock_client_instance.__enter__.return_value = mock_client_instance + mock_client_instance.get.return_value = mock_response + + api_registry_instance = ApiRegistry( + api_registry_project_id=self.project_id, location=self.location + ) + + api_registry_instance.get_toolset(mock_server_name) + + connection_params = MockMcpToolset.call_args.kwargs[ + "connection_params" + ] + self.assertEqual(connection_params.headers, expected_headers) + @patch("httpx.Client", autospec=True) async def test_get_toolset_server_not_found(self, MockHttpClient): mock_response = MagicMock() diff --git a/tests/unittests/models/test_apigee_llm.py b/tests/unittests/models/test_apigee_llm.py index 1e371e8aa13..1dd748a73de 100644 --- a/tests/unittests/models/test_apigee_llm.py +++ b/tests/unittests/models/test_apigee_llm.py @@ -14,7 +14,9 @@ from __future__ import annotations +import asyncio import os +import time from unittest import mock from unittest.mock import AsyncMock @@ -610,6 +612,128 @@ async def test_generate_content_async_dispatch_to_completions_client( mock_genai_client.assert_not_called() +@pytest.mark.asyncio +async def test_chat_completions_honors_request_timeout(): + """Chat completions use the timeout configured on the LLM request.""" + request = LlmRequest( + model='apigee/openai/gpt-4o', + contents=[], + config=types.GenerateContentConfig( + http_options=types.HttpOptions(timeout=1500) + ), + ) + response = mock.MagicMock() + response.json.return_value = { + 'choices': [{ + 'message': {'role': 'assistant', 'content': 'Done'}, + 'finish_reason': 'stop', + }] + } + http_client = mock.MagicMock() + http_client.post = AsyncMock(return_value=response) + + with mock.patch( + 'google.adk.models.apigee_llm.httpx.AsyncClient', + return_value=http_client, + ): + client = CompletionsHTTPClient(base_url=PROXY_URL) + _ = [item async for item in client.generate_content_async(request, False)] + + _, call_kwargs = http_client.post.await_args + assert call_kwargs['timeout'].read == 1.5 + # The caller's budget must not stretch the fast-failing connect phase. + assert call_kwargs['timeout'].connect == 30.0 + + +@pytest.mark.asyncio +async def test_streaming_chat_completions_honors_request_timeout(): + """Streaming chat completions use the configured request timeout.""" + request = LlmRequest( + model='apigee/openai/gpt-4o', + contents=[], + config=types.GenerateContentConfig( + http_options=types.HttpOptions(timeout=2500) + ), + ) + + async def stream_lines(): + yield 'data: [DONE]' + + response = mock.MagicMock() + response.aiter_lines = stream_lines + stream_context = mock.MagicMock() + stream_context.__aenter__ = AsyncMock(return_value=response) + stream_context.__aexit__ = AsyncMock(return_value=None) + http_client = mock.MagicMock() + http_client.stream.return_value = stream_context + + with mock.patch( + 'google.adk.models.apigee_llm.httpx.AsyncClient', + return_value=http_client, + ): + client = CompletionsHTTPClient(base_url=PROXY_URL) + _ = [item async for item in client.generate_content_async(request, True)] + + _, call_kwargs = http_client.stream.call_args + assert call_kwargs['timeout'].read == 2.5 + assert call_kwargs['timeout'].connect == 30.0 + + +@pytest.mark.asyncio +async def test_chat_completions_without_request_timeout_stays_bounded(): + """A request with no configured timeout still gets the default budget.""" + request = LlmRequest(model='apigee/openai/gpt-4o', contents=[]) + response = mock.MagicMock() + response.json.return_value = { + 'choices': [{ + 'message': {'role': 'assistant', 'content': 'Done'}, + 'finish_reason': 'stop', + }] + } + http_client = mock.MagicMock() + http_client.post = AsyncMock(return_value=response) + + with mock.patch( + 'google.adk.models.apigee_llm.httpx.AsyncClient', + return_value=http_client, + ): + client = CompletionsHTTPClient(base_url=PROXY_URL) + _ = [item async for item in client.generate_content_async(request, False)] + + _, call_kwargs = http_client.post.await_args + # A bare timeout=None here would switch every timeout back off. + assert call_kwargs['timeout'].connect == 30.0 + assert call_kwargs['timeout'].read == 600.0 + + +@pytest.mark.asyncio +async def test_streaming_chat_completions_without_request_timeout_stays_bounded(): + """A stream with no configured timeout still gets the default budget.""" + request = LlmRequest(model='apigee/openai/gpt-4o', contents=[]) + + async def stream_lines(): + yield 'data: [DONE]' + + response = mock.MagicMock() + response.aiter_lines = stream_lines + stream_context = mock.MagicMock() + stream_context.__aenter__ = AsyncMock(return_value=response) + stream_context.__aexit__ = AsyncMock(return_value=None) + http_client = mock.MagicMock() + http_client.stream.return_value = stream_context + + with mock.patch( + 'google.adk.models.apigee_llm.httpx.AsyncClient', + return_value=http_client, + ): + client = CompletionsHTTPClient(base_url=PROXY_URL) + _ = [item async for item in client.generate_content_async(request, True)] + + _, call_kwargs = http_client.stream.call_args + assert call_kwargs['timeout'].connect == 30.0 + assert call_kwargs['timeout'].read == 600.0 + + @pytest.mark.asyncio @pytest.mark.parametrize( 'model', @@ -629,6 +753,94 @@ async def test_api_key_injection_openai(model): assert client._headers['Authorization'] == 'Bearer sk-test-key' +def test_completions_http_client_bounds_requests_and_stays_on_base_url() -> ( + None +): + """Tests that the httpx client has finite timeouts and does not redirect.""" + completions_client = CompletionsHTTPClient(base_url='http://test') + try: + httpx_client = completions_client._client + timeout = httpx_client.timeout + # Pinned to the literal budgets rather than to the constants themselves, + # so that shrinking a constant to something a slow model cannot meet + # fails here. + assert timeout.connect == 30.0 + assert timeout.read == 600.0 + assert timeout.write == 600.0 + assert timeout.pool == 600.0 + assert not httpx_client.follow_redirects + finally: + completions_client.close() + + +@pytest.mark.asyncio +async def test_completions_http_client_streams_longer_than_request_timeout() -> ( + None +): + """Tests that a slow but steady stream outlives the request timeout.""" + request_timeout_seconds = 1.0 + chunk_gap_seconds = 0.25 + chunk_count = 6 + # Every gap between chunks stays well inside the budget while the whole + # generation runs past it, because httpx spends the budget per read rather + # than per request. + assert chunk_gap_seconds < request_timeout_seconds + assert chunk_gap_seconds * chunk_count > request_timeout_seconds + + async def serve_slow_stream(reader, writer): + head = await reader.readuntil(b'\r\n\r\n') + for header in head.split(b'\r\n'): + if header.lower().startswith(b'content-length:'): + await reader.readexactly(int(header.split(b':')[1])) + writer.write( + b'HTTP/1.1 200 OK\r\n' + b'Content-Type: text/event-stream\r\n' + b'Transfer-Encoding: chunked\r\n' + b'\r\n' + ) + for index in range(chunk_count): + await asyncio.sleep(chunk_gap_seconds) + body = ( + 'data: {"choices": [{"index": 0, "delta": {"content":' + f' "{index}"}}, "finish_reason": null}}]}}\n\n' + ).encode() + writer.write(f'{len(body):x}\r\n'.encode() + body + b'\r\n') + await writer.drain() + writer.write(b'0\r\n\r\n') + await writer.drain() + writer.close() + + server = await asyncio.start_server(serve_slow_stream, '127.0.0.1', 0) + port = server.sockets[0].getsockname()[1] + completions_client = CompletionsHTTPClient( + base_url=f'http://127.0.0.1:{port}' + ) + request = LlmRequest( + model='apigee/openai/gpt-4o', + contents=[Content(role='user', parts=[Part.from_text(text='hi')])], + ) + try: + with mock.patch( + 'google.adk.models.apigee_llm._REQUEST_TIMEOUT_SECONDS', + request_timeout_seconds, + ): + started = time.monotonic() + responses = [ + response + async for response in completions_client.generate_content_async( + request, stream=True + ) + ] + elapsed = time.monotonic() - started + finally: + await completions_client.aclose() + server.close() + await server.wait_closed() + + assert len(responses) == chunk_count + assert elapsed > request_timeout_seconds + + def test_parse_response_usage_metadata(): """Tests that CompletionsHTTPClient parses usage metadata correctly including reasoning tokens.""" client = CompletionsHTTPClient(base_url='http://test') diff --git a/tests/unittests/tools/computer_use/test_computer_use_toolset.py b/tests/unittests/tools/computer_use/test_computer_use_toolset.py index 27755a465fe..48278e8ae7c 100644 --- a/tests/unittests/tools/computer_use/test_computer_use_toolset.py +++ b/tests/unittests/tools/computer_use/test_computer_use_toolset.py @@ -12,10 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import socket from unittest.mock import AsyncMock from unittest.mock import MagicMock +from unittest.mock import Mock from google.adk.models.llm_request import LlmRequest +from google.adk.tools import load_web_page # Use the actual ComputerEnvironment enum from the code from google.adk.tools.computer_use.base_computer import BaseComputer from google.adk.tools.computer_use.base_computer import ComputerEnvironment @@ -32,6 +35,7 @@ class MockComputer(BaseComputer): def __init__(self): self.initialize_called = False self.close_called = False + self.navigate_calls: list[str] = [] self._screen_size = (1920, 1080) self._environment = ComputerEnvironment.ENVIRONMENT_BROWSER @@ -88,6 +92,7 @@ async def search(self) -> ComputerState: return ComputerState(screenshot=b"test", url="https://example.com") async def navigate(self, url: str) -> ComputerState: + self.navigate_calls.append(url) return ComputerState(screenshot=b"test", url=url) async def key_combination(self, keys: list[str]) -> ComputerState: @@ -610,3 +615,128 @@ async def adapted_func(): # Should not add any tools assert len(llm_request.tools_dict) == 0 + + +class TestNavigateUrlSafety: + """Test cases for the navigate() url safety guard.""" + + @pytest.fixture + def mock_computer(self): + """Fixture providing a mock computer.""" + return MockComputer() + + @pytest.fixture + def resolver(self, monkeypatch) -> Mock: + """Records every DNS lookup, resolving whatever is asked for to a public ip. + + Tests assert on this to pin down whether a url was refused before or after + its hostname was resolved. + """ + resolver = Mock( + return_value=[( + socket.AF_INET, + socket.SOCK_STREAM, + socket.IPPROTO_TCP, + "", + ("93.184.216.34", 0), + )] + ) + monkeypatch.setattr(load_web_page.socket, "getaddrinfo", resolver) + return resolver + + @staticmethod + async def _build_navigate_tool( + computer: MockComputer, **toolset_kwargs + ) -> ComputerUseTool: + """Returns the navigate tool of a toolset built over `computer`. + + Toolset settings differ per test, so they are passed in rather than fixed + by a fixture. + """ + toolset = ComputerUseToolset(computer=computer, **toolset_kwargs) + for tool in await toolset.get_tools(): + if tool.func.__name__ == "navigate": + return tool + raise AssertionError("No navigate tool in toolset") + + @pytest.mark.parametrize( + "url", + [ + ( + "http://169.254.169.254/computeMetadata/v1/instance/" + "service-accounts/default/token" + ), + # Parser-divergence regression test, not a duplicate: `urlparse` + # reads the host as 'example.com' while a browser reads + # '169.254.169.254', so the url is refused outright. + r"http://169.254.169.254\@example.com/", + "file:///etc/passwd", + "http://localhost:3000/", + ], + ids=["cloud_metadata", "backslash_authority", "file_scheme", "localhost"], + ) + @pytest.mark.asyncio + async def test_navigate_refuses_unsafe_url( + self, url, mock_computer, resolver + ): + """Unsafe urls are refused before the driver or a resolver is touched.""" + navigate_tool = await self._build_navigate_tool(mock_computer) + url_before = (await mock_computer.current_state()).url + + result = await navigate_tool.func(url=url) + + # The security-critical assertion: the driver was never asked to navigate. + assert mock_computer.navigate_calls == [] + assert result["error"] + # The computer-use model rejects a function response carrying no url, so a + # refusal reports the page the browser is still on + assert result["url"] == url_before + # Refusal is decided from the url alone, so no case reaches DNS. + resolver.assert_not_called() + + @pytest.mark.asyncio + async def test_navigate_allows_public_url_unmodified( + self, mock_computer, resolver + ): + """A public url reaches the computer exactly as the model provided it.""" + navigate_tool = await self._build_navigate_tool(mock_computer) + + result = await navigate_tool.func(url="https://example.com/search?q=adk") + + assert mock_computer.navigate_calls == ["https://example.com/search?q=adk"] + assert result.url == "https://example.com/search?q=adk" + resolver.assert_called_once() + + @pytest.mark.asyncio + async def test_navigate_allows_loopback_with_private_network_access( + self, mock_computer, resolver + ): + """allow_private_network_access=True lets loopback urls through.""" + navigate_tool = await self._build_navigate_tool( + mock_computer, allow_private_network_access=True + ) + + result = await navigate_tool.func(url="http://127.0.0.1:8000/") + + assert mock_computer.navigate_calls == ["http://127.0.0.1:8000/"] + assert result.url == "http://127.0.0.1:8000/" + # The flag short-circuits before resolution, so no lookup happens. + resolver.assert_not_called() + + @pytest.mark.asyncio + async def test_navigate_guard_preserves_tool_context( + self, mock_computer, resolver + ): + """The url guard must not cut navigate off from tool_context.""" + mock_computer.prepare = AsyncMock() + tool_context = MagicMock(tool_confirmation=None) + navigate_tool = await self._build_navigate_tool(mock_computer) + + result = await navigate_tool.run_async( + args={"url": "https://example.com"}, tool_context=tool_context + ) + + mock_computer.prepare.assert_awaited_once_with(tool_context) + assert mock_computer.navigate_calls == ["https://example.com"] + assert result["url"] == "https://example.com" + resolver.assert_called_once() diff --git a/tests/unittests/tools/test_load_web_page.py b/tests/unittests/tools/test_load_web_page.py index 5dbc70805be..ccb8b6e9b40 100644 --- a/tests/unittests/tools/test_load_web_page.py +++ b/tests/unittests/tools/test_load_web_page.py @@ -92,6 +92,86 @@ def test_load_web_page_blocks_shared_address_space_urls(monkeypatch): mock_send.assert_not_called() +def test_load_web_page_blocks_nat64_embedded_metadata_ip(monkeypatch): + _clear_proxy_env(monkeypatch) + mock_get = mock.Mock() + monkeypatch.setattr(load_web_page_module.requests, 'get', mock_get) + mock_send = mock.Mock() + monkeypatch.setattr(load_web_page_module.HTTPAdapter, 'send', mock_send) + + result = load_web_page( + 'http://[64:ff9b::169.254.169.254]/computeMetadata/v1/' + ) + + assert ( + result + == 'Failed to fetch url:' + ' http://[64:ff9b::169.254.169.254]/computeMetadata/v1/' + ) + mock_get.assert_not_called() + mock_send.assert_not_called() + + +def test_load_web_page_blocks_ipv4_compatible_embedded_private_ip(monkeypatch): + _clear_proxy_env(monkeypatch) + mock_get = mock.Mock() + monkeypatch.setattr(load_web_page_module.requests, 'get', mock_get) + mock_send = mock.Mock() + monkeypatch.setattr(load_web_page_module.HTTPAdapter, 'send', mock_send) + + result = load_web_page('http://[::169.254.169.254]/latest/meta-data/') + + assert ( + result + == 'Failed to fetch url: http://[::169.254.169.254]/latest/meta-data/' + ) + mock_get.assert_not_called() + mock_send.assert_not_called() + + +def test_load_web_page_allows_public_nat64_ip(monkeypatch): + _clear_proxy_env(monkeypatch) + mock_get = mock.Mock() + monkeypatch.setattr(load_web_page_module.requests, 'get', mock_get) + + captured_request: dict[str, object] = {} + + def _send( + self, + request, + stream=False, + timeout=None, + verify=True, + cert=None, + proxies=None, + ): + del self, stream, timeout, verify, cert, proxies + captured_request['url'] = request.url + captured_request['host_header'] = request.headers['Host'] + return _create_response( + '
This page has enough words to keep.
' + ) + + monkeypatch.setattr(load_web_page_module.HTTPAdapter, 'send', _send) + monkeypatch.setattr( + 'bs4.BeautifulSoup', + mock.Mock( + return_value=mock.Mock( + get_text=mock.Mock( + return_value='This page has enough words to keep.' + ) + ) + ), + ) + + result = load_web_page('http://[64:ff9b::8.8.8.8]/') + + assert result == 'This page has enough words to keep.' + assert captured_request['url'] == 'http://[64:ff9b::808:808]/' + assert captured_request['host_header'] == '[64:ff9b::8.8.8.8]' + mock_get.assert_not_called() + + def test_load_web_page_blocks_private_hostname_targets(monkeypatch): _clear_proxy_env(monkeypatch) monkeypatch.setattr(