veadk-python 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

veadk/a2a/registry_client.py

@@ -22,6 +22,7 @@ import os
 import re
 import time
 import uuid
+from collections.abc import Mapping
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from typing import Any
@@ -30,6 +31,7 @@ from urllib.parse import quote, urlparse, urlunparse
 import requests
 
 from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
+from veadk.utils.auth import VE_TIP_TOKEN_HEADER
 
 DEFAULT_ENDPOINT = "http://volcengineapi.byted.org/"
 DEFAULT_VERSION = "2025-10-30"
@@ -66,6 +68,8 @@ class AgentKitA2ARegistryConfig:
     top_k: int = DEFAULT_TOP_K
     timeout_ms: int = DEFAULT_TIMEOUT_MS
     poll_interval_ms: int = DEFAULT_POLL_INTERVAL_MS
+    upstream_tip_token: str = ""
+    upstream_authorization: str = ""
 
 
 @dataclass(frozen=True)
@@ -97,9 +101,41 @@ def registry_config_from_env() -> AgentKitA2ARegistryConfig:
         poll_interval_ms=_int_env(
             "REGISTRY_POLL_INTERVAL_MS", DEFAULT_POLL_INTERVAL_MS, minimum=100
         ),
+        upstream_tip_token=_first_env(
+            [
+                "REGISTRY_UPSTREAM_TIP_TOKEN",
+                "AGENTKIT_UPSTREAM_TIP_TOKEN",
+                "A2A_REGISTRY_UPSTREAM_TIP_TOKEN",
+                "VE_TIP_TOKEN",
+                "X_VE_TIP_TOKEN",
+                "TIP_TOKEN",
+            ]
+        ),
     )
 
 
+def registry_tip_token_from_headers(headers: Mapping[str, str]) -> str:
+    """Extract an inbound TIP token from HTTP headers."""
+
+    normalized = {str(key).lower(): str(value) for key, value in headers.items()}
+    return normalized.get(VE_TIP_TOKEN_HEADER.lower(), "").strip()
+
+
+def registry_authorization_from_headers(headers: Mapping[str, str]) -> str:
+    """Extract a bearer Authorization header for downstream A2A calls."""
+
+    normalized = {str(key).lower(): str(value) for key, value in headers.items()}
+    for header_name in (
+        "authorization",
+        "x-forwarded-authorization",
+        "x-original-authorization",
+    ):
+        value = _normalize_bearer_authorization(normalized.get(header_name, ""))
+        if value:
+            return value
+    return ""
+
+
 def search_agent_cards(
     prompt: str,
     top_k: int | None = None,
@@ -250,6 +286,12 @@ def _resolve_config(
                 or DEFAULT_POLL_INTERVAL_MS
             ),
         ),
+        upstream_tip_token=(
+            config.upstream_tip_token or env_config.upstream_tip_token
+        ).strip(),
+        upstream_authorization=_normalize_bearer_authorization(
+            config.upstream_authorization or env_config.upstream_authorization
+        ),
     )
 
 
