From 3e459bd2e58143a0c74d8157c2b0e2ff9d4df379 Mon Sep 17 00:00:00 2001
From: Justin Choo <justin.k.choo@gmail.com>
Date: Sun, 1 Mar 2026 17:31:58 -0800
Subject: [PATCH] fix(agent-chat): preserve draft text when pane splits

Pane splits restructure the React tree, causing ChatComposer to
unmount/remount and lose its useState text. Add a module-level
Map<paneId, string> draft store that survives React lifecycle.
ChatComposer reads from the store on mount and syncs on every
keystroke. Drafts are cleaned up when panes or tabs are closed.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/components/agent-chat/AgentChatView.tsx   |  1 +
 src/components/agent-chat/ChatComposer.tsx    | 51 ++++++++++++---
 src/components/panes/PaneContainer.tsx        |  4 +-
 src/lib/draft-store.ts                        | 25 ++++++++
 src/store/tabsSlice.ts                        |  4 +-
 .../agent-chat/ChatComposer.test.tsx          | 63 +++++++++++++++++++
 test/unit/client/lib/draft-store.test.ts      | 38 +++++++++++
 7 files changed, 175 insertions(+), 11 deletions(-)
 create mode 100644 src/lib/draft-store.ts
 create mode 100644 test/unit/client/lib/draft-store.test.ts

