From de0a8f5e41e0b055f60d2c21208f462c645f6379 Mon Sep 17 00:00:00 2001
From: hanyuanling <ushouldknowr0@gmail.com>
Date: Thu, 21 May 2026 10:59:15 +0800
Subject: [PATCH] fix(webui): keep new chat during session refresh

---
 webui/src/hooks/useSessions.ts       | 16 +++++++++-
 webui/src/tests/useSessions.test.tsx | 47 ++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/webui/src/hooks/useSessions.ts b/webui/src/hooks/useSessions.ts
index c22751c65..7b468fc89 100644
--- a/webui/src/hooks/useSessions.ts
+++ b/webui/src/hooks/useSessions.ts
@@ -27,13 +27,25 @@ export function useSessions(): {
   const [loading, setLoading] = useState(true);
   const [error, setError] = useState<string | null>(null);
   const tokenRef = useRef(token);
+  const optimisticKeysRef = useRef<Set<string>>(new Set());
   tokenRef.current = token;
 
   const refresh = useCallback(async () => {
     try {
       setLoading(true);
       const rows = await listSessions(tokenRef.current);
-      setSessions(rows);
+      const serverKeys = new Set(rows.map((row) => row.key));
+      setSessions((prev) => [
+        ...rows,
+        ...prev.filter(
+          (session) =>
+            optimisticKeysRef.current.has(session.key) &&
+            !serverKeys.has(session.key),
+        ),
+      ]);
+      for (const key of Array.from(optimisticKeysRef.current)) {
+        if (serverKeys.has(key)) optimisticKeysRef.current.delete(key);
+      }
       setError(null);
     } catch (e) {
       const msg =
@@ -57,6 +69,7 @@ export function useSessions(): {
   const createChat = useCallback(async (): Promise<string> => {
     const chatId = await client.newChat();
     const key = `websocket:${chatId}`;
+    optimisticKeysRef.current.add(key);
     // Optimistic insert; a subsequent refresh will replace it with the
     // authoritative row once the server persists the session.
     setSessions((prev) => [
@@ -77,6 +90,7 @@ export function useSessions(): {
   const deleteChat = useCallback(
     async (key: string) => {
       await apiDeleteSession(tokenRef.current, key);
+      optimisticKeysRef.current.delete(key);
       setSessions((prev) => prev.filter((s) => s.key !== key));
     },
     [],
diff --git a/webui/src/tests/useSessions.test.tsx b/webui/src/tests/useSessions.test.tsx
index 72df813e0..8e76e697e 100644
--- a/webui/src/tests/useSessions.test.tsx
+++ b/webui/src/tests/useSessions.test.tsx
@@ -157,6 +157,53 @@ describe("useSessions", () => {
     expect(api.listSessions).toHaveBeenCalledTimes(2);
   });
 
+  it("keeps a newly created chat visible until the server session list catches up", async () => {
+    vi.mocked(api.listSessions)
+      .mockResolvedValueOnce([])
+      .mockResolvedValueOnce([])
+      .mockResolvedValueOnce([
+        {
+          key: "websocket:chat-new",
+          channel: "websocket",
+          chatId: "chat-new",
+          createdAt: "2026-05-20T10:00:00Z",
+          updatedAt: "2026-05-20T10:01:00Z",
+          title: "Generated title",
+          preview: "First message",
+        },
+      ]);
+    const client = fakeClient();
+    client.newChat.mockResolvedValue("chat-new");
+
+    const { result } = renderHook(() => useSessions(), {
+      wrapper: wrap(client),
+    });
+
+    await waitFor(() => expect(result.current.loading).toBe(false));
+    expect(result.current.sessions).toEqual([]);
+
+    await act(async () => {
+      await result.current.createChat();
+    });
+
+    expect(result.current.sessions.map((s) => s.key)).toEqual(["websocket:chat-new"]);
+
+    await act(async () => {
+      await result.current.refresh();
+    });
+
+    expect(result.current.sessions.map((s) => s.key)).toEqual(["websocket:chat-new"]);
+    expect(result.current.sessions[0]?.preview).toBe("");
+
+    await act(async () => {
+      await result.current.refresh();
+    });
+
+    expect(result.current.sessions.map((s) => s.key)).toEqual(["websocket:chat-new"]);
+    expect(result.current.sessions[0]?.preview).toBe("First message");
+    expect(result.current.sessions[0]?.title).toBe("Generated title");
+  });
+
   it("passes through WebUI transcript user media as images and media", async () => {
     vi.mocked(api.fetchWebuiThread).mockResolvedValue({
       schemaVersion: 3,