Job code AI assistant: Show statuses streamed after the text answer

CHANGELOG.md

@@ -17,6 +17,12 @@ and this project adheres to
 
 ### Added
 
+- The job code AI assistant now shows the progress statuses (e.g. "Writing
+  code...") that Apollo streams _after_ the text answer while it generates code,
+  displayed below the answer in the same style as the initial "Thinking..."
+  indicator. Statuses are surfaced in whatever order Apollo sends them.
+  [#PR](https://github.com/OpenFn/lightning/pull/PR)
+
 ### Changed
 
 - Stop reporting expected credential-resolution failures (OAuth re-auth needed,
@@ -38,7 +44,8 @@ and this project adheres to
   the whole run. The search vector is now built off the insert path by a
   background `Lightning.Invocation.DataclipSearchVectorWorker` (sharing the
   `search_indexing` queue with the log-lines worker), making dataclip search
-  eventually consistent. [#4800](https://github.com/OpenFn/lightning/issues/4800)
+  eventually consistent.
+  [#4800](https://github.com/OpenFn/lightning/issues/4800)
 - Channel join crashes when multiple users open the same workflow concurrently
   [#4802](https://github.com/OpenFn/lightning/issues/4802)
 - Fix `purge_deleted` Oban job crashing when a soft-deleted project has

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

@@ -546,6 +546,25 @@ export function MessageList({
                     className={PROSE_CLASSES}
                   />
 
+                  {/* Status (e.g. "Generating code...") Apollo may stream
+                      after the text answer, while we wait for code. Same
+                      visual as the pre-text loading indicator. */}
+                  {isStreaming(message) && streamingStatus && (
+                    <div
+                      className="flex items-center gap-2"
+                      data-testid="streaming-status"
+                    >
+                      <div className="flex items-center gap-1">
+                        <span className="inline-block w-1.5 h-1.5 bg-gray-400 rounded-full animate-bounce" />
+                        <span className="inline-block w-1.5 h-1.5 bg-gray-400 rounded-full animate-bounce [animation-delay:0.15s]" />
+                        <span className="inline-block w-1.5 h-1.5 bg-gray-400 rounded-full animate-bounce [animation-delay:0.3s]" />
+                      </div>
+                      <span className="text-xs text-gray-400 italic">
+                        {streamingStatus}
+                      </span>
+                    </div>
+                  )}
+
                   {!isStreaming(message) && message.code && (
                     <div className="rounded-lg overflow-hidden border border-gray-200 bg-white">
                       <div

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

@@ -152,6 +152,14 @@ export class AIChannelRegistry {
    * Buffer a streaming text chunk and start draining word-by-word.
    */
   private bufferStreamingChunk(content: string): void {
+    // A text chunk arriving over the wire supersedes any active status
+    // (e.g. "Thinking...", "Writing code..."). Clearing here — at network
+    // arrival, not in the slow char-by-char drain — means a status Apollo
+    // streams *after* the text answer survives, while one followed by more
+    // text is correctly dismissed.
+    if (this.store.getSnapshot().streamingStatus) {
+      this.store.setStreamingStatus(null);
+    }
     this.streamingBuffer += content;
     this.startDraining();
   }

assets/js/collaborative-editor/stores/createAIAssistantStore.ts

@@ -577,13 +577,11 @@ export const createAIAssistantStore = (): AIAssistantStore => {
   const _appendStreamingChunk = (content: string) => {
     state = produce(state, draft => {
       draft.streamingContent = (draft.streamingContent || '') + content;
-      // Clear status (e.g. "Thinking...") once actual content starts arriving
-      draft.streamingStatus = null;
     });
     notify('_appendStreamingChunk');
   };
 
-  const setStreamingStatus = (text: string) => {
+  const setStreamingStatus = (text: string | null) => {
     state = produce(state, draft => {
       draft.streamingStatus = text;
     });

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

@@ -199,7 +199,7 @@ export interface AIAssistantStore {
   ) => void;
   _setProcessingState: (isProcessing: boolean) => void;
   _appendStreamingChunk: (content: string) => void;
-  setStreamingStatus: (text: string) => void;
+  setStreamingStatus: (text: string | null) => void;
   _setStreamingChanges: (changes: Record<string, unknown>) => void;
   _clearStreaming: () => void;
   _connectChannel: (channelProvider: unknown) => () => void;

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

@@ -155,6 +155,44 @@ describe('MessageList', () => {
       expect(screen.getByText('Question')).toBeInTheDocument();
       expect(screen.queryByTestId('loading-indicator')).not.toBeInTheDocument();
     });
+
+    it('should show streaming status below the text answer once content has streamed', () => {
+      const messages = [
+        createMockAIMessage({ role: 'user', content: 'Question' }),
+      ];
+
+      render(
+        <MessageList
+          messages={messages}
+          isLoading
+          streamingContent="Here is the answer"
+          streamingStatus="Generating code..."
+        />
+      );
+
+      // Pre-text loading indicator is gone once content streams in
+      expect(screen.queryByTestId('loading-indicator')).not.toBeInTheDocument();
+
+      // Status renders below the streamed text, reusing the bouncing dots
+      const status = screen.getByTestId('streaming-status');
+      expect(status).toHaveTextContent('Generating code...');
+      expect(status.querySelectorAll('.animate-bounce')).toHaveLength(3);
+    });
+
+    it('should not show streaming status when none is set', () => {
+      const messages = [
+        createMockAIMessage({ role: 'user', content: 'Question' }),
+      ];
+
+      render(
+        <MessageList
+          messages={messages}
+          streamingContent="Here is the answer"
+        />
+      );
+
+      expect(screen.queryByTestId('streaming-status')).not.toBeInTheDocument();
+    });
   });
 
   describe('Message Status', () => {

assets/test/collaborative-editor/stores/createAIAssistantStore.test.ts

@@ -336,6 +336,34 @@ describe('createAIAssistantStore', () => {
     });
   });
 
+  describe('Streaming Status', () => {
+    it('should not clear the status when content is appended', () => {
+      // Clearing on new text is the channel buffer's job (at network
+      // arrival), not the drain — so the store must leave status alone.
+      store.setStreamingStatus('Writing code...');
+      store._appendStreamingChunk('Here is the answer');
+
+      expect(store.getSnapshot().streamingStatus).toBe('Writing code...');
+      expect(store.getSnapshot().streamingContent).toBe('Here is the answer');
+    });
+
+    it('should clear the status when set to null', () => {
+      store.setStreamingStatus('Thinking...');
+      store.setStreamingStatus(null);
+
+      expect(store.getSnapshot().streamingStatus).toBeNull();
+    });
+
+    it('should clear the status once changes arrive', () => {
+      store._appendStreamingChunk('Answer');
+      store.setStreamingStatus('Writing code...');
+
+      store._setStreamingChanges({ code: 'fn(s => s)' });
+
+      expect(store.getSnapshot().streamingStatus).toBeNull();
+    });
+  });
+
   describe('State Subscriptions', () => {
     it('should notify subscribers on state changes', () => {
       const subscriber = vi.fn();