diff --git a/agent_assembly/core/assembly.py b/agent_assembly/core/assembly.py index c5e5fac..a872c8f 100644 --- a/agent_assembly/core/assembly.py +++ b/agent_assembly/core/assembly.py @@ -228,6 +228,8 @@ def init_assembly( runtime_client=runtime_client, agent_id=resolved_agent_id, enforcement_mode=enforcement_mode, + team_id=team_id, + parent_agent_id=parent_agent_id, ) registered_adapters = _register_adapters( client=client, @@ -287,6 +289,8 @@ def _register_agent_with_gateway( runtime_client: Any | None, agent_id: str, enforcement_mode: EnforcementMode | None, + team_id: str | None = None, + parent_agent_id: str | None = None, ) -> None: """Register the agent with the gateway over the native gRPC ``register``. @@ -295,6 +299,10 @@ def _register_agent_with_gateway( ``RuntimeQueryInterceptor`` later uses for ``query_policy`` — the SDK never calls a core HTTP endpoint directly (ADR 0004). + ``team_id`` and ``parent_agent_id`` are forwarded to the native register so + the gateway gets the agent's team-budget scoping and topology lineage on the + native path, restoring what the legacy REST register sent (AAASM-3415). + No native runtime (extension missing or socket unreachable) means there is nothing to register against: the call is skipped. Under ``enforce`` a native registration failure propagates so a misconfigured gateway surfaces at init; @@ -305,7 +313,13 @@ def _register_agent_with_gateway( return framework = "python" try: - register_agent(runtime_client, agent_id, framework) + register_agent( + runtime_client, + agent_id, + framework, + team_id=team_id, + parent_agent_id=parent_agent_id, + ) except Exception: if enforcement_mode == "enforce": raise diff --git a/agent_assembly/core/runtime_interceptor.py b/agent_assembly/core/runtime_interceptor.py index fe9a44a..b9194ba 100644 --- a/agent_assembly/core/runtime_interceptor.py +++ b/agent_assembly/core/runtime_interceptor.py @@ -236,6 +236,8 @@ def register_agent( agent_id: str, framework: str, gateway_endpoint: str | None = None, + team_id: str | None = None, + parent_agent_id: str | None = None, ) -> str | None: """Register ``agent_id`` with the gateway over the native ``register`` call. @@ -244,6 +246,13 @@ def register_agent( token on the shared client so later ``query_policy`` checks authenticate (ADR 0004 — the SDK never calls core HTTP endpoints directly). + ``team_id`` and ``parent_agent_id`` carry the agent's lineage/team scoping to + the gateway (AAASM-3415): ``team_id`` drives team-budget attribution and + ``parent_agent_id`` the topology graph. They are forwarded to a native build + that accepts them; an older build (whose ``register`` predates these kwargs) + is retried with the legacy positional signature so the SDK keeps working + rather than raising. + Returns the policy id the gateway assigned, or ``None`` when ``register`` is not exposed (older native build). Registration is authoritative: a native failure raises ``RuntimeError`` and is allowed to propagate so init surfaces @@ -252,7 +261,13 @@ def register_agent( register = getattr(runtime_client, "register", None) if register is None: return None - return str(register(agent_id, agent_id, framework, gateway_endpoint)) + try: + return str(register(agent_id, agent_id, framework, gateway_endpoint, team_id, parent_agent_id)) + except TypeError: + # Native build predates the team_id/parent_agent_id parameters + # (AAASM-3415). Fall back to the legacy signature so registration still + # succeeds — lineage/team are simply not forwarded against an old core. + return str(register(agent_id, agent_id, framework, gateway_endpoint)) def _native_core_available() -> bool: diff --git a/native/aa-ffi-python/Cargo.lock b/native/aa-ffi-python/Cargo.lock index e8cea70..d07222a 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.2" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=9cde1d83e4114114bd49c9546d4fa5d86d6cbd32#9cde1d83e4114114bd49c9546d4fa5d86d6cbd32" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=fdff8e402e83f01bda3fbbf1f3f7f831b391b62b#fdff8e402e83f01bda3fbbf1f3f7f831b391b62b" dependencies = [ "aa-security", "async-trait", @@ -36,7 +36,7 @@ dependencies = [ [[package]] name = "aa-proto" version = "0.0.1-beta.2" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=9cde1d83e4114114bd49c9546d4fa5d86d6cbd32#9cde1d83e4114114bd49c9546d4fa5d86d6cbd32" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=fdff8e402e83f01bda3fbbf1f3f7f831b391b62b#fdff8e402e83f01bda3fbbf1f3f7f831b391b62b" dependencies = [ "prost", "tonic", @@ -47,7 +47,7 @@ dependencies = [ [[package]] name = "aa-sdk-client" version = "0.0.1-beta.2" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=9cde1d83e4114114bd49c9546d4fa5d86d6cbd32#9cde1d83e4114114bd49c9546d4fa5d86d6cbd32" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=fdff8e402e83f01bda3fbbf1f3f7f831b391b62b#fdff8e402e83f01bda3fbbf1f3f7f831b391b62b" dependencies = [ "aa-proto", "aa-security", @@ -64,7 +64,7 @@ dependencies = [ [[package]] name = "aa-security" version = "0.0.1-beta.2" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=9cde1d83e4114114bd49c9546d4fa5d86d6cbd32#9cde1d83e4114114bd49c9546d4fa5d86d6cbd32" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=fdff8e402e83f01bda3fbbf1f3f7f831b391b62b#fdff8e402e83f01bda3fbbf1f3f7f831b391b62b" dependencies = [ "aho-corasick", "serde", diff --git a/native/aa-ffi-python/Cargo.toml b/native/aa-ffi-python/Cargo.toml index 177e602..a67716d 100644 --- a/native/aa-ffi-python/Cargo.toml +++ b/native/aa-ffi-python/Cargo.toml @@ -13,9 +13,12 @@ crate-type = ["cdylib"] # Shared crates consumed via git-SHA pin (ADR 0002, AAASM-2559). aa-core, # aa-proto, and aa-sdk-client must share one SHA so cargo resolves a single # checkout of the agent-assembly workspace. -aa-core = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "9cde1d83e4114114bd49c9546d4fa5d86d6cbd32", package = "aa-core", features = ["serde"] } -aa-proto = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "9cde1d83e4114114bd49c9546d4fa5d86d6cbd32", package = "aa-proto" } -aa-sdk-client = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "9cde1d83e4114114bd49c9546d4fa5d86d6cbd32", package = "aa-sdk-client" } +# +# 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 = "fdff8e402e83f01bda3fbbf1f3f7f831b391b62b", package = "aa-core", features = ["serde"] } +aa-proto = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "fdff8e402e83f01bda3fbbf1f3f7f831b391b62b", package = "aa-proto" } +aa-sdk-client = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "fdff8e402e83f01bda3fbbf1f3f7f831b391b62b", 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 a5d8d02..e5ef983 100644 --- a/native/aa-ffi-python/src/lib.rs +++ b/native/aa-ffi-python/src/lib.rs @@ -83,8 +83,10 @@ impl RuntimeClient { agent_id: String::new(), socket_path: Some(socket_path.clone()), // Registration is a separate explicit `register` call; connecting - // the runtime UDS does not need a gateway endpoint. + // the runtime UDS does not need a gateway endpoint or lineage. gateway_endpoint: None, + team_id: None, + parent_agent_id: None, }; let handle = spawn_ipc_thread(config.resolve_socket_path()).map_err(|error| { PyRuntimeError::new_err(format!("failed to start runtime IPC thread: {error}")) @@ -108,10 +110,19 @@ impl RuntimeClient { /// pass `None` to let the shared client resolve it from the environment or /// its default. Returns the policy id the gateway assigns this agent. /// + /// `team_id` and `parent_agent_id` carry the agent's lineage/team scoping to + /// the gateway on the native register (AAASM-3415): `team_id` drives + /// team-budget attribution and `parent_agent_id` the topology graph. Both + /// are optional — `None` leaves the agent team-unscoped / root. + /// /// Raises `RuntimeError` when the gateway is unreachable or rejects the /// registration — unlike `query_policy`, registration is not advisory, so a /// failure surfaces rather than failing open. - #[pyo3(signature = (agent_id, name, framework, gateway_endpoint=None))] + // Each parameter is a distinct optional Python kwarg crossing the PyO3 + // boundary; bundling them into a struct would force callers to build one in + // Python, complicating the call site for no benefit. + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (agent_id, name, framework, gateway_endpoint=None, team_id=None, parent_agent_id=None))] fn register( &self, py: Python<'_>, @@ -119,11 +130,15 @@ impl RuntimeClient { name: String, framework: String, gateway_endpoint: Option, + team_id: Option, + parent_agent_id: Option, ) -> PyResult { let config = AssemblyConfig { agent_id, socket_path: Some(self.socket_path.clone()), gateway_endpoint, + team_id, + parent_agent_id, }; // `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/integration/test_topology_registration.py b/test/integration/test_topology_registration.py index 59fba03..6ed07e0 100644 --- a/test/integration/test_topology_registration.py +++ b/test/integration/test_topology_registration.py @@ -2,9 +2,9 @@ The REST ``register_agent`` topology-forwarding tests were retired in AAASM-3402 when registration moved to the native gRPC ``register`` path. The -native register call does not yet carry the lineage fields, so forwarding them -over the native path is a tracked follow-up; until then this only asserts the -fields are stored at construction. +native register call now carries ``team_id`` / ``parent_agent_id`` again +(AAASM-3415) — see ``test/unit/core/test_init_registration.py`` for the +forwarding assertion. This file asserts the fields are stored at construction. """ from __future__ import annotations diff --git a/test/unit/core/_fake_core.py b/test/unit/core/_fake_core.py index 4e1fe60..702b169 100644 --- a/test/unit/core/_fake_core.py +++ b/test/unit/core/_fake_core.py @@ -22,7 +22,7 @@ class FakeRuntimeClient: def __init__(self, decision: str = "allow", reason: str = "") -> None: self._decision = decision self._reason = reason - self.register_calls: list[tuple[str, str, str, 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 @@ -32,10 +32,12 @@ def register( name: str, framework: str, gateway_endpoint: str | None = None, + team_id: str | None = None, + parent_agent_id: str | None = None, ) -> str: if self.register_should_raise is not None: raise self.register_should_raise - self.register_calls.append((agent_id, name, framework, gateway_endpoint)) + self.register_calls.append((agent_id, name, framework, gateway_endpoint, team_id, parent_agent_id)) return "policy-id-001" def query_policy( @@ -52,16 +54,42 @@ def close(self) -> None: return None +class LegacyRuntimeClient: + """Stand-in for an older native build whose ``register`` predates the + ``team_id`` / ``parent_agent_id`` parameters (AAASM-3415). + + Its ``register`` only accepts the legacy positional signature, so calling it + with the lineage kwargs raises ``TypeError`` — exercising the SDK's + backwards-compatible fallback in ``register_agent``. + """ + + def __init__(self) -> None: + self.register_calls: list[tuple[str, str, str, str | None]] = [] + + def register( + self, + agent_id: str, + name: str, + framework: str, + gateway_endpoint: str | None = None, + ) -> str: + self.register_calls.append((agent_id, name, framework, gateway_endpoint)) + return "policy-id-legacy" + + def close(self) -> None: + return None + + def install_fake_core( monkeypatch: pytest.MonkeyPatch, - runtime_client: FakeRuntimeClient, -) -> FakeRuntimeClient: + runtime_client: Any, +) -> Any: """Install a fake ``agent_assembly._core`` whose ``RuntimeClient.connect`` returns ``runtime_client``. Returns the same client for assertions.""" class _ConnectingRuntimeClient: @staticmethod - def connect(_socket_path: str) -> FakeRuntimeClient: + def connect(_socket_path: str) -> Any: return runtime_client fake_core = types.ModuleType("agent_assembly._core") diff --git a/test/unit/core/test_init_registration.py b/test/unit/core/test_init_registration.py index b777c95..8eda321 100644 --- a/test/unit/core/test_init_registration.py +++ b/test/unit/core/test_init_registration.py @@ -15,9 +15,10 @@ from agent_assembly.adapters.base import FrameworkAdapter, GovernanceInterceptor from agent_assembly.core import assembly as core_assembly from agent_assembly.core.runtime_interceptor import build_governance_interceptor +from agent_assembly.core.spawn import SpawnContext, spawn_context_scope from agent_assembly.exceptions import ConfigurationError -from ._fake_core import FakeRuntimeClient, install_fake_core +from ._fake_core import FakeRuntimeClient, LegacyRuntimeClient, install_fake_core _GW_URL = "http://gateway.test" _API_KEY = "test-key" @@ -67,7 +68,200 @@ def test_init_assembly_registers_agent_on_init(monkeypatch: pytest.MonkeyPatch) context = init_assembly(gateway_url=_GW_URL, api_key=_API_KEY, agent_id="agent-7", mode="sdk-only") try: - assert runtime_client.register_calls == [("agent-7", "agent-7", "python", None)] + assert runtime_client.register_calls == [("agent-7", "agent-7", "python", None, None, None)] + finally: + context.shutdown() + + +def test_init_assembly_forwards_team_and_parent_on_register( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """init_assembly forwards team_id/parent_agent_id to the native register (AAASM-3415).""" + runtime_client = FakeRuntimeClient(decision="allow") + install_fake_core(monkeypatch, runtime_client) + _no_network(monkeypatch) + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **_kwargs: []) + + context = init_assembly( + gateway_url=_GW_URL, + api_key=_API_KEY, + agent_id="child-1", + mode="sdk-only", + team_id="team-payments", + parent_agent_id="parent-42", + ) + try: + assert runtime_client.register_calls == [("child-1", "child-1", "python", None, "team-payments", "parent-42")] + finally: + context.shutdown() + + +def test_init_assembly_forwards_only_team_when_parent_absent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Only ``team_id`` set → team forwarded, parent stays ``None`` (AAASM-3415).""" + runtime_client = FakeRuntimeClient(decision="allow") + install_fake_core(monkeypatch, runtime_client) + _no_network(monkeypatch) + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **_kwargs: []) + + context = init_assembly( + gateway_url=_GW_URL, + api_key=_API_KEY, + agent_id="team-only", + mode="sdk-only", + team_id="team-billing", + ) + try: + assert runtime_client.register_calls == [("team-only", "team-only", "python", None, "team-billing", None)] + finally: + context.shutdown() + + +def test_init_assembly_forwards_only_parent_when_team_absent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Only ``parent_agent_id`` set → parent forwarded, team stays ``None`` (AAASM-3415).""" + runtime_client = FakeRuntimeClient(decision="allow") + install_fake_core(monkeypatch, runtime_client) + _no_network(monkeypatch) + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **_kwargs: []) + + context = init_assembly( + gateway_url=_GW_URL, + api_key=_API_KEY, + agent_id="parent-only", + mode="sdk-only", + parent_agent_id="orchestrator-1", + ) + try: + assert runtime_client.register_calls == [("parent-only", "parent-only", "python", None, None, "orchestrator-1")] + finally: + context.shutdown() + + +def test_init_assembly_no_lineage_when_neither_set(monkeypatch: pytest.MonkeyPatch) -> None: + """Neither lineage field set → both forwarded as ``None``, no crash (AAASM-3415).""" + runtime_client = FakeRuntimeClient(decision="allow") + install_fake_core(monkeypatch, runtime_client) + _no_network(monkeypatch) + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **_kwargs: []) + + context = init_assembly( + gateway_url=_GW_URL, + api_key=_API_KEY, + agent_id="solo", + mode="sdk-only", + ) + try: + assert runtime_client.register_calls == [("solo", "solo", "python", None, None, None)] + finally: + context.shutdown() + + +def test_init_assembly_forwards_ambient_spawn_parent_on_register( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Ambient spawn lineage fills ``parent_agent_id`` when config omits it (AAASM-3415). + + A spawned child that does not pass ``parent_agent_id`` explicitly inherits it + from the ambient ``_SPAWN_CTX`` set at the spawn point, and that implicit + parent is forwarded to the native register. + """ + runtime_client = FakeRuntimeClient(decision="allow") + install_fake_core(monkeypatch, runtime_client) + _no_network(monkeypatch) + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **_kwargs: []) + + ctx = SpawnContext(parent_agent_id="ambient-parent", depth=1, spawned_by_tool="delegate") + with spawn_context_scope(ctx): + context = init_assembly( + gateway_url=_GW_URL, + api_key=_API_KEY, + agent_id="spawned-child", + mode="sdk-only", + ) + try: + assert runtime_client.register_calls == [ + ("spawned-child", "spawned-child", "python", None, None, "ambient-parent") + ] + finally: + context.shutdown() + + +def test_explicit_parent_overrides_ambient_spawn_parent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Explicit ``parent_agent_id`` wins over the ambient spawn parent (AAASM-3415).""" + runtime_client = FakeRuntimeClient(decision="allow") + install_fake_core(monkeypatch, runtime_client) + _no_network(monkeypatch) + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **_kwargs: []) + + ctx = SpawnContext(parent_agent_id="ambient-parent", depth=2, spawned_by_tool="delegate") + with spawn_context_scope(ctx): + context = init_assembly( + gateway_url=_GW_URL, + api_key=_API_KEY, + agent_id="override-child", + mode="sdk-only", + parent_agent_id="explicit-parent", + ) + try: + assert runtime_client.register_calls == [ + ("override-child", "override-child", "python", None, None, "explicit-parent") + ] + finally: + context.shutdown() + + +def test_register_falls_back_on_older_native_build_without_lineage_kwargs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An older native ``register`` (no lineage kwargs) is retried with the legacy + positional signature rather than crashing (AAASM-3415).""" + legacy_client = LegacyRuntimeClient() + install_fake_core(monkeypatch, legacy_client) + _no_network(monkeypatch) + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **_kwargs: []) + + context = init_assembly( + gateway_url=_GW_URL, + api_key=_API_KEY, + agent_id="legacy-agent", + mode="sdk-only", + team_id="team-x", + parent_agent_id="parent-y", + ) + try: + # Lineage is dropped against an old core, but registration still succeeds + # via the 4-arg legacy signature — no exception. + assert legacy_client.register_calls == [("legacy-agent", "legacy-agent", "python", None)] + finally: + context.shutdown() + + +def test_init_assembly_lineage_values_round_trip_verbatim( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Unicode / long lineage ids are forwarded to register without mangling.""" + runtime_client = FakeRuntimeClient(decision="allow") + install_fake_core(monkeypatch, runtime_client) + _no_network(monkeypatch) + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **_kwargs: []) + + team = "équipe-paiements-🌐" + parent = "parent-" + ("a" * 200) + context = init_assembly( + gateway_url=_GW_URL, + api_key=_API_KEY, + agent_id="unicode-child", + mode="sdk-only", + team_id=team, + parent_agent_id=parent, + ) + try: + assert runtime_client.register_calls == [("unicode-child", "unicode-child", "python", None, team, parent)] finally: context.shutdown()