fix run duration calculation (#3427)

2025-09-13 15:57:48 +08:00
parent 5c1853b0c3
commit 863d7473a1
3 changed files with 100 additions and 45 deletions
--- a/skyvern/forge/sdk/workflow/service.py
+++ b/skyvern/forge/sdk/workflow/service.py
@@ -354,6 +354,7 @@ class WorkflowService:

        # Check if there's a related workflow script that should be used instead
        workflow_script, _ = await self._get_workflow_script(workflow, workflow_run, block_labels)
+        is_script = workflow_script is not None
        if workflow_script is not None:
            LOG.info(
                "Found related workflow script, running script instead of workflow",
@@ -393,7 +394,9 @@ class WorkflowService:
                WorkflowRunStatus.terminated,
                WorkflowRunStatus.timed_out,
            ):
-                workflow_run = await self.mark_workflow_run_as_completed(workflow_run_id=workflow_run_id)
+                workflow_run = await self.mark_workflow_run_as_completed(
+                    workflow_run_id=workflow_run_id, is_script=is_script
+                )
                await self.generate_script_if_needed(
                    workflow=workflow,
                    workflow_run=workflow_run,
@@ -413,18 +416,6 @@ class WorkflowService:
            close_browser_on_completion=close_browser_on_completion,
        )

-        # Track workflow run duration when completed
-        duration_seconds = (datetime.now(UTC) - workflow_run.created_at.replace(tzinfo=UTC)).total_seconds()
-        LOG.info(
-            "Workflow run duration metrics",
-            workflow_run_id=workflow_run_id,
-            workflow_id=workflow_run.workflow_id,
-            duration_seconds=duration_seconds,
-            workflow_run_status=workflow_run.status,
-            organization_id=organization_id,
-            is_script_run=workflow_script is not None,
-        )
-
        return workflow_run

    async def _execute_workflow_blocks(
@@ -963,77 +954,119 @@ class WorkflowService:
            browser_address=workflow_request.browser_address,
        )

-    async def mark_workflow_run_as_completed(self, workflow_run_id: str) -> WorkflowRun:
+    async def _update_workflow_run_status(
+        self,
+        workflow_run_id: str,
+        status: WorkflowRunStatus,
+        failure_reason: str | None = None,
+        is_script: bool = False,
+    ) -> WorkflowRun:
+        workflow_run = await app.DATABASE.update_workflow_run(
+            workflow_run_id=workflow_run_id,
+            status=status,
+            failure_reason=failure_reason,
+        )
+        if status in [WorkflowRunStatus.completed, WorkflowRunStatus.failed, WorkflowRunStatus.terminated]:
+            start_time = (
+                workflow_run.started_at.replace(tzinfo=UTC)
+                if workflow_run.started_at
+                else workflow_run.created_at.replace(tzinfo=UTC)
+            )
+            queued_seconds = (start_time - workflow_run.created_at.replace(tzinfo=UTC)).total_seconds()
+            duration_seconds = (datetime.now(UTC) - start_time).total_seconds()
+            LOG.info(
+                "Workflow run duration metrics",
+                workflow_run_id=workflow_run_id,
+                workflow_id=workflow_run.workflow_id,
+                queued_seconds=queued_seconds,
+                duration_seconds=duration_seconds,
+                workflow_run_status=workflow_run.status,
+                organization_id=workflow_run.organization_id,
+                is_script_run=is_script,
+            )
+        return workflow_run
+
+    async def mark_workflow_run_as_completed(self, workflow_run_id: str, is_script: bool = False) -> WorkflowRun:
        LOG.info(
            f"Marking workflow run {workflow_run_id} as completed",
            workflow_run_id=workflow_run_id,
            workflow_status="completed",
        )
-        return await app.DATABASE.update_workflow_run(
+        return await self._update_workflow_run_status(
            workflow_run_id=workflow_run_id,
            status=WorkflowRunStatus.completed,
+            is_script=is_script,
        )

-    async def mark_workflow_run_as_failed(self, workflow_run_id: str, failure_reason: str | None) -> WorkflowRun:
+    async def mark_workflow_run_as_failed(
+        self, workflow_run_id: str, failure_reason: str | None, is_script: bool = False
+    ) -> WorkflowRun:
        LOG.info(
            f"Marking workflow run {workflow_run_id} as failed",
            workflow_run_id=workflow_run_id,
            workflow_status="failed",
            failure_reason=failure_reason,
        )
-        return await app.DATABASE.update_workflow_run(
+        return await self._update_workflow_run_status(
            workflow_run_id=workflow_run_id,
            status=WorkflowRunStatus.failed,
            failure_reason=failure_reason,
+            is_script=is_script,
        )

-    async def mark_workflow_run_as_running(self, workflow_run_id: str) -> WorkflowRun:
+    async def mark_workflow_run_as_running(self, workflow_run_id: str, is_script: bool = False) -> WorkflowRun:
        LOG.info(
            f"Marking workflow run {workflow_run_id} as running",
            workflow_run_id=workflow_run_id,
            workflow_status="running",
        )
-        return await app.DATABASE.update_workflow_run(
+        return await self._update_workflow_run_status(
            workflow_run_id=workflow_run_id,
            status=WorkflowRunStatus.running,
+            is_script=is_script,
        )

-    async def mark_workflow_run_as_terminated(self, workflow_run_id: str, failure_reason: str | None) -> WorkflowRun:
+    async def mark_workflow_run_as_terminated(
+        self, workflow_run_id: str, failure_reason: str | None, is_script: bool = False
+    ) -> WorkflowRun:
        LOG.info(
            f"Marking workflow run {workflow_run_id} as terminated",
            workflow_run_id=workflow_run_id,
            workflow_status="terminated",
            failure_reason=failure_reason,
        )
-        return await app.DATABASE.update_workflow_run(
+        return await self._update_workflow_run_status(
            workflow_run_id=workflow_run_id,
            status=WorkflowRunStatus.terminated,
            failure_reason=failure_reason,
+            is_script=is_script,
        )

-    async def mark_workflow_run_as_canceled(self, workflow_run_id: str) -> WorkflowRun:
+    async def mark_workflow_run_as_canceled(self, workflow_run_id: str, is_script: bool = False) -> WorkflowRun:
        LOG.info(
            f"Marking workflow run {workflow_run_id} as canceled",
            workflow_run_id=workflow_run_id,
            workflow_status="canceled",
        )
-        return await app.DATABASE.update_workflow_run(
+        return await self._update_workflow_run_status(
            workflow_run_id=workflow_run_id,
            status=WorkflowRunStatus.canceled,
+            is_script=is_script,
        )

    async def mark_workflow_run_as_timed_out(
-        self, workflow_run_id: str, failure_reason: str | None = None
+        self, workflow_run_id: str, failure_reason: str | None = None, is_script: bool = False
    ) -> WorkflowRun:
        LOG.info(
            f"Marking workflow run {workflow_run_id} as timed out",
            workflow_run_id=workflow_run_id,
            workflow_status="timed_out",
        )
-        return await app.DATABASE.update_workflow_run(
+        return await self._update_workflow_run_status(
            workflow_run_id=workflow_run_id,
            status=WorkflowRunStatus.timed_out,
            failure_reason=failure_reason,
+            is_script=is_script,
        )

    async def get_workflow_run(self, workflow_run_id: str, organization_id: str | None = None) -> WorkflowRun:
@@ -2504,7 +2537,9 @@ class WorkflowService:
            )

            # Mark workflow run as completed
-            workflow_run = await self.mark_workflow_run_as_completed(workflow_run_id=workflow_run.workflow_run_id)
+            workflow_run = await self.mark_workflow_run_as_completed(
+                workflow_run_id=workflow_run.workflow_run_id, is_script=True
+            )

            LOG.info(
                "Successfully executed workflow script",
@@ -2526,7 +2561,7 @@ class WorkflowService:
            # Mark workflow run as failed
            failure_reason = f"Failed to execute workflow script: {str(e)}"
            workflow_run = await self.mark_workflow_run_as_failed(
-                workflow_run_id=workflow_run.workflow_run_id, failure_reason=failure_reason
+                workflow_run_id=workflow_run.workflow_run_id, failure_reason=failure_reason, is_script=True
            )

            return workflow_run