Allow templated connection IDs in agent toolsets

[kaxil]

The connection IDs of Common AI agent toolsets are now Jinja templates, rendered for each task instance just before it runs. One toolset definition can then reach a different system depending on the run:

• Per environment: the same Dag reads the staging warehouse in staging and the production one in production, with the environment in a Variable: SQLToolset(db_conn_id="warehouse_{{ var.value.environment }}").
• Per unit of work: a mapped @task.agent gives each map index its own connection, e.g. one per customer for customer-facing analytics, where each customer's rows sit behind their own database role:

@task.agent(
    llm_conn_id="pydanticai_default",
    toolsets=[SQLToolset(db_conn_id="analytics_{{ task.op_kwargs.customer }}")],
)
def report(customer: str) -> str:
    return f"Summarize this month's orders for {customer}."


report.expand(customer=customers())

Until now toolsets took their connection when the Dag was parsed. @task.llm_sql already templates db_conn_id, but an agent's toolsets could not, so every mapped agent shared one connection.

What is templated: SQLToolset.db_conn_id, MCPToolset.mcp_conn_id, and a HookToolset's hook connection ID. HookToolset reads the attribute the hook's conn_name_attr names (postgres_conn_id, ...), and falls back to conn_id for hooks such as WasbHook that keep it there. Only connection IDs are templated.

Design rationale

Why not add toolsets to AgentOperator.template_fields? Template fields are serialized into the Dag. A toolset's repr is what gets serialized, and pydantic-ai's wrappers (.prefixed(), .filtered()) are dataclasses whose repr can embed function addresses, so the Dag hash would change on every parse. toolsets stays out of template_fields, and the operator renders the toolsets itself.

Each task instance renders a copy. MappedOperator.unmap passes the partial's toolset objects straight to every unmapped task, and dag.test() runs every task in one process. Rendering in place would hand map index 0's connection to map index 1. Each opt-in leaf toolset is copied before rendering (for HookToolset, the hook too). Wrappers and Toolset capabilities are walked with pydantic-ai's visit_and_replace, and only when something in them is templated, so an untemplated custom wrapper is never rebuilt.

The opt-in attribute is agent_template_fields, not template_fields. Airflow's templater renders any object that carries template_fields in place, wherever it is nested inside another template field. agent_params is a template field, and agent_params["toolsets"] is a supported way to pass toolsets, so the familiar name would bring the leak back through that path. Third-party toolsets opt in by declaring the same attribute.

Rendering hangs off _do_render_template_fields. A mapped task never calls render_template_fields: MappedOperator renders the unmapped task through _do_render_template_fields. KubernetesPodOperator hooks the same method for the same reason.

The rendered connection is not recorded anywhere else, so each task instance logs it once, e.g. Rendered toolset sql-analytics_acme.

Screenshots

A two-customer demo on a real scheduler and API server, each customer with its own SQLite connection, and a test model that calls every tool. Map index 0 renders the acme connection for both the SQL toolset and the hook toolset; map index 1 renders globex:

Each agent's tool results come from its own customer's database:

Gotchas

• Build the connection ID from values the Dag controls (a Variable, upstream task output), not params or dag_run.conf. A task can read any connection it names, so a template driven by trigger input lets whoever triggers the Dag pick the database (for an MCP stdio connection, the command that runs on the worker). The docs say so next to each example.
• {{ customer }} does not work: the task's arguments are not template variables. Use {{ task.op_kwargs.customer }}. With AgentOperator.partial(...).expand(prompt=...), the connection has to come from the map index; the docs show that form and its ordering caveat.
• HookToolset.id now includes the connection ID (hook-PostgresHook-analytics_acme, previously hook-PostgresHook). The toolset id is part of the durable-execution step fingerprint, so a durable=True task that retries across the upgrade misses its cache once.
• Not templated: allowed_tables (validated when the toolset is created, so a template stays a literal), DataFusionToolset, a Toolset capability built from a callable, and hooks that keep their connection ID under some other attribute. A hook that looks its connection up in __init__ fails at Dag parse time, because the template is not a connection ID yet.

---

• 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.