diff --git a/agentex/src/api/routes/agents.py b/agentex/src/api/routes/agents.py index 6b6aad8f..8d4c3f97 100644 --- a/agentex/src/api/routes/agents.py +++ b/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( diff --git a/agentex/src/domain/use_cases/agents_use_case.py b/agentex/src/domain/use_cases/agents_use_case.py index ea130ac4..3d3dbae7 100644 --- a/agentex/src/domain/use_cases/agents_use_case.py +++ b/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 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..8e3b6a3b 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 ) @@ -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