From 69df0d214b73920fd0b6b4d7246041b28eae7fc0 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Fri, 19 Jun 2026 12:21:15 +0800 Subject: [PATCH 1/6] =?UTF-8?q?=E2=9C=A8=20(runtime-interceptor):=20Add=20?= =?UTF-8?q?native=20register=20+=20reusable=20client=20connect?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Expose connect_runtime_client() and register_agent() over the native RuntimeClient (AAASM-3399), and let build_governance_interceptor reuse a pre-connected client so register and query_policy share one client (and thus the issued credential token). native_available is threaded through so a missing extension (fail open) stays distinct from an unreachable socket (fail closed under enforce). Refs AAASM-3402 Co-Authored-By: Claude Opus 4.8 (1M context) --- agent_assembly/core/runtime_interceptor.py | 100 ++++++++++++++++++--- 1 file changed, 86 insertions(+), 14 deletions(-) diff --git a/agent_assembly/core/runtime_interceptor.py b/agent_assembly/core/runtime_interceptor.py index 19a98344..fe9a44a2 100644 --- a/agent_assembly/core/runtime_interceptor.py +++ b/agent_assembly/core/runtime_interceptor.py @@ -209,13 +209,74 @@ def check_tool_start(self, **_kwargs: Any) -> dict[str, str]: return {"status": "deny", "reason": self._reason} -def build_governance_interceptor(client: Any, agent_id: str, enforcement_mode: str | None = None) -> Any: +def connect_runtime_client(agent_id: str) -> Any | None: + """Connect a native ``RuntimeClient`` to the runtime UDS, or ``None``. + + Returns ``None`` when the native extension is not built or the runtime + socket is unreachable — both cases mean there is no native fast path to + register against or query. The single returned client is reused for both + :func:`register_agent` and the :class:`RuntimeQueryInterceptor`'s + ``query_policy`` calls, so the credential token stored by ``register`` is + attached to subsequent checks. + """ + try: + from agent_assembly._core import RuntimeClient + except ImportError: + return None + + socket_path = _resolve_runtime_socket_path(agent_id) + try: + return RuntimeClient.connect(socket_path) + except Exception: + return None + + +def register_agent( + runtime_client: Any, + agent_id: str, + framework: str, + gateway_endpoint: str | None = None, +) -> str | None: + """Register ``agent_id`` with the gateway over the native ``register`` call. + + Delegates to the native ``RuntimeClient.register`` (AAASM-3399), which makes + the SDK's only direct gateway gRPC call and stores the issued credential + token on the shared client so later ``query_policy`` checks authenticate + (ADR 0004 — the SDK never calls core HTTP endpoints directly). + + 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 + a misconfigured gateway rather than silently running unregistered. + """ + register = getattr(runtime_client, "register", None) + if register is None: + return None + return str(register(agent_id, agent_id, framework, gateway_endpoint)) + + +def _native_core_available() -> bool: + """Whether the native ``_core`` extension can be imported.""" + try: + from agent_assembly._core import RuntimeClient # noqa: F401 + except ImportError: + return False + return True + + +def build_governance_interceptor( + client: Any, + agent_id: str, + enforcement_mode: str | None = None, + *, + runtime_client: Any | None = None, + native_available: bool | None = None, +) -> Any: """Return the interceptor adapters should use for pre-execution checks. - When the native extension is importable and a runtime socket is reachable, - wrap ``client`` in a :class:`RuntimeQueryInterceptor` so a runtime ``deny`` - actually blocks the tool. The failure posture depends on ``enforcement_mode`` - (AAASM-3106): + When a native runtime is reachable, wrap ``client`` in a + :class:`RuntimeQueryInterceptor` so a runtime ``deny`` actually blocks the + tool. The failure posture depends on ``enforcement_mode`` (AAASM-3106): * The native extension is **missing**: return ``client`` unchanged in every mode. There is no native authority to consult, so there is nothing to fail @@ -223,18 +284,29 @@ def build_governance_interceptor(client: Any, agent_id: str, enforcement_mode: s * The native extension is **present** but the runtime socket is unreachable: under ``enforce`` return a deny-all :class:`_FailClosedInterceptor`; otherwise return ``client`` unchanged (fail open). + + :param runtime_client: A pre-connected native runtime client (e.g. the one + :func:`register_agent` registered the token on). When supplied it is + reused so register and ``query_policy`` share one client; when ``None`` + a fresh connection is established here. + :param native_available: Whether the native extension is importable. Pass + this alongside a ``runtime_client`` of ``None`` to distinguish a missing + extension (fail open in every mode) from an unreachable runtime socket + (fail closed under enforce). When ``None`` it is detected here. """ enforce = enforcement_mode == ENFORCE_MODE - try: - from agent_assembly._core import RuntimeClient - except ImportError: - # No native fast path at all — the SDK control is not engaged in any mode. - return client - socket_path = _resolve_runtime_socket_path(agent_id) - try: - runtime_client = RuntimeClient.connect(socket_path) - except Exception: + if runtime_client is None and native_available is None: + # No caller-supplied client or hint: detect + connect ourselves. + if not _native_core_available(): + return client + runtime_client = connect_runtime_client(agent_id) + native_available = True + + if runtime_client is None: + # Native missing → no authority to fail closed to, return bare client. + if not native_available: + return client # Native present but runtime unreachable: deny everything under enforce # (fail closed); proceed under observe / disabled (fail open). if enforce: From d14a3f658f63986747964149da5992c565461980 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Fri, 19 Jun 2026 12:23:11 +0800 Subject: [PATCH 2/6] =?UTF-8?q?=E2=9C=A8=20(init):=20Register=20agent=20vi?= =?UTF-8?q?a=20native=20gRPC=20on=20init=5Fassembly?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit init_assembly now connects one native runtime client, registers the agent on it (storing the credential token), and reuses that client for the RuntimeQueryInterceptor — so a deny actually blocks a tool. Registration failure propagates under enforce (fail closed) and is swallowed under observe/disabled (fail open), honoring the existing enforcement semantics. Refs AAASM-3402 Co-Authored-By: Claude Opus 4.8 (1M context) --- agent_assembly/core/assembly.py | 82 ++++++++++++++++++++++++++++----- 1 file changed, 70 insertions(+), 12 deletions(-) diff --git a/agent_assembly/core/assembly.py b/agent_assembly/core/assembly.py index 5f802254..c5e5faca 100644 --- a/agent_assembly/core/assembly.py +++ b/agent_assembly/core/assembly.py @@ -14,7 +14,12 @@ from agent_assembly.adapters.registry import AdapterRegistry from agent_assembly.client.gateway import GatewayClient from agent_assembly.core.gateway_resolver import resolve_api_key, resolve_gateway_url -from agent_assembly.core.runtime_interceptor import build_governance_interceptor +from agent_assembly.core.runtime_interceptor import ( + _native_core_available, + build_governance_interceptor, + connect_runtime_client, + register_agent, +) from agent_assembly.core.spawn import _SPAWN_CTX from agent_assembly.exceptions import AssemblyError, ConfigurationError @@ -145,17 +150,24 @@ def init_assembly( optional auto-start) per Epic 17 S-G — see ``agent_assembly.core.gateway_resolver``. + Agent registration and the pre-execution policy check go through the native + ``aa-sdk-client`` shim to the core over gRPC/UDS (ADR 0004): ``init_assembly`` + registers the agent on startup (storing the issued credential token on the + native client) and a tool call is checked via the native ``query_policy`` so + a ``deny`` blocks the tool before it runs. The SDK never calls a core HTTP + endpoint directly for registration or policy checks. + :param control_plane_url: Optional URL of the control-plane HTTP API. When - supplied, the SDK issues its HTTP routes (agent registration, policy - checks, topology edges) against it instead of ``gateway_url``. When - omitted it falls back to ``gateway_url`` — the backwards-compatible - single-host OSS dev setup. Resolution order: explicit kwarg > - ``AA_CONTROL_PLANE_URL`` env-var > unset (falls back to ``gateway_url``). - :param enforcement_mode: Per-agent governance posture sent to the gateway - at registration (see :data:`EnforcementMode`). Defaults to ``None``, - which omits the field from the registration body — the gateway then - applies its server-side default (live ``enforce``). Pass ``"observe"`` - to register the agent in dry-run / sandbox mode: every action + supplied, the SDK issues its remaining HTTP routes (topology edges, + secret dispatch) against it instead of ``gateway_url``. When omitted it + falls back to ``gateway_url`` — the backwards-compatible single-host OSS + dev setup. Resolution order: explicit kwarg > ``AA_CONTROL_PLANE_URL`` + env-var > unset (falls back to ``gateway_url``). + :param enforcement_mode: Per-agent governance posture applied to this + agent's actions (see :data:`EnforcementMode`). Defaults to ``None``, + which lets the gateway apply its server-side default (live ``enforce``). + Pass ``"observe"`` to register the agent in dry-run / sandbox mode: + every action proceeds and the gateway records would-be violations as shadow audit events. """ @@ -210,10 +222,19 @@ def init_assembly( network_mode: NetworkMode = "sdk-only" network_shutdown: Callable[[], None] = _noop_shutdown try: + native_available = _native_core_available() + runtime_client = connect_runtime_client(resolved_agent_id) if native_available else None + _register_agent_with_gateway( + runtime_client=runtime_client, + agent_id=resolved_agent_id, + enforcement_mode=enforcement_mode, + ) registered_adapters = _register_adapters( client=client, process_agent_id=resolved_agent_id, enforcement_mode=enforcement_mode, + runtime_client=runtime_client, + native_available=native_available, ) network_mode, network_shutdown = _start_network_layer(client=client, mode=mode) except Exception as error: @@ -261,10 +282,41 @@ def _validate_inputs( return gateway_url, control_plane_url +def _register_agent_with_gateway( + *, + runtime_client: Any | None, + agent_id: str, + enforcement_mode: EnforcementMode | None, +) -> None: + """Register the agent with the gateway over the native gRPC ``register``. + + Registration goes through the native runtime client (AAASM-3399) so the + issued credential token is stored on the same client the + ``RuntimeQueryInterceptor`` later uses for ``query_policy`` — the SDK never + calls a core HTTP endpoint directly (ADR 0004). + + 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; + under ``observe`` / ``disabled`` (or no mode) it is swallowed so the dry-run + / hermetic-test layer never hard-fails on registration. + """ + if runtime_client is None: + return + framework = "python" + try: + register_agent(runtime_client, agent_id, framework) + except Exception: + if enforcement_mode == "enforce": + raise + + def _register_adapters( client: GatewayClient, process_agent_id: str, enforcement_mode: EnforcementMode | None = None, + runtime_client: Any | None = None, + native_available: bool = False, ) -> list[FrameworkAdapter]: """Detect available frameworks via AdapterRegistry and register hooks. @@ -282,7 +334,13 @@ def _register_adapters( adapters = registry.get_available_adapters_by_priority() registered: list[FrameworkAdapter] = [] - interceptor: Any = build_governance_interceptor(client, process_agent_id, enforcement_mode) + interceptor: Any = build_governance_interceptor( + client, + process_agent_id, + enforcement_mode, + runtime_client=runtime_client, + native_available=native_available, + ) for adapter in adapters: adapter.set_process_agent_id(process_agent_id) From e2991b46e33ed1e029cadf144f264bb57d0b210a Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Fri, 19 Jun 2026 12:23:56 +0800 Subject: [PATCH 3/6] =?UTF-8?q?=F0=9F=97=91=EF=B8=8F=20(gateway-client):?= =?UTF-8?q?=20Retire=20legacy=20REST=20register/check=20and=20their=20test?= =?UTF-8?q?s?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit register_agent() and check_policy_compliance() issued raw HTTP to the core, which ADR 0004 forbids — registration and the policy check now go through the native gRPC path. Removes the methods plus the unit/integration tests that asserted their REST wire shape; report_edge()/dispatch_tool() (genuinely HTTP) and the GatewayClient topology-field storage tests are kept. Topology lineage over the native register call is a tracked follow-up. Refs AAASM-3402 Co-Authored-By: Claude Opus 4.8 (1M context) --- agent_assembly/client/gateway.py | 73 +-------- .../integration/test_topology_registration.py | 117 +-------------- test/unit/client/test_gateway_endpoints.py | 35 +---- test/unit/client/test_gateway_topology.py | 142 +----------------- 4 files changed, 24 insertions(+), 343 deletions(-) diff --git a/agent_assembly/client/gateway.py b/agent_assembly/client/gateway.py index ee79d96c..a2b7a015 100644 --- a/agent_assembly/client/gateway.py +++ b/agent_assembly/client/gateway.py @@ -37,19 +37,20 @@ def __init__( api_key: Optional API key for authentication timeout: Request timeout in seconds control_plane_url: Optional URL of the control-plane HTTP API. When - set, the HTTP routes (``/agents/{id}/register``, ``/policy/check``, - ``/topology/edges``) are issued against it; when ``None`` (the + set, the remaining HTTP routes (``/topology/edges``, + ``/dispatch_tool``) are issued against it; when ``None`` (the default) they fall back to ``gateway_url`` — the backwards-compatible - single-host OSS dev setup. + single-host OSS dev setup. Agent registration and policy checks + go through the native gRPC path (ADR 0004), not these HTTP routes. parent_agent_id: Parent agent ID for topology tracking team_id: Team ID this agent belongs to delegation_reason: Human-readable reason for delegation spawned_by_tool: Name of the tool that spawned this agent depth: Spawn depth in the agent lineage tree - enforcement_mode: Per-agent governance posture sent to the gateway - at registration. ``None`` (the default) omits the field from - the request body so a legacy gateway sees the same wire shape - as before; the gateway then defaults to live enforcement. + enforcement_mode: Per-agent governance posture for this agent. + ``None`` (the default) lets the gateway apply its server-side + default of live enforcement. Stored for reference; registration + itself now goes through the native gRPC path (ADR 0004). """ self.gateway_url = gateway_url.rstrip("/") self.control_plane_url = control_plane_url.rstrip("/") if control_plane_url else None @@ -96,64 +97,6 @@ def __exit__(self, *args: object) -> None: """Context manager exit.""" self.close() - async def register_agent(self) -> dict[str, Any]: - """ - Register the agent with the governance gateway. - - Returns: - Registration response data - - Raises: - GatewayError: If registration fails - """ - body: dict[str, str | int] = {} - if self.parent_agent_id is not None: - body["parent_agent_id"] = self.parent_agent_id - if self.team_id is not None: - body["team_id"] = self.team_id - if self.delegation_reason is not None: - body["delegation_reason"] = self.delegation_reason - if self.spawned_by_tool is not None: - body["spawned_by_tool"] = self.spawned_by_tool - if self.depth is not None: - body["depth"] = self.depth - if self.enforcement_mode is not None: - body["enforcement_mode"] = self.enforcement_mode - try: - response = self.client.post( - f"/agents/{self.agent_id}/register", - json=body if body else None, - ) - response.raise_for_status() - data: dict[str, Any] = response.json() - return data - except httpx.HTTPError as e: - raise GatewayError(f"Failed to register agent: {e}") from e - - async def check_policy_compliance(self, action: str) -> dict[str, Any]: - """ - Check if an action complies with governance policies. - - Args: - action: The action to check - - Returns: - Policy compliance response - - Raises: - GatewayError: If policy check fails - """ - try: - response = self.client.post( - f"/agents/{self.agent_id}/policy/check", - json={"action": action}, - ) - response.raise_for_status() - data: dict[str, Any] = response.json() - return data - except httpx.HTTPError as e: - raise GatewayError(f"Failed to check policy compliance: {e}") from e - def report_edge( self, source_agent_id: str, diff --git a/test/integration/test_topology_registration.py b/test/integration/test_topology_registration.py index 8006f2f6..59fba033 100644 --- a/test/integration/test_topology_registration.py +++ b/test/integration/test_topology_registration.py @@ -1,122 +1,17 @@ -"""Integration tests for topology field forwarding in the register_agent call (AAASM-1178). +"""Integration smoke test for GatewayClient topology field storage (AAASM-1178). -Uses httpx.MockTransport to intercept HTTP traffic so tests run without a live gateway. +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. """ from __future__ import annotations -import httpx -import pytest - -from agent_assembly import init_assembly from agent_assembly.client.gateway import GatewayClient -def _make_transport(captured: list[httpx.Request]) -> httpx.MockTransport: - def handler(request: httpx.Request) -> httpx.Response: - captured.append(request) - return httpx.Response(200, json={"registration_id": "reg-test-001"}) - - return httpx.MockTransport(handler) - - -@pytest.mark.asyncio -async def test_register_agent_topology_fields_reach_gateway() -> None: - """init_assembly + register_agent sends all 4 topology fields to the gateway.""" - captured: list[httpx.Request] = [] - transport = _make_transport(captured) - - context = init_assembly( - gateway_url="http://gateway.test", - api_key="test-api-key", - agent_id="agent-child-001", - mode="sdk-only", - parent_agent_id="agent-parent-001", - team_id="team-alpha", - delegation_reason="sub-task delegation", - spawned_by_tool="search_tool", - ) - - context.client._client = httpx.Client( - base_url="http://gateway.test", - transport=transport, - ) - - await context.client.register_agent() - context.shutdown() - - assert len(captured) == 1, f"expected 1 request, got {len(captured)}" - body = captured[0].read() - import json - - payload = json.loads(body) - - assert payload["parent_agent_id"] == "agent-parent-001" - assert payload["team_id"] == "team-alpha" - assert payload["delegation_reason"] == "sub-task delegation" - assert payload["spawned_by_tool"] == "search_tool" - - -@pytest.mark.asyncio -async def test_register_agent_no_topology_omits_body() -> None: - """When no topology params are provided, register_agent sends an empty/null body.""" - captured: list[httpx.Request] = [] - transport = _make_transport(captured) - - context = init_assembly( - gateway_url="http://gateway.test", - api_key="test-api-key", - agent_id="agent-001", - mode="sdk-only", - ) - - context.client._client = httpx.Client( - base_url="http://gateway.test", - transport=transport, - ) - - await context.client.register_agent() - context.shutdown() - - assert len(captured) == 1 - body = captured[0].read() - # No topology fields → httpx sends null body (no Content-Type: application/json) - assert body == b"" or body == b"null", f"expected empty body when no topology, got {body!r}" - - -@pytest.mark.asyncio -async def test_register_agent_partial_topology_only_sends_set_fields() -> None: - """When only some topology fields are set, only those are included in the body.""" - captured: list[httpx.Request] = [] - transport = _make_transport(captured) - - context = init_assembly( - gateway_url="http://gateway.test", - api_key="test-api-key", - agent_id="agent-partial-001", - mode="sdk-only", - team_id="team-beta", - ) - - context.client._client = httpx.Client( - base_url="http://gateway.test", - transport=transport, - ) - - await context.client.register_agent() - context.shutdown() - - import json - - assert len(captured) == 1 - payload = json.loads(captured[0].read()) - - assert payload.get("team_id") == "team-beta" - assert "parent_agent_id" not in payload - assert "delegation_reason" not in payload - assert "spawned_by_tool" not in payload - - def test_gateway_client_stores_topology_fields_at_construction() -> None: """GatewayClient stores all 4 topology fields on construction (integration smoke).""" client = GatewayClient( diff --git a/test/unit/client/test_gateway_endpoints.py b/test/unit/client/test_gateway_endpoints.py index 9902162e..9bd4ad7b 100644 --- a/test/unit/client/test_gateway_endpoints.py +++ b/test/unit/client/test_gateway_endpoints.py @@ -1,8 +1,9 @@ """Unit tests for the remaining `GatewayClient` HTTP endpoints. -Covers `check_policy_compliance`, `report_edge`, the API-key auth header, and -the `register_agent` failure branch — the success/error paths not already -exercised by the topology and dispatch_tool suites. +Covers `report_edge` and the API-key auth header — the success/error paths not +already exercised by the dispatch_tool suite. The REST ``register_agent`` / +``check_policy_compliance`` methods were retired in AAASM-3402 in favor of the +native gRPC register / query_policy path, so they are no longer tested here. """ from __future__ import annotations @@ -54,34 +55,6 @@ def test_http_client_omits_auth_header_when_no_api_key() -> None: client.close() -@pytest.mark.asyncio -async def test_register_agent_raises_gateway_error_on_http_error() -> None: - client = GatewayClient(gateway_url="http://gw.test", agent_id="a", api_key="k") - mock_post = MagicMock(return_value=_raising(httpx.ConnectError("refused"))) - with _patch_post(client, mock_post), pytest.raises(GatewayError, match="Failed to register agent"): - await client.register_agent() - - -@pytest.mark.asyncio -async def test_check_policy_compliance_returns_decision_on_success() -> None: - client = GatewayClient(gateway_url="http://gw.test", agent_id="a", api_key="k") - mock_post = MagicMock(return_value=_ok({"allowed": True, "reason": "ok"})) - with _patch_post(client, mock_post): - result = await client.check_policy_compliance("send_email") - - assert result == {"allowed": True, "reason": "ok"} - _, kwargs = mock_post.call_args - assert kwargs["json"] == {"action": "send_email"} - - -@pytest.mark.asyncio -async def test_check_policy_compliance_raises_gateway_error_on_http_error() -> None: - client = GatewayClient(gateway_url="http://gw.test", agent_id="a", api_key="k") - mock_post = MagicMock(return_value=_raising(httpx.ReadTimeout("slow"))) - with _patch_post(client, mock_post), pytest.raises(GatewayError, match="Failed to check policy compliance"): - await client.check_policy_compliance("send_email") - - def test_report_edge_serializes_metadata_and_returns_edge_id() -> None: client = GatewayClient(gateway_url="http://gw.test", agent_id="a", api_key="k") mock_post = MagicMock(return_value=_ok({"edge_id": "e-9"})) diff --git a/test/unit/client/test_gateway_topology.py b/test/unit/client/test_gateway_topology.py index 837edaa1..1bdd8eac 100644 --- a/test/unit/client/test_gateway_topology.py +++ b/test/unit/client/test_gateway_topology.py @@ -1,21 +1,15 @@ -"""Unit tests for GatewayClient topology param forwarding (AAASM-958).""" +"""Unit tests for GatewayClient topology param storage (AAASM-958). -from __future__ import annotations - -from unittest.mock import MagicMock, patch +The REST ``register_agent`` wire-shape tests were retired in AAASM-3402 when +registration moved to the native gRPC path; only the constructor-storage +contract of the topology fields is exercised here. +""" -import pytest +from __future__ import annotations from agent_assembly.client.gateway import GatewayClient -def _make_ok_response() -> MagicMock: - resp = MagicMock() - resp.json.return_value = {"registration_id": "reg-1"} - resp.raise_for_status = MagicMock() - return resp - - def test_gateway_client_stores_topology_fields() -> None: client = GatewayClient( gateway_url="http://gw.test", @@ -38,127 +32,3 @@ def test_gateway_client_topology_defaults_to_none() -> None: assert client.team_id is None assert client.delegation_reason is None assert client.spawned_by_tool is None - - -@pytest.mark.asyncio -async def test_register_agent_sends_topology_fields() -> None: - client = GatewayClient( - gateway_url="http://gw.test", - agent_id="agent-child", - api_key="key", - parent_agent_id="agent-parent", - team_id="team-alpha", - delegation_reason="sub-task delegation", - spawned_by_tool="search_tool", - ) - mock_post = MagicMock(return_value=_make_ok_response()) - with patch.object( - type(client), - "client", - new_callable=lambda: property(lambda self: MagicMock(post=mock_post)), - ): - await client.register_agent() - - _, kwargs = mock_post.call_args - body = kwargs.get("json") or {} - assert body["parent_agent_id"] == "agent-parent" - assert body["team_id"] == "team-alpha" - assert body["delegation_reason"] == "sub-task delegation" - assert body["spawned_by_tool"] == "search_tool" - - -@pytest.mark.asyncio -async def test_register_agent_omits_body_when_no_topology() -> None: - client = GatewayClient( - gateway_url="http://gw.test", - agent_id="agent-1", - api_key="key", - ) - mock_post = MagicMock(return_value=_make_ok_response()) - with patch.object( - type(client), - "client", - new_callable=lambda: property(lambda self: MagicMock(post=mock_post)), - ): - await client.register_agent() - - _, kwargs = mock_post.call_args - assert kwargs.get("json") is None - - -@pytest.mark.asyncio -async def test_register_agent_includes_depth_when_set() -> None: - client = GatewayClient( - gateway_url="http://gw.test", - agent_id="child", - api_key="key", - parent_agent_id="parent", - depth=3, - ) - mock_post = MagicMock(return_value=_make_ok_response()) - with patch.object( - type(client), - "client", - new_callable=lambda: property(lambda self: MagicMock(post=mock_post)), - ): - await client.register_agent() - - _, call_kwargs = mock_post.call_args - body = call_kwargs.get("json") or {} - assert body["depth"] == 3 - - -# ── enforcement_mode wire-shape (AAASM-1560) ───────────────────────────────── - - -@pytest.mark.asyncio -async def test_register_agent_sends_enforcement_mode_when_set() -> None: - """Registering an agent with enforcement_mode=observe puts the snake_case - string on the wire so the gateway's REST → gRPC bridge can map it to - RegisterRequest.enforcement_mode (proto enum) per AAASM-1555/1557. - """ - client = GatewayClient( - gateway_url="http://gw.test", - agent_id="experimental-agent", - api_key="key", - enforcement_mode="observe", - ) - mock_post = MagicMock(return_value=_make_ok_response()) - with patch.object( - type(client), - "client", - new_callable=lambda: property(lambda self: MagicMock(post=mock_post)), - ): - await client.register_agent() - - _, kwargs = mock_post.call_args - body = kwargs.get("json") or {} - assert body["enforcement_mode"] == "observe" - - -@pytest.mark.asyncio -async def test_register_agent_omits_enforcement_mode_when_none() -> None: - """A client constructed without enforcement_mode (the pre-feature path) - must NOT include the key in the body — keeps the wire shape identical - to before so a legacy gateway that doesn't know about the field still - accepts the registration cleanly. - """ - client = GatewayClient( - gateway_url="http://gw.test", - agent_id="legacy-agent", - api_key="key", - # no enforcement_mode kwarg - ) - mock_post = MagicMock(return_value=_make_ok_response()) - with patch.object( - type(client), - "client", - new_callable=lambda: property(lambda self: MagicMock(post=mock_post)), - ): - await client.register_agent() - - _, kwargs = mock_post.call_args - body = kwargs.get("json") - # body may be None (no fields set) or a dict that omits the key. - if body is not None: - assert "enforcement_mode" not in body From e493b0b585fd1946a134f266c05778544eb4fff9 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Fri, 19 Jun 2026 12:24:09 +0800 Subject: [PATCH 4/6] =?UTF-8?q?=F0=9F=93=9D=20(readme):=20Correct=20regist?= =?UTF-8?q?er/check=20to=20native=20gRPC=20path,=20not=20HTTP?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Public API section advertised the removed REST register_agent/ check_policy_compliance and claimed registration + policy checks go over HTTP. Restate per ADR 0004: those go SDK → aa-sdk-client → core over gRPC/UDS; only topology edges and secret dispatch remain HTTP routes. Refs AAASM-3402 Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 368080de..1298ccb5 100644 --- a/README.md +++ b/README.md @@ -126,15 +126,19 @@ What this does: ## Public API - `init_assembly(gateway_url, api_key, agent_id=None, mode="auto", *, control_plane_url=None) -> AssemblyContext` -- `async GatewayClient.register_agent() -> dict` -- `async GatewayClient.check_policy_compliance(action: str) -> dict` - Exceptions: `AssemblyError`, `AgentError`, `PolicyError`, `GatewayError`, `ConfigurationError` - Data models: `AgentConfig`, `AgentState`, `PolicyEvaluation` +Agent **registration** and the **pre-execution policy check** are not REST +calls: the SDK goes through the native `aa-sdk-client` shim to the core over +gRPC/UDS (ADR 0004) — it never calls a core HTTP endpoint directly for those. +`init_assembly` registers the agent on startup, and a tool call is checked via +the native `query_policy` so a `deny` blocks it before the tool runs. + ### Control-plane routing -By default the SDK issues its HTTP routes (agent registration, policy checks, -topology edges) against `gateway_url` — the single-host OSS dev setup. Pass +By default the SDK issues its remaining HTTP routes (topology edges, secret +dispatch) against `gateway_url` — the single-host OSS dev setup. Pass `control_plane_url` to route those HTTP calls to a separate control-plane host while `gateway_url` continues to serve the gRPC data path: From fb989beb89cdf1f1ce3c9cabbe131acc048526c2 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Fri, 19 Jun 2026 12:24:10 +0800 Subject: [PATCH 5/6] =?UTF-8?q?=E2=9C=85=20(test):=20Add=20native=20=5Fcor?= =?UTF-8?q?e=20fakes=20for=20SDK=20wiring=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit FakeRuntimeClient records register calls and returns a canned query_policy decision; install_fake_core swaps a fake agent_assembly._core into sys.modules so init_assembly exercises the native register/query path without a built extension or a running aa-runtime. Refs AAASM-3402 Co-Authored-By: Claude Opus 4.8 (1M context) --- test/unit/core/_fake_core.py | 70 ++++++++++++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) create mode 100644 test/unit/core/_fake_core.py diff --git a/test/unit/core/_fake_core.py b/test/unit/core/_fake_core.py new file mode 100644 index 00000000..4e1fe60c --- /dev/null +++ b/test/unit/core/_fake_core.py @@ -0,0 +1,70 @@ +"""Reusable fakes for the native ``agent_assembly._core`` extension. + +These let the SDK wiring tests (AAASM-3402) exercise the native register / +query_policy path without a built extension or a running ``aa-runtime``: the +``FakeRuntimeClient`` records ``register`` calls and returns a canned +``query_policy`` decision, and ``install_fake_core`` swaps a module exposing a +``RuntimeClient`` whose ``connect`` yields one into ``sys.modules``. +""" + +from __future__ import annotations + +import sys +import types +from typing import Any + +import pytest + + +class FakeRuntimeClient: + """Stand-in for the native ``RuntimeClient`` (register + query_policy).""" + + 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.query_calls: list[tuple[Any, ...]] = [] + self.register_should_raise: Exception | None = None + + def register( + self, + agent_id: str, + name: str, + framework: str, + gateway_endpoint: 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)) + return "policy-id-001" + + def query_policy( + self, + agent_id: str, + action_type: str, + tool_name: str | None = None, + tool_args_json: str | None = None, + ) -> dict[str, str]: + self.query_calls.append((agent_id, action_type, tool_name, tool_args_json)) + return {"decision": self._decision, "reason": self._reason} + + def close(self) -> None: + return None + + +def install_fake_core( + monkeypatch: pytest.MonkeyPatch, + runtime_client: FakeRuntimeClient, +) -> FakeRuntimeClient: + """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: + 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 From be29d4f2b16ddb8e3cce07f2a1b9070b5b0a5559 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Fri, 19 Jun 2026 12:24:24 +0800 Subject: [PATCH 6/6] =?UTF-8?q?=E2=9C=85=20(test):=20Add=20register-on-ini?= =?UTF-8?q?t=20and=20deny-blocks=20regression=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the unwired-enforcement gap with four cases: init_assembly calls native register; a native deny makes the adapter interceptor's check_tool_start block; observe swallows a registration failure; enforce propagates it (fail closed). Refs AAASM-3402 Co-Authored-By: Claude Opus 4.8 (1M context) --- test/unit/core/test_init_registration.py | 163 +++++++++++++++++++++++ 1 file changed, 163 insertions(+) create mode 100644 test/unit/core/test_init_registration.py diff --git a/test/unit/core/test_init_registration.py b/test/unit/core/test_init_registration.py new file mode 100644 index 00000000..b777c957 --- /dev/null +++ b/test/unit/core/test_init_registration.py @@ -0,0 +1,163 @@ +"""Tests for the native register / pre-execution check wiring (AAASM-3402). + +Closes the unwired-enforcement gap: ``init_assembly`` must register the agent +over the native gRPC path on startup, and the interceptor it hands the adapters +must block a tool call when the native ``query_policy`` returns ``deny``. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from agent_assembly import init_assembly +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.exceptions import ConfigurationError + +from ._fake_core import FakeRuntimeClient, install_fake_core + +_GW_URL = "http://gateway.test" +_API_KEY = "test-key" + + +class _CapturingAdapter(FrameworkAdapter): + """Adapter that records the interceptor it is handed at register_hooks.""" + + def __init__(self) -> None: + self.interceptor: GovernanceInterceptor | None = None + + def get_framework_name(self) -> str: + return "capturing" + + def get_supported_versions(self) -> list[str]: + return [">=0.0.0"] + + def register_hooks(self, interceptor: GovernanceInterceptor) -> None: + self.interceptor = interceptor + + def unregister_hooks(self) -> None: + return None + + +@pytest.fixture(autouse=True) +def _cleanup_active_context() -> None: + active = core_assembly._ACTIVE_CONTEXT + if active is not None and not active.is_shutdown: + active.shutdown() + core_assembly._ACTIVE_CONTEXT = None + + +def _no_network(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + core_assembly, + "_start_network_layer", + lambda **_kwargs: ("sdk-only", core_assembly._noop_shutdown), + ) + + +def test_init_assembly_registers_agent_on_init(monkeypatch: pytest.MonkeyPatch) -> None: + """init_assembly calls the native register so the gateway knows the agent.""" + 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="agent-7", mode="sdk-only") + try: + assert runtime_client.register_calls == [("agent-7", "agent-7", "python", None)] + finally: + context.shutdown() + + +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") + install_fake_core(monkeypatch, runtime_client) + _no_network(monkeypatch) + + adapter = _CapturingAdapter() + monkeypatch.setattr(core_assembly, "_register_adapters", _patched_register_adapters(adapter)) + + context = init_assembly( + gateway_url=_GW_URL, + api_key=_API_KEY, + agent_id="agent-deny", + mode="sdk-only", + enforcement_mode="enforce", + ) + try: + assert adapter.interceptor is not None + interceptor: Any = adapter.interceptor + result = interceptor.check_tool_start( + serialized={"name": "web_search"}, + input_str="q", + tool_name="web_search", + args={"q": "x"}, + ) + assert result == {"status": "deny", "reason": "policy violation"} + finally: + context.shutdown() + + +def test_observe_mode_swallows_register_failure(monkeypatch: pytest.MonkeyPatch) -> None: + """Under observe a native registration failure does not abort init.""" + runtime_client = FakeRuntimeClient(decision="allow") + runtime_client.register_should_raise = RuntimeError("gateway down") + 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="agent-observe", + mode="sdk-only", + enforcement_mode="observe", + ) + context.shutdown() + + +def test_enforce_mode_propagates_register_failure(monkeypatch: pytest.MonkeyPatch) -> None: + """Under enforce a native registration failure aborts init (fail closed).""" + runtime_client = FakeRuntimeClient(decision="allow") + runtime_client.register_should_raise = RuntimeError("gateway rejected") + install_fake_core(monkeypatch, runtime_client) + _no_network(monkeypatch) + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **_kwargs: []) + + with pytest.raises(ConfigurationError, match="Failed to initialize assembly runtime"): + init_assembly( + gateway_url=_GW_URL, + api_key=_API_KEY, + agent_id="agent-enforce", + mode="sdk-only", + enforcement_mode="enforce", + ) + + +def _patched_register_adapters(adapter: _CapturingAdapter) -> object: + """Build a stand-in for ``_register_adapters`` that drives the real + interceptor builder and registers the capturing adapter with it.""" + + def _impl( + *, + client: object, + process_agent_id: str, + enforcement_mode: str | None = None, + runtime_client: object | None = None, + native_available: bool = False, + ) -> list[FrameworkAdapter]: + interceptor = build_governance_interceptor( + client, + process_agent_id, + enforcement_mode, + runtime_client=runtime_client, + native_available=native_available, + ) + adapter.register_hooks(interceptor) + return [adapter] + + return _impl