fix: brought back inline comment drafts in ChatInput

packages/ui/src/components/chat/ChatInput.tsx

@@ -13,6 +13,8 @@ import { useConfigStore } from '@/stores/useConfigStore';
 import { useUIStore } from '@/stores/useUIStore';
 import { useMessageQueueStore, type QueuedMessage } from '@/stores/messageQueueStore';
 import type { AttachedFile } from '@/stores/types/sessionTypes';
+import { useInlineCommentDraftStore, type InlineCommentDraft } from '@/stores/useInlineCommentDraftStore';
+import { appendInlineComments } from '@/lib/messages/inlineComments';
 import { AttachedFilesList } from './FileAttachment';
 import { QueuedMessageChips } from './QueuedMessageChips';
 import { FileMentionAutocomplete, type FileMentionHandle } from './FileMentionAutocomplete';
@@ -115,6 +117,20 @@ export const ChatInput: React.FC<ChatInputProps> = ({ onOpenSettings, scrollToBo
     const addToQueue = useMessageQueueStore((state) => state.addToQueue);
     const clearQueue = useMessageQueueStore((state) => state.clearQueue);
 
+    // Inline comment drafts
+    const draftCount = useInlineCommentDraftStore(
+        React.useCallback(
+            (state) => {
+                const sessionKey = currentSessionId ?? (newSessionDraftOpen ? 'draft' : '');
+                if (!sessionKey) return 0;
+                return (state.drafts[sessionKey] ?? []).length;
+            },
+            [currentSessionId, newSessionDraftOpen]
+        )
+    );
+    const consumeDrafts = useInlineCommentDraftStore((state) => state.consumeDrafts);
+    const hasDrafts = draftCount > 0;
+
     // Session activity for auto-send on idle
     const { phase: sessionPhase } = useCurrentSessionActivity();
     const prevSessionPhaseRef = React.useRef(sessionPhase);
@@ -215,9 +231,19 @@ export const ChatInput: React.FC<ChatInputProps> = ({ onOpenSettings, scrollToBo
     // Consume pending input text (e.g., from revert action)
     React.useEffect(() => {
         if (pendingInputText !== null) {
-            const text = consumePendingInputText();
+            const pending = consumePendingInputText();
-            if (text) {
+            if (pending?.text) {
-                setMessage(text);
+                if (pending.mode === 'append') {
+                    setMessage((prev) => {
+                        const next = pending.text.trim();
+                        if (!next) return prev;
+                        const base = prev.trimEnd();
+                        if (!base.trim()) return next;
+                        return `${base} ${next}`;
+                    });
+                } else {
+                    setMessage(pending.text);
+                }
                 // Focus textarea after setting message
                 setTimeout(() => {
                     textareaRef.current?.focus();
@@ -226,7 +252,7 @@ export const ChatInput: React.FC<ChatInputProps> = ({ onOpenSettings, scrollToBo
         }
     }, [pendingInputText, consumePendingInputText]);
 
-    const hasContent = message.trim() || attachedFiles.length > 0;
+    const hasContent = message.trim() || attachedFiles.length > 0 || hasDrafts;
     const hasQueuedMessages = queuedMessages.length > 0;
     const canSend = hasContent || hasQueuedMessages;
 
@@ -239,7 +265,12 @@ export const ChatInput: React.FC<ChatInputProps> = ({ onOpenSettings, scrollToBo
     const handleQueueMessage = React.useCallback(() => {
         if (!hasContent || !currentSessionId) return;
 
-        const messageToQueue = message.replace(/^\n+|\n+$/g, '');
+        const drafts = consumeDrafts(currentSessionId);
+
+        let messageToQueue = message.replace(/^\n+|\n+$/g, '');
+        if (drafts.length > 0) {
+            messageToQueue = appendInlineComments(messageToQueue, drafts);
+        }
         const attachmentsToQueue = attachedFiles.map((file) => ({ ...file }));
 
         addToQueue(currentSessionId, {
@@ -256,7 +287,7 @@ export const ChatInput: React.FC<ChatInputProps> = ({ onOpenSettings, scrollToBo
         if (!isMobile) {
             textareaRef.current?.focus();
         }
-    }, [hasContent, currentSessionId, message, attachedFiles, addToQueue, clearAttachedFiles, isMobile]);
+    }, [hasContent, currentSessionId, message, attachedFiles, addToQueue, clearAttachedFiles, isMobile, consumeDrafts]);
 
     const handleSubmit = async (e?: React.FormEvent) => {
         e?.preventDefault();
@@ -323,6 +354,23 @@ export const ChatInput: React.FC<ChatInputProps> = ({ onOpenSettings, scrollToBo
             }
         }
 
+        const sessionKey = currentSessionId ?? (newSessionDraftOpen ? 'draft' : null);
+        let drafts: InlineCommentDraft[] = [];
+        if (sessionKey) {
+            drafts = consumeDrafts(sessionKey);
+        }
+
+        if (drafts.length > 0) {
+            if (queuedMessages.length === 0) {
+                primaryText = appendInlineComments(primaryText, drafts);
+            } else if (additionalParts.length > 0) {
+                const lastPart = additionalParts[additionalParts.length - 1];
+                lastPart.text = appendInlineComments(lastPart.text, drafts);
+            } else {
+                primaryText = appendInlineComments(primaryText, drafts);
+            }
+        }
+
         if (!primaryText && additionalParts.length === 0) return;
 
         // Clear queue and input
@@ -1525,6 +1573,22 @@ export const ChatInput: React.FC<ChatInputProps> = ({ onOpenSettings, scrollToBo
                         }, 0);
                     }}
                 />
+                {hasDrafts && (
+                    <div className="pb-2">
+                        <div
+                            className="inline-flex items-center gap-1.5 px-2.5 py-1 rounded-xl border"
+                            style={{
+                                backgroundColor: currentTheme?.colors?.surface?.elevated,
+                                borderColor: currentTheme?.colors?.interactive?.border,
+                            }}
+                        >
+                            <span className="text-xs font-medium text-muted-foreground">Review comments:</span>
+                            <span className="text-xs font-semibold" style={{ color: currentTheme?.colors?.status?.info }}>
+                                {draftCount}
+                            </span>
+                        </div>
+                    </div>
+                )}
                 <div
                     className={cn(
                         "flex flex-col relative overflow-visible",

packages/ui/src/components/chat/message/TextSelectionMenu.tsx

@@ -19,21 +19,48 @@ export const TextSelectionMenu: React.FC<TextSelectionMenuProps> = ({ containerR
   const [position, setPosition] = React.useState<MenuPosition>({ x: 0, y: 0, show: false });
   const [selectedText, setSelectedText] = React.useState('');
   const [isDragging, setIsDragging] = React.useState(false);
+  const [isClosing, setIsClosing] = React.useState(false);
   const menuRef = React.useRef<HTMLDivElement>(null);
   const pendingSelectionRef = React.useRef<{ text: string; rect: DOMRect } | null>(null);
+  const hideTimeoutRef = React.useRef<number | null>(null);
   const createSession = useSessionStore((state) => state.createSession);
   const setPendingInputText = useSessionStore((state) => state.setPendingInputText);
   const isMobile = useUIStore((state) => state.isMobile);
 
+  React.useEffect(() => {
+    return () => {
+      if (hideTimeoutRef.current !== null) {
+        window.clearTimeout(hideTimeoutRef.current);
+        hideTimeoutRef.current = null;
+      }
+    };
+  }, []);
+
   const hideMenu = React.useCallback(() => {
-    setPosition((prev) => ({ ...prev, show: false }));
+    if (hideTimeoutRef.current !== null) {
-    setSelectedText('');
+      window.clearTimeout(hideTimeoutRef.current);
-    pendingSelectionRef.current = null;
+      hideTimeoutRef.current = null;
+    }
+
+    setIsClosing(true);
+    hideTimeoutRef.current = window.setTimeout(() => {
+      setPosition((prev) => ({ ...prev, show: false }));
+      setSelectedText('');
+      pendingSelectionRef.current = null;
+      setIsClosing(false);
+      hideTimeoutRef.current = null;
+    }, 140);
   }, []);
 
   const showMenu = React.useCallback(() => {
     if (!pendingSelectionRef.current) return;
 
+    if (hideTimeoutRef.current !== null) {
+      window.clearTimeout(hideTimeoutRef.current);
+      hideTimeoutRef.current = null;
+    }
+    setIsClosing(false);
+
     const { text, rect } = pendingSelectionRef.current;
 
     // Position menu above the selection
@@ -147,13 +174,8 @@ export const TextSelectionMenu: React.FC<TextSelectionMenuProps> = ({ containerR
 
   const handleAddToChat = React.useCallback(() => {
     if (!selectedText) return;
-    
+
-    // Append to current input
+    setPendingInputText(selectedText, 'append');
-    const currentPending = useSessionStore.getState().pendingInputText || '';
-    const newText = currentPending 
-      ? `${currentPending} ${selectedText}` 
-      : selectedText;
-    setPendingInputText(newText);
     
     hideMenu();
     
@@ -166,7 +188,7 @@ export const TextSelectionMenu: React.FC<TextSelectionMenuProps> = ({ containerR
 
     const session = await createSession(undefined, null, null);
     if (session) {
-      setPendingInputText(selectedText);
+      setPendingInputText(selectedText, 'replace');
     }
 
     hideMenu();
@@ -197,18 +219,18 @@ export const TextSelectionMenu: React.FC<TextSelectionMenuProps> = ({ containerR
           'fixed left-0 right-0 bottom-0 z-50',
           'flex items-center justify-center gap-4',
           'bg-[var(--surface-elevated)] border-t border-[var(--interactive-border)]',
-          'px-4 py-3',
+          'px-3 py-2',
           'safe-area-bottom',
-          'animate-in slide-in-from-bottom duration-200'
+          isClosing ? 'animate-out fade-out-0 duration-150 pointer-events-none' : 'animate-in fade-in-0 duration-150'
         )}
         style={{
-          paddingBottom: 'calc(0.75rem + env(safe-area-inset-bottom, 0px))',
+          paddingBottom: 'calc(0.5rem + env(safe-area-inset-bottom, 0px))',
         }}
       >
         <button
           onClick={handleAddToChat}
           className={cn(
-            'flex items-center gap-2 px-4 py-2.5 rounded-lg',
+            'flex items-center gap-2 px-3 py-2 rounded-lg',
             'text-sm font-medium',
             'bg-[var(--primary-base)] text-[var(--primary-foreground)]',
             'active:opacity-80',
@@ -223,7 +245,7 @@ export const TextSelectionMenu: React.FC<TextSelectionMenuProps> = ({ containerR
         <button
           onClick={handleCreateNewSession}
           className={cn(
-            'flex items-center gap-2 px-4 py-2.5 rounded-lg',
+            'flex items-center gap-2 px-3 py-2 rounded-lg',
             'text-sm font-medium',
             'bg-[var(--interactive-selection)] text-[var(--interactive-selection-foreground)]',
             'active:opacity-80',
@@ -238,7 +260,7 @@ export const TextSelectionMenu: React.FC<TextSelectionMenuProps> = ({ containerR
         <button
           onClick={handleCopy}
           className={cn(
-            'flex items-center gap-2 px-4 py-2.5 rounded-lg',
+            'flex items-center gap-2 px-3 py-2 rounded-lg',
             'text-sm font-medium',
             'bg-[var(--surface-muted)] text-[var(--surface-foreground)]',
             'active:opacity-80',
@@ -262,8 +284,8 @@ export const TextSelectionMenu: React.FC<TextSelectionMenuProps> = ({ containerR
         'fixed z-50 flex items-center gap-1',
         'rounded-lg border border-[var(--interactive-border)]',
         'bg-[var(--surface-elevated)] shadow-lg',
-        'px-2 py-1.5',
+        'px-1.5 py-1',
-        'animate-in fade-in zoom-in-95 duration-150'
+        isClosing ? 'animate-out fade-out-0 duration-150 pointer-events-none' : 'animate-in fade-in-0 duration-150'
       )}
       style={{
         left: position.x,

packages/ui/src/stores/types/sessionTypes.ts

@@ -144,6 +144,7 @@ export interface SessionStore {
     userSummaryTitles: Map<string, { title: string; createdAt: number | null }>;
 
     pendingInputText: string | null;
+    pendingInputMode: 'replace' | 'append';
 
     newSessionDraft: NewSessionDraftState;
 
@@ -241,6 +242,6 @@ export interface SessionStore {
      handleSlashUndo: (sessionId: string) => Promise<void>;
      handleSlashRedo: (sessionId: string) => Promise<void>;
      forkFromMessage: (sessionId: string, messageId: string) => Promise<void>;
-     setPendingInputText: (text: string | null) => void;
+     setPendingInputText: (text: string | null, mode?: 'replace' | 'append') => void;
-     consumePendingInputText: () => string | null;
+     consumePendingInputText: () => { text: string; mode: 'replace' | 'append' } | null;
- }
+  }

packages/ui/src/stores/useSessionStore.ts

@@ -101,6 +101,7 @@ export const useSessionStore = create<SessionStore>()(
             sessionStatus: new Map(),
             userSummaryTitles: new Map(),
             pendingInputText: null,
+            pendingInputMode: 'replace',
             newSessionDraft: { open: true, directoryOverride: null, parentID: null },
 
                 getSessionAgentEditMode: (sessionId: string, agentName: string | undefined, defaultMode?: EditPermissionMode) => {
@@ -607,7 +608,7 @@ export const useSessionStore = create<SessionStore>()(
 
                     // Set pending input text for ChatInput to consume
                     if (messageText) {
-                        set({ pendingInputText: messageText });
+                        set({ pendingInputText: messageText, pendingInputMode: 'replace' });
                     }
                 },
 
@@ -729,7 +730,7 @@ export const useSessionStore = create<SessionStore>()(
 
                         // 4. Show fork point as pending input (will populate ChatInput)
                         if (inputText) {
-                            set({ pendingInputText: inputText });
+                            set({ pendingInputText: inputText, pendingInputMode: 'replace' });
                         }
 
                         // Load the new session's messages
@@ -744,16 +745,20 @@ export const useSessionStore = create<SessionStore>()(
                     }
                 },
 
-                setPendingInputText: (text: string | null) => {
+                setPendingInputText: (text: string | null, mode: 'replace' | 'append' = 'replace') => {
-                    set({ pendingInputText: text });
+                    set({ pendingInputText: text, pendingInputMode: mode });
                 },
 
                 consumePendingInputText: () => {
                     const text = get().pendingInputText;
+                    const mode = get().pendingInputMode;
                     if (text !== null) {
-                        set({ pendingInputText: null });
+                        set({ pendingInputText: null, pendingInputMode: 'replace' });
                     }
-                    return text;
+                    if (text === null) {
+                        return null;
+                    }
+                    return { text, mode };
                 },
             }),
         {