diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
index 79f749afeb..dd7eb4bb16 100644
--- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
+++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
@@ -7,9 +7,13 @@ import { QueryClient, QueryClientProvider } from "react-query";
import type { OneWayMessageEvent } from "utils/OneWayWebSocket";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
+ selectChatStatus,
selectOrderedMessageIDs,
selectQueuedMessages,
+ selectRetryState,
+ selectStreamError,
selectStreamState,
+ selectSubagentStatusOverrides,
useChatSelector,
useChatStore,
} from "./ChatContext";
@@ -31,6 +35,7 @@ interface MockSocket {
close: () => void;
emitData: (event: TypesGen.ChatStreamEvent) => void;
emitDataBatch: (events: readonly TypesGen.ChatStreamEvent[]) => void;
+ emitError: () => void;
}
const createMockSocket = (): MockSocket => {
@@ -89,6 +94,11 @@ const createMockSocket = (): MockSocket => {
listener(payload);
}
},
+ emitError: () => {
+ for (const listener of errorListeners) {
+ listener(new Event("error"));
+ }
+ },
};
};
@@ -150,6 +160,7 @@ const immediateAnimationFrame = (): void => {
afterEach(() => {
vi.restoreAllMocks();
+ vi.mocked(watchChat).mockReset();
});
describe("useChatStore", () => {
@@ -1373,4 +1384,483 @@ describe("useChatStore", () => {
expect(result.current.streamState).toBeNull();
});
});
+
+ it("sets chatStatus to error and populates streamError on error event", async () => {
+ immediateAnimationFrame();
+
+ const chatID = "chat-error";
+ const mockSocket = createMockSocket();
+ vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ const { result } = renderHook(
+ () => {
+ const { store } = useChatStore({
+ chatID,
+ chatMessages: [],
+ chatRecord: makeChat(chatID),
+ chatData: {
+ chat: makeChat(chatID),
+ messages: [],
+ queued_messages: [],
+ },
+ chatQueuedMessages: [],
+ setChatErrorReason,
+ clearChatErrorReason,
+ });
+ return {
+ chatStatus: useChatSelector(store, selectChatStatus),
+ streamError: useChatSelector(store, selectStreamError),
+ retryState: useChatSelector(store, selectRetryState),
+ };
+ },
+ { wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID);
+ });
+
+ act(() => {
+ mockSocket.emitData({
+ type: "error",
+ chat_id: chatID,
+ error: { message: "Rate limit exceeded" },
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.chatStatus).toBe("error");
+ });
+ expect(result.current.streamError).toBe("Rate limit exceeded");
+ expect(result.current.retryState).toBeNull();
+ expect(setChatErrorReason).toHaveBeenCalledWith(
+ chatID,
+ "Rate limit exceeded",
+ );
+ });
+
+ it("uses fallback message when error event has no message", async () => {
+ immediateAnimationFrame();
+
+ const chatID = "chat-error-empty";
+ const mockSocket = createMockSocket();
+ vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ const { result } = renderHook(
+ () => {
+ const { store } = useChatStore({
+ chatID,
+ chatMessages: [],
+ chatRecord: makeChat(chatID),
+ chatData: {
+ chat: makeChat(chatID),
+ messages: [],
+ queued_messages: [],
+ },
+ chatQueuedMessages: [],
+ setChatErrorReason,
+ clearChatErrorReason,
+ });
+ return {
+ streamError: useChatSelector(store, selectStreamError),
+ };
+ },
+ { wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID);
+ });
+
+ act(() => {
+ mockSocket.emitData({
+ type: "error",
+ chat_id: chatID,
+ error: { message: " " },
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.streamError).toBe("Chat processing failed.");
+ });
+ });
+
+ it("populates retryState on retry event", async () => {
+ immediateAnimationFrame();
+
+ const chatID = "chat-retry";
+ const mockSocket = createMockSocket();
+ vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ const { result } = renderHook(
+ () => {
+ const { store } = useChatStore({
+ chatID,
+ chatMessages: [],
+ chatRecord: makeChat(chatID),
+ chatData: {
+ chat: makeChat(chatID),
+ messages: [],
+ queued_messages: [],
+ },
+ chatQueuedMessages: [],
+ setChatErrorReason,
+ clearChatErrorReason,
+ });
+ return {
+ retryState: useChatSelector(store, selectRetryState),
+ };
+ },
+ { wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID);
+ });
+
+ act(() => {
+ mockSocket.emitData({
+ type: "retry",
+ chat_id: chatID,
+ retry: {
+ attempt: 2,
+ error: "upstream timeout",
+ delay_ms: 5000,
+ retrying_at: "2025-01-01T00:01:00.000Z",
+ },
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.retryState).toEqual({
+ attempt: 2,
+ error: "upstream timeout",
+ });
+ });
+ });
+
+ it("clears retryState when status transitions to running", async () => {
+ immediateAnimationFrame();
+
+ const chatID = "chat-retry-clear";
+ const mockSocket = createMockSocket();
+ vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ const { result } = renderHook(
+ () => {
+ const { store } = useChatStore({
+ chatID,
+ chatMessages: [],
+ chatRecord: makeChat(chatID),
+ chatData: {
+ chat: makeChat(chatID),
+ messages: [],
+ queued_messages: [],
+ },
+ chatQueuedMessages: [],
+ setChatErrorReason,
+ clearChatErrorReason,
+ });
+ return {
+ retryState: useChatSelector(store, selectRetryState),
+ chatStatus: useChatSelector(store, selectChatStatus),
+ };
+ },
+ { wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID);
+ });
+
+ // Set retry state first.
+ act(() => {
+ mockSocket.emitData({
+ type: "retry",
+ chat_id: chatID,
+ retry: {
+ attempt: 1,
+ error: "rate limited",
+ delay_ms: 3000,
+ retrying_at: "2025-01-01T00:00:30.000Z",
+ },
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.retryState).not.toBeNull();
+ });
+
+ // Transition to running — should clear retry state.
+ act(() => {
+ mockSocket.emitData({
+ type: "status",
+ chat_id: chatID,
+ status: { status: "running" },
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.chatStatus).toBe("running");
+ });
+ expect(result.current.retryState).toBeNull();
+ });
+
+ it("routes status events for other chatIDs to subagent overrides", async () => {
+ immediateAnimationFrame();
+
+ const chatID = "chat-main";
+ const subagentChatID = "chat-subagent-1";
+ const mockSocket = createMockSocket();
+ vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ const { result } = renderHook(
+ () => {
+ const { store } = useChatStore({
+ chatID,
+ chatMessages: [],
+ chatRecord: makeChat(chatID),
+ chatData: {
+ chat: makeChat(chatID),
+ messages: [],
+ queued_messages: [],
+ },
+ chatQueuedMessages: [],
+ setChatErrorReason,
+ clearChatErrorReason,
+ });
+ return {
+ chatStatus: useChatSelector(store, selectChatStatus),
+ subagentStatusOverrides: useChatSelector(
+ store,
+ selectSubagentStatusOverrides,
+ ),
+ };
+ },
+ { wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID);
+ });
+
+ act(() => {
+ mockSocket.emitData({
+ type: "status",
+ chat_id: subagentChatID,
+ status: { status: "completed" },
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.subagentStatusOverrides.get(subagentChatID)).toBe(
+ "completed",
+ );
+ });
+ // Main chat status should remain "running" from the initial
+ // chatRecord — the subagent status event must not change it.
+ expect(result.current.chatStatus).toBe("running");
+ });
+
+ it("sets streamError on WebSocket disconnect", async () => {
+ immediateAnimationFrame();
+
+ const chatID = "chat-disconnect";
+ const mockSocket = createMockSocket();
+ vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ const { result } = renderHook(
+ () => {
+ const { store } = useChatStore({
+ chatID,
+ chatMessages: [],
+ chatRecord: makeChat(chatID),
+ chatData: {
+ chat: makeChat(chatID),
+ messages: [],
+ queued_messages: [],
+ },
+ chatQueuedMessages: [],
+ setChatErrorReason,
+ clearChatErrorReason,
+ });
+ return {
+ streamError: useChatSelector(store, selectStreamError),
+ };
+ },
+ { wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID);
+ });
+
+ act(() => {
+ mockSocket.emitError();
+ });
+
+ await waitFor(() => {
+ expect(result.current.streamError).toBe("Chat stream disconnected.");
+ });
+ });
+
+ it("does not overwrite existing streamError on WebSocket disconnect", async () => {
+ immediateAnimationFrame();
+
+ const chatID = "chat-disconnect-existing";
+ const mockSocket = createMockSocket();
+ vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ const { result } = renderHook(
+ () => {
+ const { store } = useChatStore({
+ chatID,
+ chatMessages: [],
+ chatRecord: makeChat(chatID),
+ chatData: {
+ chat: makeChat(chatID),
+ messages: [],
+ queued_messages: [],
+ },
+ chatQueuedMessages: [],
+ setChatErrorReason,
+ clearChatErrorReason,
+ });
+ return {
+ streamError: useChatSelector(store, selectStreamError),
+ };
+ },
+ { wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID);
+ });
+
+ // Set an error via an error stream event first.
+ act(() => {
+ mockSocket.emitData({
+ type: "error",
+ chat_id: chatID,
+ error: { message: "Rate limit exceeded" },
+ });
+ });
+
+ await waitFor(() => {
+ expect(result.current.streamError).toBe("Rate limit exceeded");
+ });
+
+ // WebSocket disconnect should NOT overwrite the existing error.
+ act(() => {
+ mockSocket.emitError();
+ });
+
+ // The original error should be preserved.
+ await waitFor(() => {
+ expect(result.current.streamError).toBe("Rate limit exceeded");
+ });
+ });
+
+ it("clears chatErrorReason when status transitions to non-error", async () => {
+ immediateAnimationFrame();
+
+ const chatID = "chat-clear-error";
+ const mockSocket = createMockSocket();
+ vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+ const queryClient = createTestQueryClient();
+ const wrapper = ({ children }: PropsWithChildren) => (
+ {children}
+ );
+ const setChatErrorReason = vi.fn();
+ const clearChatErrorReason = vi.fn();
+
+ renderHook(
+ () => {
+ const { store } = useChatStore({
+ chatID,
+ chatMessages: [],
+ chatRecord: makeChat(chatID),
+ chatData: {
+ chat: makeChat(chatID),
+ messages: [],
+ queued_messages: [],
+ },
+ chatQueuedMessages: [],
+ setChatErrorReason,
+ clearChatErrorReason,
+ });
+ return {
+ chatStatus: useChatSelector(store, selectChatStatus),
+ };
+ },
+ { wrapper },
+ );
+
+ await waitFor(() => {
+ expect(watchChat).toHaveBeenCalledWith(chatID);
+ });
+
+ // Transition to running — should call clearChatErrorReason.
+ act(() => {
+ mockSocket.emitData({
+ type: "status",
+ chat_id: chatID,
+ status: { status: "running" },
+ });
+ });
+
+ await waitFor(() => {
+ expect(clearChatErrorReason).toHaveBeenCalledWith(chatID);
+ });
+ });
});
diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
index afec590fc4..10f2186ed5 100644
--- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
+++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
@@ -185,7 +185,7 @@ const createInitialState = (): ChatStoreState => ({
subagentStatusOverrides: new Map(),
});
-const createChatStore = (): ChatStore => {
+export const createChatStore = (): ChatStore => {
let state = createInitialState();
const listeners = new Set<() => void>();
diff --git a/site/src/pages/AgentsPage/AgentDetail/chatStore.test.ts b/site/src/pages/AgentsPage/AgentDetail/chatStore.test.ts
new file mode 100644
index 0000000000..d90fbb35aa
--- /dev/null
+++ b/site/src/pages/AgentsPage/AgentDetail/chatStore.test.ts
@@ -0,0 +1,535 @@
+import type * as TypesGen from "api/typesGenerated";
+import { describe, expect, it } from "vitest";
+import { createChatStore } from "./ChatContext";
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+/** Minimal ChatMessage factory. `created_at` is derived from `id` to make
+ * ordering deterministic in tests that care about sort order. */
+const makeMessage = (
+ id: number,
+ role: string,
+ text: string,
+ chatID = "chat-1",
+): TypesGen.ChatMessage =>
+ ({
+ id,
+ chat_id: chatID,
+ created_at: `2025-01-01T00:00:0${Math.max(Math.abs(id), 0)}.000Z`,
+ role,
+ content: [{ type: "text", text }],
+ }) as TypesGen.ChatMessage;
+
+const makeQueuedMessage = (
+ id: number,
+ text: string,
+ chatID = "chat-1",
+): TypesGen.ChatQueuedMessage =>
+ ({
+ id,
+ chat_id: chatID,
+ created_at: "2025-01-01T00:00:00Z",
+ content: [{ type: "text", text }],
+ }) as TypesGen.ChatQueuedMessage;
+
+// ---------------------------------------------------------------------------
+// replaceMessages
+// ---------------------------------------------------------------------------
+
+describe("replaceMessages", () => {
+ it("populates messagesByID and orderedMessageIDs", () => {
+ const store = createChatStore();
+ const msg1 = makeMessage(1, "user", "first");
+ const msg2 = makeMessage(2, "assistant", "second");
+
+ store.replaceMessages([msg1, msg2]);
+
+ const state = store.getSnapshot();
+ expect(state.messagesByID.size).toBe(2);
+ expect(state.messagesByID.get(1)).toBe(msg1);
+ expect(state.messagesByID.get(2)).toBe(msg2);
+ expect(state.orderedMessageIDs).toEqual([1, 2]);
+ });
+
+ it("sorts messages by created_at", () => {
+ const store = createChatStore();
+ const older = {
+ ...makeMessage(1, "user", "first"),
+ created_at: "2025-01-01T00:00:01.000Z",
+ } as TypesGen.ChatMessage;
+ const newer = {
+ ...makeMessage(2, "assistant", "second"),
+ created_at: "2025-01-01T00:00:05.000Z",
+ } as TypesGen.ChatMessage;
+
+ // Insert in reverse order.
+ store.replaceMessages([newer, older]);
+
+ expect(store.getSnapshot().orderedMessageIDs).toEqual([1, 2]);
+ });
+
+ it("treats undefined as empty array", () => {
+ const store = createChatStore();
+ store.replaceMessages([makeMessage(1, "user", "hello")]);
+
+ store.replaceMessages(undefined);
+
+ const state = store.getSnapshot();
+ expect(state.messagesByID.size).toBe(0);
+ expect(state.orderedMessageIDs).toEqual([]);
+ });
+
+ it("does not notify subscribers when content is unchanged", () => {
+ const store = createChatStore();
+ const msg = makeMessage(1, "user", "hello");
+ store.replaceMessages([msg]);
+
+ let notified = false;
+ store.subscribe(() => {
+ notified = true;
+ });
+
+ // Same object reference — maps compare equal by ref.
+ store.replaceMessages([msg]);
+
+ expect(notified).toBe(false);
+ });
+});
+
+// ---------------------------------------------------------------------------
+// upsertDurableMessage
+// ---------------------------------------------------------------------------
+
+describe("upsertDurableMessage", () => {
+ it("inserts a new message and reports isDuplicate=false, changed=true", () => {
+ const store = createChatStore();
+ const msg = makeMessage(1, "user", "hello");
+
+ const result = store.upsertDurableMessage(msg);
+
+ expect(result).toEqual({ isDuplicate: false, changed: true });
+ expect(store.getSnapshot().messagesByID.get(1)).toBe(msg);
+ expect(store.getSnapshot().orderedMessageIDs).toEqual([1]);
+ });
+
+ it("reports isDuplicate=true, changed=false for value-equal duplicate", () => {
+ const store = createChatStore();
+ const msg = makeMessage(1, "user", "hello");
+ store.upsertDurableMessage(msg);
+
+ // Different object reference, same field values.
+ const dup = makeMessage(1, "user", "hello");
+ const result = store.upsertDurableMessage(dup);
+
+ expect(result).toEqual({ isDuplicate: true, changed: false });
+ });
+
+ it("reports isDuplicate=true, changed=true when content differs", () => {
+ const store = createChatStore();
+ store.upsertDurableMessage(makeMessage(1, "assistant", "draft"));
+
+ const updated = makeMessage(1, "assistant", "final");
+ const result = store.upsertDurableMessage(updated);
+
+ expect(result).toEqual({ isDuplicate: true, changed: true });
+ expect(store.getSnapshot().messagesByID.get(1)?.content).toEqual(
+ updated.content,
+ );
+ });
+
+ it("removes optimistic (negative-ID) messages when a real message arrives", () => {
+ const store = createChatStore();
+ const optimistic = makeMessage(-1, "user", "typing...");
+ store.replaceMessages([optimistic]);
+ expect(store.getSnapshot().messagesByID.has(-1)).toBe(true);
+
+ const real = makeMessage(5, "user", "typed!");
+ store.upsertDurableMessage(real);
+
+ expect(store.getSnapshot().messagesByID.has(-1)).toBe(false);
+ expect(store.getSnapshot().messagesByID.has(5)).toBe(true);
+ });
+
+ it("only removes optimistic messages with the same role", () => {
+ const store = createChatStore();
+ const optimisticUser = makeMessage(-1, "user", "my prompt");
+ const optimisticAssistant = makeMessage(-2, "assistant", "placeholder");
+ store.replaceMessages([optimisticUser, optimisticAssistant]);
+
+ // A real "user" message arrives — only the user optimistic should
+ // be removed, not the assistant one.
+ store.upsertDurableMessage(makeMessage(5, "user", "real prompt"));
+
+ expect(store.getSnapshot().messagesByID.has(-1)).toBe(false);
+ expect(store.getSnapshot().messagesByID.has(-2)).toBe(true);
+ });
+
+ it("does not reorder when updating an existing message in place", () => {
+ const store = createChatStore();
+ store.upsertDurableMessage(makeMessage(1, "user", "first"));
+ store.upsertDurableMessage(makeMessage(2, "assistant", "second"));
+ const orderBefore = store.getSnapshot().orderedMessageIDs;
+
+ // Update content of existing message (same ID, same map size).
+ store.upsertDurableMessage(makeMessage(2, "assistant", "edited"));
+
+ // Same reference — no reorder needed because the map size
+ // didn't change and the ID already existed.
+ expect(store.getSnapshot().orderedMessageIDs).toBe(orderBefore);
+ });
+});
+
+// ---------------------------------------------------------------------------
+// setChatStatus
+// ---------------------------------------------------------------------------
+
+describe("setChatStatus", () => {
+ it("updates chatStatus", () => {
+ const store = createChatStore();
+
+ store.setChatStatus("running");
+
+ expect(store.getSnapshot().chatStatus).toBe("running");
+ });
+
+ it("accepts null to clear the status", () => {
+ const store = createChatStore();
+ store.setChatStatus("running");
+
+ store.setChatStatus(null);
+
+ expect(store.getSnapshot().chatStatus).toBeNull();
+ });
+
+ it("does not notify when setting the same status", () => {
+ const store = createChatStore();
+ store.setChatStatus("running");
+
+ let notified = false;
+ store.subscribe(() => {
+ notified = true;
+ });
+ store.setChatStatus("running");
+
+ expect(notified).toBe(false);
+ });
+});
+
+// ---------------------------------------------------------------------------
+// setStreamError / clearStreamError
+// ---------------------------------------------------------------------------
+
+describe("setStreamError / clearStreamError", () => {
+ it("stores and clears a stream error", () => {
+ const store = createChatStore();
+
+ store.setStreamError("connection lost");
+ expect(store.getSnapshot().streamError).toBe("connection lost");
+
+ store.clearStreamError();
+ expect(store.getSnapshot().streamError).toBeNull();
+ });
+
+ it("does not notify when setting the same error", () => {
+ const store = createChatStore();
+ store.setStreamError("oops");
+
+ let notified = false;
+ store.subscribe(() => {
+ notified = true;
+ });
+ store.setStreamError("oops");
+
+ expect(notified).toBe(false);
+ });
+
+ it("clearStreamError is a no-op when already null", () => {
+ const store = createChatStore();
+
+ let notified = false;
+ store.subscribe(() => {
+ notified = true;
+ });
+ store.clearStreamError();
+
+ expect(notified).toBe(false);
+ });
+});
+
+// ---------------------------------------------------------------------------
+// setRetryState / clearRetryState
+// ---------------------------------------------------------------------------
+
+describe("setRetryState / clearRetryState", () => {
+ it("stores and clears retry state", () => {
+ const store = createChatStore();
+
+ store.setRetryState({ attempt: 1, error: "rate limited" });
+ expect(store.getSnapshot().retryState).toEqual({
+ attempt: 1,
+ error: "rate limited",
+ });
+
+ store.clearRetryState();
+ expect(store.getSnapshot().retryState).toBeNull();
+ });
+
+ it("clearRetryState is a no-op when already null", () => {
+ const store = createChatStore();
+
+ let notified = false;
+ store.subscribe(() => {
+ notified = true;
+ });
+ store.clearRetryState();
+
+ expect(notified).toBe(false);
+ });
+});
+
+// ---------------------------------------------------------------------------
+// setSubagentStatusOverride
+// ---------------------------------------------------------------------------
+
+describe("setSubagentStatusOverride", () => {
+ it("stores per-chatID status overrides", () => {
+ const store = createChatStore();
+
+ store.setSubagentStatusOverride("sub-1", "running");
+ store.setSubagentStatusOverride("sub-2", "error");
+
+ const overrides = store.getSnapshot().subagentStatusOverrides;
+ expect(overrides.get("sub-1")).toBe("running");
+ expect(overrides.get("sub-2")).toBe("error");
+ });
+
+ it("does not notify when the override is unchanged", () => {
+ const store = createChatStore();
+ store.setSubagentStatusOverride("sub-1", "running");
+
+ let notified = false;
+ store.subscribe(() => {
+ notified = true;
+ });
+ store.setSubagentStatusOverride("sub-1", "running");
+
+ expect(notified).toBe(false);
+ });
+
+ it("overwrites an existing override for the same chatID", () => {
+ const store = createChatStore();
+ store.setSubagentStatusOverride("sub-1", "running");
+ store.setSubagentStatusOverride("sub-1", "completed");
+
+ expect(store.getSnapshot().subagentStatusOverrides.get("sub-1")).toBe(
+ "completed",
+ );
+ });
+});
+
+// ---------------------------------------------------------------------------
+// setQueuedMessages
+// ---------------------------------------------------------------------------
+
+describe("setQueuedMessages", () => {
+ it("stores queued messages", () => {
+ const store = createChatStore();
+ const qm = makeQueuedMessage(10, "queued");
+
+ store.setQueuedMessages([qm]);
+
+ expect(store.getSnapshot().queuedMessages).toEqual([qm]);
+ });
+
+ it("treats undefined as empty array", () => {
+ const store = createChatStore();
+ store.setQueuedMessages([makeQueuedMessage(1, "q")]);
+
+ store.setQueuedMessages(undefined);
+
+ expect(store.getSnapshot().queuedMessages).toEqual([]);
+ });
+
+ it("does not notify when queued message IDs are unchanged", () => {
+ const store = createChatStore();
+ const qm = makeQueuedMessage(10, "queued");
+ store.setQueuedMessages([qm]);
+
+ let notified = false;
+ store.subscribe(() => {
+ notified = true;
+ });
+
+ // Different object reference, same ID.
+ store.setQueuedMessages([{ ...qm }]);
+
+ expect(notified).toBe(false);
+ });
+});
+
+// ---------------------------------------------------------------------------
+// clearStreamState
+// ---------------------------------------------------------------------------
+
+describe("clearStreamState", () => {
+ it("clears stream state to null", () => {
+ const store = createChatStore();
+ // Build up some stream state via applyMessagePart.
+ store.applyMessagePart({ type: "text", text: "hello" });
+ expect(store.getSnapshot().streamState).not.toBeNull();
+
+ store.clearStreamState();
+
+ expect(store.getSnapshot().streamState).toBeNull();
+ });
+
+ it("is a no-op when stream state is already null", () => {
+ const store = createChatStore();
+
+ let notified = false;
+ store.subscribe(() => {
+ notified = true;
+ });
+ store.clearStreamState();
+
+ expect(notified).toBe(false);
+ });
+});
+
+// ---------------------------------------------------------------------------
+// applyMessagePart / applyMessageParts
+// ---------------------------------------------------------------------------
+
+describe("applyMessagePart / applyMessageParts", () => {
+ it("creates stream state from a text part", () => {
+ const store = createChatStore();
+
+ store.applyMessagePart({ type: "text", text: "hello" });
+
+ expect(store.getSnapshot().streamState?.blocks).toEqual([
+ { type: "response", text: "hello" },
+ ]);
+ });
+
+ it("appends to existing stream state", () => {
+ const store = createChatStore();
+ store.applyMessagePart({ type: "text", text: "hello" });
+ store.applyMessagePart({ type: "text", text: " world" });
+
+ expect(store.getSnapshot().streamState?.blocks).toEqual([
+ { type: "response", text: "hello world" },
+ ]);
+ });
+
+ it("applies multiple parts in a single batch", () => {
+ const store = createChatStore();
+
+ store.applyMessageParts([
+ { type: "text", text: "one" },
+ { type: "text", text: " two" },
+ ]);
+
+ expect(store.getSnapshot().streamState?.blocks).toEqual([
+ { type: "response", text: "one two" },
+ ]);
+ });
+
+ it("is a no-op for an empty parts array", () => {
+ const store = createChatStore();
+
+ let notified = false;
+ store.subscribe(() => {
+ notified = true;
+ });
+ store.applyMessageParts([]);
+
+ expect(notified).toBe(false);
+ });
+});
+
+// ---------------------------------------------------------------------------
+// resetTransientState
+// ---------------------------------------------------------------------------
+
+describe("resetTransientState", () => {
+ it("clears streamState, streamError, retryState, and subagentOverrides", () => {
+ const store = createChatStore();
+ store.applyMessagePart({ type: "text", text: "stream" });
+ store.setStreamError("oops");
+ store.setRetryState({ attempt: 2, error: "rate limit" });
+ store.setSubagentStatusOverride("sub-1", "error");
+
+ store.resetTransientState();
+
+ const state = store.getSnapshot();
+ expect(state.streamState).toBeNull();
+ expect(state.streamError).toBeNull();
+ expect(state.retryState).toBeNull();
+ expect(state.subagentStatusOverrides.size).toBe(0);
+ });
+
+ it("preserves messages and queued messages", () => {
+ const store = createChatStore();
+ store.replaceMessages([makeMessage(1, "user", "hello")]);
+ store.setQueuedMessages([makeQueuedMessage(10, "queued")]);
+ store.setStreamError("oops");
+
+ store.resetTransientState();
+
+ const state = store.getSnapshot();
+ expect(state.messagesByID.size).toBe(1);
+ expect(state.queuedMessages).toHaveLength(1);
+ });
+
+ it("is a no-op when all transient state is already clean", () => {
+ const store = createChatStore();
+
+ let notified = false;
+ store.subscribe(() => {
+ notified = true;
+ });
+ store.resetTransientState();
+
+ expect(notified).toBe(false);
+ });
+});
+
+// ---------------------------------------------------------------------------
+// subscribe
+// ---------------------------------------------------------------------------
+
+describe("subscribe", () => {
+ it("returns an unsubscribe function that prevents future notifications", () => {
+ const store = createChatStore();
+ let callCount = 0;
+ const unsubscribe = store.subscribe(() => {
+ callCount += 1;
+ });
+
+ store.setChatStatus("running");
+ expect(callCount).toBe(1);
+
+ unsubscribe();
+ store.setChatStatus("error");
+ expect(callCount).toBe(1);
+ });
+
+ it("supports multiple concurrent subscribers", () => {
+ const store = createChatStore();
+ let countA = 0;
+ let countB = 0;
+ store.subscribe(() => {
+ countA += 1;
+ });
+ store.subscribe(() => {
+ countB += 1;
+ });
+
+ store.setChatStatus("running");
+
+ expect(countA).toBe(1);
+ expect(countB).toBe(1);
+ });
+});