diff --git a/skyvern-frontend/src/routes/workflows/editor/Workspace.tsx b/skyvern-frontend/src/routes/workflows/editor/Workspace.tsx
index b5f5288b..d33d4d7e 100644
--- a/skyvern-frontend/src/routes/workflows/editor/Workspace.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/Workspace.tsx
@@ -415,6 +415,18 @@ function Workspace({
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);
 
+  useEffect(() => {
+    if (isFinalized) {
+      queryClient.invalidateQueries({
+        queryKey: ["block-scripts"],
+      });
+
+      queryClient.invalidateQueries({
+        queryKey: ["cache-key-values"],
+      });
+    }
+  }, [isFinalized, queryClient, workflowRun]);
+
   useEffect(() => {
     blockScriptStore.setScripts(blockScriptsPublished ?? {});
     // eslint-disable-next-line react-hooks/exhaustive-deps
diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/components/NodeHeader.tsx b/skyvern-frontend/src/routes/workflows/editor/nodes/components/NodeHeader.tsx
index c4f97276..21b9829e 100644
--- a/skyvern-frontend/src/routes/workflows/editor/nodes/components/NodeHeader.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/nodes/components/NodeHeader.tsx
@@ -17,6 +17,7 @@ import { useNodeLabelChangeHandler } from "@/routes/workflows/hooks/useLabelChan
 import { useDeleteNodeCallback } from "@/routes/workflows/hooks/useDeleteNodeCallback";
 import { useToggleScriptForNodeCallback } from "@/routes/workflows/hooks/useToggleScriptForNodeCallback";
 import { useDebugSessionQuery } from "@/routes/workflows/hooks/useDebugSessionQuery";
+import { useWorkflowQuery } from "@/routes/workflows/hooks/useWorkflowQuery";
 import { useWorkflowRunQuery } from "@/routes/workflows/hooks/useWorkflowRunQuery";
 import {
   debuggableWorkflowBlockTypes,
@@ -67,6 +68,7 @@ type Payload = Record<string, unknown> & {
   totp_url: string | null;
   webhook_url: string | null;
   workflow_id: string;
+  code_gen: boolean | null;
 };
 
 const getPayload = (opts: {
@@ -171,6 +173,9 @@ function NodeHeader({
   const { data: debugSession } = useDebugSessionQuery({
     workflowPermanentId,
   });
+  const { data: workflow } = useWorkflowQuery({
+    workflowPermanentId,
+  });
   const saveWorkflow = useWorkflowSave();
 
   const thisBlockIsPlaying =
@@ -458,7 +463,10 @@ function NodeHeader({
   });
 
   const handleOnPlay = () => {
-    runBlock.mutate({ codeGen: false });
+    const numBlocksInWorkflow = (workflow?.workflow_definition.blocks ?? [])
+      .length;
+
+    runBlock.mutate({ codeGen: numBlocksInWorkflow === 1 });
   };
 
   const handleOnCancel = () => {