Skip to content

Add AI Trace observability panel to the common.ai provider#69880

Draft
kaxil wants to merge 1 commit into
apache:mainfrom
astronomer:common-ai-trace-plugin
Draft

Add AI Trace observability panel to the common.ai provider#69880
kaxil wants to merge 1 commit into
apache:mainfrom
astronomer:common-ai-trace-plugin

Conversation

@kaxil

@kaxil kaxil commented Jul 14, 2026

Copy link
Copy Markdown
Member

Adds an AI Trace observability panel to the common.ai provider: an
AirflowPlugin that mounts a FastAPI sub-app on the API server and a React panel
on the UI, so a GenAI agent run (AgentOperator / @task.agent / @task.llm)
can be inspected — conversation, model, tokens, cost, latency, and the full
observation tree with a waterfall — from inside Airflow, in context on the task
instance and in a deployment-wide list.

It reads traces from two sources:

  1. A configured tracing backend (Langfuse today) — the panel resolves a task
    instance's trace via its own OTel span context and proxies the backend's read
    API. No new data is written anywhere.
  2. A backend-free ObjectStorage trace store — when [common.ai] trace_store_path is set, agent spans are written as standard OTLP JSON lines
    under an ObjectStoragePath, and the panel reads them back directly. This
    needs no Langfuse, no collector, and no core tracing — just a path. It targets
    local dev and quick testing.

Screenshots

The AI Trace panel in backend-free store mode — role-ordered conversation, the
observation tree with a duration waterfall, and per-generation cost estimated
from genai-prices (no "Open in Langfuse" link, since there is no external
backend in store mode):

ai-trace-store-list ai-trace-cost-modal image

Design rationale

  • Correlation is exact, not heuristic. common.ai agent spans nest under the
    worker's task span and inherit its trace_id, so the panel resolves "this task
    instance's trace" straight from the stored traceparent — no separate push or id
    override.
  • The store path is the correlation. The trace-store layout is keyed by
    {dag_id}/{run_id}/{task_id}/{map_index}/{try_number}.jsonl, so store mode works
    with core tracing off. Files are standard OTLP JSON, so a collector's
    otlpjsonfilereceiver can replay them into a real backend later.
  • Costs are estimated at read time via genai-prices (already a transitive
    dependency of pydantic-ai-slim), using the model, provider, and token counts
    on each span. Unlike ingest-time pricing, this stays current with the library's
    bundled price data and prices old files retroactively; unknown or self-hosted
    models simply show no cost.
  • The OTLP JSON encoder is vendored (_otlp_json.py): the upstream
    opentelemetry-exporter-otlp-json-file release is currently uninstallable from
    PyPI (its opentelemetry-proto-json dependency was never published), so a hard
    dependency would break resolution. The vendored encoder follows the OTLP
    JSON-Protobuf encoding so the files remain collector-replayable; it can be
    swapped for the upstream exporter once that ships.

Usage

[common.ai]
# Backend-free local-dev store (any ObjectStoragePath scheme):
trace_store_path = file:///tmp/airflow_ai_traces
# Include prompts / completions / tool IO (off by default):
capture_content = True

With trace_store_path unset, the panel instead reads from the connection named
by [common.ai] trace_backend_conn_id (a Langfuse connection). Open the panel
from a task instance's AI Trace tab, or the deployment-wide AI Traces nav
entry.

Security

  • Every endpoint is RBAC-gated. The task-instance panel checks TASK_INSTANCE
    access to the DAG; the bare-id lookups resolve the trace (or observation) to its
    owning task instance and require access to that DAG, so a user authorized for
    one DAG cannot read another's captured content.
  • Task-instance coordinates are validated before they reach an ObjectStoragePath
    join (it does not normalize ..), and trace ids are validated as 32-hex before
    use in a reverse lookup or file scan.

Known limitations

  • Spike-status store layout — the trace-store format is a proof of concept,
    not a stable on-disk contract. Marked as such in the config docs.
  • No retention — the store is a directory of JSONL files; clean it yourself.
  • Object stores flush at task-process exit — on s3:///gs://, a task's spans
    land when the process exits, not live (local file:// is live).
  • Cost is an estimate, not billing data; negotiated/enterprise pricing differs,
    and a shown cost can drift slightly across genai-prices upgrades.
  • Bare-id lookup requires a resolvable task instance — a trace that resolves to
    no task instance (possible in Langfuse mode) returns 404 by direct id, since there
    is no DAG to authorize against.

Adds an AI Trace AirflowPlugin: a FastAPI sub-app mounted on the API
server plus a React panel on the UI, so a GenAI agent run (AgentOperator
/ @task.agent / @task.llm) can be inspected in context -- conversation,
model, tokens, cost, latency, and the full observation tree with a
duration waterfall -- from the task instance and from a deployment-wide
list, without leaving Airflow.

The panel reads traces from two sources:

- A configured tracing backend (Langfuse today): the trace is resolved
  from the task instance's own OTel span context and read back through
  the backend's API. No new data is written anywhere.
- A backend-free ObjectStorage trace store: when [common.ai]
  trace_store_path is set, agent spans are written as standard OTLP JSON
  lines under {dag_id}/{run_id}/{task_id}/{map_index}/{try_number}.jsonl
  and the panel reads them directly. This needs no Langfuse, no
  collector, and no core tracing -- just a path -- and targets local dev
  and quick testing. Files are standard OTLP JSON, so a collector's
  otlpjsonfilereceiver can replay them into a real backend later.

Costs are estimated at read time via genai-prices (already a transitive
dependency of pydantic-ai-slim) from the model, provider, and token
counts on each span, so old files are priced retroactively and unknown
or self-hosted models simply show no cost. The OTLP JSON encoder is
vendored because the upstream opentelemetry-exporter-otlp-json-file
release is currently uninstallable from PyPI (its opentelemetry-proto-json
dependency was never published); it follows the OTLP JSON-Protobuf
encoding so files stay collector-replayable.

Security: every endpoint is RBAC-gated. The task-instance panel checks
TASK_INSTANCE access to the DAG; the bare-id lookups resolve the trace
(or observation) to its owning task instance and require access to that
DAG, so a user authorized for one DAG cannot read another's captured
content. Task-instance coordinates are validated before they reach an
ObjectStoragePath join (it does not normalize ".."), and trace ids are
validated as 32-hex before use in a reverse lookup or file scan.

Includes unit tests for the OTLP JSON encoder (wire-format round trip
against real SDK spans), the store reader and normalizer, read-time cost
estimation, and the endpoints in both modes (RBAC, path-traversal
rejection, trace-id validation, per-DAG authorization, and 404 paths).

Known limitations: the trace-store layout is a proof of concept, not a
stable on-disk contract; no retention is applied; object stores flush a
task's spans at process exit rather than live; cost is an estimate, not
billing data.
@kaxil kaxil force-pushed the common-ai-trace-plugin branch from fb9f80c to 53e7a0e Compare July 14, 2026 17:07
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.

1 participant