From 49560f45e8608f3f4fc519684ddeae0102c6a196 Mon Sep 17 00:00:00 2001
From: Douwe Osinga <douwe@squareup.com>
Date: Wed, 8 Oct 2025 13:06:21 -0400
Subject: [PATCH] make it stream

---
 ui/desktop/src/App.tsx                  | 17 ++++++++-
 ui/desktop/src/components/BaseChat2.tsx |  1 +
 ui/desktop/src/hooks/useChatStream.ts   | 48 +++++++++++++++++++------
 ui/desktop/src/updates.ts               |  2 ++
 4 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/ui/desktop/src/App.tsx b/ui/desktop/src/App.tsx
index b6c2475f02..dac9226eee 100644
--- a/ui/desktop/src/App.tsx
+++ b/ui/desktop/src/App.tsx
@@ -44,6 +44,8 @@ import {
   useAgent,
 } from './hooks/useAgent';
 import { useNavigation } from './hooks/useNavigation';
+import { USE_NEW_CHAT } from './updates';
+import Pair2 from './components/Pair2';
 
 // Route Components
 const HubRouteWrapper = ({
@@ -93,7 +95,20 @@ const PairRouteWrapper = ({
 
   const resumeSessionId = searchParams.get('resumeSessionId') ?? undefined;
 
-  return (
+  return USE_NEW_CHAT ? (
+    <Pair2
+      chat={chat}
+      setChat={setChat}
+      setView={setView}
+      agentState={agentState}
+      loadCurrentChat={loadCurrentChat}
+      setFatalError={setFatalError}
+      setAgentWaitingMessage={setAgentWaitingMessage}
+      setIsGoosehintsModalOpen={setIsGoosehintsModalOpen}
+      resumeSessionId={resumeSessionId}
+      initialMessage={initialMessage}
+    />
+  ) : (
     <Pair
       chat={chat}
       setChat={setChat}
diff --git a/ui/desktop/src/components/BaseChat2.tsx b/ui/desktop/src/components/BaseChat2.tsx
index fec42026d7..9acb4467c7 100644
--- a/ui/desktop/src/components/BaseChat2.tsx
+++ b/ui/desktop/src/components/BaseChat2.tsx
@@ -184,6 +184,7 @@ function BaseChatContent({
   const initialPrompt = messages.length == 0 && recipe?.prompt ? recipe.prompt : '';
   return (
     <div className="h-full flex flex-col min-h-0">
+      <h2>Warning: BaseChat2!</h2>
       <MainPanelLayout
         backgroundColor={'bg-background-muted'}
         removeTopPadding={true}
diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts
index 7d5c321c55..6534880430 100644
--- a/ui/desktop/src/hooks/useChatStream.ts
+++ b/ui/desktop/src/hooks/useChatStream.ts
@@ -1,6 +1,7 @@
 import { useState, useCallback, useRef } from 'react';
 import { ChatState } from '../types/chatState';
 import { Message } from '../api';
+import { getApiUrl } from '../config';
 
 const TextDecoder = globalThis.TextDecoder;
 
@@ -11,6 +12,28 @@ interface UseChatStreamProps {
   onStreamFinish?: () => void;
 }
 
+function pushMessage(currentMessages: Message[], incomingMsg: Message): Message[] {
+  const lastMsg = currentMessages[currentMessages.length - 1];
+
+  if (lastMsg?.id && lastMsg.id === incomingMsg.id) {
+    const lastContent = lastMsg.content[lastMsg.content.length - 1];
+    const newContent = incomingMsg.content[incomingMsg.content.length - 1];
+
+    if (
+      lastContent?.type === 'text' &&
+      newContent?.type === 'text' &&
+      incomingMsg.content.length === 1
+    ) {
+      lastContent.text += newContent.text;
+    } else {
+      lastMsg.content.push(...incomingMsg.content);
+    }
+    return [...currentMessages];
+  } else {
+    return [...currentMessages, incomingMsg];
+  }
+}
+
 export function useChatStream({
   sessionId,
   messages,
@@ -28,22 +51,24 @@ export function useChatStream({
         created: Date.now(),
       };
 
-      const updatedMessages = [...messages, newMessage];
-      setMessages(updatedMessages);
+      let currentMessages = [...messages, newMessage];
+      setMessages(currentMessages);
       setChatState(ChatState.Streaming);
 
       abortControllerRef.current = new AbortController();
 
       try {
-        const response = await fetch('/reply', {
+        // TODO(Douwe): this side steps our API. heyapi does support streaming though which should make
+        // this all nice & typed
+        const response = await fetch(getApiUrl('/reply'), {
           method: 'POST',
-          headers: { 'Content-Type': 'application/json' },
+          headers: {
+            'Content-Type': 'application/json',
+            'X-Secret-Key': await window.electron.getSecretKey(),
+          },
           body: JSON.stringify({
             session_id: sessionId,
-            messages: updatedMessages.map((m) => ({
-              role: m.role,
-              content: m.content,
-            })),
+            messages: currentMessages,
           }),
           signal: abortControllerRef.current.signal,
         });
@@ -72,7 +97,8 @@ export function useChatStream({
 
               if (event.message) {
                 const msg = event.message as Message;
-                setMessages([...updatedMessages, msg]);
+                currentMessages = pushMessage(currentMessages, msg);
+                setMessages(currentMessages);
               }
 
               if (event.error) {
@@ -91,8 +117,8 @@ export function useChatStream({
             }
           }
         }
-      } catch (error: any) {
-        if (error.name !== 'AbortError') {
+      } catch (error) {
+        if (error instanceof Error && error.name !== 'AbortError') {
           console.error('Stream error:', error);
         }
         setChatState(ChatState.Idle);
diff --git a/ui/desktop/src/updates.ts b/ui/desktop/src/updates.ts
index 978b762429..f98cbfc77d 100644
--- a/ui/desktop/src/updates.ts
+++ b/ui/desktop/src/updates.ts
@@ -2,3 +2,5 @@ export const UPDATES_ENABLED = true;
 export const COST_TRACKING_ENABLED = true;
 export const ANNOUNCEMENTS_ENABLED = false;
 export const CONFIGURATION_ENABLED = true;
+
+export const USE_NEW_CHAT = true;