From 896404475a1e93955eabd562caa1670364335c29 Mon Sep 17 00:00:00 2001 From: Max Parke Date: Tue, 28 Jul 2026 21:43:24 -0400 Subject: [PATCH 1/3] feat(lib): capture client-attested build provenance (#454) Co-authored-by: Claude Opus 4.8 Co-authored-by: Nitesh Dhanpal Co-authored-by: stainless-app[bot] <142633134+stainless-app[bot]@users.noreply.github.com> --- src/agentex/lib/sdk/config/agent_manifest.py | 10 +- src/agentex/lib/utils/build_provenance.py | 189 ++++++++++++++ tests/lib/test_build_provenance.py | 257 +++++++++++++++++++ uv.lock | 4 +- 4 files changed, 453 insertions(+), 7 deletions(-) create mode 100644 src/agentex/lib/utils/build_provenance.py create mode 100644 tests/lib/test_build_provenance.py diff --git a/src/agentex/lib/sdk/config/agent_manifest.py b/src/agentex/lib/sdk/config/agent_manifest.py index fd743e635..c2fe03052 100644 --- a/src/agentex/lib/sdk/config/agent_manifest.py +++ b/src/agentex/lib/sdk/config/agent_manifest.py @@ -24,6 +24,7 @@ from agentex.lib.utils.io import load_yaml_file from agentex.lib.utils.logging import make_logger from agentex.config.agent_manifest import AgentManifest # noqa: F401 +from agentex.lib.utils.build_provenance import iter_context_files logger = make_logger(__name__) @@ -189,12 +190,11 @@ def zipped(root_path: Path | None = None) -> Iterator[IO[bytes]]: tar_buffer = io.BytesIO() + # Sorted, relpath-stable enumeration (shared with the content hash) so the + # archive's member order is deterministic across machines. with tarfile.open(fileobj=tar_buffer, mode="w:gz") as tar_file: - for path in Path(root_path).rglob( - "*" - ): # Recursively add files to the tar.gz - if path.is_file(): # Ensure that we're only adding files - tar_file.add(path, arcname=path.relative_to(root_path)) + for path in iter_context_files(Path(root_path)): + tar_file.add(path, arcname=path.relative_to(root_path)) tar_buffer.seek(0) # Reset the buffer position to the beginning yield tar_buffer diff --git a/src/agentex/lib/utils/build_provenance.py b/src/agentex/lib/utils/build_provenance.py new file mode 100644 index 000000000..447980263 --- /dev/null +++ b/src/agentex/lib/utils/build_provenance.py @@ -0,0 +1,189 @@ +"""Capture client-attested source identity without failing agent builds.""" + +from __future__ import annotations + +import os +import stat +import hashlib +import subprocess +from typing import Optional +from pathlib import Path +from datetime import datetime, timezone +from dataclasses import dataclass + +from agentex.lib.utils.logging import make_logger + +logger = make_logger(__name__) + +_GIT_TIMEOUT_S = 5 +_HASH_CHUNK_BYTES = 1 << 20 + + +@dataclass(frozen=True) +class BuildProvenance: + """Source identity for one build; unavailable fields degrade to ``None``.""" + + repo: Optional[str] = None + commit: Optional[str] = None + ref: Optional[str] = None + subpath: Optional[str] = None + working_tree_hash: Optional[str] = None + dirty: Optional[bool] = None + author_name: Optional[str] = None + author_email: Optional[str] = None + build_timestamp: Optional[str] = None + + def source_fields(self) -> dict[str, object]: + """The ``source_*`` form fields for the cloud-build upload (None omitted).""" + fields = { + "source_repo": self.repo, + "source_commit": self.commit, + "source_ref": self.ref, + "source_subpath": self.subpath, + "working_tree_hash": self.working_tree_hash, + "source_dirty": self.dirty, + } + return {key: value for key, value in fields.items() if value is not None} + + def build_info(self) -> dict[str, object]: + """Return provenance using the runtime registration metadata field names.""" + info = { + "repo": self.repo, + "commit_hash": self.commit, + "branch_name": self.ref, + "subpath": self.subpath, + "working_tree_hash": self.working_tree_hash, + "dirty": self.dirty, + "author_name": self.author_name, + "author_email": self.author_email, + "build_timestamp": self.build_timestamp, + } + return {key: value for key, value in info.items() if value is not None} + + +def _git(repo_root: Path, *args: str) -> Optional[str]: + """Run a git command under ``repo_root``; return stripped stdout or None.""" + try: + proc = subprocess.run( + ("git", "-C", str(repo_root), *args), + capture_output=True, + text=True, + timeout=_GIT_TIMEOUT_S, + check=False, + ) + except (OSError, subprocess.SubprocessError): + return None + if proc.returncode != 0: + return None + return proc.stdout.strip() or None + + +def normalize_remote(url: Optional[str]) -> Optional[str]: + """Strip credentials and scheme from a remote, returning ``host/path``.""" + if not url: + return None + candidate = url.strip() + # scp-like syntax: git@host:org/repo(.git) — no scheme, host/path split on ':' + if "://" not in candidate and ":" in candidate and "/" not in candidate.split(":", 1)[0]: + candidate = candidate.split("@", 1)[-1].replace(":", "/", 1) + else: + if "://" in candidate: + candidate = candidate.split("://", 1)[1] + candidate = candidate.split("@", 1)[-1] + if candidate.endswith(".git"): + candidate = candidate[: -len(".git")] + candidate = candidate.strip("/") + if not candidate: + return None + host, slash, path = candidate.partition("/") + return f"{host.lower()}{slash}{path}" + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with open(path, "rb") as handle: + while chunk := handle.read(_HASH_CHUNK_BYTES): + digest.update(chunk) + return digest.hexdigest() + + +def iter_context_files(root: Path) -> list[Path]: + """Return files and symlinks under ``root``, sorted by POSIX relative path.""" + return sorted( + (path for path in root.rglob("*") if path.is_symlink() or path.is_file()), + key=lambda path: path.relative_to(root).as_posix(), + ) + + +def working_tree_hash(root: Path) -> str: + """Hash sorted build inputs, normalized modes, and symlink target strings.""" + lines: list[str] = [] + for path in iter_context_files(root): + relpath = path.relative_to(root).as_posix() + if path.is_symlink(): + mode = "120000" + content_digest = hashlib.sha256(os.readlink(path).encode("utf-8")).hexdigest() + else: + executable = bool(path.stat().st_mode & stat.S_IXUSR) + mode = "100755" if executable else "100644" + content_digest = _sha256_file(path) + lines.append(f"{relpath}\x00{mode}\x00{content_digest}") + return hashlib.sha256("\n".join(lines).encode("utf-8")).hexdigest() + + +def _safe_working_tree_hash(root: Path) -> Optional[str]: + """Compute the context hash without allowing capture to fail a build.""" + try: + return working_tree_hash(root) + except Exception: + logger.warning("build-provenance: content hash failed; omitting", exc_info=True) + return None + + +def capture_build_provenance( + repo_path: Path, context_root: Path, content_root: Optional[Path] = None +) -> BuildProvenance: + """Capture git coordinates and the staged build-context hash.""" + timestamp = datetime.now(timezone.utc).isoformat() + hash_root = content_root if content_root is not None else context_root + tree_hash = _safe_working_tree_hash(hash_root) + + repo_root = _git(repo_path, "rev-parse", "--show-toplevel") + if repo_root is None: + # No git — the content hash is the only identity available. + logger.info("build-provenance: %s is not a git work tree; content hash only", repo_path) + return BuildProvenance(working_tree_hash=tree_hash, build_timestamp=timestamp) + + repo_root_path = Path(repo_root) + commit = _git(repo_root_path, "rev-parse", "HEAD") + # symbolic-ref fails on a detached HEAD (→ None); fall back to an exact tag. + ref = _git(repo_root_path, "symbolic-ref", "--short", "HEAD") or _git( + repo_root_path, "describe", "--tags", "--exact-match" + ) + remote = normalize_remote(_git(repo_root_path, "remote", "get-url", "origin")) + author_name = _git(repo_root_path, "log", "-1", "--format=%an") + author_email = _git(repo_root_path, "log", "-1", "--format=%ae") + + subpath: Optional[str] = None + try: + relative = context_root.resolve().relative_to(repo_root_path.resolve()).as_posix() + subpath = relative if relative != "." else None + except ValueError: + subpath = None + + status_args = ("status", "--porcelain") + if subpath is not None: + status_args += ("--", subpath) + dirty = _git(repo_root_path, *status_args) is not None + + return BuildProvenance( + repo=remote, + commit=commit, + ref=ref, + subpath=subpath, + working_tree_hash=tree_hash, + dirty=dirty, + author_name=author_name, + author_email=author_email, + build_timestamp=timestamp, + ) diff --git a/tests/lib/test_build_provenance.py b/tests/lib/test_build_provenance.py new file mode 100644 index 000000000..9115e2804 --- /dev/null +++ b/tests/lib/test_build_provenance.py @@ -0,0 +1,257 @@ +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from agentex.lib.utils.build_provenance import ( + normalize_remote, + working_tree_hash, + iter_context_files, + capture_build_provenance, +) + + +def _git(repo: Path, *args: str) -> None: + subprocess.run(("git", "-C", str(repo), *args), check=True, capture_output=True, text=True) + + +def _init_repo(path: Path, *, remote: str | None = "git@github.com:scaleapi/demo.git") -> Path: + path.mkdir(parents=True, exist_ok=True) + _git(path, "init", "-q") + _git(path, "config", "user.email", "dev@scale.com") + _git(path, "config", "user.name", "Dev") + _git(path, "config", "commit.gpgsign", "false") + if remote: + _git(path, "remote", "add", "origin", remote) + return path + + +def _commit_all(path: Path, message: str = "init") -> None: + _git(path, "add", "-A") + _git(path, "commit", "-q", "-m", message) + _git(path, "branch", "-M", "main") + + +def _write(root: Path, rel: str, content: str = "x") -> None: + target = root / rel + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(content) + + +# --- normalize_remote --------------------------------------------------------- + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("git@github.com:scaleapi/Repo.git", "github.com/scaleapi/Repo"), + ("https://github.com/scaleapi/Repo.git", "github.com/scaleapi/Repo"), + ("https://x-token:secret@GitHub.com/scaleapi/Repo", "github.com/scaleapi/Repo"), + ("ssh://git@gitlab.com/group/sub/proj.git", "gitlab.com/group/sub/proj"), + ("", None), + (None, None), + ], +) +def test_normalize_remote(raw: str | None, expected: str | None) -> None: + assert normalize_remote(raw) == expected + + +# --- working_tree_hash -------------------------------------------------------- + + +def test_hash_is_order_independent(tmp_path: Path) -> None: + first = tmp_path / "a" + second = tmp_path / "b" + for rel in ("z.txt", "a/b.txt", "m.txt"): + _write(first, rel, rel) + # Same content, different creation order. + for rel in ("m.txt", "z.txt", "a/b.txt"): + _write(second, rel, rel) + assert working_tree_hash(first) == working_tree_hash(second) + + +def test_hash_changes_on_one_byte(tmp_path: Path) -> None: + root = tmp_path / "ctx" + _write(root, "f.txt", "hello") + before = working_tree_hash(root) + _write(root, "f.txt", "hellp") + assert working_tree_hash(root) != before + + +def test_hash_changes_when_file_added(tmp_path: Path) -> None: + root = tmp_path / "ctx" + _write(root, "f.txt", "hello") + before = working_tree_hash(root) + _write(root, "g.txt", "new") + assert working_tree_hash(root) != before + + +def test_hash_changes_on_executable_bit(tmp_path: Path) -> None: + root = tmp_path / "ctx" + script = root / "run.sh" + _write(root, "run.sh", "#!/bin/sh\n") + before = working_tree_hash(root) + script.chmod(0o755) + assert working_tree_hash(root) != before + + +def test_symlink_hashes_target_not_resolved_content(tmp_path: Path) -> None: + root = tmp_path / "ctx" + root.mkdir() + # Dangling symlinks: distinct hashes prove the target string is hashed, not + # resolved content (resolving would raise). + (root / "link").symlink_to("points/to/a") + hash_a = working_tree_hash(root) + (root / "link").unlink() + (root / "link").symlink_to("points/to/b") + assert working_tree_hash(root) != hash_a + + +def test_iter_context_files_skips_directories(tmp_path: Path) -> None: + root = tmp_path / "ctx" + _write(root, "pkg/mod.py", "x") + _write(root, "top.txt", "y") + rels = [path.relative_to(root).as_posix() for path in iter_context_files(root)] + assert rels == ["pkg/mod.py", "top.txt"] + + +# --- capture_build_provenance ------------------------------------------------- + + +def test_capture_clean_tree(tmp_path: Path) -> None: + repo = _init_repo(tmp_path / "repo") + _write(repo, "main.py", "print(1)") + _commit_all(repo) + + prov = capture_build_provenance(repo, repo) + + assert prov.repo == "github.com/scaleapi/demo" + assert prov.ref == "main" + assert prov.commit is not None and len(prov.commit) == 40 + assert prov.working_tree_hash is not None # always computed + assert prov.dirty is False + assert prov.subpath is None + assert prov.author_email == "dev@scale.com" + + +def test_capture_untracked_file_changes_hash(tmp_path: Path) -> None: + repo = _init_repo(tmp_path / "repo") + _write(repo, "main.py", "print(1)") + _commit_all(repo) + _write(repo, "scratch.py", "debug = True") # untracked + + prov = capture_build_provenance(repo, repo) + + # The stale-code guard: an untracked file is part of the build context, so it + # must move the hash (a `git diff` of tracked files alone would miss it). + assert prov.dirty is True + assert prov.working_tree_hash == working_tree_hash(repo) + assert working_tree_hash(repo) != _hash_without(repo, "scratch.py") + + +def _hash_without(repo: Path, rel: str) -> str: + removed = repo / rel + saved = removed.read_text() + removed.unlink() + try: + return working_tree_hash(repo) + finally: + removed.write_text(saved) + + +def test_capture_detached_head_has_no_ref(tmp_path: Path) -> None: + repo = _init_repo(tmp_path / "repo") + _write(repo, "main.py", "print(1)") + _commit_all(repo) + _write(repo, "main.py", "print(2)") + _git(repo, "add", "-A") + _git(repo, "commit", "-q", "-m", "second") + first = subprocess.run( + ("git", "-C", str(repo), "rev-list", "--max-parents=0", "HEAD"), + check=True, + capture_output=True, + text=True, + ).stdout.strip() + _git(repo, "checkout", "-q", first) + + prov = capture_build_provenance(repo, repo) + + assert prov.commit == first + assert prov.ref is None + + +def test_capture_detached_on_tag_uses_tag(tmp_path: Path) -> None: + repo = _init_repo(tmp_path / "repo") + _write(repo, "main.py", "print(1)") + _commit_all(repo) + _git(repo, "tag", "v1.2.3") + _git(repo, "checkout", "-q", "v1.2.3") + + assert capture_build_provenance(repo, repo).ref == "v1.2.3" + + +def test_capture_no_remote(tmp_path: Path) -> None: + repo = _init_repo(tmp_path / "repo", remote=None) + _write(repo, "main.py", "print(1)") + _commit_all(repo) + + prov = capture_build_provenance(repo, repo) + + assert prov.repo is None + assert prov.commit is not None + assert prov.working_tree_hash is not None # always computed + + +def test_capture_non_git_dir(tmp_path: Path) -> None: + plain = tmp_path / "plain" + _write(plain, "main.py", "print(1)") + + prov = capture_build_provenance(plain, plain) + + assert prov.repo is None + assert prov.commit is None + assert prov.ref is None + # No commit → the content hash is the identity; dirtiness is undefined (no VCS). + assert prov.working_tree_hash == working_tree_hash(plain) + assert prov.dirty is None + assert prov.build_timestamp is not None + + +def test_capture_never_raises_when_hash_fails(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + import agentex.lib.utils.build_provenance as bp + + plain = tmp_path / "plain" # non-git → would hash, which we force to fail + _write(plain, "main.py", "print(1)") + + def _boom(_root: Path) -> str: + raise OSError("permission denied") + + monkeypatch.setattr(bp, "working_tree_hash", _boom) + + prov = bp.capture_build_provenance(plain, plain) # must not raise + + assert prov.working_tree_hash is None + + +def test_capture_monorepo_subpath(tmp_path: Path) -> None: + repo = _init_repo(tmp_path / "repo") + _write(repo, "agents/foo/main.py", "print(1)") + _commit_all(repo) + + prov = capture_build_provenance(repo, repo / "agents" / "foo") + + assert prov.subpath == "agents/foo" + + +def test_capture_monorepo_ignores_changes_outside_context(tmp_path: Path) -> None: + repo = _init_repo(tmp_path / "repo") + _write(repo, "agents/foo/main.py", "print(1)") + _write(repo, "agents/bar/main.py", "print(2)") + _commit_all(repo) + _write(repo, "agents/bar/scratch.py", "debug = True") + + prov = capture_build_provenance(repo, repo / "agents" / "foo") + + assert prov.dirty is False diff --git a/uv.lock b/uv.lock index 43de6493c..f925c2dce 100644 --- a/uv.lock +++ b/uv.lock @@ -15,7 +15,7 @@ members = [ [[package]] name = "agentex-client" -version = "0.17.0" +version = "0.21.0" source = { editable = "." } dependencies = [ { name = "anyio" }, @@ -91,7 +91,7 @@ dev = [ [[package]] name = "agentex-sdk" -version = "0.17.0" +version = "0.21.0" source = { editable = "adk" } dependencies = [ { name = "agentex-client" }, From 2078f9fabb3099fa2d5d02f67ed5af50b8efbc09 Mon Sep 17 00:00:00 2001 From: Max Parke Date: Wed, 29 Jul 2026 11:34:39 -0400 Subject: [PATCH 2/3] refactor(lib): remove the dead build-info.json registration read-path (#455) Co-authored-by: Claude Opus 4.8 --- src/agentex/lib/environment_variables.py | 4 -- src/agentex/lib/sdk/fastacp/fastacp.py | 13 ------ src/agentex/lib/utils/registration.py | 15 +------ tests/lib/test_agent_card.py | 51 ++++++++++++------------ 4 files changed, 27 insertions(+), 56 deletions(-) diff --git a/src/agentex/lib/environment_variables.py b/src/agentex/lib/environment_variables.py index 3113b78f4..4afae3e60 100644 --- a/src/agentex/lib/environment_variables.py +++ b/src/agentex/lib/environment_variables.py @@ -37,8 +37,6 @@ class EnvVarKeys(str, Enum): HEALTH_CHECK_PORT = "HEALTH_CHECK_PORT" # Auth Configuration AUTH_PRINCIPAL_B64 = "AUTH_PRINCIPAL_B64" - # Build Information - BUILD_INFO_PATH = "BUILD_INFO_PATH" AGENT_INPUT_TYPE = "AGENT_INPUT_TYPE" # Deployment AGENTEX_DEPLOYMENT_ID = "AGENTEX_DEPLOYMENT_ID" @@ -87,8 +85,6 @@ class EnvironmentVariables(BaseModel): HEALTH_CHECK_PORT: int = 80 # Auth Configuration AUTH_PRINCIPAL_B64: str | None = None - # Build Information - BUILD_INFO_PATH: str | None = None # Deployment AGENTEX_DEPLOYMENT_ID: str | None = None # Claude Agents SDK Configuration diff --git a/src/agentex/lib/sdk/fastacp/fastacp.py b/src/agentex/lib/sdk/fastacp/fastacp.py index 42859793d..4a76c294e 100644 --- a/src/agentex/lib/sdk/fastacp/fastacp.py +++ b/src/agentex/lib/sdk/fastacp/fastacp.py @@ -1,9 +1,6 @@ from __future__ import annotations -import os -import inspect from typing import Any, Literal -from pathlib import Path from typing_extensions import deprecated from agentex.lib.types.fastacp import ( @@ -82,14 +79,6 @@ def create_agentic_acp(config: AgenticACPConfig, **kwargs) -> BaseACPServer: """ return FastACP.create_async_acp(config, **kwargs) - @staticmethod - def locate_build_info_path() -> None: - """If a build-info.json file is present, set the BUILD_INFO_PATH environment variable""" - acp_root = Path(inspect.stack()[2].filename).resolve().parents[0] - build_info_path = acp_root / "build-info.json" - if build_info_path.exists(): - os.environ["BUILD_INFO_PATH"] = str(build_info_path) - @staticmethod def create( acp_type: Literal["sync", "async", "agentic"], @@ -105,8 +94,6 @@ def create( **kwargs: Additional configuration parameters """ - FastACP.locate_build_info_path() - if acp_type == "sync": sync_config = config if isinstance(config, SyncACPConfig) else None instance = FastACP.create_sync_acp(sync_config, **kwargs) diff --git a/src/agentex/lib/utils/registration.py b/src/agentex/lib/utils/registration.py index e2bfbc00c..5fc4d4be5 100644 --- a/src/agentex/lib/utils/registration.py +++ b/src/agentex/lib/utils/registration.py @@ -20,17 +20,6 @@ def get_auth_principal(env_vars: EnvironmentVariables): except Exception: return None -def get_build_info(): - build_info_path = os.environ.get("BUILD_INFO_PATH") - logger.info(f"Getting build info from {build_info_path}") - if not build_info_path: - return None - try: - with open(build_info_path, "r") as f: - return json.load(f) - except Exception: - return None - async def register_agent(env_vars: EnvironmentVariables, agent_card=None): """Register this agent with the Agentex server""" if not env_vars.AGENTEX_BASE_URL: @@ -44,8 +33,8 @@ async def register_agent(env_vars: EnvironmentVariables, agent_card=None): or f"Generic description for agent: {env_vars.AGENT_NAME}" ) - # Build registration metadata from build-info.json + deployment env var - registration_metadata = get_build_info() or {} + # Registration metadata carries the deployment id and agent card. + registration_metadata: dict = {} if env_vars.AGENTEX_DEPLOYMENT_ID: registration_metadata["deployment_id"] = env_vars.AGENTEX_DEPLOYMENT_ID if agent_card is not None: diff --git a/tests/lib/test_agent_card.py b/tests/lib/test_agent_card.py index ccde4d33c..5d57f9e8e 100644 --- a/tests/lib/test_agent_card.py +++ b/tests/lib/test_agent_card.py @@ -358,40 +358,39 @@ async def test_agent_card_merged_into_metadata(self, mock_env_vars): card = AgentCard(input_types=["text"], data_events=["result"]) mock_client = self._make_mock_client() - with patch("agentex.lib.utils.registration.get_build_info", return_value={"version": "1.0"}): - with patch("agentex.lib.utils.registration.httpx.AsyncClient", return_value=mock_client): - from agentex.lib.utils.registration import register_agent - await register_agent(mock_env_vars, agent_card=card) + with patch("agentex.lib.utils.registration.httpx.AsyncClient", return_value=mock_client): + from agentex.lib.utils.registration import register_agent - sent_data = mock_client.post.call_args.kwargs["json"] - metadata = sent_data["registration_metadata"] + await register_agent(mock_env_vars, agent_card=card) - assert "agent_card" in metadata - assert metadata["agent_card"]["input_types"] == ["text"] - assert metadata["agent_card"]["data_events"] == ["result"] - assert metadata["version"] == "1.0" + sent_data = mock_client.post.call_args.kwargs["json"] + metadata = sent_data["registration_metadata"] - async def test_none_preserved_when_no_card_no_build_info(self, mock_env_vars): + assert "agent_card" in metadata + assert metadata["agent_card"]["input_types"] == ["text"] + assert metadata["agent_card"]["data_events"] == ["result"] + + async def test_none_preserved_when_no_card(self, mock_env_vars): mock_client = self._make_mock_client() - with patch("agentex.lib.utils.registration.get_build_info", return_value=None): - with patch("agentex.lib.utils.registration.httpx.AsyncClient", return_value=mock_client): - from agentex.lib.utils.registration import register_agent - await register_agent(mock_env_vars, agent_card=None) + with patch("agentex.lib.utils.registration.httpx.AsyncClient", return_value=mock_client): + from agentex.lib.utils.registration import register_agent + + await register_agent(mock_env_vars, agent_card=None) - sent_data = mock_client.post.call_args.kwargs["json"] - assert sent_data["registration_metadata"] is None + sent_data = mock_client.post.call_args.kwargs["json"] + assert sent_data["registration_metadata"] is None - async def test_card_creates_metadata_when_build_info_none(self, mock_env_vars): + async def test_card_creates_metadata(self, mock_env_vars): card = AgentCard(input_types=["text"]) mock_client = self._make_mock_client() - with patch("agentex.lib.utils.registration.get_build_info", return_value=None): - with patch("agentex.lib.utils.registration.httpx.AsyncClient", return_value=mock_client): - from agentex.lib.utils.registration import register_agent - await register_agent(mock_env_vars, agent_card=card) + with patch("agentex.lib.utils.registration.httpx.AsyncClient", return_value=mock_client): + from agentex.lib.utils.registration import register_agent + + await register_agent(mock_env_vars, agent_card=card) - sent_data = mock_client.post.call_args.kwargs["json"] - metadata = sent_data["registration_metadata"] - assert metadata is not None - assert "agent_card" in metadata + sent_data = mock_client.post.call_args.kwargs["json"] + metadata = sent_data["registration_metadata"] + assert metadata is not None + assert "agent_card" in metadata From 74b013a8035443c91dd760f932600f9ca1002c8c Mon Sep 17 00:00:00 2001 From: Max Parke Date: Wed, 22 Jul 2026 12:29:40 -0700 Subject: [PATCH 3/3] 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 4afae3e60..7d893e462 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" @@ -64,6 +65,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