feat(ui): Implement in-place message editing with re-response (#3798)

Co-authored-by: Zane Staggs <zane@squareup.com>

ui/desktop/src/components/BaseChat.tsx

@@ -158,6 +158,7 @@ function BaseChatContent({
     sessionMetadata,
     isUserMessage,
     clearError,
+    onMessageUpdate,
   } = useChatEngine({
     chat,
     setChat,
@@ -418,6 +419,7 @@ function BaseChatContent({
                       isUserMessage={isUserMessage}
                       onScrollToBottom={handleScrollToBottom}
                       isStreamingMessage={chatState !== ChatState.Idle}
+                      onMessageUpdate={onMessageUpdate}
                     />
                   ) : (
                     // Render messages with SearchView wrapper when search is enabled
@@ -434,6 +436,7 @@ function BaseChatContent({
                         isUserMessage={isUserMessage}
                         onScrollToBottom={handleScrollToBottom}
                         isStreamingMessage={chatState !== ChatState.Idle}
+                        onMessageUpdate={onMessageUpdate}
                       />
                     </SearchView>
                   )}

ui/desktop/src/components/MCPUIResourceRenderer.tsx

@@ -64,12 +64,11 @@ export default function MCPUIResourceRenderer({ content }: MCPUIResourceRenderer
         <UIResourceRenderer
           resource={content.resource}
           onUIAction={handleUIAction}
-          htmlProps={{
-            autoResizeIframe: {
-              height: true,
-              width: false, // set to false to allow for responsive design
-            },
-          }}
+          htmlProps={
+            {
+              // Remove autoResizeIframe as it's not in the type definition
+            }
+          }
         />
       </div>
     </div>

ui/desktop/src/components/ProgressiveMessageList.tsx

@@ -37,6 +37,7 @@ interface ProgressiveMessageListProps {
   // Custom render function for messages
   renderMessage?: (message: Message, index: number) => React.ReactNode | null;
   isStreamingMessage?: boolean; // Whether messages are currently being streamed
+  onMessageUpdate?: (messageId: string, newContent: string) => void;
 }
 
 export default function ProgressiveMessageList({
@@ -52,6 +53,7 @@ export default function ProgressiveMessageList({
   showLoadingThreshold = 50,
   renderMessage, // Custom render function
   isStreamingMessage = false, // Whether messages are currently being streamed
+  onMessageUpdate,
 }: ProgressiveMessageListProps) {
   const [renderedCount, setRenderedCount] = useState(() => {
     // Initialize with either all messages (if small) or first batch (if large)
@@ -206,7 +208,9 @@ export default function ProgressiveMessageList({
                     }}
                   />
                 ) : (
-                  !hasOnlyToolResponses(message) && <UserMessage message={message} />
+                  !hasOnlyToolResponses(message) && (
+                    <UserMessage message={message} onMessageUpdate={onMessageUpdate} />
+                  )
                 )}
               </>
             ) : (
@@ -258,6 +262,7 @@ export default function ProgressiveMessageList({
     appendMessage,
     toolCallNotifications,
     isStreamingMessage,
+    onMessageUpdate,
     hasContextHandlerContent,
     getContextHandlerType,
     onScrollToBottom,

ui/desktop/src/components/UserMessage.tsx

@@ -1,4 +1,4 @@
-import { useRef, useMemo } from 'react';
+import { useRef, useMemo, useState, useEffect, useCallback } from 'react';
 import LinkPreview from './LinkPreview';
 import ImagePreview from './ImagePreview';
 import { extractUrls } from '../utils/urlUtils';
@@ -7,13 +7,21 @@ import MarkdownContent from './MarkdownContent';
 import { Message, getTextContent } from '../types/message';
 import MessageCopyLink from './MessageCopyLink';
 import { formatMessageTimestamp } from '../utils/timeUtils';
+import Edit from './icons/Edit';
+import { Button } from './ui/button';
 
 interface UserMessageProps {
   message: Message;
+  onMessageUpdate?: (messageId: string, newContent: string) => void;
 }
 
-export default function UserMessage({ message }: UserMessageProps) {
+export default function UserMessage({ message, onMessageUpdate }: UserMessageProps) {
   const contentRef = useRef<HTMLDivElement | null>(null);
+  const textareaRef = useRef<HTMLTextAreaElement>(null);
+  const [isEditing, setIsEditing] = useState(false);
+  const [editContent, setEditContent] = useState('');
+  const [hasBeenEdited, setHasBeenEdited] = useState(false);
+  const [error, setError] = useState<string | null>(null);
 
   // Extract text content from the message
   const textContent = getTextContent(message);
@@ -21,46 +29,230 @@ export default function UserMessage({ message }: UserMessageProps) {
   // Extract image paths from the message
   const imagePaths = extractImagePaths(textContent);
 
-  // Remove image paths from text for display
-  const displayText = removeImagePathsFromText(textContent, imagePaths);
+  // Remove image paths from text for display - memoized for performance
+  const displayText = useMemo(
+    () => removeImagePathsFromText(textContent, imagePaths),
+    [textContent, imagePaths]
+  );
 
   // Memoize the timestamp
   const timestamp = useMemo(() => formatMessageTimestamp(message.created), [message.created]);
 
   // Extract URLs which explicitly contain the http:// or https:// protocol
-  const urls = extractUrls(displayText, []);
+  const urls = useMemo(() => extractUrls(displayText, []), [displayText]);
+
+  // Effect to handle message content changes and ensure persistence
+  useEffect(() => {
+    // Log content display for debugging
+    window.electron.logInfo(
+      `Displaying content for message: ${message.id} content: ${displayText}`
+    );
+
+    // If we're not editing, update the edit content to match the current message
+    if (!isEditing) {
+      setEditContent(displayText);
+    }
+  }, [message.content, displayText, message.id, isEditing]);
+
+  // Initialize edit mode with current message content
+  const initializeEditMode = useCallback(() => {
+    setEditContent(displayText);
+    setError(null);
+    window.electron.logInfo(`Entering edit mode with content: ${displayText}`);
+  }, [displayText]);
+
+  // Handle edit button click
+  const handleEditClick = useCallback(() => {
+    const newEditingState = !isEditing;
+    setIsEditing(newEditingState);
+
+    // Initialize edit content when entering edit mode
+    if (newEditingState) {
+      initializeEditMode();
+      window.electron.logInfo(`Edit interface shown for message: ${message.id}`);
+
+      // Focus the textarea after a brief delay to ensure it's rendered
+      setTimeout(() => {
+        if (textareaRef.current) {
+          textareaRef.current.focus();
+          textareaRef.current.setSelectionRange(
+            textareaRef.current.value.length,
+            textareaRef.current.value.length
+          );
+        }
+      }, 50);
+    }
+
+    window.electron.logInfo(`Edit state toggled: ${newEditingState} for message: ${message.id}`);
+  }, [isEditing, initializeEditMode, message.id]);
+
+  // Handle content changes in edit mode
+  const handleContentChange = useCallback((e: React.ChangeEvent<HTMLTextAreaElement>) => {
+    const newContent = e.target.value;
+    setEditContent(newContent);
+    setError(null); // Clear any previous errors
+    window.electron.logInfo(`Content changed: ${newContent}`);
+  }, []);
+
+  // Handle save action
+  const handleSave = useCallback(() => {
+    // Exit edit mode immediately
+    setIsEditing(false);
+
+    // Check if content has actually changed
+    if (editContent !== displayText) {
+      // Validate content
+      if (editContent.trim().length === 0) {
+        setError('Message cannot be empty');
+        return;
+      }
+
+      // Update the message content through the callback
+      if (onMessageUpdate && message.id) {
+        onMessageUpdate(message.id, editContent);
+        setHasBeenEdited(true);
+      }
+    }
+  }, [editContent, displayText, onMessageUpdate, message.id]);
+
+  // Handle cancel action
+  const handleCancel = useCallback(() => {
+    window.electron.logInfo('Cancel clicked - reverting to original content');
+    setIsEditing(false);
+    setEditContent(displayText); // Reset to original content
+    setError(null);
+  }, [displayText]);
+
+  // Handle keyboard events for accessibility
+  const handleKeyDown = useCallback(
+    (e: React.KeyboardEvent) => {
+      window.electron.logInfo(
+        `Key pressed: ${e.key}, metaKey: ${e.metaKey}, ctrlKey: ${e.ctrlKey}`
+      );
+
+      if (e.key === 'Escape') {
+        e.preventDefault();
+        handleCancel();
+      } else if (e.key === 'Enter' && (e.metaKey || e.ctrlKey)) {
+        e.preventDefault();
+        window.electron.logInfo('Cmd+Enter detected, calling handleSave');
+        handleSave();
+      }
+    },
+    [handleCancel, handleSave]
+  );
+
+  // Auto-resize textarea based on content
+  useEffect(() => {
+    if (textareaRef.current && isEditing) {
+      textareaRef.current.style.height = 'auto';
+      textareaRef.current.style.height = `${Math.min(textareaRef.current.scrollHeight, 200)}px`;
+    }
+  }, [editContent, isEditing]);
 
   return (
-    <div className="message flex justify-end mt-[16px] w-full opacity-0 animate-[appear_150ms_ease-in_forwards]">
-      <div className="flex-col max-w-[85%] w-fit">
-        <div className="flex flex-col group">
-          <div className="flex bg-background-accent text-text-on-accent rounded-xl py-2.5 px-4">
-            <div ref={contentRef}>
-              <MarkdownContent
-                content={displayText}
-                className="text-text-on-accent prose-a:text-text-on-accent prose-headings:text-text-on-accent prose-strong:text-text-on-accent prose-em:text-text-on-accent user-message"
-              />
+    <div className="w-full mt-[16px] opacity-0 animate-[appear_150ms_ease-in_forwards]">
+      <div className="flex flex-col group">
+        {isEditing ? (
+          // Truly wide, centered, in-place edit box replacing the bubble
+          <div className="w-full max-w-4xl mx-auto bg-background-light dark:bg-background-dark text-text-prominent rounded-xl border border-border-subtle shadow-lg py-4 px-4 my-2 transition-all duration-200 ease-in-out">
+            <textarea
+              ref={textareaRef}
+              value={editContent}
+              onChange={handleContentChange}
+              onKeyDown={handleKeyDown}
+              className="w-full resize-none bg-transparent text-text-prominent placeholder:text-text-subtle border border-border-subtle rounded-lg focus:outline-none focus:ring-2 focus:ring-blue-400 focus:border-blue-400 transition-all duration-200 text-base leading-relaxed"
+              style={{
+                minHeight: '120px',
+                maxHeight: '300px',
+                padding: '16px',
+                fontFamily: 'inherit',
+                lineHeight: '1.6',
+                wordBreak: 'break-word',
+                overflowWrap: 'break-word',
+              }}
+              placeholder="Edit your message..."
+              aria-label="Edit message content"
+              aria-describedby={error ? `error-${message.id}` : undefined}
+            />
+            {/* Error message */}
+            {error && (
+              <div
+                id={`error-${message.id}`}
+                className="text-red-400 text-xs mt-2 mb-2"
+                role="alert"
+                aria-live="polite"
+              >
+                {error}
+              </div>
+            )}
+            <div className="flex justify-end gap-3 mt-4">
+              <Button onClick={handleCancel} variant="ghost" aria-label="Cancel editing">
+                Cancel
+              </Button>
+              <Button onClick={handleSave} aria-label="Save changes">
+                Save
+              </Button>
             </div>
           </div>
+        ) : (
+          // Normal message display
+          <div className="message flex justify-end w-full">
+            <div className="flex-col max-w-[85%] w-fit">
+              <div className="flex flex-col group">
+                <div className="flex bg-background-accent text-text-on-accent rounded-xl py-2.5 px-4">
+                  <div ref={contentRef}>
+                    <MarkdownContent
+                      content={displayText}
+                      className="text-text-on-accent prose-a:text-text-on-accent prose-headings:text-text-on-accent prose-strong:text-text-on-accent prose-em:text-text-on-accent user-message"
+                    />
+                  </div>
+                </div>
 
-          {/* Render images if any */}
-          {imagePaths.length > 0 && (
-            <div className="flex flex-wrap gap-2 mt-2">
-              {imagePaths.map((imagePath, index) => (
-                <ImagePreview key={index} src={imagePath} alt={`Pasted image ${index + 1}`} />
-              ))}
-            </div>
-          )}
+                {/* Render images if any */}
+                {imagePaths.length > 0 && (
+                  <div className="flex flex-wrap gap-2 mt-2">
+                    {imagePaths.map((imagePath, index) => (
+                      <ImagePreview key={index} src={imagePath} alt={`Pasted image ${index + 1}`} />
+                    ))}
+                  </div>
+                )}
 
-          <div className="relative h-[22px] flex justify-end text-right">
-            <div className="absolute w-40 font-mono right-0 text-xs text-text-muted pt-1 transition-all duration-200 group-hover:-translate-y-4 group-hover:opacity-0">
-              {timestamp}
-            </div>
-            <div className="absolute right-0 pt-1">
-              <MessageCopyLink text={displayText} contentRef={contentRef} />
+                <div className="relative h-[22px] flex justify-end text-right">
+                  <div className="absolute w-40 font-mono right-0 text-xs text-text-muted pt-1 transition-all duration-200 group-hover:-translate-y-4 group-hover:opacity-0">
+                    {timestamp}
+                  </div>
+                  <div className="absolute right-0 pt-1 flex items-center gap-2">
+                    <button
+                      onClick={handleEditClick}
+                      onKeyDown={(e) => {
+                        if (e.key === 'Enter' || e.key === ' ') {
+                          e.preventDefault();
+                          handleEditClick();
+                        }
+                      }}
+                      className="flex items-center gap-1 text-xs text-text-subtle hover:cursor-pointer hover:text-text-prominent transition-all duration-200 opacity-0 group-hover:opacity-100 -translate-y-4 group-hover:translate-y-0 focus:outline-none focus:ring-2 focus:ring-blue-400 focus:ring-opacity-50 rounded"
+                      aria-label={`Edit message: ${displayText.substring(0, 50)}${displayText.length > 50 ? '...' : ''}`}
+                      aria-expanded={isEditing}
+                      title="Edit message"
+                    >
+                      <Edit className="h-3 w-3" />
+                      <span>Edit</span>
+                    </button>
+                    <MessageCopyLink text={displayText} contentRef={contentRef} />
+                  </div>
+                </div>
+              </div>
             </div>
           </div>
-        </div>
+        )}
+
+        {/* Edited indicator */}
+        {hasBeenEdited && !isEditing && (
+          <div className="text-xs text-text-subtle mt-1 text-right transition-opacity duration-200">
+            Edited
+          </div>
+        )}
 
         {/* TODO(alexhancock): Re-enable link previews once styled well again */}
         {/* eslint-disable-next-line no-constant-binary-expression */}

ui/desktop/src/hooks/useChatEngine.test.ts

@@ -0,0 +1,155 @@
+import { renderHook, act } from '@testing-library/react';
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { useChatEngine } from './useChatEngine';
+import { Message, getTextContent } from '../types/message';
+import { ChatType } from '../types/chat';
+import type { Mock } from 'vitest';
+
+// Mock the useMessageStream hook which is a dependency of useChatEngine
+vi.mock('./useMessageStream', () => ({
+  useMessageStream: vi.fn(),
+}));
+
+// Mock the sessions API which is another dependency
+vi.mock('../sessions', () => ({
+  fetchSessionDetails: vi.fn().mockResolvedValue({ metadata: {} }),
+}));
+
+describe('useChatEngine', () => {
+  let mockUseMessageStream: Mock;
+
+  beforeEach(async () => {
+    // Mock the global window object more completely for the React testing environment
+    const mockWindow = {
+      appConfig: {
+        get: vi.fn((key: string) => {
+          if (key === 'GOOSE_API_HOST') return 'http://localhost';
+          if (key === 'GOOSE_PORT') return '8000';
+          return null;
+        }),
+      },
+      electron: {
+        logInfo: vi.fn(),
+      },
+      setTimeout: vi.fn((fn: () => void) => {
+        fn(); // Execute immediately for tests
+        return 123;
+      }),
+      clearTimeout: vi.fn(),
+      dispatchEvent: vi.fn(),
+      CustomEvent: vi.fn(),
+      // Add basic browser objects required by React Testing Library
+      HTMLElement: class MockHTMLElement {},
+      Event: class MockEvent {},
+    };
+    vi.stubGlobal('window', mockWindow);
+
+    // Dynamically import the hook so we can get a reference to the mock
+    const { useMessageStream } = await import('./useMessageStream');
+    mockUseMessageStream = useMessageStream as Mock;
+
+    // Reset all mocks before each test to ensure a clean state
+    vi.clearAllMocks();
+
+    // Provide a complete, default mock implementation for useMessageStream
+    mockUseMessageStream.mockReturnValue({
+      messages: [],
+      append: vi.fn(),
+      stop: vi.fn(),
+      chatState: 'idle',
+      error: undefined,
+      setMessages: vi.fn(),
+      input: '',
+      setInput: vi.fn(),
+      handleInputChange: vi.fn(),
+      handleSubmit: vi.fn(),
+      updateMessageStreamBody: vi.fn(),
+      notifications: [],
+      sessionMetadata: undefined,
+      setError: vi.fn(),
+    });
+  });
+
+  describe('onMessageUpdate', () => {
+    it('should truncate history and append the updated message when a message is edited', () => {
+      // --- 1. ARRANGE ---
+      const initialMessages: Message[] = [
+        { id: '1', role: 'user', content: [{ type: 'text', text: 'First message' }], created: 0 },
+        {
+          id: '2',
+          role: 'assistant',
+          content: [{ type: 'text', text: 'First response' }],
+          created: 1,
+        },
+        {
+          id: '3',
+          role: 'user',
+          content: [{ type: 'text', text: 'Message to be edited' }],
+          created: 2,
+        },
+        {
+          id: '4',
+          role: 'assistant',
+          content: [{ type: 'text', text: 'Response to be deleted' }],
+          created: 3,
+        },
+      ];
+
+      const mockSetMessages = vi.fn();
+      const mockAppend = vi.fn();
+
+      // Configure the mock to return specific values for this test case
+      mockUseMessageStream.mockReturnValue({
+        messages: initialMessages,
+        append: mockAppend,
+        setMessages: mockSetMessages,
+        notifications: [],
+        stop: vi.fn(),
+        chatState: 'idle',
+        error: undefined,
+        input: '',
+        setInput: vi.fn(),
+        handleInputChange: vi.fn(),
+        handleSubmit: vi.fn(),
+        updateMessageStreamBody: vi.fn(),
+        sessionMetadata: undefined,
+        setError: vi.fn(),
+      });
+
+      const mockChat: ChatType = {
+        id: 'test-chat',
+        messages: initialMessages,
+        title: 'Test Chat',
+        messageHistoryIndex: 0,
+      };
+
+      // Render the hook with our test setup
+      const { result } = renderHook(() =>
+        useChatEngine({
+          chat: mockChat,
+          setChat: vi.fn(),
+        })
+      );
+
+      const messageIdToUpdate = '3';
+      const newContent = 'This is the edited message.';
+
+      // --- 2. ACT ---
+      // Call the function we want to test
+      act(() => {
+        result.current.onMessageUpdate(messageIdToUpdate, newContent);
+      });
+
+      // --- 3. ASSERT ---
+      // Verify that setMessages was called with the correctly truncated history
+      const expectedTruncatedHistory = initialMessages.slice(0, 2);
+      expect(mockSetMessages).toHaveBeenCalledWith(expectedTruncatedHistory);
+
+      // Verify that append was called with the new message
+      expect(mockAppend).toHaveBeenCalledTimes(1);
+      const appendedMessage = mockAppend.mock.calls[0][0];
+      expect(getTextContent(appendedMessage)).toBe(newContent);
+      expect(appendedMessage.role).toBe('user');
+    });
+  });
+});

ui/desktop/src/hooks/useChatEngine.ts

@@ -51,6 +51,9 @@ export const useChatEngine = ({
   const [localOutputTokens, setLocalOutputTokens] = useState<number>(0);
   const [powerSaveTimeoutId, setPowerSaveTimeoutId] = useState<number | null>(null);
 
+  // Track pending edited message
+  const [pendingEdit, setPendingEdit] = useState<{ id: string; content: string } | null>(null);
+
   // Store message in global history when it's added (if enabled)
   const storeMessageInHistory = useCallback(
     (message: Message) => {
@@ -408,6 +411,34 @@ export const useChatEngine = ({
     }, new Map());
   }, [notifications]);
 
+  // Handle message updates from the UI
+  const onMessageUpdate = useCallback(
+    (messageId: string, newContent: string) => {
+      const messageIndex = messages.findIndex((msg) => msg.id === messageId);
+
+      if (messageIndex !== -1) {
+        // Truncate the history to the point *before* the edited message.
+        const history = messages.slice(0, messageIndex);
+
+        // Set the truncated history.
+        setMessages(history);
+
+        // Instead of setTimeout, set pendingEdit which will be handled in useEffect
+        setPendingEdit({ id: messageId, content: newContent });
+      }
+    },
+    [messages, setMessages, setPendingEdit]
+  );
+
+  // Listen for pending edit and append message after messages updated
+  useEffect(() => {
+    if (pendingEdit) {
+      const updatedMessage = createUserMessage(pendingEdit.content);
+      append(updatedMessage);
+      setPendingEdit(null); // Reset after processing
+    }
+  }, [pendingEdit, append]);
+
   return {
     // Core message data
     messages,
@@ -451,5 +482,8 @@ export const useChatEngine = ({
 
     // Error management
     clearError: () => setError(undefined),
+
+    // New functions for message editing
+    onMessageUpdate,
   };
 };

ui/desktop/src/test/setup.ts

@@ -1,5 +1,12 @@
 import '@testing-library/jest-dom';
-import { vi } from 'vitest';
+import { vi, afterEach } from 'vitest';
+import { cleanup } from '@testing-library/react';
+
+// This is the standard setup to ensure that React Testing Library's
+// automatic cleanup runs after each test.
+afterEach(() => {
+  cleanup();
+});
 
 // Mock console methods to avoid noise in tests
 // eslint-disable-next-line no-undef