Skip to content

Add a vendor-neutral managed-agent hook contract to Common AI - #73532

Draft
kaxil wants to merge 2 commits into
apache:mainfrom
astronomer:managed-agent-hook-contract
Draft

kaxil wants to merge 2 commits into
apache:mainfrom
astronomer:managed-agent-hook-contract

Conversation

@kaxil

@kaxil kaxil commented Sep 22, 2026

Copy link
Copy Markdown
Member

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. BaseManagedAgentToolset from #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 DbApiHook put it: the hook. airflow.providers.common.ai.managed_agents defines BaseManagedAgentHook, a mixin a vendor hook adopts next to its own base. The agent is an argument, like a statement is to DbApiHook.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) and FailoverManagedAgentClient (clients composed across clouds).

claims = BedrockAgentCoreManagedAgentHook(aws_conn_id="aws_prod", region_name="us-east-1").agent(RUNTIME_ARN)
AgentOperator(
    task_id="triage",
    llm_conn_id="anthropic_default",
    prompt="Review claim 4411 and decide whether to pay it.",
    toolsets=[ManagedAgentToolset(claims, tool_name="ask_claims_agent", description="Reviews an insurance claim.")],
)

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 raising AirflowOptionalProviderFeatureException, an extra in pyproject.toml, pytest.importorskip in the tests. The two adoptions here follow aws/queues/sqs.py. Plain amazon installs are untouched; pydantic-ai only arrives with amazon[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_options in and raw out. Each hook rejects options that would point the call somewhere else (another ARN, account, engine). The model sees text. Python callers get raw.

Why ManagedAgentRejected instead of ModelRetry. 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's INVALID_ARGUMENT is 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 FailoverManagedAgentToolset took toolsets as members but only ever called invoke on them, so a member's tool_name, description and max_retries did 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 a session_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.

replayable now does something. #71946 declared it; nothing read it; durable=True replayed managed-agent calls from cache regardless. CachingToolset now runs the tool instead of replaying when the wrapped toolset says replayable=False, and it looks through prefixed, filtered and CombinedToolset, 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:

  • FailoverManagedAgentToolset is gone. Use ManagedAgentToolset(FailoverManagedAgentClient([...]), ...).
  • BaseManagedAgentToolset.agent_ref returns a ManagedAgentRef rather than a dict[str, str].
  • managed_agent.served loses role and position, gains tool. managed_agent.failover loses tool. served counts answers, failover counts 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 version marker for the release manager.

Gotchas

  • Bedrock: one boto3 client per distinct request timeout, reused, so the toolset's timeout becomes connect/read timeout without a new session per tool call. Retries off unless the connection or caller set them. The connection's config_kwargs apply 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.
  • Agent Engine: the query path keeps no conversation state, so a session_id is refused rather than silently sent as a fresh call. Query jobs stay on RunQueryJobOperator, which can defer.
  • Not touched: Azure Foundry, Cortex, Anthropic. Cortex has PRs in flight and is the obvious next adoption. Anthropic's session-shaped agent wants an operator, not a toolset.

Supersedes the toolset-only design of #71946.


  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {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.

@boring-cyborg boring-cyborg Bot added area:providers kind:documentation provider:amazon AWS/Amazon - related issues provider:common-ai provider:google Google (including GCP) related issues labels Sep 22, 2026
@kaxil
kaxil force-pushed the managed-agent-hook-contract branch from 959f796 to f71663e Compare September 22, 2026 10:14
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

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:providers kind:documentation provider:amazon AWS/Amazon - related issues provider:common-ai provider:google Google (including GCP) related issues

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant