Merge remote-tracking branch 'origin/main' into fix/final-response-match-full-response

Author: he-yufeng

src/google/adk/telemetry/google_cloud.py

@@ -14,13 +14,17 @@
 
 from __future__ import annotations
 
+import enum
 import logging
 import os
+from typing import Any
+from typing import Callable
 from typing import cast
 from typing import Optional
 from typing import TYPE_CHECKING
 
 import google.auth
+from google.auth.transport import mtls
 from opentelemetry.sdk._logs import LogRecordProcessor
 from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
 from opentelemetry.sdk.metrics.export import MetricReader
@@ -40,6 +44,19 @@
 _GCP_LOG_NAME_ENV_VARIABLE_NAME = 'GOOGLE_CLOUD_DEFAULT_LOG_NAME'
 _DEFAULT_LOG_NAME = 'adk-otel'
 
+_DEFAULT_TELEMETRY_TRACES_ENPOINT = 'https://telemetry.googleapis.com/v1/traces'
+_DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT = (
+    'https://telemetry.mtls.googleapis.com/v1/traces'
+)
+
+
+class _MtlsEndpoint(enum.Enum):
+  """The mTLS endpoint setting."""
+
+  AUTO = 'auto'
+  ALWAYS = 'always'
+  NEVER = 'never'
+
 
 def get_gcp_exporters(
     enable_cloud_tracing: bool = False,
@@ -100,10 +117,24 @@ def _get_gcp_span_exporter(credentials: Credentials) -> SpanProcessor:
   from google.auth.transport.requests import AuthorizedSession
   from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 
+  session = AuthorizedSession(credentials=credentials)
+
+  use_client_cert = _use_client_cert_effective()
+  if use_client_cert:
+    client_cert_source = (
+        mtls.default_client_cert_source()
+        if mtls.has_default_client_cert_source()
+        else None
+    )
+    session.configure_mtls_channel()
+    endpoint = _get_api_endpoint(client_cert_source)
+  else:
+    endpoint = _DEFAULT_TELEMETRY_TRACES_ENPOINT
+
   return BatchSpanProcessor(
       OTLPSpanExporter(
-          session=AuthorizedSession(credentials=credentials),
-          endpoint='https://telemetry.googleapis.com/v1/traces',
+          session=session,
+          endpoint=endpoint,
       )
   )
 
@@ -158,3 +189,62 @@ def get_gcp_resource(project_id: Optional[str] = None) -> Resource:
         ' GCE, GKE or CloudRun related resource attributes may be missing'
     )
   return resource
+
+
+def _get_api_endpoint(
+    client_cert_source: Callable[[], tuple[bytes, bytes]] | None = None,
+) -> str:
+  """Returns API endpoint based on mTLS configuration and cert availability.
+
+  Args:
+      client_cert_source: A callable that returns the client certificate and
+        key, or None.
+
+  Returns:
+      str: The API endpoint to be used.
+  """
+  use_mtls_endpoint_str = os.getenv(
+      'GOOGLE_API_USE_MTLS_ENDPOINT', _MtlsEndpoint.AUTO.value
+  ).lower()
+
+  try:
+    use_mtls_endpoint = _MtlsEndpoint(use_mtls_endpoint_str)
+  except ValueError:
+    logger.warning(
+        'Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be one of '
+        '%s. Defaulting to %s.',
+        [e.value for e in _MtlsEndpoint],
+        _MtlsEndpoint.AUTO.value,
+    )
+    use_mtls_endpoint = _MtlsEndpoint.AUTO
+
+  if (use_mtls_endpoint is _MtlsEndpoint.ALWAYS) or (
+      use_mtls_endpoint is _MtlsEndpoint.AUTO and client_cert_source
+  ):
+    return _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT
+
+  return _DEFAULT_TELEMETRY_TRACES_ENPOINT
+
+
+def _use_client_cert_effective() -> bool:
+  """Returns whether client certificate should be used for mTLS.
+
+  This checks if the google-auth version supports should_use_client_cert
+  automatic mTLS enablement. Alternatively, it reads from the
+  GOOGLE_API_USE_CLIENT_CERTIFICATE env var.
+
+  Returns:
+      bool: whether client certificate should be used for mTLS.
+  """
+  try:
+    return bool(mtls.should_use_client_cert())
+  except (ImportError, AttributeError):
+    use_client_cert_str = os.getenv(
+        'GOOGLE_API_USE_CLIENT_CERTIFICATE', 'false'
+    ).lower()
+    if use_client_cert_str not in ('true', 'false'):
+      logger.warning(
+          'Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be'
+          ' either `true` or `false`'
+      )
+    return use_client_cert_str == 'true'

