diff --git a/skyvern/forge/agent_functions.py b/skyvern/forge/agent_functions.py index 192ad10a..43ee875f 100644 --- a/skyvern/forge/agent_functions.py +++ b/skyvern/forge/agent_functions.py @@ -9,7 +9,7 @@ from playwright.async_api import Frame, Page from skyvern.config import settings from skyvern.constants import SKYVERN_ID_ATTR -from skyvern.exceptions import StepUnableToExecuteError, TaskAlreadyTimeout +from skyvern.exceptions import DisabledBlockExecutionError, StepUnableToExecuteError, TaskAlreadyTimeout from skyvern.forge import app from skyvern.forge.async_operations import AsyncOperation from skyvern.forge.prompts import prompt_engine @@ -503,3 +503,7 @@ class AgentFunction: return element_tree return cleanup_element_tree_func + + async def validate_code_block(self, organization_id: str | None = None) -> None: + if not settings.ENABLE_CODE_BLOCK: + raise DisabledBlockExecutionError("CodeBlock is disabled") diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py index f84005cf..8b7a234c 100644 --- a/skyvern/forge/sdk/workflow/models/block.py +++ b/skyvern/forge/sdk/workflow/models/block.py @@ -30,7 +30,6 @@ from skyvern.config import settings from skyvern.constants import GET_DOWNLOADED_FILES_TIMEOUT, MAX_UPLOAD_FILE_COUNT from skyvern.exceptions import ( ContextParameterValueNotFound, - DisabledBlockExecutionError, MissingBrowserState, MissingBrowserStatePage, SkyvernException, @@ -1192,8 +1191,7 @@ async def wrapper(): browser_session_id: str | None = None, **kwargs: dict, ) -> BlockResult: - if not settings.ENABLE_CODE_BLOCK: - raise DisabledBlockExecutionError("CodeBlock is disabled") + await app.AGENT_FUNCTION.validate_code_block(organization_id=organization_id) # TODO: only support to use code block to manupilate the browser page # support browser context in the future