From f230d7a61d281bc3d6db950099bb6988287516fe Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 6 Mar 2026 19:16:56 +0530 Subject: [PATCH 1/8] feat(adk): Support ToolContext for aut_token_getters in toolbox-adk --- packages/toolbox-adk/README.md | 8 +++- packages/toolbox-adk/src/toolbox_adk/tool.py | 25 ++++++++++- .../toolbox-adk/src/toolbox_adk/toolset.py | 14 +++++-- packages/toolbox-adk/tests/unit/test_tool.py | 42 +++++++++++++++++++ .../toolbox-adk/tests/unit/test_toolset.py | 13 ++---- 5 files changed, 86 insertions(+), 16 deletions(-) diff --git a/packages/toolbox-adk/README.md b/packages/toolbox-adk/README.md index bc9b3ea7..581814f1 100644 --- a/packages/toolbox-adk/README.md +++ b/packages/toolbox-adk/README.md @@ -210,6 +210,9 @@ creds = CredentialStrategy.from_adk_credentials(auth_credential, scheme) Some tools may define their own authentication requirements (e.g., Salesforce OAuth, GitHub PAT) via `authSources` in their schema. You can provide a mapping of getters to resolve these tokens at runtime. +> [!TIP] +> Getters can optionally accept the ADK `ToolContext` as a single argument. This enables seamless integration of dynamic, end-user tokens that are tied to the current agent execution state. + ```python async def get_salesforce_token(): # Fetch token from secret manager or reliable source @@ -218,8 +221,9 @@ async def get_salesforce_token(): toolset = ToolboxToolset( server_url="...", auth_token_getters={ - "salesforce-auth": get_salesforce_token, # Async callable - "github-pat": lambda: "my-pat-token" # Sync callable or static lambda + "salesforce-auth": get_salesforce_token, # Async callable + "github-pat": lambda: "my-pat-token", # Sync callable or static lambda + "oauth-user": lambda ctx: ctx.state.get("auth_token") # Dynamic context-aware callable } ) ``` diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index 720407bf..992e6437 100644 --- a/packages/toolbox-adk/src/toolbox_adk/tool.py +++ b/packages/toolbox-adk/src/toolbox_adk/tool.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import inspect import logging from typing import Any, Awaitable, Callable, Dict, Optional @@ -68,11 +69,13 @@ def __init__( self, core_tool: CoreToolboxTool, auth_config: Optional[CredentialConfig] = None, + adk_token_getters: Optional[Mapping[str, Any]] = 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. """ # We act as a proxy. # We need to extract metadata from the core tool to satisfy BaseTool's contract. @@ -95,6 +98,7 @@ def __init__( ) self._core_tool = core_tool self._auth_config = auth_config + self._adk_token_getters = adk_token_getters or {} def _param_type_to_schema_type(self, param_type: str) -> Type: type_map = { @@ -260,12 +264,30 @@ async def run_async( "Falling back to request_credential.", exc_info=True, ) - # Fallback to request logic tool_context.request_credential(auth_config_adk) return { "error": f"OAuth2 Credentials required for {self.name}. A consent link has been generated for the user. Do NOT attempt to run this tool again until the user confirms they have logged in." } + if self._adk_token_getters: + needed_services = set( + list(self._core_tool._required_authn_params.values()) + + list(self._core_tool._required_authz_tokens) + ) + + for service, getter in self._adk_token_getters.items(): + if service in needed_services: + sig = inspect.signature(getter) + + if len(sig.parameters) == 1: + bound_getter = lambda t=getter, ctx=tool_context: t(ctx) + else: + bound_getter = getter + + self._core_tool = self._core_tool.add_auth_token_getter( + service, bound_getter + ) + result: Optional[Any] = None error: Optional[Exception] = None @@ -288,4 +310,5 @@ def bind_params(self, bounded_params: Dict[str, Any]) -> "ToolboxTool": return ToolboxTool( core_tool=new_core_tool, auth_config=self._auth_config, + adk_token_getters=self._adk_token_getters, ) diff --git a/packages/toolbox-adk/src/toolbox_adk/toolset.py b/packages/toolbox-adk/src/toolbox_adk/toolset.py index 8c7c340f..763de2a7 100644 --- a/packages/toolbox-adk/src/toolbox_adk/toolset.py +++ b/packages/toolbox-adk/src/toolbox_adk/toolset.py @@ -41,7 +41,15 @@ def __init__( ] = None, bound_params: Optional[Mapping[str, Union[Callable[[], Any], Any]]] = None, auth_token_getters: Optional[ - Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]] + Mapping[ + str, + Union[ + Callable[[], str], + Callable[[], Awaitable[str]], + Callable[[ToolContext], str], + Callable[[ToolContext], Awaitable[str]], + ], + ] ] = None, **kwargs: Any, ): @@ -91,7 +99,6 @@ async def get_tools( core_tools = await self.client.load_toolset( self.__toolset_name, bound_params=self.__bound_params or {}, - auth_token_getters=self.__auth_token_getters or {}, ) tools.extend(core_tools) @@ -101,7 +108,6 @@ async def get_tools( core_tool = await self.client.load_tool( name, bound_params=self.__bound_params or {}, - auth_token_getters=self.__auth_token_getters or {}, ) tools.append(core_tool) @@ -110,7 +116,6 @@ async def get_tools( core_tools = await self.client.load_toolset( None, bound_params=self.__bound_params or {}, - auth_token_getters=self.__auth_token_getters or {}, ) tools.extend(core_tools) @@ -119,6 +124,7 @@ async def get_tools( ToolboxTool( core_tool=t, auth_config=self.client.credential_config, + adk_token_getters=self.__auth_token_getters, ) 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 0f3c5120..d04d0161 100644 --- a/packages/toolbox-adk/tests/unit/test_tool.py +++ b/packages/toolbox-adk/tests/unit/test_tool.py @@ -76,6 +76,48 @@ async def test_bind_params(self): assert new_tool._core_tool == new_core_mock mock_core.bind_params.assert_called_with({"a": 1}) + @pytest.mark.asyncio + async def test_dynamic_adk_token_getters(self): + core_tool = AsyncMock() + core_tool.__name__ = "mock" + core_tool.__doc__ = "mock doc" + core_tool._required_authn_params = {"param1": "service1"} + core_tool._required_authz_tokens = ["service2"] + core_tool.add_auth_token_getter = MagicMock(return_value=core_tool) + + def getter1(): + return "token1" + + def getter2(ctx): + return ctx.state.get("token2") + + adk_getters = { + "service1": getter1, + "service2": getter2, + } + + tool = ToolboxTool(core_tool, adk_token_getters=adk_getters) + + ctx = MagicMock() + ctx.state = {"token2": "dynamic_token2"} + + await tool.run_async({}, ctx) + + assert core_tool.add_auth_token_getter.call_count == 2 + + args1 = core_tool.add_auth_token_getter.call_args_list[0][0] + args2 = core_tool.add_auth_token_getter.call_args_list[1][0] + + # Because we iterate over items(), order might be dependent. + # Check that both services were processed and bound correctly + bound_getters = {args1[0]: args1[1], args2[0]: args2[1]} + + assert "service1" in bound_getters + assert bound_getters["service1"]() == "token1" + + assert "service2" in bound_getters + assert bound_getters["service2"]() == "dynamic_token2" + @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 54962f5e..143af0bc 100644 --- a/packages/toolbox-adk/tests/unit/test_toolset.py +++ b/packages/toolbox-adk/tests/unit/test_toolset.py @@ -49,12 +49,8 @@ async def test_get_tools_load_set_and_list(self, mock_client_cls): assert isinstance(tools[0], ToolboxTool) assert isinstance(tools[1], ToolboxTool) - mock_client.load_toolset.assert_awaited_with( - "set1", bound_params={"p": 1}, auth_token_getters={} - ) - mock_client.load_tool.assert_awaited_with( - "toolA", bound_params={"p": 1}, auth_token_getters={} - ) + mock_client.load_toolset.assert_awaited_with("set1", bound_params={"p": 1}) + mock_client.load_tool.assert_awaited_with("toolA", bound_params={"p": 1}) @patch("toolbox_adk.toolset.ToolboxClient") @pytest.mark.asyncio @@ -75,9 +71,8 @@ async def test_get_tools_with_auth_token_getters(self, mock_client_cls): tools = await toolset.get_tools() assert len(tools) == 1 - mock_client.load_tool.assert_awaited_with( - "toolA", bound_params={}, auth_token_getters=auth_getters - ) + 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 From d41d46ece36cd9961c625d9f38786388cecfd87c Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 6 Mar 2026 20:05:46 +0530 Subject: [PATCH 2/8] chore: Add clarifying comment --- packages/toolbox-adk/src/toolbox_adk/tool.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index 992e6437..634216f3 100644 --- a/packages/toolbox-adk/src/toolbox_adk/tool.py +++ b/packages/toolbox-adk/src/toolbox_adk/tool.py @@ -270,6 +270,8 @@ async def run_async( } if self._adk_token_getters: + # Pre-filter toolset getters to avoid unused-token errors from the core tool. + # This deferred loop also enables dynamic 1-arity `tool_context` injection. needed_services = set( list(self._core_tool._required_authn_params.values()) + list(self._core_tool._required_authz_tokens) From b6c6027add1c51687d01fbb824257e5c0eb5b1bd Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 6 Mar 2026 20:10:46 +0530 Subject: [PATCH 3/8] fix: fix lint and tests --- packages/toolbox-adk/src/toolbox_adk/tool.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index 634216f3..314bb271 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, Optional +from typing import Any, Awaitable, Callable, Dict, Mapping, Optional import google.adk.auth.exchanger.oauth2_credential_exchanger as oauth2_credential_exchanger import google.adk.auth.oauth2_credential_util as oauth2_credential_util From cb0492b8480dd50c2dfbd5a44f1ff96d1df29d9e Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 6 Mar 2026 21:02:53 +0530 Subject: [PATCH 4/8] chore: refactor tool strict validation logic --- .../toolbox-adk/src/toolbox_adk/toolset.py | 24 ++++++ .../tests/integration/test_integration.py | 40 +++++++++- .../toolbox-adk/tests/unit/test_toolset.py | 40 ++++++++++ .../toolbox-core/src/toolbox_core/client.py | 76 ++++++++----------- .../toolbox-core/src/toolbox_core/utils.py | 42 ++++++++++ 5 files changed, 175 insertions(+), 47 deletions(-) diff --git a/packages/toolbox-adk/src/toolbox_adk/toolset.py b/packages/toolbox-adk/src/toolbox_adk/toolset.py index 763de2a7..7db52c81 100644 --- a/packages/toolbox-adk/src/toolbox_adk/toolset.py +++ b/packages/toolbox-adk/src/toolbox_adk/toolset.py @@ -23,6 +23,7 @@ from .client import ToolboxClient from .credentials import CredentialConfig from .tool import ToolboxTool +from toolbox_core.utils import validate_unused_requirements class ToolboxToolset(BaseToolset): @@ -119,6 +120,29 @@ async def get_tools( ) tools.extend(core_tools) + # 4. Strictly validate unused toolset auth token getters using core logic + if self.__auth_token_getters: + overall_used_auth_keys = set() + for t in tools: + for reqs in t._required_authn_params.values(): + overall_used_auth_keys.update(reqs) + overall_used_auth_keys.update(t._required_authz_tokens) + + # Generate intuitive name for the error string if a specific toolset wasn't used + validation_name = self.__toolset_name + if not validation_name: + validation_name = ", ".join(self.__tool_names) if self.__tool_names else "default" + + validate_unused_requirements( + provided_auth_keys=set(self.__auth_token_getters.keys()), + provided_bound_keys=set(), + used_auth_keys=overall_used_auth_keys, + used_bound_keys=set(), + name=validation_name, + is_toolset=True, + target_type="list of tools" if not self.__toolset_name else None, + ) + # Wrap all core tools in ToolboxTool return [ ToolboxTool( diff --git a/packages/toolbox-adk/tests/integration/test_integration.py b/packages/toolbox-adk/tests/integration/test_integration.py index 524df61e..dd6e5897 100644 --- a/packages/toolbox-adk/tests/integration/test_integration.py +++ b/packages/toolbox-adk/tests/integration/test_integration.py @@ -545,7 +545,45 @@ async def test_run_tool_unauth_with_auth(self, auth_token2: str): try: with pytest.raises( ValueError, - match=rf"Validation failed for tool 'get-row-by-id': unused auth tokens: my-test-auth", + match=rf"Validation failed for list of tools 'get-row-by-id': unused auth tokens could not be applied to any tool: my-test-auth", + ): + await toolset.get_tools() + finally: + await toolset.close() + + async def test_run_multiple_tools_unauth_with_auth(self, auth_token2: str): + """Tests running multiple tools that don't require auth, verifying formatting of tool lists.""" + toolset = ToolboxToolset( + server_url="http://localhost:5000", + tool_names=["get-row-by-id", "search-rows"], + auth_token_getters={"my-test-auth": lambda: auth_token2}, + credentials=CredentialStrategy.toolbox_identity(), + ) + try: + with pytest.raises( + ValueError, + match=rf"Validation failed for list of tools 'get-row-by-id, search-rows': unused auth tokens could not be applied to any tool: my-test-auth", + ): + await toolset.get_tools() + finally: + await toolset.close() + + async def test_run_multiple_tools_partial_auth_usage(self, auth_token2: str): + """Tests that when some tokens are used and some aren't across diverse tools, only the truly unused tokens appear in the error.""" + toolset = ToolboxToolset( + server_url="http://localhost:5000", + tool_names=["get-row-by-id-auth", "search-rows"], # first requires 'my-test-auth', second requires nothing + auth_token_getters={ + "my-test-auth": lambda: auth_token2, + "extra-token": lambda: "fake" + }, + credentials=CredentialStrategy.toolbox_identity(), + ) + try: + with pytest.raises( + ValueError, + # 'my-test-auth' should be cleanly consumed and absent from the final error string. + match=r"Validation failed for list of tools 'get-row-by-id-auth, search-rows': unused auth tokens could not be applied to any tool: extra-token", ): await toolset.get_tools() finally: diff --git a/packages/toolbox-adk/tests/unit/test_toolset.py b/packages/toolbox-adk/tests/unit/test_toolset.py index 143af0bc..537345a7 100644 --- a/packages/toolbox-adk/tests/unit/test_toolset.py +++ b/packages/toolbox-adk/tests/unit/test_toolset.py @@ -61,6 +61,8 @@ async def test_get_tools_with_auth_token_getters(self, mock_client_cls): t1 = MagicMock() t1.__name__ = "tool1" t1.__doc__ = "desc1" + t1._required_authn_params = {"param1": ["service"]} + t1._required_authz_tokens = [] mock_client.load_tool = AsyncMock(return_value=t1) auth_getters = {"service": lambda: "token"} @@ -74,6 +76,44 @@ 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 + @pytest.mark.parametrize( + "authn,authz,should_raise", + [ + ({}, [], True), # No requirements, token is completely unused + ({"param1": ["service"]}, [], False), # authn natively consumes it + ({}, ["service"], False), # authz natively consumes it + ({"param1": ["other"]}, ["service"], False), # unused by authn, but authz consumes it + ({"param1": ["service"]}, ["other"], False), # authn consumes it, authz doesn't + ({"param1": ["other"]}, ["other"], True), # Requirements exist, but token is unused by both + ], + ) + async def test_get_tools_auth_validation( + self, mock_client_cls, authn, authz, should_raise + ): + mock_client = mock_client_cls.return_value + + t1 = MagicMock() + t1.__name__ = "tool1" + t1._required_authn_params = authn + t1._required_authz_tokens = authz + mock_client.load_tool = AsyncMock(return_value=t1) + + auth_getters = {"service": lambda: "token"} + toolset = ToolboxToolset( + "url", tool_names=["toolA"], auth_token_getters=auth_getters + ) + + if should_raise: + with pytest.raises( + ValueError, match="unused auth tokens could not be applied to any tool: service" + ): + await toolset.get_tools() + else: + tools = await toolset.get_tools() + assert len(tools) == 1 + @patch("toolbox_adk.toolset.ToolboxClient") @pytest.mark.asyncio async def test_close(self, mock_client_cls): diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 79e598e6..d3a0009d 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -30,7 +30,12 @@ ) from .protocol import Protocol, ToolSchema from .tool import ToolboxTool -from .utils import identify_auth_requirements, resolve_value, warn_if_http_and_headers +from .utils import ( + identify_auth_requirements, + resolve_value, + validate_unused_requirements, + warn_if_http_and_headers, +) class ToolboxClient: @@ -238,20 +243,14 @@ async def load_tool( provided_auth_keys = set(auth_token_getters.keys()) provided_bound_keys = set(bound_params.keys()) - unused_auth = provided_auth_keys - used_auth_keys - unused_bound = provided_bound_keys - used_bound_keys - - if unused_auth or unused_bound: - error_messages = [] - if unused_auth: - error_messages.append(f"unused auth tokens: {', '.join(unused_auth)}") - if unused_bound: - error_messages.append( - f"unused bound parameters: {', '.join(unused_bound)}" - ) - raise ValueError( - f"Validation failed for tool '{name}': { '; '.join(error_messages) }." - ) + validate_unused_requirements( + provided_auth_keys, + provided_bound_keys, + used_auth_keys, + used_bound_keys, + name, + is_toolset=False, + ) return tool @@ -318,41 +317,26 @@ async def load_toolset( tools.append(tool) if strict: - unused_auth = provided_auth_keys - used_auth_keys - unused_bound = provided_bound_keys - used_bound_keys - if unused_auth or unused_bound: - error_messages = [] - if unused_auth: - error_messages.append( - f"unused auth tokens: {', '.join(unused_auth)}" - ) - if unused_bound: - error_messages.append( - f"unused bound parameters: {', '.join(unused_bound)}" - ) - raise ValueError( - f"Validation failed for tool '{tool_name}': { '; '.join(error_messages) }." - ) + validate_unused_requirements( + provided_auth_keys, + provided_bound_keys, + used_auth_keys, + used_bound_keys, + tool_name, + is_toolset=False, + ) else: overall_used_auth_keys.update(used_auth_keys) overall_used_bound_params.update(used_bound_keys) - unused_auth = provided_auth_keys - overall_used_auth_keys - unused_bound = provided_bound_keys - overall_used_bound_params - - if unused_auth or unused_bound: - error_messages = [] - if unused_auth: - error_messages.append( - f"unused auth tokens could not be applied to any tool: {', '.join(unused_auth)}" - ) - if unused_bound: - error_messages.append( - f"unused bound parameters could not be applied to any tool: {', '.join(unused_bound)}" - ) - raise ValueError( - f"Validation failed for toolset '{name or 'default'}': { '; '.join(error_messages) }." - ) + validate_unused_requirements( + provided_auth_keys, + provided_bound_keys, + overall_used_auth_keys, + overall_used_bound_params, + name or "default", + is_toolset=True, + ) return tools diff --git a/packages/toolbox-core/src/toolbox_core/utils.py b/packages/toolbox-core/src/toolbox_core/utils.py index 00c00157..2ed43c8a 100644 --- a/packages/toolbox-core/src/toolbox_core/utils.py +++ b/packages/toolbox-core/src/toolbox_core/utils.py @@ -163,3 +163,45 @@ async def resolve_value( elif callable(source): return source() return source + + +def validate_unused_requirements( + provided_auth_keys: set[str], + provided_bound_keys: set[str], + used_auth_keys: set[str], + used_bound_keys: set[str], + name: str, + is_toolset: bool = False, + target_type: str | None = None, +) -> None: + """ + Validates that no provided authentication tokens or bound parameters went unused. + Raises a ValueError if any unused requirements are found, formatted appropriately + for either a single tool or a full toolset. + """ + unused_auth = provided_auth_keys - used_auth_keys + unused_bound = provided_bound_keys - used_bound_keys + + if unused_auth or unused_bound: + error_messages = [] + if unused_auth: + if is_toolset: + error_messages.append( + f"unused auth tokens could not be applied to any tool: {', '.join(unused_auth)}" + ) + else: + error_messages.append(f"unused auth tokens: {', '.join(unused_auth)}") + if unused_bound: + if is_toolset: + error_messages.append( + f"unused bound parameters could not be applied to any tool: {', '.join(unused_bound)}" + ) + else: + error_messages.append( + f"unused bound parameters: {', '.join(unused_bound)}" + ) + + final_target_type = target_type if target_type else ("toolset" if is_toolset else "tool") + raise ValueError( + f"Validation failed for {final_target_type} '{name}': {'; '.join(error_messages)}." + ) From 9233619de5d7803860dac67fd3d230b77439710f Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 6 Mar 2026 21:03:06 +0530 Subject: [PATCH 5/8] chore: delint --- .../toolbox-adk/src/toolbox_adk/toolset.py | 6 ++++-- .../tests/integration/test_integration.py | 9 +++++--- .../toolbox-adk/tests/unit/test_toolset.py | 21 +++++++++++++++---- .../toolbox-core/src/toolbox_core/utils.py | 4 +++- 4 files changed, 30 insertions(+), 10 deletions(-) diff --git a/packages/toolbox-adk/src/toolbox_adk/toolset.py b/packages/toolbox-adk/src/toolbox_adk/toolset.py index 7db52c81..7689f96d 100644 --- a/packages/toolbox-adk/src/toolbox_adk/toolset.py +++ b/packages/toolbox-adk/src/toolbox_adk/toolset.py @@ -18,12 +18,12 @@ 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.utils import validate_unused_requirements from typing_extensions import override from .client import ToolboxClient from .credentials import CredentialConfig from .tool import ToolboxTool -from toolbox_core.utils import validate_unused_requirements class ToolboxToolset(BaseToolset): @@ -131,7 +131,9 @@ async def get_tools( # Generate intuitive name for the error string if a specific toolset wasn't used validation_name = self.__toolset_name if not validation_name: - validation_name = ", ".join(self.__tool_names) if self.__tool_names else "default" + validation_name = ( + ", ".join(self.__tool_names) if self.__tool_names else "default" + ) validate_unused_requirements( provided_auth_keys=set(self.__auth_token_getters.keys()), diff --git a/packages/toolbox-adk/tests/integration/test_integration.py b/packages/toolbox-adk/tests/integration/test_integration.py index dd6e5897..6fed882e 100644 --- a/packages/toolbox-adk/tests/integration/test_integration.py +++ b/packages/toolbox-adk/tests/integration/test_integration.py @@ -572,10 +572,13 @@ async def test_run_multiple_tools_partial_auth_usage(self, auth_token2: str): """Tests that when some tokens are used and some aren't across diverse tools, only the truly unused tokens appear in the error.""" toolset = ToolboxToolset( server_url="http://localhost:5000", - tool_names=["get-row-by-id-auth", "search-rows"], # first requires 'my-test-auth', second requires nothing + tool_names=[ + "get-row-by-id-auth", + "search-rows", + ], # first requires 'my-test-auth', second requires nothing auth_token_getters={ - "my-test-auth": lambda: auth_token2, - "extra-token": lambda: "fake" + "my-test-auth": lambda: auth_token2, + "extra-token": lambda: "fake", }, credentials=CredentialStrategy.toolbox_identity(), ) diff --git a/packages/toolbox-adk/tests/unit/test_toolset.py b/packages/toolbox-adk/tests/unit/test_toolset.py index 537345a7..ff632e1e 100644 --- a/packages/toolbox-adk/tests/unit/test_toolset.py +++ b/packages/toolbox-adk/tests/unit/test_toolset.py @@ -84,9 +84,21 @@ async def test_get_tools_with_auth_token_getters(self, mock_client_cls): ({}, [], True), # No requirements, token is completely unused ({"param1": ["service"]}, [], False), # authn natively consumes it ({}, ["service"], False), # authz natively consumes it - ({"param1": ["other"]}, ["service"], False), # unused by authn, but authz consumes it - ({"param1": ["service"]}, ["other"], False), # authn consumes it, authz doesn't - ({"param1": ["other"]}, ["other"], True), # Requirements exist, but token is unused by both + ( + {"param1": ["other"]}, + ["service"], + False, + ), # unused by authn, but authz consumes it + ( + {"param1": ["service"]}, + ["other"], + False, + ), # authn consumes it, authz doesn't + ( + {"param1": ["other"]}, + ["other"], + True, + ), # Requirements exist, but token is unused by both ], ) async def test_get_tools_auth_validation( @@ -107,7 +119,8 @@ async def test_get_tools_auth_validation( if should_raise: with pytest.raises( - ValueError, match="unused auth tokens could not be applied to any tool: service" + ValueError, + match="unused auth tokens could not be applied to any tool: service", ): await toolset.get_tools() else: diff --git a/packages/toolbox-core/src/toolbox_core/utils.py b/packages/toolbox-core/src/toolbox_core/utils.py index 2ed43c8a..fabde6c1 100644 --- a/packages/toolbox-core/src/toolbox_core/utils.py +++ b/packages/toolbox-core/src/toolbox_core/utils.py @@ -201,7 +201,9 @@ def validate_unused_requirements( f"unused bound parameters: {', '.join(unused_bound)}" ) - final_target_type = target_type if target_type else ("toolset" if is_toolset else "tool") + final_target_type = ( + target_type if target_type else ("toolset" if is_toolset else "tool") + ) raise ValueError( f"Validation failed for {final_target_type} '{name}': {'; '.join(error_messages)}." ) From 5da8f218862171d10af9778daafb0ecab02c6907 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 6 Mar 2026 21:17:52 +0530 Subject: [PATCH 6/8] chore: delint --- packages/toolbox-adk/src/toolbox_adk/tool.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index 314bb271..59e813b2 100644 --- a/packages/toolbox-adk/src/toolbox_adk/tool.py +++ b/packages/toolbox-adk/src/toolbox_adk/tool.py @@ -272,10 +272,10 @@ async def run_async( if self._adk_token_getters: # Pre-filter toolset getters to avoid unused-token errors from the core tool. # This deferred loop also enables dynamic 1-arity `tool_context` injection. - needed_services = set( - list(self._core_tool._required_authn_params.values()) - + list(self._core_tool._required_authz_tokens) - ) + needed_services = set() + for reqs in self._core_tool._required_authn_params.values(): + needed_services.update(reqs) + needed_services.update(self._core_tool._required_authz_tokens) for service, getter in self._adk_token_getters.items(): if service in needed_services: From c406b599582dfbf2f037e0b270aa237d8757f483 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 6 Mar 2026 22:18:51 +0530 Subject: [PATCH 7/8] chore: fix tests and type check --- packages/toolbox-adk/tests/unit/test_tool.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/toolbox-adk/tests/unit/test_tool.py b/packages/toolbox-adk/tests/unit/test_tool.py index d04d0161..2d0fb9d5 100644 --- a/packages/toolbox-adk/tests/unit/test_tool.py +++ b/packages/toolbox-adk/tests/unit/test_tool.py @@ -81,7 +81,7 @@ async def test_dynamic_adk_token_getters(self): core_tool = AsyncMock() core_tool.__name__ = "mock" core_tool.__doc__ = "mock doc" - core_tool._required_authn_params = {"param1": "service1"} + core_tool._required_authn_params = {"param1": ["service1"]} core_tool._required_authz_tokens = ["service2"] core_tool.add_auth_token_getter = MagicMock(return_value=core_tool) From 7f35c0cb90ebb391ae810c684d09719e3506b8d3 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 6 Mar 2026 22:28:54 +0530 Subject: [PATCH 8/8] fix: fix CI/CD pipeline for toolbox-adk for local toolbox-core --- .github/workflows/lint-toolbox-adk.yaml | 6 +++--- .github/workflows/lint-toolbox-core.yaml | 6 +++--- .github/workflows/lint-toolbox-langchain.yaml | 6 +++--- .github/workflows/lint-toolbox-llamaindex.yaml | 6 +++--- packages/toolbox-adk/integration.cloudbuild.yaml | 3 ++- packages/toolbox-core/integration.cloudbuild.yaml | 3 ++- packages/toolbox-langchain/integration.cloudbuild.yaml | 3 ++- packages/toolbox-llamaindex/integration.cloudbuild.yaml | 3 ++- 8 files changed, 20 insertions(+), 16 deletions(-) diff --git a/.github/workflows/lint-toolbox-adk.yaml b/.github/workflows/lint-toolbox-adk.yaml index 6bb8e3bc..4e41544e 100644 --- a/.github/workflows/lint-toolbox-adk.yaml +++ b/.github/workflows/lint-toolbox-adk.yaml @@ -67,12 +67,12 @@ jobs: with: python-version: "3.13" + - name: Install test requirements + run: pip install -e .[test] + - name: Install library requirements run: pip install -r requirements.txt - - name: Install test requirements - run: pip install .[test] - - name: Run linters run: | black --check . diff --git a/.github/workflows/lint-toolbox-core.yaml b/.github/workflows/lint-toolbox-core.yaml index 05994176..100f2e8e 100644 --- a/.github/workflows/lint-toolbox-core.yaml +++ b/.github/workflows/lint-toolbox-core.yaml @@ -67,12 +67,12 @@ jobs: with: python-version: "3.13" + - name: Install test requirements + run: pip install -e .[test] + - name: Install library requirements run: pip install -r requirements.txt - - name: Install test requirements - run: pip install .[test] - - name: Run linters run: | black --check . diff --git a/.github/workflows/lint-toolbox-langchain.yaml b/.github/workflows/lint-toolbox-langchain.yaml index 79d46ada..dae263fa 100644 --- a/.github/workflows/lint-toolbox-langchain.yaml +++ b/.github/workflows/lint-toolbox-langchain.yaml @@ -67,12 +67,12 @@ jobs: with: python-version: "3.13" + - name: Install test requirements + run: pip install -e .[test] + - name: Install library requirements run: pip install -r requirements.txt - - name: Install test requirements - run: pip install .[test] - - name: Run linters run: | black --check . diff --git a/.github/workflows/lint-toolbox-llamaindex.yaml b/.github/workflows/lint-toolbox-llamaindex.yaml index c6f12ea5..7523413a 100644 --- a/.github/workflows/lint-toolbox-llamaindex.yaml +++ b/.github/workflows/lint-toolbox-llamaindex.yaml @@ -67,12 +67,12 @@ jobs: with: python-version: "3.13" + - name: Install test requirements + run: pip install -e .[test] + - name: Install library requirements run: pip install -r requirements.txt - - name: Install test requirements - run: pip install .[test] - - name: Run linters run: | black --check . diff --git a/packages/toolbox-adk/integration.cloudbuild.yaml b/packages/toolbox-adk/integration.cloudbuild.yaml index 441b53f1..5b7cddbb 100644 --- a/packages/toolbox-adk/integration.cloudbuild.yaml +++ b/packages/toolbox-adk/integration.cloudbuild.yaml @@ -28,7 +28,8 @@ steps: # Use $$ to escape shell variable for Cloud Build CORE_VERSION=$$(python -c "v={}; exec(open('../toolbox-core/src/toolbox_core/version.py').read(), v); print(v['__version__'])") sed -i "s/toolbox-core==[0-9.]*/toolbox-core==$$CORE_VERSION/g" pyproject.toml - uv pip install -r requirements.txt -e '.[test]' + uv pip install -r requirements.txt + uv pip install -e '.[test]' entrypoint: /bin/bash - id: Run integration tests name: 'python:${_VERSION}' diff --git a/packages/toolbox-core/integration.cloudbuild.yaml b/packages/toolbox-core/integration.cloudbuild.yaml index cc28fc9f..e6dc930d 100644 --- a/packages/toolbox-core/integration.cloudbuild.yaml +++ b/packages/toolbox-core/integration.cloudbuild.yaml @@ -24,7 +24,8 @@ steps: uv venv /workspace/venv source /workspace/venv/bin/activate uv pip install uv - uv pip install -r requirements.txt -e '.[test]' + uv pip install -e '.[test]' + uv pip install -r requirements.txt entrypoint: /bin/bash - id: Run integration tests name: 'python:${_VERSION}' diff --git a/packages/toolbox-langchain/integration.cloudbuild.yaml b/packages/toolbox-langchain/integration.cloudbuild.yaml index 32e87863..ac26005d 100644 --- a/packages/toolbox-langchain/integration.cloudbuild.yaml +++ b/packages/toolbox-langchain/integration.cloudbuild.yaml @@ -28,7 +28,8 @@ steps: # Use $$ to escape shell variable for Cloud Build CORE_VERSION=$$(python -c "v={}; exec(open('../toolbox-core/src/toolbox_core/version.py').read(), v); print(v['__version__'])") sed -i "s/toolbox-core==[0-9.]*/toolbox-core==$$CORE_VERSION/g" pyproject.toml - uv pip install -r requirements.txt -e '.[test]' + uv pip install -r requirements.txt + uv pip install -e '.[test]' entrypoint: /bin/bash - id: Run integration tests name: 'python:${_VERSION}' diff --git a/packages/toolbox-llamaindex/integration.cloudbuild.yaml b/packages/toolbox-llamaindex/integration.cloudbuild.yaml index 5d7b5dc8..d3586bf2 100644 --- a/packages/toolbox-llamaindex/integration.cloudbuild.yaml +++ b/packages/toolbox-llamaindex/integration.cloudbuild.yaml @@ -28,7 +28,8 @@ steps: # Use $$ to escape shell variable for Cloud Build CORE_VERSION=$$(python -c "v={}; exec(open('../toolbox-core/src/toolbox_core/version.py').read(), v); print(v['__version__'])") sed -i "s/toolbox-core==[0-9.]*/toolbox-core==$$CORE_VERSION/g" pyproject.toml - uv pip install -r requirements.txt -e '.[test]' + uv pip install -r requirements.txt + uv pip install -e '.[test]' entrypoint: /bin/bash - id: Run integration tests name: 'python:${_VERSION}'