[AlejandroMorgante]

Add DatabricksAgentInvokeOperator and DatabricksAgentHook to invoke agents hosted with DurableAgentServer on Databricks Apps. Supports service principal OAuth, background submission, synchronous polling and deferrable waiting, preserving the full invocation response in XCom.

The hook implements BaseManagedAgentHook for Common AI and toolset use, with a default session per invocation. Request, state and error handling were checked against the runtime and templates in Databricks' databricks-ai-bridge repository, tag databricks-agentbricks-v0.4.0. Stored failures are distinguished from transient HTTP errors, and interruptions nested in the output are returned by the operator and rejected by Common AI.

Invocation IDs remain stable across retries and clears with unchanged input and session. Stored failures stop task retries and require a new ID to run again. GET timeouts use the minimum of the hook's HTTP timeout and the remaining budget; synchronous POST requests retain the full remaining budget. Deadline checks reject late results, but network activity and OAuth refresh can delay timeout reporting. Background hook methods accept the app URL per call.

Includes documentation, focused unit tests and a deterministic system-test fixture requiring no model endpoint.

Validation:

• After the GET timeout and per-call URL fixes, 174 hook, operator and trigger tests passed in Breeze, plus five local timeout-budget probes.
• After adding an async HTTP-options regression test and strengthening the decreasing polling-budget test, the affected modules passed again: 120 hook tests and 46 operator tests.
• The system test was rerun successfully against a real Databricks workspace with OAuth, covering synchronous, Common AI and deferrable invocation, before the final test-only changes.
• The Databricks documentation build passed after clarifying timeout guarantees. Ruff and commit hooks passed for the final test changes. Broader checks are left to CI.

Earlier validation included the full Databricks suite (1087 passed, 12 skipped), 142 agent tests with the minimum tenacity==8.3.0, five local checks against the Databricks 0.4.0 runtime with simulated OAuth, targeted mypy, static/manual checks, spellcheck and the dependency lower-bound checker. Those broader and minimum-dependency checks have not been repeated for the latest changes.

See the fixture README for deployment, configuration and cleanup. The fixture uses an in-memory store and does not test recovery across app restarts. With the app and configuration prepared:

SYSTEM_TESTS_ENV_ID=your-unique-id \
BREEZE_INIT_COMMAND='set -a; . /files/databricks-agent-system-test.env; set +a' \
breeze testing system-tests \
  --backend sqlite \
  --forward-credentials \
  --test-timeout 2400 \
  providers/databricks/tests/system/databricks/example_databricks_agent.py \
  -q

---

Was generative AI tooling used to co-author this PR?

• Yes — Codex (GPT-6)

Generated-by: Codex (GPT-6) following the guidelines

[AlejandroMorgante]

@SameerMesiah97 could you please help me with a review when you have some time? Thanks! 🙏

[AlejandroMorgante]

End-to-end validation

Validated commit 68fc3d9b60d8c6f24ef4e074d6bf052d09fd5ff2 against a real Databricks App running the deterministic agent fixture.

Both synchronous and deferrable invocations completed successfully. Verification tasks confirmed the response content, original input, and session ID. The deferrable task resumed successfully after the trigger fired.

Environment: Airflow 3.3.0, Python 3.12, PostgreSQL, and LocalExecutor. The test loaded the provider source from the PR commit. No model endpoint was used.

Successful Airflow run

Synchronous invocation

Deferrable invocation

Agent input and output

Databricks App

Databricks invocation logs

The temporary OAuth secret was revoked and the Databricks App stopped after validation.

[kaxil]

Approving, the earlier points are all sorted. None of these need to hold the merge:

• In the text extraction at hooks/agent.py:253-267, one non-text block in the last assistant message hits the break, and text ends up as the JSON dump of the whole transcript. ChatDatabricks with use_responses_api=True puts reasoning and mcp_call blocks right next to the answer, so this can show up with ordinary agents. I'd skip those blocks and only fall back when no text is found.
• agent.rst:155-157 should say a literal invocation_id has to be unique per Dag run. Otherwise the next run gets the stored answer back, or a 409 if the prompt changed.
• hooks/agent.py:223-226 is only reachable because test_reraised_http_error_remains_transient sets retry_args["reraise"] itself; the constructor doesn't accept retry_args. I'd drop the branch and the test.
• The operator fixture's submit reply has no status, so test_defers never hits the != "completed" check with a value the server actually sends (queued/active).
• The hook tests sleep a real second per retry, since wait_exponential ignores retry_delay=0 on the first wait. Stubbing retry_args["sleep"] in the fixture covers the whole file.
• caller isn't passed through, so every request goes out as operator/Unknown.

Also curious why this is limited to client-secret service principals. _get_token already gets workspace OAuth tokens for the federated modes from the same endpoint.

[AlejandroMorgante]

@kaxil

Thanks for the thorough review and all the suggestions!

Applied all the follow-ups: mixed-content text extraction, invocation ID documentation, removal of the reraise=True branch and test, realistic deferral states, mocked retry sleeps, caller propagation, and OAuth federation through the base hook.

Validated client-secret OAuth, AWS federation, a custom OIDC token provider, and Kubernetes federation using a local service-account token. The system test passed again after rebasing, covering synchronous, managed-agent, and deferrable execution. Pre-commit checks passed.