From 035001f841e7008a5fbfcdea92ed36f64fa3facf Mon Sep 17 00:00:00 2001
From: suyao <sy20010504@gmail.com>
Date: Mon, 22 Sep 2025 22:53:57 +0800
Subject: [PATCH] Ensure session messages are properly restored when switching
 sessions

- Add Redux selector to check for existing messages in store
- Always reload messages to Redux when session data is available
- Add effect to restore messages when component mounts if missing from Redux
---
 src/renderer/src/hooks/agents/useSession.ts | 32 +++++++++++++++------
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/src/renderer/src/hooks/agents/useSession.ts b/src/renderer/src/hooks/agents/useSession.ts
index ba54447c52..22092eb681 100644
--- a/src/renderer/src/hooks/agents/useSession.ts
+++ b/src/renderer/src/hooks/agents/useSession.ts
@@ -1,6 +1,6 @@
-import { useAppDispatch } from '@renderer/store'
-import { removeManyBlocks, upsertManyBlocks } from '@renderer/store/messageBlock'
-import { newMessagesActions } from '@renderer/store/newMessage'
+import { useAppDispatch, useAppSelector } from '@renderer/store'
+import { upsertManyBlocks } from '@renderer/store/messageBlock'
+import { newMessagesActions, selectMessagesForTopic } from '@renderer/store/newMessage'
 import { AgentPersistedMessage, UpdateSessionForm } from '@renderer/types'
 import { buildAgentSessionTopicId } from '@renderer/utils/agentSession'
 import { useCallback, useEffect, useMemo, useRef } from 'react'
@@ -17,6 +17,9 @@ export const useSession = (agentId: string, sessionId: string) => {
   const sessionTopicId = useMemo(() => buildAgentSessionTopicId(sessionId), [sessionId])
   const blockIdsRef = useRef<string[]>([])
 
+  // Check if messages are already in Redux
+  const messagesInRedux = useAppSelector((state) => selectMessagesForTopic(state, sessionTopicId))
+
   const fetcher = async () => {
     const data = await client.getSession(agentId, sessionId)
     return data
@@ -25,6 +28,9 @@ export const useSession = (agentId: string, sessionId: string) => {
 
   useEffect(() => {
     const messages = data?.messages ?? []
+
+    // Always reload messages to Redux when data is available
+    // This ensures messages are restored when switching back to a session
     if (!messages.length) {
       dispatch(newMessagesActions.messagesReceived({ topicId: sessionTopicId, messages: [] }))
       blockIdsRef.current = []
@@ -46,14 +52,24 @@ export const useSession = (agentId: string, sessionId: string) => {
     dispatch(newMessagesActions.messagesReceived({ topicId: sessionTopicId, messages: messageRecords }))
   }, [data?.messages, dispatch, sessionTopicId])
 
+  // Also ensure messages are reloaded when component mounts if they're missing from Redux
   useEffect(() => {
-    return () => {
-      if (blockIdsRef.current.length > 0) {
-        dispatch(removeManyBlocks(blockIdsRef.current))
+    // If we have data but no messages in Redux, reload them
+    if (data?.messages && data.messages.length > 0 && messagesInRedux.length === 0) {
+      const messages = data.messages
+      const persistedEntries = messages
+        .map((entity) => entity.content as AgentPersistedMessage | undefined)
+        .filter((entry): entry is AgentPersistedMessage => Boolean(entry))
+
+      const allBlocks = persistedEntries.flatMap((entry) => entry.blocks)
+      if (allBlocks.length > 0) {
+        dispatch(upsertManyBlocks(allBlocks))
       }
-      dispatch(newMessagesActions.clearTopicMessages(sessionTopicId))
+
+      const messageRecords = persistedEntries.map((entry) => entry.message)
+      dispatch(newMessagesActions.messagesReceived({ topicId: sessionTopicId, messages: messageRecords }))
     }
-  }, [dispatch, sessionTopicId])
+  }, [data?.messages, dispatch, messagesInRedux.length, sessionTopicId])
 
   const updateSession = useCallback(
     async (form: UpdateSessionForm) => {