From 9b32a77c59231324c39516fde2aa3f74a3df6516 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Mon, 14 Apr 2025 09:41:59 -0700 Subject: [PATCH] set task_id in skyvern_context when task block starts (#2114) --- skyvern/forge/sdk/workflow/models/block.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py index 8b7a234c..53caa843 100644 --- a/skyvern/forge/sdk/workflow/models/block.py +++ b/skyvern/forge/sdk/workflow/models/block.py @@ -48,6 +48,7 @@ from skyvern.forge.sdk.api.files import ( ) from skyvern.forge.sdk.api.llm.api_handler_factory import LLMAPIHandlerFactory from skyvern.forge.sdk.artifact.models import ArtifactType +from skyvern.forge.sdk.core import skyvern_context from skyvern.forge.sdk.db.enums import TaskType from skyvern.forge.sdk.schemas.files import FileInfo from skyvern.forge.sdk.schemas.task_v2 import TaskV2Status @@ -612,6 +613,8 @@ class BaseTaskBlock(Block): raise e try: + current_context = skyvern_context.ensure_context() + current_context.task_id = task.task_id await app.agent.execute_step( organization=organization, task=task, @@ -630,6 +633,8 @@ class BaseTaskBlock(Block): failure_reason=str(e), ) raise e + finally: + current_context.task_id = None # Check task status updated_task = await app.DATABASE.get_task(