[AAASM-3415] 🔧 (native): Forward team_id/parent_agent_id on native register

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

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:

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"] }

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,22 +110,35 @@ 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<'_>,
         agent_id: String,
         name: String,
         framework: String,
         gateway_endpoint: Option<String>,
+        team_id: Option<String>,
+        parent_agent_id: Option<String>,
     ) -> PyResult<String> {
         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

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

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")