Global Chat: Receive job code for multiple steps

CHANGELOG.md

@@ -19,11 +19,16 @@ and this project adheres to
 
 ### Changed
 
+- Global chat can now change multiple workflow steps in a single response. It
+  receives a full workflow YAML from the Apollo AI server with each step's job
+  code embedded, and applies the changes together. When a step is open in the
+  editor, its diff is previewed before applying; previewing several step diffs
+  at once is a follow-up.
+  [#4890](https://github.com/OpenFn/lightning/issues/4890)
 - Consolidated run and work order state definitions into single source of truth
   by adding `Run.active_states/0`, `WorkOrder.states/0`, and
   `WorkOrder.active_states/0` and replacing all hardcoded state lists across the
-  codebase
-  [#4589](https://github.com/OpenFn/lightning/issues/4589)
+  codebase [#4589](https://github.com/OpenFn/lightning/issues/4589)
 
 ### Fixed
 

assets/js/collaborative-editor/components/AIAssistantPanelWrapper.tsx

@@ -560,36 +560,55 @@ export function AIAssistantPanelWrapper({
   );
 
   // Hook to handle workflow/job code application logic
-  const { handleApplyWorkflow, handlePreviewJobCode, handleApplyJobCode } =
-    useAIWorkflowApplications({
-      sessionId,
-      page: aiMode?.page || 'workflow_template',
-      currentSession:
-        sessionId && messages.length > 0
-          ? {
-              messages,
-              workflowTemplateContext,
-            }
-          : null,
-      currentUserId: user?.id,
-      aiMode,
-      workflowActions: {
-        importWorkflow,
-        startApplyingWorkflow,
-        doneApplyingWorkflow,
-        startApplyingJobCode,
-        doneApplyingJobCode,
-        updateJob,
-      },
-      monacoRef,
-      jobs,
-      canApplyChanges,
-      connectionState,
-      setPreviewingMessageId,
-      previewingMessageId,
-      setApplyingMessageId,
-      appliedMessageIdsRef,
-    });
+  const {
+    handleApplyWorkflow,
+    handlePreviewJobCode,
+    handlePreviewGlobalStep,
+    handleApplyJobCode,
+  } = useAIWorkflowApplications({
+    sessionId,
+    page: aiMode?.page || 'workflow_template',
+    currentSession:
+      sessionId && messages.length > 0
+        ? {
+            messages,
+            workflowTemplateContext,
+          }
+        : null,
+    currentUserId: user?.id,
+    aiMode,
+    workflowActions: {
+      importWorkflow,
+      startApplyingWorkflow,
+      doneApplyingWorkflow,
+      startApplyingJobCode,
+      doneApplyingJobCode,
+      updateJob,
+    },
+    monacoRef,
+    jobs,
+    canApplyChanges,
+    connectionState,
+    setPreviewingMessageId,
+    previewingMessageId,
+    setApplyingMessageId,
+    appliedMessageIdsRef,
+  });
+
+  // Route auto-preview to the right handler: global messages carry a full
+  // workflow YAML (the open step's diff is extracted from it), job-code
+  // messages carry the job body directly.
+  const handleAutoPreview = useCallback(
+    (code: string, messageId: string) => {
+      const message = messages.find(m => m.id === messageId);
+      if (message?.from_global) {
+        handlePreviewGlobalStep(code, messageId);
+      } else {
+        handlePreviewJobCode(code, messageId);
+      }
+    },
+    [messages, handlePreviewGlobalStep, handlePreviewJobCode]
+  );
 
   // Auto-preview job code when AI responds with code
   // Only for the user who authored the triggering message
@@ -599,7 +618,7 @@ export function AIAssistantPanelWrapper({
       ? { id: sessionId, session_type: 'workflow_template', messages }
       : null,
     currentUserId: user?.id,
-    onPreview: handlePreviewJobCode,
+    onPreview: handleAutoPreview,
   });
 
   // Auto-apply streaming changes as soon as they arrive (before text finishes)
@@ -707,7 +726,12 @@ export function AIAssistantPanelWrapper({
                 messages={messages}
                 isLoading={isLoading}
                 onApplyWorkflow={
-                  aiMode?.page === 'workflow_template' && !isApplyingWorkflow
+                  (aiMode?.page === 'workflow_template' ||
+                    // Global messages apply the full workflow even while a
+                    // job is open; handleApplyWorkflow still no-ops for
+                    // non-global messages outside workflow_template mode.
+                    messages.some(m => m.from_global && m.code)) &&
+                  !isApplyingWorkflow
                     ? (yaml, messageId) => {
                         void handleApplyWorkflow(yaml, messageId);
                       }
@@ -723,6 +747,8 @@ export function AIAssistantPanelWrapper({
                 onPreviewJobCode={
                   aiMode?.page === 'job_code' ? handlePreviewJobCode : undefined
                 }
+                onPreviewGlobalStep={handlePreviewGlobalStep}
+                canPreviewGlobalStep={aiMode?.page === 'job_code'}
                 applyingMessageId={
                   // If anyone is applying (including other users), pass the message ID
                   // to show "APPLYING..." state. Prioritize stored message ID from store,

assets/js/collaborative-editor/components/MessageList.tsx

@@ -384,6 +384,10 @@ interface MessageListProps {
   onApplyWorkflow?: ((yaml: string, messageId: string) => void) | undefined;
   onApplyJobCode?: ((code: string, messageId: string) => void) | undefined;
   onPreviewJobCode?: ((code: string, messageId: string) => void) | undefined;
+  /** Per-step diff preview for global messages (extracts the open job's body from the YAML) */
+  onPreviewGlobalStep?: ((yaml: string, messageId: string) => void) | undefined;
+  /** Whether a job is open in the IDE, enabling preview for global messages */
+  canPreviewGlobalStep?: boolean;
   applyingMessageId?: string | null | undefined;
   previewingMessageId?: string | null | undefined;
   showAddButtons?: boolean;
@@ -400,6 +404,8 @@ export function MessageList({
   onApplyWorkflow,
   onApplyJobCode,
   onPreviewJobCode,
+  onPreviewGlobalStep,
+  canPreviewGlobalStep = false,
   applyingMessageId,
   previewingMessageId,
   showAddButtons = false,
@@ -589,7 +595,10 @@ export function MessageList({
                           code={message.code}
                           showAdd={showAddButtons}
                           showApply={showApplyButton}
-                          showPreview={!!message.job_id}
+                          showPreview={
+                            !!message.job_id ||
+                            (!!message.from_global && canPreviewGlobalStep)
+                          }
                           onApply={() => {
                             if (message.job_id) {
                               onApplyJobCode?.(message.code!, message.id);
@@ -598,7 +607,12 @@ export function MessageList({
                             }
                           }}
                           onPreview={() => {
-                            onPreviewJobCode?.(message.code!, message.id);
+                            if (message.from_global) {
+                              // Per-step diff from the full workflow YAML
+                              onPreviewGlobalStep?.(message.code!, message.id);
+                            } else {
+                              onPreviewJobCode?.(message.code!, message.id);
+                            }
                           }}
                           isApplying={!!applyingMessageId}
                           isPreviewActive={previewingMessageId === message.id}

assets/js/collaborative-editor/hooks/useAIWorkflowApplications.ts

@@ -170,7 +170,14 @@ export function useAIWorkflowApplications({
    */
   const handleApplyWorkflow = useCallback(
     async (yaml: string, messageId: string) => {
-      if (!aiMode || aiMode.page !== 'workflow_template') {
+      if (!aiMode) return;
+      // Global messages carry a full workflow YAML and may be applied even
+      // while a job is open (job_code mode). Non-global workflow chat keeps
+      // the workflow_template-only guard so its Apply stays a no-op when a
+      // job is open.
+      const isGlobal = !!currentSession?.messages.find(m => m.id === messageId)
+        ?.from_global;
+      if (aiMode.page !== 'workflow_template' && !isGlobal) {
         console.error(
           '[AI Assistant] Cannot apply workflow - not in workflow mode',
           {
@@ -179,6 +186,15 @@ export function useAIWorkflowApplications({
         );
         return;
       }
+
+      // A global message applied while a step is open leaves an active diff in
+      // the open step. Clear it so the editor returns to an editable state.
+      const monaco = monacoRef?.current;
+      if (previewingMessageId && monaco) {
+        monaco.clearDiff();
+        setPreviewingMessageId(null);
+      }
+
       setApplyingMessageId(messageId);
 
       // Signal to all collaborators that we're starting to apply
@@ -218,12 +234,16 @@ export function useAIWorkflowApplications({
       }
     },
     [
-      aiMode?.page,
+      aiMode,
+      currentSession,
       importWorkflow,
       startApplyingWorkflow,
       doneApplyingWorkflow,
       jobs,
       setApplyingMessageId,
+      monacoRef,
+      previewingMessageId,
+      setPreviewingMessageId,
     ]
   );
 
@@ -295,6 +315,83 @@ export function useAIWorkflowApplications({
     [aiMode, jobs, previewingMessageId, monacoRef, setPreviewingMessageId]
   );
 
+  /**
+   * Preview the open job's diff from a global full-workflow YAML message
+   *
+   * Mirrors handlePreviewJobCode, but extracts the open job's body from the
+   * workflow YAML (global messages carry the whole workflow in `code`).
+   * Shows a diff only when the open step's body actually changed; clears any
+   * stale diff otherwise.
+   */
+  const handlePreviewGlobalStep = useCallback(
+    (yaml: string, messageId: string) => {
+      if (!aiMode || aiMode.page !== 'job_code') return; // only when a step is open
+      const jobId = (aiMode.context as JobCodeContext).job_id;
+      if (!jobId) return;
+
+      // Same dedup guards as handlePreviewJobCode
+      if (previewingMessageId === messageId) return;
+      if (previewingMessageId === '__streaming__') {
+        setPreviewingMessageId(messageId);
+        return;
+      }
+
+      const currentJob = jobs.find(j => j.id === jobId);
+      const currentBody = currentJob?.body ?? '';
+
+      let newBody: string | undefined;
+      try {
+        const spec = parseWorkflowYAML(yaml);
+        // ids from the YAML are preserved, so we match the open step by id
+        const state = convertWorkflowSpecToState(spec);
+        newBody = state.jobs.find(j => j.id === jobId)?.body;
+      } catch (error) {
+        console.error(
+          '[AI Assistant] Failed to parse global workflow YAML:',
+          error
+        );
+        notifications.alert({
+          title: 'Could not preview step',
+          description:
+            error instanceof Error
+              ? error.message
+              : 'The AI server returned invalid workflow YAML.',
+        });
+        return;
+      }
+
+      if (newBody === undefined) {
+        // Open step's id wasn't in the YAML, so the server likely didn't preserve it
+        console.warn(
+          '[AI Assistant] Open step not found in global workflow YAML',
+          { jobId }
+        );
+        notifications.warning({
+          title: 'Could not preview this step',
+          description: `Step "${
+            currentJob?.name ?? jobId
+          }" was not found in the AI response (id: ${jobId}). Its ID may not have been preserved by the server.`,
+        });
+        if (previewingMessageId) monacoRef?.current?.clearDiff();
+        return;
+      }
+
+      if (newBody === currentBody) {
+        // open step genuinely unchanged -> ensure no stale diff is shown
+        if (previewingMessageId) monacoRef?.current?.clearDiff();
+        return;
+      }
+
+      const monaco = monacoRef?.current;
+      if (previewingMessageId && monaco) monaco.clearDiff();
+      if (monaco) {
+        monaco.showDiff(currentBody, newBody);
+        setPreviewingMessageId(messageId);
+      }
+    },
+    [aiMode, jobs, previewingMessageId, monacoRef, setPreviewingMessageId]
+  );
+
   /**
    * Apply job code to Y.Doc
    *
@@ -453,6 +550,7 @@ export function useAIWorkflowApplications({
   return {
     handleApplyWorkflow,
     handlePreviewJobCode,
+    handlePreviewGlobalStep,
     handleApplyJobCode,
   };
 }

assets/js/collaborative-editor/hooks/useAutoPreview.ts

@@ -94,9 +94,10 @@ export function useAutoPreview({
           new Date(b.inserted_at).getTime() - new Date(a.inserted_at).getTime()
       )[0];
 
-    if (!latestCodeMessage || !latestCodeMessage.job_id) {
-      return;
-    }
+    if (!latestCodeMessage) return;
+    // Global messages have no job_id but carry a full workflow YAML; the
+    // onPreview callback handles extracting the open step's body from it.
+    if (!latestCodeMessage.job_id && !latestCodeMessage.from_global) return;
 
     // Skip if we've already auto-previewed this message
     if (stateRef.current.lastAutoPreviewedMessageId === latestCodeMessage.id) {

assets/js/collaborative-editor/lib/AIChannelRegistry.ts

@@ -907,14 +907,23 @@ export class AIChannelRegistry {
       if (context.workflow_id) {
         params['workflow_id'] = context.workflow_id;
       }
-      if (context.code) {
-        params['code'] = context.code;
-      }
       if (context.content) {
         params['content'] = context.content;
       }
     }
 
+    // Workflow YAML (applicable to both session types). For global chat the
+    // `code` slot carries the FULL serialized workflow YAML (every step body
+    // embedded), not a single job's code. When a step is open the context is
+    // JobCodeContext-shaped and took the branch above, which does not forward
+    // `code`. Forwarding it here (outside the branch) ensures the YAML reaches
+    // Apollo on the *first* turn — the only message sent via the channel join.
+    // Later turns go through `new_message` and were never affected. Plain job
+    // chat never sets `context.code`, so this is a no-op there.
+    if ('code' in context && context.code) {
+      params['code'] = context.code;
+    }
+
     // Global assistant flags (applicable to both session types)
     if ('use_global_assistant' in context && context.use_global_assistant) {
       params['use_global_assistant'] = true;

assets/js/collaborative-editor/types/ai-assistant.ts

@@ -54,6 +54,11 @@ export interface Message {
   user_id?: string;
   user?: MessageUser | null;
   job_id?: string;
+  /**
+   * True when this message came from the global AI assistant. Global
+   * messages carry a full workflow YAML in `code` and never a `job_id`.
+   */
+  from_global?: boolean;
 }
 
 /**
@@ -72,6 +77,10 @@ export interface JobCodeContext {
   job_body?: string;
   job_adaptor?: string;
   workflow_id?: string;
+
+  // Full serialized workflow YAML attached by global chat (sent as the message
+  // `code`). Present even with a step open, so it lives on the job context too.
+  code?: string;
 }
 
 /**
@@ -101,6 +110,10 @@ export type WorkflowTemplateContext =
 
       workflow_id?: string;
       project_id: string;
+
+      // Full serialized workflow YAML attached by global chat (sent as the
+      // message `code`), present even when a step is open.
+      code?: string;
     };
 
 /**