chore: merge Identity and IdentityClient

zhxie · zhxie · commit d4347512339d · 2025-11-14T11:09:09.000+08:00
Signed-off-by: Xie Zhihao &lt;zhihao.xie@bytedance.com&gt;
diff --git a/veadk/cloud/cloud_agent_engine.py b/veadk/cloud/cloud_agent_engine.py
@@ -26,7 +26,7 @@
 from veadk.config import getenv, veadk_environments
 from veadk.integrations.ve_apig.ve_apig import APIGateway
 from veadk.integrations.ve_faas.ve_faas import VeFaaS
-from veadk.integrations.ve_identity.ve_identity import Identity
+from veadk.integrations.ve_identity.identity_client import IdentityClient
 from veadk.utils.logger import get_logger
 from veadk.utils.misc import formatted_timestamp
 
@@ -47,7 +47,7 @@ class CloudAgentEngine(BaseModel):
         region (str): Region for Volcengine services. Defaults to "cn-beijing".
         _vefaas_service (VeFaaS): Internal VeFaaS client instance, initialized post-creation.
         _veapig_service (APIGateway): Internal VeAPIG client instance, initialized post-creation.
-        _veidentity_service (Identity): Internal Identity client instance, initialized post-creation.
+        _veidentity_service (IdentityClient): Internal Identity client instance, initialized post-creation.
 
     Note:
         Credentials must be set via environment variables for default behavior.
@@ -91,7 +91,7 @@ def model_post_init(self, context: Any, /) -> None:
             secret_key=self.volcengine_secret_key,
             region=self.region,
         )
