diff --git a/agentex/src/adapters/authorization/adapter_agentex_authz_proxy.py b/agentex/src/adapters/authorization/adapter_agentex_authz_proxy.py index 0e44479d..c8ba4f3a 100644 --- a/agentex/src/adapters/authorization/adapter_agentex_authz_proxy.py +++ b/agentex/src/adapters/authorization/adapter_agentex_authz_proxy.py @@ -85,6 +85,34 @@ async def list_resources( ) return response["items"] + async def register_resource( + self, + principal: AgentexAuthPrincipalContext, + resource: AgentexResource, + parent: AgentexResource | None = None, + ) -> None: + payload = { + "principal": principal, + "resource": resource.model_dump(), + "parent": parent.model_dump() if parent is not None else None, + } + await HttpRequestHandler.post_with_error_handling( + self.agentex_auth_url, "/v1/authz/register", json=payload + ) + + async def deregister_resource( + self, + principal: AgentexAuthPrincipalContext, + resource: AgentexResource, + ) -> None: + payload = { + "principal": principal, + "resource": resource.model_dump(), + } + await HttpRequestHandler.post_with_error_handling( + self.agentex_auth_url, "/v1/authz/deregister", json=payload + ) + DAgentexAuthorization = Annotated[ AgentexAuthorizationProxy, Depends(AgentexAuthorizationProxy) diff --git a/agentex/src/adapters/authorization/port.py b/agentex/src/adapters/authorization/port.py index 80c05200..f4f33a0a 100644 --- a/agentex/src/adapters/authorization/port.py +++ b/agentex/src/adapters/authorization/port.py @@ -49,3 +49,29 @@ async def list_resources( filter_operation: AuthorizedOperationType = AuthorizedOperationType.read, ) -> Iterable[str]: """List resource_ids for a given principal""" + + @abstractmethod + async def register_resource( + self, + principal: PrincipalT, + resource: AgentexResource, + parent: AgentexResource | None = None, + ) -> None: + """Register a newly created resource in SpiceDB with the principal as + owner. Optionally writes a lifecycle parent edge. + + Use this on resource create instead of ``grant`` when the resource + type's SpiceDB definition has a parent relation that permission + checks cascade through (e.g. ``agent_api_key`` declares + ``parent_agent``). Without writing that edge here the cascade fails + closed. + """ + + @abstractmethod + async def deregister_resource( + self, + principal: PrincipalT, + resource: AgentexResource, + ) -> None: + """Deregister a deleted resource and all of its relationships + (owner, parent, grantees) in a single atomic call.""" diff --git a/agentex/src/api/routes/agent_api_keys.py b/agentex/src/api/routes/agent_api_keys.py index dec3e00f..e041531f 100644 --- a/agentex/src/api/routes/agent_api_keys.py +++ b/agentex/src/api/routes/agent_api_keys.py @@ -8,6 +8,7 @@ CreateAPIKeyResponse, ) from src.domain.entities.agent_api_keys import AgentAPIKeyType +from src.domain.services.authorization_service import DAuthorizationService from src.domain.use_cases.agent_api_keys_use_case import DAgentAPIKeysUseCase from src.domain.use_cases.agents_use_case import DAgentsUseCase from src.utils.logging import make_logger @@ -28,6 +29,7 @@ async def create_api_key( request: CreateAPIKeyRequest, agent_api_key_use_case: DAgentAPIKeysUseCase, agent_use_case: DAgentsUseCase, + authorization_service: DAuthorizationService, ) -> CreateAPIKeyResponse: if not request.agent_id and not request.agent_name: raise HTTPException( @@ -52,11 +54,13 @@ async def create_api_key( raise HTTPException(status_code=409, detail=error_msg) new_api_key = request.api_key or secrets.token_hex(32) + account_id = getattr(authorization_service.principal_context, "account_id", None) agent_api_key_entity = await agent_api_key_use_case.create( agent_id=agent.id, api_key=str(new_api_key), name=request.name, api_key_type=request.api_key_type, + account_id=account_id, ) return CreateAPIKeyResponse( id=agent_api_key_entity.id, @@ -161,8 +165,10 @@ async def get_agent_api_key( async def delete_agent_api_key( id: str, agent_api_key_use_case: DAgentAPIKeysUseCase, + authorization_service: DAuthorizationService, ) -> str: - await agent_api_key_use_case.delete(id=id) + account_id = getattr(authorization_service.principal_context, "account_id", None) + await agent_api_key_use_case.delete(id=id, account_id=account_id) return f"Agent API key with ID {id} deleted" @@ -176,6 +182,7 @@ async def delete_agent_api_key_by_name( api_key_name: str, agent_api_key_use_case: DAgentAPIKeysUseCase, agent_use_case: DAgentsUseCase, + authorization_service: DAuthorizationService, agent_id: str | None = None, agent_name: str | None = None, api_key_type: AgentAPIKeyType = AgentAPIKeyType.EXTERNAL, @@ -191,8 +198,12 @@ async def delete_agent_api_key_by_name( detail="Only one of 'agent_id' or 'agent_name' should be provided to delete an agent api_key.", ) agent = await agent_use_case.get(id=agent_id, name=agent_name) + account_id = getattr(authorization_service.principal_context, "account_id", None) await agent_api_key_use_case.delete_by_agent_id_and_key_name( - agent_id=agent.id, key_name=api_key_name, api_key_type=api_key_type + agent_id=agent.id, + key_name=api_key_name, + api_key_type=api_key_type, + account_id=account_id, ) return f"Agent api_key '{api_key_name}' deleted" diff --git a/agentex/src/api/schemas/authorization_types.py b/agentex/src/api/schemas/authorization_types.py index 585fa3a3..7e98fdae 100644 --- a/agentex/src/api/schemas/authorization_types.py +++ b/agentex/src/api/schemas/authorization_types.py @@ -14,6 +14,7 @@ class AuthorizedOperationType(StrEnum): class AgentexResourceType(StrEnum): agent = "agent" task = "task" + api_key = "api_key" # Resources that inherit permissions from their parent task @@ -37,6 +38,10 @@ def agent(cls, selector: str) -> "AgentexResource": def task(cls, selector: str) -> "AgentexResource": return cls(type=AgentexResourceType.task, selector=selector) + @classmethod + def api_key(cls, selector: str) -> "AgentexResource": + return cls(type=AgentexResourceType.api_key, selector=selector) + class AgentexResourceOptionalSelector(BaseModel): type: AgentexResourceType @@ -49,3 +54,7 @@ def agent(cls, selector: str | None = None) -> "AgentexResourceOptionalSelector" @classmethod def task(cls, selector: str | None = None) -> "AgentexResourceOptionalSelector": return cls(type=AgentexResourceType.task, selector=selector) + + @classmethod + def api_key(cls, selector: str | None = None) -> "AgentexResourceOptionalSelector": + return cls(type=AgentexResourceType.api_key, selector=selector) diff --git a/agentex/src/domain/services/authorization_service.py b/agentex/src/domain/services/authorization_service.py index 8400603b..06125d57 100644 --- a/agentex/src/domain/services/authorization_service.py +++ b/agentex/src/domain/services/authorization_service.py @@ -188,5 +188,62 @@ async def list_resources( ) return result + async def register_resource( + self, + resource: AgentexResource, + parent: AgentexResource | None = None, + *, + principal_context=..., + ) -> None: + """Register a newly created resource with the principal as owner. + + Prefer this over ``grant`` when the resource's SpiceDB schema has + a parent relation that permissions cascade through (e.g. + ``agent_api_key`` declares ``parent_agent``). Pass ``parent`` to + link the child to its parent atomically; without it the cascade + fails closed. + """ + if self._bypass(): + logger.info(f"Authorization bypassed for register_resource on {resource}") + return None + + effective_principal = ( + principal_context + if principal_context is not ... + else self.principal_context + ) + logger.info( + "[authorization_service] Registering %s:%s for principal %s (parent=%s)", + resource.type, + resource.selector, + effective_principal, + f"{parent.type}:{parent.selector}" if parent is not None else None, + ) + await self.gateway.register_resource(effective_principal, resource, parent) + + async def deregister_resource( + self, + resource: AgentexResource, + *, + principal_context=..., + ) -> None: + """Deregister a deleted resource and all of its relationships.""" + if self._bypass(): + logger.info(f"Authorization bypassed for deregister_resource on {resource}") + return None + + effective_principal = ( + principal_context + if principal_context is not ... + else self.principal_context + ) + logger.info( + "[authorization_service] Deregistering %s:%s for principal %s", + resource.type, + resource.selector, + effective_principal, + ) + await self.gateway.deregister_resource(effective_principal, resource) + DAuthorizationService = Annotated[AuthorizationService, Depends(AuthorizationService)] diff --git a/agentex/src/domain/use_cases/agent_api_keys_use_case.py b/agentex/src/domain/use_cases/agent_api_keys_use_case.py index bf455aee..94ea6de9 100644 --- a/agentex/src/domain/use_cases/agent_api_keys_use_case.py +++ b/agentex/src/domain/use_cases/agent_api_keys_use_case.py @@ -13,6 +13,7 @@ ) from src.adapters.crud_store.exceptions import ItemDoesNotExist from src.api.middleware_utils import get_request_headers_to_forward, verify_auth_gateway +from src.api.schemas.authorization_types import AgentexResource from src.config.dependencies import ( DHttpxClient, resolve_environment_variable_dependency, @@ -27,6 +28,7 @@ DAgentAPIKeyRepository, ) from src.domain.repositories.agent_repository import DAgentRepository +from src.domain.services.authorization_service import DAuthorizationService from src.utils.ids import orm_id from src.utils.logging import make_logger @@ -39,10 +41,12 @@ def __init__( agent_api_key_repository: DAgentAPIKeyRepository, agent_repository: DAgentRepository, client: DHttpxClient, + authorization_service: DAuthorizationService, ): self.agent_api_key_repo = agent_api_key_repository self.agent_repo = agent_repository self.client = client + self.authorization_service = authorization_service self.auth_gateway_enabled = bool( resolve_environment_variable_dependency(EnvVarKeys.AGENTEX_AUTH_URL) ) @@ -76,6 +80,7 @@ async def create( agent_id: str, api_key_type: AgentAPIKeyType, api_key: str, + account_id: str | None = None, ) -> AgentAPIKeyEntity: agent = await self.get_agent(agent_id=agent_id) if not agent: @@ -83,10 +88,21 @@ async def create( status_code=404, detail=f"Agent ID {agent_id} not found.", ) + + api_key_id = orm_id() + + # Unconditional — agentex-auth decides per-account whether the call + # routes to Spark or the legacy backend. + await self._register_api_key_in_auth( + api_key_id=api_key_id, + agent_id=agent.id, + account_id=account_id, + ) + # TODO: encrypt API key before storing it # Initialize a new agent api_key agent_api_key = AgentAPIKeyEntity( - id=orm_id(), + id=api_key_id, name=name, agent_id=agent.id, api_key_type=api_key_type, @@ -94,6 +110,80 @@ async def create( ) return await self.agent_api_key_repo.create(item=agent_api_key) + async def _register_api_key_in_auth( + self, + *, + api_key_id: str, + agent_id: str, + account_id: str | None, + ) -> None: + """Register a new agent_api_key with the auth service, including the + parent_agent edge so permissions cascade from the owning agent. + + Called BEFORE the Postgres write — a failure raises and prevents the + row from being persisted, so there is no compensating action. Skipped + with a warning when no usable creator identity is available on the + principal context (e.g. internal-key creation paths without an + authenticated user). + """ + principal_context = self.authorization_service.principal_context + 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: + logger.warning( + "Skipping auth registration for api_key: no creator resolvable", + extra={ + "api_key_id": api_key_id, + "agent_id": agent_id, + "account_id": account_id, + }, + ) + return + try: + await self.authorization_service.register_resource( + resource=AgentexResource.api_key(api_key_id), + parent=AgentexResource.agent(agent_id), + ) + except Exception as exc: + # Fail closed: log + re-raise so the Postgres row is never written. + logger.exception( + "Auth register_resource failed for agent_api_key; aborting create", + extra={ + "api_key_id": api_key_id, + "agent_id": agent_id, + "account_id": account_id, + "error_type": type(exc).__name__, + }, + ) + raise + + async def _deregister_api_key_from_auth( + self, *, api_key_id: str, account_id: str | None + ) -> None: + """Best-effort deregistration of an api_key's auth tuples on delete. + + ``deregister_resource`` removes the resource and all of its + relationships (owner, parent, grantees) atomically. Always invoked; + agentex-auth's per-account routing (#353) decides whether the call + targets Spark or the legacy backend. Failures are logged but do not + block the delete. + """ + try: + await self.authorization_service.deregister_resource( + resource=AgentexResource.api_key(api_key_id), + ) + except Exception as exc: + # Best-effort: log and continue. Postgres row already deleted. + logger.warning( + "Auth deregister failed for agent_api_key; tuple may be orphaned", + extra={ + "api_key_id": api_key_id, + "account_id": account_id, + "error_type": type(exc).__name__, + }, + exc_info=True, + ) + async def get(self, id: str) -> AgentAPIKeyEntity: return await self.agent_api_key_repo.get(id=id) @@ -123,22 +213,54 @@ async def get_external_by_agent_id_and_key( agent_id=agent_id, api_key=api_key ) - async def delete(self, id: str) -> None: - return await self.agent_api_key_repo.delete(id=id) + async def delete(self, id: str, account_id: str | None = None) -> None: + # Pre-fetch so we skip deregister when the row never existed, matching + # the delete_by_* methods. + try: + existing = await self.agent_api_key_repo.get(id=id) + except ItemDoesNotExist: + existing = None + await self.agent_api_key_repo.delete(id=id) + if existing is not None: + await self._deregister_api_key_from_auth( + api_key_id=id, account_id=account_id + ) async def delete_by_agent_id_and_key_name( - self, agent_id: str, key_name: str, api_key_type: AgentAPIKeyType + self, + agent_id: str, + key_name: str, + api_key_type: AgentAPIKeyType, + account_id: str | None = None, ) -> None: - return await self.agent_api_key_repo.delete_by_agent_id_and_key_name( + existing = await self.agent_api_key_repo.get_by_agent_id_and_name( + agent_id=agent_id, name=key_name, api_key_type=api_key_type + ) + await self.agent_api_key_repo.delete_by_agent_id_and_key_name( agent_id=agent_id, key_name=key_name, api_key_type=api_key_type ) + if existing is not None: + await self._deregister_api_key_from_auth( + api_key_id=existing.id, account_id=account_id + ) async def delete_by_agent_name_and_key_name( - self, agent_name: str, key_name: str, api_key_type: AgentAPIKeyType + self, + agent_name: str, + key_name: str, + api_key_type: AgentAPIKeyType, + account_id: str | None = None, ) -> None: - return await self.agent_api_key_repo.delete_by_agent_name_and_key_name( + existing = await self.agent_api_key_repo.get_by_agent_name_and_key_name( + agent_name, key_name, api_key_type + ) + await self.agent_api_key_repo.delete_by_agent_name_and_key_name( agent_name=agent_name, key_name=key_name, api_key_type=api_key_type ) + if existing is not None: + await self._deregister_api_key_from_auth( + api_key_id=existing.id, account_id=account_id + ) async def list( self, agent_id: str, limit: int, page_number: int diff --git a/agentex/tests/integration/fixtures/integration_client.py b/agentex/tests/integration/fixtures/integration_client.py index f1396d57..d07fc4a1 100644 --- a/agentex/tests/integration/fixtures/integration_client.py +++ b/agentex/tests/integration/fixtures/integration_client.py @@ -6,7 +6,7 @@ import asyncio import os import uuid -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pymongo import pytest @@ -398,10 +398,17 @@ def create_agents_use_case(): ) def create_agent_api_keys_use_case(): + noop_authorization_service = Mock() + noop_authorization_service.principal_context = None + noop_authorization_service.grant = AsyncMock(return_value={}) + noop_authorization_service.revoke = AsyncMock(return_value=None) + noop_authorization_service.register_resource = AsyncMock(return_value=None) + noop_authorization_service.deregister_resource = AsyncMock(return_value=None) return AgentAPIKeysUseCase( agent_api_key_repository=isolated_repositories["agent_api_key_repository"], agent_repository=isolated_repositories["agent_repository"], client=isolated_api_key_http_client, # Use mock client for forwarding requests + authorization_service=noop_authorization_service, ) def create_deployment_history_use_case(): diff --git a/agentex/tests/integration/services/__init__.py b/agentex/tests/integration/services/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/agentex/tests/integration/services/test_agent_api_key_service_dual_write.py b/agentex/tests/integration/services/test_agent_api_key_service_dual_write.py new file mode 100644 index 00000000..02b8ecf5 --- /dev/null +++ b/agentex/tests/integration/services/test_agent_api_key_service_dual_write.py @@ -0,0 +1,316 @@ +"""Integration tests for AgentAPIKeysUseCase dual-write to Spark AuthZ. + +These cover the AGX1-272 dual-write path. scale-agentex calls +``register_resource`` / ``deregister_resource`` unconditionally; per-account +routing (Spark vs legacy SGP) is owned by agentex-auth (scaleapi/agentex#353) +so scale-agentex does NOT couple to egp-api-backend's feature-flag service. + +- Create calls register_resource with parent=agent (the parent_agent edge + is load-bearing for the SpiceDB cascade). +- Delete calls deregister_resource after the Postgres row is gone. +- Spark failure prevents row: when register_resource raises, the api_key + is NOT persisted. +- Deregister failure does not block delete: when deregister_resource + raises, the DB delete still completes and the failure is logged. +- No creator → no register: if neither user_id nor service_account_id is + resolvable, the dual-write is a no-op (logged) and the row still lands. + +The tests intentionally mock the repository, authorization service, agent +repository, and HTTP client. The behaviour under test is the call sequencing +inside ``AgentAPIKeysUseCase`` — not Postgres or Spark itself. + +Note on structural divergence from the task PR (AGX1-274): tasks live behind +``AgentTaskService``; agent_api_keys have no service layer, so the dual-write +logic is colocated in ``AgentAPIKeysUseCase``. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from src.adapters.crud_store.exceptions import ItemDoesNotExist +from src.api.schemas.authorization_types import AgentexResource, AgentexResourceType +from src.domain.entities.agent_api_keys import AgentAPIKeyEntity, AgentAPIKeyType +from src.domain.entities.agents import ACPType, AgentEntity, AgentStatus +from src.domain.use_cases.agent_api_keys_use_case import AgentAPIKeysUseCase +from src.utils.ids import orm_id + + +def _principal(user_id: str | None, account_id: str | None) -> SimpleNamespace: + """Minimal stand-in for AgentexAuthPrincipalContext.""" + return SimpleNamespace( + user_id=user_id, service_account_id=None, account_id=account_id + ) + + +def _agent() -> AgentEntity: + agent_id = orm_id() + return AgentEntity( + id=agent_id, + name=f"agent-{agent_id[:8]}", + description="dual-write test agent", + status=AgentStatus.READY, + acp_type=ACPType.SYNC, + acp_url="http://test-acp", + ) + + +def _stub_api_key(id: str) -> AgentAPIKeyEntity: + """Minimal entity stand-in for the repo.get default in _build_use_case.""" + return AgentAPIKeyEntity( + id=id, + name="stub", + agent_id=orm_id(), + api_key_type=AgentAPIKeyType.EXTERNAL, + api_key="stub", + ) + + +def _build_use_case( + *, + principal: SimpleNamespace | None, + register_resource: AsyncMock | None = None, + deregister_resource: AsyncMock | None = None, + agent: AgentEntity | None = None, + create_raises: Exception | None = None, + monkeypatch: pytest.MonkeyPatch, +) -> tuple[AgentAPIKeysUseCase, Mock, AsyncMock, AsyncMock]: + sample_agent = agent or _agent() + + agent_repository = Mock() + agent_repository.get = AsyncMock(return_value=sample_agent) + + agent_api_key_repository = Mock() + if create_raises is None: + agent_api_key_repository.create = AsyncMock(side_effect=lambda item: item) + else: + agent_api_key_repository.create = AsyncMock(side_effect=create_raises) + agent_api_key_repository.delete = AsyncMock(return_value=None) + # Default get() returns a sentinel "exists" so delete() flows through the + # deregister path; tests covering "row doesn't exist" override this. + agent_api_key_repository.get = AsyncMock(side_effect=lambda id: _stub_api_key(id)) + agent_api_key_repository.get_by_agent_id_and_name = AsyncMock(return_value=None) + agent_api_key_repository.get_by_agent_name_and_key_name = AsyncMock( + return_value=None + ) + agent_api_key_repository.delete_by_agent_id_and_key_name = AsyncMock( + return_value=None + ) + agent_api_key_repository.delete_by_agent_name_and_key_name = AsyncMock( + return_value=None + ) + + authorization_service = Mock() + authorization_service.principal_context = principal + authorization_service.register_resource = register_resource or AsyncMock( + return_value=None + ) + authorization_service.deregister_resource = deregister_resource or AsyncMock( + return_value=None + ) + + # Patch env var lookup inside UseCase __init__ so we don't depend on real + # env configuration to instantiate. + monkeypatch.setenv("AGENTEX_AUTH_URL", "") + monkeypatch.setenv("ENVIRONMENT", "test") + monkeypatch.setenv("WEBHOOK_REQUEST_TIMEOUT", "10") + + use_case = AgentAPIKeysUseCase( + agent_api_key_repository=agent_api_key_repository, + agent_repository=agent_repository, + client=Mock(), + authorization_service=authorization_service, + ) + return ( + use_case, + agent_api_key_repository, + authorization_service.register_resource, + authorization_service.deregister_resource, + ) + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_create_api_key_calls_register_resource_with_parent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = _agent() + use_case, repo, register, _ = _build_use_case( + principal=_principal(user_id="user-A", account_id="acct-1"), + agent=agent, + monkeypatch=monkeypatch, + ) + + api_key = await use_case.create( + name="k1", + agent_id=agent.id, + api_key_type=AgentAPIKeyType.EXTERNAL, + api_key="secret", + account_id="acct-1", + ) + + register.assert_awaited_once() + registered_resource: AgentexResource = register.await_args.kwargs["resource"] + assert registered_resource.type == AgentexResourceType.api_key + assert registered_resource.selector == api_key.id + # parent_agent edge is load-bearing — without it the SpiceDB cascade + # `read = ... & parent_agent->read & ...` fails closed for every reader. + registered_parent: AgentexResource = register.await_args.kwargs["parent"] + assert registered_parent is not None + assert registered_parent.type == AgentexResourceType.agent + assert registered_parent.selector == agent.id + repo.create.assert_awaited_once() + # Sanity: the persisted entity itself; we don't persist creator audit + # columns in OSS scale-agentex (Harvey's review feedback on #248). + assert api_key.id is not None + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_delete_api_key_calls_deregister_resource( + monkeypatch: pytest.MonkeyPatch, +) -> None: + use_case, repo, _, deregister = _build_use_case( + principal=_principal(user_id="user-A", account_id="acct-1"), + monkeypatch=monkeypatch, + ) + + api_key_id = orm_id() + await use_case.delete(id=api_key_id, account_id="acct-1") + + repo.delete.assert_awaited_once_with(id=api_key_id) + deregister.assert_awaited_once() + deregistered_resource: AgentexResource = deregister.await_args.kwargs["resource"] + assert deregistered_resource.type == AgentexResourceType.api_key + assert deregistered_resource.selector == api_key_id + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_create_api_key_grant_failure_prevents_db_row( + monkeypatch: pytest.MonkeyPatch, +) -> None: + register_resource = AsyncMock(side_effect=RuntimeError("spark unavailable")) + agent = _agent() + use_case, repo, _, _ = _build_use_case( + principal=_principal(user_id="user-A", account_id="acct-1"), + register_resource=register_resource, + agent=agent, + monkeypatch=monkeypatch, + ) + + with pytest.raises(RuntimeError, match="spark unavailable"): + await use_case.create( + name="k1", + agent_id=agent.id, + api_key_type=AgentAPIKeyType.EXTERNAL, + api_key="secret", + account_id="acct-1", + ) + + repo.create.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_delete_api_key_revoke_failure_does_not_block_delete( + monkeypatch: pytest.MonkeyPatch, +) -> None: + deregister = AsyncMock(side_effect=RuntimeError("spark unavailable")) + use_case, repo, _, deregister_ref = _build_use_case( + principal=_principal(user_id="user-A", account_id="acct-1"), + deregister_resource=deregister, + monkeypatch=monkeypatch, + ) + + # Should NOT raise. + await use_case.delete(id=orm_id(), account_id="acct-1") + + repo.delete.assert_awaited_once() + deregister_ref.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_create_api_key_skips_grant_when_no_creator_resolvable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """If neither user_id nor service_account_id is available on the principal, + the dual-write is a no-op (logged) and the row still lands without a tuple.""" + agent = _agent() + use_case, repo, register, _ = _build_use_case( + principal=_principal(user_id=None, account_id="acct-1"), + agent=agent, + monkeypatch=monkeypatch, + ) + + api_key = await use_case.create( + name="k1", + agent_id=agent.id, + api_key_type=AgentAPIKeyType.EXTERNAL, + api_key="secret", + account_id="acct-1", + ) + + register.assert_not_called() + repo.create.assert_awaited_once() + # Sanity: the row landed even though we skipped the auth-side registration. + assert api_key.id is not None + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_delete_by_agent_id_and_key_name_revokes_existing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = _agent() + existing_id = orm_id() + use_case, repo, _, deregister = _build_use_case( + principal=_principal(user_id="user-A", account_id="acct-1"), + agent=agent, + monkeypatch=monkeypatch, + ) + repo.get_by_agent_id_and_name = AsyncMock( + return_value=AgentAPIKeyEntity( + id=existing_id, + agent_id=agent.id, + name="k1", + api_key_type=AgentAPIKeyType.EXTERNAL, + api_key="secret", + ) + ) + + await use_case.delete_by_agent_id_and_key_name( + agent_id=agent.id, + key_name="k1", + api_key_type=AgentAPIKeyType.EXTERNAL, + account_id="acct-1", + ) + + repo.delete_by_agent_id_and_key_name.assert_awaited_once() + deregister.assert_awaited_once() + deregistered_resource: AgentexResource = deregister.await_args.kwargs["resource"] + assert deregistered_resource.selector == existing_id + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_delete_api_key_skips_deregister_when_row_does_not_exist( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """When the api_key id doesn't exist, the pre-fetch raises and the + deregister call is skipped — matches the delete_by_* pattern and avoids + a wasted auth round-trip on a no-op delete.""" + use_case, repo, _, deregister = _build_use_case( + principal=_principal(user_id="user-A", account_id="acct-1"), + monkeypatch=monkeypatch, + ) + # Override the default "row exists" sentinel. + repo.get = AsyncMock(side_effect=ItemDoesNotExist("not found")) + + await use_case.delete(id=orm_id(), account_id="acct-1") + + repo.delete.assert_awaited_once() + deregister.assert_not_called() diff --git a/agentex/tests/unit/use_cases/test_agents_api_keys_use_case.py b/agentex/tests/unit/use_cases/test_agents_api_keys_use_case.py index 27d3d4d9..c93e16d4 100644 --- a/agentex/tests/unit/use_cases/test_agents_api_keys_use_case.py +++ b/agentex/tests/unit/use_cases/test_agents_api_keys_use_case.py @@ -1,4 +1,4 @@ -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock from uuid import uuid4 import pytest @@ -35,10 +35,17 @@ def agent_api_keys_use_case( agent_api_key_repository, agent_repository, mock_http_client ): """Real AgentAPIKeysUseCase instance with real repositories""" + authorization_service = Mock() + authorization_service.principal_context = None + authorization_service.grant = AsyncMock(return_value={}) + authorization_service.revoke = AsyncMock(return_value=None) + authorization_service.register_resource = AsyncMock(return_value=None) + authorization_service.deregister_resource = AsyncMock(return_value=None) return AgentAPIKeysUseCase( agent_api_key_repository=agent_api_key_repository, agent_repository=agent_repository, client=mock_http_client, + authorization_service=authorization_service, )