[AAASM-6444] ✨ (init): Add org_id and registration_key to init_assembly

agent_assembly/core/assembly.py

@@ -186,19 +186,21 @@
 
 
 def init_assembly(
     gateway_url: str | None = None,
     api_key: str | None = None,
     agent_id: str | None = None,
     mode: RuntimeMode = "auto",
     *,
     control_plane_url: str | None = None,
     parent_agent_id: str | None = None,
     team_id: str | None = None,
+    org_id: str | None = None,
+    registration_key: str | None = None,
     delegation_reason: str | None = None,
     spawned_by_tool: str | None = None,
     depth: int | None = None,
     enforcement_mode: EnforcementMode | None = None,
     allow_insecure: bool = False,
 ) -> AssemblyContext:
     """Initialize the Agent Assembly SDK runtime for this process.
 
@@ -245,6 +247,16 @@
         the agent in dry-run / sandbox mode: every action proceeds (local checks
         fail open) and the gateway records would-be violations as shadow audit
         events.
+    :param org_id: Organization this agent claims to belong to. It is a claim,
+        not an authorization (AAASM-6444): the gateway admits it only when the
+        ``registration_key`` presented with it is bound to the same
+        organization, and rejects a mismatch or an unauthenticated claim. Omit
+        it to take the organization from the key's binding. Audit entries and
+        ``GET /api/v1/logs?org_id=`` then carry that organization.
+    :param registration_key: Operator API key (``aa_...``) the gateway binds to
+        an organization. Never sent on any other route; falls back to the
+        ``AA_REGISTRATION_KEY`` environment variable on the native side. Distinct
+        from ``api_key``, which authenticates the SDK's HTTP calls.
     :param allow_insecure: Opt into a plaintext (non-TLS) native ``register``
         channel when the resolved gateway host is non-loopback (AAASM-4664).
         Defaults to ``False`` — secure-by-default (AAASM-4655): a derived
@@ -281,6 +293,11 @@
             spawned_by_tool = _spawn.spawned_by_tool
 
     resolved_agent_id = _validate_agent_id(agent_id or _DEFAULT_AGENT_ID)
+    # A blank value is "not given", never an empty claim or an empty key: the native client would
+    # otherwise send an empty header that overrides AA_REGISTRATION_KEY and the gateway would
+    # refuse a registration the caller meant to leave to the environment.
+    org_id = (org_id or "").strip() or None
+    registration_key = (registration_key or "").strip() or None
 
     global _ACTIVE_CONTEXT
     with _INIT_LOCK:
@@ -323,6 +340,8 @@
                 native_available=native_available,
                 team_id=team_id,
                 parent_agent_id=parent_agent_id,
+                org_id=org_id,
+                registration_key=registration_key,
             )
             registered_adapters, audit_sink = _register_adapters(
                 client=client,
@@ -513,6 +532,8 @@
     native_available: bool,
     team_id: str | None = None,
     parent_agent_id: str | None = None,
+    org_id: str | None = None,
+    registration_key: str | None = None,
 ) -> bool:
     """Register the agent with the gateway over the native gRPC ``register``.
 
@@ -559,6 +580,8 @@
             gateway_endpoint=gateway_endpoint,
             team_id=team_id,
             parent_agent_id=parent_agent_id,
+            org_id=org_id,
+            registration_key=registration_key,
         )
     except Exception as error:
         if _register_failure_is_fatal(mode=mode, enforcement_mode=enforcement_mode):

agent_assembly/core/runtime_interceptor.py

