Add Databricks agent invocation operator and hook

providers/databricks/README.rst

@@ -57,6 +57,7 @@ PIP package                                 Version required
 ``apache-airflow-providers-common-compat``  ``>=1.21.0``
 ``apache-airflow-providers-common-sql``     ``>=1.32.0``
 ``requests``                                ``>=2.32.0,<3``
+``tenacity``                                ``>=8.3.0``
 ``databricks-sql-connector``                ``>=4.4.0``
 ``aiohttp``                                 ``>=3.14.0,<4``
 ``mergedeep``                               ``>=1.3.4``

providers/databricks/docs/index.rst

@@ -105,6 +105,7 @@ PIP package                                 Version required
 ``apache-airflow-providers-common-compat``  ``>=1.21.0``
 ``apache-airflow-providers-common-sql``     ``>=1.32.0``
 ``requests``                                ``>=2.32.0,<3``
+``tenacity``                                ``>=8.3.0``
 ``databricks-sql-connector``                ``>=4.4.0``
 ``aiohttp``                                 ``>=3.14.0,<4``
 ``mergedeep``                               ``>=1.3.4``
@@ -133,6 +134,7 @@ You can install such cross-provider dependencies when installing from PyPI. For
 Dependent package                                                                                               Extra
 ==============================================================================================================  ===============
 `apache-airflow-providers-amazon <https://airflow.apache.org/docs/apache-airflow-providers-amazon>`_            ``amazon``
