From 0361006c3e8d14486855b259ee77c0380df3ffe6 Mon Sep 17 00:00:00 2001 From: Max Parke Date: Wed, 22 Jul 2026 12:29:40 -0700 Subject: [PATCH] feat(lineage): capture tool data-source refs and agent build version in span data MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements the trace data-source-ref convention (SGP-6513): tools declare which data sources they touch — statically, via an args resolver, or by name-keyed registry for MCP/unowned tools — and every tool-span path merges the resolved refs into span data under sgp.lineage.refs, which the SGP tracing processor already ships as span metadata. Capture is decoupled from lineage derivation so agents instrument from day one and edges backfill later. Also stamps __agent_version__ from a new AGENT_VERSION env var (same mechanism as __agent_name__), completing the trace-side join-key set: span -> agent version snapshot is the runtime half of SGP-6132. Convention spec: scaleapi packages/sgp-lineage/docs/specs/ 2026-07-22-sgp-6513-trace-data-source-ref-convention.md (PR #153026). Co-Authored-By: Claude Fable 5 --- src/agentex/lib/adk/__init__.py | 8 + .../adk/providers/_modules/sync_provider.py | 7 + src/agentex/lib/core/harness/tracer.py | 16 ++ .../lib/core/services/adk/providers/openai.py | 49 +++-- .../models/temporal_streaming_model.py | 4 + src/agentex/lib/core/tracing/lineage.py | 174 ++++++++++++++++++ .../processors/sgp_tracing_processor.py | 2 + src/agentex/lib/environment_variables.py | 3 + tests/lib/core/harness/test_tracer_lineage.py | 53 ++++++ .../processors/test_sgp_tracing_processor.py | 57 ++++-- tests/lib/core/tracing/test_lineage.py | 147 +++++++++++++++ 11 files changed, 484 insertions(+), 36 deletions(-) create mode 100644 src/agentex/lib/core/tracing/lineage.py create mode 100644 tests/lib/core/harness/test_tracer_lineage.py create mode 100644 tests/lib/core/tracing/test_lineage.py diff --git a/src/agentex/lib/adk/__init__.py b/src/agentex/lib/adk/__init__.py index 25b858485..d5be0ac52 100644 --- a/src/agentex/lib/adk/__init__.py +++ b/src/agentex/lib/adk/__init__.py @@ -29,6 +29,10 @@ from agentex.lib.adk._modules.tasks import TasksModule from agentex.lib.adk._modules.tracing import TracingModule, TurnSpan +# Data-source refs for lineage (SGP-6513); implementation lives in core.tracing +from agentex.lib.core.tracing import lineage +from agentex.lib.core.tracing.lineage import DataSourceRef, data_sources + # Unified harness surface (AGX1-375) from agentex.lib.core.harness import ( UnifiedEmitter, @@ -67,6 +71,10 @@ "events", "agent_task_tracker", "TurnSpan", + # Lineage data-source refs (SGP-6513) + "lineage", + "DataSourceRef", + "data_sources", # Checkpointing / LangGraph "create_checkpointer", "stream_langgraph_events", diff --git a/src/agentex/lib/adk/providers/_modules/sync_provider.py b/src/agentex/lib/adk/providers/_modules/sync_provider.py index 86696a2b5..120915eec 100644 --- a/src/agentex/lib/adk/providers/_modules/sync_provider.py +++ b/src/agentex/lib/adk/providers/_modules/sync_provider.py @@ -19,6 +19,7 @@ from agentex import AsyncAgentex from agentex.lib.utils.logging import make_logger from agentex.lib.core.tracing.tracer import AsyncTracer +from agentex.lib.core.tracing.lineage import merge_refs_into_data, resolve_refs_from_items logger = make_logger(__name__) @@ -185,6 +186,9 @@ async def get_response( "new_items": new_items, "final_output": final_output, } + lineage_refs = resolve_refs_from_items(new_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return response else: @@ -303,6 +307,9 @@ async def stream_response( "new_items": new_items, "final_output": final_response_text if final_response_text else None, } + lineage_refs = resolve_refs_from_items(new_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) finally: # End the span after all events have been yielded await trace.end_span(span) diff --git a/src/agentex/lib/core/harness/tracer.py b/src/agentex/lib/core/harness/tracer.py index bf37bad30..34cd95616 100644 --- a/src/agentex/lib/core/harness/tracer.py +++ b/src/agentex/lib/core/harness/tracer.py @@ -6,6 +6,17 @@ from agentex.lib.core.harness.types import OpenSpan, CloseSpan, SpanSignal +try: + from agentex.lib.core.tracing.lineage import resolve_refs, merge_refs_into_data +except Exception: # keep the harness importable without optional tracing deps + + def resolve_refs(tool_name: str, arguments: dict[str, Any] | None) -> list[dict[str, Any]]: # noqa: ARG001 + return [] + + def merge_refs_into_data(data: dict[str, Any] | None, refs: list[dict[str, Any]]) -> dict[str, Any]: # noqa: ARG001 + return dict(data or {}) + + try: from agentex.lib.utils.logging import make_logger @@ -80,6 +91,11 @@ async def handle(self, signal: SpanSignal) -> None: task_id=self.task_id, ) if span is not None: + if signal.kind == "tool": + refs = resolve_refs(signal.name, signal.input if isinstance(signal.input, dict) else {}) + if refs: + data = span.data if isinstance(span.data, dict) else {} + span.data = merge_refs_into_data(data, refs) self._open[signal.key] = span elif isinstance(signal, CloseSpan): span = self._open.pop(signal.key, None) diff --git a/src/agentex/lib/core/services/adk/providers/openai.py b/src/agentex/lib/core/services/adk/providers/openai.py index a2513ea01..cc411dc30 100644 --- a/src/agentex/lib/core/services/adk/providers/openai.py +++ b/src/agentex/lib/core/services/adk/providers/openai.py @@ -25,6 +25,7 @@ from agentex.lib.utils.temporal import heartbeat_if_in_workflow from agentex.lib.core.tracing.tracer import AsyncTracer from agentex.lib.core.harness.emitter import UnifiedEmitter +from agentex.lib.core.tracing.lineage import merge_refs_into_data, resolve_refs_from_items from agentex.types.task_message_update import StreamTaskMessageFull from agentex.types.task_message_content import ( TextContent, @@ -286,13 +287,17 @@ async def run_agent( result = await Runner.run(starting_agent=agent, input=input_list) if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return result @@ -431,13 +436,17 @@ async def run_agent_auto_send( result = await Runner.run(starting_agent=agent, input=input_list) if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) tool_call_map: dict[str, Any] = {} @@ -646,13 +655,17 @@ async def run_agent_streamed( result = Runner.run_streamed(starting_agent=agent, input=input_list) if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return result @@ -906,12 +919,16 @@ async def run_agent_streamed_auto_send( raise if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return result diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py index 7c8690f21..c985d5e65 100644 --- a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py +++ b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py @@ -64,6 +64,7 @@ from agentex.lib import adk from agentex.lib.utils.logging import make_logger from agentex.lib.core.tracing.tracer import AsyncTracer +from agentex.lib.core.tracing.lineage import merge_refs_into_data, resolve_refs_from_items from agentex.types.task_message_delta import TextDelta, ToolRequestDelta, ReasoningContentDelta, ReasoningSummaryDelta from agentex.types.task_message_update import StreamTaskMessageFull, StreamTaskMessageDelta from agentex.types.task_message_content import TextContent, ReasoningContent, ToolRequestContent, ToolResponseContent @@ -1257,6 +1258,9 @@ async def get_response( output_data["tool_outputs"] = tool_outputs span.output = output_data + lineage_refs = resolve_refs_from_items(new_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) # Streaming-only metrics. Token counters and the success request # counter are emitted by LLMMetricsHooks.on_llm_end so they fire diff --git a/src/agentex/lib/core/tracing/lineage.py b/src/agentex/lib/core/tracing/lineage.py new file mode 100644 index 000000000..75eaffdc0 --- /dev/null +++ b/src/agentex/lib/core/tracing/lineage.py @@ -0,0 +1,174 @@ +"""Data-source reference capture for lineage: tools declare which sources they +touch and the refs land in span data under the ``sgp.lineage.refs`` key.""" + +from __future__ import annotations + +import re +import json +from typing import Any, Literal, Callable, Iterable + +from pydantic import Field, BaseModel, field_validator + +try: + from agentex.lib.utils.logging import make_logger + + logger = make_logger(__name__) +except Exception: # ddtrace may be absent in some envs; fall back to stdlib + import logging + + logger = logging.getLogger(__name__) + +LINEAGE_REFS_KEY = "sgp.lineage.refs" + +# The URI arm of the lineage namespace identifier rule (namespace-conventions.md): +# lowercase scheme and host (dots/hyphens only — normalize `_` to `-`), one optional path segment. +_URI_NAMESPACE_RE = re.compile(r"^[a-z][a-z0-9._-]*://[a-z0-9.-]+(/[a-zA-Z0-9._-]*)?$") + +RefResolver = Callable[[dict[str, Any]], "list[DataSourceRef]"] + + +class DataSourceRef(BaseModel): + """One data source a tool call touched, as a lineage coordinate.""" + + namespace: str = Field(max_length=512) + name: str = Field(min_length=1, max_length=512) + version: str | None = Field(default=None, max_length=256) + role: Literal["input", "output"] = "input" + + def __init__(self, namespace: str | None = None, name: str | None = None, **kwargs: Any) -> None: + if namespace is not None: + kwargs["namespace"] = namespace + if name is not None: + kwargs["name"] = name + super().__init__(**kwargs) + + @field_validator("namespace") + @classmethod + def _namespace_is_uri_form(cls, value: str) -> str: + if not _URI_NAMESPACE_RE.match(value): + raise ValueError(f"namespace must be URI-form (scheme://system), got: {value!r}") + return value + + +class _ToolSources(BaseModel): + refs: list[DataSourceRef] = Field(default_factory=list) + resolver: RefResolver | None = None + + model_config = {"arbitrary_types_allowed": True} + + +_tool_sources: dict[str, _ToolSources] = {} + + +def register_tool_sources( + tool_name: str, + refs: Iterable[DataSourceRef] | None = None, + resolver: RefResolver | None = None, +) -> None: + """Declare the data sources a tool touches, keyed by its tool name. + + Use for tools the agent does not own (e.g. MCP proxy tools). Static refs and + a resolver over the tool's parsed arguments may be combined; repeated + registration for the same name replaces the prior entry. The registry is + process-wide: co-located agents sharing a tool name share (and overwrite) + one entry, so disambiguate shared names before co-locating agent types. + """ + _tool_sources[tool_name] = _ToolSources(refs=list(refs or []), resolver=resolver) + + +def data_sources(*refs: DataSourceRef, resolver: RefResolver | None = None) -> Callable[[Any], Any]: + """Decorator form of ``register_tool_sources`` for tools the agent owns. + + Works below or above ``@function_tool``: the tool name is taken from the + decorated object's ``name`` attribute when present, else ``__name__``. + """ + + def _register(obj: Any) -> Any: + tool_name = getattr(obj, "name", None) or getattr(obj, "__name__", None) + if isinstance(tool_name, str) and tool_name: + register_tool_sources(tool_name, refs=refs, resolver=resolver) + else: + logger.warning("data_sources could not determine a tool name for %r; refs not registered", obj) + return obj + + return _register + + +def clear_tool_sources() -> None: + """Reset the registry (test isolation).""" + _tool_sources.clear() + + +def resolve_refs(tool_name: str, arguments: dict[str, Any] | None) -> list[dict[str, Any]]: + """Resolve registered refs for one tool call to serialized, deduplicated dicts. + + Resolver failures are logged and swallowed: ref capture must never break a + tool call or its tracing. + """ + entry = _tool_sources.get(tool_name) + if entry is None: + return [] + refs = list(entry.refs) + if entry.resolver is not None: + try: + refs.extend(entry.resolver(arguments or {})) + except Exception: + logger.warning("data-source resolver for tool %s failed; static refs kept", tool_name, exc_info=True) + return _dedupe(refs) + + +def resolve_refs_from_items(items: Iterable[Any]) -> list[dict[str, Any]]: + """Resolve refs across serialized run items, matching ``function_call`` entries. + + Accepts the item dicts the providers already build for span output; string + ``arguments`` are parsed as JSON for resolver-based registrations. + """ + refs: list[dict[str, Any]] = [] + for item in items: + if not isinstance(item, dict) or item.get("type") != "function_call": + continue + tool_name = item.get("name") + if not isinstance(tool_name, str) or not tool_name: + continue + arguments = item.get("arguments") + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except (ValueError, TypeError): + arguments = {} + refs.extend(resolve_refs(tool_name, arguments if isinstance(arguments, dict) else {})) + return _dedupe_dicts(refs) + + +def record(span: Any, refs: Iterable[DataSourceRef]) -> None: + """Attach refs to a manually managed span (no-op when the span is None).""" + if span is None: + return + merged = merge_refs_into_data(getattr(span, "data", None), _dedupe(list(refs))) + span.data = merged + + +def merge_refs_into_data(data: dict[str, Any] | None, refs: list[dict[str, Any]]) -> dict[str, Any]: + """Merge serialized refs into a span data dict, deduplicating with any present.""" + out = dict(data) if isinstance(data, dict) else {} + if refs: + existing = out.get(LINEAGE_REFS_KEY) + combined = list(existing) if isinstance(existing, list) else [] + combined.extend(refs) + out[LINEAGE_REFS_KEY] = _dedupe_dicts(combined) + return out + + +def _dedupe(refs: list[DataSourceRef]) -> list[dict[str, Any]]: + return _dedupe_dicts([ref.model_dump(exclude_none=True) for ref in refs]) + + +def _dedupe_dicts(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: + seen: set[tuple[Any, ...]] = set() + out: list[dict[str, Any]] = [] + for ref in refs: + key = (ref.get("namespace"), ref.get("name"), ref.get("version"), ref.get("role")) + if key not in seen: + seen.add(key) + out.append(ref) + return out diff --git a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index 6d186de5f..32b7bae73 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -65,6 +65,8 @@ def _add_source_to_span(span: Span, env_vars: EnvironmentVariables) -> None: span.data["__agent_name__"] = env_vars.AGENT_NAME if env_vars.AGENT_ID is not None: span.data["__agent_id__"] = env_vars.AGENT_ID + if env_vars.AGENT_VERSION is not None: + span.data["__agent_version__"] = env_vars.AGENT_VERSION def _build_sgp_span(span: Span, env_vars: EnvironmentVariables) -> SGPSpan: diff --git a/src/agentex/lib/environment_variables.py b/src/agentex/lib/environment_variables.py index 3113b78f4..cbad0f2d8 100644 --- a/src/agentex/lib/environment_variables.py +++ b/src/agentex/lib/environment_variables.py @@ -24,6 +24,7 @@ class EnvVarKeys(str, Enum): AGENT_NAME = "AGENT_NAME" AGENT_DESCRIPTION = "AGENT_DESCRIPTION" AGENT_ID = "AGENT_ID" + AGENT_VERSION = "AGENT_VERSION" AGENT_API_KEY = "AGENT_API_KEY" # ACP Configuration ACP_URL = "ACP_URL" @@ -66,6 +67,8 @@ class EnvironmentVariables(BaseModel): AGENT_NAME: str AGENT_DESCRIPTION: str | None = None AGENT_ID: str | None = None + # Build/version discriminator (image tag or git sha), set by the deployment + AGENT_VERSION: str | None = None AGENT_API_KEY: str | None = None ACP_TYPE: str | None = "async" AGENT_INPUT_TYPE: str | None = None diff --git a/tests/lib/core/harness/test_tracer_lineage.py b/tests/lib/core/harness/test_tracer_lineage.py new file mode 100644 index 000000000..75799caee --- /dev/null +++ b/tests/lib/core/harness/test_tracer_lineage.py @@ -0,0 +1,53 @@ +"""SpanTracer stamps registered data-source refs onto tool spans (SGP-6513).""" + +import pytest + +from agentex.lib.core.harness.types import OpenSpan, CloseSpan +from agentex.lib.core.harness.tracer import SpanTracer +from agentex.lib.core.tracing.lineage import ( + LINEAGE_REFS_KEY, + DataSourceRef, + clear_tool_sources, + register_tool_sources, +) + +from ._fakes import FakeTracing + + +@pytest.fixture(autouse=True) +def _clean_registry(): + clear_tool_sources() + yield + clear_tool_sources() + + +@pytest.mark.asyncio +async def test_tool_open_span_carries_registered_refs(): + register_tool_sources( + "query_guidance", + refs=[DataSourceRef("databricks://ey-tax", "guidance.rulings")], + resolver=lambda args: [DataSourceRef("elasticsearch://ey", args["index"])], + ) + fake = FakeTracing() + tracer = SpanTracer(trace_id="t1", parent_span_id="p1", tracing=fake) + + await tracer.handle(OpenSpan(key="c1", kind="tool", name="query_guidance", input={"index": "filings"})) + await tracer.handle(CloseSpan(key="c1", output={"ok": True}, is_complete=True)) + + (span,) = fake.ended_spans + namespaces = {ref["namespace"] for ref in span.data[LINEAGE_REFS_KEY]} + assert namespaces == {"databricks://ey-tax", "elasticsearch://ey"} + + +@pytest.mark.asyncio +async def test_unregistered_tool_and_reasoning_spans_carry_no_refs(): + fake = FakeTracing() + tracer = SpanTracer(trace_id="t1", parent_span_id=None, tracing=fake) + + await tracer.handle(OpenSpan(key="c1", kind="tool", name="unregistered", input={})) + await tracer.handle(CloseSpan(key="c1", output=None, is_complete=True)) + await tracer.handle(OpenSpan(key="reasoning:0", kind="reasoning", name="reasoning", input={})) + await tracer.handle(CloseSpan(key="reasoning:0", output="thought", is_complete=True)) + + for span in fake.ended_spans: + assert not (isinstance(span.data, dict) and LINEAGE_REFS_KEY in span.data) diff --git a/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py b/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py index dc8bab127..4a233fb72 100644 --- a/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py +++ b/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py @@ -39,6 +39,30 @@ def _make_mock_sgp_span() -> MagicMock: return sgp_span +class TestSourceStamps: + def test_agent_identity_and_version_stamped_into_span_data(self): + from agentex.lib.core.tracing.processors.sgp_tracing_processor import _add_source_to_span + + env = MagicMock(ACP_TYPE="async", AGENT_NAME="emu-tax", AGENT_ID="a1", AGENT_VERSION="sha-abc123") + span = _make_span() + _add_source_to_span(span, env) + assert span.data == { + "__source__": "agentex", + "__acp_type__": "async", + "__agent_name__": "emu-tax", + "__agent_id__": "a1", + "__agent_version__": "sha-abc123", + } + + def test_unset_identity_fields_are_omitted(self): + from agentex.lib.core.tracing.processors.sgp_tracing_processor import _add_source_to_span + + env = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) + span = _make_span() + _add_source_to_span(span, env) + assert span.data == {"__source__": "agentex"} + + # --------------------------------------------------------------------------- # Sync processor tests # --------------------------------------------------------------------------- @@ -48,7 +72,7 @@ class TestSGPSyncTracingProcessor: @staticmethod def _make_processor(): mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) mock_create_span = MagicMock(side_effect=lambda **kwargs: _make_mock_sgp_span()) with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.SGPClient"), patch( @@ -150,7 +174,7 @@ class TestSGPAsyncTracingProcessor: @staticmethod def _make_processor(): mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) mock_create_span = MagicMock(side_effect=lambda **kwargs: _make_mock_sgp_span()) mock_async_client = MagicMock() @@ -319,11 +343,9 @@ async def test_get_client_caches_per_event_loop(self): keepalive instead of paying a TLS handshake per span. """ mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) - with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch( - f"{MODULE}.AsyncSGPClient" - ) as mock_sgp_cls: + with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient") as mock_sgp_cls: mock_sgp_cls.side_effect = lambda **kwargs: MagicMock() from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( @@ -365,11 +387,11 @@ def capture_limits(*args, **kwargs): return original_async_client(*args, **kwargs) mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) - with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch( - f"{MODULE}.AsyncSGPClient" - ), patch("httpx.AsyncClient", side_effect=capture_limits): + with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient"), patch( + "httpx.AsyncClient", side_effect=capture_limits + ): from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( SGPAsyncTracingProcessor, ) @@ -380,8 +402,7 @@ def capture_limits(*args, **kwargs): assert len(captured_limits) == 1 max_keepalive = captured_limits[0].max_keepalive_connections assert max_keepalive is not None and max_keepalive > 0, ( - f"SGP async client should have keepalive enabled, got " - f"max_keepalive_connections={max_keepalive}" + f"SGP async client should have keepalive enabled, got max_keepalive_connections={max_keepalive}" ) def test_cache_is_weakkeydict_and_evicts_dead_loops(self): @@ -395,7 +416,7 @@ def test_cache_is_weakkeydict_and_evicts_dead_loops(self): import weakref mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient"): from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( @@ -428,18 +449,14 @@ async def test_disabled_processor_returns_none_client(self): from agentex.lib.types.tracing import SGPTracingProcessorConfig mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) - with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch( - f"{MODULE}.AsyncSGPClient" - ) as mock_sgp_cls: + with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient") as mock_sgp_cls: from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( SGPAsyncTracingProcessor, ) - processor = SGPAsyncTracingProcessor( - SGPTracingProcessorConfig(sgp_api_key="", sgp_account_id="") - ) + processor = SGPAsyncTracingProcessor(SGPTracingProcessorConfig(sgp_api_key="", sgp_account_id="")) assert processor._get_client() is None assert mock_sgp_cls.call_count == 0 diff --git a/tests/lib/core/tracing/test_lineage.py b/tests/lib/core/tracing/test_lineage.py new file mode 100644 index 000000000..c0fc3ebb9 --- /dev/null +++ b/tests/lib/core/tracing/test_lineage.py @@ -0,0 +1,147 @@ +"""Unit tests for the data-source ref module (sgp.lineage.refs capture).""" + +import json + +import pytest +from pydantic import ValidationError + +from agentex.lib.core.tracing.lineage import ( + LINEAGE_REFS_KEY, + DataSourceRef, + record, + data_sources, + resolve_refs, + clear_tool_sources, + merge_refs_into_data, + register_tool_sources, + resolve_refs_from_items, +) + + +@pytest.fixture(autouse=True) +def _clean_registry(): + clear_tool_sources() + yield + clear_tool_sources() + + +ES_REF = DataSourceRef("elasticsearch://ey-embryonic", "companies_v3") +DBX_REF = DataSourceRef("databricks://ey-tax", "guidance.rulings", role="input") + + +class TestDataSourceRef: + def test_positional_construction(self): + ref = DataSourceRef("s3://bucket", "key", version="v1", role="output") + assert ref.namespace == "s3://bucket" + assert ref.name == "key" + assert ref.version == "v1" + assert ref.role == "output" + + def test_non_uri_namespace_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("not-a-uri", "name") + + def test_underscore_host_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("mcp://ey_tax_server", "competitive-edge") + + def test_host_with_path_segment_allowed(self): + DataSourceRef("confluence://ey-tax/TAX", "page-123") + + def test_empty_name_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("s3://bucket", "") + + def test_bad_role_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("s3://bucket", "key", role="sideways") + + +class TestRegistryAndResolve: + def test_unregistered_tool_resolves_empty(self): + assert resolve_refs("unknown_tool", {}) == [] + + def test_static_refs(self): + register_tool_sources("search", refs=[ES_REF]) + refs = resolve_refs("search", {"q": "acme"}) + assert refs == [{"namespace": "elasticsearch://ey-embryonic", "name": "companies_v3", "role": "input"}] + + def test_resolver_refs_combined_with_static(self): + register_tool_sources( + "query_table", + refs=[ES_REF], + resolver=lambda args: [DataSourceRef("databricks://ey-tax", args["table"])], + ) + refs = resolve_refs("query_table", {"table": "guidance.rulings"}) + assert {r["namespace"] for r in refs} == {"elasticsearch://ey-embryonic", "databricks://ey-tax"} + + def test_resolver_failure_keeps_static_refs(self): + register_tool_sources("flaky", refs=[ES_REF], resolver=lambda args: args["missing"]) + refs = resolve_refs("flaky", {}) + assert len(refs) == 1 + + def test_reregistration_replaces(self): + register_tool_sources("search", refs=[ES_REF]) + register_tool_sources("search", refs=[DBX_REF]) + assert resolve_refs("search", {})[0]["namespace"] == "databricks://ey-tax" + + def test_dedupe(self): + register_tool_sources("search", refs=[ES_REF, ES_REF]) + assert len(resolve_refs("search", {})) == 1 + + +class TestDecorator: + def test_registers_by_function_name(self): + @data_sources(ES_REF) + def search_companies(q: str) -> str: + return q + + assert search_companies("x") == "x" + assert resolve_refs("search_companies", {}) != [] + + def test_registers_by_name_attribute(self): + class FakeFunctionTool: + name = "mcp_search" + + data_sources(DBX_REF)(FakeFunctionTool()) + assert resolve_refs("mcp_search", {}) != [] + + +class TestResolveFromItems: + def test_matches_function_call_items_and_parses_string_arguments(self): + register_tool_sources( + "query_table", + resolver=lambda args: [DataSourceRef("databricks://ey-tax", args["table"])], + ) + items = [ + {"type": "message", "content": []}, + {"type": "function_call", "name": "query_table", "arguments": json.dumps({"table": "t1"})}, + {"type": "function_call", "name": "unregistered", "arguments": "{}"}, + "not-a-dict", + ] + refs = resolve_refs_from_items(items) + assert refs == [{"namespace": "databricks://ey-tax", "name": "t1", "role": "input"}] + + def test_malformed_arguments_fall_back_to_static(self): + register_tool_sources("search", refs=[ES_REF]) + items = [{"type": "function_call", "name": "search", "arguments": "{not json"}] + assert len(resolve_refs_from_items(items)) == 1 + + +class TestRecordAndMerge: + def test_record_on_none_span_is_noop(self): + record(None, [ES_REF]) + + def test_record_merges_into_span_data(self): + class Span: + data = {"__span_type__": "CUSTOM"} + + span = Span() + record(span, [ES_REF]) + assert span.data["__span_type__"] == "CUSTOM" + assert span.data[LINEAGE_REFS_KEY][0]["name"] == "companies_v3" + + def test_merge_dedupes_against_existing(self): + data = merge_refs_into_data(None, [ES_REF.model_dump(exclude_none=True)]) + data = merge_refs_into_data(data, [ES_REF.model_dump(exclude_none=True)]) + assert len(data[LINEAGE_REFS_KEY]) == 1