@@ -659,6 +659,8 @@ def register_agent(
     gateway_endpoint: str | None = None,
     team_id: str | None = None,
     parent_agent_id: str | None = None,
+    org_id: str | None = None,
+    registration_key: str | None = None,
 ) -> str | None:
     """Register ``agent_id`` with the gateway over the native ``register`` call.
 
@@ -674,6 +676,14 @@ def register_agent(
     is retried with the legacy positional signature so the SDK keeps working
     rather than raising.
 
+    ``org_id`` is a *claim*: the gateway admits it only when ``registration_key``
+    (an operator API key bound to that organization) is presented with it, and
+    derives the organization from the key when ``org_id`` is omitted
+    (AAASM-6444). Unlike ``team_id`` these are never silently dropped: a native
+    build that cannot carry them raises instead of registering the agent
+    org-less, because an agent the caller believes is org-scoped but is not
+    would write its audit entries outside that organization.
+
     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
@@ -682,6 +692,25 @@ def register_agent(
     register = getattr(runtime_client, "register", None)
     if register is None:
         return None
+    if org_id is not None or registration_key is not None:
+        try:
+            return str(
+                register(
+                    agent_id,
+                    agent_id,
+                    framework,
+                    gateway_endpoint,
+                    team_id,
+                    parent_agent_id,
+                    org_id,
+                    registration_key,
+                )
+            )
+        except TypeError as error:
+            raise RuntimeError(
+                "this native agent_assembly._core build cannot carry org_id/registration_key; "
+                "upgrade the SDK instead of registering the agent without its organization"
+            ) from error
     try:
         return str(register(agent_id, agent_id, framework, gateway_endpoint, team_id, parent_agent_id))
     except TypeError:

native/aa-ffi-python/Cargo.toml

@@ -14,14 +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 agent-assembly AAASM-5835's merge commit, which fixes
-# AssemblyClient::query_policy sending the caller's un-normalized human label
-# instead of the registered did:key DID (every governed tool call failed with
-# a spurious "credential token registered to a different agent" error).
-# AAASM-5839.
-aa-core = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "0044d515d020863767d8011e18e266ec37344e5c", package = "aa-core", features = ["serde"] }
-aa-proto = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "0044d515d020863767d8011e18e266ec37344e5c", package = "aa-proto" }
-aa-sdk-client = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "0044d515d020863767d8011e18e266ec37344e5c", package = "aa-sdk-client" }
+# Pinned to the agent-assembly merge commit of AAASM-6451 (main d4093e89b). It carries
+# AAASM-6447/6451: `AssemblyConfig::org_id` and `registration_key` (the org claim and the
+# operator key the gateway binds it to) and the gateway-restart re-bind, and it keeps the
+# AAASM-5839 fix (query_policy sends the registered did:key DID, not the human label).
+aa-core = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "d4093e89b07d8ba1c38e6caa92637d2e8edb65d7", package = "aa-core", features = ["serde"] }
+aa-proto = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "d4093e89b07d8ba1c38e6caa92637d2e8edb65d7", package = "aa-proto" }
+aa-sdk-client = { git = "https://github.com/ai-agent-assembly/agent-assembly.git", rev = "d4093e89b07d8ba1c38e6caa92637d2e8edb65d7", 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

@@ -108,6 +108,8 @@ impl RuntimeClient {
             // AASM_STATE_DIR / ~/.aasm/identity, matching this shim's pre-AAASM-5332
             // ambient-resolution behavior (AAASM-5839).
             identity_dir: None,
+            org_id: None,
+            registration_key: None,
         };
         let handle = spawn_ipc_thread(
             config.resolve_socket_path(),
@@ -141,14 +143,18 @@ impl RuntimeClient {
     /// team-budget attribution and `parent_agent_id` the topology graph. Both
     /// are optional — `None` leaves the agent team-unscoped / root.
     ///
+    /// `org_id` is a claim and `registration_key` the operator API key the
+    /// gateway binds to an organization (AAASM-6444); see
+    /// [`aa_sdk_client::AssemblyConfig::org_id`].
+    ///
     /// 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.
     // 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))]
+    #[pyo3(signature = (agent_id, name, framework, gateway_endpoint=None, team_id=None, parent_agent_id=None, org_id=None, registration_key=None))]
     fn register(
         &self,
         py: Python<'_>,
@@ -158,6 +164,8 @@ impl RuntimeClient {
         gateway_endpoint: Option<String>,
         team_id: Option<String>,
         parent_agent_id: Option<String>,
+        org_id: Option<String>,
+        registration_key: Option<String>,
     ) -> PyResult<String> {
         let config = AssemblyConfig {
             agent_id,
@@ -170,6 +178,10 @@ impl RuntimeClient {
             sdk_version: None,
             // See the `connect` config above (AAASM-5839).
             identity_dir: None,
+            // AAASM-6444: the org is a claim the gateway validates against the
+            // key; both are optional and absent means an org-less registration.
+            org_id,
+            registration_key: registration_key.map(aa_sdk_client::RegistrationKey::new),
         };
         // `register` is async (tonic). Release the GIL and drive the future on a
         // private current-thread runtime so the blocking gRPC round-trip never