diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py index a2386c70..e4167df8 100644 --- a/skyvern/forge/sdk/workflow/models/block.py +++ b/skyvern/forge/sdk/workflow/models/block.py @@ -2989,6 +2989,26 @@ class HttpRequestBlock(Block): ) +def get_all_blocks(blocks: list[BlockTypeVar]) -> list[BlockTypeVar]: + """ + Recursively get "all blocks" in a workflow definition. + + At time of writing, blocks can be nested via the ForLoop block. This function + returns all blocks, flattened. + """ + + all_blocks: list[BlockTypeVar] = [] + + for block in blocks: + all_blocks.append(block) + + if block.block_type == BlockType.FOR_LOOP: + nested_blocks = get_all_blocks(block.loop_blocks) + all_blocks.extend(nested_blocks) + + return all_blocks + + BlockSubclasses = Union[ ForLoopBlock, TaskBlock, diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index 63613dd6..c26b501d 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -63,6 +63,7 @@ from skyvern.forge.sdk.workflow.models.block import ( UrlBlock, ValidationBlock, WaitBlock, + get_all_blocks, ) from skyvern.forge.sdk.workflow.models.parameter import ( PARAMETER_TYPE, @@ -330,12 +331,12 @@ class WorkflowService: ) return workflow_run - all_blocks = workflow.workflow_definition.blocks + top_level_blocks = workflow.workflow_definition.blocks + all_blocks = get_all_blocks(top_level_blocks) if block_labels and len(block_labels): blocks: list[BlockTypeVar] = [] all_labels = {block.label: block for block in all_blocks} - for label in block_labels: if label not in all_labels: raise BlockNotFound(block_label=label) @@ -350,7 +351,7 @@ class WorkflowService: ) else: - blocks = all_blocks + blocks = top_level_blocks if not blocks: raise SkyvernException(f"No blocks found for the given block labels: {block_labels}")