api: add agent, agentEndpoint and agentThread grants (#859)

Author: theomonnom

livekit-api/livekit/api/__init__.py

@@ -47,6 +47,9 @@
 from .access_token import (
     InferenceGrants,
     ObservabilityGrants,
+    AgentGrants,
+    AgentEndpointGrants,
+    AgentThreadGrants,
     VideoGrants,
     SIPGrants,
     AccessToken,
@@ -65,6 +68,9 @@
     "connector_service",
     "InferenceGrants",
     "ObservabilityGrants",
+    "AgentGrants",
+    "AgentEndpointGrants",
+    "AgentThreadGrants",
     "VideoGrants",
     "SIPGrants",
     "AccessToken",

livekit-api/livekit/api/access_token.py

@@ -19,7 +19,7 @@
 import jwt
 import dataclasses
 from dataclasses import dataclass
-from typing import Optional, List, Literal
+from typing import Optional, List, Literal, Type, TypeVar
 from google.protobuf.json_format import MessageToDict, ParseDict
 
 from livekit.protocol.room import RoomConfiguration
@@ -92,6 +92,43 @@ class ObservabilityGrants:
     write: bool = False
 
 
+@dataclass
+class AgentGrants:
+    # create/update/delete Cloud Agents
+    admin: bool = False
+    # manage simulations and scenarios for evaluating agents
+    simulation_admin: bool = False
+    # access the project's agent databases
+    database_admin: bool = False
+    # manage the project's agent dispatch queue
+    dispatch_admin: bool = False
+
+
+@dataclass
+class AgentEndpointGrants:
+    # invoke an agent's non-public HTTP endpoints
+    call: bool = False
+    # restrict to one agent; empty grants every agent in the project
+    agent_name: str = ""
+    # restrict to one deployment; empty grants every deployment
+    deployment: str = ""
+
+
+@dataclass
+class AgentThreadGrants:
+    # every thread and scope of the project; the action flags still apply
+    all: bool = False
+    # scope in which threads may be listed and created
+    scope: str = ""
+    # threads that may be read, written and deleted
+    thread_ids: Optional[List[str]] = None
+    list: bool = False
+    create: bool = False
+    write: bool = False
+    read: bool = False
+    delete: bool = False
+
+
 @dataclass
 class Claims:
     identity: str = ""
@@ -102,6 +139,9 @@ class Claims:
     sip: Optional[SIPGrants] = None
     inference: Optional[InferenceGrants] = None
     observability: Optional[ObservabilityGrants] = None
+    agent: Optional[AgentGrants] = None
+    agent_endpoint: Optional[AgentEndpointGrants] = None
+    agent_thread: Optional[AgentThreadGrants] = None
     attributes: Optional[dict[str, str]] = None
     sha256: Optional[str] = None
     room_preset: Optional[str] = None
@@ -162,6 +202,18 @@ def with_inference_grants(self, grants: InferenceGrants) -> "AccessToken":
         self.claims.inference = grants
         return self
 
+    def with_agent_grants(self, grants: AgentGrants) -> "AccessToken":
+        self.claims.agent = grants
+        return self
+
+    def with_agent_endpoint_grants(self, grants: AgentEndpointGrants) -> "AccessToken":
+        self.claims.agent_endpoint = grants
+        return self
+
+    def with_agent_thread_grants(self, grants: AgentThreadGrants) -> "AccessToken":
+        self.claims.agent_thread = grants
+        return self
+
     def with_identity(self, identity: str) -> "AccessToken":
         self.identity = identity
         return self
@@ -279,6 +331,9 @@ def verify(self, token: str, *, verify_signature: bool = True) -> Claims:
             attributes=claims.get("attributes", {}),
             metadata=claims.get("metadata", ""),
             sha256=claims.get("sha256", ""),
+            agent=_parse_grants(claims, "agent", AgentGrants),
+            agent_endpoint=_parse_grants(claims, "agentEndpoint", AgentEndpointGrants),
+            agent_thread=_parse_grants(claims, "agentThread", AgentThreadGrants),
         )
 
         if claims.get("roomPreset"):
@@ -294,6 +349,17 @@ def verify(self, token: str, *, verify_signature: bool = True) -> Claims:
         return grant_claims
 
 
+_G = TypeVar("_G")
+
+
+def _parse_grants(claims: dict, key: str, cls: Type[_G]) -> Optional[_G]:
+    grants = claims.get(key)
+    if grants is None:
+        return None
+    fields = cls.__dataclass_fields__  # type: ignore[attr-defined]
+    return cls(**{camel_to_snake(k): v for k, v in grants.items() if camel_to_snake(k) in fields})
+
+
 def camel_to_snake(t: str) -> str:
     return re.sub(r"(?<!^)(?=[A-Z])", "_", t).lower()
 

tests/api/test_access_token.py

@@ -3,7 +3,15 @@
 
 import jwt
 import pytest
-from livekit.api import AccessToken, TokenVerifier, VideoGrants, SIPGrants
+from livekit.api import (
+    AccessToken,
+    AgentEndpointGrants,
+    AgentGrants,
+    AgentThreadGrants,
+    SIPGrants,
+    TokenVerifier,
+    VideoGrants,
+)
 from livekit.protocol.room import RoomConfiguration
 from livekit.protocol.agent_dispatch import RoomAgentDispatch
 
@@ -37,6 +45,30 @@ def test_verify_token() -> None:
     assert claims.attributes["key2"] == "value2"
 
 
+def test_agent_grants() -> None:
+    agent = AgentGrants(database_admin=True)
+    endpoint = AgentEndpointGrants(call=True, agent_name="my-agent")
+    thread = AgentThreadGrants(thread_ids=["t1"], read=True, write=True)
+
+    token = (
+        AccessToken(TEST_API_KEY, TEST_API_SECRET)
+        .with_agent_grants(agent)
+        .with_agent_endpoint_grants(endpoint)
+        .with_agent_thread_grants(thread)
+        .to_jwt()
+    )
+
+    decoded = jwt.decode(token, options={"verify_signature": False})
+    assert decoded["agentThread"]["threadIds"] == ["t1"]
+    assert decoded["agentEndpoint"]["agentName"] == "my-agent"
+    assert decoded["agent"]["databaseAdmin"] is True
+
+    claims = TokenVerifier(TEST_API_KEY, TEST_API_SECRET).verify(token)
+    assert claims.agent == agent
+    assert claims.agent_endpoint == endpoint
+    assert claims.agent_thread == thread
+
+
 def test_agent_config() -> None:
     token = (
         AccessToken(TEST_API_KEY, TEST_API_SECRET)