Skip to content

Stamp Airflow run identity onto agent runs and traces - #73275

Merged
pankajkoti merged 2 commits into
apache:mainfrom
astronomer:stamp-airflow-run-identity-onto-agent-runs
Sep 17, 2026
Merged

pankajkoti merged 2 commits into
apache:mainfrom
astronomer:stamp-airflow-run-identity-onto-agent-runs

Conversation

@pankajkoti

@pankajkoti pankajkoti commented Sep 17, 2026 •

Copy link
Copy Markdown
Member

Give an agent run and the Airflow task instance that produced it a shared
identifier, so a task's output can be joined to its GenAI trace without
reconstructing it from timestamps.

AgentOperator (and @task.agent) now:

  • pass the task-instance id (unique per attempt, since Airflow regenerates it
    on each retry) as the pydantic-ai run_id, so it surfaces on
    AgentRunResult.run_id, the run's messages, and gen_ai.agent.call.id on
    the agent-run span
  • stamp the Airflow identity on every emitted GenAI span (airflow.dag_id,
    airflow.task_id, airflow.dag_run.run_id, airflow.task_instance.try_number,
    airflow.task_instance.map_index, airflow.task_instance.id), so a span is
    filterable by those directly rather than by walking up to the parent task span
  • expose the run_id and token usage on XCom for a downstream task

Relies on the pydantic-ai run_id= parameter (pydantic-ai#6615, added in
v2.16.0). The provider already floors pydantic-ai-slim at >=2.23.0, which
includes it, so no dependency change is needed.


Was generative AI tooling used to co-author this PR?
  • Yes — Claude Code (Opus 4.8)

Generated-by: Claude Code (Opus 4.8) following the guidelines


🤖 Generated with Claude Code

An agent run and the task instance that produced it shared no identifier,
so a GenAI trace in an OTel backend could not be tied back to its task and
a task had no handle on its own run. Observability surfaces had to
reconstruct that join from timestamps. Giving the two a shared, per-attempt
key makes the join exact and unblocks a downstream AI traces UI.
@pankajkoti
pankajkoti force-pushed the stamp-airflow-run-identity-onto-agent-runs branch from 593b003 to 4a6cd33 Compare September 17, 2026 10:13
@pankajkoti

Copy link
Copy Markdown
Member Author

Snapshot from Jaeger UI

Screenshot 2026-09-17 at 4 22 57 PM

@pankajkoti
pankajkoti marked this pull request as ready for review September 17, 2026 11:02
Comment thread providers/common/ai/src/airflow/providers/common/ai/operators/agent.py Outdated
Comment thread providers/common/ai/src/airflow/providers/common/ai/observability.py Outdated
The first pass mutated a possibly caller-owned InstrumentationSettings object
in place, so a settings instance shared across agent tasks or reused across HITL
re-runs would carry the wrong run's identity and accumulate nested tracer
wrappers. It also wrote run metadata to XCom regardless of do_xcom_push, guarded
a task-instance id that is non-nullable at runtime, and omitted the run cost the
docstring already promised. These changes keep the stamping per-run, honour the
operator's XCom suppression flag, drop the dead guard, and surface cost.
@pankajkoti
pankajkoti merged commit 833b61f into apache:main Sep 17, 2026
83 checks passed
@pankajkoti
pankajkoti deleted the stamp-airflow-run-identity-onto-agent-runs branch September 17, 2026 14:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants