fix: Add credential_key to ApplicationIntegrationToolset and IntegrationConnectorTool

This change allows a credential key to be passed through the toolset and connector, enabling the ToolAuthHandler to use a specific key when looking up credentials in the tool context.

Fixes #4553

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 941375909

Author: xuanyang15

src/google/adk/tools/application_integration_tool/application_integration_toolset.py

@@ -103,6 +103,7 @@ def __init__(
       auth_scheme: Optional[AuthScheme] = None,
       auth_credential: Optional[AuthCredential] = None,
       tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
+      credential_key: Optional[str] = None,
   ):
     """Args:
 
@@ -146,12 +147,14 @@ def __init__(
     self._service_account_json = service_account_json
     self._auth_scheme = auth_scheme
     self._auth_credential = auth_credential
+    self._credential_key = credential_key
     # Store auth config as instance variable so ADK can populate
     # exchanged_auth_credential in-place before calling get_tools()
     self._auth_config: Optional[AuthConfig] = (
         AuthConfig(
             auth_scheme=auth_scheme,
             raw_auth_credential=auth_credential,
+            credential_key=credential_key,
         )
         if auth_scheme
         else None
@@ -221,6 +224,7 @@ def _parse_spec_to_toolset(
           spec_dict=spec_dict,
           auth_credential=auth_credential,
           auth_scheme=auth_scheme,
+          credential_key=self._credential_key,
           tool_filter=self.tool_filter,
       )
       return
@@ -274,6 +278,7 @@ def _parse_spec_to_toolset(
               rest_api_tool=rest_api_tool,
               auth_scheme=connector_auth_scheme,
               auth_credential=connector_auth_credential,
+              credential_key=self._credential_key,
           )
       )
 

src/google/adk/tools/application_integration_tool/integration_connector_tool.py

@@ -80,6 +80,7 @@ def __init__(
       rest_api_tool: RestApiTool,
       auth_scheme: Optional[Union[AuthScheme, str]] = None,
       auth_credential: Optional[Union[AuthCredential, str]] = None,
+      credential_key: Optional[str] = None,
   ):
     """Initializes the ApplicationIntegrationTool.
 
@@ -115,6 +116,7 @@ def __init__(
     self._rest_api_tool = rest_api_tool
     self._auth_scheme = auth_scheme
     self._auth_credential = auth_credential
+    self._credential_key = credential_key
 
   @override
   def _get_declaration(self) -> FunctionDeclaration:
@@ -156,7 +158,10 @@ async def run_async(
   ) -> Dict[str, Any]:
 
     tool_auth_handler = ToolAuthHandler.from_tool_context(
-        tool_context, self._auth_scheme, self._auth_credential
+        tool_context,
+        self._auth_scheme,
+        self._auth_credential,
+        credential_key=self._credential_key,
     )
     auth_result = await tool_auth_handler.prepare_auth_credentials()
 

tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py

@@ -567,6 +567,7 @@ async def test_init_with_connection_and_custom_auth(
       tool_instructions=tool_instructions,
       auth_scheme=oauth2_scheme,
       auth_credential=auth_credential,
+      credential_key="test-key",
   )
   mock_integration_client.assert_called_once_with(
       project,
@@ -594,6 +595,7 @@ async def test_init_with_connection_and_custom_auth(
   assert (await toolset.get_tools())[0]._operation == "EXECUTE_ACTION"
   assert (await toolset.get_tools())[0]._auth_scheme == oauth2_scheme
   assert (await toolset.get_tools())[0]._auth_credential == auth_credential
+  assert (await toolset.get_tools())[0]._credential_key == "test-key"
 
 
 @pytest.mark.asyncio

tests/unittests/tools/application_integration_tool/test_integration_connector_tool.py

@@ -23,6 +23,7 @@
 from google.adk.tools.application_integration_tool.integration_connector_tool import IntegrationConnectorTool
 from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
 from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import AuthPreparationResult
+from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import ToolAuthHandler
 from google.genai.types import FunctionDeclaration
 from google.genai.types import Schema
 from google.genai.types import Type
@@ -96,6 +97,7 @@ def integration_tool_with_auth(mock_rest_api_tool):
               credentials=HttpCredentials(token="mocked_token"),
           ),
       ),
+      credential_key="test-key",
   )
 
 
@@ -190,8 +192,8 @@ async def test_run_with_auth_async_none_token(
       "sortByColumns": ["a", "b"],
   }
 
-  with mock.patch(
-      "google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.ToolAuthHandler.from_tool_context"
+  with mock.patch.object(
+      ToolAuthHandler, "from_tool_context", autospec=True
   ) as mock_from_tool_context:
     mock_tool_auth_handler_instance = mock.MagicMock()
     # Simulate an AuthCredential that would cause _prepare_dynamic_euc to return None
@@ -214,6 +216,12 @@ async def test_run_with_auth_async_none_token(
     result = await integration_tool_with_auth.run_async(
         args=input_args, tool_context={}
     )
+    mock_from_tool_context.assert_called_once_with(
+        {},
+        None,
+        integration_tool_with_auth._auth_credential,
+        credential_key="test-key",
+    )
 
     mock_rest_api_tool.call.assert_called_once_with(
         args=expected_call_args, tool_context={}
@@ -241,8 +249,8 @@ async def test_run_with_auth_async(
       "action": "TestAction",
   }
 
-  with mock.patch(
-      "google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.ToolAuthHandler.from_tool_context"
+  with mock.patch.object(
+      ToolAuthHandler, "from_tool_context", autospec=True
   ) as mock_from_tool_context:
     mock_tool_auth_handler_instance = mock.MagicMock()
 
@@ -262,6 +270,12 @@ async def test_run_with_auth_async(
     result = await integration_tool_with_auth.run_async(
         args=input_args, tool_context={}
     )
+    mock_from_tool_context.assert_called_once_with(
+        {},
+        None,
+        integration_tool_with_auth._auth_credential,
+        credential_key="test-key",
+    )
     mock_rest_api_tool.call.assert_called_once_with(
         args=expected_call_args, tool_context={}
     )