+`apache-airflow-providers-common-ai <https://airflow.apache.org/docs/apache-airflow-providers-common-ai>`_      ``common.ai``
 `apache-airflow-providers-google <https://airflow.apache.org/docs/apache-airflow-providers-google>`_            ``google``
 `apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_  ``openlineage``
 ==============================================================================================================  ===============
@@ -145,12 +147,13 @@ Install them when installing from PyPI. For example:
 
 .. code-block:: bash
 
-    pip install apache-airflow-providers-databricks[avro]
+    pip install apache-airflow-providers-databricks[common.ai]
 
 
 ==================  ================================================================================================================================================================
 Extra               Dependencies
 ==================  ================================================================================================================================================================
+``common.ai``       ``apache-airflow-providers-common-ai>=0.10.0``
 ``avro``            ``fastavro>=1.9.0; python_version<"3.14"``, ``fastavro>=1.10.0; python_version>="3.12" and python_version<"3.14"``, ``fastavro>=1.12.1; python_version>="3.14"``
 ``amazon``          ``apache-airflow-providers-amazon>=9.22.0``
 ``azure-identity``  ``azure-identity>=1.25.3``

providers/databricks/docs/operators/agent.rst

@@ -0,0 +1,151 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+ ..   http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+
+
+.. _howto/operator:DatabricksAgentInvokeOperator:
+
+.. spelling:word-list::
+
+    MLflow
+
+Invoke a Databricks agent
+============================
+
+Use :class:`~airflow.providers.databricks.operators.agent.DatabricksAgentInvokeOperator`
+to invoke an agent deployed on Databricks Apps using ``DurableAgentServer``.
+The operator submits a background request to the invocation API and returns the
+invocation response, including its status and output, through XCom.
+
+Authentication
+--------------
+
+Configure a Databricks connection with the workspace URL in ``host``, the service
+principal's client ID in ``login``, its client secret in ``password``, and
+``{"service_principal_oauth": true}`` in ``extra``. The service principal needs
+permission to use the deployed app. Databricks Apps requires OAuth; personal
+access tokens are not supported. Supply the app's HTTPS base URL as ``app_url``.
+
+Invoke an agent
+---------------
+
+.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_agent.py
+    :language: python
+    :start-after: [START howto_operator_databricks_agent_invoke]
+    :end-before: [END howto_operator_databricks_agent_invoke]
+
+The agent defines the schema of ``input``. Agents generated by the Agent Bricks
+CLI templates require a ``session_id``. Reuse the session ID across invocations
+to continue a conversation. The hook sends it as ``X-Routing-Key`` when supplied.
+
+By default, the operator waits for the result. Set ``deferrable=True`` to release
+the worker while waiting, or ``wait_for_termination=False`` to return the initial
+submission response immediately. ``polling_period_seconds`` controls polling.
+The timeout limits the wait after submission; it does not cancel the remote run.
+Synchronous polling bounds each request and its retries by the remaining wait
+and raises :class:`~airflow.providers.databricks.exceptions.DatabricksAgentInvocationTimeout`
+when that budget expires. OAuth refresh has a separate HTTP timeout and can delay
+reporting that the wait expired. In deferrable mode, Airflow enforces the timeout on the deferred task.
+
+.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_agent.py
+    :language: python
+    :start-after: [START howto_operator_databricks_agent_invoke_deferrable]
+    :end-before: [END howto_operator_databricks_agent_invoke_deferrable]
+
+Retries and results
+-------------------
+
+An invocation requires a UUID. Unless ``invocation_id`` is provided, the operator
+generates a stable UUID from the app URL, Dag ID, task ID, run ID, map index,
+rendered input and session ID. App URLs with and without a trailing slash produce
+the same UUID. Retries and cleared tasks with unchanged input and session in the
+same Dag run reuse the invocation. Changing the rendered input or session generates
+a new UUID. To deliberately repeat the same request in that run, supply a new UUID.
+Databricks retains idempotency while it retains the invocation record. Reusing
+an explicit ID with different input or session returns an HTTP 409 error.
+
+A stored ``failed`` invocation raises ``AirflowFailException``, preventing automatic
+task retries: submitting the same UUID returns that failure without running the
+agent again. Clearing the task with unchanged input and session also returns the
+stored failure. Supply a new ``invocation_id`` to run it again.
+
+The runtime states are ``queued``, ``active``, ``completed`` and ``failed``.
+CLI template agents report a pause as ``status="completed"`` with
+``output.status="interrupted"``. The operator returns this response successfully
+so downstream tasks can handle human input. To resume it, submit a new invocation
+with the same session ID and the appropriate ``resume`` input.
+The returned response preserves the agent's output without assuming its schema.
+
+The operator targets ``DurableAgentServer``'s ``/api/invocations`` API. Legacy
+MLflow agent servers using ``/responses`` and Model Serving endpoints are not
+supported by this operator.
+
+For direct calls, use
+:class:`~airflow.providers.databricks.hooks.agent.DatabricksAgentHook` and its
+``create_invocation`` and ``get_invocation`` methods.
+See `Query agents deployed on Databricks
+<https://docs.databricks.com/aws/en/agents/query-llms>`_ for the API contract.
+
+Common AI managed-agent interface
+---------------------------------
+
+Install ``apache-airflow-providers-databricks[common.ai]`` to use the hook with
+Common AI consumers such as ``ManagedAgentToolset``. Bind the deployed app's URL
+to a hook configured with the same OAuth connection:
+
+Import ``ManagedAgentRequest`` from ``airflow.providers.common.ai.managed_agents``
+and ``DatabricksAgentHook`` from ``airflow.providers.databricks.hooks.agent``:
+
+.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_agent.py
+    :language: python
+    :start-after: [START howto_databricks_managed_agent]
+    :end-before: [END howto_databricks_managed_agent]
+
+This interface uses the synchronous invocation API. It converts a prompt to a
+user message, or passes supplied messages under ``input.messages``. For an agent
+that expects a prompt under another key, set ``vendor_options={"input_key": "question"}``.
+Other input schemas remain available through ``create_invocation`` and the operator.
+
+Each call generates a new invocation UUID. For one direct ``agent.invoke(...)``
+per task, a stable UUID in ``vendor_options={"invocation_id": "YOUR_UUID"}`` can
+reuse that invocation across task retries. Do not set a fixed invocation ID in
+``ManagedAgentToolset`` vendor options: the toolset forwards those options on every
+call, so a different prompt would conflict and a repeated prompt would return the
+stored result.
+
+HTTP retries within a call reuse the same UUID. ``session_id`` continues the
+conversation and is sent as the routing key. When omitted, the hook uses the
+invocation UUID as a one-shot session ID, including for toolset calls. A hook can bind multiple app URLs;
+vendor options cannot override the app URL or connection.
+
+``ManagedAgentResponse.raw`` preserves the full invocation response. String outputs
+become ``text``; an object with a string ``output`` field uses that field as text;
+other outputs become JSON text. Non-string outputs are also available in
+``structured``. The invocation ID is returned as ``trace_ref``.
+
+The managed-agent interface requires a ``completed`` invocation whose output is
+not marked ``interrupted`` by the CLI template. Stored failures, interruptions,
+missing or unexpected statuses raise ``ManagedAgentInvocationError``; use the
+operator to handle an interrupted invocation's response directly.
+Terminal HTTP errors raise the same exception, while transient HTTP and connection
+errors propagate after the hook's configured retries. A synchronous agent failure
+returns HTTP 500. After HTTP retries, the hook queries the known invocation ID:
+a stored ``failed`` status becomes ``ManagedAgentInvocationError``. If the lookup
+cannot confirm a stored failure, the original HTTP error propagates as transient.
+The request timeout bounds the HTTP call, its retries and the status lookup;
+OAuth refresh retains its separate timeout.
+The operator and background hook methods do not require the Common AI extra.

providers/databricks/provider.yaml

@@ -122,6 +122,7 @@ integrations:
       - /docs/apache-airflow-providers-databricks/operators/submit_run.rst
       - /docs/apache-airflow-providers-databricks/operators/run_now.rst
       - /docs/apache-airflow-providers-databricks/operators/task.rst
