preserve temp chat mid-stream and scope Stop to active thread (#2971)

Author: karthikjeeyar

workspaces/lightspeed/.changeset/selfish-adults-obey.md

@@ -0,0 +1,5 @@
+---
+'@red-hat-developer-hub/backstage-plugin-lightspeed': patch
+---
+
+Fix new chat streams when switching threads and scope Stop to the streaming conversation (RHDHBUGS-3040).

workspaces/lightspeed/plugins/lightspeed/src/components/LightSpeedChat.tsx

@@ -835,18 +835,27 @@ export const LightspeedChat = ({
     setNewChatCreated(false);
   };
 
-  const { conversationMessages, handleInputPrompt, scrollToBottomRef } =
-    useConversationMessages(
-      viewConversationId,
-      userName,
-      selectedModel,
-      selectedProvider,
-      avatar,
-      onComplete,
-      onStart,
-      undefined,
-      onRequestIdReady,
-    );
+  const {
+    conversationMessages,
+    handleInputPrompt,
+    scrollToBottomRef,
+    streamingConversationId,
+  } = useConversationMessages(
+    viewConversationId,
+    userName,
+    selectedModel,
+    selectedProvider,
+    avatar,
+    onComplete,
+    onStart,
+    undefined,
+    onRequestIdReady,
+  );
+
+  const streamingUiMatchesView =
+    isSendButtonDisabled &&
+    streamingConversationId !== null &&
+    viewConversationId === streamingConversationId;
 
   const [messages, setMessages] =
     useState<MessageProps[]>(conversationMessages);
@@ -1463,7 +1472,7 @@ export const LightspeedChat = ({
             ref={scrollToBottomRef}
             welcomePrompts={welcomePrompts}
             conversationId={conversationId}
-            isStreaming={isSendButtonDisabled}
+            isStreaming={streamingUiMatchesView}
             topicRestrictionEnabled={topicRestrictionEnabled}
             displayMode={displayMode}
           />
@@ -1487,8 +1496,10 @@ export const LightspeedChat = ({
           hasMicrophoneButton
           value={draftMessage}
           onChange={handleDraftMessage}
-          hasStopButton={isSendButtonDisabled}
-          handleStopButton={isSendButtonDisabled ? handleStopButton : undefined}
+          hasStopButton={streamingUiMatchesView}
+          handleStopButton={
+            streamingUiMatchesView ? handleStopButton : undefined
+          }
           buttonProps={{
             attach: {
               inputTestId: 'attachment-input',

workspaces/lightspeed/plugins/lightspeed/src/hooks/__tests__/useConversationMessages.test.tsx

@@ -840,4 +840,91 @@ data: {"event": "token", "data": {"id": 2, "token": ""}}\n
       ).toBeUndefined();
     });
   });
+
+  it('preserves temp thread when switching conversations mid-stream and returning (RHDHBUGS-3040)', async () => {
+    mockLightspeedApi.getConversationMessages.mockResolvedValue([]);
+    const onComplete = jest.fn();
+
+    let resolveSecondRead!: (value: IteratorResult<Uint8Array | null>) => void;
+    const secondReadPromise = new Promise<IteratorResult<Uint8Array | null>>(
+      resolve => {
+        resolveSecondRead = resolve;
+      },
+    );
+
+    const firstChunk = createSSEStream([
+      { event: 'start', data: { conversation_id: 'persisted-after-stream' } },
+      { event: 'token', data: { id: 0, token: 'Hello ', role: 'inference' } },
+    ]);
+    const secondChunk = createSSEStream([
+      { event: 'token', data: { id: 1, token: 'world!', role: 'inference' } },
+    ]);
+
+    const read = jest
+      .fn()
+      .mockResolvedValueOnce({
+        done: false,
+        value: new TextEncoder().encode(firstChunk),
+      })
+      .mockImplementationOnce(() => secondReadPromise)
+      .mockResolvedValueOnce({ done: true, value: null });
+
+    const lightSpeedApi = {
+      ...mockLightspeedApi,
+      createMessage: jest.fn().mockResolvedValue({ read }),
+    };
+    (useApi as jest.Mock).mockReturnValue(lightSpeedApi);
+
+    const { result, rerender } = renderHook(
+      ({ conversationId }) =>
+        useConversationMessages(
+          conversationId,
+          'test-user',
+          'gpt-3',
+          'openai',
+          'user.png',
+          onComplete,
+        ),
+      {
+        initialProps: { conversationId: TEMP_CONVERSATION_ID },
+        wrapper,
+      },
+    );
+
+    await act(async () => {
+      void result.current.handleInputPrompt('Hi');
+    });
+
+    await waitFor(() => {
+      const msgs = result.current.conversations[TEMP_CONVERSATION_ID];
+      expect(msgs?.[1]?.content).toContain('Hello ');
+    });
+
+    expect(result.current.streamingConversationId).toBe(TEMP_CONVERSATION_ID);
+
+    rerender({ conversationId: 'other-conv-id' });
+
+    expect(result.current.streamingConversationId).toBe(TEMP_CONVERSATION_ID);
+
+    rerender({ conversationId: TEMP_CONVERSATION_ID });
+
+    await waitFor(() => {
+      const msgs = result.current.conversations[TEMP_CONVERSATION_ID] ?? [];
+      expect(msgs.length).toBeGreaterThanOrEqual(2);
+      expect(String(msgs[1]?.content ?? '')).toContain('Hello ');
+    });
+
+    await act(async () => {
+      resolveSecondRead({
+        done: false,
+        value: new TextEncoder().encode(secondChunk),
+      });
+    });
+
+    await waitFor(() => {
+      expect(onComplete).toHaveBeenCalled();
+    });
+
+    expect(result.current.streamingConversationId).toBeNull();
+  });
 });