diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py index 3836cca5..022dd1c8 100644 --- a/skyvern/forge/sdk/workflow/models/block.py +++ b/skyvern/forge/sdk/workflow/models/block.py @@ -1,8 +1,10 @@ import abc +import asyncio import csv import json import os import smtplib +import textwrap import uuid from dataclasses import dataclass from email.message import EmailMessage @@ -501,8 +503,13 @@ class CodeBlock(Block): async def execute(self, workflow_run_id: str, **kwargs: dict) -> BlockResult: # get workflow run context workflow_run_context = self.get_workflow_run_context(workflow_run_id) + # get all parameters into a dictionary parameter_values = {} + maybe_browser_state = await app.BROWSER_MANAGER.get_for_workflow_run(workflow_run_id) + if maybe_browser_state and maybe_browser_state.page: + parameter_values["skyvern_page"] = maybe_browser_state.page + for parameter in self.parameters: value = workflow_run_context.get_value(parameter.key) secret_value = workflow_run_context.get_original_secret_value_or_none(value) @@ -511,9 +518,25 @@ class CodeBlock(Block): else: parameter_values[parameter.key] = value + # Add asyncio and the current event loop to the parameter_values + parameter_values["asyncio"] = asyncio + parameter_values["__builtins__"] = __builtins__ # Include builtins for exec context + local_variables: dict[str, Any] = {} - exec(self.code, parameter_values, local_variables) - result = {"result": local_variables.get("result")} + result_container: dict[str, Any] = {} + # Define the user_code function and return it + user_code = textwrap.indent(self.code, " ") + full_code = f""" +async def user_code(): +{user_code} + result_container['result'] = locals().get('result') +""" + + exec(full_code, {**parameter_values, "result_container": result_container}, local_variables) + # Await the returned user_code function + await local_variables["user_code"]() + + result = {"result": result_container.get("result")} await self.record_output_parameter_value(workflow_run_context, workflow_run_id, result) return self.build_block_result(success=True, output_parameter_value=result) diff --git a/skyvern/webeye/browser_manager.py b/skyvern/webeye/browser_manager.py index 4a0a3c35..1d527081 100644 --- a/skyvern/webeye/browser_manager.py +++ b/skyvern/webeye/browser_manager.py @@ -85,6 +85,11 @@ class BrowserManager: self.pages[workflow_run.workflow_run_id] = browser_state return browser_state + async def get_for_workflow_run(self, workflow_run_id: str) -> BrowserState | None: + if workflow_run_id in self.pages: + return self.pages[workflow_run_id] + return None + def set_video_artifact_for_task(self, task: Task, artifact_id: str) -> None: if task.workflow_run_id and task.workflow_run_id in self.pages: if self.pages[task.workflow_run_id].browser_artifacts.video_artifact_id: