diff --git a/packages/react-headless/src/hooks/useIsThreadStreaming.ts b/packages/react-headless/src/hooks/useIsThreadStreaming.ts new file mode 100644 index 000000000..90a0765ad --- /dev/null +++ b/packages/react-headless/src/hooks/useIsThreadStreaming.ts @@ -0,0 +1,16 @@ +import { useStore } from "zustand"; +import { useChatStore } from "../store/ChatContext"; + +/** + * Whether `threadId` is currently streaming in the **background** — i.e. it has a + * run in flight while NOT being the active thread. Background streams are exactly + * the keys of {@link ChatStore.inFlightThreads} (the active thread lives in the + * flat store fields and is never here, and an entry is dropped the moment its run + * completes), so this is the direct signal for the sidebar loader. + * + * @category Hooks + */ +export function useIsThreadStreaming(threadId: string | null): boolean { + const store = useChatStore(); + return useStore(store, (s) => threadId != null && threadId in s.inFlightThreads); +} diff --git a/packages/react-headless/src/hooks/useThread.ts b/packages/react-headless/src/hooks/useThread.ts index 55c601250..876f23ce6 100644 --- a/packages/react-headless/src/hooks/useThread.ts +++ b/packages/react-headless/src/hooks/useThread.ts @@ -12,6 +12,7 @@ import type { type ThreadSlice = ThreadState & ThreadActions; type ThreadListSlice = ThreadListState & ThreadListActions; +// The active thread lives directly in the store's flat fields. const threadSelector = (s: ChatStore): ThreadSlice => ({ messages: s.messages, isRunning: s.isRunning, @@ -32,6 +33,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/index.ts b/packages/react-headless/src/index.ts index 6bd05b30a..1bcb87d64 100644 --- a/packages/react-headless/src/index.ts +++ b/packages/react-headless/src/index.ts @@ -4,6 +4,7 @@ export type { ArtifactListFilter } from "./hooks/useArtifactList"; export { useArtifactRenderer } from "./hooks/useArtifactRenderer"; export { useDetailedView } from "./hooks/useDetailedView"; export { useDetailedViewPortalTarget } from "./hooks/useDetailedViewPortalTarget"; +export { useIsThreadStreaming } from "./hooks/useIsThreadStreaming"; export { MessageContext, MessageProvider, useMessage } from "./hooks/useMessage"; export { useThread, useThreadList } from "./hooks/useThread"; export { useToolActivities } from "./hooks/useToolActivities"; @@ -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..cd47085fb 100644 --- a/packages/react-headless/src/store/__tests__/createChatStore.test.ts +++ b/packages/react-headless/src/store/__tests__/createChatStore.test.ts @@ -1,5 +1,5 @@ -import { beforeEach, describe, expect, it, vi } from "vitest"; -import type { Message, Thread, UserMessage } from "../types"; +import { describe, expect, it, vi } from "vitest"; +import type { Message, Thread, ThreadStateEntry } from "../types"; import { makeStore } from "./__helpers/makeStore"; // ── Helpers ── @@ -15,6 +15,24 @@ const makeMessage = (id: string, role: "user" | "assistant" = "user"): Message = const flushPromises = () => new Promise((r) => setTimeout(r, 0)); +type Store = ReturnType; + +/** The active thread's flat state — what `useThread` shows. */ +const active = (store: Store) => { + const s = store.getState(); + return { + messages: s.messages, + isRunning: s.isRunning, + isLoadingMessages: s.isLoadingMessages, + threadError: s.threadError, + executingToolCallIds: s.executingToolCallIds, + }; +}; + +/** A thread's BACKGROUND entry (undefined unless it's streaming while not active). */ +const bgOf = (store: Store, id: string): ThreadStateEntry | undefined => + store.getState().inFlightThreads[id]; + // ── Test suite ── describe("createChatStore", () => { @@ -26,292 +44,218 @@ describe("createChatStore", () => { it("fetches threads and sets state", async () => { const threads = [makeThread("t1"), makeThread("t2", 1)]; const listThreads = vi.fn().mockResolvedValue({ threads }); - const store = makeStore({ listThreads }); - expect(store.getState().isLoadingThreads).toBe(false); store.getState().loadThreads(); expect(store.getState().isLoadingThreads).toBe(true); - await flushPromises(); expect(store.getState().isLoadingThreads).toBe(false); expect(store.getState().threads).toHaveLength(2); - expect(store.getState().hasMoreThreads).toBe(false); - expect(listThreads).toHaveBeenCalledWith(undefined); }); it("sets threadListError on failure", async () => { const error = new Error("network"); - const listThreads = vi.fn().mockRejectedValue(error); - - const store = makeStore({ listThreads }); + const store = makeStore({ listThreads: vi.fn().mockRejectedValue(error) }); store.getState().loadThreads(); await flushPromises(); - - expect(store.getState().isLoadingThreads).toBe(false); expect(store.getState().threadListError).toBe(error); }); - - it("handles pagination cursor", async () => { - const listThreads = vi.fn().mockResolvedValue({ - threads: [makeThread("t1")], - nextCursor: "page2", - }); - - const store = makeStore({ listThreads }); - store.getState().loadThreads(); - await flushPromises(); - - expect(store.getState().hasMoreThreads).toBe(true); - expect(listThreads).toHaveBeenCalledWith(undefined); - }); }); describe("loadMoreThreads", () => { it("appends threads using cursor", async () => { - const page1 = [makeThread("t1")]; - const page2 = [makeThread("t2", 1)]; const listThreads = vi .fn() - .mockResolvedValueOnce({ threads: page1, nextCursor: "c2" }) - .mockResolvedValueOnce({ threads: page2 }); - + .mockResolvedValueOnce({ threads: [makeThread("t1")], nextCursor: "c2" }) + .mockResolvedValueOnce({ threads: [makeThread("t2", 1)] }); const store = makeStore({ listThreads }); store.getState().loadThreads(); await flushPromises(); - expect(store.getState().threads).toHaveLength(1); - store.getState().loadMoreThreads(); await flushPromises(); expect(store.getState().threads).toHaveLength(2); - expect(store.getState().hasMoreThreads).toBe(false); expect(listThreads).toHaveBeenCalledWith("c2"); }); - - it("no-ops when no more pages", async () => { - const listThreads = vi.fn().mockResolvedValue({ threads: [makeThread("t1")] }); - - const store = makeStore({ listThreads }); - store.getState().loadThreads(); - await flushPromises(); - - store.getState().loadMoreThreads(); - await flushPromises(); - - expect(listThreads).toHaveBeenCalledTimes(1); - }); }); + // ──────────────────────────────────────────── + // selectThread — loads into the active (flat) fields + // ──────────────────────────────────────────── + describe("selectThread", () => { - it("sets selectedThreadId, loads messages, clears previous", async () => { - const messages: Message[] = [makeMessage("m1"), makeMessage("m2", "assistant")]; + it("loads the thread's messages into the active fields", async () => { + const messages = [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).isLoadingMessages).toBe(true); + await flushPromises(); + expect(active(store).messages).toEqual(messages); + expect(active(store).isLoadingMessages).toBe(false); + expect(getMessages).toHaveBeenCalledWith("t1"); + }); + + it("discards an idle thread on switch (no background entry) and reloads on return", async () => { + const getMessages = vi.fn().mockResolvedValue([makeMessage("stored")]); + const store = makeStore({ getMessages }); + + store.getState().selectThread("t1"); + await flushPromises(); + store.getState().selectThread("t2"); // leaving idle t1 await flushPromises(); - expect(store.getState().messages).toEqual(messages); - expect(store.getState().isLoadingMessages).toBe(false); + expect(bgOf(store, "t1")).toBeUndefined(); // not kept in the background + // Returning to t1 reloads it from storage (fresh, no stale in-memory copy). + store.getState().selectThread("t1"); + await flushPromises(); expect(getMessages).toHaveBeenCalledWith("t1"); + expect(getMessages).toHaveBeenCalledTimes(3); // t1, t2, t1 again }); it("sets threadError on load failure", async () => { const error = new Error("load failed"); - const getMessages = vi.fn().mockRejectedValue(error); - - const store = makeStore({ getMessages }); + const store = makeStore({ getMessages: vi.fn().mockRejectedValue(error) }); 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 the active fields to a blank chat", () => { const store = makeStore(); - - store.setState({ - selectedThreadId: "t1", - messages: [makeMessage("m1")], - threadError: new Error("old"), - }); + store.setState({ selectedThreadId: "t1", messages: [makeMessage("m1")] }); store.getState().switchToNewThread(); expect(store.getState().selectedThreadId).toBeNull(); - expect(store.getState().messages).toEqual([]); - expect(store.getState().threadError).toBeNull(); + expect(active(store).messages).toEqual([]); }); - }); - describe("createThread", () => { - it("adds thread to list", async () => { - const newThread = makeThread("t-new"); - const createThread = vi.fn().mockResolvedValue(newThread); + 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().mockImplementation(() => new Promise(() => {})); + const store = makeStore({ + createThread, + send, + streamProtocol: { parse: async function* () {} }, + }); - const store = makeStore({ createThread }); - store.setState({ threads: [makeThread("t-existing")] }); + store.getState().processMessage({ role: "user", content: "hi" }); + await flushPromises(); + expect(store.getState().isCreatingThread).toBe(true); - const result = await store.getState().createThread({ - id: "m1", - role: "user", - content: "hello", - } as UserMessage); + store.getState().switchToNewThread(); // ignored mid-creation + expect(active(store).isRunning).toBe(true); - expect(result).toEqual(newThread); - expect(store.getState().threads).toHaveLength(2); - expect(store.getState().threads.map((t) => t.id)).toContain("t-new"); + resolveCreate(makeThread("t-real")); + await flushPromises(); + expect(store.getState().isCreatingThread).toBe(false); }); }); describe("deleteThread", () => { - it("removes thread from list", async () => { - const deleteThread = vi.fn().mockResolvedValue(undefined); - const store = makeStore({ deleteThread }); - + it("removes the thread from the list", async () => { + const store = makeStore({ deleteThread: vi.fn().mockResolvedValue(undefined) }); store.setState({ threads: [makeThread("t1"), makeThread("t2", 1)] }); - store.getState().deleteThread("t1"); await flushPromises(); - - expect(store.getState().threads).toHaveLength(1); - expect(store.getState().threads[0].id).toBe("t2"); + expect(store.getState().threads.map((t) => t.id)).toEqual(["t2"]); }); - it("switches to new thread if deleted thread was selected", async () => { - const deleteThread = vi.fn().mockResolvedValue(undefined); - const store = makeStore({ deleteThread }); - + it("clears the active view when the deleted thread was selected", async () => { + const store = makeStore({ deleteThread: vi.fn().mockResolvedValue(undefined) }); store.setState({ threads: [makeThread("t1")], selectedThreadId: "t1", messages: [makeMessage("m1")], }); - store.getState().deleteThread("t1"); await flushPromises(); - expect(store.getState().selectedThreadId).toBeNull(); - expect(store.getState().messages).toEqual([]); + expect(active(store).messages).toEqual([]); }); - it("sets isPending during operation", async () => { - let resolveDelete: () => void; - const deleteThread = vi.fn().mockImplementation( - () => - new Promise((r) => { - resolveDelete = r; - }), - ); - - const store = makeStore({ deleteThread }); - store.setState({ threads: [makeThread("t1")] }); - - store.getState().deleteThread("t1"); - - expect(store.getState().threads[0].isPending).toBe(true); - - resolveDelete!(); + it("aborts and drops a background streaming thread", async () => { + let capturedSignal: AbortSignal | undefined; + const send = vi.fn().mockImplementation(({ signal }) => { + capturedSignal = signal; + return new Promise(() => {}); + }); + const store = makeStore({ + deleteThread: vi.fn().mockResolvedValue(undefined), + send, + streamProtocol: { parse: async function* () {} }, + }); + // Start a run on t1, switch away so it streams in the background. + store.setState({ threads: [makeThread("t1")], selectedThreadId: "t1" }); + store.getState().processMessage({ role: "user", content: "hi" }); await flushPromises(); - - expect(store.getState().threads).toHaveLength(0); - }); - }); - - describe("updateThread", () => { - it("updates thread in list", async () => { - const updated = { ...makeThread("t1"), title: "Renamed" }; - const updateThread = vi.fn().mockResolvedValue(updated); - - const store = makeStore({ updateThread }); - store.setState({ threads: [makeThread("t1")] }); - - store.getState().updateThread(updated); + store.getState().selectThread("t2"); await flushPromises(); + expect(bgOf(store, "t1")?.isRunning).toBe(true); - expect(store.getState().threads[0].title).toBe("Renamed"); + store.getState().deleteThread("t1"); + await flushPromises(); + expect(capturedSignal?.aborted).toBe(true); + expect(bgOf(store, "t1")).toBeUndefined(); }); }); // ──────────────────────────────────────────── - // Message CRUD + // Message CRUD (on the active thread) // ──────────────────────────────────────────── describe("message CRUD", () => { - let store: ReturnType; - - beforeEach(() => { - store = makeStore(); + it("appends, replaces, updates and deletes on the active fields", () => { + const store = makeStore(); store.setState({ 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); - it("setMessages replaces all", () => { - store.getState().setMessages([makeMessage("new")]); - expect(store.getState().messages).toHaveLength(1); - expect(store.getState().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"); - }); + store.getState().updateMessage({ ...makeMessage("m1"), content: "edited" } as Message); + 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.map((m) => m.id)).toEqual(["m2", "m3"]); + + store.getState().setMessages([makeMessage("only")]); + expect(active(store).messages.map((m) => m.id)).toEqual(["only"]); }); }); // ──────────────────────────────────────────── - // processMessage (calls llm.send) + // processMessage // ──────────────────────────────────────────── describe("processMessage", () => { - it("appends optimistic user message and calls llm.send", async () => { + it("appends the 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 () => { - const newThread = makeThread("t-auto"); - const createThread = vi.fn().mockResolvedValue(newThread); + it("creates a thread when none is selected and follows into it", async () => { + const createThread = vi.fn().mockResolvedValue(makeThread("t-auto")); const send = vi.fn().mockResolvedValue(new Response("", { status: 200 })); - const store = makeStore({ createThread, send, @@ -322,36 +266,27 @@ describe("createChatStore", () => { expect(createThread).toHaveBeenCalledOnce(); expect(store.getState().selectedThreadId).toBe("t-auto"); + expect(active(store).messages).toHaveLength(1); + expect(store.getState().isCreatingThread).toBe(false); }); - it("no-ops when already running", async () => { + it("no-ops when the active 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* () {} } }); + store.setState({ selectedThreadId: "t1", isRunning: true }); await store.getState().processMessage({ role: "user", content: "hello" }); - expect(send).not.toHaveBeenCalled(); }); 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?.message).toBe("api down"); + expect(active(store).isRunning).toBe(false); }); }); @@ -360,46 +295,59 @@ describe("createChatStore", () => { // ──────────────────────────────────────────── describe("cancelMessage", () => { - it("aborts in-flight request", async () => { + it("aborts the active thread's run", 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* () {} }, + return new Promise(() => {}); }); + 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: "hi" }); 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 background thread's run", async () => { + const signals: AbortSignal[] = []; + const send = vi.fn().mockImplementation(({ signal }) => { + signals.push(signal); + return new Promise(() => {}); + }); + 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"); // t1 → background + await flushPromises(); + + store.getState().cancelMessage(); // targets the active view (t2), not t1 + await flushPromises(); + expect(signals[0].aborted).toBe(false); + expect(bgOf(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("background streaming", () => { + it("moves a streaming thread to inFlightThreads on switch, promotes it back on return", async () => { + let capturedSignal: AbortSignal | undefined; const send = vi.fn().mockImplementation(({ signal }) => { capturedSignal = signal; - return new Promise(() => {}); // never resolves + return new Promise(() => {}); }); - const newMessages = [makeMessage("new-m1")]; + const newMessages = [makeMessage("t2-m1")]; const getMessages = vi.fn().mockResolvedValue(newMessages); - const store = makeStore({ send, getMessages, @@ -407,22 +355,94 @@ describe("createChatStore", () => { }); store.setState({ selectedThreadId: "t1" }); - // Start streaming store.getState().processMessage({ role: "user", content: "hello" }); await flushPromises(); - expect(store.getState().isRunning).toBe(true); + expect(active(store).isRunning).toBe(true); - // Switch thread mid-stream + // Switch to t2 — t1 keeps streaming in the background (not aborted). store.getState().selectThread("t2"); + expect(capturedSignal?.aborted).toBe(false); + expect(bgOf(store, "t1")?.isRunning).toBe(true); + await flushPromises(); + expect(active(store).messages).toEqual(newMessages); // t2 loaded into active - expect(capturedSignal!.aborted).toBe(true); - expect(store.getState().selectedThreadId).toBe("t2"); - expect(store.getState().isLoadingMessages).toBe(true); + // Switch back to t1 — promoted out of the background into the active fields (no reload). + store.getState().selectThread("t1"); + expect(active(store).isRunning).toBe(true); + expect(active(store).messages).toHaveLength(1); // optimistic user message + expect(bgOf(store, "t1")).toBeUndefined(); + expect(getMessages).toHaveBeenCalledTimes(1); // only t2 was loaded + }); + + it("runs a background thread and the active thread concurrently", async () => { + const signals: AbortSignal[] = []; + const send = vi.fn().mockImplementation(({ signal }) => { + signals.push(signal); + return new Promise(() => {}); + }); + 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(bgOf(store, "t1")?.isRunning).toBe(true); // background + expect(active(store).isRunning).toBe(true); // active (t2) + expect(send).toHaveBeenCalledTimes(2); + expect(signals[0]).not.toBe(signals[1]); + }); + + it("drops a background run from inFlightThreads when it completes", async () => { + let resolveSend!: (r: Response) => void; + const send = vi + .fn() + .mockImplementation(() => new Promise((r) => (resolveSend = r))); + const store = makeStore({ send, streamProtocol: { parse: async function* () {} } }); + + store.setState({ selectedThreadId: "t1" }); + store.getState().processMessage({ role: "user", content: "hi" }); + await flushPromises(); + store.getState().selectThread("t2"); // t1 → background (still awaiting send) + await flushPromises(); + expect(bgOf(store, "t1")?.isRunning).toBe(true); + + resolveSend(new Response("", { status: 200 })); // stream resolves + completes (empty) + await flushPromises(); + expect(bgOf(store, "t1")).toBeUndefined(); // dropped on completion + }); + + it("backgrounds a draft run 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().mockImplementation(() => new Promise(() => {})); // never resolves + const getMessages = vi.fn().mockResolvedValue([]); + const store = makeStore({ + createThread, + send, + getMessages, + streamProtocol: { parse: async function* () {} }, + }); + + store.getState().processMessage({ role: "user", content: "hello" }); + await flushPromises(); + expect(store.getState().isCreatingThread).toBe(true); + + store.getState().selectThread("saved"); // leave the draft mid-creation + await flushPromises(); + resolveCreate(makeThread("t-real")); await flushPromises(); - expect(store.getState().messages).toEqual(newMessages); - expect(store.getState().isLoadingMessages).toBe(false); + expect(store.getState().selectedThreadId).toBe("saved"); // did not follow + expect(bgOf(store, "t-real")?.isRunning).toBe(true); // re-keyed draft runs in background + expect(bgOf(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..17f20a011 100644 --- a/packages/react-headless/src/store/createChatStore.ts +++ b/packages/react-headless/src/store/createChatStore.ts @@ -2,244 +2,369 @@ 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 type { ChatStore, Message, Thread, ThreadStateEntry, UserMessage } from "./types"; export interface CreateChatStoreConfig { storage: ChatStorage; llm: ChatLLM; } +/** Key for a brand-new chat's run before its thread has a real id (see the re-key in `processMessage`). */ +const DRAFT_KEY = "__draft__"; + const mergeThreadList = (existing: Thread[], incoming: Thread[]): Thread[] => Array.from(new Map([...existing, ...incoming].map((t) => [t.id, t])).values()).sort( (a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime(), ); +/** Reset the flat active-thread fields to empty (fresh/blank view). */ +const emptyActive = () => ({ + messages: [] as Message[], + isRunning: false, + isLoadingMessages: false, + threadError: null, + executingToolCallIds: new Set(), + _abortController: null, +}); + export const createChatStore = (config: CreateChatStoreConfig) => { const { storage, llm } = config; 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, + subscribeWithSelector((set, get) => { + /** Snapshot the flat active-thread fields as a background {@link ThreadStateEntry}. */ + const snapshotActive = (s: ChatStore): ThreadStateEntry => ({ + messages: s.messages, + isRunning: s.isRunning, + isLoadingMessages: s.isLoadingMessages, + threadError: s.threadError, + executingToolCallIds: s.executingToolCallIds, + abortController: s._abortController, + }); + + /** Turn a background entry into a flat active-fields patch (`abortController` → `_abortController`). */ + const activeFrom = (e: ThreadStateEntry) => ({ + messages: e.messages, + isRunning: e.isRunning, + isLoadingMessages: e.isLoadingMessages, + threadError: e.threadError, + executingToolCallIds: e.executingToolCallIds, + _abortController: e.abortController, + }); + + /** Read a running thread's live state wherever it lives (active flat fields or the background map). */ + const readRun = (s: ChatStore, key: string): ThreadStateEntry | undefined => + key === (s.selectedThreadId ?? DRAFT_KEY) ? snapshotActive(s) : s.inFlightThreads[key]; + + /** + * Route a run's write to wherever its thread currently lives: the flat active + * fields if it's the on-screen thread, else its `inFlightThreads[key]` entry. + * A run that gets backgrounded mid-flight (user switched away) keeps writing to + * the map; guarded so a dropped background thread swallows late callbacks. + */ + const writeRun = ( + key: string, + fn: (cur: ThreadStateEntry) => Partial | null, + ) => + set((s) => { + if (key === (s.selectedThreadId ?? DRAFT_KEY)) { + const patch = fn(snapshotActive(s)); + if (!patch) return s; + const { abortController, ...rest } = patch; + return "abortController" in patch + ? { ...rest, _abortController: abortController } + : rest; + } + const cur = s.inFlightThreads[key]; + if (!cur) return s; // dropped background thread — ignore late writes + const patch = fn(cur); + if (!patch) return s; + return { inFlightThreads: { ...s.inFlightThreads, [key]: { ...cur, ...patch } } }; + }); + + /** Move the current active thread into the background map (used when switching away mid-stream). */ + const backgroundActive = (key: string) => + set((s) => ({ inFlightThreads: { ...s.inFlightThreads, [key]: snapshotActive(s) } })); + + /** Remove a thread's background entry. */ + const dropInFlight = (key: string) => + set((s) => { + if (!(key in s.inFlightThreads)) return s; + const next = { ...s.inFlightThreads }; + delete next[key]; + return { inFlightThreads: next }; + }); + + return { + // ── Thread List State ── + threads: [], + isLoadingThreads: false, + threadListError: null, + selectedThreadId: null, + hasMoreThreads: false, + isCreatingThread: false, + _nextCursor: undefined, + + // ── Active Thread State (flat, like the single-thread store) ── + ...emptyActive(), + + // ── Background streaming threads (NOT the active one) ── + inFlightThreads: {}, + + // ── 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 }); }); - }) - .catch((e) => { - set({ isLoadingThreads: false, threadListError: e }); - }); - }, - - loadMoreThreads: () => { - const cursor = get()._nextCursor; - if (cursor === undefined) return; - threadStorage - .listThreads(cursor) - .then(({ threads = [], nextCursor }) => { + }, + + 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: () => { + // Guard the create window — a draft still becoming a real thread owns the + // flat fields; resetting now would collide with its re-key. + if (get().isCreatingThread) return; + const s = get(); + // A still-streaming active thread keeps running in the background; an idle + // one is simply discarded (no stale cache). + if (s.isRunning) backgroundActive(s.selectedThreadId ?? DRAFT_KEY); + set({ selectedThreadId: null, ...emptyActive() }); + }, + + createThread: async (firstMessage: UserMessage) => { + const thread = await threadStorage.createThread(firstMessage); + set((s) => ({ threads: mergeThreadList(s.threads, [thread]) })); + return thread; + }, + + selectThread: (threadId: string) => { + const s = get(); + const leftKey = s.selectedThreadId ?? DRAFT_KEY; + if (leftKey === threadId) return; // re-selecting the active thread — no-op + + // Keep the thread we're leaving alive in the background only if it's still + // streaming; otherwise let it go (it will reload from storage on return). + if (s.isRunning) backgroundActive(leftKey); + + set({ selectedThreadId: threadId }); + + // If the target is streaming in the background, promote it into the flat + // active fields (live — no reload). + const bg = get().inFlightThreads[threadId]; + if (bg) { + set((st) => { + const next = { ...st.inFlightThreads }; + delete next[threadId]; + return { ...activeFrom(bg), inFlightThreads: next }; + }); + return; + } + + // Otherwise load a fresh copy from storage. + set({ ...emptyActive(), isLoadingMessages: true }); + threadStorage + .getMessages(threadId) + .then((messages) => { + if ((get().selectedThreadId ?? DRAFT_KEY) !== threadId) return; // navigated away + set({ messages, isLoadingMessages: false }); + }) + .catch((e) => { + if ((get().selectedThreadId ?? DRAFT_KEY) !== threadId) return; + set({ threadError: e, isLoadingMessages: false }); + }); + }, + + updateThread: (thread: Thread) => { + const setPending = (id: string, isPending: boolean) => set((s) => ({ - threads: mergeThreadList(s.threads, threads), - _nextCursor: nextCursor, - hasMoreThreads: nextCursor !== undefined, + threads: s.threads.map((t) => (t.id === id ? { ...t, isPending } : t)), })); - }) - .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(), - }); - 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) => { + 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 === updated.id ? updated : t)), + threads: s.threads.map((t) => (t.id === id ? { ...t, isPending } : 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(), - }); - set((s) => ({ messages: [...s.messages, optimisticMessage] })); + setPending(threadId, true); + threadStorage + .deleteThread(threadId) + .then(() => { + const state = get(); + const isActive = state.selectedThreadId === threadId; + // Abort its run whether it's the active thread or a background stream. + if (isActive) state._abortController?.abort(); + else state.inFlightThreads[threadId]?.abortController?.abort(); + dropInFlight(threadId); + set((s) => ({ threads: s.threads.filter((t) => t.id !== threadId) })); + // If we deleted the active thread, clear the flat fields to a blank chat + // (don't background it — it's gone). + if (isActive) set({ selectedThreadId: null, ...emptyActive() }); + }) + .catch(() => setPending(threadId, false)); + }, - abortController.signal.addEventListener("abort", () => { - set({ _abortController: null, isRunning: false }); - }); + // ── Thread Actions ── - try { - let threadId = get().selectedThreadId; + processMessage: async (message) => { + // A run only ever starts on the active thread; block if it's already running. + if (get().isRunning) return; - if (!threadId) { - const created = await get().createThread(optimisticMessage); - threadId = created.id; - set({ selectedThreadId: threadId }); - } + const isNewChat = !get().selectedThreadId; + // The run's key — a mutable local so post-re-key callbacks target the real id. + let runKey = get().selectedThreadId ?? DRAFT_KEY; + const abortController = new AbortController(); + const optimisticMessage: UserMessage = { + ...message, + id: crypto.randomUUID(), + role: "user", + }; + + // Start the run on the active (flat) fields. + set((s) => ({ + messages: [...s.messages, optimisticMessage], + isRunning: true, + threadError: null, + executingToolCallIds: new Set(), + _abortController: abortController, + ...(isNewChat ? { isCreatingThread: true } : null), + })); - const response = await llm.send({ - threadId, - messages: get().messages, - signal: abortController.signal, + // On abort, flip the run off wherever it now lives. + abortController.signal.addEventListener("abort", () => { + writeRun(runKey, () => ({ isRunning: false, abortController: null })); }); - if (response instanceof Response && !response.ok) { - throw new Error(`Request failed: ${response.status} ${response.statusText}`); - } + try { + if (isNewChat) { + try { + const created = await get().createThread(optimisticMessage); + // Re-key the draft to its real id. Two cases: + // A) still on the draft (flat active) → just relabel selectedThreadId. + // B) navigated away mid-creation → the draft was moved to + // inFlightThreads[DRAFT_KEY]; re-key it there (stays background). + set((s) => { + if (DRAFT_KEY in s.inFlightThreads) { + const draft = s.inFlightThreads[DRAFT_KEY]!; + const next = { ...s.inFlightThreads }; + delete next[DRAFT_KEY]; + next[created.id] = draft; + return { inFlightThreads: next }; + } + return s.selectedThreadId === null ? { selectedThreadId: created.id } : s; + }); + runKey = created.id; + } finally { + set({ isCreatingThread: false }); + } + } - await processStreamedMessage({ - response, - createMessage: (msg) => set((s) => ({ messages: [...s.messages, msg] })), - updateMessage: (msg) => - 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) }, - ), - // 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)) }); + const response = await llm.send({ + threadId: runKey, + messages: readRun(get(), 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) => + writeRun(runKey, (cur) => ({ messages: [...cur.messages, msg] })), + updateMessage: (msg) => + writeRun(runKey, (cur) => ({ + messages: cur.messages.map((m) => (m.id === msg.id ? msg : m)), + })), + // Keep the Set reference stable (return null) when membership is + // unchanged so `useToolActivities` doesn't re-run needlessly. + markToolExecuting: (id) => + writeRun(runKey, (cur) => + cur.executingToolCallIds.has(id) + ? null + : { executingToolCallIds: new Set(cur.executingToolCallIds).add(id) }, + ), + clearToolExecuting: (id) => + writeRun(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) { + writeRun(runKey, () => ({ + threadError: e instanceof Error ? e : new Error(String(e)), + })); + } + } finally { + writeRun(runKey, () => ({ + isRunning: false, + abortController: null, + executingToolCallIds: new Set(), + })); + // If the run finished on a BACKGROUND thread, drop it — it's no longer + // active or streaming, so a return trip reloads it fresh from storage. + if (runKey !== (get().selectedThreadId ?? DRAFT_KEY)) dropInFlight(runKey); } - } 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, - executingToolCallIds: new Set(), - }); - } - }, - - 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(); - }, - })), + }, + + 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(); + }, + }; + }), ); return store; diff --git a/packages/react-headless/src/store/types.ts b/packages/react-headless/src/store/types.ts index 4ddc3c870..dd8ab7a25 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 = { @@ -39,15 +40,15 @@ export type ThreadState = { isRunning: boolean; isLoadingMessages: boolean; threadError: Error | null; + executingToolCallIds: Set; +}; + +export type ThreadStateEntry = ThreadState & { /** - * Tool calls whose arguments have closed (`TOOL_CALL_END` seen) but whose - * result message has not yet arrived — i.e. the tool is currently executing. - * Drives the `"executing"` status in {@link ToolActivity}; reset to an empty - * set when a new message run starts or the thread switches. The reference is - * stable across unrelated store updates (a new `Set` is created only when the - * membership changes), so selector consumers don't re-render needlessly. + * @internal the in-flight run's controller, or `null` when idle. + * Intentionally not exposed to avoid usage from public contract */ - executingToolCallIds: Set; + abortController: AbortController | null; }; export type ThreadActions = { @@ -65,9 +66,16 @@ export type ChatStore = ThreadListState & ThreadListActions & ThreadState & ThreadActions & { + /** + * Threads streaming in the BACKGROUND — i.e. runs on threads that are not the + * active view. The active thread lives in the flat {@link ThreadState} fields + * above; a background entry is dropped the moment its run completes (a return + * trip reloads it from storage), so this holds only genuinely in-flight threads. + */ + inFlightThreads: Record; /** @internal */ _nextCursor?: string | undefined; - /** @internal */ + /** @internal the active thread's in-flight run controller, or `null` when idle. */ _abortController: AbortController | null; }; 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" > ); }; diff --git a/packages/react-ui/src/components/AgentInterface/threadlist.scss b/packages/react-ui/src/components/AgentInterface/threadlist.scss index f4ef47d82..be54baf71 100644 --- a/packages/react-ui/src/components/AgentInterface/threadlist.scss +++ b/packages/react-ui/src/components/AgentInterface/threadlist.scss @@ -201,6 +201,45 @@ $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; + 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__icon { + animation: none; + } +} + .openui-agent-thread-button-dropdown-menu { display: flex; flex-direction: column;