@@ -534,15 +576,25 @@ def _send_message(
     if task_id:
         message["taskId"] = task_id
 
+    headers: dict[str, str] = {}
     try:
+        headers = _agent_auth_headers(card, config)
         return _a2a_jsonrpc(
             card["url"],
             "message/send",
             {"message": message, "configuration": {"blocking": False}},
-            _agent_auth_headers(card, config),
+            headers,
             config,
         )
     except RegistryError as exc:
+        if _should_retry_with_m2m(exc, card, config, headers):
+            return _a2a_jsonrpc(
+                card["url"],
+                "message/send",
+                {"message": message, "configuration": {"blocking": False}},
+                _agent_auth_headers(card, config, force_m2m=True),
+                config,
+            )
         if exc.code in {
             "A2A_HTTP_FAILED",
             "A2A_RESPONSE_PARSE_FAILED",
@@ -563,13 +615,25 @@ def _poll_card(
     started: float | None = None,
 ) -> dict[str, Any]:
     started = started or time.monotonic()
-    a2a_result = _a2a_jsonrpc(
-        card["url"],
-        "tasks/get",
-        {"id": task_id.strip(), "historyLength": max(0, int(history_length))},
-        _agent_auth_headers(card, config),
-        config,
-    )
+    headers = _agent_auth_headers(card, config)
+    try:
+        a2a_result = _a2a_jsonrpc(
+            card["url"],
+            "tasks/get",
+            {"id": task_id.strip(), "historyLength": max(0, int(history_length))},
+            headers,
+            config,
+        )
+    except RegistryError as exc:
+        if not _should_retry_with_m2m(exc, card, config, headers):
+            raise
+        a2a_result = _a2a_jsonrpc(
+            card["url"],
+            "tasks/get",
+            {"id": task_id.strip(), "historyLength": max(0, int(history_length))},
+            _agent_auth_headers(card, config, force_m2m=True),
+            config,
+        )
     state = _task_state(a2a_result)
     is_terminal = state in TERMINAL_STATES
     payload: dict[str, Any] = {
@@ -815,11 +879,18 @@ def _sanitize_get_agent_result(
 
 
 def _agent_auth_headers(
-    card: dict[str, Any], config: AgentKitA2ARegistryConfig | None = None
+    card: dict[str, Any],
+    config: AgentKitA2ARegistryConfig | None = None,
+    *,
+    force_m2m: bool = False,
 ) -> dict[str, str]:
+    resolved_config = _resolve_config(config)
     security = card.get("security") or []
     schemes = card.get("securitySchemes") or {}
     headers: dict[str, str] = {}
+    upstream_authorization = _normalize_bearer_authorization(
+        resolved_config.upstream_authorization
+    )
 
     for requirement in security:
         if not isinstance(requirement, dict):
@@ -837,9 +908,17 @@ def _agent_auth_headers(
                 if isinstance(token, str) and token:
                     headers[header_name] = token
             elif scheme_type == "oauth2":
-                headers["Authorization"] = "Bearer " + _oauth2_client_credentials_token(
-                    scheme, _resolve_config(config)
-                )
+                if upstream_authorization and not force_m2m:
+                    headers["Authorization"] = upstream_authorization
+                else:
+                    headers["Authorization"] = (
+                        "Bearer "
+                        + _oauth2_client_credentials_token(scheme, resolved_config)
+                    )
+
+    tip_token = resolved_config.upstream_tip_token
+    if tip_token:
+        headers[VE_TIP_TOKEN_HEADER] = tip_token
 
     if security and not headers:
         raise RegistryError(
@@ -849,6 +928,38 @@ def _agent_auth_headers(
     return headers
 
 
+def _should_retry_with_m2m(
+    exc: RegistryError,
+    card: dict[str, Any],
+    config: AgentKitA2ARegistryConfig | None,
+    headers: dict[str, str],
+) -> bool:
+    if exc.code != "A2A_HTTP_FAILED" or exc.diagnostics.get("status_code") != 401:
+        return False
+
+    resolved_config = _resolve_config(config)
+    upstream_authorization = _normalize_bearer_authorization(
+        resolved_config.upstream_authorization
+    )
+    if not upstream_authorization:
+        return False
+    if headers.get("Authorization") != upstream_authorization:
+        return False
+    return _has_oauth2_security(card)
+
+
+def _has_oauth2_security(card: dict[str, Any]) -> bool:
+    schemes = card.get("securitySchemes") or {}
+    for requirement in card.get("security") or []:
+        if not isinstance(requirement, dict):
+            continue
+        for scheme_name in requirement:
+            scheme = schemes.get(scheme_name) or {}
+            if str(scheme.get("type") or "").lower() == "oauth2":
+                return True
+    return False
+
+
 def _oauth2_client_credentials_token(
     scheme: dict[str, Any], config: AgentKitA2ARegistryConfig
 ) -> str:
@@ -957,6 +1068,16 @@ def _clean_config_url(value: Any) -> str:
     return cleaned
 
 
+def _normalize_bearer_authorization(value: Any) -> str:
+    cleaned = _clean_config_url(value)
+    if not cleaned:
+        return ""
+    parts = cleaned.split(None, 1)
+    if len(parts) == 2 and parts[0].lower() == "bearer" and parts[1].strip():
+        return f"Bearer {parts[1].strip()}"
+    return ""
+
+
 def _user_pool_id_from_token_url(token_url: str) -> str:
     host = urlparse(token_url).netloc
     match = re.match(r"userpool-([^.]+)\.userpool\.auth\.id\.", host)

veadk/cli/_frontend_deploy_iam.py

@@ -0,0 +1,122 @@
+# 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.
+
+"""Ensure the IAM role the VeFaaS-hosted frontend runs as.
+
+Volcengine IAM has no inline role policy, so the role is built as: CreatePolicy
+(custom) + CreateRole (trust ``vefaas``) + AttachRolePolicy. The whole thing is
+idempotent — a fixed role/policy name is reused across re-deploys.
+"""
+
+import json
+
+from veadk.cli._frontend_deploy_policy import (
+    FRONTEND_DEPLOY_POLICY,
+    FRONTEND_DEPLOY_TRUST_POLICY,
+)
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+DEFAULT_ROLE_NAME = "VeADKFrontendServiceRole"
+DEFAULT_POLICY_NAME = "VeADKFrontendPolicy"
+
+
+def _result(resp: dict) -> dict:
+    """Extract ``Result`` from a Volcengine response, raising on an API error."""
+    meta = (resp or {}).get("ResponseMetadata", {}) or {}
+    if meta.get("Error"):
+        raise RuntimeError(meta["Error"].get("Message") or str(meta["Error"]))
+    return (resp or {}).get("Result", {}) or {}
+
+
+def _role_trn(result: dict) -> str | None:
+    """Pull the role TRN out of a GetRole/CreateRole Result (shape varies)."""
+    role = result.get("Role") or result
+    return role.get("Trn") or role.get("trn")
+
+
+def ensure_frontend_role(
+    access_key: str,
+    secret_key: str,
+    role_name: str = DEFAULT_ROLE_NAME,
+    policy_name: str = DEFAULT_POLICY_NAME,
+) -> str:
+    """Get-or-create the frontend's IAM role and return its TRN.
+
+    IAM is a global service, so region is irrelevant here. Safe to call
+    repeatedly: an existing role is reused, and create/attach errors for
+    already-existing resources are tolerated.
+    """
+    from volcengine.iam.IamService import IamService
+
+    svc = IamService()
+    svc.set_ak(access_key)
+    svc.set_sk(secret_key)
+
+    # Reuse an existing role if present.
+    try:
+        existing = _result(svc.get_role({"RoleName": role_name}))
+        trn = _role_trn(existing)
+        if trn:
+            logger.info(f"Reusing existing IAM role {role_name} ({trn})")
+            return trn
+    except Exception as e:
+        logger.info(f"Role {role_name} not found, creating it: {e}")
+
+    # Create the custom policy (tolerate "already exists").
+    try:
+        svc.create_policy(
+            {
+                "PolicyName": policy_name,
+                "PolicyDocument": json.dumps(FRONTEND_DEPLOY_POLICY),
+                "Description": "VeADK frontend deploy permissions",
+            }
+        )
+        logger.info(f"Created IAM policy {policy_name}")
+    except Exception as e:
+        logger.info(f"CreatePolicy {policy_name} skipped/failed (may exist): {e}")
+
+    # Create the role with the vefaas trust relationship.
+    created = _result(
+        svc.create_role(
+            {
+                "RoleName": role_name,
+                "TrustPolicyDocument": json.dumps(FRONTEND_DEPLOY_TRUST_POLICY),
+                "Description": "VeADK frontend VeFaaS runtime role",
+            }
+        )
+    )
+
+    # Attach the policy to the role (tolerate "already attached").
+    try:
+        svc.attach_role_policy(
+            {
+                "RoleName": role_name,
+                "PolicyName": policy_name,
+                "PolicyType": "Custom",
+            }
+        )
+        logger.info(f"Attached policy {policy_name} to role {role_name}")
+    except Exception as e:
+        logger.info(f"AttachRolePolicy skipped/failed (may be attached): {e}")
+
+    trn = _role_trn(created)
+    if not trn:
+        # CreateRole didn't echo the TRN — read it back.
+        trn = _role_trn(_result(svc.get_role({"RoleName": role_name})))
+    if not trn:
+        raise RuntimeError(f"Could not resolve TRN for role {role_name}")
+    logger.info(f"Ensured IAM role {role_name} ({trn})")
+    return trn

veadk/cli/_frontend_deploy_policy.py

@@ -0,0 +1,226 @@
+# 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.
+
+"""IAM policy for the VeFaaS-hosted frontend (`veadk frontend deploy`).
+
+The deployed function assumes this role (trust principal ``vefaas``) and reads
+its STS credentials from the injected credential file, so the frontend's
+``/web/*`` endpoints can manage AgentKit runtimes on the user's behalf. The
+document grants the union of permissions the frontend + AgentKit deploy/runtime
+flows exercise (agentkit, vefaas, cr, cp, tos, iam, ark, identity, vikingdb, …).
+"""
+
+# The trust relationship that lets a VeFaaS function assume the role.
+FRONTEND_DEPLOY_TRUST_POLICY: dict = {
+    "Statement": [
+        {
+            "Effect": "Allow",
+            "Action": ["sts:AssumeRole"],
+            "Principal": {"Service": ["vefaas"]},
+        }
+    ]
+}
+
+# The permission policy attached to the role (provided by the platform team).
+FRONTEND_DEPLOY_POLICY: dict = {
+    "Statement": [
+        {
+            "Effect": "Allow",
+            "Action": [
+                "agentkit:*",
+                "apmplus_server:Get*",
+                "apmplus_server:List*",
+                "Volc_Observe:Get*",
+                "Volc_Observe:*Describe*",
+                "Volc_Observe:*List*",
+                "cloudmonitor:Get*",
+                "cloudmonitor:*Describe*",
+                "cloudmonitor:*List*",
+                "tls:Get*",
+                "tls:Describe*",
+                "tls:List*",
+                "apig:ConvertMcpConfig",
+                "apmplus_server:CreateOrder",
+                "ark:CreateApiKey",
+                "ark:GetEndpoint",
+                "ark:GetRawApiKey",
+                "ark:InnerDescribeModelEndpoints",
+                "ark:ListApiKeys",
+                "ark:ListEndpoints",
+                "ark:ListFoundationModelVersions",
+                "ark:ListModelChargeItems",
+                "ark:OpenModelChargeItem",
+                "ark:ListFoundationModels",
+                "content_customization:GetQuickOpenSearchInfo",
+                "cp:CreatePipeline",
+                "cp:CreateWorkspace",
+                "cp:DeletePipeline",
+                "cp:GetDefaultWorkspaceInner",
+                "cp:GetInstance",
+                "cp:GetTaskRunLog",
+                "cp:GetTaskRunLogDownloadURI",
+                "cp:ListPipelineRuns",
+                "cp:ListPipelineRunStagesInner",
+                "cp:ListPipelines",
+                "cp:ListWorkspaces",
+                "cp:RunPipeline",
+                "cr:AddVpcEndpointLinkedVpc",
+                "cr:CreateEndpointAclPolicies",
+                "cr:CreateNamespace",
+                "cr:CreateRegistry",
+                "cr:CreateRepository",
+                "cr:DeleteEndpointAclPolicies",
+                "cr:DeleteNamespace",
+                "cr:DeleteRegistry",
+                "cr:DeleteTags",
+                "cr:GetAuthorizationToken",
+                "cr:GetPublicEndpoint",
+                "cr:GetVpcEndpoint",
+                "cr:ListDomains",
+                "cr:ListNamespaces",
+                "cr:ListRegistries",
+                "cr:ListRepositories",
+                "cr:ListTags",
+                "cr:RemoveVpcEndpointLinkedVpc",
+                "cr:StartRegistry",
+                "cr:UpdateNamespace",
+                "cr:UpdatePublicEndpoint",
+                "cr:UpdateVpcEndpoint",
+                "ESCloud:GetMLApiKey",
+                "ESCloud:GetScene",
+                "ESCloud:GetSceneInstance",
+                "ESCloud:ListMLApiKeys",
+                "iam:AttachRolePolicy",
+                "iam:CheckServiceLinkedRole",
+                "iam:CreateRole",
+                "iam:CreateServiceLinkedRole",
+                "iam:DeleteRole",
+                "iam:DetachRolePolicy",
+                "iam:GetPolicy",
+                "iam:GetRole",
+                "iam:GetServiceLinkedRoleTemplate",
+                "iam:GetUser",
+                "iam:ListProjects",
+                "iam:ListRoles",
+                "iam:ListUsers",
+                "iam:UpdateRole",
+                "iam:ListAttachedRolePolicies",
+                "id:CreateApiKey",
+                "id:CreateApiKeyCredentialProvider",
+                "id:CreateInboundAuthConfig",
+                "id:DeleteApiKey",
+                "id:DeleteInboundAuthConfig",
+                "id:GetCallerIdentity",
+                "id:GetInboundAuthConfig",
+                "id:GetResourceApiKey",
+                "id:GetTenantServiceStatus",
+                "id:GetUserPool",
+                "id:GetUserPoolClient",
+                "id:ListCredentialProviders",
+                "id:ListNamespaces",
+                "id:ListServices",
+                "id:ListUserPoolClients",
+                "id:ListUserPools",
+                "id:UpdateInboundAuthConfig",
+                "location:DescribeRegions",
+                "mem0:CreateMemoryProject",
+                "mem0:DeleteMemoryProject",
+                "mem0:DescribeAPIKeyDetail",
+                "mem0:DescribeMemoryProjectDetail",
+                "mem0:DescribeMemoryProjects",
+                "mem0:ModifyMemoryProjectLongTermStrategy",
+                "mem0:SearchMemory",
+                "mem0:SearchMemoryAllInfo",
+                "ml_platform:GetKnowledgeBaseServiceInfo",
+                "ml_platform:ListCollections",
+                "sts:AssumeRole",
+                "tls:ActiveTlsAccount",
+                "tls:SearchLogs",
+                "tos:CopyObject",
+                "tos:CreateBucket",
+                "tos:DeleteBucket",
+                "tos:DeleteObject",
+                "tos:GetAccountStatus",
+                "tos:GetBucketInfo",
+                "tos:HeadBucket",
+                "tos:HeadObject",
+                "tos:ListBuckets",
+                "tos:ListObjectsV2",
+                "tos:PostObject",
+                "tos:PutObject",
+                "vefaas:CreateFunction",
+                "vefaas:CreateSandbox",
+                "vefaas:DeleteFunction",
+                "vefaas:DescribeSandbox",
+                "vefaas:EnableUserCrVpcTunnel",
+                "vefaas:GetAvailabilityZones",
+                "vefaas:GetFunction",
+                "vefaas:GetImageConfig",
+                "vefaas:GetImageSyncStatus",
+                "vefaas:GetReleaseStatus",
+                "vefaas:KillSandbox",
+                "vefaas:ListFunctions",
+                "vefaas:ListRuntimes",
+                "vefaas:ListSandboxes",
+                "vefaas:QueryUserCrVpcTunnel",
+                "vefaas:Release",
+                "vefaas:SetSandboxTimeout",
+                "vefaas:UpdateFunction",
+                "vikingdb:GetKnowledgeBaseServiceInfo",
+                "vikingdb:GetMemorydbInstanceDetail",
+                "vikingdb:ListCollections",
+                "vikingdb:MemoryCollectionInfo",
+                "vikingdb:MemoryCollectionList",
+                "vikingdb:MemorySearch",
+                "vikingdb:MemorySearchUser",
+                "vikingdb:GetCollectionInfo",
+                "vikingdb:CreateCollection",
+                "vikingdb:CreateKnowledgedbInstance",
+                "vikingdb:GetKnowledgedbInstanceStatus",
+                "vikingdb:AddDocument",
+                "vikingdb:UpdateCollection",
+                "vikingdb:DeleteCollection",
+                "vikingdb:CreateMemorydbInstance",
+                "vpc:DescribeSecurityGroups",
+                "vpc:DescribeSubnets",
+                "vpc:DescribeVpcAttributes",
+                "vpc:DescribeVpcs",
+                "vke:ListClusters",
+                "vke:ForwardKubernetesApi",
+                "vke:ListServices",
+                "ecs:DescribeInstances",
+                "llmshield:ListInstances",
+                "llmshield:ListApps",
+                "llmshield:CreateOuterApp",
+                "ark:ListSubscribeTrade",
+                "ark:EstimateSubscribePrice",
+                "ark:CreateSubscribeTrade",
+                "ark:GetCodingPlanUsage",
+                "rds_mysql:DescribeDBInstances",
+                "rds_mysql:DescribeDBInstanceEndpoints",
+                "rds_mysql:DescribeDBAccounts",
+                "rds_mysql:DescribeDatabases",
+                "rds_postgresql:DescribeDBInstances",
+                "rds_postgresql:DescribeDBAccounts",
+                "rds_postgresql:DescribeDatabases",
+                "aidap:DescribeWorkspaces",
+                "aidap:DescribeBranches",
+                "aidap:DescribeComputes",
+                "aidap:DescribeDBAccounts",
+                "aidap:DescribeDatabases",
+            ],
+            "Resource": ["*"],
+        }
+    ]
+}

veadk/cli/cli.py

@@ -20,7 +20,7 @@ from veadk.cli.cli_clean import clean
 from veadk.cli.cli_create import create
 from veadk.cli.cli_deploy import deploy
 from veadk.cli.cli_eval import eval
-from veadk.cli.cli_frontend import frontend
+from veadk.cli.cli_frontend import frontend, studio
 from veadk.cli.cli_harness import harness
 from veadk.cli.cli_init import init
 from veadk.cli.cli_kb import kb
@@ -52,6 +52,7 @@ veadk.add_command(create)
 veadk.add_command(prompt)
 veadk.add_command(web)
 veadk.add_command(frontend)
+veadk.add_command(studio)
 veadk.add_command(pipeline)
 veadk.add_command(eval)
 veadk.add_command(kb)