fix(authz): skip cached checks for api key principals

Author: deepthi-rao-scale

agentex/src/api/authentication_cache.py

@@ -319,6 +319,12 @@ async def get_authorization_check(
         principal_context: Any,
     ) -> bool | None:
         """Get cached authorization check result."""
+        if self._contains_api_key(principal_context):
+            logger.debug(
+                "Skipping authorization check cache lookup for API key principal"
+            )
+            return None
+
         cache_key = self._create_authorization_cache_key(
             resource_type, resource_selector, operation, principal_context
         )
@@ -339,6 +345,12 @@ async def set_authorization_check(
         allowed: bool,
     ) -> None:
         """Cache authorization check result."""
+        if self._contains_api_key(principal_context):
+            logger.debug(
+                "Skipping authorization check cache write for API key principal"
+            )
+            return
+
         cache_key = self._create_authorization_cache_key(
             resource_type, resource_selector, operation, principal_context
         )
@@ -358,6 +370,11 @@ async def clear_all(self) -> None:
         await self.authorization_check_cache.clear()
         logger.info("All authentication and authorization caches cleared")
 
+    async def clear_authorization_checks(self) -> None:
+        """Clear cached authorization check results."""
+        await self.authorization_check_cache.clear()
+        logger.info("Authorization check cache cleared")
+
     async def cleanup_expired(self) -> None:
         """Remove expired entries from all caches."""
         await self.agent_identity_cache.remove_expired()

agentex/src/domain/services/authorization_service.py

@@ -39,6 +39,10 @@ def _bypass(self) -> bool:
     def is_enabled(self) -> bool:
         return self.enabled
 
+    async def _clear_authorization_cache(self) -> None:
+        auth_cache = await get_auth_cache()
+        await auth_cache.clear_authorization_checks()
+
     async def grant(
         self, resource: AgentexResource, *, commit: bool = True, principal_context=...
     ) -> None:
@@ -61,6 +65,7 @@ async def grant(
             resource,
             AuthorizedOperationType.create,
         )
+        await self._clear_authorization_cache()
         return result
 
     async def revoke(
@@ -87,6 +92,7 @@ async def revoke(
         logger.info(
             f"Revoked {AuthorizedOperationType.delete} permission on {resource.type}:{resource.selector}"
         )
+        await self._clear_authorization_cache()
         return result
 
     async def check(
@@ -214,6 +220,7 @@ async def register_resource(
             f"{parent.type}:{parent.selector}" if parent is not None else None,
         )
         await self.gateway.register_resource(effective_principal, resource, parent)
+        await self._clear_authorization_cache()
 
     async def deregister_resource(
         self,
@@ -237,6 +244,7 @@ async def deregister_resource(
             resource.selector,
         )
         await self.gateway.deregister_resource(effective_principal, resource)
+        await self._clear_authorization_cache()
 
 
 DAuthorizationService = Annotated[AuthorizationService, Depends(AuthorizationService)]

agentex/tests/unit/api/test_authentication_cache_metrics.py

@@ -102,3 +102,53 @@ async def test_auth_gateway_response_with_api_key_principal_is_not_cached():
     await cache.set_auth_gateway_response(headers, principal)
 
     assert await cache.get_auth_gateway_response(headers) is None
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_authorization_check_with_api_key_principal_is_not_cached():
+    cache = AuthenticationCache()
+    principal = {"user_id": "user-1", "account_id": "acct-1", "api_key": "secret-key"}
+
+    await cache.set_authorization_check(
+        resource_type="agent",
+        resource_selector="agent-1",
+        operation="execute",
+        principal_context=principal,
+        allowed=True,
+    )
+
+    assert (
+        await cache.get_authorization_check(
+            resource_type="agent",
+            resource_selector="agent-1",
+            operation="execute",
+            principal_context=principal,
+        )
+        is None
+    )
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_authorization_check_without_api_key_principal_is_cached():
+    cache = AuthenticationCache()
+    principal = {"user_id": "user-1", "account_id": "acct-1"}
+
+    await cache.set_authorization_check(
+        resource_type="agent",
+        resource_selector="agent-1",
+        operation="read",
+        principal_context=principal,
+        allowed=True,
+    )
+
+    assert (
+        await cache.get_authorization_check(
+            resource_type="agent",
+            resource_selector="agent-1",
+            operation="read",
+            principal_context=principal,
+        )
+        is True
+    )

agentex/tests/unit/services/test_authorization_service_cache.py

@@ -0,0 +1,89 @@
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+from src.api.authentication_cache import reset_auth_cache
+from src.api.schemas.authorization_types import AgentexResource, AuthorizedOperationType
+from src.domain.services.authorization_service import AuthorizationService
+
+
+def _request_with_principal(principal_context):
+    return SimpleNamespace(
+        state=SimpleNamespace(
+            principal_context=principal_context,
+            agent_identity=None,
+        )
+    )
+
+
+def _service(principal_context, gateway):
+    return AuthorizationService(
+        enabled=True,
+        gateway=gateway,
+        request=_request_with_principal(principal_context),
+    )
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_api_key_principal_authorization_check_calls_gateway_each_time():
+    await reset_auth_cache()
+    try:
+        gateway = AsyncMock()
+        gateway.check.return_value = True
+        service = _service(
+            {"user_id": "user-1", "account_id": "acct-1", "api_key": "secret-key"},
+            gateway,
+        )
+        resource = AgentexResource.agent("agent-1")
+
+        assert await service.check(resource, AuthorizedOperationType.execute) is True
+        assert await service.check(resource, AuthorizedOperationType.execute) is True
+
+        assert gateway.check.await_count == 2
+    finally:
+        await reset_auth_cache()
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_non_api_key_principal_authorization_check_uses_cache():
+    await reset_auth_cache()
+    try:
+        gateway = AsyncMock()
+        gateway.check.return_value = True
+        service = _service({"user_id": "user-1", "account_id": "acct-1"}, gateway)
+        resource = AgentexResource.agent("agent-1")
+
+        assert await service.check(resource, AuthorizedOperationType.read) is True
+        assert await service.check(resource, AuthorizedOperationType.read) is True
+
+        assert gateway.check.await_count == 1
+    finally:
+        await reset_auth_cache()
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "mutation",
+    ["grant", "revoke", "register_resource", "deregister_resource"],
+)
+async def test_authorization_mutations_clear_cached_authorization_checks(mutation):
+    await reset_auth_cache()
+    try:
+        gateway = AsyncMock()
+        gateway.check.return_value = True
+        service = _service({"user_id": "user-1", "account_id": "acct-1"}, gateway)
+        resource = AgentexResource.agent("agent-1")
+
+        assert await service.check(resource, AuthorizedOperationType.read) is True
+        assert await service.check(resource, AuthorizedOperationType.read) is True
+        assert gateway.check.await_count == 1
+
+        await getattr(service, mutation)(resource)
+
+        assert await service.check(resource, AuthorizedOperationType.read) is True
+        assert gateway.check.await_count == 2
+    finally:
+        await reset_auth_cache()