Added support for dynamic auth in integration connector tool

PiperOrigin-RevId: 759676602

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

@@ -12,9 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List
-from typing import Optional
-from typing import Union
+import logging
+from typing import List, Optional, Union
 
 from fastapi.openapi.models import HTTPBearer
 from typing_extensions import override
@@ -24,6 +23,7 @@ from ...auth.auth_credential import AuthCredential
 from ...auth.auth_credential import AuthCredentialTypes
 from ...auth.auth_credential import ServiceAccount
 from ...auth.auth_credential import ServiceAccountCredential
+from ...auth.auth_schemes import AuthScheme
 from ..base_toolset import BaseToolset
 from ..base_toolset import ToolPredicate
 from ..openapi_tool.auth.auth_helpers import service_account_scheme_credential
@@ -35,6 +35,9 @@ from .clients.integration_client import IntegrationClient
 from .integration_connector_tool import IntegrationConnectorTool
 
 
+logger = logging.getLogger(__name__)
+
+
 # TODO(cheliu): Apply a common toolset interface
 class ApplicationIntegrationToolset(BaseToolset):
   """ApplicationIntegrationToolset generates tools from a given Application
@@ -93,6 +96,8 @@ class ApplicationIntegrationToolset(BaseToolset):
       # tool/python function description.
       tool_instructions: Optional[str] = "",
       service_account_json: Optional[str] = None,
+      auth_scheme: Optional[AuthScheme] = None,
+      auth_credential: Optional[AuthCredential] = None,
       tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
   ):
     """Args:
@@ -132,6 +137,8 @@ class ApplicationIntegrationToolset(BaseToolset):
     self._tool_name_prefix = tool_name_prefix
     self._tool_instructions = tool_instructions
     self._service_account_json = service_account_json
+    self._auth_scheme = auth_scheme
+    self._auth_credential = auth_credential
     self.tool_filter = tool_filter
 
     integration_client = IntegrationClient(
@@ -212,6 +219,27 @@ class ApplicationIntegrationToolset(BaseToolset):
         rest_api_tool.configure_auth_scheme(auth_scheme)
       if auth_credential:
         rest_api_tool.configure_auth_credential(auth_credential)
+
+      auth_override_enabled = connection_details.get(
+          "authOverrideEnabled", False
+      )
+
+      if (
+          self._auth_scheme
+          and self._auth_credential
+          and not auth_override_enabled
+      ):
+        # Case: Auth provided, but override is OFF. Don't use provided auth.
+        logger.warning(
+            "Authentication schema and credentials are not used because"
+            " authOverrideEnabled is not enabled in the connection."
+        )
+        connector_auth_scheme = None
+        connector_auth_credential = None
+      else:
+        connector_auth_scheme = self._auth_scheme
+        connector_auth_credential = self._auth_credential
+
       self._tools.append(
           IntegrationConnectorTool(
               name=rest_api_tool.name,
@@ -223,6 +251,8 @@ class ApplicationIntegrationToolset(BaseToolset):
               action=action,
               operation=operation,
               rest_api_tool=rest_api_tool,
+              auth_scheme=connector_auth_scheme,
+              auth_credential=connector_auth_credential,
           )
       )
 

src/google/adk/tools/application_integration_tool/clients/connections_client.py

@@ -554,6 +554,9 @@ class ConnectionsClient:
             "serviceName": {"$ref": "#/components/schemas/serviceName"},
             "host": {"$ref": "#/components/schemas/host"},
             "entity": {"$ref": "#/components/schemas/entity"},
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
         },
     }
 
@@ -580,6 +583,9 @@ class ConnectionsClient:
             "serviceName": {"$ref": "#/components/schemas/serviceName"},
             "host": {"$ref": "#/components/schemas/host"},
             "entity": {"$ref": "#/components/schemas/entity"},
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
             "filterClause": {"$ref": "#/components/schemas/filterClause"},
         },
     }
@@ -603,6 +609,9 @@ class ConnectionsClient:
             "serviceName": {"$ref": "#/components/schemas/serviceName"},
             "host": {"$ref": "#/components/schemas/host"},
             "entity": {"$ref": "#/components/schemas/entity"},
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
         },
     }
 
