test(site): add comprehensive ChatContext store and integration tests (#22444)

## Summary

Export `createChatStore`, `ChatStore`, and `ChatStoreState` from
`ChatContext.ts` so the pure store logic can be unit tested directly
without React rendering overhead.

## Changes

### Production code (3-line change)
- Added `export` to `ChatStoreState`, `ChatStore`, and `createChatStore`
in `ChatContext.ts`

### chatStore.test.ts — 35 pure store unit tests (runs in ~6ms)
Covers every store method directly with synchronous, zero-React tests:
- `replaceMessages`: population, ordering, undefined handling,
dedup/no-emit
- `upsertDurableMessage`: insert, duplicate detection, value-change
update, optimistic placeholder removal (negative IDs), role-scoped
cleanup, in-place update without reorder
- `setChatStatus`: set, null clear, idempotency
- `setStreamError` / `clearStreamError`: set, clear, no-op guards, dedup
- `setRetryState` / `clearRetryState`: set, clear, no-op guard
- `setSubagentStatusOverride`: single, accumulation, dedup, overwrite
- `setQueuedMessages`: set, undefined handling, ID-based dedup
- `clearStreamState`: clear, no-op guard
- `applyMessagePart` / `applyMessageParts`: text, append, batch, empty
no-op
- `resetTransientState`: clears all transient state, preserves messages,
no-op guard
- `subscribe`: unsubscribe lifecycle, multiple subscribers

### ChatContext.test.tsx — 8 new integration tests
WebSocket event handling that was previously untested:
- **Error events**: sets chatStatus to error, populates streamError,
clears retryState, calls setChatErrorReason; uses fallback when error
has blank text
- **Retry events**: populates retryState; status transition to running
clears retryState
- **Subagent status overrides**: status events with different chat_id go
to subagentStatusOverrides, not main chatStatus
- **WebSocket disconnect**: sets streamError; preserves existing error
on disconnect
- **Status transitions**: clears chatErrorReason on non-error status

### Test infrastructure improvements
- Added `emitError()` helper to MockSocket for testing WebSocket
disconnect
- Added `vi.mocked(watchChat).mockReset()` to `afterEach` for reliable
test isolation between tests that use `mockReturnValueOnce`

## Test results
```
✓ chatStore.test.ts (35 tests) 6ms
✓ ChatContext.test.tsx (23 tests) 107ms
  58 passed (58)
```

---------

Co-authored-by: Coder <coder@users.noreply.github.com>

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) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+		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) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+		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) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+		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) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+		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) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+		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) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+		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) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+		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) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+		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);
+		});
+	});
 });

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>();
 

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);
+	});
+});