From f8103d6859653c4011f4615084632bad7f372e8e Mon Sep 17 00:00:00 2001
From: chliny <chliny11@gmail.com>
Date: Tue, 18 Aug 2026 13:38:56 +0800
Subject: [PATCH] fix: preserve session-specific chat drafts

---
 app/session/[id].tsx   | 25 ++++++++++++++++++++-----
 src/stores/sessions.ts | 17 +++++++++++++++++
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/app/session/[id].tsx b/app/session/[id].tsx
index 97f885f0..67c38096 100644
--- a/app/session/[id].tsx
+++ b/app/session/[id].tsx
@@ -82,7 +82,9 @@ export default function SessionScreen() {
   const flatListRef = useRef<FlatList>(null)
   const modelSheetRef = useRef<BottomSheet>(null)
   const variantSheetRef = useRef<BottomSheet>(null)
-  const [input, setInput] = useState("")
+  const [input, setInputState] = useState(() => (id ? useSessions.getState().drafts[id] || "" : ""))
+  const inputRef = useRef(input)
+  inputRef.current = input
   const [attachments, setAttachments] = useState<Attachment[]>([])
   const [showInfo, setShowInfo] = useState(false)
 
@@ -94,6 +96,8 @@ export default function SessionScreen() {
     loadingMore,
     hasMore,
     selectSession,
+    setDraft,
+    clearDraft,
     sendMessage,
     abortSession,
     loadOlderMessages,
@@ -101,6 +105,16 @@ export default function SessionScreen() {
     unrevertSession,
   } = useSessions()
 
+  const setInput = useCallback(
+    (value: string | ((current: string) => string)) => {
+      const next = typeof value === "function" ? value(inputRef.current) : value
+      inputRef.current = next
+      setInputState(next)
+      if (id) setDraft(id, next)
+    },
+    [id, setDraft],
+  )
+
   // Derive sending state for this specific session
   const isSending = useSessions((s) => !!(currentSession && s.sending[currentSession.id]))
 
@@ -194,9 +208,6 @@ export default function SessionScreen() {
   // handleMessageLongPress's deps — kept as a plain ref assignment (not
   // state) so the callback below stays referentially stable across
   // keystrokes for MessageBubble's custom memo comparator.
-  const inputRef = useRef(input)
-  inputRef.current = input
-
   const applyRevertResult = useCallback((result: Awaited<ReturnType<typeof revertToMessage>>) => {
     if (!result.ok) {
       if (result.reason === "unsupported") {
@@ -265,6 +276,9 @@ export default function SessionScreen() {
   useFocusEffect(
     useCallback(() => {
       if (!id) return
+      const draft = useSessions.getState().drafts[id] || ""
+      inputRef.current = draft
+      setInputState(draft)
       selectSession(id, directory).then(() => {
         // Re-fetch pending permissions/questions from the server to recover from
         // missed SSE events or failed optimistic removals
@@ -272,7 +286,7 @@ export default function SessionScreen() {
         const c = directory ? (connState.clientForDirectory(directory) ?? connState.client) : connState.client
         if (c) refreshPending(c, id)
       })
-    }, [id, directory]),
+    }, [id, directory, selectSession]),
   )
 
   // Sync model chip from latest assistant message
@@ -417,6 +431,7 @@ export default function SessionScreen() {
     const text = input.trim()
     const files = [...attachments]
     setInput("")
+    if (id) clearDraft(id)
     setAttachments([])
 
     // Server slash commands (no attachments for commands)
diff --git a/src/stores/sessions.ts b/src/stores/sessions.ts
index ccd21f33..26372bf2 100644
--- a/src/stores/sessions.ts
+++ b/src/stores/sessions.ts
@@ -33,6 +33,8 @@ interface SessionsState {
   isLoading: boolean
   // Per-session optimistic sending flag — bridging gap between user tap and SSE busy
   sending: Record<string, boolean>
+  // Unsent composer text, isolated by session ID.
+  drafts: Record<string, string>
   loadingMore: boolean
   hasMore: boolean
   error: string | null
@@ -40,6 +42,8 @@ interface SessionsState {
   // Actions
   loadSessions: () => Promise<void>
   selectSession: (sessionID: string, directory?: string) => Promise<void>
+  setDraft: (sessionID: string, text: string) => void
+  clearDraft: (sessionID: string) => void
   loadOlderMessages: () => Promise<void>
   createSession: (title?: string) => Promise<Session | null>
   deleteSession: (sessionID: string) => Promise<void>
@@ -93,10 +97,22 @@ export const useSessions = create<SessionsState>((set, get) => ({
   parts: {},
   isLoading: false,
   sending: {},
+  drafts: {},
   loadingMore: false,
   hasMore: false,
   error: null,
 
+  setDraft: (sessionID, text) =>
+    set((state) => ({ drafts: { ...state.drafts, [sessionID]: text } })),
+
+  clearDraft: (sessionID) =>
+    set((state) => {
+      if (!(sessionID in state.drafts)) return state
+      const drafts = { ...state.drafts }
+      delete drafts[sessionID]
+      return { drafts }
+    }),
+
   loadSessions: async () => {
     const connState = useConnections.getState()
     // Use a directory-less client so the server returns sessions from ALL projects,
@@ -245,6 +261,7 @@ export const useSessions = create<SessionsState>((set, get) => ({
       await client.session.delete(sessionID)
       set((state) => ({
         sessions: state.sessions.filter((s) => s.id !== sessionID),
+        drafts: Object.fromEntries(Object.entries(state.drafts).filter(([id]) => id !== sessionID)),
         currentSession: state.currentSession?.id === sessionID ? null : state.currentSession,
         messages: state.currentSession?.id === sessionID ? [] : state.messages,
         parts: state.currentSession?.id === sessionID ? {} : state.parts,