src/google/adk/tools/skill_toolset.py

@@ -19,6 +19,7 @@
 from __future__ import annotations
 
 import asyncio
+import collections
 import json
 import logging
 import mimetypes
@@ -915,10 +916,11 @@ def __init__(
     self._script_timeout = script_timeout
     # Needed for mid-turn reloading of skill tools.
     self._use_invocation_cache = False
-    self._invocation_cache: dict[
+    # Cache fetched remote skill definitions per turn to reduce requests to registry
+    self._fetched_skill_cache: collections.OrderedDict[
         str,
         dict[str, models.Skill | asyncio.Future[models.Skill | None] | None],
-    ] = {}
+    ] = collections.OrderedDict()
     self._max_cache_turns = 16
 
     self._provided_tools_by_name = {}
@@ -1023,14 +1025,13 @@ async def _get_or_fetch_skill(
       return None
 
     if invocation_id:
-      if invocation_id not in self._invocation_cache:
+      if invocation_id not in self._fetched_skill_cache:
         # Enforce bounded cache (FIFO eviction)
-        if len(self._invocation_cache) >= self._max_cache_turns:
-          oldest = next(iter(self._invocation_cache))
-          self._invocation_cache.pop(oldest)
-        self._invocation_cache[invocation_id] = {}
+        if len(self._fetched_skill_cache) >= self._max_cache_turns:
+          self._fetched_skill_cache.popitem(last=False)
+        self._fetched_skill_cache[invocation_id] = {}
 
-      turn_cache = self._invocation_cache[invocation_id]
+      turn_cache = self._fetched_skill_cache[invocation_id]
       if skill_name in turn_cache:
         cached = turn_cache[skill_name]
         if isinstance(cached, asyncio.Future):
@@ -1080,6 +1081,16 @@ async def process_llm_request(
 
     llm_request.append_instructions(instructions)
 
+  @override
+  async def close(self) -> None:
+    """Performs cleanup and releases resources held by the toolset."""
+    for turn_cache in self._fetched_skill_cache.values():
+      for cached in turn_cache.values():
+        if isinstance(cached, asyncio.Future) and not cached.done():
+          cached.cancel()
+    self._fetched_skill_cache.clear()
+    await super().close()
+
 
 def __getattr__(name: str) -> Any:
   if name == "DEFAULT_SKILL_SYSTEM_INSTRUCTION":

tests/unittests/telemetry/test_google_cloud.py

@@ -16,8 +16,18 @@
 from typing import Optional
 from unittest import mock
 
+from google.adk.telemetry import google_cloud
+from google.adk.telemetry.google_cloud import _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT
+from google.adk.telemetry.google_cloud import _DEFAULT_TELEMETRY_TRACES_ENPOINT
+from google.adk.telemetry.google_cloud import _get_api_endpoint
+from google.adk.telemetry.google_cloud import _get_gcp_span_exporter
+from google.adk.telemetry.google_cloud import _use_client_cert_effective
 from google.adk.telemetry.google_cloud import get_gcp_exporters
 from google.adk.telemetry.google_cloud import get_gcp_resource
+import google.auth.credentials
+from google.auth.transport import mtls
+from google.auth.transport import requests
+from opentelemetry.exporter.otlp.proto.http import trace_exporter
 import pytest
 
 
@@ -89,3 +99,108 @@ def test_get_gcp_resource(
       otel_resource.attributes.get("gcp.project_id", None)
       == expected_project_id
   )
+
+
+@mock.patch.object(mtls, "should_use_client_cert", autospec=True)
+def test_use_client_cert_effective_from_mtls(mock_should_use):
+  mock_should_use.return_value = True
+  assert _use_client_cert_effective()
+
+  mock_should_use.return_value = False
+  assert not _use_client_cert_effective()
+
+
+def test_use_client_cert_effective_from_env(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+):
+  with mock.patch.object(
+      mtls,
+      "should_use_client_cert",
+      autospec=True,
+      side_effect=AttributeError,
+  ):
+    monkeypatch.setenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "true")
+    assert _use_client_cert_effective()
+
+    monkeypatch.setenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false")
+    assert not _use_client_cert_effective()
+
+    # Test invalid value defaults to False
+    monkeypatch.setenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "maybe")
+    assert not _use_client_cert_effective()
+    assert (
+        "Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be"
+        " either `true` or `false`"
+        in caplog.text
+    )
+
+
+@pytest.mark.parametrize(
+    "env_val, cert_source, expected",
+    [
+        ("auto", lambda: b"cert", _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT),
+        ("auto", None, _DEFAULT_TELEMETRY_TRACES_ENPOINT),
+        ("always", None, _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT),
+        ("never", lambda: b"cert", _DEFAULT_TELEMETRY_TRACES_ENPOINT),
+        ("invalid", None, _DEFAULT_TELEMETRY_TRACES_ENPOINT),
+    ],
+)
+def test_get_api_endpoint(
+    env_val,
+    cert_source,
+    expected,
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+):
+  monkeypatch.setenv("GOOGLE_API_USE_MTLS_ENDPOINT", env_val)
+  if env_val == "invalid":
+    assert _get_api_endpoint(cert_source) == expected
+    assert (
+        "Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be one of"
+        in caplog.text
+    )
+  else:
+    assert _get_api_endpoint(cert_source) == expected
+
+
+@mock.patch.object(requests, "AuthorizedSession", autospec=True)
+@mock.patch(
+    "opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter",
+    autospec=True,
+)
+@mock.patch(
+    "google.adk.telemetry.google_cloud.BatchSpanProcessor", autospec=True
+)
+@mock.patch(
+    "google.adk.telemetry.google_cloud._use_client_cert_effective",
+    autospec=True,
+)
+@mock.patch(
+    "google.auth.transport.mtls.has_default_client_cert_source", autospec=True
+)
+@mock.patch(
+    "google.auth.transport.mtls.default_client_cert_source", autospec=True
+)
+def test_get_gcp_span_exporter_mtls(
+    mock_default_cert: mock.MagicMock,
+    mock_has_cert: mock.MagicMock,
+    mock_use_cert: mock.MagicMock,
+    mock_batch: mock.MagicMock,
+    mock_exporter: mock.MagicMock,
+    mock_session: mock.MagicMock,
+):
+  credentials = mock.create_autospec(
+      google.auth.credentials.Credentials, instance=True
+  )
+  mock_use_cert.return_value = True
+  mock_has_cert.return_value = True
+  mock_default_cert.return_value = b"cert"
+
+  _get_gcp_span_exporter(credentials)
+
+  mock_session.assert_called_once_with(credentials=credentials)
+  mock_session.return_value.configure_mtls_channel.assert_called_once()
+  mock_exporter.assert_called_once_with(
+      session=mock_session.return_value,
+      endpoint=_DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT,
+  )

tests/unittests/tools/test_skill_toolset.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import collections
 import logging
 import sys
 from unittest import mock
@@ -1859,14 +1860,14 @@ async def test_turn_scoped_skill_cache_eviction(mock_registry, mock_skill1):
   for i in range(16):
     await toolset._get_or_fetch_skill("skill1", f"turn-{i}")
 
-  assert len(toolset._invocation_cache) == 16
-  assert "turn-0" in toolset._invocation_cache
+  assert len(toolset._fetched_skill_cache) == 16
+  assert "turn-0" in toolset._fetched_skill_cache
 
   # Next turn should evict oldest (turn-0)
   await toolset._get_or_fetch_skill("skill1", "turn-16")
-  assert len(toolset._invocation_cache) == 16
-  assert "turn-0" not in toolset._invocation_cache
-  assert "turn-1" in toolset._invocation_cache
+  assert len(toolset._fetched_skill_cache) == 16
+  assert "turn-0" not in toolset._fetched_skill_cache
+  assert "turn-1" in toolset._fetched_skill_cache
 
 
 @pytest.mark.asyncio
@@ -1891,3 +1892,36 @@ async def delayed_get_skill(name):
 
   # Registry should have been called exactly once
   mock_registry.get_skill.assert_called_once_with(name="skill1")
+
+
+def test_skill_toolset_disables_invocation_cache():
+  """Verify SkillToolset disables tool invocation caching to allow dynamic tools."""
+  toolset = skill_toolset.SkillToolset()
+  assert toolset._use_invocation_cache is False
+
+
+@pytest.mark.asyncio
+async def test_close_cancels_futures_and_clears_cache():
+  # pylint: disable=protected-access
+  toolset = skill_toolset.SkillToolset()
+
+  # Create mock futures for testing close() behavior
+  loop = asyncio.get_running_loop()
+  fut1 = loop.create_future()
+  fut2 = loop.create_future()
+  fut2.set_result(None)  # Already done future
+
+  toolset._fetched_skill_cache = collections.OrderedDict(
+      {
+          "turn1": {
+              "skill1": fut1,
+              "skill2": fut2,
+          }
+      }
+  )
+
+  await toolset.close()
+
+  assert fut1.cancelled()
+  assert not fut2.cancelled()  # Done futures shouldn't/can't be cancelled
+  assert not toolset._fetched_skill_cache