Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions src/google/adk/agents/llm_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
20 changes: 19 additions & 1 deletion src/google/adk/integrations/api_registry/api_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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."""

Expand Down Expand Up @@ -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,
Expand Down
62 changes: 55 additions & 7 deletions src/google/adk/models/apigee_llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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'

Expand All @@ -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()
Expand All @@ -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():
Expand Down
59 changes: 59 additions & 0 deletions src/google/adk/tools/computer_use/computer_use_toolset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)


Expand All @@ -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

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
36 changes: 35 additions & 1 deletion src/google/adk/tools/load_web_page.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, ...]:
Expand Down
25 changes: 25 additions & 0 deletions tests/unittests/agents/test_llm_agent_fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
Loading
Loading