From 4d432097e62a8bcae05ec378c438c981e51fa39c Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Sat, 7 Feb 2026 03:00:05 +0530 Subject: [PATCH 1/7] fix(adk): resolve OAuth persistence, ID token extraction, and auth flow crashes This PR addresses several authentication bugs within the `toolbox-adk` wrapper that were preventing successful OIDC flows, causing repeated login prompts, and crashing during tool execution. `adk-python` natively drops the `id_token` during OAuth2 exchange (`update_credential_with_tokens`), which is strictly required by the MCP Toolbox for `USER_IDENTITY` authentication. #### Solution Added a temporary monkey patch to hook into the `adk-python` credential update process and explicitly preserve `tokens.get("id_token")`. > [!NOTE] > Tracked by `TODO`s to remove once upstream PR https://github.com/google/adk-python/pull/4402 is merged. Tools repeatedly prompted users to authenticate because the exchanged credentials were not being appropriately persisted across sessions (`exchanged_auth_credential` was left unpopulated/None). Explicitly assigned the newly fetched user credentials to `auth_config_adk.exchanged_auth_credential` and synced them to storage using `tool_context._invocation_context.credential_service.save_credential()`. Passing an empty list of scopes triggered an `AttributeError` deep inside the `adk-python` `auth_handler.py` due to a falsy chain evaluation on empty dictionaries. Added default OIDC fallback scopes (`["openid", "profile", "email"]`) for `USER_IDENTITY` flows when no explicit scopes are provided. If a tool required the same authentication service across multiple parameters, the `add_auth_token_getter` logic would attempt to register it twice, causing a `ValueError: already registered`. Collected `needed_services` into a deduplicated `set` and added a protective check before registration (`if not hasattr(...) or s not in ...`) to ensure idempotency. - [x] Verified that submitting empty scopes successfully defaults to `openid profile email` and resolves without crashing `adk-python`. - [x] Verified that the `id_token` is successfully propagated bounds and retrieved via `getattr(creds.oauth2, "id_token", ...)` during tool execution. - [x] Verified that subsequent tool executions load the saved credentials seamlessly from the `CredentialService` without triggering continuous Google OAuth consent screens. --- packages/toolbox-adk/src/toolbox_adk/tool.py | 150 ++++++++++++------ .../tests/integration/test_integration.py | 15 +- packages/toolbox-adk/tests/unit/test_tool.py | 57 +++++-- 3 files changed, 164 insertions(+), 58 deletions(-) diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index 1b0476f0..80eb86ff 100644 --- a/packages/toolbox-adk/src/toolbox_adk/tool.py +++ b/packages/toolbox-adk/src/toolbox_adk/tool.py @@ -36,6 +36,22 @@ from .client import USER_TOKEN_CONTEXT_VAR from .credentials import CredentialConfig, CredentialType +# --- Monkey Patch ADK OAuth2 Exchange to Retain ID Tokens --- +# TODO(id_token): Remove this monkey patch once the PR https://github.com/google/adk-python/pull/4402 is merged. +# Google's ID Token is required by MCP Toolbox but ADK's `update_credential_with_tokens` natively drops the `id_token`. +import google.adk.auth.oauth2_credential_util as oauth2_credential_util +import google.adk.auth.exchanger.oauth2_credential_exchanger as oauth2_credential_exchanger +_orig_update_cred = oauth2_credential_util.update_credential_with_tokens + +def _patched_update_credential_with_tokens(auth_credential, tokens): + _orig_update_cred(auth_credential, tokens) + if tokens and "id_token" in tokens and auth_credential and auth_credential.oauth2: + # Pydantic's `extra="allow"` config preserves this dynamically set attribute + setattr(auth_credential.oauth2, "id_token", tokens["id_token"]) + +oauth2_credential_util.update_credential_with_tokens = _patched_update_credential_with_tokens +oauth2_credential_exchanger.update_credential_with_tokens = _patched_update_credential_with_tokens +# ------------------------------------------------------------- class ToolboxTool(BaseTool): """ @@ -135,56 +151,98 @@ async def run_async( reset_token = None if self._auth_config and self._auth_config.type == CredentialType.USER_IDENTITY: - if not self._auth_config.client_id or not self._auth_config.client_secret: - raise ValueError("USER_IDENTITY requires client_id and client_secret") - - # Construct ADK AuthConfig - scopes = self._auth_config.scopes or [ - "https://www.googleapis.com/auth/cloud-platform" - ] - scope_dict = {s: "" for s in scopes} - - auth_config_adk = AuthConfig( - auth_scheme=OAuth2( - flows=OAuthFlows( - authorizationCode=OAuthFlowAuthorizationCode( - authorizationUrl="https://accounts.google.com/o/oauth2/auth", - tokenUrl="https://oauth2.googleapis.com/token", - scopes=scope_dict, + requires_auth = ( + len(self._core_tool._required_authn_params) > 0 + or len(self._core_tool._required_authz_tokens) > 0 + ) + + if requires_auth: + if not self._auth_config.client_id or not self._auth_config.client_secret: + raise ValueError("USER_IDENTITY requires client_id and client_secret") + + # Construct ADK AuthConfig + scopes = self._auth_config.scopes or ["openid", "profile", "email"] + scope_dict = {s: "" for s in scopes} + + auth_config_adk = AuthConfig( + auth_scheme=OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://accounts.google.com/o/oauth2/auth", + tokenUrl="https://oauth2.googleapis.com/token", + scopes=scope_dict, + ) ) - ) - ), - raw_auth_credential=AuthCredential( - auth_type=AuthCredentialTypes.OAUTH2, - oauth2=OAuth2Auth( - client_id=self._auth_config.client_id, - client_secret=self._auth_config.client_secret, ), - ), - ) + raw_auth_credential=AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id=self._auth_config.client_id, + client_secret=self._auth_config.client_secret, + ), + ), + ) - # Check if we already have credentials from a previous exchange - try: - # get_auth_response returns AuthCredential if found - creds = tool_context.get_auth_response(auth_config_adk) - if creds and creds.oauth2 and creds.oauth2.access_token: - reset_token = USER_TOKEN_CONTEXT_VAR.set(creds.oauth2.access_token) - else: - # Request credentials and pause execution + # Check if we already have credentials from a previous exchange + try: + # Try to load credential from credential service first (persists across sessions) + creds = None + try: + if tool_context._invocation_context.credential_service: + creds = await tool_context._invocation_context.credential_service.load_credential( + auth_config=auth_config_adk, + callback_context=tool_context + ) + except ValueError: + # Credential service might not be initialized + pass + + if not creds: + # Fallback to session state (get_auth_response returns AuthCredential if found) + creds = tool_context.get_auth_response(auth_config_adk) + + if creds and creds.oauth2 and creds.oauth2.access_token: + reset_token = USER_TOKEN_CONTEXT_VAR.set(creds.oauth2.access_token) + + # Bind the token to the underlying core_tool so it constructs headers properly + needed_services = set() + for requested_service in (list(self._core_tool._required_authn_params.values()) + list(self._core_tool._required_authz_tokens)): + if isinstance(requested_service, list): + needed_services.update(requested_service) + else: + needed_services.add(requested_service) + + for s in needed_services: + # Only add if not already registered (prevents ValueError on duplicate params or subsequent runs) + if not hasattr(self._core_tool, '_auth_token_getters') or s not in self._core_tool._auth_token_getters: + # TODO(id_token): Uncomment this line and remove the `getattr` fallback below once PR https://github.com/google/adk-python/pull/4402 is merged. + # self._core_tool = self._core_tool.add_auth_token_getter(s, lambda t=creds.oauth2.id_token or creds.oauth2.access_token: t) + self._core_tool = self._core_tool.add_auth_token_getter(s, lambda t=getattr(creds.oauth2, "id_token", creds.oauth2.access_token): t) + # Once we use it from get_auth_response, save it to the auth service for future use + try: + if tool_context._invocation_context.credential_service: + auth_config_adk.exchanged_auth_credential = creds + await tool_context._invocation_context.credential_service.save_credential( + auth_config=auth_config_adk, + callback_context=tool_context + ) + except Exception as e: + logging.debug(f"Failed to save credential to service: {e}") + else: + 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."} + except Exception as e: + if "credential" in str(e).lower() or isinstance(e, ValueError): + raise e + + logging.warning( + f"Unexpected error in get_auth_response during User Identity (OAuth2) retrieval: {e}. " + "Falling back to request_credential.", + exc_info=True + ) + # Fallback to request logic tool_context.request_credential(auth_config_adk) - return None - except Exception as e: - if "credential" in str(e).lower() or isinstance(e, ValueError): - raise e - - logging.warning( - f"Unexpected error in get_auth_response during User Identity (OAuth2) retrieval: {e}. " - "Falling back to request_credential.", - exc_info=True - ) - # Fallback to request logic - tool_context.request_credential(auth_config_adk) - return None + 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."} result: Optional[Any] = None error: Optional[Exception] = None diff --git a/packages/toolbox-adk/tests/integration/test_integration.py b/packages/toolbox-adk/tests/integration/test_integration.py index c7088e08..4c0a26f6 100644 --- a/packages/toolbox-adk/tests/integration/test_integration.py +++ b/packages/toolbox-adk/tests/integration/test_integration.py @@ -203,30 +203,41 @@ async def test_3lo_flow_simulation(self): # The wrapper should catch the missing creds and request them. assert result_first is None, "Tool should return None sig for auth requirement" mock_ctx_first.request_credential.assert_called_once() - + # Inspect the requested config auth_config = mock_ctx_first.request_credential.call_args[0][0] assert auth_config.raw_auth_credential.oauth2.client_id == "test-client-id" + # Verify the default fallback scopes were assigned correctly to avoid upstream crashes + assert auth_config.auth_scheme.flows.authorizationCode.scopes == {"openid": "", "profile": "", "email": ""} mock_ctx_second = MagicMock() # Simulate "Auth Response Found" mock_creds = AuthCredential( auth_type=AuthCredentialTypes.OAUTH2, - oauth2=OAuth2Auth(access_token="fake-access-token"), + oauth2=OAuth2Auth(access_token="fake-access-token", id_token="fake-id-token"), ) mock_ctx_second.get_auth_response.return_value = mock_creds + # Setup the credential service mock to verify credential persistence across sessions + mock_cred_service = AsyncMock() + mock_ctx_second._invocation_context = MagicMock() + mock_ctx_second._invocation_context.credential_service = mock_cred_service + print("Running tool second time (expecting success or server error)...") try: result_second = await tool.run_async({"num_rows": "1"}, mock_ctx_second) assert result_second is not None + # Verify that the tool saved the credentials to the storage service backends locally + mock_cred_service.save_credential.assert_called_once() except Exception as e: mock_ctx_second.request_credential.assert_not_called() err_msg = str(e).lower() assert any(x in err_msg for x in ["401", "403", "unauthorized", "forbidden"]), f"Caught UNEXPECTED exception: {type(e).__name__}: {e}" print(f"Caught expected server exception with fake token: {e}") + # Verify that the tool AT LEAST triggered save_credential before failing via core_tool inner HTTP req + mock_cred_service.save_credential.assert_called_once() finally: await toolset.close() diff --git a/packages/toolbox-adk/tests/unit/test_tool.py b/packages/toolbox-adk/tests/unit/test_tool.py index 6c7c7a14..3b3ce19a 100644 --- a/packages/toolbox-adk/tests/unit/test_tool.py +++ b/packages/toolbox-adk/tests/unit/test_tool.py @@ -155,6 +155,8 @@ async def test_3lo_missing_client_secret(self): core_tool = AsyncMock() core_tool.__name__ = "mock_tool" core_tool.__doc__ = "mock doc" + core_tool._required_authn_params = {"mock_param": "mock_service"} + core_tool._required_authz_tokens = [] auth_config = CredentialConfig(type=CredentialType.USER_IDENTITY) # Missing client_id/secret @@ -174,6 +176,8 @@ async def test_3lo_request_credential_when_missing(self): core_tool.__doc__ = "mock" core_tool.__name__ = "mock_tool" core_tool.__doc__ = "mock doc" + core_tool._required_authn_params = {"mock_param": "mock_service"} + core_tool._required_authz_tokens = [] auth_config = CredentialConfig( type=CredentialType.USER_IDENTITY, client_id="cid", client_secret="csec" @@ -187,8 +191,8 @@ async def test_3lo_request_credential_when_missing(self): result = await tool.run_async({}, ctx) - # Verify result is None (signal pause) - assert result is None + # Verify result is error/stop + assert isinstance(result, dict) and "error" in result # Verify request_credential was called ctx.request_credential.assert_called_once() # Verify core tool was NOT called @@ -198,10 +202,12 @@ async def test_3lo_request_credential_when_missing(self): async def test_3lo_uses_existing_credential(self): # Test that if creds exist, they are used and injected core_tool = AsyncMock(return_value="success") - core_tool.__name__ = "mock" - core_tool.__doc__ = "mock" core_tool.__name__ = "mock_tool" core_tool.__doc__ = "mock doc" + # Setup overlapping needed services to test deduplication + core_tool._required_authn_params = {"mock_param": "mock_service", "another_param": "mock_service"} + core_tool._required_authz_tokens = ["mock_service"] + core_tool.add_auth_token_getter = MagicMock(return_value=core_tool) auth_config = CredentialConfig( type=CredentialType.USER_IDENTITY, client_id="cid", client_secret="csec" @@ -210,11 +216,19 @@ async def test_3lo_uses_existing_credential(self): tool = ToolboxTool(core_tool, auth_config=auth_config) ctx = MagicMock() - # Mock get_auth_response returning valid creds + # Mock get_auth_response returning valid creds with both access & id tokens mock_creds = MagicMock() - mock_creds.oauth2.access_token = "valid_token" + mock_creds.oauth2.access_token = "valid_access_token" + mock_creds.oauth2.id_token = "valid_id_token" ctx.get_auth_response.return_value = mock_creds + # Set up invocation context and credential service mock to verify saving and avoid await errors + 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 + result = await tool.run_async({}, ctx) # Verify result is success @@ -223,15 +237,31 @@ async def test_3lo_uses_existing_credential(self): ctx.request_credential.assert_not_called() # Verify core tool WAS called core_tool.assert_called_once() + + # Verify deduplication: add_auth_token_getter should only be called ONCE for "mock_service" + core_tool.add_auth_token_getter.assert_called_once() + call_args_getter = core_tool.add_auth_token_getter.call_args[0] + assert call_args_getter[0] == "mock_service" + # Evaluate the getter lambda to ensure it prefers id_token + token_getter_lambda = call_args_getter[1] + assert token_getter_lambda() == "valid_id_token" + + # Verify save_credential was called with the exchanged credential + mock_cred_service.save_credential.assert_called_once() + call_args = mock_cred_service.save_credential.call_args[1] + assert call_args["auth_config"].exchanged_auth_credential == mock_creds + + # Verify safe scope fallback to ["openid", "profile", "email"] when scopes is None + assert call_args["auth_config"].auth_scheme.flows.authorizationCode.scopes == {"openid": "", "profile": "", "email": ""} @pytest.mark.asyncio async def test_3lo_exception_reraise(self): # Test that specific credential errors are re-raised core_tool = AsyncMock() - core_tool.__name__ = "mock" - core_tool.__doc__ = "mock" core_tool.__name__ = "mock_tool" core_tool.__doc__ = "mock doc" + core_tool._required_authn_params = {"mock_param": "mock_service"} + core_tool._required_authz_tokens = [] auth_config = CredentialConfig( type=CredentialType.USER_IDENTITY, client_id="cid", client_secret="csec" @@ -239,6 +269,11 @@ async def test_3lo_exception_reraise(self): tool = ToolboxTool(core_tool, auth_config=auth_config) ctx = MagicMock() + mock_cred_service = MagicMock() + mock_cred_service.load_credential = AsyncMock(return_value=None) + ctx._invocation_context = MagicMock() + ctx._invocation_context.credential_service = mock_cred_service + # Mock get_auth_response raising ValueError ctx.get_auth_response.side_effect = ValueError("Invalid Credential") @@ -253,6 +288,8 @@ async def test_3lo_exception_fallback(self): core_tool.__doc__ = "mock" core_tool.__name__ = "mock_tool" core_tool.__doc__ = "mock doc" + core_tool._required_authn_params = {"mock_param": "mock_service"} + core_tool._required_authz_tokens = [] auth_config = CredentialConfig( type=CredentialType.USER_IDENTITY, client_id="cid", client_secret="csec" @@ -265,8 +302,8 @@ async def test_3lo_exception_fallback(self): result = await tool.run_async({}, ctx) - # Should catch RuntimeError, call request_credential, and return None - assert result is None + # Should catch RuntimeError, call request_credential, and return error map + assert isinstance(result, dict) and "error" in result ctx.request_credential.assert_called_once() def test_param_type_to_schema_type(self): From fdc40ee7d4da5e4447edd8dd2c3ffbbd12e8d90c Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Sat, 7 Feb 2026 04:34:35 +0530 Subject: [PATCH 2/7] fix(adk): force proxy auth requirement in integration test to trigger 3lo simulation --- packages/toolbox-adk/tests/integration/test_integration.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/packages/toolbox-adk/tests/integration/test_integration.py b/packages/toolbox-adk/tests/integration/test_integration.py index 4c0a26f6..49511a58 100644 --- a/packages/toolbox-adk/tests/integration/test_integration.py +++ b/packages/toolbox-adk/tests/integration/test_integration.py @@ -188,6 +188,10 @@ async def test_3lo_flow_simulation(self): assert declaration.name == "get-n-rows" assert "num_rows" in declaration.parameters.properties + # Force the proxy tool to require auth to properly simulate the 3LO flow branches + tool._core_tool._required_authn_params = {"mock_param": "mock_service"} + tool._core_tool._required_authz_tokens = [] + # Create a mock context that behaves like ADK's ReadonlyContext mock_ctx_first = MagicMock() # Simulate "No Auth Response Found" @@ -201,7 +205,7 @@ async def test_3lo_flow_simulation(self): result_first = await tool.run_async({"num_rows": "1"}, mock_ctx_first) # The wrapper should catch the missing creds and request them. - assert result_first is None, "Tool should return None sig for auth requirement" + assert isinstance(result_first, dict) and "error" in result_first, "Tool should return error sig for auth requirement" mock_ctx_first.request_credential.assert_called_once() # Inspect the requested config From 937b385540032e631326f2caad20c49b370dc8c1 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Sat, 7 Feb 2026 04:45:12 +0530 Subject: [PATCH 3/7] fix(adk): use patch to safely inject proxy auth requirements for integration tests --- .../tests/integration/test_integration.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/packages/toolbox-adk/tests/integration/test_integration.py b/packages/toolbox-adk/tests/integration/test_integration.py index 49511a58..cc52dc3a 100644 --- a/packages/toolbox-adk/tests/integration/test_integration.py +++ b/packages/toolbox-adk/tests/integration/test_integration.py @@ -189,9 +189,6 @@ async def test_3lo_flow_simulation(self): assert "num_rows" in declaration.parameters.properties # Force the proxy tool to require auth to properly simulate the 3LO flow branches - tool._core_tool._required_authn_params = {"mock_param": "mock_service"} - tool._core_tool._required_authz_tokens = [] - # Create a mock context that behaves like ADK's ReadonlyContext mock_ctx_first = MagicMock() # Simulate "No Auth Response Found" @@ -202,7 +199,12 @@ async def test_3lo_flow_simulation(self): mock_ctx_first._invocation_context.credential_service = mock_cred_service_first print("Running tool first time (expecting auth request)...") - result_first = await tool.run_async({"num_rows": "1"}, mock_ctx_first) + from unittest.mock import patch, PropertyMock + with patch.object(type(tool._core_tool), "_required_authn_params", new_callable=PropertyMock) as mock_auth, \ + patch.object(type(tool._core_tool), "_required_authz_tokens", new_callable=PropertyMock) as mock_authz: + mock_auth.return_value = {"mock_param": "mock_service"} + mock_authz.return_value = [] + result_first = await tool.run_async({"num_rows": "1"}, mock_ctx_first) # The wrapper should catch the missing creds and request them. assert isinstance(result_first, dict) and "error" in result_first, "Tool should return error sig for auth requirement" @@ -231,7 +233,11 @@ async def test_3lo_flow_simulation(self): print("Running tool second time (expecting success or server error)...") try: - result_second = await tool.run_async({"num_rows": "1"}, mock_ctx_second) + with patch.object(type(tool._core_tool), "_required_authn_params", new_callable=PropertyMock) as mock_auth, \ + patch.object(type(tool._core_tool), "_required_authz_tokens", new_callable=PropertyMock) as mock_authz: + mock_auth.return_value = {"mock_param": "mock_service"} + mock_authz.return_value = [] + result_second = await tool.run_async({"num_rows": "1"}, mock_ctx_second) assert result_second is not None # Verify that the tool saved the credentials to the storage service backends locally mock_cred_service.save_credential.assert_called_once() From 1ba6e495f8f405842df88b119b4f594714b167b1 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Sat, 7 Feb 2026 04:51:41 +0530 Subject: [PATCH 4/7] fix(adk): use name-mangled assignment to properly mock toolbox_core proxy private auth params --- .../tests/integration/test_integration.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/packages/toolbox-adk/tests/integration/test_integration.py b/packages/toolbox-adk/tests/integration/test_integration.py index cc52dc3a..c6b8e11b 100644 --- a/packages/toolbox-adk/tests/integration/test_integration.py +++ b/packages/toolbox-adk/tests/integration/test_integration.py @@ -189,6 +189,9 @@ async def test_3lo_flow_simulation(self): assert "num_rows" in declaration.parameters.properties # Force the proxy tool to require auth to properly simulate the 3LO flow branches + tool._core_tool._ToolboxTool__required_authn_params = {"mock_param": "mock_service"} + tool._core_tool._ToolboxTool__required_authz_tokens = [] + # Create a mock context that behaves like ADK's ReadonlyContext mock_ctx_first = MagicMock() # Simulate "No Auth Response Found" @@ -199,12 +202,7 @@ async def test_3lo_flow_simulation(self): mock_ctx_first._invocation_context.credential_service = mock_cred_service_first print("Running tool first time (expecting auth request)...") - from unittest.mock import patch, PropertyMock - with patch.object(type(tool._core_tool), "_required_authn_params", new_callable=PropertyMock) as mock_auth, \ - patch.object(type(tool._core_tool), "_required_authz_tokens", new_callable=PropertyMock) as mock_authz: - mock_auth.return_value = {"mock_param": "mock_service"} - mock_authz.return_value = [] - result_first = await tool.run_async({"num_rows": "1"}, mock_ctx_first) + result_first = await tool.run_async({"num_rows": "1"}, mock_ctx_first) # The wrapper should catch the missing creds and request them. assert isinstance(result_first, dict) and "error" in result_first, "Tool should return error sig for auth requirement" @@ -233,11 +231,7 @@ async def test_3lo_flow_simulation(self): print("Running tool second time (expecting success or server error)...") try: - with patch.object(type(tool._core_tool), "_required_authn_params", new_callable=PropertyMock) as mock_auth, \ - patch.object(type(tool._core_tool), "_required_authz_tokens", new_callable=PropertyMock) as mock_authz: - mock_auth.return_value = {"mock_param": "mock_service"} - mock_authz.return_value = [] - result_second = await tool.run_async({"num_rows": "1"}, mock_ctx_second) + result_second = await tool.run_async({"num_rows": "1"}, mock_ctx_second) assert result_second is not None # Verify that the tool saved the credentials to the storage service backends locally mock_cred_service.save_credential.assert_called_once() From 2faeaa9d4028fc30d11cb4163f33b7b4f63b1273 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Sat, 7 Feb 2026 04:58:03 +0530 Subject: [PATCH 5/7] fix(adk): use authz tokens instead of authn params to avoid toolbox-core bindings validation error --- packages/toolbox-adk/tests/integration/test_integration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/toolbox-adk/tests/integration/test_integration.py b/packages/toolbox-adk/tests/integration/test_integration.py index c6b8e11b..cdee0886 100644 --- a/packages/toolbox-adk/tests/integration/test_integration.py +++ b/packages/toolbox-adk/tests/integration/test_integration.py @@ -189,8 +189,8 @@ async def test_3lo_flow_simulation(self): assert "num_rows" in declaration.parameters.properties # Force the proxy tool to require auth to properly simulate the 3LO flow branches - tool._core_tool._ToolboxTool__required_authn_params = {"mock_param": "mock_service"} - tool._core_tool._ToolboxTool__required_authz_tokens = [] + tool._core_tool._ToolboxTool__required_authn_params = {} + tool._core_tool._ToolboxTool__required_authz_tokens = ["mock_service"] # Create a mock context that behaves like ADK's ReadonlyContext mock_ctx_first = MagicMock() From 7df32769b5c2cdbad2331eb813979d06f6edf68f Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Sat, 7 Feb 2026 05:03:48 +0530 Subject: [PATCH 6/7] test(adk): correctly mock credential persistence check to evaluate get_auth_response successfully --- packages/toolbox-adk/tests/integration/test_integration.py | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/toolbox-adk/tests/integration/test_integration.py b/packages/toolbox-adk/tests/integration/test_integration.py index cdee0886..fa4fdc1a 100644 --- a/packages/toolbox-adk/tests/integration/test_integration.py +++ b/packages/toolbox-adk/tests/integration/test_integration.py @@ -225,6 +225,7 @@ async def test_3lo_flow_simulation(self): # Setup the credential service mock to verify credential persistence across sessions mock_cred_service = AsyncMock() + mock_cred_service.load_credential.return_value = None mock_ctx_second._invocation_context = MagicMock() mock_ctx_second._invocation_context.credential_service = mock_cred_service From 145792867544cf636171ebcf3317b21104a7f493 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Mon, 9 Feb 2026 14:12:24 +0530 Subject: [PATCH 7/7] chore: Improve comments --- packages/toolbox-adk/src/toolbox_adk/tool.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index 80eb86ff..dde37ce4 100644 --- a/packages/toolbox-adk/src/toolbox_adk/tool.py +++ b/packages/toolbox-adk/src/toolbox_adk/tool.py @@ -37,8 +37,8 @@ from .credentials import CredentialConfig, CredentialType # --- Monkey Patch ADK OAuth2 Exchange to Retain ID Tokens --- -# TODO(id_token): Remove this monkey patch once the PR https://github.com/google/adk-python/pull/4402 is merged. # Google's ID Token is required by MCP Toolbox but ADK's `update_credential_with_tokens` natively drops the `id_token`. +# TODO(id_token): Remove this monkey patch once the PR https://github.com/google/adk-python/pull/4402 is merged. import google.adk.auth.oauth2_credential_util as oauth2_credential_util import google.adk.auth.exchanger.oauth2_credential_exchanger as oauth2_credential_exchanger _orig_update_cred = oauth2_credential_util.update_credential_with_tokens @@ -46,7 +46,6 @@ def _patched_update_credential_with_tokens(auth_credential, tokens): _orig_update_cred(auth_credential, tokens) if tokens and "id_token" in tokens and auth_credential and auth_credential.oauth2: - # Pydantic's `extra="allow"` config preserves this dynamically set attribute setattr(auth_credential.oauth2, "id_token", tokens["id_token"]) oauth2_credential_util.update_credential_with_tokens = _patched_update_credential_with_tokens