feat(lineage): capture tool data-source refs and agent build version in span data

[max-parke-scale]

Implements the SDK half of the trace data-source-ref convention (SGP-6513; convention spec in scaleapi#153026): tools declare which data sources they actually touch, and the refs land in span data under sgp.lineage.refs as lineage coordinates. The SGP tracing processor already ships span data as sgp-traces span metadata, so refs are immediately filterable via the spans-search extra_metadata DSL — no schema, processor, or service changes. Capture is decoupled from lineage derivation: agents instrument from day one and graph edges backfill later from materialized traces.

Surface (agentex.lib.adk.lineage, implementation in core/tracing/lineage.py)

from agentex.lib.adk import DataSourceRef, data_sources, lineage

@function_tool
@data_sources(DataSourceRef("databricks://ey-tax", "guidance.rulings"))
async def query_guidance(q: str) -> str: ...          # owned tool, static refs

@data_sources(resolver=lambda args: [DataSourceRef("databricks://ey-tax", args["table"])])
                                                       # dynamic refs from validated args

lineage.register_tool_sources("competitive_edge_search",  # MCP/unowned tools,
    [DataSourceRef("elasticsearch://ey-embryonic", "companies_v3")])  # name-keyed

lineage.record(span, [DataSourceRef(...)])             # manual spans

DataSourceRef validates URI-form namespaces per the lineage namespace conventions and rejects payload-shaped values; refs deduplicate by (namespace, name, version, role). Resolver failures are logged and swallowed — ref capture never breaks a tool call or its tracing.

Where refs get stamped

• Unified harness: SpanTracer resolves refs for every OpenSpan(kind="tool") signal — one point covering all harness turns. Import is guarded so the harness stays importable without optional tracing deps (same pattern as its logger).
• run_agent* service paths (core/services/adk/providers/openai.py, 4 methods): refs resolve from the serialized function_call items and merge onto the run span.
• SyncStreamingProvider and the Temporal streaming model: same item-based merge on their model spans.

Agent build version stamp (__agent_version__)

The processor already env-stamps __agent_name__ / __agent_id__ into every span; this adds __agent_version__ from a new optional AGENT_VERSION env var (image tag or git sha, set by the deployment). This completes the trace-side join-key set — span → exact agent-version snapshot instead of a temporal join against deploy history — the runtime half of SGP-6132. Unset means no stamp, so nothing changes until a deployment wires the env var; the deploy-chart wiring is a follow-up in the deployment repo.

Deliberately deferred

MCP-server-side ref attachment via result _meta with agent-side harvest (the exact-grain path for server-private dynamic sources) — needs a cross-process contract; name-keyed agent-side declaration covers current consumers.

Tests

tests/lib/core/tracing/test_lineage.py (validation, registry, decorator both placements, item resolution, merge/dedupe), tests/lib/core/harness/test_tracer_lineage.py (SpanTracer stamps refs on tool spans, leaves reasoning/unregistered spans untouched), and TestSourceStamps in the SGP processor suite (__agent_version__ stamped when set, omitted when unset). Full harness + tracing + temporal + adk suites green locally; ruff check / ruff format clean.

🧑‍💻🤖 — posted via Claude Code

Greptile Summary

This PR implements the SDK side of the trace data-source-ref convention (SGP-6513): tools declare which data sources they touch via DataSourceRef / @data_sources / register_tool_sources, and the refs are stamped into span data under sgp.lineage.refs at every capture point. It also adds AGENT_VERSION stamping to the SGP tracing processor from a new optional AGENT_VERSION env var.

• lineage.py: New module with DataSourceRef (URI-form namespace validation, dedup by (namespace, name, version, role)), a process-wide registry, a @data_sources decorator, and helpers (resolve_refs, resolve_refs_from_items, merge_refs_into_data, record) consumed by all integration points.
• Integration points: SpanTracer (harness, tool spans), all four run_agent* methods in openai.py, SyncStreamingModel, and TemporalStreamingModel — each resolves refs from serialized function_call items and merges them onto the relevant span; resolver failures are swallowed so ref capture never breaks a tool call.
• AGENT_VERSION: Optional env var added to EnvironmentVariables and stamped in _add_source_to_span alongside the existing __agent_name__ / __agent_id__ fields; no-op until a deployment wires the env var.

Confidence Score: 5/5

Safe to merge; the change is additive and self-contained — ref capture never affects tool execution or tracing correctness, and the new env var is a no-op until wired in deployment.

All integration points follow the same defensive pattern: resolver failures are logged and swallowed, refs are only merged when non-empty, and the harness import is guarded so the module stays importable without optional tracing deps. The only finding is a misleading inline comment in the regex.

Files Needing Attention: No files require special attention.

Important Files Changed

Filename	Overview
src/agentex/lib/core/tracing/lineage.py	New module implementing DataSourceRef, the registry, decorator, resolver, and merge helpers; solid validation, dedup, and swallowed-resolver-failure pattern throughout
src/agentex/lib/core/harness/tracer.py	Adds guarded lineage import + stamps refs on tool spans at OpenSpan time; fallback stubs ensure harness stays importable without tracing deps
src/agentex/lib/core/services/adk/providers/openai.py	Extracts serialized_items before span.output assignment in all four run_agent* methods, then merges refs from function_call items onto the run span
src/agentex/lib/adk/providers/_modules/sync_provider.py	Merges lineage refs onto the model span after the run in both get_response and stream_response paths; new_items are already serialized dicts so resolve_refs_from_items is appropriate
src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py	Merges lineage refs from model response output items onto span data; new_items are populated from response_output via _serialize_item so function_call items are properly captured
src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py	Adds AGENT_VERSION stamp to span data in _add_source_to_span following the same optional guard pattern as AGENT_ID
src/agentex/lib/environment_variables.py	Adds AGENT_VERSION as an optional env var with None default; no behaviour change until a deployment sets the env var
tests/lib/core/tracing/test_lineage.py	Comprehensive unit tests covering validation, registry, decorator, item resolution, merge/dedup, and record(); autouse fixture clears the module-level registry for isolation
tests/lib/core/harness/test_tracer_lineage.py	Validates that SpanTracer stamps refs on tool spans and leaves reasoning/unregistered spans untouched
tests/lib/core/tracing/processors/test_sgp_tracing_processor.py	Adds TestSourceStamps class covering AGENT_VERSION stamping; updates existing test mocks to include AGENT_VERSION=None
src/agentex/lib/adk/init.py	Re-exports lineage module, DataSourceRef, and data_sources decorator from the public adk surface

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["@data_sources decorator\nor register_tool_sources()"] --> R[("_tool_sources\nprocess-wide registry")]
    B["lineage.record(span, refs)"] --> M

    subgraph Capture["Ref capture sites"]
        C["SpanTracer\nOpenSpan(kind='tool')"] -->|resolve_refs| M
        D["OpenAIService\nrun_agent*"] -->|resolve_refs_from_items| M
        E["SyncStreamingModel\nget_response / stream_response"] -->|resolve_refs_from_items| M
        F["TemporalStreamingModel\nstream_response"] -->|resolve_refs_from_items| M
    end

    R -->|lookup by tool name| C
    R -->|lookup by tool name| D
    R -->|lookup by tool name| E
    R -->|lookup by tool name| F

    M["merge_refs_into_data()\ndedup by (ns, name, ver, role)"] --> S["span.data\nsgp.lineage.refs"]

    subgraph SGP["SGP processor"]
        G["_add_source_to_span"] -->|"__agent_version__\n(AGENT_VERSION env var)"| S
    end

Loading

_{Reviews (3): Last reviewed commit: "feat(lineage): capture tool data-source ..." | Re-trigger Greptile}