From 54dadcf3706f1b697185d17a0cee975453e16be8 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Tue, 25 Mar 2025 02:47:02 -0700
Subject: [PATCH] code block validation (#2017)

---
 skyvern/forge/agent_functions.py           | 6 +++++-
 skyvern/forge/sdk/workflow/models/block.py | 4 +---
 2 files changed, 6 insertions(+), 4 deletions(-)

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