feat: fine grained access control policy

This allows a set of rules to be defined for determining access to resources.

Signed-off-by: Gordon Sim <gsim@redhat.com>

llama_stack/providers/inline/agents/meta_reference/__init__.py

@@ -6,12 +6,12 @@
 
 from typing import Any
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.distribution.datatypes import AccessRule, Api
 
 from .config import MetaReferenceAgentsImplConfig
 
 
-async def get_provider_impl(config: MetaReferenceAgentsImplConfig, deps: dict[Api, Any]):
+async def get_provider_impl(config: MetaReferenceAgentsImplConfig, deps: dict[Api, Any], policy: list[AccessRule]):
     from .agents import MetaReferenceAgentsImpl
 
     impl = MetaReferenceAgentsImpl(
@@ -21,6 +21,7 @@ async def get_provider_impl(config: MetaReferenceAgentsImplConfig, deps: dict[Ap
         deps[Api.safety],
         deps[Api.tool_runtime],
         deps[Api.tool_groups],
+        policy,
     )
     await impl.initialize()
     return impl

llama_stack/providers/inline/agents/meta_reference/agent_instance.py

@@ -60,6 +60,7 @@ from llama_stack.apis.inference import (
 from llama_stack.apis.safety import Safety
 from llama_stack.apis.tools import ToolGroups, ToolInvocationResult, ToolRuntime
 from llama_stack.apis.vector_io import VectorIO
+from llama_stack.distribution.datatypes import AccessRule
 from llama_stack.log import get_logger
 from llama_stack.models.llama.datatypes import (
     BuiltinTool,
@@ -96,13 +97,14 @@ class ChatAgent(ShieldRunnerMixin):
         vector_io_api: VectorIO,
         persistence_store: KVStore,
         created_at: str,
+        policy: list[AccessRule],
     ):
         self.agent_id = agent_id
         self.agent_config = agent_config
         self.inference_api = inference_api
         self.safety_api = safety_api
         self.vector_io_api = vector_io_api
-        self.storage = AgentPersistence(agent_id, persistence_store)
+        self.storage = AgentPersistence(agent_id, persistence_store, policy)
         self.tool_runtime_api = tool_runtime_api
         self.tool_groups_api = tool_groups_api
         self.created_at = created_at

llama_stack/providers/inline/agents/meta_reference/agents.py

@@ -40,6 +40,7 @@ from llama_stack.apis.inference import (
 from llama_stack.apis.safety import Safety
 from llama_stack.apis.tools import ToolGroups, ToolRuntime
 from llama_stack.apis.vector_io import VectorIO
+from llama_stack.distribution.datatypes import AccessRule
 from llama_stack.providers.utils.kvstore import InmemoryKVStoreImpl, kvstore_impl
 from llama_stack.providers.utils.pagination import paginate_records
 from llama_stack.providers.utils.responses.responses_store import ResponsesStore
@@ -61,6 +62,7 @@ class MetaReferenceAgentsImpl(Agents):
         safety_api: Safety,
         tool_runtime_api: ToolRuntime,
         tool_groups_api: ToolGroups,
+        policy: list[AccessRule],
     ):
         self.config = config
         self.inference_api = inference_api
@@ -71,6 +73,7 @@ class MetaReferenceAgentsImpl(Agents):
 
         self.in_memory_store = InmemoryKVStoreImpl()
         self.openai_responses_impl: OpenAIResponsesImpl | None = None
+        self.policy = policy
 
     async def initialize(self) -> None:
         self.persistence_store = await kvstore_impl(self.config.persistence_store)
@@ -129,6 +132,7 @@ class MetaReferenceAgentsImpl(Agents):
                 self.persistence_store if agent_info.enable_session_persistence else self.in_memory_store
             ),
             created_at=agent_info.created_at,
+            policy=self.policy,
         )
 
     async def create_agent_session(

llama_stack/providers/inline/agents/meta_reference/persistence.py

@@ -10,9 +10,9 @@ import uuid
 from datetime import datetime, timezone
 
 from llama_stack.apis.agents import AgentConfig, Session, ToolExecutionStep, Turn
-from llama_stack.distribution.access_control import check_access
-from llama_stack.distribution.datatypes import AccessAttributes
-from llama_stack.distribution.request_headers import get_auth_attributes
+from llama_stack.distribution.access_control.access_control import AccessDeniedError, is_action_allowed
+from llama_stack.distribution.access_control.datatypes import AccessAttributes, AccessRule
+from llama_stack.distribution.request_headers import get_authenticated_user
 from llama_stack.providers.utils.kvstore import KVStore
 
 log = logging.getLogger(__name__)
@@ -23,6 +23,8 @@ class AgentSessionInfo(Session):
     vector_db_id: str | None = None
     started_at: datetime
     access_attributes: AccessAttributes | None = None
+    identifier: str | None = None
+    type: str = "session"
 
 
 class AgentInfo(AgentConfig):
@@ -30,15 +32,17 @@ class AgentInfo(AgentConfig):
 
 
 class AgentPersistence:
-    def __init__(self, agent_id: str, kvstore: KVStore):
+    def __init__(self, agent_id: str, kvstore: KVStore, policy: list[AccessRule]):
         self.agent_id = agent_id
         self.kvstore = kvstore
+        self.policy = policy
 
     async def create_session(self, name: str) -> str:
         session_id = str(uuid.uuid4())
 
         # Get current user's auth attributes for new sessions
-        auth_attributes = get_auth_attributes()
+        user = get_authenticated_user()
+        auth_attributes = user and user.attributes
         access_attributes = AccessAttributes(**auth_attributes) if auth_attributes else None
 
         session_info = AgentSessionInfo(
@@ -47,7 +51,10 @@ class AgentPersistence:
             started_at=datetime.now(timezone.utc),
             access_attributes=access_attributes,
             turns=[],
+            identifier=name,  # should this be qualified in any way?
         )
+        if not is_action_allowed(self.policy, "create", session_info, user):
+            raise AccessDeniedError()
 
         await self.kvstore.set(
             key=f"session:{self.agent_id}:{session_id}",
@@ -76,7 +83,7 @@ class AgentPersistence:
         if not hasattr(session_info, "access_attributes"):
             return True
 
-        return check_access(session_info.session_id, session_info.access_attributes, get_auth_attributes())
+        return is_action_allowed(self.policy, "read", session_info, get_authenticated_user())
 
     async def get_session_if_accessible(self, session_id: str) -> AgentSessionInfo | None:
         """Get session info if the user has access to it. For internal use by sub-session methods."""