diff --git a/src/components/agent-chat/AgentChatView.tsx b/src/components/agent-chat/AgentChatView.tsx
index 926ee031..c1b36c3c 100644
--- a/src/components/agent-chat/AgentChatView.tsx
+++ b/src/components/agent-chat/AgentChatView.tsx
@@ -614,6 +614,7 @@ export default function AgentChatView({ tabId, paneId, paneContent, hidden }: Ag
       {/* Composer */}
       <ChatComposer
         ref={composerRef}
+        paneId={paneId}
         onSend={handleSend}
         onInterrupt={handleInterrupt}
         disabled={!isInteractive && !isRunning}
diff --git a/src/components/agent-chat/ChatComposer.tsx b/src/components/agent-chat/ChatComposer.tsx
index dbcc54c0..eebae0ff 100644
--- a/src/components/agent-chat/ChatComposer.tsx
+++ b/src/components/agent-chat/ChatComposer.tsx
@@ -1,12 +1,14 @@
-import { forwardRef, useCallback, useImperativeHandle, useRef, useState, type KeyboardEvent } from 'react'
+import { forwardRef, useCallback, useEffect, useImperativeHandle, useRef, useState, type KeyboardEvent } from 'react'
 import { Send, Square } from 'lucide-react'
 import { cn } from '@/lib/utils'
+import { getDraft, setDraft, clearDraft } from '@/lib/draft-store'
 
 export interface ChatComposerHandle {
   focus: () => void
 }
 
 interface ChatComposerProps {
+  paneId?: string
   onSend: (text: string) => void
   onInterrupt: () => void
   disabled?: boolean
@@ -15,10 +17,28 @@ interface ChatComposerProps {
   autoFocus?: boolean
 }
 
-const ChatComposer = forwardRef<ChatComposerHandle, ChatComposerProps>(function ChatComposer({ onSend, onInterrupt, disabled, isRunning, placeholder, autoFocus }, ref) {
-  const [text, setText] = useState('')
+const ChatComposer = forwardRef<ChatComposerHandle, ChatComposerProps>(function ChatComposer({ paneId, onSend, onInterrupt, disabled, isRunning, placeholder, autoFocus }, ref) {
+  const [text, setText] = useState(() => (paneId ? getDraft(paneId) : ''))
   const textareaRef = useRef<HTMLTextAreaElement | null>(null)
 
+  const resizeTextarea = useCallback((el: HTMLTextAreaElement) => {
+    el.style.height = 'auto'
+    el.style.height = `${Math.min(el.scrollHeight, 200)}px`
+  }, [])
+
+  // Resync text state and textarea height if paneId changes (component reused for a different pane)
+  const prevPaneIdRef = useRef(paneId)
+  useEffect(() => {
+    if (paneId !== prevPaneIdRef.current) {
+      prevPaneIdRef.current = paneId
+      setText(paneId ? getDraft(paneId) : '')
+      // Schedule resize after React paints the new text
+      requestAnimationFrame(() => {
+        if (textareaRef.current) resizeTextarea(textareaRef.current)
+      })
+    }
+  }, [paneId, resizeTextarea])
+
   useImperativeHandle(ref, () => ({
     focus: () => textareaRef.current?.focus(),
   }), [])
@@ -40,16 +60,23 @@ const ChatComposer = forwardRef<ChatComposerHandle, ChatComposerProps>(function
     }
   }, [])
 
+  // Sync draft store on every text change
+  const handleTextChange = useCallback((value: string) => {
+    setText(value)
+    if (paneId) setDraft(paneId, value)
+  }, [paneId])
+
   const handleSend = useCallback(() => {
     const trimmed = text.trim()
     if (!trimmed) return
     onSend(trimmed)
     setText('')
+    if (paneId) clearDraft(paneId)
     // Reset height
     if (textareaRef.current) {
       textareaRef.current.style.height = 'auto'
     }
-  }, [text, onSend])
+  }, [text, onSend, paneId])
 
   const handleKeyDown = useCallback((e: KeyboardEvent<HTMLTextAreaElement>) => {
     if (e.key === 'Enter' && !e.shiftKey) {
@@ -63,10 +90,16 @@ const ChatComposer = forwardRef<ChatComposerHandle, ChatComposerProps>(function
   }, [handleSend, isRunning, onInterrupt])
 
   const handleInput = useCallback(() => {
-    const el = textareaRef.current
-    if (!el) return
-    el.style.height = 'auto'
-    el.style.height = `${Math.min(el.scrollHeight, 200)}px`
+    if (textareaRef.current) resizeTextarea(textareaRef.current)
+  }, [resizeTextarea])
+
+  // Restore textarea height when mounting with a saved draft
+  useEffect(() => {
+    if (text && textareaRef.current) {
+      resizeTextarea(textareaRef.current)
+    }
+    // Only run on mount — text is intentionally excluded
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [])
 
   return (
@@ -75,7 +108,7 @@ const ChatComposer = forwardRef<ChatComposerHandle, ChatComposerProps>(function
         <textarea
           ref={textareaCallbackRef}
           value={text}
-          onChange={(e) => setText(e.target.value)}
+          onChange={(e) => handleTextChange(e.target.value)}
           onKeyDown={handleKeyDown}
           onInput={handleInput}
           disabled={disabled}
diff --git a/src/components/panes/PaneContainer.tsx b/src/components/panes/PaneContainer.tsx
index a3fdc01b..7cffa080 100644
--- a/src/components/panes/PaneContainer.tsx
+++ b/src/components/panes/PaneContainer.tsx
@@ -14,6 +14,7 @@ import PanePicker, { type PanePickerType } from './PanePicker'
 import DirectoryPicker from './DirectoryPicker'
 import { getProviderLabel, isCodingCliProviderName } from '@/lib/coding-cli-utils'
 import { isAgentChatProviderName, getAgentChatProviderConfig } from '@/lib/agent-chat-utils'
+import { clearDraft } from '@/lib/draft-store'
 import { getTerminalActions } from '@/lib/pane-action-registry'
 import { cn } from '@/lib/utils'
 import { getWsClient } from '@/lib/ws-client'
@@ -209,8 +210,9 @@ export default function PaneContainer({ tabId, node, hidden }: PaneContainerProp
         dispatch(updateTab({ id: tabId, updates: { terminalId: undefined } }))
       }
     }
-    // Clean up SDK session if this pane has one
+    // Clean up agent-chat resources
     if (content.kind === 'agent-chat') {
+      clearDraft(paneId)
       const sessionId = content.sessionId || sdkPendingCreates[content.createRequestId]
       if (sessionId) {
         ws.send({ type: 'sdk.kill', sessionId })
diff --git a/src/lib/draft-store.ts b/src/lib/draft-store.ts
new file mode 100644
index 00000000..52aeee2d
--- /dev/null
+++ b/src/lib/draft-store.ts
@@ -0,0 +1,25 @@
+/**
+ * Module-level store for chat composer draft text.
+ *
+ * Survives React component unmount/remount (e.g. pane splits) without
+ * involving Redux or localStorage.  Keyed by paneId so each pane keeps
+ * its own independent draft.
+ */
+
+const drafts = new Map<string, string>()
+
+export function getDraft(paneId: string): string {
+  return drafts.get(paneId) ?? ''
+}
+
+export function setDraft(paneId: string, text: string): void {
+  if (text) {
+    drafts.set(paneId, text)
+  } else {
+    drafts.delete(paneId)
+  }
+}
+
+export function clearDraft(paneId: string): void {
+  drafts.delete(paneId)
+}
diff --git a/src/store/tabsSlice.ts b/src/store/tabsSlice.ts
index d7688e78..5e0d5a9a 100644
--- a/src/store/tabsSlice.ts
+++ b/src/store/tabsSlice.ts
@@ -9,6 +9,7 @@ import { getProviderLabel } from '@/lib/coding-cli-utils'
 import { buildResumeContent } from '@/lib/session-type-utils'
 import { isAgentChatProviderName, getAgentChatProviderConfig, getAgentChatProviderLabel } from '@/lib/agent-chat-utils'
 import { recordClosedTabSnapshot } from './tabRegistrySlice'
+import { clearDraft } from '@/lib/draft-store'
 import {
   buildClosedTabRegistryRecord,
   countPaneLeaves,
@@ -282,10 +283,11 @@ export const closeTab = createAsyncThunk(
     dispatch(removeTab(tabId))
     dispatch(removeLayout({ tabId }))
 
-    // Clean up attention for the tab and all its panes
+    // Clean up attention and drafts for the tab and all its panes
     dispatch(clearTabAttention({ tabId }))
     for (const paneId of paneIds) {
       dispatch(clearPaneAttention({ paneId }))
+      clearDraft(paneId)
     }
   }
 )
diff --git a/test/unit/client/components/agent-chat/ChatComposer.test.tsx b/test/unit/client/components/agent-chat/ChatComposer.test.tsx
index 0ed7e739..888d03e0 100644
--- a/test/unit/client/components/agent-chat/ChatComposer.test.tsx
+++ b/test/unit/client/components/agent-chat/ChatComposer.test.tsx
@@ -2,10 +2,14 @@ import { describe, it, expect, vi, afterEach } from 'vitest'
 import { render, screen, cleanup } from '@testing-library/react'
 import userEvent from '@testing-library/user-event'
 import ChatComposer from '../../../../../src/components/agent-chat/ChatComposer'
+import { getDraft, clearDraft } from '@/lib/draft-store'
 
 describe('ChatComposer', () => {
   afterEach(() => {
     cleanup()
+    clearDraft('test-pane')
+    clearDraft('pane-a')
+    clearDraft('pane-b')
   })
   it('renders textarea and send button', () => {
     render(<ChatComposer onSend={() => {}} onInterrupt={() => {}} />)
@@ -74,4 +78,63 @@ describe('ChatComposer', () => {
     await user.keyboard('{Escape}')
     expect(onInterrupt).not.toHaveBeenCalled()
   })
+
+  describe('draft preservation', () => {
+    it('starts empty when no saved draft exists', () => {
+      render(<ChatComposer paneId="test-pane" onSend={() => {}} onInterrupt={() => {}} />)
+      expect(screen.getByRole('textbox')).toHaveValue('')
+    })
+
+    it('restores draft text after unmount and remount', async () => {
+      const user = userEvent.setup()
+      const { unmount } = render(
+        <ChatComposer paneId="test-pane" onSend={() => {}} onInterrupt={() => {}} />
+      )
+      await user.type(screen.getByRole('textbox'), 'work in progress')
+      unmount()
+
+      // Remount with the same paneId
+      render(<ChatComposer paneId="test-pane" onSend={() => {}} onInterrupt={() => {}} />)
+      expect(screen.getByRole('textbox')).toHaveValue('work in progress')
+    })
+
+    it('clears draft after sending a message', async () => {
+      const user = userEvent.setup()
+      render(<ChatComposer paneId="test-pane" onSend={() => {}} onInterrupt={() => {}} />)
+      await user.type(screen.getByRole('textbox'), 'sent message{Enter}')
+      expect(getDraft('test-pane')).toBe('')
+    })
+
+    it('keeps independent drafts per paneId', async () => {
+      const user = userEvent.setup()
+      const { unmount: unmountA } = render(
+        <ChatComposer paneId="pane-a" onSend={() => {}} onInterrupt={() => {}} />
+      )
+      await user.type(screen.getByRole('textbox'), 'draft A')
+      unmountA()
+
+      const { unmount: unmountB } = render(
+        <ChatComposer paneId="pane-b" onSend={() => {}} onInterrupt={() => {}} />
+      )
+      await user.type(screen.getByRole('textbox'), 'draft B')
+      unmountB()
+
+      // Remount A — should have its own draft
+      render(<ChatComposer paneId="pane-a" onSend={() => {}} onInterrupt={() => {}} />)
+      expect(screen.getByRole('textbox')).toHaveValue('draft A')
+    })
+
+    it('works without paneId (no persistence, backwards compatible)', async () => {
+      const user = userEvent.setup()
+      const { unmount } = render(
+        <ChatComposer onSend={() => {}} onInterrupt={() => {}} />
+      )
+      await user.type(screen.getByRole('textbox'), 'no pane id')
+      unmount()
+
+      // Remount without paneId — starts empty (no crash)
+      render(<ChatComposer onSend={() => {}} onInterrupt={() => {}} />)
+      expect(screen.getByRole('textbox')).toHaveValue('')
+    })
+  })
 })
diff --git a/test/unit/client/lib/draft-store.test.ts b/test/unit/client/lib/draft-store.test.ts
new file mode 100644
index 00000000..e57d75e2
--- /dev/null
+++ b/test/unit/client/lib/draft-store.test.ts
@@ -0,0 +1,38 @@
+import { describe, it, expect, beforeEach } from 'vitest'
+import { getDraft, setDraft, clearDraft } from '@/lib/draft-store'
+
+describe('draft-store', () => {
+  beforeEach(() => {
+    // Clean slate — clear any drafts from previous tests
+    clearDraft('pane-a')
+    clearDraft('pane-b')
+  })
+
+  it('returns empty string when no draft exists', () => {
+    expect(getDraft('nonexistent')).toBe('')
+  })
+
+  it('stores and retrieves a draft', () => {
+    setDraft('pane-a', 'hello world')
+    expect(getDraft('pane-a')).toBe('hello world')
+  })
+
+  it('keeps independent drafts per paneId', () => {
+    setDraft('pane-a', 'draft A')
+    setDraft('pane-b', 'draft B')
+    expect(getDraft('pane-a')).toBe('draft A')
+    expect(getDraft('pane-b')).toBe('draft B')
+  })
+
+  it('clears a specific draft', () => {
+    setDraft('pane-a', 'some text')
+    clearDraft('pane-a')
+    expect(getDraft('pane-a')).toBe('')
+  })
+
+  it('removes draft when set to empty string', () => {
+    setDraft('pane-a', 'some text')
+    setDraft('pane-a', '')
+    expect(getDraft('pane-a')).toBe('')
+  })
+})