-        self._veidentity_service = Identity(
+        self._veidentity_service = IdentityClient(
             access_key=self.volcengine_access_key,
             secret_key=self.volcengine_secret_key,
             region=self.region,
@@ -293,14 +293,9 @@ def deploy(
             )
             _ = function_id  # for future use
 
-            app = self._vefaas_service.get_application_details(app_id=app_id)
-            cloud_resource = json.loads(app["CloudResource"])
-            veapig_gateway_id = cloud_resource["framework"]["triggers"][0][
-                "DetailedConfig"
-            ]["GatewayId"]
-            veapig_route_id = cloud_resource["framework"]["triggers"][0]["Routes"][0][
-                "Id"
-            ]
+            veapig_gateway_id, _, veapig_route_id = (
+                self._vefaas_service.get_application_route(app_id=app_id)
+            )
 
             if auth_method == "oauth2":
                 # Get or create the Identity user pool.
@@ -311,6 +306,7 @@ def deploy(
                     identity_user_pool_id = self._veidentity_service.create_user_pool(
                         name=identity_user_pool_name,
                     )
+                issuer = f"https://auth.id.{self.region}.volces.com/userpool/{identity_user_pool_id}"
 
                 # Create APIG upstream for Identity.
                 identity_domain = f"auth.id.{self.region}.volces.com"
@@ -364,9 +360,9 @@ def deploy(
 
                     plugin_name = "wasm-oauth2-sso"
                     plugin_config = {
-                        "AuthorizationUrl": f"https://auth.id.{self.region}.volces.com/userpool/{identity_user_pool_id}/authorize",
+                        "AuthorizationUrl": f"{issuer}/authorize",
                         "UpstreamId": veapig_identity_upstream_id,
-                        "TokenUrl": f"https://auth.id.{self.region}.volces.com/userpool/{identity_user_pool_id}/oauth/token",
+                        "TokenUrl": f"{issuer}/oauth/token",
                         "RedirectPath": "/callback",
                         "SignoutPath": "/signout",
                         "ClientId": identity_client_id,
@@ -377,12 +373,11 @@ def deploy(
                     plugin_config = {
                         "RemoteJwks": {
                             "UpstreamId": veapig_identity_upstream_id,
-                            "Url": f"auth.id.{self.region}.volces.com/userpool/{identity_user_pool_id}/keys",
+                            "Url": f"{issuer}/keys",
                         },
-                        "Issuer": f"https://auth.id.{self.region}.volces.com/userpool/{identity_user_pool_id}",
+                        "Issuer": issuer,
                         "ValidateConsumer": False,
                     }
-
                 self._vefaas_service.apig_client.create_plugin_binding(
                     scope="ROUTE",
                     target=veapig_route_id,
diff --git a/veadk/integrations/ve_faas/ve_faas.py b/veadk/integrations/ve_faas/ve_faas.py
@@ -206,9 +206,11 @@ def _release_application(self, app_id: str):
             logs = "\n".join(self._get_application_logs(app_id=app_id))
             log_text = re.sub(
                 r'([{"\']?(key|secret|token|pass|auth|credential|access|api|ak|sk|doubao|volces|coze)[^"\'\s]*["\']?\s*[:=]\s*)(["\']?)([^"\'\s]+)(["\']?)|([A-Za-z0-9+/=]{20,})',
-                lambda m: f"{m.group(1)}{m.group(3)}******{m.group(5)}"
-                if m.group(1)
-                else "******",
+                lambda m: (
+                    f"{m.group(1)}{m.group(3)}******{m.group(5)}"
+                    if m.group(1)
+                    else "******"
+                ),
                 logs,
                 flags=re.IGNORECASE,
             )
@@ -232,9 +234,11 @@ def _list_application(self, app_id: str = None, app_name: str = None):
         request_body = {
             "OrderBy": {"Key": "CreateTime", "Ascend": False},
             "FunctionId": app_id if app_id else None,
-            "Filters": [{"Item": {"Key": "Name", "Value": [app_name]}}]
-            if app_name and not app_id
-            else None,
+            "Filters": (
+                [{"Item": {"Key": "Name", "Value": [app_name]}}]
+                if app_name and not app_id
+                else None
+            ),
         }
         # remove None
         request_body = {k: v for k, v in request_body.items() if v is not None}
@@ -354,6 +358,26 @@ def get_application_details(self, app_id: str = None, app_name: str = None):
                 if app["Name"] == app_name:
                     return app
 
+    def get_application_route(
+        self, app_id: str = None, app_name: str = None
+    ) -> tuple[str, str, str] | None:
+        app = self.get_application_details(
+            app_id=app_id,
+            app_name=app_name,
+        )
+        if not app:
+            return None
+
+        cloud_resource = json.loads(app["CloudResource"])
+        gateway_id = cloud_resource["framework"]["triggers"][0]["DetailedConfig"][
+            "GatewayId"
+        ]
+        service_id = cloud_resource["framework"]["triggers"][0]["Routes"][0][
+            "ServiceId"
+        ]
+        route_id = cloud_resource["framework"]["triggers"][0]["Routes"][0]["Id"]
+        return gateway_id, service_id, route_id
+
     def find_app_id_by_name(self, name: str):
         apps = self._list_application(app_name=name)
         for app in apps:
diff --git a/veadk/integrations/ve_identity/identity_client.py b/veadk/integrations/ve_identity/identity_client.py
@@ -533,3 +533,133 @@ async def create_oauth2_credential_provider_with_dcr(
 
         # Create the credential provider with updated config
         return self.create_oauth2_credential_provider(request_params)
+
+    def create_user_pool(self, name: str) -> str:
+        from volcenginesdkid import CreateUserPoolRequest, CreateUserPoolResponse
+
+        request = CreateUserPoolRequest(
+            name=name,
+        )
+        response: CreateUserPoolResponse = self._api_client.create_user_pool(request)
+
+        return response.uid
+
+    def get_user_pool(self, name: str) -> str | None:
+        from volcenginesdkid import (
+            ListUserPoolsRequest,
+            ListUserPoolsResponse,
+            FilterForListUserPoolsInput,
+            DataForListUsersOutput,
+        )
+
+        request = ListUserPoolsRequest(
+            page_number=1,
+            page_size=1,
+            filter=FilterForListUserPoolsInput(
+                name=name,
+            ),
+        )
+        response: ListUserPoolsResponse = self._api_client.list_user_pools(request)
+        if response.total_count == 0:
+            return None
+
+        user_pool: DataForListUsersOutput = response.data[0]
+        return user_pool.uid
+
+    def create_user_pool_client(
+        self, user_pool_uid: str, name: str, client_type: str
+    ) -> tuple[str, str]:
+        from volcenginesdkid import (
+            CreateUserPoolClientRequest,
+            CreateUserPoolClientResponse,
+        )
+
+        request = CreateUserPoolClientRequest(
+            user_pool_uid=user_pool_uid,
+            name=name,
+            client_type=client_type,
+        )
+        response: CreateUserPoolClientResponse = (
+            self._api_client.create_user_pool_client(request)
+        )
+        return response.uid, response.client_secret
+
+    def register_callback_for_user_pool_client(
+        self,
+        user_pool_uid: str,
+        client_uid: str,
+        callback_url: str,
+        web_origin: str,
+    ):
+        from volcenginesdkid import (
+            GetUserPoolClientRequest,
+            GetUserPoolClientResponse,
+            UpdateUserPoolClientRequest,
+        )
+
+        request = GetUserPoolClientRequest(
+            user_pool_uid=user_pool_uid,
+            client_uid=client_uid,
+        )
+        response: GetUserPoolClientResponse = self._api_client.get_user_pool_client(
+            request
+        )
+
+        allowed_callback_urls = response.allowed_callback_urls
+        if not allowed_callback_urls:
+            allowed_callback_urls = []
+        allowed_callback_urls.append(callback_url)
+        allowed_web_origins = response.allowed_web_origins
+        if not allowed_web_origins:
+            allowed_web_origins = []
+        allowed_web_origins.append(web_origin)
+
+        request2 = UpdateUserPoolClientRequest(
+            user_pool_uid=user_pool_uid,
+            client_uid=client_uid,
+            name=response.name,
+            description=response.description,
+            allowed_callback_urls=allowed_callback_urls,
+            allowed_logout_urls=response.allowed_logout_urls,
+            allowed_web_origins=allowed_web_origins,
+            allowed_cors=response.allowed_cors,
+            id_token=response.id_token,
+            refresh_token=response.refresh_token,
+        )
+        self._api_client.update_user_pool_client(request2)
+
+    def get_user_pool_client(
+        self, user_pool_uid: str, name: str
+    ) -> tuple[str, str] | None:
+        from volcenginesdkid import (
+            ListUserPoolClientsRequest,
+            ListUserPoolClientsResponse,
+            FilterForListUserPoolClientsInput,
+            DataForListUserPoolClientsOutput,
+            GetUserPoolClientRequest,
+            GetUserPoolClientResponse,
+        )
+
+        request = ListUserPoolClientsRequest(
+            user_pool_uid=user_pool_uid,
+            page_number=1,
+            page_size=1,
+            filter=FilterForListUserPoolClientsInput(
+                name=name,
+            ),
+        )
+        response: ListUserPoolClientsResponse = self._api_client.list_user_pool_clients(
+            request
+        )
+        if response.total_count == 0:
+            return None
+
+        client: DataForListUserPoolClientsOutput = response.data[0]
+        request2 = GetUserPoolClientRequest(
+            user_pool_uid=user_pool_uid,
+            client_uid=client.uid,
+        )
+        response2: GetUserPoolClientResponse = self._api_client.get_user_pool_client(
+            request2
+        )
+        return response2.uid, response2.client_secret
diff --git a/veadk/integrations/ve_identity/ve_identity.py b/veadk/integrations/ve_identity/ve_identity.py
