diff --git a/agent_assembly/core/assembly.py b/agent_assembly/core/assembly.py index ffe90465..96812afd 100644 --- a/agent_assembly/core/assembly.py +++ b/agent_assembly/core/assembly.py @@ -194,6 +194,8 @@ def init_assembly( 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, @@ -245,6 +247,16 @@ def init_assembly( 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 @@ def init_assembly( 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 @@ def init_assembly( 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 @@ def _register_agent_with_gateway( 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 @@ def _register_agent_with_gateway( 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): diff --git a/agent_assembly/core/runtime_interceptor.py b/agent_assembly/core/runtime_interceptor.py index 50532f36..f52e162a 100644 --- a/agent_assembly/core/runtime_interceptor.py +++ b/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: diff --git a/native/aa-ffi-python/Cargo.lock b/native/aa-ffi-python/Cargo.lock index f8cd6fe1..efed1518 100644 --- a/native/aa-ffi-python/Cargo.lock +++ b/native/aa-ffi-python/Cargo.lock @@ -4,20 +4,22 @@ version = 4 [[package]] name = "aa-core" -version = "0.0.1-rc.7" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=0044d515d020863767d8011e18e266ec37344e5c#0044d515d020863767d8011e18e266ec37344e5c" +version = "0.0.1-rc.8" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=d4093e89b07d8ba1c38e6caa92637d2e8edb65d7#d4093e89b07d8ba1c38e6caa92637d2e8edb65d7" dependencies = [ "aa-security", "async-trait", "cfg-if", "chrono", "dirs", + "libc", "serde", "serde_json", - "serde_yaml", + "serde_norway", "sha2", "thiserror", "toml", + "toml_edit", ] [[package]] @@ -36,8 +38,8 @@ dependencies = [ [[package]] name = "aa-proto" -version = "0.0.1-rc.7" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=0044d515d020863767d8011e18e266ec37344e5c#0044d515d020863767d8011e18e266ec37344e5c" +version = "0.0.1-rc.8" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=d4093e89b07d8ba1c38e6caa92637d2e8edb65d7#d4093e89b07d8ba1c38e6caa92637d2e8edb65d7" dependencies = [ "prost", "tonic", @@ -47,8 +49,8 @@ dependencies = [ [[package]] name = "aa-sdk-client" -version = "0.0.1-rc.7" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=0044d515d020863767d8011e18e266ec37344e5c#0044d515d020863767d8011e18e266ec37344e5c" +version = "0.0.1-rc.8" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=d4093e89b07d8ba1c38e6caa92637d2e8edb65d7#d4093e89b07d8ba1c38e6caa92637d2e8edb65d7" dependencies = [ "aa-proto", "aa-security", @@ -65,12 +67,12 @@ dependencies = [ [[package]] name = "aa-security" -version = "0.0.1-rc.7" -source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=0044d515d020863767d8011e18e266ec37344e5c#0044d515d020863767d8011e18e266ec37344e5c" +version = "0.0.1-rc.8" +source = "git+https://github.com/ai-agent-assembly/agent-assembly.git?rev=d4093e89b07d8ba1c38e6caa92637d2e8edb65d7#d4093e89b07d8ba1c38e6caa92637d2e8edb65d7" dependencies = [ "aho-corasick", "serde", - "serde_yaml", + "serde_norway", ] [[package]] @@ -1116,25 +1118,25 @@ dependencies = [ ] [[package]] -name = "serde_spanned" -version = "1.1.1" +name = "serde_norway" +version = "0.9.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +checksum = "e408f29489b5fd500fab51ff1484fc859bb655f32c671f307dcd733b72e8168c" dependencies = [ - "serde_core", + "indexmap", + "itoa", + "ryu", + "serde", + "unsafe-libyaml-norway", ] [[package]] -name = "serde_yaml" -version = "0.9.34+deprecated" +name = "serde_spanned" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" dependencies = [ - "indexmap", - "itoa", - "ryu", - "serde", - "unsafe-libyaml", + "serde_core", ] [[package]] @@ -1344,6 +1346,19 @@ dependencies = [ "serde_core", ] +[[package]] +name = "toml_edit" +version = "0.25.15+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1340ea94a5856333492c9064b02c778b191dd2c853778d9609debdcdfea3a614" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "toml_writer", + "winnow", +] + [[package]] name = "toml_parser" version = "1.1.3+spec-1.1.0" @@ -1520,10 +1535,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" [[package]] -name = "unsafe-libyaml" -version = "0.2.11" +name = "unsafe-libyaml-norway" +version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" +checksum = "b39abd59bf32521c7f2301b52d05a6a2c975b6003521cbd0c6dc1582f0a22104" [[package]] name = "want" @@ -1710,6 +1725,9 @@ name = "winnow" version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] [[package]] name = "wit-bindgen" diff --git a/native/aa-ffi-python/Cargo.toml b/native/aa-ffi-python/Cargo.toml index bb6d49c0..c88c2aa7 100644 --- a/native/aa-ffi-python/Cargo.toml +++ b/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"] } diff --git a/native/aa-ffi-python/src/lib.rs b/native/aa-ffi-python/src/lib.rs index 473f3d6c..4db76935 100644 --- a/native/aa-ffi-python/src/lib.rs +++ b/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,6 +143,10 @@ 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. @@ -148,7 +154,7 @@ impl RuntimeClient { // 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, team_id: Option, parent_agent_id: Option, + org_id: Option, + registration_key: Option, ) -> PyResult { 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 diff --git a/test/unit/core/test_init_org_identity.py b/test/unit/core/test_init_org_identity.py new file mode 100644 index 00000000..05507294 --- /dev/null +++ b/test/unit/core/test_init_org_identity.py @@ -0,0 +1,140 @@ +"""``init_assembly(org_id=..., registration_key=...)`` reaches the native register (AAASM-6444). + +The organization is a *claim* the gateway validates against the registration +key, so the SDK must (a) carry both to the native client, (b) leave the +pre-existing call shape untouched when neither is given, and (c) never fall back +to registering the agent org-less on a native build that cannot carry them. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from agent_assembly import init_assembly +from agent_assembly.core import assembly as core_assembly +from agent_assembly.core.audit_sink import AUDIT_SINK_ABSENT +from agent_assembly.exceptions import ConfigurationError + +from ._fake_core import FakeRuntimeClient, install_fake_core +from .test_init_registration import _API_KEY, _GRPC_ENDPOINT, _GW_URL, _no_network + + +class _OrgAwareRuntimeClient(FakeRuntimeClient): + """A native client whose ``register`` accepts the org claim and key.""" + + def __init__(self) -> None: + super().__init__(decision="allow") + self.org_register_calls: list[tuple[Any, ...]] = [] + + def register( + self, + agent_id: str, + name: str, + framework: str, + 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: + self.org_register_calls.append( + (agent_id, name, framework, gateway_endpoint, team_id, parent_agent_id, org_id, registration_key) + ) + if self.register_should_raise is not None: + raise self.register_should_raise + return "policy-id-001" + + +def _init(monkeypatch: pytest.MonkeyPatch, client: FakeRuntimeClient, **kwargs: Any) -> Any: + install_fake_core(monkeypatch, client) + _no_network(monkeypatch) + monkeypatch.setattr(core_assembly, "_register_adapters", lambda **_kw: ([], AUDIT_SINK_ABSENT)) + return init_assembly(gateway_url=_GW_URL, api_key=_API_KEY, mode="sdk-only", **kwargs) + + +def test_init_assembly_forwards_org_claim_and_registration_key(monkeypatch: pytest.MonkeyPatch) -> None: + client = _OrgAwareRuntimeClient() + context = _init( + monkeypatch, + client, + agent_id="acme-bot", + team_id="platform", + org_id="acme", + registration_key="aa_00112233445566778899aabbccddeeff", + ) + try: + assert client.org_register_calls == [ + ( + "acme-bot", + "acme-bot", + "python", + _GRPC_ENDPOINT, + "platform", + None, + "acme", + "aa_00112233445566778899aabbccddeeff", + ) + ] + finally: + context.shutdown() + + +def test_init_assembly_key_alone_lets_the_gateway_derive_the_org(monkeypatch: pytest.MonkeyPatch) -> None: + client = _OrgAwareRuntimeClient() + context = _init(monkeypatch, client, agent_id="derived", registration_key="aa_k") + try: + assert client.org_register_calls[0][6:] == (None, "aa_k") + finally: + context.shutdown() + + +def test_init_assembly_without_org_keeps_the_existing_register_call_shape(monkeypatch: pytest.MonkeyPatch) -> None: + client = FakeRuntimeClient(decision="allow") + context = _init(monkeypatch, client, agent_id="plain") + try: + assert client.register_calls == [("plain", "plain", "python", _GRPC_ENDPOINT, None, None)] + finally: + context.shutdown() + + +def test_init_assembly_never_registers_org_less_when_native_cannot_carry_the_org( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Negative control: the legacy 6-arg client has no org parameters. Silently + dropping the claim would register the agent outside the organization the + caller asked for, so init fails instead and nothing is registered.""" + client = FakeRuntimeClient(decision="allow") # register() has no org_id kwarg + with pytest.raises(ConfigurationError, match="cannot carry org_id"): + _init(monkeypatch, client, agent_id="acme-bot", org_id="acme", enforcement_mode="enforce") + assert client.register_calls == [] + + +def test_init_assembly_treats_blank_org_and_key_as_absent(monkeypatch: pytest.MonkeyPatch) -> None: + """A blank value must not become an empty claim or an empty key header.""" + client = FakeRuntimeClient(decision="allow") # legacy 6-arg register: any org kwarg would raise TypeError + context = _init(monkeypatch, client, agent_id="blank", org_id=" ", registration_key="") + try: + assert client.register_calls == [("blank", "blank", "python", _GRPC_ENDPOINT, None, None)] + finally: + context.shutdown() + + +def test_registration_failure_output_never_contains_the_registration_key( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], caplog: pytest.LogCaptureFixture +) -> None: + """The SDK adds nothing to a failed registration's output: the key is neither + in the stderr warning nor in any log record, whatever the native layer raised.""" + secret = "aa_" + "9" * 32 + client = _OrgAwareRuntimeClient() + client.register_should_raise = RuntimeError("gateway rejected registration: no valid key") + with caplog.at_level("DEBUG"): + context = _init(monkeypatch, client, agent_id="leaky", org_id="acme", registration_key=secret) + try: + captured = capsys.readouterr() + assert "NOT registered" in captured.err # the failure path really ran + assert secret not in captured.err + captured.out + assert secret not in caplog.text + finally: + context.shutdown()