Repository navigation
Add Databricks agent invocation operator and hook - #74386
Open
AlejandroMorgante wants to merge 18 commits into
Open
AlejandroMorgante wants to merge 18 commits into
AlejandroMorgante wants to merge 18 commits into
Conversation
Contributor
Author
|
@SameerMesiah97 could you please help me with a review when you have some time? Thanks! 🙏 |
AlejandroMorgante
marked this pull request as draft
October 9, 2026 19:31
AlejandroMorgante
marked this pull request as ready for review
October 9, 2026 19:31
kaxil
reviewed
Oct 9, 2026
kaxil
reviewed
Oct 9, 2026
AlejandroMorgante
force-pushed
the
add-databricks-agent-invoke
branch
from
October 9, 2026 23:25
152283d to
e13d7dd
Compare
kaxil
reviewed
Oct 10, 2026
AlejandroMorgante
marked this pull request as draft
October 10, 2026 15:25
AlejandroMorgante
force-pushed
the
add-databricks-agent-invoke
branch
from
October 10, 2026 15:28
7266b14 to
7f3b704
Compare
Contributor
Author
End-to-end validationValidated commit 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 runSynchronous invocationDeferrable invocationAgent input and outputDatabricks AppDatabricks invocation logsThe temporary OAuth secret was revoked and the Databricks App stopped after validation. |
AlejandroMorgante
marked this pull request as ready for review
October 11, 2026 01:05
kaxil
reviewed
Oct 13, 2026
AlejandroMorgante
force-pushed
the
add-databricks-agent-invoke
branch
from
October 13, 2026 20:22
68fc3d9 to
5416249
Compare
kaxil
reviewed
Oct 17, 2026
kaxil
reviewed
Oct 17, 2026
AlejandroMorgante
force-pushed
the
add-databricks-agent-invoke
branch
2 times, most recently
from
October 20, 2026 01:23
975f4c8 to
f3522fd
Compare
kaxil
reviewed
Oct 21, 2026
AlejandroMorgante
force-pushed
the
add-databricks-agent-invoke
branch
from
October 22, 2026 02:44
f3522fd to
1983021
Compare
Allow Airflow tasks to invoke DurableAgentServer agents deployed on Databricks Apps and wait without occupying a worker. Generated-by: Codex (GPT-6)
Keep the new operator documentation valid for Sphinx and recognize the official MLflow spelling. Generated-by: Codex (GPT-6)
Keep the operator guide focused on user configuration and invocation, following the existing Databricks documentation conventions. Generated-by: Codex (GPT-6)
A completed response must not turn an expired invocation wait into success. Recoverable async disconnects should not fail a task immediately, and polling failures need safe context for diagnosis. Generated-by: Codex (GPT-6)
Common AI consumers need to consult Databricks agents through the same contract as Vertex AI and Bedrock. Existing background and deferrable workflows must remain usable without installing the optional Common AI provider. Generated-by: Codex (GPT-6)
Replacement mocks should follow Airflow's testing conventions so they remain tied to the API they represent. Generated-by: Codex (GPT-6)
Mocked responses must reflect the Databricks runtime contract so session, interruption, failure and idempotency defects cannot hide behind passing tests. Non-default settings and unconditional assertions make routing and polling regressions observable. This draft PR step intentionally exposes failures pending implementation fixes in the following commits. Generated-by: Codex (GPT-6)
Stored failures and template interruptions must not be treated as transient errors or completed answers. Task retries need to respect the server's idempotency rules, and toolset calls need a session. Generated-by: Codex (GPT-6)
Conditional expectations made individual regression cases harder to assess during review. Generated-by: Codex (GPT-6)
Make agent task progress traceable from submission through completion without exposing request or response contents. Generated-by: Codex (GPT-6)
Completed background submissions omit the agent output, explicit request budgets can exceed the default HTTP timeout, and netrc credentials must not replace service principal authentication. Generated-by: Codex (GPT-6)
Common AI consumers need the final assistant answer rather than a JSON envelope containing intermediate tool calls and their results. The official OpenAI and LangGraph templates serialize assistant messages differently. Generated-by: Codex (GPT-6)
Known failed invocations must fail the task even when waiting is disabled. Common AI users also need the supported dependency floor and guidance on agent tools and timeouts to avoid unintended duplicate invocations. Generated-by: Codex (GPT-6)
The OpenAI template permits SDK versions that reject Common AI text blocks. String message content matches the documented Databricks input format and is accepted by both OpenAI and LangGraph templates. Generated-by: Codex (GPT-6)
Status requests should not consume the entire invocation budget, and connection-only hooks should support background invocations across apps. Generated-by: Codex (GPT-6)
Network activity and OAuth refresh can delay deadline detection, so callers should not interpret the invocation budget as a strict wall-clock cutoff. Generated-by: Codex (GPT-6)
Guard against losing async HTTP options or restarting the polling budget between status requests. Generated-by: Codex (GPT-6)
AlejandroMorgante
force-pushed
the
add-databricks-agent-invoke
branch
from
October 22, 2026 02:58
1983021 to
89fdc43
Compare
Ensure alternate Tenacity error propagation preserves transient HTTP failures instead of treating them as terminal agent failures. Generated-by: Codex (GPT-6)
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.






Add
DatabricksAgentInvokeOperatorandDatabricksAgentHookto invoke agents hosted withDurableAgentServeron Databricks Apps. Supports service principal OAuth, background submission, synchronous polling and deferrable waiting, preserving the full invocation response in XCom.The hook implements
BaseManagedAgentHookfor 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-bridgerepository, tagdatabricks-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:
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 \ -qWas generative AI tooling used to co-author this PR?
Generated-by: Codex (GPT-6) following the guidelines