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
         )

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,63 @@
+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()