Refactor agent authorization to use delegation chain from JWT

Replaces custom JWT parsing with a utility function to extract user_id and delegation chain from the JWT. Updates the agent authorization logic to use the extracted actors for original_callers and improves error handling and logging. Also updates the IdentityClient interface to accept original_callers and adds license headers to jwt.py.

Author: loveyana

veadk/integrations/ve_identity/identity_client.py

@@ -695,6 +695,7 @@ def check_permission(
         principal: Dict[str, str],
         operation: Dict[str, str],
         resource: Dict[str, str],
+        original_callers: Optional[List[Dict[str, str]]] = None,
         namespace: str = "default",
     ) -> bool:
         """Check if the principal has permission to perform the operation on the resource.
@@ -703,6 +704,7 @@ def check_permission(
             principal: Principal information, e.g., {"Type": "User", "Id": "user123"}
             operation: Operation to check, e.g., {"Type": "Action", "Id": "invoke"}
             resource: Resource information, e.g., {"Type": "Agent", "Id": "agent456"}
+            original_callers: Optional list of original callers.
             namespace: Namespace of the resource. Defaults to "default".
 
         Returns:
@@ -721,6 +723,7 @@ def check_permission(
             operation=operation,
             principal=principal,
             resource=resource,
+            original_callers=original_callers,
         )
 
         response: volcenginesdkid.CheckPermissionResponse = (

veadk/tools/builtin_tools/agent_authorization.py

@@ -23,6 +23,7 @@
 from veadk.integrations.ve_identity.identity_client import IdentityClient
 from veadk.integrations.ve_identity.token_manager import get_workload_token
 from veadk.utils.logger import get_logger
+from veadk.utils.jwt import extract_delegation_chain_from_jwt
 
 logger = get_logger(__name__)
 
@@ -31,85 +32,61 @@
 identity_client = IdentityClient(region=region)
 
 
-def _strip_bearer_prefix(token: str) -> str:
-    """Remove 'Bearer ' prefix from token if present.
-    Args:
-        token: Token string that may contain "Bearer " prefix
-    Returns:
-        Token without "Bearer " prefix
-    """
-    return token[7:] if token.lower().startswith("bearer ") else token
-
-
-def _extract_role_id_from_jwt(token: str) -> Optional[str]:
-    """Extract role_id (sub field) from JWT token.
-    Args:
-        token: JWT token string (with or without "Bearer " prefix)
-    Returns:
-        Role ID from sub field, or None if parsing fails
-    """
-    try:
-        # Remove "Bearer " prefix if present
-        token = _strip_bearer_prefix(token)
-
-        # JWT token has 3 parts separated by dots: header.payload.signature
-        parts = token.split(".")
-        if len(parts) != 3:
-            logger.error("Invalid JWT format: expected 3 parts")
-            return None
-
-        # Decode payload (second part)
-        payload_part = parts[1]
-
-        # Add padding for base64url decoding (JWT doesn't use padding)
-        missing_padding = len(payload_part) % 4
-        if missing_padding:
-            payload_part += "=" * (4 - missing_padding)
-
-        # Decode base64 and parse JSON
-        decoded_bytes = base64.urlsafe_b64decode(payload_part)
-        payload = json.loads(decoded_bytes.decode("utf-8"))
-
-        # Extract sub field as role_id
-        return payload.get("act").get("sub")
-
-    except (ValueError, json.JSONDecodeError) as e:
-        logger.error(f"Failed to parse JWT token: {e}")
-        return None
-    except Exception as e:
-        logger.error(f"Unexpected error parsing JWT: {e}")
-        return None
-
-
 async def check_agent_authorization(
     callback_context: CallbackContext,
 ) -> Optional[types.Content]:
-    """Check if the agent is authorized to run using VeIdentity."""
-    user_id = callback_context._invocation_context.user_id
-
+    """Check if the agent is authorized to run using Agent Identity."""
     try:
         workload_token = await get_workload_token(
             tool_context=callback_context, identity_client=identity_client
         )
 
-        # Parse role_id from workload_token
-        role_id = _extract_role_id_from_jwt(workload_token)
+        # Parse user_id and actors from workload_token
+        user_id, actors = extract_delegation_chain_from_jwt(workload_token)
 
-        principal = {"Type": "User", "Id": user_id}
-        operation = {"Type": "Action", "Id": "invoke"}
-        resource = {"Type": "Agent", "Id": role_id}
+        if not user_id:
+            logger.warning("Failed to extract user_id from JWT token")
+            return types.Content(
+                parts=[types.Part(text="Failed to verify agent authorization.")],
+                role="model",
+            )
+
+        if len(actors) == 0:
+            logger.warning("Failed to extract actors from JWT token")
+            return types.Content(
+                parts=[types.Part(text="Failed to verify agent authorization.")],
+                role="model",
+            )
+
+        # The first actor in the chain is the agent itself
+        role_id = actors[0]
+
+        principal = {"Type": "user", "Id": user_id}
+        operation = {"Type": "action", "Id": "invoke"}
+        resource = {"Type": "agent", "Id": role_id}
+        original_callers = [{"Type": "agent", "Id": actor} for actor in actors[1:]]
 
         allowed = identity_client.check_permission(
-            principal=principal, operation=operation, resource=resource
+            principal=principal,
+            operation=operation,
+            resource=resource,
+            original_callers=original_callers,
         )
 
         if allowed:
-            logger.info("Agent is authorized to run.")
+            logger.info(f"Agent {role_id} is authorized to run by user {user_id}.")
             return None
         else:
-            logger.warning("Agent is not authorized to run.")
+            logger.warning(
+                f"Agent {role_id} is not authorized to run by user {user_id}."
+            )
             return types.Content(
-                parts=[types.Part(text="Agent is not authorized to run.")], role="model"
+                parts=[
+                    types.Part(
+                        text=f"Agent {role_id} is not authorized to run by user {user_id}."
+                    )
+                ],
+                role="model",
             )
 
     except Exception as e:

veadk/utils/jwt.py

@@ -1,3 +1,17 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import base64
 import json
 from typing import Optional