append complete action (for validation) to the end of every task block (#3726)

skyvern/core/script_generations/constants.py

@@ -6,3 +6,8 @@ SCRIPT_TASK_BLOCKS = {
     "extraction",
     "login",
 }
+SCRIPT_TASK_BLOCKS_WITH_COMPLETE_ACTION = {
+    "task",
+    "navigation",
+    "login",
+}

skyvern/core/script_generations/generate_script.py

@@ -17,7 +17,7 @@ import structlog
 from libcst import Attribute, Call, Dict, DictElement, FunctionDef, Name, Param
 
 from skyvern.config import settings
-from skyvern.core.script_generations.constants import SCRIPT_TASK_BLOCKS
+from skyvern.core.script_generations.constants import SCRIPT_TASK_BLOCKS, SCRIPT_TASK_BLOCKS_WITH_COMPLETE_ACTION
 from skyvern.core.script_generations.generate_workflow_parameters import (
     generate_workflow_parameters_schema,
     hydrate_input_text_actions_with_field_names,
@@ -92,6 +92,7 @@ ACTION_MAP = {
     "verification_code": "verification_code",
     "wait": "wait",
     "extract": "extract",
+    "complete": "complete",
 }
 ACTIONS_WITH_XPATH = [
     "click",
@@ -402,26 +403,34 @@ def _action_to_stmt(act: dict[str, Any], task: dict[str, Any], assign_to_output:
                     comma=cst.Comma(),
                 )
             )
+    intention = act.get("intention") or act.get("reasoning") or ""
+    if intention:
+        args.extend(
+            [
+                cst.Arg(
+                    keyword=cst.Name("intention"),
+                    value=_value(intention),
+                    whitespace_after_arg=cst.ParenthesizedWhitespace(indent=True),
+                    comma=cst.Comma(),
+                ),
+            ]
+        )
 
-    args.extend(
-        [
-            cst.Arg(
-                keyword=cst.Name("intention"),
-                value=_value(act.get("intention") or act.get("reasoning") or ""),
-                whitespace_after_arg=cst.ParenthesizedWhitespace(indent=True),
-                comma=cst.Comma(),
+    # Only use indented parentheses if we have arguments
+    if args:
+        call = cst.Call(
+            func=cst.Attribute(value=cst.Name("page"), attr=cst.Name(method)),
+            args=args,
+            whitespace_before_args=cst.ParenthesizedWhitespace(
+                indent=True,
+                last_line=cst.SimpleWhitespace(INDENT),
             ),
-        ]
-    )
-
-    call = cst.Call(
-        func=cst.Attribute(value=cst.Name("page"), attr=cst.Name(method)),
-        args=args,
-        whitespace_before_args=cst.ParenthesizedWhitespace(
-            indent=True,
-            last_line=cst.SimpleWhitespace(INDENT),
-        ),
-    )
+        )
+    else:
+        call = cst.Call(
+            func=cst.Attribute(value=cst.Name("page"), attr=cst.Name(method)),
+            args=args,
+        )
 
     # await page.method(...)
     await_expr = cst.Await(call)
@@ -456,6 +465,12 @@ def _build_block_fn(block: dict[str, Any], actions: list[dict[str, Any]]) -> Fun
         assign_to_output = is_extraction_block and act["action_type"] == "extract"
         body_stmts.append(_action_to_stmt(act, block, assign_to_output=assign_to_output))
 
+    # add complete action
+    block_type = block.get("block_type")
+    if block_type in SCRIPT_TASK_BLOCKS_WITH_COMPLETE_ACTION:
+        complete_action = {"action_type": "complete"}
+        body_stmts.append(_action_to_stmt(complete_action, block, assign_to_output=assign_to_output))
+
     # For extraction blocks, add return output statement if we have actions
     if is_extraction_block and any(
         act["action_type"] == "extract"

skyvern/core/script_generations/skyvern_page.py

@@ -14,7 +14,7 @@ from playwright.async_api import Page
 
 from skyvern.config import settings
 from skyvern.constants import SPECIAL_FIELD_VERIFICATION_CODE
-from skyvern.exceptions import WorkflowRunNotFound
+from skyvern.exceptions import ScriptTerminationException, WorkflowRunNotFound
 from skyvern.forge import app
 from skyvern.forge.prompts import prompt_engine
 from skyvern.forge.sdk.api.files import download_file
@@ -28,12 +28,18 @@ from skyvern.webeye.actions.action_types import ActionType
 from skyvern.webeye.actions.actions import (
     Action,
     ActionStatus,
+    CompleteAction,
     ExtractAction,
     InputTextAction,
     SelectOption,
     SolveCaptchaAction,
 )
-from skyvern.webeye.actions.handler import ActionHandler, handle_input_text_action, handle_select_option_action
+from skyvern.webeye.actions.handler import (
+    ActionHandler,
+    handle_complete_action,
+    handle_input_text_action,
+    handle_select_option_action,
+)
 from skyvern.webeye.actions.parse_actions import parse_actions
 from skyvern.webeye.browser_factory import BrowserState
 from skyvern.webeye.scraper.scraper import ScrapedPage, scrape_website
@@ -791,11 +797,33 @@ class SkyvernPage:
         return
 
     @action_wrap(ActionType.COMPLETE)
-    async def complete(
-        self, data_extraction_goal: str, intention: str | None = None, data: str | dict[str, Any] | None = None
-    ) -> None:
-        # TODO: update the workflow run status to completed
-        return
+    async def complete(self, intention: str | None = None, data: str | dict[str, Any] | None = None) -> None:
+        # TODO: add validation here. if it doesn't pass the validation criteria:
+        #  1. terminate the workflow run if fallback to ai is false
+        #  2. fallback to ai if fallback to ai is true
+        context = skyvern_context.current()
+        if (
+            not context
+            or not context.organization_id
+            or not context.workflow_run_id
+            or not context.task_id
+            or not context.step_id
+        ):
+            return
+        task = await app.DATABASE.get_task(context.task_id, context.organization_id)
+        step = await app.DATABASE.get_step(context.step_id, context.organization_id)
+        if task and step:
+            action = CompleteAction(
+                organization_id=context.organization_id,
+                task_id=context.task_id,
+                step_id=context.step_id,
+                step_order=step.order,
+                action_order=context.action_order,
+            )
+            # result = await ActionHandler.handle_action(self.scraped_page, task, step, self.page, action)
+            result = await handle_complete_action(action, self.page, self.scraped_page, task, step)
+            if result and result[-1].success is False:
+                raise ScriptTerminationException(result[-1].exception_message)
 
     @action_wrap(ActionType.RELOAD_PAGE)
     async def reload_page(self, intention: str | None = None, data: str | dict[str, Any] | None = None) -> None:

skyvern/forge/api_app.py

@@ -1,7 +1,6 @@
 import uuid
-from contextlib import asynccontextmanager
 from datetime import datetime
-from typing import Any, AsyncGenerator, Awaitable, Callable
+from typing import Awaitable, Callable
 
 import structlog
 from fastapi import FastAPI, Response, status
@@ -51,20 +50,12 @@ def custom_openapi() -> dict:
     return app.openapi_schema
 
 
-@asynccontextmanager
-async def lifespan(_: FastAPI) -> AsyncGenerator[None, Any]:
-    """Lifespan context manager for FastAPI app startup and shutdown."""
-    LOG.info("Server started")
-    yield
-    LOG.info("Server shutting down")
-
-
 def get_agent_app() -> FastAPI:
     """
     Start the agent server.
     """
 
-    app = FastAPI(lifespan=lifespan)
+    app = FastAPI()
 
     # Add CORS middleware
     app.add_middleware(

skyvern/forge/sdk/workflow/service.py

@@ -23,6 +23,7 @@ from skyvern.exceptions import (
     FailedToSendWebhook,
     InvalidCredentialId,
     MissingValueForParameter,
+    ScriptTerminationException,
     SkyvernException,
     WorkflowNotFound,
     WorkflowRunNotFound,
@@ -678,9 +679,17 @@ class WorkflowService:
 
                         LOG.debug("Executing run_signature wrapper", wrapper_code=wrapper_code)
 
-                        exec_code = compile(wrapper_code, "<run_signature>", "exec")
-                        exec(exec_code, exec_globals)
-                        output_value = await exec_globals["__run_signature_wrapper"]()
+                        try:
+                            exec_code = compile(wrapper_code, "<run_signature>", "exec")
+                            exec(exec_code, exec_globals)
+                            output_value = await exec_globals["__run_signature_wrapper"]()
+                        except ScriptTerminationException as e:
+                            LOG.warning(
+                                "Script termination",
+                                block_label=block.label,
+                                error=str(e),
+                                exc_info=True,
+                            )
 
                         # Execution succeeded - get the block result from the workflow run blocks
                         # The script execution should have created the workflow run block

skyvern/services/script_service.py

@@ -807,12 +807,16 @@ async def _fallback_to_ai_run(
         # Update block status to completed if workflow block was created
         if workflow_run_block_id:
             # refresh the task
+            failure_reason = None
             refreshed_task = await app.DATABASE.get_task(task_id=task_id, organization_id=organization_id)
             if refreshed_task:
                 task = refreshed_task
+            if task.status in [TaskStatus.terminated, TaskStatus.failed]:
+                failure_reason = task.failure_reason
             await _update_workflow_block(
                 workflow_run_block_id,
                 BlockStatus(task.status.value),
+                failure_reason=failure_reason,
                 label=cache_key,
             )