diff --git a/src/google/adk/tools/application_integration_tool/application_integration_toolset.py b/src/google/adk/tools/application_integration_tool/application_integration_toolset.py index d904de49..523843e2 100644 --- a/src/google/adk/tools/application_integration_tool/application_integration_toolset.py +++ b/src/google/adk/tools/application_integration_tool/application_integration_toolset.py @@ -245,6 +245,11 @@ def _parse_spec_to_tools(self, spec_dict, connection_details): action=action, operation=operation, rest_api_tool=rest_api_tool, + dynamic_auth_config=( + {"userToken": None} + if connection_details.get("authOverrideEnabled") + else None + ), ) self.generated_tools[tool.name] = tool diff --git a/src/google/adk/tools/application_integration_tool/clients/connections_client.py b/src/google/adk/tools/application_integration_tool/clients/connections_client.py index 1561a8eb..212c7aed 100644 --- a/src/google/adk/tools/application_integration_tool/clients/connections_client.py +++ b/src/google/adk/tools/application_integration_tool/clients/connections_client.py @@ -857,3 +857,40 @@ def _poll_operation(self, operation_id: str) -> Dict[str, Any]: operation_done = operation_response.get("done", False) time.sleep(1) return operation_response + def execute_connection( + self, + connection_name: str, + operation_id: str, + inputs: Dict[str, Any], + dynamic_auth_config: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Executes an operation on the specified connection, with optional auth override. + + Args: + connection_name: Full resource name of the connection. + operation_id: Operation to perform (e.g. EXECUTE_ACTION). + inputs: Payload for 'connectorInputPayload'. + dynamic_auth_config: If provided, sent under 'dynamicAuthConfig' for override. + + Returns: + The JSON-decoded response from the connector API. + + Raises: + requests.exceptions.RequestException: If the HTTP request fails. + """ + url = f"{self.connector_url}/v1/{connection_name}:execute" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {self._get_access_token()}", + } + body: Dict[str, Any] = { + "connectionName": connection_name, + "operation": operation_id, + "connectorInputPayload": inputs, + } + if dynamic_auth_config is not None: + body["dynamicAuthConfig"] = dynamic_auth_config + + response = requests.post(url, headers=headers, json=body) + response.raise_for_status() + return response.json() diff --git a/src/google/adk/tools/application_integration_tool/integration_connector_tool.py b/src/google/adk/tools/application_integration_tool/integration_connector_tool.py index 2513da53..e95c4c86 100644 --- a/src/google/adk/tools/application_integration_tool/integration_connector_tool.py +++ b/src/google/adk/tools/application_integration_tool/integration_connector_tool.py @@ -14,6 +14,7 @@ import logging +import asyncio from typing import Any from typing import Dict from typing import Optional @@ -75,6 +76,7 @@ def __init__( operation: str, action: str, rest_api_tool: RestApiTool, + dynamic_auth_config: Optional[Dict[str, Any]] = None, ): """Initializes the ApplicationIntegrationTool. @@ -108,6 +110,7 @@ def __init__( self.operation = operation self.action = action self.rest_api_tool = rest_api_tool + self.dynamic_auth_config = dynamic_auth_config @override def _get_declaration(self) -> FunctionDeclaration: @@ -136,8 +139,16 @@ async def run_async( args['entity'] = self.entity args['operation'] = self.operation args['action'] = self.action + + # if an override config was set on this tool, include it + if self.dynamic_auth_config is not None: + args['dynamicAuthConfig'] = self.dynamic_auth_config + logger.info('Running tool: %s with args: %s', self.name, args) - return self.rest_api_tool.call(args=args, tool_context=tool_context) + result = self.rest_api_tool.call(args=args, tool_context=tool_context) + if asyncio.iscoroutine(result): + result = await result + return result def __str__(self): return ( diff --git a/tests/unittests/tools/application_integration_tool/clients/test_connections_client.py b/tests/unittests/tools/application_integration_tool/clients/test_connections_client.py index d60938dd..c7ce46a6 100644 --- a/tests/unittests/tools/application_integration_tool/clients/test_connections_client.py +++ b/tests/unittests/tools/application_integration_tool/clients/test_connections_client.py @@ -618,3 +618,49 @@ def test_get_access_token_refreshes_expired_token( token = client._get_access_token() assert token == "new_token" mock_refresh.assert_called_once() + + +def test_execute_connection_includes_dynamic_auth_config(monkeypatch): + # Arrange: a client with fake credentials + client = ConnectionsClient( + project="proj-id", + location="us-central1", + connection="conn-id", + ) + # Stub out token fetch + monkeypatch.setattr(client, "_get_access_token", lambda: "fake-auth-token") + + captured = {} + + def fake_post(url, headers=None, json=None): + captured["url"] = url + captured["headers"] = headers + captured["json"] = json + + class FakeResponse: + + def raise_for_status(self): + pass + + def json(self): + return {"success": True} + + return FakeResponse() + + # Replace requests.post + monkeypatch.setattr(requests, "post", fake_post) + + # Act: call execute_connection with a dynamicAuthConfig + result = client.execute_connection( + connection_name=( + "projects/proj-id/locations/us-central1/connections/conn-id" + ), + operation_id="EXECUTE_ACTION", + inputs={"foo": "bar"}, + dynamic_auth_config={"token": "end-user-token"}, + ) + + # Assert: that the JSON body included our dynamicAuthConfig + assert "dynamicAuthConfig" in captured["json"] + assert captured["json"]["dynamicAuthConfig"] == {"token": "end-user-token"} + assert result == {"success": True} diff --git a/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py b/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py index 28dbb9da..3fe320f4 100644 --- a/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py +++ b/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py @@ -391,3 +391,66 @@ def test_initialization_with_connection_details( mock_integration_client.return_value.get_openapi_spec_for_connection.assert_called_once_with( tool_name, tool_instructions ) + + +@pytest.mark.usefixtures("mock_openapi_entity_spec_parser") +def test_tool_has_dynamic_auth_config_when_override_enabled( + project, + location, + mock_integration_client, + mock_connections_client, + mock_openapi_toolset, + connection_details, +): + # simulate authOverrideEnabled = True coming from ConnectionsClient + connection_details["authOverrideEnabled"] = True + mock_connections_client.return_value.get_connection_details.return_value = ( + connection_details + ) + + toolset = ApplicationIntegrationToolset( + project, + location, + connection=connection_details["name"], + entity_operations=["E"], + ) + tool = toolset.get_tools()[0] + # should have picked up the flag and defaulted to {"userToken": None} + assert isinstance(tool, IntegrationConnectorTool) + assert tool.dynamic_auth_config == {"userToken": None} + + +@pytest.mark.asyncio +async def test_integration_connector_injects_dynamic_authConfig(monkeypatch): + # build a single IntegrationConnectorTool manually + parser = get_mocked_parsed_operation( + "op", {"x-action": "A", "x-operation": "EXECUTE_ACTION"} + ) + parsed_op = parser.parse.return_value[0] + rest = rest_api_tool.RestApiTool.from_parsed_operation(parsed_op) + tool = IntegrationConnectorTool( + name="t", + description="d", + connection_name="c", + connection_host="h", + connection_service_name="s", + entity="Issues", + operation="EXECUTE_ACTION", + action="A", + rest_api_tool=rest, + dynamic_auth_config={"userToken": "XYZ"}, + ) + + # stub out RestApiTool.call to capture args + seen = {} + + async def fake_call(self, args, tool_context): + seen.update(args) + return {"ok": True} + + monkeypatch.setattr(rest_api_tool.RestApiTool, "call", fake_call) + + out = await tool.run_async(args={}, tool_context=None) + assert out == {"ok": True} + # verifies the camel-case key went into the final args + assert seen["dynamicAuthConfig"] == {"userToken": "XYZ"}