Fix 'Edit In Place' and 'Fork Session' features (#6970)

Signed-off-by: Bright Zheng <bright@squareup.com>
Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
Co-authored-by: Zane Staggs <zane@squareup.com>

Author: 3 people

ui/desktop/src/App.tsx

@@ -593,7 +593,7 @@ export function AppInner() {
         console.log('[App] Processing initial message from launcher:', initialMessage);
         navigate('/pair', {
           state: {
-            initialMessage,
+            initialMessage: { msg: initialMessage, images: [] },
           },
         });
         setTimeout(() => {

ui/desktop/src/components/BaseChat.tsx

@@ -279,7 +279,7 @@ export default function BaseChat({
       navigate(`/pair?${params.toString()}`, {
         state: {
           disableAnimation: true,
-          initialMessage: editedMessage,
+          initialMessage: editedMessage ? { msg: editedMessage, images: [] } : undefined,
         },
       });
     };

ui/desktop/src/hooks/useAutoSubmit.ts

@@ -76,9 +76,12 @@ export function useAutoSubmit({
 
     // Scenario 2: Forked session with edited message
     if (shouldStartAgent && initialMessage) {
-      hasAutoSubmittedRef.current = true;
-      handleSubmit(initialMessage);
-      clearInitialMessage();
+      if (messages.length > 0) {
+        hasAutoSubmittedRef.current = true;
+        handleSubmit(initialMessage);
+        clearInitialMessage();
+        return;
+      }
       return;
     }
 

ui/desktop/src/hooks/useChatStream.ts

@@ -369,9 +369,6 @@ export function useChatStream({
 
       const isNewSession = sessionId && sessionId.match(/^\d{8}_\d{6}$/);
       if (isNewSession) {
-        console.log(
-          'useChatStream: Message stream finished for new session, emitting message-stream-finished event'
-        );
         window.dispatchEvent(new CustomEvent(AppEvents.MESSAGE_STREAM_FINISHED));
       }
 
@@ -591,6 +588,7 @@ export function useChatStream({
           body: {
             session_id: sessionId,
             user_message: newMessage,
+            ...(hasExistingMessages && { conversation_so_far: currentState.messages }),
           },
           throwOnError: true,
           signal: abortControllerRef.current.signal,
@@ -751,55 +749,40 @@ export function useChatStream({
           });
 
           if (sessionResponse.data?.conversation) {
-            const updatedMessages = [...sessionResponse.data.conversation];
-
-            if (updatedMessages.length > 0) {
-              const lastMessage = updatedMessages[updatedMessages.length - 1];
-              if (lastMessage.role === 'user') {
-                lastMessage.content = [{ type: 'text', text: newContent }];
-
-                for (const content of message.content) {
-                  if (content.type === 'image') {
-                    lastMessage.content.push(content);
-                  }
-                }
-
-                dispatch({ type: 'SET_MESSAGES', payload: updatedMessages });
-                dispatch({ type: 'START_STREAMING' });
-                abortControllerRef.current = new AbortController();
-
-                try {
-                  const placeholderMessage = createUserMessage(newContent);
-
-                  const { stream } = await reply({
-                    body: {
-                      session_id: targetSessionId,
-                      user_message: placeholderMessage,
-                    },
-                    throwOnError: true,
-                    signal: abortControllerRef.current.signal,
-                  });
-
-                  await streamFromResponse(
-                    stream,
-                    updatedMessages,
-                    dispatch,
-                    onFinish,
-                    targetSessionId
-                  );
-                } catch (error) {
-                  if (error instanceof Error && error.name === 'AbortError') {
-                    dispatch({ type: 'SET_CHAT_STATE', payload: ChatState.Idle });
-                  } else {
-                    throw error;
-                  }
-                }
-                return;
+            const truncatedMessages = [...sessionResponse.data.conversation];
+            const updatedUserMessage = createUserMessage(newContent);
+
+            for (const content of message.content) {
+              if (content.type === 'image') {
+                updatedUserMessage.content.push(content);
               }
             }
 
-            dispatch({ type: 'SET_MESSAGES', payload: sessionResponse.data.conversation });
-            await handleSubmit({ msg: newContent, images: [] });
+            const messagesForUI = [...truncatedMessages, updatedUserMessage];
+            dispatch({ type: 'SET_MESSAGES', payload: messagesForUI });
+            dispatch({ type: 'START_STREAMING' });
+
+            abortControllerRef.current = new AbortController();
+
+            try {
+              const { stream } = await reply({
+                body: {
+                  session_id: targetSessionId,
+                  user_message: updatedUserMessage,
+                  conversation_so_far: truncatedMessages,
+                },
+                throwOnError: true,
+                signal: abortControllerRef.current.signal,
+              });
+
+              await streamFromResponse(stream, messagesForUI, dispatch, onFinish, targetSessionId);
+            } catch (error) {
+              if (error instanceof Error && error.name === 'AbortError') {
+                dispatch({ type: 'SET_CHAT_STATE', payload: ChatState.Idle });
+              } else {
+                throw error;
+              }
+            }
           } else {
             await handleSubmit({ msg: newContent, images: [] });
           }