From 555a9377aa099589189fdd9e6c2dfbe1d635d99a Mon Sep 17 00:00:00 2001 From: Parth Ajmera Date: Wed, 15 Jul 2026 11:21:15 -0700 Subject: [PATCH 1/3] feat(adk): pass telemetry attributes --- packages/toolbox-adk/src/toolbox_adk/tool.py | 60 ++++++++++++- .../toolbox-adk/src/toolbox_adk/toolset.py | 17 ++++ packages/toolbox-adk/tests/unit/test_tool.py | 87 +++++++++++++++++++ .../toolbox-adk/tests/unit/test_toolset.py | 24 ++++- 4 files changed, 184 insertions(+), 4 deletions(-) diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index 0d94821b6..9e27e0807 100644 --- a/packages/toolbox-adk/src/toolbox_adk/tool.py +++ b/packages/toolbox-adk/src/toolbox_adk/tool.py @@ -14,7 +14,7 @@ import inspect import logging -from typing import Any, Awaitable, Callable, Dict, Mapping, Optional +from typing import Any, Awaitable, Callable, Dict, Mapping, Optional, Union import toolbox_core from fastapi.openapi.models import OAuth2, OAuthFlowAuthorizationCode, OAuthFlows @@ -27,7 +27,11 @@ from google.adk.tools.base_tool import BaseTool from google.adk.tools.tool_context import ToolContext from google.genai.types import FunctionDeclaration, Schema, Type -from toolbox_core.protocol import AdditionalPropertiesSchema, ParameterSchema +from toolbox_core.protocol import ( + AdditionalPropertiesSchema, + ParameterSchema, + TelemetryAttributes, +) from toolbox_core.tool import ToolboxTool as CoreToolboxTool from typing_extensions import override @@ -45,12 +49,26 @@ def __init__( core_tool: CoreToolboxTool, auth_config: Optional[CredentialConfig] = None, adk_token_getters: Optional[Mapping[str, Any]] = None, + telemetry_attributes: Optional[ + Union[ + TelemetryAttributes, + Callable[[], TelemetryAttributes], + Callable[[ToolContext], TelemetryAttributes], + Callable[[], Awaitable[TelemetryAttributes]], + Callable[[ToolContext], Awaitable[TelemetryAttributes]], + ] + ] = None, ): """ Args: core_tool: The underlying toolbox_core.py tool instance. auth_config: Credential configuration to handle interactive flows. adk_token_getters: Tool-specific auth token getters. + telemetry_attributes: Telemetry attributes (model, user id, agent + id) sent to the server with each invocation. Either a static + TelemetryAttributes instance, or a callable resolved on every + invocation; callables taking one argument receive the live + ToolContext, mirroring the auth token getter pattern. """ # We act as a proxy. # We need to extract metadata from the core tool to satisfy BaseTool's contract. @@ -74,6 +92,7 @@ def __init__( self._core_tool = core_tool self._auth_config = auth_config self._adk_token_getters = adk_token_getters or {} + self._telemetry_attributes = telemetry_attributes def _param_type_to_schema_type(self, param_type: str) -> Type: type_map = { @@ -288,8 +307,20 @@ async def run_async( error: Optional[Exception] = None try: + # Resolve telemetry attributes onto a per-invocation copy of the + # core tool. Unlike the auth getters above, this must not be + # assigned back to self._core_tool: attributes resolved from one + # ToolContext (e.g. a user id) would otherwise persist and leak + # into later invocations served by this same tool instance. + core_tool = self._core_tool + telemetry_attributes = await self._resolve_telemetry_attributes( + tool_context + ) + if telemetry_attributes is not None: + core_tool = core_tool.add_telemetry_attributes(telemetry_attributes) + # Execute the core tool - result = await self._core_tool(**args) + result = await core_tool(**args) return result except Exception as e: @@ -299,6 +330,28 @@ async def run_async( if reset_token: USER_TOKEN_CONTEXT_VAR.reset(reset_token) + async def _resolve_telemetry_attributes( + self, tool_context: ToolContext + ) -> Optional[TelemetryAttributes]: + """Resolves the configured telemetry attributes for one invocation. + + Static instances are returned as-is. Callables are invoked on every + call; 1-arity callables receive the live tool_context (same arity + sniffing as adk_token_getters) and may be sync or async. + """ + attributes = self._telemetry_attributes + if attributes is None or isinstance(attributes, TelemetryAttributes): + return attributes + + sig = inspect.signature(attributes) + if len(sig.parameters) == 1: + resolved = attributes(tool_context) + else: + resolved = attributes() + if inspect.isawaitable(resolved): + resolved = await resolved + return resolved + def bind_params(self, bounded_params: Dict[str, Any]) -> "ToolboxTool": """Allows runtime binding of parameters, delegating to core tool.""" new_core_tool = self._core_tool.bind_params(bounded_params) @@ -307,4 +360,5 @@ def bind_params(self, bounded_params: Dict[str, Any]) -> "ToolboxTool": core_tool=new_core_tool, auth_config=self._auth_config, adk_token_getters=self._adk_token_getters, + telemetry_attributes=self._telemetry_attributes, ) diff --git a/packages/toolbox-adk/src/toolbox_adk/toolset.py b/packages/toolbox-adk/src/toolbox_adk/toolset.py index 7689f96d9..956e54b7b 100644 --- a/packages/toolbox-adk/src/toolbox_adk/toolset.py +++ b/packages/toolbox-adk/src/toolbox_adk/toolset.py @@ -18,6 +18,7 @@ from google.adk.tools.base_tool import BaseTool from google.adk.tools.base_toolset import BaseToolset from google.adk.tools.tool_context import ToolContext +from toolbox_core.protocol import TelemetryAttributes from toolbox_core.utils import validate_unused_requirements from typing_extensions import override @@ -52,6 +53,15 @@ def __init__( ], ] ] = None, + telemetry_attributes: Optional[ + Union[ + TelemetryAttributes, + Callable[[], TelemetryAttributes], + Callable[[ToolContext], TelemetryAttributes], + Callable[[], Awaitable[TelemetryAttributes]], + Callable[[ToolContext], Awaitable[TelemetryAttributes]], + ] + ] = None, **kwargs: Any, ): """ @@ -63,6 +73,11 @@ def __init__( additional_headers: Extra headers (static or dynamic). bound_params: Parameters to bind globally to loaded tools. auth_token_getters: Mapping of auth service names to token getters. + telemetry_attributes: Telemetry attributes (model, user id, agent + id) sent to the server with each tool invocation. Either a + static TelemetryAttributes instance, or a callable resolved on + every invocation; callables taking one argument receive the + live ToolContext. """ super().__init__() self.__server_url = server_url @@ -75,6 +90,7 @@ def __init__( self.__tool_names = tool_names self.__bound_params = bound_params self.__auth_token_getters = auth_token_getters + self.__telemetry_attributes = telemetry_attributes @property def client(self) -> ToolboxClient: @@ -151,6 +167,7 @@ async def get_tools( core_tool=t, auth_config=self.client.credential_config, adk_token_getters=self.__auth_token_getters, + telemetry_attributes=self.__telemetry_attributes, ) for t in tools ] diff --git a/packages/toolbox-adk/tests/unit/test_tool.py b/packages/toolbox-adk/tests/unit/test_tool.py index 3b55290ea..d944f6c85 100644 --- a/packages/toolbox-adk/tests/unit/test_tool.py +++ b/packages/toolbox-adk/tests/unit/test_tool.py @@ -16,7 +16,9 @@ import pytest from google.genai.types import Type +from toolbox_core.protocol import TelemetryAttributes +from toolbox_adk.client import USER_TOKEN_CONTEXT_VAR from toolbox_adk.credentials import CredentialConfig, CredentialType from toolbox_adk.tool import ToolboxTool @@ -118,6 +120,49 @@ def getter2(ctx): assert "service2" in bound_getters assert bound_getters["service2"]() == "dynamic_token2" + @pytest.mark.asyncio + async def test_dynamic_telemetry_attributes_resolved_per_invocation(self): + core_tool = AsyncMock() + core_tool.__name__ = "mock" + core_tool.__doc__ = "mock doc" + + derived_tools = [] + + def add_telemetry_attributes(attrs): + derived_tool = AsyncMock(return_value=f"ok-{attrs.user_id}") + derived_tool.__name__ = "mock" + derived_tool.__doc__ = "mock doc" + derived_tools.append(derived_tool) + return derived_tool + + core_tool.add_telemetry_attributes = MagicMock( + side_effect=add_telemetry_attributes + ) + + def telemetry_getter(ctx): + return TelemetryAttributes( + llm_model="gemini-2.5-pro", + user_id=ctx.state["user_id"], + agent_id="agent-1", + ) + + tool = ToolboxTool(core_tool, telemetry_attributes=telemetry_getter) + + ctx1 = MagicMock() + ctx1.state = {"user_id": "user-1"} + ctx2 = MagicMock() + ctx2.state = {"user_id": "user-2"} + + assert await tool.run_async({}, ctx1) == "ok-user-1" + assert await tool.run_async({}, ctx2) == "ok-user-2" + + calls = core_tool.add_telemetry_attributes.call_args_list + assert [c.args[0].user_id for c in calls] == ["user-1", "user-2"] + assert all(c.args[0].llm_model == "gemini-2.5-pro" for c in calls) + assert core_tool.await_count == 0 + assert [t.await_count for t in derived_tools] == [1, 1] + assert tool._core_tool is core_tool + @pytest.mark.asyncio async def test_3lo_missing_client_secret(self): # Test ValueError when client_id/secret missing @@ -230,6 +275,48 @@ async def test_3lo_uses_existing_credential(self): "email": "", } + @pytest.mark.asyncio + async def test_3lo_resets_user_token_when_telemetry_getter_raises(self): + core_tool = AsyncMock(return_value="success") + core_tool.__name__ = "mock_tool" + core_tool.__doc__ = "mock doc" + core_tool._required_authn_params = {"mock_param": "mock_service"} + core_tool._required_authz_tokens = [] + core_tool.add_auth_token_getter = MagicMock(return_value=core_tool) + + auth_config = CredentialConfig( + type=CredentialType.USER_IDENTITY, client_id="cid", client_secret="csec" + ) + + def telemetry_getter(ctx): + raise RuntimeError("telemetry boom") + + tool = ToolboxTool( + core_tool, auth_config=auth_config, telemetry_attributes=telemetry_getter + ) + + ctx = MagicMock() + mock_creds = MagicMock() + mock_creds.oauth2.access_token = "valid_access_token" + mock_creds.oauth2.id_token = "valid_id_token" + ctx.get_auth_response.return_value = mock_creds + + mock_cred_service = MagicMock() + mock_cred_service.load_credential = AsyncMock(return_value=None) + mock_cred_service.save_credential = AsyncMock(return_value=None) + ctx._invocation_context = MagicMock() + ctx._invocation_context.credential_service = mock_cred_service + + previous_token = USER_TOKEN_CONTEXT_VAR.set("previous_token") + try: + with pytest.raises(RuntimeError, match="telemetry boom"): + await tool.run_async({}, ctx) + + assert USER_TOKEN_CONTEXT_VAR.get() == "previous_token" + core_tool.assert_not_awaited() + finally: + USER_TOKEN_CONTEXT_VAR.reset(previous_token) + @pytest.mark.asyncio async def test_3lo_exception_reraise(self): # Test that specific credential errors are re-raised diff --git a/packages/toolbox-adk/tests/unit/test_toolset.py b/packages/toolbox-adk/tests/unit/test_toolset.py index ff632e1ea..9c9a9b707 100644 --- a/packages/toolbox-adk/tests/unit/test_toolset.py +++ b/packages/toolbox-adk/tests/unit/test_toolset.py @@ -16,7 +16,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from toolbox_core.protocol import Protocol +from toolbox_core.protocol import Protocol, TelemetryAttributes from toolbox_adk.tool import ToolboxTool from toolbox_adk.toolset import ToolboxToolset @@ -76,6 +76,28 @@ async def test_get_tools_with_auth_token_getters(self, mock_client_cls): mock_client.load_tool.assert_awaited_with("toolA", bound_params={}) assert tools[0]._adk_token_getters == auth_getters + @patch("toolbox_adk.toolset.ToolboxClient") + @pytest.mark.asyncio + async def test_get_tools_passes_telemetry_attributes_to_wrappers( + self, mock_client_cls + ): + mock_client = mock_client_cls.return_value + + t1 = MagicMock() + t1.__name__ = "tool1" + t1.__doc__ = "desc1" + mock_client.load_tool = AsyncMock(return_value=t1) + + attrs = TelemetryAttributes(llm_model="gemini-2.5-pro") + toolset = ToolboxToolset( + "url", tool_names=["toolA"], telemetry_attributes=attrs + ) + + tools = await toolset.get_tools() + + assert len(tools) == 1 + assert tools[0]._telemetry_attributes is attrs + @patch("toolbox_adk.toolset.ToolboxClient") @pytest.mark.asyncio @pytest.mark.parametrize( From 18582910ef1f64666d44f75b396aaaba66b10ead Mon Sep 17 00:00:00 2001 From: Parth Ajmera Date: Mon, 20 Jul 2026 23:44:40 -0700 Subject: [PATCH 2/3] fix(adk): narrow telemetry getter callables --- packages/toolbox-adk/src/toolbox_adk/tool.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index 9e27e0807..d7bc88b83 100644 --- a/packages/toolbox-adk/src/toolbox_adk/tool.py +++ b/packages/toolbox-adk/src/toolbox_adk/tool.py @@ -14,7 +14,7 @@ import inspect import logging -from typing import Any, Awaitable, Callable, Dict, Mapping, Optional, Union +from typing import Any, Awaitable, Callable, Dict, Mapping, Optional, Union, cast import toolbox_core from fastapi.openapi.models import OAuth2, OAuthFlowAuthorizationCode, OAuthFlows @@ -345,9 +345,23 @@ async def _resolve_telemetry_attributes( sig = inspect.signature(attributes) if len(sig.parameters) == 1: - resolved = attributes(tool_context) + context_getter = cast( + Union[ + Callable[[ToolContext], TelemetryAttributes], + Callable[[ToolContext], Awaitable[TelemetryAttributes]], + ], + attributes, + ) + resolved = context_getter(tool_context) else: - resolved = attributes() + no_arg_getter = cast( + Union[ + Callable[[], TelemetryAttributes], + Callable[[], Awaitable[TelemetryAttributes]], + ], + attributes, + ) + resolved = no_arg_getter() if inspect.isawaitable(resolved): resolved = await resolved return resolved From 611a197c4a90a1945ec2e86e4964c91ce493b66e Mon Sep 17 00:00:00 2001 From: Parth Ajmera Date: Tue, 21 Jul 2026 01:22:57 -0700 Subject: [PATCH 3/3] fix(adk): handle telemetry getter review cases --- packages/toolbox-adk/src/toolbox_adk/tool.py | 23 ++++++++- packages/toolbox-adk/tests/unit/test_tool.py | 47 ++++++++++++++++++- .../toolbox-adk/tests/unit/test_toolset.py | 2 +- 3 files changed, 67 insertions(+), 5 deletions(-) diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index d7bc88b83..ae9eca015 100644 --- a/packages/toolbox-adk/src/toolbox_adk/tool.py +++ b/packages/toolbox-adk/src/toolbox_adk/tool.py @@ -344,7 +344,7 @@ async def _resolve_telemetry_attributes( return attributes sig = inspect.signature(attributes) - if len(sig.parameters) == 1: + if self._accepts_tool_context_arg(sig, tool_context): context_getter = cast( Union[ Callable[[ToolContext], TelemetryAttributes], @@ -364,7 +364,26 @@ async def _resolve_telemetry_attributes( resolved = no_arg_getter() if inspect.isawaitable(resolved): resolved = await resolved - return resolved + return self._validate_resolved_telemetry_attributes(resolved) + + def _validate_resolved_telemetry_attributes( + self, resolved: Any + ) -> Optional[TelemetryAttributes]: + if resolved is None or isinstance(resolved, TelemetryAttributes): + return resolved + raise TypeError( + "telemetry_attributes callable must return TelemetryAttributes or None" + ) + + def _accepts_tool_context_arg( + self, sig: inspect.Signature, tool_context: ToolContext + ) -> bool: + """Returns whether a callable signature can receive ToolContext.""" + try: + sig.bind(tool_context) + return True + except TypeError: + return False def bind_params(self, bounded_params: Dict[str, Any]) -> "ToolboxTool": """Allows runtime binding of parameters, delegating to core tool.""" diff --git a/packages/toolbox-adk/tests/unit/test_tool.py b/packages/toolbox-adk/tests/unit/test_tool.py index d944f6c85..b507c3498 100644 --- a/packages/toolbox-adk/tests/unit/test_tool.py +++ b/packages/toolbox-adk/tests/unit/test_tool.py @@ -141,7 +141,7 @@ def add_telemetry_attributes(attrs): def telemetry_getter(ctx): return TelemetryAttributes( - llm_model="gemini-2.5-pro", + llm_model="gemini-3.5-flash", user_id=ctx.state["user_id"], agent_id="agent-1", ) @@ -158,11 +158,54 @@ def telemetry_getter(ctx): calls = core_tool.add_telemetry_attributes.call_args_list assert [c.args[0].user_id for c in calls] == ["user-1", "user-2"] - assert all(c.args[0].llm_model == "gemini-2.5-pro" for c in calls) + assert all(c.args[0].llm_model == "gemini-3.5-flash" for c in calls) assert core_tool.await_count == 0 assert [t.await_count for t in derived_tools] == [1, 1] assert tool._core_tool is core_tool + @pytest.mark.asyncio + async def test_telemetry_getter_with_default_argument_receives_tool_context(self): + core_tool = AsyncMock() + core_tool.__name__ = "mock" + core_tool.__doc__ = "mock doc" + + derived_tool = AsyncMock(return_value="ok") + derived_tool.__name__ = "mock" + derived_tool.__doc__ = "mock doc" + core_tool.add_telemetry_attributes = MagicMock(return_value=derived_tool) + + def telemetry_getter(ctx, default_model="gemini-3.5-flash"): + return TelemetryAttributes( + llm_model=default_model, + user_id=ctx.state["user_id"], + agent_id="agent-1", + ) + + tool = ToolboxTool(core_tool, telemetry_attributes=telemetry_getter) + + ctx = MagicMock() + ctx.state = {"user_id": "user-1"} + + assert await tool.run_async({}, ctx) == "ok" + + attrs = core_tool.add_telemetry_attributes.call_args.args[0] + assert attrs.llm_model == "gemini-3.5-flash" + assert attrs.user_id == "user-1" + + @pytest.mark.asyncio + async def test_invalid_telemetry_getter_return_type_raises(self): + core_tool = AsyncMock() + core_tool.__name__ = "mock" + core_tool.__doc__ = "mock doc" + + tool = ToolboxTool(core_tool, telemetry_attributes=lambda: {"user_id": "u1"}) + + with pytest.raises( + TypeError, + match="telemetry_attributes callable must return TelemetryAttributes or None", + ): + await tool.run_async({}, MagicMock()) + @pytest.mark.asyncio async def test_3lo_missing_client_secret(self): # Test ValueError when client_id/secret missing diff --git a/packages/toolbox-adk/tests/unit/test_toolset.py b/packages/toolbox-adk/tests/unit/test_toolset.py index 9c9a9b707..440a8b5c2 100644 --- a/packages/toolbox-adk/tests/unit/test_toolset.py +++ b/packages/toolbox-adk/tests/unit/test_toolset.py @@ -88,7 +88,7 @@ async def test_get_tools_passes_telemetry_attributes_to_wrappers( t1.__doc__ = "desc1" mock_client.load_tool = AsyncMock(return_value=t1) - attrs = TelemetryAttributes(llm_model="gemini-2.5-pro") + attrs = TelemetryAttributes(llm_model="gemini-3.5-flash") toolset = ToolboxToolset( "url", tool_names=["toolA"], telemetry_attributes=attrs )