From 16f61af6cfe7dbab820604b323297150ffe1ebc7 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Wed, 5 Nov 2025 08:46:03 +0800 Subject: [PATCH] add workflow_run_id and workflow_run_block_id in create/update script_blocks code (#3904) --- skyvern/core/script_generations/generate_script.py | 14 ++++++++++++++ .../script_generations/transform_workflow_run.py | 5 +++++ skyvern/forge/sdk/db/client.py | 10 ++++++++++ skyvern/forge/sdk/db/utils.py | 2 ++ skyvern/schemas/scripts.py | 2 ++ skyvern/services/script_service.py | 2 ++ 6 files changed, 35 insertions(+) diff --git a/skyvern/core/script_generations/generate_script.py b/skyvern/core/script_generations/generate_script.py index 65d8067e..a16e8a55 100644 --- a/skyvern/core/script_generations/generate_script.py +++ b/skyvern/core/script_generations/generate_script.py @@ -1781,6 +1781,8 @@ async def generate_workflow_script_python_code( block_label=block_name, update=pending, run_signature=run_signature, + workflow_run_id=task.get("workflow_run_id"), + workflow_run_block_id=task.get("workflow_run_block_id"), ) except Exception as e: LOG.error("Failed to create script block", error=str(e), exc_info=True) @@ -1837,6 +1839,8 @@ async def generate_workflow_script_python_code( block_label=block_name, update=pending, run_signature=run_signature, + workflow_run_id=task_v2.get("workflow_run_id"), + workflow_run_block_id=task_v2.get("workflow_run_block_id"), ) except Exception as e: LOG.error("Failed to create task_v2 script block", error=str(e), exc_info=True) @@ -1926,6 +1930,8 @@ async def create_or_update_script_block( block_label: str, update: bool = False, run_signature: str | None = None, + workflow_run_id: str | None = None, + workflow_run_block_id: str | None = None, ) -> None: """ Create a script block in the database and save the block code to a script file. @@ -1939,6 +1945,8 @@ async def create_or_update_script_block( block_label: Optional custom name for the block (defaults to function name) update: Whether to update the script block instead of creating a new one run_signature: The function call code to execute this block (e.g., "await skyvern.action(...)") + workflow_run_id: The workflow run that generated this cached block + workflow_run_block_id: The workflow run block that generated this cached block """ block_code_bytes = block_code if isinstance(block_code, bytes) else block_code.encode("utf-8") try: @@ -1955,6 +1963,8 @@ async def create_or_update_script_block( organization_id=organization_id, script_block_label=block_label, run_signature=run_signature, + workflow_run_id=workflow_run_id, + workflow_run_block_id=workflow_run_block_id, ) elif run_signature: # Update the run_signature if provided @@ -1962,6 +1972,8 @@ async def create_or_update_script_block( script_block_id=script_block.script_block_id, organization_id=organization_id, run_signature=run_signature, + workflow_run_id=workflow_run_id, + workflow_run_block_id=workflow_run_block_id, ) # Step 4: Create script file for the block @@ -2011,6 +2023,8 @@ async def create_or_update_script_block( script_block_id=script_block.script_block_id, organization_id=organization_id, script_file_id=script_file.file_id, + workflow_run_id=workflow_run_id, + workflow_run_block_id=workflow_run_block_id, ) except Exception as e: diff --git a/skyvern/core/script_generations/transform_workflow_run.py b/skyvern/core/script_generations/transform_workflow_run.py index 8b46ee63..2a73cff1 100644 --- a/skyvern/core/script_generations/transform_workflow_run.py +++ b/skyvern/core/script_generations/transform_workflow_run.py @@ -155,6 +155,11 @@ async def transform_workflow_run_to_code_gen_input(workflow_run_id: str, organiz else: LOG.warning(f"Task v2 block {run_block.label} does not have a child workflow run id") + final_dump["workflow_run_id"] = workflow_run_id + if run_block: + final_dump["workflow_run_block_id"] = run_block.workflow_run_block_id + else: + final_dump["workflow_run_block_id"] = None workflow_block_dump.append(final_dump) return CodeGenInput( diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 515c2138..fc6fac92 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -4499,6 +4499,8 @@ class AgentDB: script_block_label: str, script_file_id: str | None = None, run_signature: str | None = None, + workflow_run_id: str | None = None, + workflow_run_block_id: str | None = None, ) -> ScriptBlock: """Create a script block.""" async with self.Session() as session: @@ -4509,6 +4511,8 @@ class AgentDB: script_block_label=script_block_label, script_file_id=script_file_id, run_signature=run_signature, + workflow_run_id=workflow_run_id, + workflow_run_block_id=workflow_run_block_id, ) session.add(script_block) await session.commit() @@ -4521,6 +4525,8 @@ class AgentDB: organization_id: str, script_file_id: str | None = None, run_signature: str | None = None, + workflow_run_id: str | None = None, + workflow_run_block_id: str | None = None, ) -> ScriptBlock: async with self.Session() as session: script_block = ( @@ -4535,6 +4541,10 @@ class AgentDB: script_block.script_file_id = script_file_id if run_signature is not None: script_block.run_signature = run_signature + if workflow_run_id is not None: + script_block.workflow_run_id = workflow_run_id + if workflow_run_block_id is not None: + script_block.workflow_run_block_id = workflow_run_block_id await session.commit() await session.refresh(script_block) return convert_to_script_block(script_block) diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index f5c8ca6d..bc2e71fc 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/skyvern/forge/sdk/db/utils.py @@ -584,6 +584,8 @@ def convert_to_script_block(script_block_model: ScriptBlockModel) -> ScriptBlock script_block_label=script_block_model.script_block_label, script_file_id=script_block_model.script_file_id, run_signature=script_block_model.run_signature, + workflow_run_id=script_block_model.workflow_run_id, + workflow_run_block_id=script_block_model.workflow_run_block_id, created_at=script_block_model.created_at, modified_at=script_block_model.modified_at, deleted_at=script_block_model.deleted_at, diff --git a/skyvern/schemas/scripts.py b/skyvern/schemas/scripts.py index 1c62e18d..2481e8a4 100644 --- a/skyvern/schemas/scripts.py +++ b/skyvern/schemas/scripts.py @@ -135,6 +135,8 @@ class ScriptBlock(BaseModel): script_block_label: str script_file_id: str | None = None run_signature: str | None = None # The function call code to execute this block + workflow_run_id: str | None = None + workflow_run_block_id: str | None = None created_at: datetime modified_at: datetime deleted_at: datetime | None = None diff --git a/skyvern/services/script_service.py b/skyvern/services/script_service.py index dd2cdefa..bec6ccda 100644 --- a/skyvern/services/script_service.py +++ b/skyvern/services/script_service.py @@ -1034,6 +1034,8 @@ async def _regenerate_script_block_after_ai_fallback( script_id=new_script.script_id, organization_id=organization_id, block_label=existing_block.script_block_label, + workflow_run_id=existing_block.workflow_run_id, + workflow_run_block_id=existing_block.workflow_run_block_id, ) block_file_content_bytes = ( block_file_content if isinstance(block_file_content, bytes) else block_file_content.encode("utf-8")