fix(authz): avoid stale authorization checks

Author: deepthi-rao-scale

agentex/src/api/authentication_cache.py

@@ -85,6 +85,14 @@ async def clear(self) -> None:
         async with self._lock:
             self.cache.clear()
 
+    async def delete_by_prefix(self, prefix: str) -> int:
+        """Delete cache entries whose keys start with the given prefix."""
+        async with self._lock:
+            keys = [key for key in self.cache if key.startswith(prefix)]
+            for key in keys:
+                del self.cache[key]
+            return len(keys)
+
     async def remove_expired(self) -> None:
         """Remove all expired entries from cache."""
         async with self._lock:
@@ -258,17 +266,20 @@ def _contains_api_key(principal_context: Any) -> bool:
     # Authorization Check Cache Methods (Async)
 
     @staticmethod
-    def _create_authorization_cache_key(
+    def _create_authorization_resource_cache_prefix(
         resource_type: str,
         resource_selector: str,
-        operation: str,
-        principal_context: Any,
     ) -> str:
-        """
-        Create a cache key for authorization checks.
+        resource_key = AuthenticationCache._hash_dict(
+            {
+                "resource_type": resource_type,
+                "resource_selector": resource_selector,
+            }
+        )
+        return f"authz:{resource_key}:"
 
-        Combines resource info, operation, and principal context into a unique key.
-        """
+    @staticmethod
+    def _authorization_principal_key_data(principal_context: Any) -> dict[str, Any]:
         # Extract relevant fields from principal context for cache key
         principal_key_data = {}
         if principal_context:
@@ -301,15 +312,48 @@ def _create_authorization_cache_key(
                     "context_hash": AuthenticationCache._hash_dict(context_dict)
                 }
 
-        # Create the cache key components
-        cache_data = {
-            "resource_type": resource_type,
-            "resource_selector": resource_selector,
-            "operation": operation,
-            "principal": principal_key_data,
-        }
+        return principal_key_data
+
+    @staticmethod
+    def _create_authorization_principal_cache_key(principal_context: Any) -> str:
+        return AuthenticationCache._hash_dict(
+            AuthenticationCache._authorization_principal_key_data(principal_context)
+        )
 
-        return f"authz:{AuthenticationCache._hash_dict(cache_data)}"
+    @staticmethod
+    def _create_authorization_resource_principal_cache_prefix(
+        resource_type: str,
+        resource_selector: str,
+        principal_context: Any,
+    ) -> str:
+        return (
+            AuthenticationCache._create_authorization_resource_cache_prefix(
+                resource_type, resource_selector
+            )
+            + AuthenticationCache._create_authorization_principal_cache_key(
+                principal_context
+            )
+            + ":"
+        )
+
+    @staticmethod
+    def _create_authorization_cache_key(
+        resource_type: str,
+        resource_selector: str,
+        operation: str,
+        principal_context: Any,
+    ) -> str:
+        """
+        Create a cache key for authorization checks.
+
+        Combines resource info, operation, and principal context into a unique key.
+        """
+        return (
+            AuthenticationCache._create_authorization_resource_principal_cache_prefix(
+                resource_type, resource_selector, principal_context
+            )
+            + AuthenticationCache._hash_dict({"operation": operation})
+        )
 
     async def get_authorization_check(
         self,
@@ -319,6 +363,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 +389,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 +414,24 @@ 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_for_resource_principal(
+        self,
+        resource_type: str,
+        resource_selector: str,
+        principal_context: Any,
+    ) -> None:
+        """Clear cached authorization check results for one resource/principal."""
+        prefix = self._create_authorization_resource_principal_cache_prefix(
+            resource_type, resource_selector, principal_context
+        )
+        deleted = await self.authorization_check_cache.delete_by_prefix(prefix)
+        logger.info(
+            "Authorization check cache cleared for %s:%s matched_entries=%d",
+            resource_type,
+            resource_selector,
+            deleted,
+        )
+
     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,16 @@ def _bypass(self) -> bool:
     def is_enabled(self) -> bool:
         return self.enabled
 
+    async def _clear_authorization_cache(
+        self, resource: AgentexResource, principal_context
+    ) -> None:
+        auth_cache = await get_auth_cache()
+        await auth_cache.clear_authorization_checks_for_resource_principal(
+            resource_type=str(resource.type),
+            resource_selector=resource.selector,
+            principal_context=principal_context,
+        )
+
     async def grant(
         self, resource: AgentexResource, *, commit: bool = True, principal_context=...
     ) -> None:
@@ -54,13 +64,17 @@ async def grant(
             resource.type,
             resource.selector,
         )
-        result = await self.gateway.grant(
+        effective_principal = (
             principal_context
             if principal_context is not ...
-            else self.principal_context,
+            else self.principal_context
+        )
+        result = await self.gateway.grant(
+            effective_principal,
             resource,
             AuthorizedOperationType.create,
         )
+        await self._clear_authorization_cache(resource, effective_principal)
         return result
 
     async def revoke(
@@ -77,16 +91,20 @@ async def revoke(
             resource.selector,
         )
 
-        result = await self.gateway.revoke(
+        effective_principal = (
             principal_context
             if principal_context is not ...
-            else self.principal_context,
+            else self.principal_context
+        )
+        result = await self.gateway.revoke(
+            effective_principal,
             resource,
             AuthorizedOperationType.delete,
         )
         logger.info(
             f"Revoked {AuthorizedOperationType.delete} permission on {resource.type}:{resource.selector}"
         )
+        await self._clear_authorization_cache(resource, effective_principal)
         return result
 
     async def check(
@@ -214,6 +232,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(resource, effective_principal)
 
     async def deregister_resource(
         self,
@@ -237,6 +256,7 @@ async def deregister_resource(
             resource.selector,
         )
         await self.gateway.deregister_resource(effective_principal, resource)
+        await self._clear_authorization_cache(resource, effective_principal)
 
 
 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
+    )