@@ -625,6 +634,9 @@ class ConnectionsClient:
             "serviceName": {"$ref": "#/components/schemas/serviceName"},
             "host": {"$ref": "#/components/schemas/host"},
             "entity": {"$ref": "#/components/schemas/entity"},
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
             "filterClause": {"$ref": "#/components/schemas/filterClause"},
         },
     }
@@ -649,6 +661,9 @@ class ConnectionsClient:
             "serviceName": {"$ref": "#/components/schemas/serviceName"},
             "host": {"$ref": "#/components/schemas/host"},
             "entity": {"$ref": "#/components/schemas/entity"},
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
         },
     }
 
@@ -673,6 +688,9 @@ class ConnectionsClient:
             "connectorInputPayload": {
                 "$ref": f"#/components/schemas/connectorInputPayload_{action}"
             },
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
         },
     }
 

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

@@ -12,20 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import logging
-from typing import Any
-from typing import Dict
-from typing import Optional
+from typing import Any, Dict, Optional, Union
 
-from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
-from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import to_gemini_schema
 from google.genai.types import FunctionDeclaration
 from typing_extensions import override
 
+from ...auth.auth_credential import AuthCredential
+from ...auth.auth_schemes import AuthScheme
 from .. import BaseTool
+from ..openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
+from ..openapi_tool.openapi_spec_parser.rest_api_tool import to_gemini_schema
+from ..openapi_tool.openapi_spec_parser.tool_auth_handler import ToolAuthHandler
 from ..tool_context import ToolContext
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -56,6 +57,7 @@ class IntegrationConnectorTool(BaseTool):
       'entity',
       'operation',
       'action',
+      'dynamic_auth_config',
   ]
 
   OPTIONAL_FIELDS = [
@@ -75,6 +77,8 @@ class IntegrationConnectorTool(BaseTool):
       operation: str,
       action: str,
       rest_api_tool: RestApiTool,
+      auth_scheme: Optional[Union[AuthScheme, str]] = None,
+      auth_credential: Optional[Union[AuthCredential, str]] = None,
   ):
     """Initializes the ApplicationIntegrationTool.
 
@@ -108,6 +112,8 @@ class IntegrationConnectorTool(BaseTool):
     self._operation = operation
     self._action = action
     self._rest_api_tool = rest_api_tool
+    self._auth_scheme = auth_scheme
+    self._auth_credential = auth_credential
 
   @override
   def _get_declaration(self) -> FunctionDeclaration:
@@ -126,10 +132,45 @@ class IntegrationConnectorTool(BaseTool):
     )
     return function_decl
 
+  def _prepare_dynamic_euc(self, auth_credential: AuthCredential) -> str:
+    if (
+        auth_credential
+        and auth_credential.http
+        and auth_credential.http.credentials
+        and auth_credential.http.credentials.token
+    ):
+      return auth_credential.http.credentials.token
+    return None
+
   @override
   async def run_async(
       self, *, args: dict[str, Any], tool_context: Optional[ToolContext]
   ) -> Dict[str, Any]:
+
+    tool_auth_handler = ToolAuthHandler.from_tool_context(
+        tool_context, self._auth_scheme, self._auth_credential
+    )
+    auth_result = tool_auth_handler.prepare_auth_credentials()
+
+    if auth_result.state == 'pending':
+      return {
+          'pending': True,
+          'message': 'Needs your authorization to access your data.',
+      }
+
+    # Attach parameters from auth into main parameters list
+    if auth_result.auth_credential:
+      # Attach parameters from auth into main parameters list
+      auth_credential_token = self._prepare_dynamic_euc(
+          auth_result.auth_credential
+      )
+      if auth_credential_token:
+        args['dynamic_auth_config'] = {
+            'oauth2_auth_code_flow.access_token': auth_credential_token
+        }
+      else:
+        args['dynamic_auth_config'] = {'oauth2_auth_code_flow.access_token': {}}
+
     args['connection_name'] = self._connection_name
     args['service_name'] = self._connection_service_name
     args['host'] = self._connection_host