This is an automated email from the ASF dual-hosted git repository.
pankajkoti pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 833b61fadc1 Stamp Airflow run identity onto agent runs and traces
(#73275)
833b61fadc1 is described below
commit 833b61fadc1ff4ec9b3a915bccbd4d35c7e3df2e
Author: Pankaj Koti <[email protected]>
AuthorDate: Thu Sep 17 19:59:04 2026 +0530
Stamp Airflow run identity onto agent runs and traces (#73275)
* Stamp Airflow run identity onto agent runs and traces
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.
* Address review feedback on agent run identity stamping
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.
---
providers/common/ai/docs/observability.rst | 28 +++-
.../airflow/providers/common/ai/observability.py | 85 ++++++++++-
.../airflow/providers/common/ai/operators/agent.py | 52 ++++++-
.../tests/unit/common/ai/decorators/test_agent.py | 42 ++++--
.../tests/unit/common/ai/operators/test_agent.py | 156 +++++++++++++++++----
.../ai/tests/unit/common/ai/test_observability.py | 88 ++++++++++++
6 files changed, 401 insertions(+), 50 deletions(-)
diff --git a/providers/common/ai/docs/observability.rst
b/providers/common/ai/docs/observability.rst
index 5715b58be5b..918c6f7ffd3 100644
--- a/providers/common/ai/docs/observability.rst
+++ b/providers/common/ai/docs/observability.rst
@@ -41,13 +41,35 @@ How it works
configured under ``[traces]`` / the standard ``OTEL_EXPORTER_OTLP_*``
environment variables. If core tracing is not enabled in the worker process,
no GenAI spans are emitted.
-* **Correlation is automatic.** The worker opens a task span before the
operator
- runs, so the agent's spans nest under it and inherit the task's ``trace_id``
- and ``airflow.*`` attributes (dag id, run id, task id, try number, map
index).
+* **Correlation.** The worker opens a task span before the operator runs, so
the
+ agent's spans nest under it and share its ``trace_id``.
+ :class:`~airflow.providers.common.ai.operators.agent.AgentOperator` (and
+ ``@task.agent``) additionally stamps the task-instance identity on every
GenAI
+ span it emits: the five keys core tracing already puts on the task span
+ (``airflow.dag_id``, ``airflow.task_id``, ``airflow.dag_run.run_id``,
+ ``airflow.task_instance.try_number``, ``airflow.task_instance.map_index``)
plus
+ ``airflow.task_instance.id`` as the per-attempt run join key. So a span is
+ filterable by dag, task, run, attempt, or map index directly, without walking
+ up to the parent span (OpenTelemetry children inherit trace context, not
+ attributes).
An automatic retry reuses the task instance's persisted trace context, so all
attempts share one trace and appear as repeated task-run spans on it,
distinguished by ``try number``. Only a manual clear or rerun regenerates the
context and starts a new trace.
+* **Run join key.** For an ``AgentOperator`` run, the task-instance id (unique
+ per attempt, since Airflow regenerates it on each retry) is passed to
+ pydantic-ai as the run's ``run_id``. It surfaces on the run's GenAI spans as
+ ``gen_ai.agent.call.id``, and the operator also exposes it, alongside the
run's
+ token usage, on XCom under the ``run_id`` and ``usage`` keys. A downstream
task
+ can then reference the run
+ (``ti.xcom_pull(task_ids="my_agent", key="run_id")``) and a trace backend can
+ join a task's output to its agent trace without parsing logs. With
+ ``enable_hitl_review`` the ``run_id`` and ``usage`` reflect the initial model
+ run, not the human-feedback regenerations.
+* **Scope.** The ``airflow.*`` identity attributes and the ``run_id`` /
``usage``
+ XComs come only from ``AgentOperator`` and ``@task.agent``. The other LLM
+ operators still emit GenAI spans correlated to the task span by nesting, but
+ without the identity attributes or the run join key.
* **Content is off by default.** Only token counts, model id, latency, tool
names, and finish reason are recorded. Prompt and completion text is never
emitted unless you opt in (see below).
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/observability.py
b/providers/common/ai/src/airflow/providers/common/ai/observability.py
index 08a59c40595..25e4a8712b7 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/observability.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/observability.py
@@ -31,11 +31,17 @@ or not configured in this process) no spans are emitted.
from __future__ import annotations
-from typing import TYPE_CHECKING, Literal
+import copy
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Literal, cast
from airflow.providers.common.compat.sdk import conf
if TYPE_CHECKING:
+ from collections.abc import Iterator
+
+ from opentelemetry.trace import Span, Tracer
+ from pydantic_ai import Agent
from pydantic_ai.models.instrumented import InstrumentationSettings
SECTION = "common.ai"
@@ -106,3 +112,80 @@ def genai_instrumentation_settings() ->
InstrumentationSettings | None:
include_binary_content=False,
tracer_provider=provider,
)
+
+
+def build_run_identity_attributes(ti: Any) -> dict[str, Any]:
+ """
+ Build the Airflow identity attributes to stamp on a run's GenAI spans.
+
+ Reuses core's task-span attribute keys (see ``_make_task_span``) so agent
+ spans filter identically to the task span they nest under, plus the
+ per-attempt task-instance id as the run join key carried on every span.
+ """
+ return {
+ "airflow.dag_id": ti.dag_id,
+ "airflow.task_id": ti.task_id,
+ "airflow.dag_run.run_id": ti.run_id,
+ "airflow.task_instance.try_number": ti.try_number,
+ "airflow.task_instance.map_index": ti.map_index if ti.map_index is not
None else -1,
+ "airflow.task_instance.id": str(ti.id),
+ }
+
+
+def stamp_identity_on_agent_spans(agent: Agent, attributes: dict[str, Any]) ->
None:
+ """
+ Stamp *attributes* on every GenAI span *agent* emits during its run.
+
+ pydantic-ai opens all of a run's agent/model/tool spans from
+ ``InstrumentationSettings.tracer``, so wrapping that one tracer reaches
them
+ all without touching the shared core ``TracerProvider``. No-op when the
+ agent is not instrumented with an ``InstrumentationSettings`` (tracing off,
+ or the caller supplied its own non-settings ``instrument`` value).
+
+ The settings object may be caller-owned and shared across agents (a single
+ module-level ``InstrumentationSettings`` handed to several tasks) or reused
+ across HITL re-runs. Mutating it in place would leak one run's identity
into
+ another and nest ``_IdentityTracer`` wrappers on each stamp, so we wrap a
copy
+ and swap it onto this agent, leaving the original untouched.
+ """
+ from pydantic_ai.models.instrumented import InstrumentationSettings
+
+ instrument = agent.instrument
+ if isinstance(instrument, InstrumentationSettings):
+ # ``tracer`` is not a constructor arg, so copy then set the attribute.
+ settings = copy.copy(instrument)
+ # _IdentityTracer implements the Tracer surface structurally (it cannot
+ # subclass Tracer without importing opentelemetry at module load).
+ settings.tracer = cast("Tracer", _IdentityTracer(instrument.tracer,
attributes))
+ agent.instrument = settings
+
+
+class _IdentityTracer:
+ """OpenTelemetry ``Tracer`` wrapper that stamps fixed attributes on every
span it starts."""
+
+ def __init__(self, tracer: Tracer, attributes: dict[str, Any]) -> None:
+ self._tracer = tracer
+ self._attributes = attributes
+
+ def start_span(self, *args: Any, **kwargs: Any) -> Span:
+ span = self._tracer.start_span(*args, **kwargs)
+ span.set_attributes(self._attributes)
+ return span
+
+ @contextmanager
+ def start_as_current_span(self, *args: Any, **kwargs: Any) ->
Iterator[Span]:
+ with self._tracer.start_as_current_span(*args, **kwargs) as span:
+ span.set_attributes(self._attributes)
+ yield span
+
+ def __getattr__(self, name: str) -> Any:
+ # Delegate any other Tracer method to the wrapped tracer. Dunders are
not
+ # forwarded (so copy/pickle of the settings does not silently unwrap
the
+ # stamping), and the ``__dict__`` lookup guards against recursion
before
+ # ``_tracer`` is set.
+ if name.startswith("__"):
+ raise AttributeError(name)
+ tracer = self.__dict__.get("_tracer")
+ if tracer is None:
+ raise AttributeError(name)
+ return getattr(tracer, name)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
index 612b7027d3b..e5d1212dad5 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
@@ -29,6 +29,10 @@ from pydantic import BaseModel
from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
from airflow.providers.common.ai.mixins.hitl_review import HITLReviewMixin
+from airflow.providers.common.ai.observability import (
+ build_run_identity_attributes,
+ stamp_identity_on_agent_spans,
+)
from airflow.providers.common.ai.utils.logging import log_run_summary,
wrap_toolsets_for_logging
from airflow.providers.common.ai.utils.output_type import
rehydrate_pydantic_output
from airflow.providers.common.ai.utils.usage import coerce_usage_limits
@@ -126,6 +130,13 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
and run the agent. The agent reasons about the prompt, calls tools in a
multi-turn loop, and returns a final answer.
+ Alongside the returned agent output, the run's ``run_id`` and token
``usage``
+ are pushed to XCom under the ``run_id`` and ``usage`` keys, so a downstream
+ task can reference the run and its cost. The ``run_id`` also ties the task
to
+ its GenAI trace (see the provider's observability docs). With
+ ``enable_hitl_review``, these reflect the initial model run, not the
+ human-feedback regenerations.
+
:param prompt: The prompt to send to the agent.
:param llm_conn_id: Connection ID for the LLM provider.
:param model_id: Model identifier (e.g. ``"openai:gpt-5"``).
@@ -470,7 +481,14 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
agent = self._build_agent()
- run_kwargs: dict[str, Any] = {"usage_limits": usage_limits}
+ ti = context["task_instance"]
+ self._run_identity_attrs = build_run_identity_attributes(ti)
+ stamp_identity_on_agent_spans(agent, self._run_identity_attrs)
+
+ # The task-instance id is non-nullable and regenerated on each retry,
so it
+ # is a unique, reverse-resolvable join key. It lands on result.run_id,
the
+ # run's messages, and the ``gen_ai.agent.call.id`` span attribute.
+ run_kwargs: dict[str, Any] = {"usage_limits": usage_limits, "run_id":
str(ti.id)}
history = self._resolve_message_history()
if history is not None:
run_kwargs["message_history"] = history
@@ -492,6 +510,7 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
result = agent.run_sync(self.prompt, **run_kwargs)
log_run_summary(self.log, result)
+ self._emit_run_metadata(context, result)
if self._durable_counter is not None:
c = self._durable_counter
@@ -535,10 +554,10 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
output = output.model_dump()
# Clean up the durable cache only after the run and every post-run step
- # that can still fail (the message-history XCom push above and output
- # serialization) has succeeded. Cleaning up earlier and then raising
- # would leave the Airflow retry with an empty cache, re-executing every
- # already-completed model and tool step.
+ # that can still fail (the run-metadata and message-history XCom pushes
+ # above and output serialization) has succeeded. Cleaning up earlier
and
+ # then raising would leave the Airflow retry with an empty cache,
+ # re-executing every already-completed model and tool step.
if self._durable_storage is not None:
self._durable_storage.cleanup()
return output
@@ -573,10 +592,33 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
transcript =
ModelMessagesTypeAdapter.dump_json(result.all_messages()).decode()
context["task_instance"].xcom_push(key="message_history",
value=transcript)
+ def _emit_run_metadata(self, context: Context, result: Any) -> None:
+ """Expose the pydantic-ai run id and token usage on XCom for
downstream tasks."""
+ if not self.do_xcom_push:
+ return
+ usage = result.usage
+ ti = context["task_instance"]
+ ti.xcom_push(key="run_id", value=result.run_id)
+ ti.xcom_push(
+ key="usage",
+ value={
+ "requests": usage.requests,
+ "input_tokens": usage.input_tokens,
+ "output_tokens": usage.output_tokens,
+ "total_tokens": usage.total_tokens,
+ "tool_calls": usage.tool_calls,
+ # Decimal | None, stringified so XCom serialization stays
lossless.
+ "cost": str(usage.cost) if usage.cost is not None else None,
+ },
+ )
+
def regenerate_with_feedback(self, *, feedback: str, message_history: Any)
-> tuple[str, Any]:
"""Re-run the agent with *feedback* appended to the conversation
history."""
usage_limits = coerce_usage_limits(self.usage_limits)
agent = self._build_agent()
+ identity = getattr(self, "_run_identity_attrs", None)
+ if identity:
+ stamp_identity_on_agent_spans(agent, identity)
messages = message_history or []
result = agent.run_sync(
feedback,
diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
index 1731805fe11..00fd431b66f 100644
--- a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
@@ -40,6 +40,16 @@ class Summary(BaseModel):
text: str
+def _make_context():
+ """A context whose ``task_instance`` carries a stable id for run-identity
stamping.
+
+ A real dict (not a mock) so the decorator's ``context_merge`` of
``op_kwargs`` works.
+ """
+ ti = MagicMock()
+ ti.configure_mock(id="ti-1", dag_id="dag", task_id="task", run_id="run",
map_index=-1, try_number=1)
+ return {"task_instance": ti}
+
+
class TestAgentDecoratedOperator:
def test_custom_operator_name(self):
assert _AgentDecoratedOperator.custom_operator_name == "@task.agent"
@@ -47,7 +57,7 @@ class TestAgentDecoratedOperator:
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_calls_callable_and_returns_output(self, mock_hook_cls,
make_mock_run_result):
"""The callable's return value becomes the agent prompt."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("The top
customer is Acme Corp.")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -55,11 +65,13 @@ class TestAgentDecoratedOperator:
return "Who is our top customer?"
op = _AgentDecoratedOperator(task_id="test",
python_callable=my_prompt, llm_conn_id="my_llm")
- result = op.execute(context={})
+ result = op.execute(context=_make_context())
assert result == "The top customer is Acme Corp."
assert op.prompt == "Who is our top customer?"
- mock_agent.run_sync.assert_called_once_with("Who is our top
customer?", usage_limits=None)
+ mock_agent.run_sync.assert_called_once_with(
+ "Who is our top customer?", usage_limits=None, run_id="ti-1"
+ )
@pytest.mark.parametrize(
"return_value",
@@ -79,7 +91,7 @@ class TestAgentDecoratedOperator:
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_accepts_sequence_prompt(self, mock_hook_cls,
make_mock_run_result):
"""A non-empty Sequence[UserContent] return value is forwarded to
run_sync as-is."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("ok")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -90,10 +102,10 @@ class TestAgentDecoratedOperator:
return prompt
op = _AgentDecoratedOperator(task_id="test",
python_callable=my_prompt, llm_conn_id="my_llm")
- op.execute(context={})
+ op.execute(context=_make_context())
assert op.prompt == prompt
- mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None)
+ mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None,
run_id="ti-1")
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_sequence_prompt_with_hitl_review_raises_before_run_sync(self,
mock_hook_cls):
@@ -103,7 +115,7 @@ class TestAgentDecoratedOperator:
if not AIRFLOW_V_3_1_PLUS:
pytest.skip("enable_hitl_review requires Airflow >= 3.1.0")
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = _AgentDecoratedOperator(
@@ -120,7 +132,7 @@ class TestAgentDecoratedOperator:
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_merges_op_kwargs_into_callable(self, mock_hook_cls,
make_mock_run_result):
"""op_kwargs are resolved by the callable to build the prompt."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("done")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -133,15 +145,17 @@ class TestAgentDecoratedOperator:
llm_conn_id="my_llm",
op_kwargs={"topic": "revenue trends"},
)
- op.execute(context={"task_instance": MagicMock()})
+ op.execute(context=_make_context())
assert op.prompt == "Analyze revenue trends"
- mock_agent.run_sync.assert_called_once_with("Analyze revenue trends",
usage_limits=None)
+ mock_agent.run_sync.assert_called_once_with(
+ "Analyze revenue trends", usage_limits=None, run_id="ti-1"
+ )
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_passes_toolsets_through(self, mock_hook_cls,
make_mock_run_result):
"""Toolsets passed to the decorator are forwarded to the agent."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("result")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -153,7 +167,7 @@ class TestAgentDecoratedOperator:
llm_conn_id="my_llm",
toolsets=[mock_toolset],
)
- op.execute(context={})
+ op.execute(context=_make_context())
create_call =
mock_hook_cls.get_hook.return_value.create_agent.call_args
passed_toolsets = create_call[1]["toolsets"]
@@ -165,7 +179,7 @@ class TestAgentDecoratedOperator:
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_structured_output(self, mock_hook_cls,
make_mock_run_result):
"""BaseModel output flows through XCom as the Pydantic instance."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value =
make_mock_run_result(Summary(text="Great results"))
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -175,7 +189,7 @@ class TestAgentDecoratedOperator:
llm_conn_id="my_llm",
output_type=Summary,
)
- result = op.execute(context={})
+ result = op.execute(context=_make_context())
assert isinstance(result, Summary)
assert result.text == "Great results"
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index 462f5fc70c9..da8360f24c5 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -66,13 +66,30 @@ class Summary(BaseModel):
score: float = 0.0
-def _make_mock_agent(output, make_mock_run_result):
+def _make_mock_agent(output, make_mock_run_result, *, cost=None):
"""Create a mock agent that returns the given output."""
- mock_agent = MagicMock(spec=["run_sync"])
- mock_agent.run_sync.return_value = make_mock_run_result(output)
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ mock_agent.run_sync.return_value = make_mock_run_result(output, cost=cost)
return mock_agent
+def _make_ti(*, id="ti-1", dag_id="dag", task_id="task", run_id="run",
map_index=-1, try_number=1):
+ """Return a task-instance double carrying the identity fields execute()
reads."""
+ ti = MagicMock()
+ ti.configure_mock(
+ id=id, dag_id=dag_id, task_id=task_id, run_id=run_id,
map_index=map_index, try_number=try_number
+ )
+ return ti
+
+
+def _make_context(ti=None):
+ """A context whose ``task_instance`` is a configured ti. Other keys stay
generic mocks."""
+ ti = ti if ti is not None else _make_ti()
+ ctx = MagicMock()
+ ctx.__getitem__.side_effect = lambda key: ti if key == "task_instance"
else MagicMock()
+ return ctx
+
+
PRICED_COST = Decimal("0.10")
@@ -201,9 +218,9 @@ class TestAgentOperatorExecute:
llm_conn_id="my_llm",
usage_limits=limits,
)
- op.execute(context=MagicMock())
+ op.execute(context=_make_context())
- mock_agent.run_sync.assert_called_once_with("run", usage_limits=limits)
+ mock_agent.run_sync.assert_called_once_with("run",
usage_limits=limits, run_id="ti-1")
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_coerces_usage_limits_dict_before_run_sync(self,
mock_hook_cls, make_mock_run_result):
@@ -301,7 +318,7 @@ class TestAgentOperatorExecute:
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_raises_valueerror_for_unparsable_usage_limits_value(self,
mock_hook_cls):
"""An unparsable templated ``usage_limits`` value surfaces as
``ValueError`` before any model call."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = AgentOperator(
@@ -360,14 +377,14 @@ class TestAgentOperatorExecute:
llm_conn_id="my_llm",
system_prompt="You are helpful.",
)
- result = op.execute(context=MagicMock())
+ result = op.execute(context=_make_context())
assert result == "The answer is 42."
mock_hook_cls.get_hook.assert_called_once_with("my_llm",
hook_params={"model_id": None})
mock_hook_cls.get_hook.return_value.create_agent.assert_called_once_with(
output_type=str, instructions="You are helpful."
)
- mock_agent.run_sync.assert_called_once_with("What is the answer?",
usage_limits=None)
+ mock_agent.run_sync.assert_called_once_with("What is the answer?",
usage_limits=None, run_id="ti-1")
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_passes_toolsets_in_agent_kwargs(self, mock_hook_cls,
make_mock_run_result):
@@ -567,7 +584,7 @@ class TestAgentOperatorExecute:
msg_history = [MagicMock()]
mock_result = make_mock_run_result("Initial output")
mock_result.all_messages.return_value = msg_history
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = mock_result
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
mock_run_hitl.return_value = "Approved output"
@@ -596,7 +613,7 @@ class TestAgentOperatorExecute:
):
"""When enable_hitl_review=True and output_type is BaseModel, execute
returns the model instance."""
mock_result = make_mock_run_result(Summary(text="Approved summary",
score=0.9))
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = mock_result
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
# run_hitl_review returns JSON string (as stored in
session.current_output)
@@ -627,7 +644,7 @@ class TestAgentOperatorExecute:
):
"""When enable_hitl_review=True and output_type is str, execute
returns string as-is."""
mock_result = make_mock_run_result("Initial output")
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = mock_result
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
mock_run_hitl.return_value = "Approved output"
@@ -657,7 +674,7 @@ class TestAgentOperatorExecute:
from airflow.providers.common.ai.exceptions import
HITLMaxIterationsError
mock_result = make_mock_run_result("Initial output")
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = mock_result
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
mock_run_hitl.side_effect = HITLMaxIterationsError("Task exceeded max
iterations.")
@@ -731,7 +748,7 @@ class TestAgentOperatorRegenerateWithFeedback:
msg_history = [MagicMock()]
mock_result = make_mock_run_result("Revised output")
mock_result.all_messages.return_value = msg_history + [MagicMock()]
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = mock_result
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -758,7 +775,7 @@ class TestAgentOperatorRegenerateWithFeedback:
"""regenerate_with_feedback returns JSON string for BaseModel
output."""
mock_result = make_mock_run_result(Summary(text="Revised"))
mock_result.all_messages.return_value = []
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = mock_result
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -854,10 +871,10 @@ class TestAgentOperatorDurable:
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = AgentOperator(task_id="test", prompt="test", llm_conn_id="my_llm")
- op.execute(context=MagicMock())
+ op.execute(context=_make_context())
# run_sync called directly, no override
- mock_agent.run_sync.assert_called_once_with("test", usage_limits=None)
+ mock_agent.run_sync.assert_called_once_with("test", usage_limits=None,
run_id="ti-1")
def test_build_durable_capabilities_wraps_toolset_capability(self):
"""A ``Toolset`` capability's inner toolset is wrapped with
CachingToolset;
@@ -940,7 +957,7 @@ class TestAgentOperatorDurable:
storage = MagicMock(spec=DurableStorageProtocol)
mock_build_storage.return_value = storage
- mock_agent = MagicMock(spec=["run_sync", "model", "override"])
+ mock_agent = MagicMock(spec=["run_sync", "model", "override",
"instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("ok")
mock_agent.model = "test-model"
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -948,7 +965,7 @@ class TestAgentOperatorDurable:
op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c",
durable=True, message_history="[]")
with patch.object(op, "_emit_message_history",
side_effect=RuntimeError("xcom down")):
with pytest.raises(RuntimeError, match="xcom down"):
- op.execute(context={})
+ op.execute(context=_make_context())
storage.cleanup.assert_not_called()
@@ -966,7 +983,7 @@ class TestAgentOperatorMultimodalPromptGuard:
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_rejects_sequence_prompt_with_hitl_review(self,
mock_hook_cls):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = AgentOperator(
@@ -1030,11 +1047,14 @@ class TestAgentOperatorMessageHistory:
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = AgentOperator(task_id="t", prompt="run", llm_conn_id="c")
- context = MagicMock()
+ context = _make_context()
op.execute(context=context)
assert "message_history" not in mock_agent.run_sync.call_args.kwargs
- context["task_instance"].xcom_push.assert_not_called()
+ # The transcript is not emitted without history, but run id + usage
always are.
+ pushed_keys = {c.kwargs["key"] for c in
context["task_instance"].xcom_push.call_args_list}
+ assert "message_history" not in pushed_keys
+ assert pushed_keys == {"run_id", "usage"}
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_transcript_emitted_to_xcom_when_history_set(self, mock_hook_cls,
make_mock_run_result):
@@ -1044,14 +1064,13 @@ class TestAgentOperatorMessageHistory:
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = AgentOperator(task_id="t", prompt="run", llm_conn_id="c",
message_history=[])
- context = MagicMock()
+ context = _make_context()
op.execute(context=context)
ti = context["task_instance"]
- ti.xcom_push.assert_called_once()
- push_kwargs = ti.xcom_push.call_args.kwargs
- assert push_kwargs["key"] == "message_history"
- restored = ModelMessagesTypeAdapter.validate_json(push_kwargs["value"])
+ pushes = {c.kwargs["key"]: c.kwargs["value"] for c in
ti.xcom_push.call_args_list}
+ assert set(pushes) == {"run_id", "usage", "message_history"}
+ restored =
ModelMessagesTypeAdapter.validate_json(pushes["message_history"])
assert len(restored) == 2
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
@@ -1096,7 +1115,7 @@ class TestAgentOperatorMessageHistory:
mock_build_storage.return_value =
MagicMock(spec=DurableStorageProtocol)
- mock_agent = MagicMock(spec=["run_sync", "model", "override"])
+ mock_agent = MagicMock(spec=["run_sync", "model", "override",
"instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("ok")
mock_agent.model = "test-model"
mock_agent.override.return_value.__enter__ =
MagicMock(return_value=None)
@@ -1148,3 +1167,86 @@ class TestAgentOperatorHITLArgumentChecks:
enable_hitl_review=True,
**conflicting_kwargs,
)
+
+
+class TestAgentOperatorRunIdentity:
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_run_id_and_usage_pushed_to_xcom(self, mock_hook_cls,
make_mock_run_result):
+ """The pydantic-ai run id and token usage are exposed on XCom for
downstream tasks."""
+ mock_agent = _make_mock_agent("ok", make_mock_run_result)
+ mock_agent.run_sync.return_value.run_id = "the-run-id"
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = AgentOperator(task_id="t", prompt="run", llm_conn_id="c")
+ context = _make_context()
+ op.execute(context=context)
+
+ pushes = {
+ c.kwargs["key"]: c.kwargs["value"] for c in
context["task_instance"].xcom_push.call_args_list
+ }
+ assert pushes["run_id"] == "the-run-id"
+ assert pushes["usage"] == {
+ "requests": 1,
+ "input_tokens": 0,
+ "output_tokens": 0,
+ "total_tokens": 0,
+ "tool_calls": 0,
+ "cost": None,
+ }
+
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_usage_cost_is_stringified_on_xcom(self, mock_hook_cls,
make_mock_run_result):
+ """A non-None run cost (Decimal) is stringified before it goes to
XCom."""
+ mock_agent = _make_mock_agent("ok", make_mock_run_result,
cost=PRICED_COST)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = AgentOperator(task_id="t", prompt="run", llm_conn_id="c")
+ context = _make_context()
+ op.execute(context=context)
+
+ pushes = {
+ c.kwargs["key"]: c.kwargs["value"] for c in
context["task_instance"].xcom_push.call_args_list
+ }
+ assert pushes["usage"]["cost"] == str(PRICED_COST)
+
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_run_metadata_not_pushed_when_do_xcom_push_false(self,
mock_hook_cls, make_mock_run_result):
+ """do_xcom_push=False suppresses the run_id/usage pushes like any
other operator XCom."""
+ mock_agent = _make_mock_agent("ok", make_mock_run_result)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = AgentOperator(task_id="t", prompt="run", llm_conn_id="c",
do_xcom_push=False)
+ context = _make_context()
+ op.execute(context=context)
+
+ pushed_keys = {c.kwargs["key"] for c in
context["task_instance"].xcom_push.call_args_list}
+ assert "run_id" not in pushed_keys
+ assert "usage" not in pushed_keys
+
+
@patch("airflow.providers.common.ai.operators.agent.stamp_identity_on_agent_spans",
autospec=True)
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_execute_stamps_identity_on_agent_spans(self, mock_hook_cls,
mock_stamp, make_mock_run_result):
+ """execute() derives the identity from the task instance and stamps it
on the agent's spans."""
+ mock_agent = _make_mock_agent("ok", make_mock_run_result)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = AgentOperator(task_id="t", prompt="run", llm_conn_id="c")
+ op.execute(context=_make_context(_make_ti(id="ti-9")))
+
+ mock_stamp.assert_called_once()
+ agent_arg, attrs = mock_stamp.call_args.args
+ assert agent_arg is mock_agent
+ assert attrs["airflow.task_instance.id"] == "ti-9"
+
+
@patch("airflow.providers.common.ai.operators.agent.stamp_identity_on_agent_spans",
autospec=True)
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_regenerate_with_feedback_stamps_identity(self, mock_hook_cls,
mock_stamp, make_mock_run_result):
+ """A HITL re-run stamps the same identity the initial run resolved."""
+ mock_agent = _make_mock_agent("revised", make_mock_run_result)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = AgentOperator(task_id="t", prompt="run", llm_conn_id="c")
+ op._run_identity_attrs = {"airflow.task_instance.id": "ti-9"}
+ op.regenerate_with_feedback(feedback="more", message_history=[])
+
+ mock_stamp.assert_called_once_with(mock_agent,
{"airflow.task_instance.id": "ti-9"})
diff --git a/providers/common/ai/tests/unit/common/ai/test_observability.py
b/providers/common/ai/tests/unit/common/ai/test_observability.py
index 4b0c84f7e50..91f4e8c91df 100644
--- a/providers/common/ai/tests/unit/common/ai/test_observability.py
+++ b/providers/common/ai/tests/unit/common/ai/test_observability.py
@@ -17,11 +17,15 @@
from __future__ import annotations
import json
+from types import SimpleNamespace
from unittest.mock import MagicMock, patch
+import pytest
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import
InMemorySpanExporter
+from pydantic_ai import Agent
+from pydantic_ai.models.instrumented import InstrumentationSettings
from pydantic_ai.models.test import TestModel
from airflow.providers.common.ai import observability
@@ -141,3 +145,87 @@ class TestEndToEndSpanEmission:
assert genai
# With the opt-in, the prompt text is present on the spans.
assert self._PROMPT in attrs_blob
+
+
+class TestBuildRunIdentityAttributes:
+ @pytest.mark.parametrize(("map_index", "expected_map_index"), [(3, 3),
(None, -1)])
+ def test_builds_expected_attributes(self, map_index, expected_map_index):
+ ti = SimpleNamespace(
+ id="ti-1", dag_id="d", task_id="t", run_id="r", try_number=2,
map_index=map_index
+ )
+ assert observability.build_run_identity_attributes(ti) == {
+ "airflow.dag_id": "d",
+ "airflow.task_id": "t",
+ "airflow.dag_run.run_id": "r",
+ "airflow.task_instance.try_number": 2,
+ "airflow.task_instance.map_index": expected_map_index,
+ "airflow.task_instance.id": "ti-1",
+ }
+
+
+class TestStampIdentityOnAgentSpans:
+ _ATTRS = {
+ "airflow.dag_id": "d",
+ "airflow.task_id": "t",
+ "airflow.dag_run.run_id": "r",
+ "airflow.task_instance.try_number": 2,
+ "airflow.task_instance.map_index": -1,
+ "airflow.task_instance.id": "ti-xyz",
+ }
+
+ def test_attributes_and_run_id_land_on_genai_spans(self):
+ exporter = InMemorySpanExporter()
+ provider = TracerProvider()
+ provider.add_span_processor(SimpleSpanProcessor(exporter))
+
+ hook = PydanticAIHook(llm_conn_id="c", model_id="test")
+ with (
+ patch.object(observability, "conf", _conf(enabled=True)),
+ patch.object(observability, "_live_tracer_provider",
return_value=provider),
+ patch.object(hook, "get_conn", return_value=TestModel()),
+ ):
+ agent = hook.create_agent(instructions="be helpful")
+ observability.stamp_identity_on_agent_spans(agent, self._ATTRS)
+ agent.run_sync("hi", run_id="ti-xyz")
+
+ genai = [
+ s
+ for s in exporter.get_finished_spans()
+ if s.attributes and any(k.startswith("gen_ai.") for k in
s.attributes)
+ ]
+ assert genai, "expected gen_ai spans to be emitted"
+ # Every GenAI span carries the full Airflow identity, not just the
agent-run span.
+ for span in genai:
+ for key, value in self._ATTRS.items():
+ assert span.attributes.get(key) == value
+ # run_id passed to the run surfaces as the OTel agent-call id.
+ assert any(s.attributes.get("gen_ai.agent.call.id") == "ti-xyz" for s
in genai)
+
+ def test_noop_when_agent_not_instrumented(self):
+ # A non-InstrumentationSettings ``instrument`` (here ``False``) is left
+ # untouched: stamping must not wrap it or raise.
+ agent = Agent(TestModel())
+ agent.instrument = False
+ observability.stamp_identity_on_agent_spans(agent, self._ATTRS)
+ assert agent.instrument is False
+
+ def test_wraps_a_copy_without_mutating_or_nesting_shared_settings(self):
+ # A caller-owned settings object handed to several agents (or reused
across
+ # HITL re-runs) must stay untouched: each agent gets its own wrapped
copy,
+ # and every wrapper wraps the real tracer directly rather than nesting.
+ provider = TracerProvider()
+ shared = InstrumentationSettings(tracer_provider=provider)
+ original_tracer = shared.tracer
+
+ agent1, agent2 = Agent(TestModel()), Agent(TestModel())
+ agent1.instrument = shared
+ agent2.instrument = shared
+ observability.stamp_identity_on_agent_spans(agent1, self._ATTRS)
+ observability.stamp_identity_on_agent_spans(agent2, self._ATTRS)
+
+ assert shared.tracer is original_tracer
+ assert agent1.instrument is not shared
+ assert agent1.instrument is not agent2.instrument
+ assert isinstance(agent1.instrument.tracer,
observability._IdentityTracer)
+ assert agent1.instrument.tracer._tracer is original_tracer
+ assert agent2.instrument.tracer._tracer is original_tracer
