diff --git a/agentex/src/api/routes/agents.py b/agentex/src/api/routes/agents.py index abd10b60..e6004343 100644 --- a/agentex/src/api/routes/agents.py +++ b/agentex/src/api/routes/agents.py @@ -220,6 +220,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( diff --git a/agentex/src/domain/use_cases/agents_use_case.py b/agentex/src/domain/use_cases/agents_use_case.py index ea130ac4..d1851edf 100644 --- a/agentex/src/domain/use_cases/agents_use_case.py +++ b/agentex/src/domain/use_cases/agents_use_case.py @@ -57,7 +57,11 @@ async def _safe_deregister(self, agent_id: str) -> None: "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, + ) -> bool: """Register a newly created agent in the authorization graph. Called before the row is persisted so a failure aborts the create with no @@ -67,29 +71,50 @@ async def _register_in_auth(self, agent_id: str) -> bool: 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. + + 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 await self.authorization_service.register_resource( - AgentexResource.agent(agent_id) + AgentexResource.agent(agent_id), + principal_context=principal_context, ) return True + @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) + ) + async def register_agent( self, name: str, @@ -99,6 +124,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 +234,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 = 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: diff --git a/agentex/tests/integration/use_cases/test_agent_authz_dual_write.py b/agentex/tests/integration/use_cases/test_agent_authz_dual_write.py index 00a1b0e0..6b98346a 100644 --- a/agentex/tests/integration/use_cases/test_agent_authz_dual_write.py +++ b/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 ) @@ -264,6 +264,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