fix workflow_runs.code_gen to enforce code generate only for task v2 block run (#3597)

skyvern/services/block_service.py

@@ -45,6 +45,7 @@ async def ensure_workflow_run(
         max_steps=x_max_steps_override,
         request_id=context.request_id,
         debug_session_id=block_run_request.debug_session_id,
+        code_gen=block_run_request.code_gen,
     )
 
     return workflow_run
@@ -62,7 +63,6 @@ async def execute_blocks(
     user_id: str,
     browser_session_id: str | None = None,
     block_outputs: dict[str, t.Any] | None = None,
-    code_gen: bool | None = None,
 ) -> None:
     """
     Runs one or more blocks of a workflow.
@@ -115,5 +115,4 @@ async def execute_blocks(
         api_key=api_key,
         block_labels=block_labels,
         block_outputs=block_outputs,
-        code_gen=code_gen,
     )

skyvern/services/task_v2_service.py

@@ -306,7 +306,6 @@ async def run_task_v2(
     request_id: str | None = None,
     max_steps_override: str | int | None = None,
     browser_session_id: str | None = None,
-    code_gen: bool | None = None,
 ) -> TaskV2:
     organization_id = organization.organization_id
     try:
@@ -335,7 +334,6 @@ async def run_task_v2(
             request_id=request_id,
             max_steps_override=max_steps_override,
             browser_session_id=browser_session_id,
-            code_gen=code_gen,
         )
     except TaskTerminationError as e:
         task_v2 = await mark_task_v2_as_terminated(
@@ -392,7 +390,6 @@ async def run_task_v2_helper(
     request_id: str | None = None,
     max_steps_override: str | int | None = None,
     browser_session_id: str | None = None,
-    code_gen: bool | None = None,
 ) -> tuple[Workflow, WorkflowRun, TaskV2] | tuple[None, None, TaskV2]:
     organization_id = organization.organization_id
     task_v2_id = task_v2.observer_cruise_id
@@ -800,7 +797,6 @@ async def run_task_v2_helper(
         block_result = await block.execute_safe(
             workflow_run_id=workflow_run_id,
             organization_id=organization_id,
-            code_gen=code_gen,
         )
         task_history_record["status"] = str(block_result.status)
         if block_result.failure_reason:
@@ -904,11 +900,10 @@ async def run_task_v2_helper(
                     context=context,
                     screenshots=completion_screenshots,
                 )
-                if task_v2.run_with == "code" or code_gen:  # TODO(jdo): not sure about this one...
+                if task_v2.run_with == "code":
                     await app.WORKFLOW_SERVICE.generate_script_if_needed(
                         workflow=workflow,
                         workflow_run=workflow_run,
-                        code_gen=code_gen,
                     )
                 break
 

skyvern/services/workflow_service.py

@@ -23,6 +23,7 @@ async def prepare_workflow(
     max_steps: int | None = None,
     request_id: str | None = None,
     debug_session_id: str | None = None,
+    code_gen: bool | None = None,
 ) -> WorkflowRun:
     """
     Prepare a workflow to be run.
@@ -40,6 +41,7 @@ async def prepare_workflow(
         max_steps_override=max_steps,
         is_template_workflow=template,
         debug_session_id=debug_session_id,
+        code_gen=code_gen,
     )
 
     workflow = await app.WORKFLOW_SERVICE.get_workflow_by_permanent_id(