Conversation
kaxil
force-pushed
the
managed-agent-hook-contract
branch
from
September 22, 2026 10:14
959f796 to
f71663e
Compare
Move the Bedrock hook's text_key and max_response_bytes from the constructor into vendor_options, as AWS hooks must keep a thin constructor. Patch Stats.incr rather than the Stats class so the tests hold on Airflow 3.0, type the toolset test RunContext, and add common.ai to the amazon and google dependency lists that the selective-checks tests expect.
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.
Airflow has five ways to invoke a vendor-hosted agent (Bedrock AgentCore, Vertex AI Agent Engine, Azure AI Foundry, Snowflake Cortex, Anthropic Managed Agents) and none of them share an argument name, a result shape or an error class.
BaseManagedAgentToolsetfrom #71946 was supposed to fix that. A month later it has zero implementations. The reason is structural: adopting it means subclassing a pydantic-ai toolset from inside a provider that still supports Airflow 2, which nobody can do.So this moves the common surface to where
DbApiHookput it: the hook.airflow.providers.common.ai.managed_agentsdefinesBaseManagedAgentHook, a mixin a vendor hook adopts next to its own base. The agent is an argument, like a statement is toDbApiHook.run. Three methods (resolve_agent,agent_capabilities,invoke_agent), typed request/response dataclasses, no pydantic-ai import. Amazon adopts it for AgentCore runtimes, Google for Agent Engine deployments. On top, written once:ManagedAgentToolset(any client of the contract becomes one tool for a calling model) andFailoverManagedAgentClient(clients composed across clouds).Design rationale
Vendors take common.ai as an optional extra, exactly as they take common.messaging today. common.ai needs Airflow 3.0; amazon and google still floor 2.11. There is already a pattern in the repo for a 3.x-only common base adopted by 2.x providers:
BaseMessageQueueProvider. Guarded import raisingAirflowOptionalProviderFeatureException, an extra inpyproject.toml,pytest.importorskipin the tests. The two adoptions here followaws/queues/sqs.py. Plainamazoninstalls are untouched; pydantic-ai only arrives withamazon[common.ai].Type the intersection, pass the rest through. Identity, prompt or messages, an optional session, a timeout, a text answer: those mean the same thing everywhere and are typed. The rest rides in
vendor_optionsin andrawout. Each hook rejects options that would point the call somewhere else (another ARN, account, engine). The model seestext. Python callers getraw.Why
ManagedAgentRejectedinstead ofModelRetry. The contract has to stay free of pydantic-ai or the guarded import stops being cheap. The toolset translates at the boundary. As it turns out neither adopted platform has a rephrase-class error anyway: AgentCore returns a container's complaints inside a 200, and Agent Engine'sINVALID_ARGUMENTis an author mistake (wrong input key), not something a model can fix by rewording. Both hooks raise terminal errors or let transient ones through.Failover over clients, not toolsets. The removed
FailoverManagedAgentToolsettook toolsets as members but only ever calledinvokeon them, so a member'stool_name,descriptionandmax_retriesdid nothing. It also resolved every member's identity before the call, so a standby with a bad connection failed calls the primary would have answered. The new group takes clients, resolves identity only for logs and metric tags, logs which member answered, and refuses any request carrying asession_id(a failover starts a fresh conversation on the standby). A bound agent applies the same rule for a single hook whose capabilities do not include sessions.replayablenow does something. #71946 declared it; nothing read it;durable=Truereplayed managed-agent calls from cache regardless.CachingToolsetnow runs the tool instead of replaying when the wrapped toolset saysreplayable=False, and it looks throughprefixed,filteredandCombinedToolset, which do not carry the attribute.Migration path
common.ai is 0.x and the old base had no subclasses in tree. Still, for anyone who tried it:
FailoverManagedAgentToolsetis gone. UseManagedAgentToolset(FailoverManagedAgentClient([...]), ...).BaseManagedAgentToolset.agent_refreturns aManagedAgentRefrather than adict[str, str].managed_agent.servedlosesroleandposition, gainstool.managed_agent.failoverlosestool.servedcounts answers,failovercounts transitions; read them next to each other. Which member answered is in the task log.The new amazon and google extras carry the
# use next versionmarker for the release manager.Gotchas
timeoutbecomes connect/read timeout without a new session per tool call. Retries off unless the connection or caller set them. The connection'sconfig_kwargsapply as for any AWS hook. Session ids are checked against AgentCore's 33 to 256 character bound before the call. Non-JSON, malformed or oversize bodies are terminal errors that name the agent and connection.session_idis refused rather than silently sent as a fresh call. Query jobs stay onRunQueryJobOperator, which can defer.Supersedes the toolset-only design of #71946.
{pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.