feat(adk): pass telemetry attributes through toolbox-adk

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, cast
 
 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,61 @@ 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 self._accepts_tool_context_arg(sig, tool_context):
+            context_getter = cast(
+                Union[
+                    Callable[[ToolContext], TelemetryAttributes],
+                    Callable[[ToolContext], Awaitable[TelemetryAttributes]],
+                ],
+                attributes,
+            )
+            resolved = context_getter(tool_context)
+        else:
+            no_arg_getter = cast(
+                Union[
+                    Callable[[], TelemetryAttributes],
+                    Callable[[], Awaitable[TelemetryAttributes]],
+                ],
+                attributes,
+            )
+            resolved = no_arg_getter()
+        if inspect.isawaitable(resolved):
+            resolved = await 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."""
         new_core_tool = self._core_tool.bind_params(bounded_params)
@@ -307,4 +393,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,
         )

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
         ]

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,92 @@ 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-3.5-flash",
+                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-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
@@ -230,6 +318,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

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-3.5-flash")
+        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(