+      - /docs/apache-airflow-providers-databricks/operators/agent.rst
     logo: /docs/integration-logos/Databricks.png
     tags: [service]
   - integration-name: Databricks SQL
@@ -149,6 +150,7 @@ operators:
   - integration-name: Databricks
     python-modules:
       - airflow.providers.databricks.operators.databricks
+      - airflow.providers.databricks.operators.agent
   - integration-name: Databricks SQL
     python-modules:
       - airflow.providers.databricks.operators.databricks_sql
@@ -179,6 +181,7 @@ hooks:
     python-modules:
       - airflow.providers.databricks.hooks.databricks
       - airflow.providers.databricks.hooks.databricks_base
+      - airflow.providers.databricks.hooks.agent
   - integration-name: Databricks SQL
     python-modules:
       - airflow.providers.databricks.hooks.databricks_sql
@@ -187,6 +190,7 @@ triggers:
   - integration-name: Databricks
     python-modules:
       - airflow.providers.databricks.triggers.databricks
+      - airflow.providers.databricks.triggers.agent
 
 sensors:
   - integration-name: Databricks

providers/databricks/pyproject.toml

@@ -62,6 +62,7 @@ dependencies = [
     "apache-airflow-providers-common-compat>=1.21.0",
     "apache-airflow-providers-common-sql>=1.32.0",
     "requests>=2.32.0,<3",
+    "tenacity>=8.3.0",
     "databricks-sql-connector>=4.4.0",
     "aiohttp>=3.14.0, <4",
     "mergedeep>=1.3.4",
@@ -76,6 +77,9 @@ dependencies = [
 # The optional dependencies should be modified in place in the generated file
 # Any change in the dependencies is preserved when the file is regenerated
 [project.optional-dependencies]
+"common.ai" = [
+    "apache-airflow-providers-common-ai>=0.10.0",  # use next version
+]
 # pip install apache-airflow-providers-databricks[sdk]
 "avro" = [
     'fastavro>=1.9.0; python_version<"3.14"',
@@ -113,6 +117,7 @@ dev = [
     "apache-airflow-task-sdk",
     "apache-airflow-devel-common",
     "apache-airflow-providers-amazon",
+    "apache-airflow-providers-common-ai",
     "apache-airflow-providers-common-compat",
     "apache-airflow-providers-common-sql",
     "apache-airflow-providers-google",

providers/databricks/src/airflow/providers/databricks/exceptions.py

@@ -40,6 +40,14 @@ class DatabricksWarehouseError(AirflowException):
     """Raised when a SQL warehouse fails to reach or times out waiting for a target state."""
 
 
+class DatabricksAgentInvocationError(AirflowException):
+    """Raised when an agent invocation fails or cannot be completed."""
+
+
+class DatabricksAgentInvocationTimeout(DatabricksAgentInvocationError):
+    """Raised when waiting for an agent invocation times out."""
+
+
 class DatabricksApiError(AirflowException):
     """Raised when a Databricks REST API call returns an error response."""
 

providers/databricks/src/airflow/providers/databricks/get_provider_info.py

@@ -36,6 +36,7 @@ def get_provider_info():
                     "/docs/apache-airflow-providers-databricks/operators/submit_run.rst",
                     "/docs/apache-airflow-providers-databricks/operators/run_now.rst",
                     "/docs/apache-airflow-providers-databricks/operators/task.rst",
+                    "/docs/apache-airflow-providers-databricks/operators/agent.rst",
                 ],
                 "logo": "/docs/integration-logos/Databricks.png",
                 "tags": ["service"],
@@ -71,7 +72,10 @@ def get_provider_info():
         "operators": [
             {
                 "integration-name": "Databricks",
-                "python-modules": ["airflow.providers.databricks.operators.databricks"],
+                "python-modules": [
+                    "airflow.providers.databricks.operators.databricks",
+                    "airflow.providers.databricks.operators.agent",
+                ],
             },
             {
                 "integration-name": "Databricks SQL",
@@ -111,6 +115,7 @@ def get_provider_info():
                 "python-modules": [
                     "airflow.providers.databricks.hooks.databricks",
                     "airflow.providers.databricks.hooks.databricks_base",
+                    "airflow.providers.databricks.hooks.agent",
                 ],
             },
             {
@@ -121,7 +126,10 @@ def get_provider_info():
         "triggers": [
             {
                 "integration-name": "Databricks",
-                "python-modules": ["airflow.providers.databricks.triggers.databricks"],
+                "python-modules": [
+                    "airflow.providers.databricks.triggers.databricks",
+                    "airflow.providers.databricks.triggers.agent",
+                ],
             }
         ],
         "sensors": [