From b3e409ba6e042cdfaacc6033604cf110ee26a635 Mon Sep 17 00:00:00 2001 From: Abhin Rustagi Date: Tue, 21 Jul 2026 16:23:27 +0530 Subject: [PATCH 1/4] feat: add a loader on sidebar for streaming threads --- .../components/AgentInterface/ThreadList.tsx | 84 +++++++++++-------- .../components/AgentInterface/threadlist.scss | 12 +++ 2 files changed, 61 insertions(+), 35 deletions(-) diff --git a/packages/react-ui/src/components/AgentInterface/ThreadList.tsx b/packages/react-ui/src/components/AgentInterface/ThreadList.tsx index e62c619dd..8a29c7a4c 100644 --- a/packages/react-ui/src/components/AgentInterface/ThreadList.tsx +++ b/packages/react-ui/src/components/AgentInterface/ThreadList.tsx @@ -1,5 +1,5 @@ import type { Thread } from "@openuidev/react-headless"; -import { useThreadList } from "@openuidev/react-headless"; +import { useThread, useThreadList } from "@openuidev/react-headless"; import * as DropdownMenu from "@radix-ui/react-dropdown-menu"; import clsx from "clsx"; import { EllipsisIcon, Trash2Icon } from "lucide-react"; @@ -22,6 +22,8 @@ export const ThreadButton = ({ const selectThread = useThreadList((s) => s.selectThread); const deleteThread = useThreadList((s) => s.deleteThread); const selectedThreadId = useThreadList((s) => s.selectedThreadId); + const isRunning = useThread((s) => s.isRunning); + const isStreaming = selectedThreadId === id && isRunning; const { isSidebarOpen, setIsSidebarOpen } = useAgentInterfaceStore((state) => ({ isSidebarOpen: state.isSidebarOpen, setIsSidebarOpen: state.setIsSidebarOpen, @@ -56,42 +58,54 @@ export const ThreadButton = ({ > {title} - - - } - size="2-extra-small" - variant="tertiary" - aria-label="Thread actions" - /> - - - - { - deleteThread(id); - }} + {isStreaming && ( + +
+ + )} + {!isStreaming && ( + + + } + size="2-extra-small" + variant="tertiary" + aria-label="Thread actions" + /> + + + - - - - - + + + + + + )}
); }; diff --git a/packages/react-ui/src/components/AgentInterface/threadlist.scss b/packages/react-ui/src/components/AgentInterface/threadlist.scss index c8308649d..c7ffaf7fe 100644 --- a/packages/react-ui/src/components/AgentInterface/threadlist.scss +++ b/packages/react-ui/src/components/AgentInterface/threadlist.scss @@ -91,6 +91,18 @@ } } +.openui-agent-thread-button-loader { + display: flex; + align-items: center; + margin-right: cssUtils.$space-2xs; +} + +@media (prefers-reduced-motion: reduce) { + .openui-agent-thread-button-loader .openui-message-loading { + animation: none; + } +} + .openui-agent-thread-button-dropdown-menu { display: flex; flex-direction: column; From 2ad4e6b2a0d9149c216c5d7909359a4abc4176ef Mon Sep 17 00:00:00 2001 From: Abhin Rustagi Date: Tue, 21 Jul 2026 16:26:21 +0530 Subject: [PATCH 2/4] fix: switch to div --- .../react-ui/src/components/AgentInterface/ThreadList.tsx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/react-ui/src/components/AgentInterface/ThreadList.tsx b/packages/react-ui/src/components/AgentInterface/ThreadList.tsx index 8a29c7a4c..6550931e6 100644 --- a/packages/react-ui/src/components/AgentInterface/ThreadList.tsx +++ b/packages/react-ui/src/components/AgentInterface/ThreadList.tsx @@ -59,14 +59,14 @@ export const ThreadButton = ({ {title} {isStreaming && ( -
- +
)} {!isStreaming && ( From e4a63379f4f166e0471c3e3ae89157c1845b1f17 Mon Sep 17 00:00:00 2001 From: Abhin Rustagi Date: Wed, 22 Jul 2026 00:41:19 +0530 Subject: [PATCH 3/4] fix: update styles --- .../components/AgentInterface/ThreadList.tsx | 5 +-- .../components/AgentInterface/threadlist.scss | 31 +++++++++++++++++-- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/packages/react-ui/src/components/AgentInterface/ThreadList.tsx b/packages/react-ui/src/components/AgentInterface/ThreadList.tsx index bb7080b92..7adc3911a 100644 --- a/packages/react-ui/src/components/AgentInterface/ThreadList.tsx +++ b/packages/react-ui/src/components/AgentInterface/ThreadList.tsx @@ -1,7 +1,7 @@ import { useThread, useThreadList } from "@openuidev/react-headless"; import * as DropdownMenu from "@radix-ui/react-dropdown-menu"; import clsx from "clsx"; -import { EllipsisIcon, Trash2Icon } from "lucide-react"; +import { EllipsisIcon, Loader2, Trash2Icon } from "lucide-react"; import { useCallback, useEffect, useRef, useState } from "react"; import { useLayoutContext } from "../../context/LayoutContext"; import { Button } from "../Button"; @@ -63,6 +63,7 @@ export const ThreadButton = ({ { "openui-agent-thread-button--selected": selectedThreadId === id, "openui-agent-thread-button--actions-open": isActionsOpen, + "openui-agent-thread-button--streaming": isStreaming, }, className, )} @@ -89,7 +90,7 @@ export const ThreadButton = ({ aria-live="polite" aria-label="Generating response" > -
+
)} {!isStreaming && ( diff --git a/packages/react-ui/src/components/AgentInterface/threadlist.scss b/packages/react-ui/src/components/AgentInterface/threadlist.scss index 53104923c..be54baf71 100644 --- a/packages/react-ui/src/components/AgentInterface/threadlist.scss +++ b/packages/react-ui/src/components/AgentInterface/threadlist.scss @@ -201,14 +201,41 @@ $thread-list-mask-height: cssUtils.$space-xl; } } +// Streaming loader occupies the same trailing slot as the (absolute) dropdown +// trigger, so the sidebar row shows exactly one control at a time. .openui-agent-thread-button-loader { + position: absolute; + top: 50%; + right: cssUtils.$space-2xs; + transform: translateY(-50%); display: flex; align-items: center; - margin-right: cssUtils.$space-2xs; + justify-content: center; + width: 24px; + pointer-events: none; +} + +// Simple spinner (lucide Loader2). The SVG is inlined by the icon component; only +// the spin keyframe needs to live here, so it always ships in this CSS bundle. +.openui-agent-thread-button-loader__icon { + color: cssUtils.$text-neutral-tertiary; + animation: openui-agent-thread-button-spin 0.9s linear infinite; +} + +@keyframes openui-agent-thread-button-spin { + to { + transform: rotate(360deg); + } +} + +// Match the hover behaviour: reserve room for the trailing loader so the title +// ellipsis-truncates before it instead of running underneath. +.openui-agent-thread-button--streaming .openui-agent-thread-button-title { + padding-right: calc(24px + cssUtils.$space-s-m); } @media (prefers-reduced-motion: reduce) { - .openui-agent-thread-button-loader .openui-message-loading { + .openui-agent-thread-button-loader__icon { animation: none; } } From 9f55ba2f791d49b9eea54a46fb41e3a7c5b50da8 Mon Sep 17 00:00:00 2001 From: Abhin Rustagi Date: Wed, 22 Jul 2026 15:01:12 +0530 Subject: [PATCH 4/4] fix: allow background thread streaming --- .../react-headless/src/hooks/useThread.ts | 27 +- .../src/hooks/useThreadState.ts | 32 ++ packages/react-headless/src/index.ts | 2 + .../store/__tests__/createChatStore.test.ts | 358 ++++++++++--- .../src/store/createChatStore.ts | 507 ++++++++++-------- .../src/store/threadStateEntry.ts | 48 ++ packages/react-headless/src/store/types.ts | 10 +- .../AgentInterface/MobileHeader.tsx | 2 + .../AgentInterface/NewChatButton.tsx | 3 + .../components/AgentInterface/ThreadList.tsx | 6 +- 10 files changed, 683 insertions(+), 312 deletions(-) create mode 100644 packages/react-headless/src/hooks/useThreadState.ts create mode 100644 packages/react-headless/src/store/threadStateEntry.ts diff --git a/packages/react-headless/src/hooks/useThread.ts b/packages/react-headless/src/hooks/useThread.ts index 55c601250..1b53c9991 100644 --- a/packages/react-headless/src/hooks/useThread.ts +++ b/packages/react-headless/src/hooks/useThread.ts @@ -1,6 +1,7 @@ import { useStore } from "zustand"; import { useShallow } from "zustand/react/shallow"; import { useChatStore } from "../store/ChatContext"; +import { deriveThreadState, EMPTY_THREAD_STATE, resolveViewKey } from "../store/threadStateEntry"; import type { ChatStore, ThreadActions, @@ -12,19 +13,18 @@ import type { type ThreadSlice = ThreadState & ThreadActions; type ThreadListSlice = ThreadListState & ThreadListActions; -const threadSelector = (s: ChatStore): ThreadSlice => ({ - messages: s.messages, - isRunning: s.isRunning, - isLoadingMessages: s.isLoadingMessages, - threadError: s.threadError, - executingToolCallIds: s.executingToolCallIds, - processMessage: s.processMessage, - appendMessages: s.appendMessages, - updateMessage: s.updateMessage, - setMessages: s.setMessages, - deleteMessage: s.deleteMessage, - cancelMessage: s.cancelMessage, -}); +const threadSelector = (s: ChatStore): ThreadSlice => { + const entry = s.threadStates[resolveViewKey(s.selectedThreadId)] ?? EMPTY_THREAD_STATE; + return { + ...deriveThreadState(entry), + processMessage: s.processMessage, + appendMessages: s.appendMessages, + updateMessage: s.updateMessage, + setMessages: s.setMessages, + deleteMessage: s.deleteMessage, + cancelMessage: s.cancelMessage, + }; +}; const threadListSelector = (s: ChatStore): ThreadListSlice => ({ threads: s.threads, @@ -32,6 +32,7 @@ const threadListSelector = (s: ChatStore): ThreadListSlice => ({ threadListError: s.threadListError, selectedThreadId: s.selectedThreadId, hasMoreThreads: s.hasMoreThreads, + isCreatingThread: s.isCreatingThread, loadThreads: s.loadThreads, loadMoreThreads: s.loadMoreThreads, switchToNewThread: s.switchToNewThread, diff --git a/packages/react-headless/src/hooks/useThreadState.ts b/packages/react-headless/src/hooks/useThreadState.ts new file mode 100644 index 000000000..c8ccbf365 --- /dev/null +++ b/packages/react-headless/src/hooks/useThreadState.ts @@ -0,0 +1,32 @@ +import { useStore } from "zustand"; +import { useShallow } from "zustand/react/shallow"; +import { useChatStore } from "../store/ChatContext"; +import { deriveThreadState, EMPTY_THREAD_STATE } from "../store/threadStateEntry"; +import type { ChatStore, ThreadState } from "../store/types"; + +/** + * Read-only {@link ThreadState} for a thread **by id** — the by-id sibling of + * {@link useThread} (which is hardwired to the *selected* thread). Reflects a run + * even while it streams in the **background** (the thread isn't the one on screen), + * so a sidebar row can show a loader / error / loading state for any thread. + * + * Returns actions-free state only; actions (`processMessage`, `cancelMessage`, …) + * always target the selected thread, so use {@link useThread} for those. + * + * @example + * const isStreaming = useThreadState(threadId, (s) => s.isRunning); + * const { messages, threadError } = useThreadState(threadId); + * + * @category Hooks + */ +export function useThreadState(threadId: string | null): ThreadState; +export function useThreadState(threadId: string | null, selector: (state: ThreadState) => T): T; +export function useThreadState(threadId: string | null, selector?: (state: ThreadState) => T) { + const store = useChatStore(); + const pick = (s: ChatStore): ThreadState => + deriveThreadState((threadId != null && s.threadStates[threadId]) || EMPTY_THREAD_STATE); + if (selector) { + return useStore(store, (s) => selector(pick(s))); + } + return useStore(store, useShallow(pick)); +} diff --git a/packages/react-headless/src/index.ts b/packages/react-headless/src/index.ts index 6bd05b30a..83c079938 100644 --- a/packages/react-headless/src/index.ts +++ b/packages/react-headless/src/index.ts @@ -6,6 +6,7 @@ export { useDetailedView } from "./hooks/useDetailedView"; export { useDetailedViewPortalTarget } from "./hooks/useDetailedViewPortalTarget"; export { MessageContext, MessageProvider, useMessage } from "./hooks/useMessage"; export { useThread, useThreadList } from "./hooks/useThread"; +export { useThreadState } from "./hooks/useThreadState"; export { useToolActivities } from "./hooks/useToolActivities"; export { defineArtifactCategories } from "./store/artifactCategories"; @@ -78,6 +79,7 @@ export type { ThreadListActions, ThreadListState, ThreadState, + ThreadStateEntry, } from "./store/types"; export type { diff --git a/packages/react-headless/src/store/__tests__/createChatStore.test.ts b/packages/react-headless/src/store/__tests__/createChatStore.test.ts index 1d28dfbab..7b549d501 100644 --- a/packages/react-headless/src/store/__tests__/createChatStore.test.ts +++ b/packages/react-headless/src/store/__tests__/createChatStore.test.ts @@ -1,5 +1,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; -import type { Message, Thread, UserMessage } from "../types"; +import { EventType } from "../../types/stream"; +import { DRAFT_KEY, makeThreadState } from "../threadStateEntry"; +import type { Message, Thread, ThreadStateEntry, UserMessage } from "../types"; import { makeStore } from "./__helpers/makeStore"; // ── Helpers ── @@ -15,6 +17,25 @@ const makeMessage = (id: string, role: "user" | "assistant" = "user"): Message = const flushPromises = () => new Promise((r) => setTimeout(r, 0)); +type Store = ReturnType; + +// The key whose entry the view currently shows. +const viewKeyOf = (store: Store) => store.getState().selectedThreadId ?? DRAFT_KEY; +// The selected thread's entry (empty defaults when absent), i.e. what `useThread` shows. +const active = (store: Store): ThreadStateEntry => + store.getState().threadStates[viewKeyOf(store)] ?? makeThreadState(); +// A specific thread's raw entry (may be undefined). +const entryOf = (store: Store, key: string): ThreadStateEntry | undefined => + store.getState().threadStates[key]; +// Seed/patch one thread's entry, creating it if absent (replaces the old `setState({ messages })`). +const seed = (store: Store, key: string, patch: Partial) => + store.setState((s) => ({ + threadStates: { + ...s.threadStates, + [key]: { ...(s.threadStates[key] ?? makeThreadState()), ...patch }, + }, + })); + // ── Test suite ── describe("createChatStore", () => { @@ -106,27 +127,39 @@ describe("createChatStore", () => { }); describe("selectThread", () => { - it("sets selectedThreadId, loads messages, clears previous", async () => { + it("sets selectedThreadId and loads messages into that thread's entry", async () => { const messages: Message[] = [makeMessage("m1"), makeMessage("m2", "assistant")]; const getMessages = vi.fn().mockResolvedValue(messages); const store = makeStore({ getMessages }); - store.setState({ messages: [makeMessage("old")] }); - store.getState().selectThread("t1"); expect(store.getState().selectedThreadId).toBe("t1"); - expect(store.getState().messages).toEqual([]); - expect(store.getState().isLoadingMessages).toBe(true); + expect(active(store).messages).toEqual([]); + expect(active(store).isLoadingMessages).toBe(true); await flushPromises(); - expect(store.getState().messages).toEqual(messages); - expect(store.getState().isLoadingMessages).toBe(false); + expect(active(store).messages).toEqual(messages); + expect(active(store).isLoadingMessages).toBe(false); expect(getMessages).toHaveBeenCalledWith("t1"); }); + it("does not reload a thread that already has an in-memory entry", async () => { + const getMessages = vi.fn().mockResolvedValue([makeMessage("stored")]); + const store = makeStore({ getMessages }); + + // Thread already loaded/streamed this session. + seed(store, "t1", { messages: [makeMessage("in-memory")] }); + + store.getState().selectThread("t1"); + await flushPromises(); + + expect(active(store).messages.map((m) => m.id)).toEqual(["in-memory"]); + expect(getMessages).not.toHaveBeenCalled(); + }); + it("sets threadError on load failure", async () => { const error = new Error("load failed"); const getMessages = vi.fn().mockRejectedValue(error); @@ -135,26 +168,54 @@ describe("createChatStore", () => { store.getState().selectThread("t1"); await flushPromises(); - expect(store.getState().threadError).toBe(error); - expect(store.getState().isLoadingMessages).toBe(false); + expect(active(store).threadError).toBe(error); + expect(active(store).isLoadingMessages).toBe(false); }); }); describe("switchToNewThread", () => { - it("clears selection, messages, and errors", () => { + it("clears selection and shows an empty draft view", () => { const store = makeStore(); - store.setState({ - selectedThreadId: "t1", - messages: [makeMessage("m1")], - threadError: new Error("old"), - }); + seed(store, "t1", { messages: [makeMessage("m1")], threadError: new Error("old") }); + store.setState({ selectedThreadId: "t1" }); store.getState().switchToNewThread(); expect(store.getState().selectedThreadId).toBeNull(); - expect(store.getState().messages).toEqual([]); - expect(store.getState().threadError).toBeNull(); + expect(active(store).messages).toEqual([]); + expect(active(store).threadError).toBeNull(); + }); + + it("is a no-op while a new thread is being created", async () => { + let resolveCreate!: (t: Thread) => void; + const createThread = vi.fn().mockImplementation( + () => + new Promise((r) => { + resolveCreate = r; + }), + ); + const send = vi.fn().mockResolvedValue(new Response("", { status: 200 })); + const store = makeStore({ createThread, send, streamProtocol: { parse: async function* () {} } }); + + // Brand-new chat: first message kicks off createThread (still pending). + store.getState().processMessage({ role: "user", content: "hello" }); + await flushPromises(); + + expect(store.getState().isCreatingThread).toBe(true); + expect(entryOf(store, DRAFT_KEY)?.isRunning).toBe(true); + + // Attempt to start a fresh chat mid-creation → ignored, draft run preserved. + store.getState().switchToNewThread(); + expect(entryOf(store, DRAFT_KEY)?.isRunning).toBe(true); + + resolveCreate(makeThread("t-real")); + await flushPromises(); + + // Draft re-keyed to the real thread; DRAFT slot freed. + expect(entryOf(store, DRAFT_KEY)).toBeUndefined(); + expect(entryOf(store, "t-real")).toBeDefined(); + expect(store.getState().isCreatingThread).toBe(false); }); }); @@ -196,17 +257,41 @@ describe("createChatStore", () => { const deleteThread = vi.fn().mockResolvedValue(undefined); const store = makeStore({ deleteThread }); - store.setState({ - threads: [makeThread("t1")], - selectedThreadId: "t1", - messages: [makeMessage("m1")], - }); + seed(store, "t1", { messages: [makeMessage("m1")] }); + store.setState({ threads: [makeThread("t1")], selectedThreadId: "t1" }); store.getState().deleteThread("t1"); await flushPromises(); expect(store.getState().selectedThreadId).toBeNull(); - expect(store.getState().messages).toEqual([]); + expect(active(store).messages).toEqual([]); + expect(entryOf(store, "t1")).toBeUndefined(); + }); + + it("aborts an in-flight run on the deleted thread and drops its entry (no ghost)", async () => { + let capturedSignal: AbortSignal | undefined; + const send = vi.fn().mockImplementation(({ signal }) => { + capturedSignal = signal; + return new Promise(() => {}); // never resolves + }); + const deleteThread = vi.fn().mockResolvedValue(undefined); + const store = makeStore({ send, deleteThread, streamProtocol: { parse: async function* () {} } }); + + store.setState({ threads: [makeThread("t1")], selectedThreadId: "t1" }); + store.getState().processMessage({ role: "user", content: "hello" }); + await flushPromises(); + expect(entryOf(store, "t1")?.isRunning).toBe(true); + + store.getState().deleteThread("t1"); + await flushPromises(); + + expect(capturedSignal?.aborted).toBe(true); + expect(entryOf(store, "t1")).toBeUndefined(); + expect(store.getState().threads).toHaveLength(0); + + // Late `finally`/abort callbacks for the dropped key must not resurrect it. + await flushPromises(); + expect(entryOf(store, "t1")).toBeUndefined(); }); it("sets isPending during operation", async () => { @@ -248,39 +333,40 @@ describe("createChatStore", () => { }); // ──────────────────────────────────────────── - // Message CRUD + // Message CRUD (operate on the active view entry) // ──────────────────────────────────────────── describe("message CRUD", () => { - let store: ReturnType; + let store: Store; beforeEach(() => { store = makeStore(); - store.setState({ messages: [makeMessage("m1"), makeMessage("m2", "assistant")] }); + // selectedThreadId is null → the active view is the DRAFT entry. + seed(store, DRAFT_KEY, { messages: [makeMessage("m1"), makeMessage("m2", "assistant")] }); }); it("appendMessages adds to end", () => { store.getState().appendMessages(makeMessage("m3")); - expect(store.getState().messages).toHaveLength(3); - expect(store.getState().messages[2].id).toBe("m3"); + expect(active(store).messages).toHaveLength(3); + expect(active(store).messages[2].id).toBe("m3"); }); it("setMessages replaces all", () => { store.getState().setMessages([makeMessage("new")]); - expect(store.getState().messages).toHaveLength(1); - expect(store.getState().messages[0].id).toBe("new"); + expect(active(store).messages).toHaveLength(1); + expect(active(store).messages[0].id).toBe("new"); }); it("updateMessage replaces by id", () => { const updated = { ...makeMessage("m1"), content: "edited" } as Message; store.getState().updateMessage(updated); - expect((store.getState().messages[0] as any).content).toBe("edited"); + expect((active(store).messages[0] as any).content).toBe("edited"); }); it("deleteMessage removes by id", () => { store.getState().deleteMessage("m1"); - expect(store.getState().messages).toHaveLength(1); - expect(store.getState().messages[0].id).toBe("m2"); + expect(active(store).messages).toHaveLength(1); + expect(active(store).messages[0].id).toBe("m2"); }); }); @@ -292,46 +378,39 @@ describe("createChatStore", () => { it("appends optimistic user message and calls llm.send", async () => { const send = vi.fn().mockResolvedValue(new Response("", { status: 200 })); - const store = makeStore({ - send, - streamProtocol: { parse: async function* () {} }, - }); - + const store = makeStore({ send, streamProtocol: { parse: async function* () {} } }); store.setState({ selectedThreadId: "t1" }); await store.getState().processMessage({ role: "user", content: "hello" }); - expect(store.getState().messages).toHaveLength(1); - expect(store.getState().messages[0].role).toBe("user"); - expect(store.getState().isRunning).toBe(false); + expect(active(store).messages).toHaveLength(1); + expect(active(store).messages[0].role).toBe("user"); + expect(active(store).isRunning).toBe(false); expect(send).toHaveBeenCalledOnce(); }); - it("creates thread when none selected", async () => { + it("creates thread when none selected and follows into it", async () => { const newThread = makeThread("t-auto"); const createThread = vi.fn().mockResolvedValue(newThread); const send = vi.fn().mockResolvedValue(new Response("", { status: 200 })); - const store = makeStore({ - createThread, - send, - streamProtocol: { parse: async function* () {} }, - }); + const store = makeStore({ createThread, send, streamProtocol: { parse: async function* () {} } }); await store.getState().processMessage({ role: "user", content: "hello" }); expect(createThread).toHaveBeenCalledOnce(); expect(store.getState().selectedThreadId).toBe("t-auto"); + expect(entryOf(store, DRAFT_KEY)).toBeUndefined(); + expect(entryOf(store, "t-auto")?.messages).toHaveLength(1); + expect(store.getState().isCreatingThread).toBe(false); }); - it("no-ops when already running", async () => { + it("no-ops when the same thread is already running", async () => { const send = vi.fn().mockResolvedValue(new Response("", { status: 200 })); - const store = makeStore({ - send, - streamProtocol: { parse: async function* () {} }, - }); - store.setState({ isRunning: true, selectedThreadId: "t1" }); + const store = makeStore({ send, streamProtocol: { parse: async function* () {} } }); + seed(store, "t1", { isRunning: true }); + store.setState({ selectedThreadId: "t1" }); await store.getState().processMessage({ role: "user", content: "hello" }); @@ -341,17 +420,14 @@ describe("createChatStore", () => { it("sets threadError on failure", async () => { const send = vi.fn().mockRejectedValue(new Error("api down")); - const store = makeStore({ - send, - streamProtocol: { parse: async function* () {} }, - }); + const store = makeStore({ send, streamProtocol: { parse: async function* () {} } }); store.setState({ selectedThreadId: "t1" }); await store.getState().processMessage({ role: "user", content: "hello" }); - expect(store.getState().threadError).toBeInstanceOf(Error); - expect(store.getState().threadError?.message).toBe("api down"); - expect(store.getState().isRunning).toBe(false); + expect(active(store).threadError).toBeInstanceOf(Error); + expect(active(store).threadError?.message).toBe("api down"); + expect(active(store).isRunning).toBe(false); }); }); @@ -360,39 +436,59 @@ describe("createChatStore", () => { // ──────────────────────────────────────────── describe("cancelMessage", () => { - it("aborts in-flight request", async () => { + it("aborts the selected thread's in-flight request", async () => { let capturedSignal: AbortSignal; const send = vi.fn().mockImplementation(({ signal }) => { capturedSignal = signal; return new Promise(() => {}); // never resolves }); - const store = makeStore({ - send, - streamProtocol: { parse: async function* () {} }, - }); + const store = makeStore({ send, streamProtocol: { parse: async function* () {} } }); store.setState({ selectedThreadId: "t1" }); - const _promise = store.getState().processMessage({ role: "user", content: "hello" }); + store.getState().processMessage({ role: "user", content: "hello" }); await flushPromises(); - expect(store.getState().isRunning).toBe(true); + expect(active(store).isRunning).toBe(true); store.getState().cancelMessage(); await flushPromises(); - expect(store.getState().isRunning).toBe(false); + expect(active(store).isRunning).toBe(false); expect(capturedSignal!.aborted).toBe(true); }); + + it("does not abort a run on a non-selected (background) thread", async () => { + const signals: AbortSignal[] = []; + const send = vi.fn().mockImplementation(({ signal }) => { + signals.push(signal); + return new Promise(() => {}); // never resolves + }); + const store = makeStore({ send, streamProtocol: { parse: async function* () {} } }); + + // Run on t1, then switch to t2 (t1 keeps running in background). + store.setState({ selectedThreadId: "t1" }); + store.getState().processMessage({ role: "user", content: "one" }); + await flushPromises(); + store.getState().selectThread("t2"); + await flushPromises(); + + // Cancelling from t2 must not touch t1's run. + store.getState().cancelMessage(); + await flushPromises(); + + expect(signals[0].aborted).toBe(false); + expect(entryOf(store, "t1")?.isRunning).toBe(true); + }); }); // ──────────────────────────────────────────── - // Thread switch during stream + // Background streaming across thread switches // ──────────────────────────────────────────── - describe("selectThread while streaming", () => { - it("cancels current stream and loads new thread", async () => { - let capturedSignal: AbortSignal; + describe("thread switch while streaming", () => { + it("keeps the run alive in the background and loads the new thread", async () => { + let capturedSignal: AbortSignal | undefined; const send = vi.fn().mockImplementation(({ signal }) => { capturedSignal = signal; return new Promise(() => {}); // never resolves @@ -400,29 +496,123 @@ describe("createChatStore", () => { const newMessages = [makeMessage("new-m1")]; const getMessages = vi.fn().mockResolvedValue(newMessages); + const store = makeStore({ send, getMessages, streamProtocol: { parse: async function* () {} } }); + store.setState({ selectedThreadId: "t1" }); + + // Start streaming on t1. + store.getState().processMessage({ role: "user", content: "hello" }); + await flushPromises(); + expect(active(store).isRunning).toBe(true); + + // Switch to t2 mid-stream — the run must NOT be aborted (behaviour change). + store.getState().selectThread("t2"); + + expect(capturedSignal?.aborted).toBe(false); + expect(store.getState().selectedThreadId).toBe("t2"); + expect(active(store).isLoadingMessages).toBe(true); + expect(entryOf(store, "t1")?.isRunning).toBe(true); // t1 still streaming + + await flushPromises(); + + expect(active(store).messages).toEqual(newMessages); + expect(active(store).isLoadingMessages).toBe(false); + + // Switch back to t1 → shows its in-session messages, NO storage reload. + store.getState().selectThread("t1"); + expect(active(store).isRunning).toBe(true); + expect(active(store).messages).toHaveLength(1); // the optimistic user message + expect(getMessages).toHaveBeenCalledTimes(1); // only t2 was loaded + expect(getMessages).toHaveBeenCalledWith("t2"); + expect(capturedSignal?.aborted).toBe(false); + }); + + it("runs two threads concurrently, each writing only its own entry", async () => { + const signals: AbortSignal[] = []; + const send = vi.fn().mockImplementation(({ signal }) => { + signals.push(signal); + return new Promise(() => {}); // never resolves + }); + const store = makeStore({ send, streamProtocol: { parse: async function* () {} } }); + + store.setState({ selectedThreadId: "t1" }); + store.getState().processMessage({ role: "user", content: "one" }); + await flushPromises(); + + store.getState().selectThread("t2"); + await flushPromises(); + store.getState().processMessage({ role: "user", content: "two" }); + await flushPromises(); + + expect(entryOf(store, "t1")?.isRunning).toBe(true); + expect(entryOf(store, "t2")?.isRunning).toBe(true); + expect(send).toHaveBeenCalledTimes(2); + // Distinct runs, distinct controllers. + expect(signals[0]).not.toBe(signals[1]); + expect(entryOf(store, "t1")?.messages).toHaveLength(1); + expect(entryOf(store, "t2")?.messages).toHaveLength(1); + }); + + it("a completed background run's messages survive a switch away and back", async () => { + const send = vi.fn().mockResolvedValue(new Response("", { status: 200 })); + // One text chunk → the assistant message is created synchronously (isFirst path). + const streamProtocol = { + parse: async function* () { + yield { type: EventType.TEXT_MESSAGE_CONTENT, delta: "hi there" } as never; + }, + }; + const getMessages = vi.fn().mockResolvedValue([]); + const store = makeStore({ send, streamProtocol, getMessages }); + + store.setState({ selectedThreadId: "t1" }); + await store.getState().processMessage({ role: "user", content: "hello" }); + + // optimistic user + streamed assistant message + expect(entryOf(store, "t1")?.messages).toHaveLength(2); + expect(entryOf(store, "t1")?.isRunning).toBe(false); + + store.getState().selectThread("t2"); + await flushPromises(); + store.getState().selectThread("t1"); + + expect(active(store).messages).toHaveLength(2); + expect(getMessages).not.toHaveBeenCalledWith("t1"); // never reloaded from storage + }); + + it("re-keys a draft run to the background when the user navigates away mid-creation", async () => { + let resolveCreate!: (t: Thread) => void; + const createThread = vi.fn().mockImplementation( + () => + new Promise((r) => { + resolveCreate = r; + }), + ); + const send = vi.fn().mockResolvedValue(new Response("", { status: 200 })); + const getMessages = vi.fn().mockResolvedValue([]); const store = makeStore({ + createThread, send, getMessages, streamProtocol: { parse: async function* () {} }, }); - store.setState({ selectedThreadId: "t1" }); - // Start streaming + // New chat, first message → createThread pending, run under DRAFT. store.getState().processMessage({ role: "user", content: "hello" }); await flushPromises(); - expect(store.getState().isRunning).toBe(true); - - // Switch thread mid-stream - store.getState().selectThread("t2"); + expect(store.getState().isCreatingThread).toBe(true); - expect(capturedSignal!.aborted).toBe(true); - expect(store.getState().selectedThreadId).toBe("t2"); - expect(store.getState().isLoadingMessages).toBe(true); + // User selects a saved thread before createThread resolves. + store.getState().selectThread("saved"); + await flushPromises(); + resolveCreate(makeThread("t-real")); await flushPromises(); - expect(store.getState().messages).toEqual(newMessages); - expect(store.getState().isLoadingMessages).toBe(false); + // Did NOT follow — still viewing "saved"; the run backgrounds under its real id. + expect(store.getState().selectedThreadId).toBe("saved"); + expect(entryOf(store, DRAFT_KEY)).toBeUndefined(); + expect(entryOf(store, "t-real")).toBeDefined(); + expect(entryOf(store, "t-real")?.messages).toHaveLength(1); + expect(store.getState().isCreatingThread).toBe(false); }); }); }); diff --git a/packages/react-headless/src/store/createChatStore.ts b/packages/react-headless/src/store/createChatStore.ts index ef26cfa0c..e2c0d9d41 100644 --- a/packages/react-headless/src/store/createChatStore.ts +++ b/packages/react-headless/src/store/createChatStore.ts @@ -2,7 +2,8 @@ import { createStore } from "zustand"; import { subscribeWithSelector } from "zustand/middleware"; import type { ChatLLM, ChatStorage } from "../adapters/types"; import { processStreamedMessage } from "../stream/processStreamedMessage"; -import type { ChatStore, Message, Thread, UserMessage } from "./types"; +import { DRAFT_KEY, makeThreadState, resolveViewKey } from "./threadStateEntry"; +import type { ChatStore, Message, Thread, ThreadStateEntry, UserMessage } from "./types"; export interface CreateChatStoreConfig { storage: ChatStorage; @@ -19,227 +20,315 @@ export const createChatStore = (config: CreateChatStoreConfig) => { const { thread: threadStorage } = storage; const store = createStore()( - subscribeWithSelector((set, get) => ({ - // Thread List State - threads: [], - isLoadingThreads: false, - threadListError: null, - selectedThreadId: null, - hasMoreThreads: false, - _nextCursor: undefined, - - // Thread State - messages: [], - isRunning: false, - isLoadingMessages: false, - threadError: null, - executingToolCallIds: new Set(), - _abortController: null, - - // ── Thread List Actions ── - - loadThreads: () => { - set({ isLoadingThreads: true, threadListError: null }); - threadStorage - .listThreads(undefined) - .then(({ threads = [], nextCursor }) => { - set({ - threads, - isLoadingThreads: false, - _nextCursor: nextCursor, - hasMoreThreads: nextCursor !== undefined, - }); - }) - .catch((e) => { - set({ isLoadingThreads: false, threadListError: e }); - }); - }, - - loadMoreThreads: () => { - const cursor = get()._nextCursor; - if (cursor === undefined) return; - threadStorage - .listThreads(cursor) - .then(({ threads = [], nextCursor }) => { - set((s) => ({ - threads: mergeThreadList(s.threads, threads), - _nextCursor: nextCursor, - hasMoreThreads: nextCursor !== undefined, - })); - }) - .catch((e) => { - set({ threadListError: e }); - }); - }, - - switchToNewThread: () => { - get().cancelMessage(); - set({ - selectedThreadId: null, - messages: [], - threadError: null, - executingToolCallIds: new Set(), - }); - }, - - createThread: async (firstMessage: UserMessage) => { - const thread = await threadStorage.createThread(firstMessage); - set((s) => ({ threads: mergeThreadList(s.threads, [thread]) })); - return thread; - }, - - selectThread: (threadId: string) => { - get().cancelMessage(); - set({ - selectedThreadId: threadId, - messages: [], - isLoadingMessages: true, - threadError: null, - executingToolCallIds: new Set(), + subscribeWithSelector((set, get) => { + /** The key whose entry the view currently shows. */ + const viewKey = () => resolveViewKey(get().selectedThreadId); + + /** + * Guarded functional update of one thread's entry. No-ops when the entry + * is absent (thread deleted) or `fn` returns `null`, so a dropped thread's + * late `finally`/abort/straggler callbacks can't resurrect a ghost. + * Spreads preserve every other entry's identity. + */ + const withThreadState = ( + key: string, + fn: (cur: ThreadStateEntry) => Partial | null, + ) => + set((s) => { + const cur = s.threadStates[key]; + if (!cur) return s; + const patch = fn(cur); + if (!patch) return s; + return { threadStates: { ...s.threadStates, [key]: { ...cur, ...patch } } }; }); - threadStorage - .getMessages(threadId) - .then((messages) => set({ messages, isLoadingMessages: false })) - .catch((e) => set({ threadError: e, isLoadingMessages: false })); - }, - - updateThread: (thread: Thread) => { - const setPending = (id: string, isPending: boolean) => - set((s) => ({ threads: s.threads.map((t) => (t.id === id ? { ...t, isPending } : t)) })); - setPending(thread.id, true); - threadStorage - .updateThread(thread) - .then((updated) => { - set((s) => ({ - threads: s.threads.map((t) => (t.id === updated.id ? updated : t)), - })); - }) - .catch(() => setPending(thread.id, false)); - }, - - deleteThread: (threadId: string) => { - const setPending = (id: string, isPending: boolean) => - set((s) => ({ threads: s.threads.map((t) => (t.id === id ? { ...t, isPending } : t)) })); - setPending(threadId, true); - threadStorage - .deleteThread(threadId) - .then(() => { - const state = get(); - set({ threads: state.threads.filter((t) => t.id !== threadId) }); - if (state.selectedThreadId === threadId) { - state.switchToNewThread(); - } - }) - .catch(() => setPending(threadId, false)); - }, - - // ── Thread Actions ── - - processMessage: async (message) => { - const state = get(); - if (state.isRunning) return; - - const abortController = new AbortController(); - const optimisticMessage: UserMessage = { - ...message, - id: crypto.randomUUID(), - role: "user", - }; - - set({ - _abortController: abortController, - isRunning: true, - threadError: null, - executingToolCallIds: new Set(), + + /** Guarded merge of a static patch into one thread's entry. */ + const patchThreadState = (key: string, patch: Partial) => + withThreadState(key, () => patch); + + /** Update one thread's entry, creating it if absent. */ + const upsertThreadState = ( + key: string, + fn: (cur: ThreadStateEntry) => Partial, + ) => + set((s) => { + const cur = s.threadStates[key] ?? makeThreadState(); + return { threadStates: { ...s.threadStates, [key]: { ...cur, ...fn(cur) } } }; }); - set((s) => ({ messages: [...s.messages, optimisticMessage] })); - abortController.signal.addEventListener("abort", () => { - set({ _abortController: null, isRunning: false }); + /** Drop a thread's entry entirely. */ + const dropThreadState = (key: string) => + set((s) => { + if (!(key in s.threadStates)) return s; + const next = { ...s.threadStates }; + delete next[key]; + return { threadStates: next }; }); - try { - let threadId = get().selectedThreadId; + return { + // ── Thread List State ── + threads: [], + isLoadingThreads: false, + threadListError: null, + selectedThreadId: null, + hasMoreThreads: false, + isCreatingThread: false, + _nextCursor: undefined, - if (!threadId) { - const created = await get().createThread(optimisticMessage); - threadId = created.id; - set({ selectedThreadId: threadId }); - } + // ── Per-thread State ── + threadStates: {}, - const response = await llm.send({ - threadId, - messages: get().messages, - signal: abortController.signal, - }); + // ── Thread List Actions ── - if (response instanceof Response && !response.ok) { - throw new Error(`Request failed: ${response.status} ${response.statusText}`); - } + loadThreads: () => { + set({ isLoadingThreads: true, threadListError: null }); + threadStorage + .listThreads(undefined) + .then(({ threads = [], nextCursor }) => { + set({ + threads, + isLoadingThreads: false, + _nextCursor: nextCursor, + hasMoreThreads: nextCursor !== undefined, + }); + }) + .catch((e) => { + set({ isLoadingThreads: false, threadListError: e }); + }); + }, - await processStreamedMessage({ - response, - createMessage: (msg) => set((s) => ({ messages: [...s.messages, msg] })), - updateMessage: (msg) => + loadMoreThreads: () => { + const cursor = get()._nextCursor; + if (cursor === undefined) return; + threadStorage + .listThreads(cursor) + .then(({ threads = [], nextCursor }) => { set((s) => ({ - messages: s.messages.map((m) => (m.id === msg.id ? msg : m)), - })), - // A tool's args have closed (TOOL_CALL_END) → it is now executing. - markToolExecuting: (id) => - set((s) => - s.executingToolCallIds.has(id) - ? s - : { executingToolCallIds: new Set(s.executingToolCallIds).add(id) }, + threads: mergeThreadList(s.threads, threads), + _nextCursor: nextCursor, + hasMoreThreads: nextCursor !== undefined, + })); + }) + .catch((e) => { + set({ threadListError: e }); + }); + }, + + switchToNewThread: () => { + // A brand-new chat whose first message is still becoming a real thread + // owns DRAFT_KEY; resetting it now would collide with that run's re-key. + if (get().isCreatingThread) return; + // No abort — a run on a saved thread keeps streaming in the background. + set({ selectedThreadId: null }); + dropThreadState(DRAFT_KEY); + }, + + createThread: async (firstMessage: UserMessage) => { + const thread = await threadStorage.createThread(firstMessage); + set((s) => ({ threads: mergeThreadList(s.threads, [thread]) })); + return thread; + }, + + selectThread: (threadId: string) => { + // No abort — the previously-selected thread's run (if any) keeps + // streaming into its own entry in the background. + set({ selectedThreadId: threadId }); + + // An in-memory entry (live/finished run, or an already-loaded thread) + // is shown as-is. Never reload: storage may not yet have messages + // streamed this session. + if (get().threadStates[threadId]) return; + + upsertThreadState(threadId, () => ({ isLoadingMessages: true })); + threadStorage + .getMessages(threadId) + .then((messages) => + // A run started (or the thread was deleted) meanwhile → don't clobber. + withThreadState(threadId, (cur) => + cur.isRunning ? null : { messages, isLoadingMessages: false }, ), - // Its result landed (or it errored) → no longer executing. - clearToolExecuting: (id) => - set((s) => { - if (!s.executingToolCallIds.has(id)) return s; - const next = new Set(s.executingToolCallIds); - next.delete(id); - return { executingToolCallIds: next }; - }), - adapter: llm.streamProtocol, - }); - } catch (e) { - if (!abortController.signal.aborted) { - set({ threadError: e instanceof Error ? e : new Error(String(e)) }); - } - } finally { - // Clear any tool calls still flagged "executing" — adapters that emit - // TOOL_CALL_END without a matching TOOL_CALL_RESULT (e.g. client-side - // tool calls in the OpenAI adapters) would otherwise leave them stuck - // in the executing set after the run ends. - set({ - _abortController: null, - isRunning: false, + ) + .catch((e) => patchThreadState(threadId, { threadError: e, isLoadingMessages: false })); + }, + + updateThread: (thread: Thread) => { + const setPending = (id: string, isPending: boolean) => + set((s) => ({ + threads: s.threads.map((t) => (t.id === id ? { ...t, isPending } : t)), + })); + setPending(thread.id, true); + threadStorage + .updateThread(thread) + .then((updated) => { + set((s) => ({ + threads: s.threads.map((t) => (t.id === updated.id ? updated : t)), + })); + }) + .catch(() => setPending(thread.id, false)); + }, + + deleteThread: (threadId: string) => { + const setPending = (id: string, isPending: boolean) => + set((s) => ({ + threads: s.threads.map((t) => (t.id === id ? { ...t, isPending } : t)), + })); + setPending(threadId, true); + threadStorage + .deleteThread(threadId) + .then(() => { + const state = get(); + // Stop any in-flight run on this thread, then drop its entry. Late + // stream callbacks for this key then no-op (withThreadState is guarded). + state.threadStates[threadId]?.abortController?.abort(); + dropThreadState(threadId); + set((s) => ({ threads: s.threads.filter((t) => t.id !== threadId) })); + if (state.selectedThreadId === threadId) { + get().switchToNewThread(); + } + }) + .catch(() => setPending(threadId, false)); + }, + + // ── Thread Actions ── + + processMessage: async (message) => { + const startState = get(); + // The run's key: the selected thread, or DRAFT_KEY for a brand-new chat. + // A mutable local so post-re-key callbacks target the real threadId. + let runKey = startState.selectedThreadId ?? DRAFT_KEY; + + // Per-thread concurrency guard: only block if THIS thread is already + // running. Different threads run concurrently. + if (startState.threadStates[runKey]?.isRunning) return; + + const isNewChat = !startState.selectedThreadId; + const abortController = new AbortController(); + const optimisticMessage: UserMessage = { + ...message, + id: crypto.randomUUID(), + role: "user", + }; + + // Start the run on this thread's entry, preserving any already-loaded + // messages and appending the optimistic user message. + if (isNewChat) set({ isCreatingThread: true }); + upsertThreadState(runKey, (cur) => ({ + messages: [...cur.messages, optimisticMessage], + isRunning: true, + threadError: null, executingToolCallIds: new Set(), + abortController, + })); + + // On abort, flip the run off on its own entry. + abortController.signal.addEventListener("abort", () => { + patchThreadState(runKey, { isRunning: false, abortController: null }); }); - } - }, - - appendMessages: (...newMessages: Message[]) => { - set((s) => ({ messages: [...s.messages, ...newMessages] })); - }, - - updateMessage: (message: Message) => { - set((s) => ({ - messages: s.messages.map((m) => (m.id === message.id ? message : m)), - })); - }, - - setMessages: (messages: Message[]) => { - set({ messages }); - }, - - deleteMessage: (messageId: string) => { - set((s) => ({ messages: s.messages.filter((m) => m.id !== messageId) })); - }, - - cancelMessage: () => { - get()._abortController?.abort(); - }, - })), + + try { + if (isNewChat) { + try { + const created = await get().createThread(optimisticMessage); + // Re-key the draft entry to the real threadId, carrying the live + // run. Follow into the new thread ONLY if the user is still on the + // draft view — a background re-key must not write selectedThreadId, + // or ChatProvider would reset the ephemeral stores of the thread + // the user is now viewing. + set((s) => { + const draft = s.threadStates[DRAFT_KEY]; + const nextStates = { ...s.threadStates }; + if (draft) { + delete nextStates[DRAFT_KEY]; + nextStates[created.id] = draft; + } + const follow = s.selectedThreadId === null; + return { + threadStates: nextStates, + ...(follow ? { selectedThreadId: created.id } : null), + }; + }); + runKey = created.id; + } finally { + set({ isCreatingThread: false }); + } + } + + const response = await llm.send({ + threadId: runKey, + messages: get().threadStates[runKey]?.messages ?? [], + signal: abortController.signal, + }); + + if (response instanceof Response && !response.ok) { + throw new Error(`Request failed: ${response.status} ${response.statusText}`); + } + + await processStreamedMessage({ + response, + createMessage: (msg) => + withThreadState(runKey, (cur) => ({ messages: [...cur.messages, msg] })), + updateMessage: (msg) => + withThreadState(runKey, (cur) => ({ + messages: cur.messages.map((m) => (m.id === msg.id ? msg : m)), + })), + // A tool's args have closed (TOOL_CALL_END) → it is now executing. + // The `null` no-op keeps the Set reference stable when membership is + // unchanged so `useToolActivities` doesn't re-run needlessly. + markToolExecuting: (id) => + withThreadState(runKey, (cur) => + cur.executingToolCallIds.has(id) + ? null + : { executingToolCallIds: new Set(cur.executingToolCallIds).add(id) }, + ), + // Its result landed (or it errored) → no longer executing. + clearToolExecuting: (id) => + withThreadState(runKey, (cur) => { + if (!cur.executingToolCallIds.has(id)) return null; + const next = new Set(cur.executingToolCallIds); + next.delete(id); + return { executingToolCallIds: next }; + }), + adapter: llm.streamProtocol, + }); + } catch (e) { + if (!abortController.signal.aborted) { + patchThreadState(runKey, { + threadError: e instanceof Error ? e : new Error(String(e)), + }); + } + } finally { + // Clear run flags + any tool calls still flagged "executing" — adapters + // that emit TOOL_CALL_END without a matching TOOL_CALL_RESULT (e.g. + // client-side tool calls in the OpenAI adapters) would otherwise leave + // them stuck. + patchThreadState(runKey, { + isRunning: false, + abortController: null, + executingToolCallIds: new Set(), + }); + } + }, + + appendMessages: (...newMessages: Message[]) => + upsertThreadState(viewKey(), (cur) => ({ messages: [...cur.messages, ...newMessages] })), + + updateMessage: (message: Message) => + withThreadState(viewKey(), (cur) => ({ + messages: cur.messages.map((m) => (m.id === message.id ? message : m)), + })), + + setMessages: (messages: Message[]) => upsertThreadState(viewKey(), () => ({ messages })), + + deleteMessage: (messageId: string) => + withThreadState(viewKey(), (cur) => ({ + messages: cur.messages.filter((m) => m.id !== messageId), + })), + + cancelMessage: () => { + get().threadStates[viewKey()]?.abortController?.abort(); + }, + }; + }), ); return store; diff --git a/packages/react-headless/src/store/threadStateEntry.ts b/packages/react-headless/src/store/threadStateEntry.ts new file mode 100644 index 000000000..cb3faf5f6 --- /dev/null +++ b/packages/react-headless/src/store/threadStateEntry.ts @@ -0,0 +1,48 @@ +import type { Message } from "../types/message"; +import type { ThreadState, ThreadStateEntry } from "./types"; + +/** Key for the single unsaved "new chat" slot in {@link ChatStore.threadStates}. */ +export const DRAFT_KEY = "__draft__"; + +/** The map key whose entry the view currently shows. */ +export const resolveViewKey = (selectedThreadId: string | null): string => + selectedThreadId ?? DRAFT_KEY; + +/** + * The public {@link ThreadState} view of an entry — drops the internal + * `abortController`. Field references are preserved, so `useShallow` consumers + * don't re-render unless a field actually changes. Shared by `useThread` (the + * selected thread) and `useThreadState` (any thread by id). + */ +export const deriveThreadState = (entry: ThreadStateEntry): ThreadState => ({ + messages: entry.messages, + isRunning: entry.isRunning, + isLoadingMessages: entry.isLoadingMessages, + threadError: entry.threadError, + executingToolCallIds: entry.executingToolCallIds, +}); + +// Frozen singletons so an absent/empty thread yields *stable* references +const EMPTY_MESSAGES = Object.freeze([]) as unknown as Message[]; +const EMPTY_EXECUTING = Object.freeze(new Set()) as Set; + +/** The state shown for a thread that has no entry yet (blank/never-loaded). */ +export const EMPTY_THREAD_STATE: ThreadStateEntry = Object.freeze({ + messages: EMPTY_MESSAGES, + isRunning: false, + isLoadingMessages: false, + threadError: null, + executingToolCallIds: EMPTY_EXECUTING, + abortController: null, +}); + +/** Build a fresh, independently-mutable entry, applying an optional patch. */ +export const makeThreadState = (patch?: Partial): ThreadStateEntry => ({ + messages: [], + isRunning: false, + isLoadingMessages: false, + threadError: null, + executingToolCallIds: new Set(), + abortController: null, + ...patch, +}); diff --git a/packages/react-headless/src/store/types.ts b/packages/react-headless/src/store/types.ts index 4ddc3c870..46fa3b9be 100644 --- a/packages/react-headless/src/store/types.ts +++ b/packages/react-headless/src/store/types.ts @@ -20,6 +20,7 @@ export type ThreadListState = { threadListError: Error | null; selectedThreadId: string | null; hasMoreThreads: boolean; + isCreatingThread: boolean; }; export type ThreadListActions = { @@ -50,6 +51,11 @@ export type ThreadState = { executingToolCallIds: Set; }; +export type ThreadStateEntry = ThreadState & { + /** @internal — the in-flight run's controller, or `null` when idle. */ + abortController: AbortController | null; +}; + export type ThreadActions = { processMessage: (message: CreateMessage) => Promise; appendMessages: (...messages: Message[]) => void; @@ -63,12 +69,10 @@ export type ThreadActions = { export type ChatStore = ThreadListState & ThreadListActions & - ThreadState & ThreadActions & { + threadStates: Record; /** @internal */ _nextCursor?: string | undefined; - /** @internal */ - _abortController: AbortController | null; }; // ── Provider props ── diff --git a/packages/react-ui/src/components/AgentInterface/MobileHeader.tsx b/packages/react-ui/src/components/AgentInterface/MobileHeader.tsx index c47b05b6f..53d2c2c54 100644 --- a/packages/react-ui/src/components/AgentInterface/MobileHeader.tsx +++ b/packages/react-ui/src/components/AgentInterface/MobileHeader.tsx @@ -25,6 +25,7 @@ export const MobileHeader = ({ children, }: MobileHeaderProps) => { const switchToNewThread = useThreadList((s) => s.switchToNewThread); + const isCreatingThread = useThreadList((s) => s.isCreatingThread); const { agentName: ctxAgentName, setIsSidebarOpen } = useAgentInterfaceStore((state) => ({ agentName: state.agentName, setIsSidebarOpen: state.setIsSidebarOpen, @@ -66,6 +67,7 @@ export const MobileHeader = ({ size="medium" icon={} onClick={switchToNewThread} + disabled={isCreatingThread} variant="secondary" aria-label="New chat" /> diff --git a/packages/react-ui/src/components/AgentInterface/NewChatButton.tsx b/packages/react-ui/src/components/AgentInterface/NewChatButton.tsx index 26ae5db95..77453ef33 100644 --- a/packages/react-ui/src/components/AgentInterface/NewChatButton.tsx +++ b/packages/react-ui/src/components/AgentInterface/NewChatButton.tsx @@ -11,6 +11,7 @@ import { useAgentInterfaceStore } from "./_shared/store"; export const NewChatButton = ({ className }: { className?: string }) => { const switchToNewThread = useThreadList((s) => s.switchToNewThread); + const isCreatingThread = useThreadList((s) => s.isCreatingThread); const { isSidebarOpen } = useAgentInterfaceStore((state) => ({ isSidebarOpen: state.isSidebarOpen, })); @@ -38,6 +39,7 @@ export const NewChatButton = ({ className }: { className?: string }) => { iconLeft={} className={clsx("openui-agent-new-chat-floating-button", className)} onClick={handleNewChat} + disabled={isCreatingThread} aria-label="New chat" > New Chat @@ -55,6 +57,7 @@ export const NewChatButton = ({ className }: { className?: string }) => { className, )} onClick={handleNewChat} + disabled={isCreatingThread} aria-label="New chat" >