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 introduces the SDK half of the trace data-source-ref convention: tools declare which data sources they touch via DataSourceRef, the @data_sources decorator (or register_tool_sources for MCP/unowned tools), and refs land in span data under sgp.lineage.refs. It also stamps __agent_version__ from a new optional AGENT_VERSION env var, completing the span join-key set for trace-to-deploy-history linkage.

• New lineage.py module: Implements DataSourceRef (Pydantic-validated URI-form namespace), a process-wide registry, resolve_refs/resolve_refs_from_items helpers, and merge_refs_into_data — all with resolver-failure isolation so ref capture can never break a tool call.
• Integration points: SpanTracer resolves refs on every OpenSpan(kind="tool") signal; four OpenAIService methods and SyncStreamingModel/TemporalStreamingModel merge refs from serialized function_call items onto their run spans.
• AGENT_VERSION env var: Added to EnvironmentVariables and stamped by the SGP tracing processor when set; backwards-compatible (unset = no stamp).

Confidence Score: 5/5

Safe to merge — changes are additive, all failure paths are swallowed by design, and no existing span data is overwritten.

The feature is purely additive: refs are merged into span.data via a copy, resolver exceptions are caught and logged, the guarded import in tracer.py keeps the harness functional if the module ever fails to import, and AGENT_VERSION is opt-in via a new optional env var. Tests cover validation, registry, decorator ordering, item resolution, merge/dedupe, and the processor stamp in both set and unset states. No existing behaviour is changed.

No files require special attention.

Important Files Changed

Filename	Overview
src/agentex/lib/core/tracing/lineage.py	New module implementing the full lineage ref capture surface — validation, registry, decorator, resolution, merge/dedupe. Well-structured with defensive error handling throughout.
src/agentex/lib/core/harness/tracer.py	Adds lineage ref stamping to tool spans in SpanTracer.handle; uses a guarded import following the existing optional-deps pattern to keep the harness importable without tracing deps.
src/agentex/lib/core/services/adk/providers/openai.py	Four run_agent* methods refactored to extract serialized_items into a local variable, then resolve and merge lineage refs onto the run span — clean, symmetric changes.
src/agentex/lib/adk/providers/_modules/sync_provider.py	Adds lineage ref merge to two SyncStreamingModel spans; new_items is already a list of serialized dicts so resolve_refs_from_items works correctly.
src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py	Merges lineage refs from serialized response_output items onto the model span; new_items is built from _serialize_item dicts so function_call entries are correctly matched.
src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py	Stamps agent_version from AGENT_VERSION env var into span data when set; backwards-compatible minimal change.
src/agentex/lib/environment_variables.py	Adds AGENT_VERSION as an optional string field and enum key; no behavioral impact until the env var is set in deployments.
tests/lib/core/tracing/test_lineage.py	Comprehensive unit tests covering validation, registry, decorator placement, item resolution, merge/dedupe — good use of autouse fixture for registry isolation.
tests/lib/core/harness/test_tracer_lineage.py	Verifies SpanTracer stamps refs on tool spans and leaves reasoning/unregistered spans untouched; clean fixture-based registry isolation.
tests/lib/core/tracing/processors/test_sgp_tracing_processor.py	Adds TestSourceStamps covering agent_version stamping and omission when unset; updates existing processor mocks to include AGENT_VERSION=None.
src/agentex/lib/adk/init.py	Re-exports lineage, DataSourceRef, and data_sources from the adk public surface; straightforward addition to all.

Sequence Diagram

sequenceDiagram
    participant Tool as Tool / MCP Tool
    participant Registry as lineage._tool_sources
    participant Tracer as SpanTracer / Provider
    participant Span as Span.data
    participant Processor as SGPTracingProcessor

    Note over Tool,Registry: Registration (import / decorator time)
    Tool->>Registry: register_tool_sources(name, refs, resolver)

    Note over Tracer,Span: Runtime — OpenSpan(kind="tool") or function_call item
    Tracer->>Registry: resolve_refs(tool_name, arguments)
    Registry-->>Tracer: [DataSourceRef dicts] (resolver failure swallowed)
    Tracer->>Span: merge_refs_into_data(span.data, refs)
    Note over Span: sgp.lineage.refs = deduplicated refs

    Note over Processor,Span: On span end
    Processor->>Span: _add_source_to_span → __source__, __agent_name__, __agent_id__, __agent_version__
    Processor->>Processor: ship span as SGP metadata (sgp-traces)

Loading

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