diff --git a/agentex/openapi.yaml b/agentex/openapi.yaml index b122020e..100fb9b8 100644 --- a/agentex/openapi.yaml +++ b/agentex/openapi.yaml @@ -5221,6 +5221,12 @@ components: acp_type: $ref: '#/components/schemas/ACPType' description: The type of ACP to use for the agent. + principal_context: + anyOf: + - {} + - type: 'null' + title: Principal Context + description: Principal used for authorization registration_metadata: anyOf: - additionalProperties: true diff --git a/agentex/src/api/routes/agents.py b/agentex/src/api/routes/agents.py index 245c865e..b7d370d3 100644 --- a/agentex/src/api/routes/agents.py +++ b/agentex/src/api/routes/agents.py @@ -195,11 +195,15 @@ async def register_agent( If agent_id is not provided, the system will look for an existing agent by name and update it, or create a new one if it doesn't exist. """ - enforce_ownership = _has_resolvable_creator(authorization_service.principal_context) + principal_context = authorization_service.principal_context + if not _has_resolvable_creator(principal_context): + principal_context = request.principal_context + enforce_ownership = _has_resolvable_creator(principal_context) if enforce_ownership: await authorization_service.check( AgentexResource.agent("*"), AuthorizedOperationType.create, + principal_context=principal_context, ) logger.info( "Registering agent name=%s agent_id=%s acp_type=%s", @@ -220,6 +224,7 @@ async def register_agent( if enforce_ownership: await authorization_service.grant( AgentexResource.agent(agent_entity.id), + principal_context=principal_context, ) response_fields = agent_entity.model_dump() existing_key = await api_keys_use_case.get_internal_api_key_by_agent_id( diff --git a/agentex/src/api/schemas/agents.py b/agentex/src/api/schemas/agents.py index 40fd4273..79d92532 100644 --- a/agentex/src/api/schemas/agents.py +++ b/agentex/src/api/schemas/agents.py @@ -88,6 +88,9 @@ class RegisterAgentRequest(BaseModel): description="Optional agent ID if the agent already exists and needs to be updated.", ) acp_type: ACPType = Field(..., description="The type of ACP to use for the agent.") + principal_context: Any | None = Field( + default=None, description="Principal used for authorization" + ) registration_metadata: dict[str, Any] | None = Field( default=None, description="The metadata for the agent's registration.", diff --git a/agentex/tests/unit/api/test_agents_authz.py b/agentex/tests/unit/api/test_agents_authz.py index 650df233..581d79d5 100644 --- a/agentex/tests/unit/api/test_agents_authz.py +++ b/agentex/tests/unit/api/test_agents_authz.py @@ -392,12 +392,13 @@ def _mocks(principal_context): return authorization, agents_use_case, api_keys_use_case @staticmethod - def _request(): + def _request(principal_context=None): return RegisterAgentRequest( name="my-agent", description="d", acp_url="http://agent:5000", acp_type=ACPType.ASYNC, + principal_context=principal_context, ) @pytest.mark.parametrize("principal_context", [None, {}, {"account_id": "acct"}]) @@ -413,6 +414,21 @@ async def test_unresolvable_creator_skips_check_and_grant(self, principal_contex use_case.register_agent.assert_awaited_once() assert resp.agent_api_key == "internal-key" + async def test_body_principal_enforces_check_and_grant(self): + # Legacy deployed pods send the deploy principal in the body because + # /agents/register is whitelisted and has no request-state principal. + body_principal = {"user_id": "u", "account_id": "acct"} + authz, use_case, api_keys = self._mocks(principal_context=None) + + await register_agent( + self._request(principal_context=body_principal), use_case, authz, api_keys + ) + + authz.check.assert_awaited_once() + assert authz.check.await_args.kwargs["principal_context"] == body_principal + authz.grant.assert_awaited_once() + assert authz.grant.await_args.kwargs["principal_context"] == body_principal + async def test_dict_principal_enforces_check_and_grant(self): authz, use_case, api_keys = self._mocks( principal_context={"user_id": "u", "account_id": "acct"}