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(