feat(server): add attribute based access control for resources

llama_stack/distribution/request_headers.py

@@ -7,21 +7,29 @@
 import contextvars
 import json
 import logging
-from typing import Any, ContextManager, Dict, Optional
+from typing import Any, ContextManager, Dict, List, Optional
 
 from .utils.dynamic import instantiate_class_type
 
 log = logging.getLogger(__name__)
 
-# Context variable for request provider data
+# Context variable for request provider data and auth attributes
 PROVIDER_DATA_VAR = contextvars.ContextVar("provider_data", default=None)
 
 
 class RequestProviderDataContext(ContextManager):
     """Context manager for request provider data"""
 
-    def __init__(self, provider_data: Optional[Dict[str, Any]] = None):
-        self.provider_data = provider_data
+    def __init__(
+        self, provider_data: Optional[Dict[str, Any]] = None, auth_attributes: Optional[Dict[str, List[str]]] = None
+    ):
+        # Initialize with either provider_data or create a new dict
+        self.provider_data = provider_data or {}
+
+        # Add auth attributes under a special key if provided
+        if auth_attributes:
+            self.provider_data["__auth_attributes"] = auth_attributes
+
         self.token = None
 
     def __enter__(self):
@@ -80,7 +88,17 @@ def parse_request_provider_data(headers: Dict[str, str]) -> Optional[Dict[str, A
         return None
 
 
-def request_provider_data_context(headers: Dict[str, str]) -> ContextManager:
-    """Context manager that sets request provider data from headers for the duration of the context"""
+def request_provider_data_context(
+    headers: Dict[str, str], auth_attributes: Optional[Dict[str, List[str]]] = None
+) -> ContextManager:
+    """Context manager that sets request provider data from headers and auth attributes for the duration of the context"""
     provider_data = parse_request_provider_data(headers)
-    return RequestProviderDataContext(provider_data)
+    return RequestProviderDataContext(provider_data, auth_attributes)
+
+
+def get_auth_attributes() -> Optional[Dict[str, List[str]]]:
+    """Helper to retrieve auth attributes from the provider data context"""
+    provider_data = PROVIDER_DATA_VAR.get()
+    if not provider_data:
+        return None
+    return provider_data.get("__auth_attributes")