fix(authz): fall back to body principal in _register_in_auth

agentex/src/api/routes/agents.py

@@ -223,6 +223,9 @@ async def register_agent(
             acp_type=request.acp_type,
             registration_metadata=request.registration_metadata,
             agent_input_type=request.agent_input_type,
+            # Same fallback as check/grant above: the use case prefers the
+            # middleware principal and only reads this on the whitelisted path.
+            body_principal_context=request.principal_context,
         )
         if enforce_ownership:
             await authorization_service.grant(

agentex/src/domain/use_cases/agents_use_case.py

@@ -41,54 +41,85 @@ def __init__(
         self.temporal_adapter = temporal_adapter
         self.authorization_service = authorization_service
 
-    async def _safe_deregister(self, agent_id: str) -> None:
+    async def _safe_deregister(
+        self, agent_id: str, principal_context: Any = None
+    ) -> None:
         """Best-effort removal of an agent from the authorization graph.
 
         Swallows and logs any failure so a compensating/post-delete deregister
         never masks the in-flight error (create) or fails a delete that already
-        succeeded.
+        succeeded. ``principal_context`` should be the same resolved principal
+        that was used for the preceding ``register_resource`` call so that
+        register and deregister are symmetric on all paths (including the
+        whitelisted pod path where the middleware principal is None).
         """
         try:
             await self.authorization_service.deregister_resource(
-                AgentexResource.agent(agent_id)
+                AgentexResource.agent(agent_id),
+                principal_context=principal_context,
             )
         except Exception:
             logger.exception(
                 "authorization deregister failed for agent %s; swallowed", agent_id
             )
 
-    async def _register_in_auth(self, agent_id: str) -> bool:
+    async def _register_in_auth(
+        self,
+        agent_id: str,
+        body_principal_context: Any = None,
+    ) -> tuple[bool, Any]:
         """Register a newly created agent in the authorization graph.
 
         Called before the row is persisted so a failure aborts the create with no
         orphaned row; the _safe_deregister compensation undoes it. Skipped with a
         warning when no creator identity is resolvable on the principal context:
         an agent has no parent edge, so the principal is the sole anchor for
-        ownership and there is nothing to attribute it to. Returns whether a
-        register call was actually made so compensation can avoid deregistering
-        a resource it never registered.
+        ownership and there is nothing to attribute it to. Returns a tuple of
+        (registered, resolved_principal_context) so the caller can pass the same
+        principal to _safe_deregister — keeping register and deregister symmetric.
+
+        Prefer the middleware-derived principal (an authenticated user or service
+        account). The whitelisted ``/agents/register`` path clears the middleware
+        principal; the pod SDK ships the manifest-declared identity in the request
+        body instead, and ``body_principal_context`` carries that through so
+        ownership can still be minted from the manifest identity. Mirrors the
+        route-layer fallback the ``register_agent`` route applies to
+        ``check``/``grant``.
         """
         principal_context = self.authorization_service.principal_context
-        # principal_context is `Any` (a dict from /v1/authn), not a typed model,
-        # so attribute access via getattr always yields None and silently skips
-        # the Spark resource registration. Read from the dict (fall back to attr
-        # access for any object-shaped principal).
-        if isinstance(principal_context, dict):
-            user_id = principal_context.get("user_id")
-            service_account_id = principal_context.get("service_account_id")
-        else:
-            user_id = getattr(principal_context, "user_id", None)
-            service_account_id = getattr(principal_context, "service_account_id", None)
-        if user_id is None and service_account_id is None:
+        if not self._has_resolvable_creator(principal_context):
+            principal_context = body_principal_context
+        if not self._has_resolvable_creator(principal_context):
             logger.warning(
                 "Skipping authorization registration for agent: no creator resolvable",
                 extra={"agent_id": agent_id},
             )
-            return False
+            return False, None
         await self.authorization_service.register_resource(
-            AgentexResource.agent(agent_id)
+            AgentexResource.agent(agent_id),
+            principal_context=principal_context,
+        )
+        return True, principal_context
+
+    @staticmethod
+    def _has_resolvable_creator(principal_context: Any) -> bool:
+        """Whether a creator identity (user or service account) is present.
+
+        ``principal_context`` is ``Any`` (dict from ``/v1/authn`` or an object
+        for tests), so attribute access via getattr on a dict always yields
+        None. Read from the dict when applicable, fall back to attr access.
+        """
+        if principal_context is None:
+            return False
+        if isinstance(principal_context, dict):
+            return bool(
+                principal_context.get("user_id")
+                or principal_context.get("service_account_id")
+            )
+        return bool(
+            getattr(principal_context, "user_id", None)
+            or getattr(principal_context, "service_account_id", None)
         )
-        return True
 
     async def register_agent(
         self,
@@ -99,6 +130,7 @@ async def register_agent(
         acp_type: ACPType = ACPType.ASYNC,
         registration_metadata: dict[str, Any] | None = None,
         agent_input_type: AgentInputType | None = None,
+        body_principal_context: Any = None,
     ) -> AgentEntity:
         deployment_id = (registration_metadata or {}).get("deployment_id")
 
@@ -208,7 +240,9 @@ async def register_agent(
             # failure here aborts the create with no orphaned row. Only the
             # genuine-create path registers — the update paths above must not,
             # or re-registering would rewrite the owner to the current caller.
-            registered_in_auth = await self._register_in_auth(agent.id)
+            registered_in_auth, resolved_principal = await self._register_in_auth(
+                agent.id, body_principal_context=body_principal_context
+            )
             # This is a problem only if multiple pods spin up and then make a request all at the same time.
             # In that case, the first pod will create the agent and the rest should succeed silently
             try:
@@ -221,11 +255,15 @@ async def register_agent(
                 # registration and re-fetch the persisted agent so downstream
                 # code (complete_deployment_registration) uses the correct agent_id
                 if registered_in_auth:
-                    await self._safe_deregister(agent.id)
+                    await self._safe_deregister(
+                        agent.id, principal_context=resolved_principal
+                    )
                 agent = await self.agent_repo.get(name=name)
             except Exception:
                 if registered_in_auth:
-                    await self._safe_deregister(agent.id)
+                    await self._safe_deregister(
+                        agent.id, principal_context=resolved_principal
+                    )
                 raise
         await self.complete_deployment_registration(
             agent, acp_url, registration_metadata
@@ -276,7 +314,7 @@ async def register_build(
         # Record ownership before persisting, same as register_agent's genuine
         # create branch. The early return above for an existing agent means we
         # only ever register on a true first-time build create.
-        registered_in_auth = await self._register_in_auth(agent.id)
+        registered_in_auth, resolved_principal = await self._register_in_auth(agent.id)
         # If multiple builds for the same new agent race, the first wins and the
         # rest re-fetch the persisted row instead of erroring.
         try:
@@ -287,11 +325,15 @@ async def register_build(
             )
             # undo our ownership registration from _register_in_auth
             if registered_in_auth:
-                await self._safe_deregister(agent.id)
+                await self._safe_deregister(
+                    agent.id, principal_context=resolved_principal
+                )
             agent = await self.agent_repo.get(name=name)
         except Exception:
             if registered_in_auth:
-                await self._safe_deregister(agent.id)
+                await self._safe_deregister(
+                    agent.id, principal_context=resolved_principal
+                )
             raise
         return agent
 

agentex/tests/integration/use_cases/test_agent_authz_dual_write.py

@@ -96,7 +96,7 @@ async def test_create_registers_before_persist_with_no_parent(
         # what makes a registration failure abort the request cleanly.
         observed = {}
 
-        async def _record_existence(resource, parent=None):
+        async def _record_existence(resource, parent=None, *, principal_context=None):
             observed["row_exists_at_register"] = await _agent_exists(
                 agent_repo, resource.selector
             )
@@ -199,6 +199,37 @@ async def test_duplicate_compensates_then_adopts_existing_row(self):
         compensated = authorization_service.deregister_resource.call_args.args[0]
         assert compensated.selector == registered.selector
 
+    async def test_duplicate_compensates_with_body_principal_on_whitelisted_path(self):
+        # Whitelisted path: middleware principal is None, body principal used for
+        # register. Compensation deregister must forward the same body principal —
+        # not None — so register and deregister are symmetric.
+        existing = _existing_agent(f"dw-dup-body-{uuid4().hex[:8]}")
+        agent_repo = Mock()
+        agent_repo.get = AsyncMock(side_effect=[ItemDoesNotExist("absent"), existing])
+        agent_repo.create = AsyncMock(side_effect=DuplicateItemError("exists"))
+        use_case, authorization_service = _build_use_case(
+            agent_repository=agent_repo,
+            principal=_principal(user_id=None, service_account_id=None),
+        )
+
+        body_principal = {
+            "service_account_id": "sa-from-manifest",
+            "account_id": "acct-1",
+        }
+
+        result = await use_case.register_agent(
+            name=existing.name,
+            description="dup on whitelisted path",
+            acp_url="http://new-acp",
+            body_principal_context=body_principal,
+        )
+
+        assert result.id == existing.id
+        authorization_service.register_resource.assert_awaited_once()
+        authorization_service.deregister_resource.assert_awaited_once()
+        deregister_call = authorization_service.deregister_resource.call_args
+        assert deregister_call.kwargs.get("principal_context") == body_principal
+
     async def test_update_path_does_not_register(self):
         # register_agent called with an agent_id is an update, not a create: it
         # must NOT register, or ownership would be rewritten to the caller.
@@ -264,6 +295,72 @@ async def test_persist_failure_without_resolvable_creator_skips_compensation(sel
         authorization_service.grant.assert_not_awaited()
         authorization_service.revoke.assert_not_awaited()
 
+    async def test_create_falls_back_to_body_principal_when_middleware_missing(
+        self,
+    ):
+        # Pod self-registration via whitelisted /agents/register clears the
+        # middleware principal, so the SDK ships the manifest-declared identity
+        # in the request body. The use case must accept that body principal so
+        # ownership is still minted from the manifest identity — otherwise the
+        # agent registers but has no owner tuple and is invisible to the UI.
+        agent_repo = Mock()
+        agent_repo.get = AsyncMock(side_effect=ItemDoesNotExist("absent"))
+        agent_repo.create = AsyncMock(side_effect=lambda item: item)
+        use_case, authorization_service = _build_use_case(
+            agent_repository=agent_repo,
+            principal=_principal(user_id=None, service_account_id=None),
+        )
+
+        # Dict shape matches what /v1/authn returns and what the SDK decodes
+        # from AUTH_PRINCIPAL_B64 (see scale-agentex-python registration.py).
+        body_principal = {
+            "service_account_id": "sa-from-manifest",
+            "account_id": "acct-1",
+        }
+
+        agent = await use_case.register_agent(
+            name=f"dw-body-fallback-{uuid4().hex[:8]}",
+            description="body principal fallback",
+            acp_url="http://new-acp",
+            body_principal_context=body_principal,
+        )
+
+        authorization_service.register_resource.assert_awaited_once()
+        call = authorization_service.register_resource.call_args
+        registered_resource: AgentexResource = call.args[0]
+        assert registered_resource.type == AgentexResourceType.agent
+        assert registered_resource.selector == agent.id
+        # The body principal must be forwarded to the gateway so it doesn't
+        # fall back to the (unset) middleware principal.
+        assert call.kwargs.get("principal_context") == body_principal
+
+    async def test_middleware_principal_takes_precedence_over_body(self):
+        # An authenticated caller (via CLI/UI) sets the middleware principal.
+        # Even if a body principal is also sent, the authenticated identity
+        # from the middleware must win — the body is only a fallback for the
+        # whitelisted pod path.
+        agent_repo = Mock()
+        agent_repo.get = AsyncMock(side_effect=ItemDoesNotExist("absent"))
+        agent_repo.create = AsyncMock(side_effect=lambda item: item)
+        middleware_principal = _principal(user_id="user-from-middleware")
+        use_case, authorization_service = _build_use_case(
+            agent_repository=agent_repo,
+            principal=middleware_principal,
+        )
+
+        body_principal = {"service_account_id": "sa-should-not-win"}
+
+        await use_case.register_agent(
+            name=f"dw-middleware-wins-{uuid4().hex[:8]}",
+            description="middleware principal wins",
+            acp_url="http://new-acp",
+            body_principal_context=body_principal,
+        )
+
+        authorization_service.register_resource.assert_awaited_once()
+        call = authorization_service.register_resource.call_args
+        assert call.kwargs.get("principal_context") is middleware_principal
+
 
 @pytest.mark.integration
 @pytest.mark.asyncio