From fc6ddb7bf65526186ddc3e26a9be08cf8b193317 Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Tue, 4 Jun 2024 08:27:04 -0700
Subject: [PATCH] fix workflow run get with workflow_permanent_id (#415)

---
 skyvern/forge/sdk/routes/agent_protocol.py |  2 +-
 skyvern/forge/sdk/workflow/service.py      | 25 +++++++---------------
 2 files changed, 9 insertions(+), 18 deletions(-)

diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py
index 6f32ea9a..a72485d6 100644
--- a/skyvern/forge/sdk/routes/agent_protocol.py
+++ b/skyvern/forge/sdk/routes/agent_protocol.py
@@ -533,7 +533,7 @@ async def get_workflow_run(
 ) -> WorkflowRunStatusResponse:
     analytics.capture("skyvern-oss-agent-workflow-run-get")
     return await app.WORKFLOW_SERVICE.build_workflow_run_status_response(
-        workflow_id=workflow_id,
+        workflow_permanent_id=workflow_id,
         workflow_run_id=workflow_run_id,
         organization_id=current_org.organization_id,
     )
diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py
index 2f1cf4b5..479f399f 100644
--- a/skyvern/forge/sdk/workflow/service.py
+++ b/skyvern/forge/sdk/workflow/service.py
@@ -6,13 +6,7 @@ import requests
 import structlog
 
 from skyvern import analytics
-from skyvern.exceptions import (
-    FailedToSendWebhook,
-    MissingValueForParameter,
-    WorkflowNotFound,
-    WorkflowOrganizationMismatch,
-    WorkflowRunNotFound,
-)
+from skyvern.exceptions import FailedToSendWebhook, MissingValueForParameter, WorkflowNotFound, WorkflowRunNotFound
 from skyvern.forge import app
 from skyvern.forge.sdk.artifact.models import ArtifactType
 from skyvern.forge.sdk.core import skyvern_context
@@ -527,17 +521,14 @@ class WorkflowService:
 
     async def build_workflow_run_status_response(
         self,
-        workflow_id: str,
+        workflow_permanent_id: str,
         workflow_run_id: str,
         organization_id: str,
     ) -> WorkflowRunStatusResponse:
-        workflow = await self.get_workflow(workflow_id=workflow_id, organization_id=organization_id)
+        workflow = await self.get_workflow_by_permanent_id(workflow_permanent_id, organization_id=organization_id)
         if workflow is None:
-            LOG.error(f"Workflow {workflow_id} not found")
-            raise WorkflowNotFound(workflow_id=workflow_id)
-        if workflow.organization_id != organization_id:
-            LOG.error(f"Workflow {workflow_id} does not belong to organization {organization_id}")
-            raise WorkflowOrganizationMismatch(workflow_id=workflow_id, organization_id=organization_id)
+            LOG.error(f"Workflow {workflow_permanent_id} not found")
+            raise WorkflowNotFound(workflow_permanent_id=workflow_permanent_id)
 
         workflow_run = await self.get_workflow_run(workflow_run_id=workflow_run_id)
         workflow_run_tasks = await app.DATABASE.get_tasks_by_workflow_run_id(workflow_run_id=workflow_run_id)
@@ -574,14 +565,14 @@ class WorkflowService:
         output_parameter_tuples: list[
             tuple[OutputParameter, WorkflowRunOutputParameter]
         ] = await self.get_output_parameter_workflow_run_output_parameter_tuples(
-            workflow_id=workflow_id, workflow_run_id=workflow_run_id
+            workflow_id=workflow_run.workflow_id, workflow_run_id=workflow_run_id
         )
         if output_parameter_tuples:
             outputs = {output_parameter.key: output.value for output_parameter, output in output_parameter_tuples}
         else:
             LOG.error(
                 "No output parameters found for workflow run",
-                workflow_id=workflow_id,
+                workflow_permanent_id=workflow_permanent_id,
                 workflow_run_id=workflow_run_id,
             )
             outputs = None
@@ -621,7 +612,7 @@ class WorkflowService:
         await app.ARTIFACT_MANAGER.wait_for_upload_aiotasks_for_tasks(all_workflow_task_ids)
 
         workflow_run_status_response = await self.build_workflow_run_status_response(
-            workflow_id=workflow.workflow_id,
+            workflow_permanent_id=workflow.workflow_permanent_id,
             workflow_run_id=workflow_run.workflow_run_id,
             organization_id=workflow.organization_id,
         )