diff --git a/agent_assembly/core/runtime_interceptor.py b/agent_assembly/core/runtime_interceptor.py index ce1928db..d6790ee4 100644 --- a/agent_assembly/core/runtime_interceptor.py +++ b/agent_assembly/core/runtime_interceptor.py @@ -34,6 +34,7 @@ import json import os +from importlib import metadata from typing import Any from agent_assembly.exceptions import OpTerminatedError @@ -288,12 +289,44 @@ def connect_runtime_client(agent_id: str) -> Any | None: return None socket_path = _resolve_runtime_socket_path(agent_id) + sdk_version = _sdk_version() try: - return RuntimeClient.connect(socket_path) + # Forward the agent identity (signed into the runtime handshake, + # AAASM-3587) and the installed PyPI package version (signed into the + # handshake so downgrade detection reflects the real SDK release, + # AAASM-3683). + return RuntimeClient.connect(socket_path, agent_id, sdk_version) + except TypeError: + # Native build predates the agent_id / sdk_version parameters; fall back + # to the legacy positional signature so connecting still works against an + # older core (the version simply is not forwarded). + try: + return RuntimeClient.connect(socket_path) + except Exception: + return None except Exception: return None +def _sdk_version() -> str | None: + """Return the installed ``agent-assembly`` distribution version, or ``None``. + + Prefers ``importlib.metadata.version`` (the authoritative installed-package + version) and falls back to the in-tree ``agent_assembly.__version__`` when the + distribution metadata is unavailable (e.g. an editable checkout that was never + installed). ``None`` lets the native shim fall back to the crate version. + """ + try: + return metadata.version("agent-assembly") + except metadata.PackageNotFoundError: + try: + from agent_assembly import __version__ + + return __version__ + except Exception: + return None + + def register_agent( runtime_client: Any, agent_id: str, diff --git a/native/aa-ffi-python/Cargo.lock b/native/aa-ffi-python/Cargo.lock index e7361f8a..1c01df21 100644 --- a/native/aa-ffi-python/Cargo.lock +++ b/native/aa-ffi-python/Cargo.lock @@ -5,7 +5,7 @@ version = 4 [[package]] name = "aa-core" version = "0.0.1-beta.3" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=ebc4d7dc9da81cdf4f265d681c138500ca6e98e7#ebc4d7dc9da81cdf4f265d681c138500ca6e98e7" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=b0576a72422bad7c6f380e666aa6ad70d1939822#b0576a72422bad7c6f380e666aa6ad70d1939822" dependencies = [ "aa-security", "async-trait", @@ -36,7 +36,7 @@ dependencies = [ [[package]] name = "aa-proto" version = "0.0.1-beta.3" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=ebc4d7dc9da81cdf4f265d681c138500ca6e98e7#ebc4d7dc9da81cdf4f265d681c138500ca6e98e7" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=b0576a72422bad7c6f380e666aa6ad70d1939822#b0576a72422bad7c6f380e666aa6ad70d1939822" dependencies = [ "prost", "tonic", @@ -47,7 +47,7 @@ dependencies = [ [[package]] name = "aa-sdk-client" version = "0.0.1-beta.3" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=ebc4d7dc9da81cdf4f265d681c138500ca6e98e7#ebc4d7dc9da81cdf4f265d681c138500ca6e98e7" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=b0576a72422bad7c6f380e666aa6ad70d1939822#b0576a72422bad7c6f380e666aa6ad70d1939822" dependencies = [ "aa-proto", "aa-security", @@ -59,15 +59,17 @@ dependencies = [ "tokio", "tonic", "tracing", + "zeroize", ] [[package]] name = "aa-security" version = "0.0.1-beta.3" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=ebc4d7dc9da81cdf4f265d681c138500ca6e98e7#ebc4d7dc9da81cdf4f265d681c138500ca6e98e7" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=b0576a72422bad7c6f380e666aa6ad70d1939822#b0576a72422bad7c6f380e666aa6ad70d1939822" dependencies = [ "aho-corasick", "serde", + "serde_yaml", ] [[package]] @@ -1863,6 +1865,20 @@ name = "zeroize" version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] [[package]] name = "zmij" diff --git a/native/aa-ffi-python/Cargo.toml b/native/aa-ffi-python/Cargo.toml index 80343910..72cbc2f1 100644 --- a/native/aa-ffi-python/Cargo.toml +++ b/native/aa-ffi-python/Cargo.toml @@ -14,11 +14,13 @@ crate-type = ["cdylib"] # aa-proto, and aa-sdk-client must share one SHA so cargo resolves a single # checkout of the agent-assembly workspace. # -# Pinned to the merge commit of agent-assembly PR #1160, which adds the -# `team_id` / `parent_agent_id` fields to `AssemblyConfig` the shim now sets. -aa-core = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "ebc4d7dc9da81cdf4f265d681c138500ca6e98e7", package = "aa-core", features = ["serde"] } -aa-proto = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "ebc4d7dc9da81cdf4f265d681c138500ca6e98e7", package = "aa-proto" } -aa-sdk-client = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "ebc4d7dc9da81cdf4f265d681c138500ca6e98e7", package = "aa-sdk-client" } +# Pinned to agent-assembly AAASM-3683, which adds the `sdk_version` field to +# `AssemblyConfig` and threads it through `spawn_ipc_thread` into the signed IPC +# handshake — the version this shim now forwards from the PyPI package version. +# Re-pin to the squash-merge commit once PR #1226 lands on master. +aa-core = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "b0576a72422bad7c6f380e666aa6ad70d1939822", package = "aa-core", features = ["serde"] } +aa-proto = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "b0576a72422bad7c6f380e666aa6ad70d1939822", package = "aa-proto" } +aa-sdk-client = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "b0576a72422bad7c6f380e666aa6ad70d1939822", package = "aa-sdk-client" } prost = "0.14" pyo3 = { version = "0.29", features = ["extension-module"] } serde = { version = "1.0", features = ["derive"] } diff --git a/native/aa-ffi-python/src/lib.rs b/native/aa-ffi-python/src/lib.rs index e5ef983f..b31dfe2b 100644 --- a/native/aa-ffi-python/src/lib.rs +++ b/native/aa-ffi-python/src/lib.rs @@ -77,18 +77,35 @@ impl RuntimeClient { /// Spawns the shared client's background IPC thread and returns a handle. /// Event reporting is fire-and-forget; a failed connection surfaces on a /// later `send_event` rather than here. + /// + /// `agent_id` is the agent identity the background thread signs the runtime + /// session handshake with (AAASM-3587); pass `None` only when no identity is + /// available yet. + /// + /// `sdk_version` is the user-facing PyPI package version (`agent_assembly.__version__`) + /// the Python layer forwards so it — not the shared `aa-sdk-client` crate + /// version — is what gets signed into the handshake (AAASM-3683). `None` + /// lets the shared client fall back to the crate version (no regression vs + /// AAASM-3666). #[staticmethod] - fn connect(socket_path: String) -> PyResult { + #[pyo3(signature = (socket_path, agent_id=None, sdk_version=None))] + fn connect(socket_path: String, agent_id: Option, sdk_version: Option) -> PyResult { let config = AssemblyConfig { - agent_id: String::new(), + agent_id: agent_id.unwrap_or_default(), socket_path: Some(socket_path.clone()), // Registration is a separate explicit `register` call; connecting // the runtime UDS does not need a gateway endpoint or lineage. gateway_endpoint: None, team_id: None, parent_agent_id: None, + sdk_version, }; - let handle = spawn_ipc_thread(config.resolve_socket_path()).map_err(|error| { + let handle = spawn_ipc_thread( + config.resolve_socket_path(), + config.agent_id.clone(), + config.resolved_sdk_version(), + ) + .map_err(|error| { PyRuntimeError::new_err(format!("failed to start runtime IPC thread: {error}")) })?; Ok(Self { @@ -139,6 +156,9 @@ impl RuntimeClient { gateway_endpoint, team_id, parent_agent_id, + // The version is signed at IPC-handshake time (`connect`), not on the + // gateway register, so it is not needed for this config. + sdk_version: None, }; // `register` is async (tonic). Release the GIL and drive the future on a // private current-thread runtime so the blocking gRPC round-trip never diff --git a/test/unit/core/_fake_core.py b/test/unit/core/_fake_core.py index 702b169c..16c028d3 100644 --- a/test/unit/core/_fake_core.py +++ b/test/unit/core/_fake_core.py @@ -11,6 +11,7 @@ import sys import types +from collections.abc import Callable from typing import Any import pytest @@ -25,6 +26,9 @@ def __init__(self, decision: str = "allow", reason: str = "") -> None: self.register_calls: list[tuple[str, str, str, str | None, str | None, str | None]] = [] self.query_calls: list[tuple[Any, ...]] = [] self.register_should_raise: Exception | None = None + # Set by install_fake_core's connect to the (socket_path, agent_id, + # sdk_version) it was called with (AAASM-3683). + self.connect_args: tuple[str, str | None, str | None] | None = None def register( self, @@ -85,14 +89,43 @@ def install_fake_core( runtime_client: Any, ) -> Any: """Install a fake ``agent_assembly._core`` whose ``RuntimeClient.connect`` - returns ``runtime_client``. Returns the same client for assertions.""" + returns ``runtime_client``. Returns the same client for assertions. + + ``connect`` accepts the AAASM-3683 ``agent_id`` / ``sdk_version`` arguments + and records them on ``runtime_client.connect_args`` so callers can assert the + installed package version is forwarded into the handshake. + """ class _ConnectingRuntimeClient: @staticmethod - def connect(_socket_path: str) -> Any: + def connect(_socket_path: str, agent_id: str | None = None, sdk_version: str | None = None) -> Any: + runtime_client.connect_args = (_socket_path, agent_id, sdk_version) return runtime_client fake_core = types.ModuleType("agent_assembly._core") fake_core.RuntimeClient = _ConnectingRuntimeClient # type: ignore[attr-defined] monkeypatch.setitem(sys.modules, "agent_assembly._core", fake_core) return runtime_client + + +def install_fake_core_with_connect( + monkeypatch: pytest.MonkeyPatch, + connect: Callable[..., Any], +) -> None: + """Install a fake ``agent_assembly._core`` whose ``RuntimeClient.connect`` is + the supplied callable. + + Lets a test drive the connect path against a native build that, for example, + predates the AAASM-3683 ``sdk_version`` parameter (``connect`` raising + ``TypeError`` for the extra argument) or fails outright. + """ + + runtime_client_cls = type( + "_CustomConnectRuntimeClient", + (), + {"connect": staticmethod(connect)}, + ) + + fake_core = types.ModuleType("agent_assembly._core") + fake_core.RuntimeClient = runtime_client_cls # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "agent_assembly._core", fake_core) diff --git a/test/unit/core/test_init_registration.py b/test/unit/core/test_init_registration.py index 8eda321b..460cce7e 100644 --- a/test/unit/core/test_init_registration.py +++ b/test/unit/core/test_init_registration.py @@ -18,7 +18,12 @@ from agent_assembly.core.spawn import SpawnContext, spawn_context_scope from agent_assembly.exceptions import ConfigurationError -from ._fake_core import FakeRuntimeClient, LegacyRuntimeClient, install_fake_core +from ._fake_core import ( + FakeRuntimeClient, + LegacyRuntimeClient, + install_fake_core, + install_fake_core_with_connect, +) _GW_URL = "http://gateway.test" _API_KEY = "test-key" @@ -266,6 +271,165 @@ def test_init_assembly_lineage_values_round_trip_verbatim( context.shutdown() +def test_connect_forwards_agent_id_and_installed_package_version( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """connect_runtime_client forwards the agent id and the installed PyPI + package version into the native ``connect`` so the language-package version + (not the crate version) is signed into the handshake (AAASM-3683).""" + from agent_assembly.core import runtime_interceptor + + runtime_client = FakeRuntimeClient(decision="allow") + install_fake_core(monkeypatch, runtime_client) + # Pin the resolved package version deterministically. + monkeypatch.setattr(runtime_interceptor, "_sdk_version", lambda: "7.8.9") + + result = runtime_interceptor.connect_runtime_client("agent-vers") + + assert result is runtime_client + assert runtime_client.connect_args is not None + _socket, agent_id, sdk_version = runtime_client.connect_args + assert agent_id == "agent-vers" + assert sdk_version == "7.8.9" + + +def test_connect_passes_none_version_when_unresolvable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """When the package version cannot be resolved, ``None`` is forwarded so the + native shim falls back to the crate version (no regression, AAASM-3683).""" + from agent_assembly.core import runtime_interceptor + + runtime_client = FakeRuntimeClient(decision="allow") + install_fake_core(monkeypatch, runtime_client) + monkeypatch.setattr(runtime_interceptor, "_sdk_version", lambda: None) + + runtime_interceptor.connect_runtime_client("agent-novers") + + assert runtime_client.connect_args is not None + _socket, _agent_id, sdk_version = runtime_client.connect_args + assert sdk_version is None + + +def test_sdk_version_reads_installed_distribution_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """_sdk_version prefers importlib.metadata.version('agent-assembly').""" + import importlib.metadata + + from agent_assembly.core import runtime_interceptor + + monkeypatch.setattr(importlib.metadata, "version", lambda _name: "3.2.1") + assert runtime_interceptor._sdk_version() == "3.2.1" + + +def test_sdk_version_falls_back_to_module_version_when_metadata_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """When the distribution is not installed (e.g. an editable checkout that was + never ``pip install``-ed), ``_sdk_version`` falls back to the in-tree + ``agent_assembly.__version__`` rather than returning ``None`` (AAASM-3683).""" + import importlib.metadata + + import agent_assembly + from agent_assembly.core import runtime_interceptor + + def _raise_not_found(_name: str) -> str: + raise importlib.metadata.PackageNotFoundError(_name) + + monkeypatch.setattr(importlib.metadata, "version", _raise_not_found) + monkeypatch.setattr(agent_assembly, "__version__", "9.9.9-editable", raising=False) + + assert runtime_interceptor._sdk_version() == "9.9.9-editable" + + +def test_sdk_version_returns_none_when_no_version_available( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """If neither the distribution metadata nor ``agent_assembly.__version__`` can + be resolved, ``_sdk_version`` returns ``None`` so the native shim falls back + to the crate version instead of raising (AAASM-3683).""" + import importlib.metadata + + import agent_assembly + from agent_assembly.core import runtime_interceptor + + def _raise_not_found(_name: str) -> str: + raise importlib.metadata.PackageNotFoundError(_name) + + monkeypatch.setattr(importlib.metadata, "version", _raise_not_found) + # Remove the in-tree fallback so ``from agent_assembly import __version__`` fails. + monkeypatch.delattr(agent_assembly, "__version__", raising=False) + + assert runtime_interceptor._sdk_version() is None + + +def test_connect_falls_back_to_legacy_signature_on_typeerror( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Against an older native build whose ``connect`` predates the + ``agent_id`` / ``sdk_version`` parameters, the three-arg call raises + ``TypeError`` and the SDK retries with the legacy single-arg signature so a + connection is still established (AAASM-3683).""" + from agent_assembly.core import runtime_interceptor + + legacy_client = object() + calls: list[tuple[Any, ...]] = [] + + def _legacy_connect(*args: Any) -> Any: + calls.append(args) + # Old native build only accepts the socket path. + if len(args) != 1: + raise TypeError("connect() takes 1 positional argument") + return legacy_client + + install_fake_core_with_connect(monkeypatch, _legacy_connect) + monkeypatch.setattr(runtime_interceptor, "_sdk_version", lambda: "1.2.3") + + result = runtime_interceptor.connect_runtime_client("legacy-agent") + + # The three-arg attempt is made first, then the one-arg legacy retry. + assert len(calls) == 2 + assert len(calls[0]) == 3 + assert len(calls[1]) == 1 + assert result is legacy_client + + +def test_connect_returns_none_when_legacy_retry_also_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """If the legacy single-arg ``connect`` retry also raises (e.g. the socket is + unreachable), ``connect_runtime_client`` returns ``None`` rather than + propagating, so there is simply no native fast path (AAASM-3683).""" + from agent_assembly.core import runtime_interceptor + + def _always_failing_connect(*args: Any) -> Any: + if len(args) != 1: + raise TypeError("connect() takes 1 positional argument") + raise OSError("socket unreachable") + + install_fake_core_with_connect(monkeypatch, _always_failing_connect) + monkeypatch.setattr(runtime_interceptor, "_sdk_version", lambda: "1.2.3") + + assert runtime_interceptor.connect_runtime_client("legacy-agent") is None + + +def test_connect_returns_none_on_generic_connect_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A non-``TypeError`` failure from the native ``connect`` (e.g. the runtime + socket is unreachable) yields ``None`` rather than raising (AAASM-3683).""" + from agent_assembly.core import runtime_interceptor + + def _failing_connect(*_args: Any) -> Any: + raise OSError("runtime socket unreachable") + + install_fake_core_with_connect(monkeypatch, _failing_connect) + monkeypatch.setattr(runtime_interceptor, "_sdk_version", lambda: "1.2.3") + + assert runtime_interceptor.connect_runtime_client("agent-x") is None + + def test_init_assembly_deny_blocks_tool_via_interceptor(monkeypatch: pytest.MonkeyPatch) -> None: """A native ``deny`` makes the adapter interceptor's check_tool_start block.""" runtime_client = FakeRuntimeClient(decision="deny", reason="policy violation")