From b308c6b05db54b347dc8d4781c462518333c4125 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Wed, 23 Oct 2024 23:20:51 -0700
Subject: [PATCH] enforce collection id for login parameters (#1043)

---
 skyvern/forge/sdk/routes/agent_protocol.py | 10 +++++++++-
 skyvern/forge/sdk/workflow/exceptions.py   |  8 ++++++++
 skyvern/forge/sdk/workflow/service.py      |  7 +++++++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py
index fdd2eb9f..7fa0f855 100644
--- a/skyvern/forge/sdk/routes/agent_protocol.py
+++ b/skyvern/forge/sdk/routes/agent_protocol.py
@@ -51,7 +51,11 @@ from skyvern.forge.sdk.schemas.tasks import (
 )
 from skyvern.forge.sdk.services import org_auth_service
 from skyvern.forge.sdk.settings_manager import SettingsManager
-from skyvern.forge.sdk.workflow.exceptions import FailedToCreateWorkflow, FailedToUpdateWorkflow
+from skyvern.forge.sdk.workflow.exceptions import (
+    FailedToCreateWorkflow,
+    FailedToUpdateWorkflow,
+    WorkflowParameterMissingRequiredValue,
+)
 from skyvern.forge.sdk.workflow.models.workflow import (
     RunWorkflowResponse,
     Workflow,
@@ -707,6 +711,8 @@ async def create_workflow(
         return await app.WORKFLOW_SERVICE.create_workflow_from_request(
             organization=current_org, request=workflow_create_request
         )
+    except WorkflowParameterMissingRequiredValue as e:
+        raise e
     except Exception as e:
         LOG.error("Failed to create workflow", exc_info=True, organization_id=current_org.organization_id)
         raise FailedToCreateWorkflow(str(e))
@@ -753,6 +759,8 @@ async def update_workflow(
             request=workflow_create_request,
             workflow_permanent_id=workflow_permanent_id,
         )
+    except WorkflowParameterMissingRequiredValue as e:
+        raise e
     except Exception as e:
         LOG.exception("Failed to update workflow", workflow_permanent_id=workflow_permanent_id)
         raise FailedToUpdateWorkflow(workflow_permanent_id, f"<{type(e).__name__}: {str(e)}>")
diff --git a/skyvern/forge/sdk/workflow/exceptions.py b/skyvern/forge/sdk/workflow/exceptions.py
index f9e64def..96e09813 100644
--- a/skyvern/forge/sdk/workflow/exceptions.py
+++ b/skyvern/forge/sdk/workflow/exceptions.py
@@ -99,3 +99,11 @@ class InvalidFileType(BaseWorkflowHTTPException):
             f"File URL {file_url} is not a valid {file_type} file. Error: {error}",
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
         )
+
+
+class WorkflowParameterMissingRequiredValue(BaseWorkflowHTTPException):
+    def __init__(self, workflow_parameter_type: str, workflow_parameter_key: str, required_value: str) -> None:
+        super().__init__(
+            f"Missing required value for workflow parameter. Workflow parameter type: {workflow_parameter_type}. workflow_parameter_key: {workflow_parameter_key}. Required value: {required_value}",
+            status_code=status.HTTP_400_BAD_REQUEST,
+        )
diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py
index 6404cd37..4d9b4016 100644
--- a/skyvern/forge/sdk/workflow/service.py
+++ b/skyvern/forge/sdk/workflow/service.py
@@ -19,6 +19,7 @@ from skyvern.forge.sdk.workflow.exceptions import (
     InvalidWorkflowDefinition,
     WorkflowDefinitionHasDuplicateParameterKeys,
     WorkflowDefinitionHasReservedParameterKeys,
+    WorkflowParameterMissingRequiredValue,
 )
 from skyvern.forge.sdk.workflow.models.block import (
     BlockStatus,
@@ -957,6 +958,12 @@ class WorkflowService:
                         description=parameter.description,
                     )
                 elif parameter.parameter_type == ParameterType.BITWARDEN_LOGIN_CREDENTIAL:
+                    if not parameter.bitwarden_collection_id:
+                        raise WorkflowParameterMissingRequiredValue(
+                            workflow_parameter_type=ParameterType.BITWARDEN_LOGIN_CREDENTIAL,
+                            workflow_parameter_key=parameter.key,
+                            required_value="bitwarden_collection_id",
+                        )
                     parameters[parameter.key] = await self.create_bitwarden_login_credential_parameter(
                         workflow_id=workflow.workflow_id,
                         bitwarden_client_id_aws_secret_key=parameter.bitwarden_client_id_aws_secret_key,