feat: deploy with OAuth2

Signed-off-by: Xie Zhihao <zhihao.xie@bytedance.com>

Author: zhxie

veadk/cli/cli_deploy.py

@@ -49,13 +49,23 @@
     type=click.Choice(["local", "mysql"]),
     help="Backend for short-term memory",
 )
+@click.option("--use-adk-web", is_flag=True, help="Whether to use ADK Web")
 @click.option(
     "--auth-method",
     default="none",
-    type=click.Choice(["none", "api-key"]),
+    type=click.Choice(["none", "api-key", "oauth2"]),
     help="=Authentication method for agent",
 )
-@click.option("--use-adk-web", is_flag=True, help="Whether to use ADK Web")
+@click.option(
+    "--veidentity-user-pool-name",
+    default="",
+    help="Expected Volcengine Identity user pool name",
+)
+@click.option(
+    "--veidentity-client-name",
+    default="",
+    help="Expected Volcengine Identity client name",
+)
 @click.option("--path", default=".", help="Local project path")
 def deploy(
     access_key: str,
@@ -67,6 +77,8 @@ def deploy(
     short_term_memory_backend: str,
     use_adk_web: bool,
     auth_method: str,
+    veidentity_user_pool_name: str,
+    veidentity_client_name: str,
     path: str,
 ) -> None:
     """Deploy a user project to Volcengine FaaS application.
@@ -99,6 +111,8 @@ def deploy(
         use_adk_web: Flag to enable ADK Web interface for the deployed agent
         auth_method: Authentication for the agent.
             Choices are 'none', 'api-key' or 'oauth2'.
+        veidentity_user_pool_name: Optional Volcengine Identity user pool name
+        veidentity_client_name: Optional Volcengine Identity client name
         path: Local directory path containing the VeADK project to deploy
 
     Note:
@@ -141,6 +155,8 @@ def deploy(
         "veapig_upstream_name": veapig_upstream_name,
         "use_adk_web": use_adk_web,
         "auth_method": auth_method,
+        "veidentity_user_pool_name": veidentity_user_pool_name,
+        "veidentity_client_name": veidentity_client_name,
         "veadk_version": VERSION,
     }
 

veadk/cli/cli_init.py

@@ -66,10 +66,12 @@ def _render_prompts() -> dict[str, Any]:
     auth_method_options = {
         "1": "None",
         "2": "API key",
+        "3": "OAuth2",
     }
     auth_methods = {
         "1": "none",
         "2": "api-key",
+        "3": "oauth2",
     }
 
     click.echo("Choose an authentication method:")
@@ -81,13 +83,27 @@ def _render_prompts() -> dict[str, Any]:
     )
     auth_method = auth_methods[auth_method_idx]
 
+    veidentity_user_pool_name = ""
+    veidentity_client_name = ""
+    if auth_method == "oauth2":
+        veidentity_user_pool_name = click.prompt(
+            "Volcengine Identity user pool name", default="", show_default=True
+        )
+
+        if use_adk_web:
+            veidentity_client_name = click.prompt(
+                "Volcengine Identity client name", default="", show_default=True
+            )
+
     return {
         "vefaas_application_name": vefaas_application_name,
         "veapig_instance_name": veapig_instance_name,
         "veapig_service_name": veapig_service_name,
         "veapig_upstream_name": veapig_upstream_name,
         "use_adk_web": use_adk_web,
         "auth_method": auth_method,
+        "veidentity_user_pool_name": veidentity_user_pool_name,
+        "veidentity_client_name": veidentity_client_name,
         "veadk_version": VERSION,
     }
 

veadk/cloud/cloud_agent_engine.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import os
 import socket
 import subprocess
@@ -23,7 +24,9 @@
 
 from veadk.cloud.cloud_app import CloudApp
 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.utils.logger import get_logger
 from veadk.utils.misc import formatted_timestamp
 
@@ -43,6 +46,8 @@ class CloudAgentEngine(BaseModel):
             Defaults to VOLCENGINE_SECRET_KEY environment variable.
         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.
 
     Note:
         Credentials must be set via environment variables for default behavior.
@@ -81,6 +86,16 @@ def model_post_init(self, context: Any, /) -> None:
             secret_key=self.volcengine_secret_key,
             region=self.region,
         )
+        self._veapig_service = APIGateway(
+            access_key=self.volcengine_access_key,
+            secret_key=self.volcengine_secret_key,
+            region=self.region,
+        )
+        self._veidentity_service = Identity(
+            access_key=self.volcengine_access_key,
+            secret_key=self.volcengine_secret_key,
+            region=self.region,
+        )
 
     def _prepare(self, path: str, name: str):
         """Prepares the local project for deployment by validating path and name.
@@ -197,8 +212,10 @@ def deploy(
         gateway_name: str = "",
         gateway_service_name: str = "",
         gateway_upstream_name: str = "",
-        auth_method: str = "none",
         use_adk_web: bool = False,
+        auth_method: str = "none",
+        identity_user_pool_name: str = "",
+        identity_client_name: str = "",
         local_test: bool = False,
     ) -> CloudApp:
         """Deploys a local agent project to Volcengine FaaS, creating necessary resources.
@@ -211,8 +228,10 @@ def deploy(
             gateway_name (str, optional): Custom gateway resource name. Defaults to timestamped.
             gateway_service_name (str, optional): Custom service name. Defaults to timestamped.
             gateway_upstream_name (str, optional): Custom upstream name. Defaults to timestamped.
-            auth_method (str, optional): Authentication for the agent. Defaults to none.
             use_adk_web (bool): Enable ADK Web configuration. Defaults to False.
+            auth_method (str, optional): Authentication for the agent. Defaults to none.
+            identity_user_pool_name (str, optional): Custom user pool name. Defaults to timestamped.
+            identity_client_name (str, optional): Custom client name. Defaults to timestamped.
             local_test (bool): Perform FastAPI server test before deploy. Defaults to False.
 
         Returns:
@@ -256,6 +275,12 @@ def deploy(
             gateway_service_name = f"{application_name}-gw-svr-{formatted_timestamp()}"
         if not gateway_upstream_name:
             gateway_upstream_name = f"{application_name}-gw-us-{formatted_timestamp()}"
+        if not identity_user_pool_name:
+            identity_user_pool_name = (
+                f"{application_name}-id-up-{formatted_timestamp()}"
+            )
+        if not identity_client_name:
+            identity_client_name = f"{application_name}-id-cli-{formatted_timestamp()}"
 
         try:
             vefaas_application_url, app_id, function_id = self._vefaas_service.deploy(
@@ -268,6 +293,103 @@ 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"
+            ]
+
+            if auth_method == "oauth2":
+                # Get or create the Identity user pool.
+                identity_user_pool_id = self._veidentity_service.get_user_pool(
+                    name=identity_user_pool_name,
+                )
+                if not identity_user_pool_id:
+                    identity_user_pool_id = self._veidentity_service.create_user_pool(
+                        name=identity_user_pool_name,
+                    )
+
+                # Create APIG upstream for Identity.
+                identity_domain = f"auth.id.{self.region}.volces.com"
+                veapig_identity_upstream_id = (
+                    self._veapig_service.check_domain_upstream_exist(
+                        domain=identity_domain,
+                        port=443,
+                        gateway_id=veapig_gateway_id,
+                    )
+                )
+                if not veapig_identity_upstream_id:
+                    veapig_identity_upstream_id = (
+                        self._veapig_service.create_domain_upstream(
+                            domain=f"auth.id.{self.region}.volces.com",
+                            port=443,
+                            is_https=True,
+                            gateway_id=veapig_gateway_id,
+                            upstream_name=f"id-{formatted_timestamp()}",
+                        )
+                    )
+
+                # Create plugin binding.
+                plugin_name = ""
+                plugin_config = {}
+                if use_adk_web:
+                    # Get or create the Identity client.
+                    identity_client_id = ""
+                    identity_client_secret = ""
+                    identity_client = self._veidentity_service.get_user_pool_client(
+                        user_pool_uid=identity_user_pool_id,
+                        name=identity_client_name,
+                    )
+                    if identity_client:
+                        identity_client_id = identity_client[0]
+                        identity_client_secret = identity_client[1]
+                    else:
+                        identity_client_id, identity_client_secret = (
+                            self._veidentity_service.create_user_pool_client(
+                                user_pool_uid=identity_user_pool_id,
+                                name=identity_client_name,
+                                client_type="WEB_APPLICATION",
+                            )
+                        )
+
+                    self._veidentity_service.register_callback_for_user_pool_client(
+                        user_pool_uid=identity_user_pool_id,
+                        client_uid=identity_client_id,
+                        callback_url=f"{vefaas_application_url}/callback",
+                        web_origin=vefaas_application_url,
+                    )
+
+                    plugin_name = "wasm-oauth2-sso"
+                    plugin_config = {
+                        "AuthorizationUrl": f"https://auth.id.{self.region}.volces.com/userpool/{identity_user_pool_id}/authorize",
+                        "UpstreamId": veapig_identity_upstream_id,
+                        "TokenUrl": f"https://auth.id.{self.region}.volces.com/userpool/{identity_user_pool_id}/oauth/token",
+                        "RedirectPath": "/callback",
+                        "SignoutPath": "/signout",
+                        "ClientId": identity_client_id,
+                        "ClientSecret": identity_client_secret,
+                    }
+                else:
+                    plugin_name = "wasm-jwt-auth"
+                    plugin_config = {
+                        "RemoteJwks": {
+                            "UpstreamId": veapig_identity_upstream_id,
+                            "Url": f"auth.id.{self.region}.volces.com/userpool/{identity_user_pool_id}/keys",
+                        },
+                        "Issuer": f"https://auth.id.{self.region}.volces.com/userpool/{identity_user_pool_id}",
+                        "ValidateConsumer": False,
+                    }
+
+                self._vefaas_service.apig_client.create_plugin_binding(
+                    scope="ROUTE",
+                    target=veapig_route_id,
+                    plugin_name=plugin_name,
+                    plugin_config=json.dumps(plugin_config),
+                )
+
             return CloudApp(
                 vefaas_application_name=application_name,
                 vefaas_endpoint=vefaas_application_url,