feat(ai-chat): resumable SSE — клиент + удаление поллинга/латчей (#381 PR 2) #387

Merged
vvzvlad merged 2 commits from feat/381-resumable-sse-pr2 into feat/381-resumable-sse-pr1 2026-07-06 16:25:12 +03:00
12 changed files with 1094 additions and 929 deletions
@@ -41,17 +41,10 @@ import { extractPageSlugId } from "@/lib";
import {
AI_CHATS_RQ_KEY,
AI_CHAT_MESSAGES_RQ_KEY,
AI_CHAT_RUN_RQ_KEY,
useAiChatMessagesQuery,
useAiChatRunQuery,
useAiChatsQuery,
useAiRolesQuery,
} from "@/features/ai-chat/queries/ai-chat-query.ts";
import {
shouldClearLatchOnQueryError,
shouldClearStoppingLatch,
shouldObserveRun,
} from "@/features/ai-chat/utils/run-polling.ts";
import { workspaceAtom } from "@/features/user/atoms/current-user-atom";
import ConversationList from "@/features/ai-chat/components/conversation-list.tsx";
import ChatThread from "@/features/ai-chat/components/chat-thread.tsx";
@@ -85,6 +78,12 @@ const MIN_HEIGHT = 400;
// Margin kept between the window and the viewport edges while dragging.
const EDGE_MARGIN = 8;
// #184 phase 1.5: hard cap on the degraded-poll fallback. The poll is armed when
// a resume attempt could not attach to the live run and disarmed by the thread on
// settle / local stream; this cap is the ONLY backstop against an endless tick
// (a stuck 'streaming' row before the boot-sweep, or a user-tail 204 with no run).
const DEGRADED_POLL_MAX_MS = 10 * 60_000;
/** Compact token formatter: 1.2M / 3.4k / 950. */
function formatTokens(n: number): string {
if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
@@ -242,150 +241,62 @@ export default function AiChatWindow() {
[roles],
);
// #184 phase 1.5: degraded-poll fallback (replaces the F4/F5/F7 latches). When
// ChatThread could not attach to a still-running run it arms this via
// onResumeFallback(true); the thread disarms it on settle / local stream. The
// window only OWNS the timer (armedAtRef stamps when it was armed for the cap).
const [degradedPoll, setDegradedPoll] = useState(false);
const armedAtRef = useRef(0);
const onResumeFallback = useCallback((active: boolean): void => {
if (active) armedAtRef.current = Date.now();
setDegradedPoll(active);
}, []);
// Reset the degraded poll whenever the open chat changes: it is scoped to the
// resume attempt of the previously-open chat (invariant 8).
useEffect(() => {
setDegradedPoll(false);
}, [activeChatId]);
const { data: messageRows, isLoading: messagesLoading } =
useAiChatMessagesQuery(activeChatId ?? undefined);
useAiChatMessagesQuery(
activeChatId ?? undefined,
// DELIBERATELY DUMB (invariant 8 / task 2.4): poll every 2.5s while armed
// and under the 10-min cap; otherwise off. NO error checks (TanStack v5
// resets fetchFailureCount each fetch, so consecutive errors are not
// expressible — and the poll must survive a server restart) and NO tail
// checks (the settled/local-stream semantics live in ChatThread, which
// disarms via onResumeFallback(false)). The time cap is the only backstop.
() =>
degradedPoll === true &&
Date.now() - armedAtRef.current < DEGRADED_POLL_MAX_MS
? 2500
: false,
);
// #184 reconnect-and-live-follow. Whether detached agent runs are enabled for
// this workspace. The reconnect endpoint itself is NOT flag-gated server-side
// (it is only owner-gated and returns `{ run: null }` when the chat has no
// run); but when the feature is off no runs are ever created, so polling it
// would always come back empty — we gate it off here to avoid pointless polls.
// this workspace. When the feature is off no runs are ever created, so the
// resume attempt would only ever 204; gating ChatThread's resume on it avoids a
// pointless attach round-trip.
const workspace = useAtomValue(workspaceAtom);
const autonomousRunsEnabled =
workspace?.settings?.ai?.autonomousRuns === true;
// Whether THIS tab is the one actively streaming the open chat's run locally
// (it started the run here and holds the SSE). Reported up from ChatThread. We
// are the STREAMER while true and a passive OBSERVER while false — the basis of
// the observer-vs-streamer detection. Reset to false by the fresh ChatThread's
// mount effect on every chat switch.
const [localStreaming, setLocalStreaming] = useState(false);
const onStreamingChange = useCallback((streaming: boolean) => {
setLocalStreaming(streaming);
}, []);
// #184 Stop wiring. While a detached run is being stopped we SUPPRESS the
// observer merge so the stopping run's still-persisting output does not
// re-stream back into view between the moment the user pressed Stop and the run
// actually settling as 'aborted' server-side. Polling itself keeps running (so
// the terminal transition is still detected) — only the visual merge is gated.
// Cleared when the run is observed terminal (below) or the chat is switched.
const [stoppingRun, setStoppingRun] = useState(false);
// Reset the stopping latch whenever the open chat changes: it is scoped to the
// run of the previously-open chat.
useEffect(() => {
setStoppingRun(false);
}, [activeChatId]);
// Authoritative stop of the open chat's detached run (the Stop button in
// autonomous mode). Latch "stopping" first (suppresses the re-stream flash),
// then request the server stop — the ONLY thing that ends a detached run; a mere
// local SSE abort is a client disconnect the server ignores. On failure we
// release the latch so the observer resumes (better to show the live run than to
// freeze the view) and surface the error.
// autonomous mode). Request the server stop — the ONLY thing that ends a
// detached run; a mere local SSE abort is a client disconnect the server
// ignores. On failure surface the error.
const handleServerStop = useCallback(
(chatId: string): void => {
setStoppingRun(true);
// #234 F4: drop the PREVIOUS turn's run from the cache so `run` becomes null
// until the CURRENT turn's run is fetched fresh. Without this, once the local
// stream aborts (localStreaming -> false) the run query re-enables and
// react-query SYNCHRONOUSLY returns the still-cached prior terminal run; the
// terminal effect would then clear the stopping latch against that STALE run
// before the current turn's (still-running, detached, growing) run is ever
// observed — re-opening the observer merge and flashing the growing output
// over the frozen row. With the cache cleared the terminal effect's
// `if (!run) return` holds the latch until the current run itself is observed
// terminal (see shouldClearStoppingLatch).
queryClient.removeQueries({ queryKey: AI_CHAT_RUN_RQ_KEY(chatId) });
void stopRun(chatId).catch(() => {
setStoppingRun(false);
notifications.show({
message: t("Failed to stop the run"),
color: "red",
});
});
},
[t, queryClient],
[t],
);
// Poll the latest run of the open chat ONLY when we are a passive observer:
// feature on, a chat is open, and we are NOT the local streamer (the streamer
// already has the live SSE — polling/merging too would double-render). The
// query's own status-keyed refetchInterval stops once the run is terminal.
const { data: runData, isError: runQueryFailed } = useAiChatRunQuery(
activeChatId ?? undefined,
autonomousRunsEnabled && !localStreaming,
);
const run = runData?.run ?? null;
// Safety net (#234 F4 review): after handleServerStop clears the run cache,
// `run` is null until the current turn's run is fetched fresh, and the terminal
// effect below holds the latch via `if (!run) return`. If that refetch instead
// ERRORS PERMANENTLY (the GET-run keeps failing) while we are no longer the
// streamer, the run stays null, its status-keyed refetchInterval is off, and
// nothing would ever observe a terminal run — freezing the view with the
// observer merge suppressed. Release the latch on that error so the live view
// resumes rather than stays stuck (the local stopRun may already have succeeded
// independently).
//
// #234 F7: this must NOT fire on a TRANSIENT error while `run` is still an
// ACTIVE held run. In TanStack Query v5 (retry:false) the query's `data` is
// RETAINED on error, so `runQueryFailed` can be true while `run` is still
// pending/running — releasing then would re-open the observer merge and flash
// the growing detached run over the frozen row (the very flash F4 prevents). The
// decision is the pure, unit-tested `shouldClearLatchOnQueryError`, which gates
// on the run NOT being active: it cures only the genuine permanent-null-freeze
// (`run === null`) and never releases against an active run.
useEffect(() => {
if (
shouldClearLatchOnQueryError({
stoppingRun,
isLocalStreaming: localStreaming,
runQueryFailed,
run,
})
)
setStoppingRun(false);
}, [stoppingRun, localStreaming, runQueryFailed, run]);
// The run's incrementally-persisted assistant message to merge into the thread,
// but only while we are an observer (never when we are the streamer — guards
// against a stale poll fighting the live stream). Includes a terminal run so the
// final persisted output is shown on reopen.
const observedRow =
shouldObserveRun(run, localStreaming) && !stoppingRun
? (runData?.message ?? null)
: null;
// When the observed run reaches a terminal status, do a final messages refetch
// so the persisted final state (token/context badge, export source) is shown,
// then the query's refetchInterval has already stopped polling. Deduped per run
// id so it fires exactly once per run, not on every subsequent poll-less render.
const finalizedRunIdRef = useRef<string | null>(null);
useEffect(() => {
if (!run || !activeChatId) return;
if (run.status === "pending" || run.status === "running") {
// Active again (a new run) — re-arm so its terminal transition fires once.
finalizedRunIdRef.current = null;
return;
}
// Terminal: a stop we requested has landed (or the run finished on its own),
// so release the stopping latch — the observer merge can now show the final
// persisted (aborted/finished) output without any live re-stream. The decision
// is the pure, unit-tested `shouldClearStoppingLatch` (run-polling.ts): release
// ONLY when we requested a stop, this tab is no longer the streamer, AND the
// CURRENT run is terminal. The #234 F4 cache removal in handleServerStop makes
// `run` null (this branch's `if (!run) return` above holds) until the current
// turn's run is fetched fresh, so the latch can never clear against a stale
// cached run.
if (shouldClearStoppingLatch({ stoppingRun, run, isLocalStreaming: localStreaming }))
setStoppingRun(false);
if (finalizedRunIdRef.current === run.id) return;
finalizedRunIdRef.current = run.id;
queryClient.invalidateQueries({
queryKey: AI_CHAT_MESSAGES_RQ_KEY(activeChatId),
});
}, [run, activeChatId, queryClient, stoppingRun, localStreaming]);
// The page the user is currently viewing. AiChatWindow lives in a pathless
// parent layout route, so useParams() can't see :pageSlug. Match the full
// pathname against the authenticated page route instead so "the current page"
@@ -1034,16 +945,13 @@ export default function AiChatWindow() {
assistantName={currentRole?.name}
onTurnFinished={onTurnFinished}
onServerChatId={onServerChatId}
// #184: live-follow a still-running run when we reopened the chat as
// a passive observer; null when there is nothing to observe or this
// tab is the streamer. onStreamingChange lets the window stop polling
// while we are the streamer.
observedRow={observedRow}
onStreamingChange={onStreamingChange}
// #184 phase 1.5: arm/disarm the degraded-poll fallback when a
// resume attempt could not attach to the live run; the thread
// disarms it on settle / local stream.
onResumeFallback={onResumeFallback}
// #184: in autonomous mode the Stop button must hit the authoritative
// server stop (a local SSE abort is a client disconnect the server
// ignores). onServerStop also arms the "stopping" latch above so the
// stopped run's output does not re-stream via the observer merge.
// ignores).
autonomousRunsEnabled={autonomousRunsEnabled}
onServerStop={handleServerStop}
/>
@@ -1,6 +1,13 @@
import { describe, it, expect, beforeEach, vi } from "vitest";
import { render, screen, fireEvent, act, cleanup } from "@testing-library/react";
import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import {
render,
screen,
fireEvent,
act,
cleanup,
} from "@testing-library/react";
import { MantineProvider } from "@mantine/core";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
// Shared, hoisted mock state so the @ai-sdk/react and "ai" module mocks (hoisted
// above the imports) can expose the captured useChat callbacks / transport and
@@ -12,50 +19,61 @@ const h = vi.hoisted(() => ({
sendMessage: vi.fn(),
stop: vi.fn(),
setMessages: vi.fn(),
resumeStream: vi.fn(),
// The messages array useChat was seeded with (to assert strip/seed behavior).
seededMessages: null as null | unknown[],
transport: null as null | {
prepareSendMessagesRequest: (arg: {
prepareSendMessagesRequest?: (arg: {
messages: unknown[];
body: Record<string, unknown>;
}) => { body: Record<string, unknown> };
prepareReconnectToStreamRequest?: () => { api?: string };
fetch?: (input: unknown, init?: { method?: string }) => Promise<unknown>;
},
},
}));
// Mock useChat: capture onFinish, return the spies and the controllable status.
// Mock useChat: capture onFinish + seeded messages, return the spies and the
// controllable status.
vi.mock("@ai-sdk/react", () => ({
useChat: (opts: { onFinish?: (arg: Record<string, unknown>) => void }) => {
useChat: (opts: {
messages?: unknown[];
onFinish?: (arg: Record<string, unknown>) => void;
}) => {
h.state.onFinish = opts.onFinish ?? null;
h.state.seededMessages = opts.messages ?? null;
return {
messages: [],
sendMessage: h.state.sendMessage,
status: h.state.status,
stop: h.state.stop,
error: null,
// #184: ChatThread reads setMessages to merge a polled observer run.
setMessages: h.state.setMessages,
resumeStream: h.state.resumeStream,
};
},
}));
// Mock "ai": deterministic ids + a transport that records its options so the test
// can invoke prepareSendMessagesRequest and assert the `interrupted` flag.
// can invoke prepareSendMessagesRequest / prepareReconnectToStreamRequest / fetch.
vi.mock("ai", () => {
let counter = 0;
return {
generateId: () => `gid-${counter++}`,
DefaultChatTransport: class {
constructor(opts: {
prepareSendMessagesRequest: (arg: {
messages: unknown[];
body: Record<string, unknown>;
}) => { body: Record<string, unknown> };
}) {
h.state.transport = opts;
constructor(opts: Record<string, unknown>) {
h.state.transport = opts as never;
}
},
};
});
// Keep the ai-chat-query import light: ChatThread only needs the messages RQ key,
// so stub the module to avoid pulling axios / i18n transitively.
vi.mock("@/features/ai-chat/queries/ai-chat-query.ts", () => ({
AI_CHAT_MESSAGES_RQ_KEY: (chatId: string) => ["ai-chat-messages", chatId],
}));
// Stub the heavy children: MessageList (markdown/render) and ChatInput (the
// composer). The ChatInput stub exposes a button that queues a message, the only
// interaction this test needs to populate the queue while "streaming".
@@ -63,49 +81,90 @@ vi.mock("@/features/ai-chat/components/message-list.tsx", () => ({
default: () => <div data-testid="message-list" />,
}));
vi.mock("@/features/ai-chat/components/chat-input.tsx", () => ({
default: ({ onQueue }: { onQueue: (text: string) => void }) => (
<button data-testid="queue-btn" onClick={() => onQueue("queued text")}>
queue
</button>
default: ({
onQueue,
onStop,
}: {
onQueue: (text: string) => void;
onStop: () => void;
}) => (
<>
<button data-testid="queue-btn" onClick={() => onQueue("queued text")}>
queue
</button>
<button aria-label="Stop" onClick={() => onStop()}>
stop
</button>
</>
),
}));
import ChatThread from "./chat-thread";
import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.ts";
function renderThread() {
function row(
id: string,
role: string,
status?: string,
text = "",
): IAiChatMessageRow {
return { id, role, content: text, status, createdAt: "2026-01-01T00:00:00Z" };
}
function renderThread(props?: {
chatId?: string | null;
initialRows?: IAiChatMessageRow[];
autonomousRunsEnabled?: boolean;
}) {
const onTurnFinished = vi.fn();
render(
<MantineProvider>
<ChatThread chatId="c1" initialRows={[]} onTurnFinished={onTurnFinished} />
</MantineProvider>,
const onResumeFallback = vi.fn();
const onServerStop = vi.fn();
const queryClient = new QueryClient({
defaultOptions: { queries: { retry: false } },
});
const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries");
const { unmount } = render(
<QueryClientProvider client={queryClient}>
<MantineProvider>
<ChatThread
chatId={props?.chatId === undefined ? "c1" : props.chatId}
initialRows={props?.initialRows ?? []}
autonomousRunsEnabled={props?.autonomousRunsEnabled}
onTurnFinished={onTurnFinished}
onResumeFallback={onResumeFallback}
onServerStop={onServerStop}
/>
</MantineProvider>
</QueryClientProvider>,
);
return { onTurnFinished };
return { onTurnFinished, onResumeFallback, onServerStop, invalidateSpy, unmount };
}
function resetState() {
h.state.status = "streaming";
h.state.onFinish = null;
h.state.seededMessages = null;
h.state.transport = null;
h.state.sendMessage.mockClear();
h.state.stop.mockClear();
h.state.setMessages.mockClear();
h.state.resumeStream.mockClear();
}
describe("ChatThread — send now (#198)", () => {
beforeEach(() => {
h.state.status = "streaming";
h.state.onFinish = null;
h.state.sendMessage.mockClear();
h.state.stop.mockClear();
h.state.transport = null;
});
beforeEach(resetState);
it("aborts the current turn and resends the queued message on the abort", () => {
renderThread();
// Queue a message while the turn is streaming.
fireEvent.click(screen.getByTestId("queue-btn"));
const sendNowBtn = screen.getByLabelText("Send now");
expect(sendNowBtn).toBeTruthy();
// "Send now" interrupts the current turn (stop), but does NOT send yet —
// the resend happens once the abort lands in onFinish.
fireEvent.click(sendNowBtn);
expect(h.state.stop).toHaveBeenCalledTimes(1);
expect(h.state.sendMessage).not.toHaveBeenCalled();
// The abort we triggered reaches onFinish: the promoted head is flushed.
act(() => {
h.state.onFinish?.({
message: { id: "a", role: "assistant", parts: [] },
@@ -122,10 +181,8 @@ describe("ChatThread — send now (#198)", () => {
fireEvent.click(screen.getByTestId("queue-btn"));
fireEvent.click(screen.getByLabelText("Send now"));
const prep = h.state.transport!.prepareSendMessagesRequest;
// The send right after "send now" carries interrupted: true...
const prep = h.state.transport!.prepareSendMessagesRequest!;
expect(prep({ messages: [], body: {} }).body.interrupted).toBe(true);
// ...and only that one (the flag is read-and-cleared).
expect(prep({ messages: [], body: {} }).body.interrupted).toBe(false);
});
@@ -136,42 +193,24 @@ describe("ChatThread — send now (#198)", () => {
fireEvent.click(screen.getByTestId("queue-btn"));
fireEvent.click(screen.getByLabelText("Send now"));
// No turn to interrupt: sent straight away, no abort, not flagged.
expect(h.state.stop).not.toHaveBeenCalled();
expect(h.state.sendMessage).toHaveBeenCalledWith({ text: "queued text" });
const prep = h.state.transport!.prepareSendMessagesRequest;
const prep = h.state.transport!.prepareSendMessagesRequest!;
expect(prep({ messages: [], body: {} }).body.interrupted).toBe(false);
});
});
// The turn-end decision lives in the `onFinish` handler: given the terminal
// outcome of a turn (`isAbort` / `isDisconnect` / `isError`, or none = clean),
// it decides whether to CONTINUE (flush the next queued message) or END (leave
// the queue intact for the user), and which stop notice — if any — to show.
// `sendNow` is exercised above; these tests pin down the plain outcomes.
describe("ChatThread — turn-end decision (onFinish)", () => {
beforeEach(() => {
h.state.status = "streaming";
h.state.onFinish = null;
h.state.sendMessage.mockClear();
h.state.stop.mockClear();
h.state.transport = null;
});
beforeEach(resetState);
// Drive a fresh onFinish with the given terminal flags after queueing a
// message, and report both what the parent was told and whether the queue was
// flushed (a resend to the sendMessage spy).
function finishWith(flags: {
isAbort?: boolean;
isDisconnect?: boolean;
isError?: boolean;
}) {
// Tear down any prior render so the loop-driven "every outcome" case does
// not leave duplicate queue buttons in the DOM.
cleanup();
h.state.sendMessage.mockClear();
const { onTurnFinished } = renderThread();
// Populate the queue while the turn is streaming.
fireEvent.click(screen.getByTestId("queue-btn"));
act(() => {
h.state.onFinish?.({
@@ -187,16 +226,12 @@ describe("ChatThread — turn-end decision (onFinish)", () => {
it("CONTINUES — flushes the next queued message on a clean finish", () => {
finishWith({});
// Clean finish (no terminal flag): the queued message is auto-sent.
expect(h.state.sendMessage).toHaveBeenCalledWith({ text: "queued text" });
// A clean finish shows no stop notice.
expect(screen.queryByText("Response stopped.")).toBeNull();
});
it("ENDS — keeps the queue intact on a user abort and shows the stopped notice", () => {
finishWith({ isAbort: true });
// A plain Stop (not the sendNow interrupt path) must NOT auto-resend: the
// queue is preserved for the user to decide.
expect(h.state.sendMessage).not.toHaveBeenCalled();
expect(screen.getByText("Response stopped.")).toBeTruthy();
});
@@ -211,15 +246,11 @@ describe("ChatThread — turn-end decision (onFinish)", () => {
it("ENDS — keeps the queue intact on a stream error (no auto-retry, no stopped notice)", () => {
finishWith({ isError: true });
// Blindly retrying after a failure would be wrong; the queue is left alone.
expect(h.state.sendMessage).not.toHaveBeenCalled();
// isError clears the neutral notice (the error banner covers this case).
expect(screen.queryByText("Response stopped.")).toBeNull();
});
it("notifies the parent on EVERY terminal outcome", () => {
// The chat-list refresh / new-chat id adoption must run on success and on
// every failure path alike.
for (const flags of [
{},
{ isAbort: true },
@@ -232,55 +263,411 @@ describe("ChatThread — turn-end decision (onFinish)", () => {
});
});
// #184 passive-observer merge: when reconnecting to a still-running run, the
// parent feeds the polled run message via `observedRow`; ChatThread merges it via
// setMessages — but ONLY when this tab is NOT itself streaming (the streamer's
// SSE owns the view, so a stale observedRow must never overwrite it).
describe("ChatThread — observer run merge (#184)", () => {
beforeEach(() => {
h.state.onFinish = null;
h.state.setMessages.mockReset();
// #184 phase 1.5: the resumable-SSE client. A reopened tab resumes the live run
// via the SDK's reconnect transport (attach: replay + tail) instead of polling.
describe("ChatThread — resume (attach) machinery (#184)", () => {
const streamingTail = () => [
row("u1", "user", undefined, "hi"),
row("a1", "assistant", "streaming", "partial"),
];
const settledTail = () => [
row("u1", "user", undefined, "hi"),
row("a1", "assistant", "succeeded", "done"),
];
const userTail = () => [row("u1", "user", undefined, "hi")];
const visibleMsg = {
id: "a1",
role: "assistant",
parts: [{ type: "text", text: "streamed answer" }],
};
const emptyMsg = { id: "a1", role: "assistant", parts: [] };
beforeEach(resetState);
// NOTE: do NOT vi.unstubAllGlobals() here — vitest.setup.ts installs
// matchMedia/localStorage via vi.stubGlobal and unstubbing wipes them for the
// rest of the file. Fetch is re-stubbed per test that needs it.
afterEach(cleanup);
it("resumes on mount only when the flag is on, chatId is set, and the tail is not a settled assistant", () => {
// streaming tail -> resume
renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() });
expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
// user tail -> resume (the assistant row may not be seeded yet)
cleanup();
h.state.resumeStream.mockClear();
renderThread({ autonomousRunsEnabled: true, initialRows: userTail() });
expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
// settled assistant tail -> NO resume
cleanup();
h.state.resumeStream.mockClear();
renderThread({ autonomousRunsEnabled: true, initialRows: settledTail() });
expect(h.state.resumeStream).not.toHaveBeenCalled();
// flag off -> NO resume
cleanup();
h.state.resumeStream.mockClear();
renderThread({ autonomousRunsEnabled: false, initialRows: streamingTail() });
expect(h.state.resumeStream).not.toHaveBeenCalled();
// no chatId -> NO resume
cleanup();
h.state.resumeStream.mockClear();
renderThread({
autonomousRunsEnabled: true,
chatId: null,
initialRows: streamingTail(),
});
expect(h.state.resumeStream).not.toHaveBeenCalled();
});
const observedRow = {
id: "a-run",
role: "assistant",
content: "step 1\nstep 2",
metadata: {
parts: [{ type: "text", text: "step 1\nstep 2" }],
},
createdAt: "2026-01-01T00:00:00Z",
} as const;
it("strips the streaming tail from the seed, but keeps a user tail whole", () => {
renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() });
// 2 rows in, streaming tail stripped -> 1 seeded message.
expect(h.state.seededMessages).toHaveLength(1);
function renderObserver(status: string) {
h.state.status = status;
render(
cleanup();
renderThread({ autonomousRunsEnabled: true, initialRows: userTail() });
// user tail is not stripped.
expect(h.state.seededMessages).toHaveLength(1);
});
it("builds the attach URL with expect=live&anchor only when the streaming tail was stripped", () => {
renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() });
expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe(
"/api/ai-chat/runs/c1/stream?expect=live&anchor=a1",
);
cleanup();
renderThread({ autonomousRunsEnabled: true, initialRows: userTail() });
expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe(
"/api/ai-chat/runs/c1/stream",
);
});
async function fetch204() {
vi.stubGlobal(
"fetch",
vi.fn().mockResolvedValue({ status: 204, ok: false }),
);
await act(async () => {
await h.state.transport!.fetch!("http://x", { method: "GET" });
});
}
it("204 on a user tail: no crash, no restore, reconcile+invalidate, onResumeFallback(true)", async () => {
const { onResumeFallback, invalidateSpy } = renderThread({
autonomousRunsEnabled: true,
initialRows: userTail(),
});
await fetch204();
// No stripped row -> no restore merge.
expect(h.state.setMessages).not.toHaveBeenCalled();
expect(invalidateSpy).toHaveBeenCalledWith({
queryKey: ["ai-chat-messages", "c1"],
});
expect(onResumeFallback).toHaveBeenCalledWith(true);
});
it("204 on a streaming tail: restore + invalidate + onResumeFallback(true)", async () => {
const { onResumeFallback, invalidateSpy } = renderThread({
autonomousRunsEnabled: true,
initialRows: streamingTail(),
});
await fetch204();
// Stripped row is restored to the store.
expect(h.state.setMessages).toHaveBeenCalledTimes(1);
expect(invalidateSpy).toHaveBeenCalledWith({
queryKey: ["ai-chat-messages", "c1"],
});
expect(onResumeFallback).toHaveBeenCalledWith(true);
});
it("F7 restart-survival: a 500 attach failure restores the stripped row AND arms the poll (not lost)", async () => {
const { onResumeFallback, invalidateSpy } = renderThread({
autonomousRunsEnabled: true,
initialRows: streamingTail(),
});
vi.stubGlobal(
"fetch",
vi.fn().mockResolvedValue({ status: 500, ok: false }),
);
await act(async () => {
await h.state.transport!.fetch!("http://x", { method: "GET" });
});
expect(h.state.setMessages).toHaveBeenCalledTimes(1); // stripped row restored
expect(invalidateSpy).toHaveBeenCalledWith({
queryKey: ["ai-chat-messages", "c1"],
});
expect(onResumeFallback).toHaveBeenCalledWith(true); // degraded poll armed
});
it("F7 restart-survival: a network throw restores the stripped row AND arms the poll", async () => {
const { onResumeFallback, invalidateSpy } = renderThread({
autonomousRunsEnabled: true,
initialRows: streamingTail(),
});
vi.stubGlobal(
"fetch",
vi.fn().mockRejectedValue(new Error("network down")),
);
await act(async () => {
await h.state
.transport!.fetch!("http://x", { method: "GET" })
.catch(() => undefined); // the wrapper rethrows; swallow here
});
expect(h.state.setMessages).toHaveBeenCalledTimes(1);
expect(invalidateSpy).toHaveBeenCalledWith({
queryKey: ["ai-chat-messages", "c1"],
});
expect(onResumeFallback).toHaveBeenCalledWith(true);
});
it("unmount during a pending attach aborts the controller and gates late callbacks", async () => {
const { onResumeFallback, invalidateSpy, unmount } = renderThread({
autonomousRunsEnabled: true,
initialRows: streamingTail(),
});
let abortSeen = false;
let resolveFetch!: (v: unknown) => void;
vi.stubGlobal(
"fetch",
vi.fn().mockImplementation((_input: unknown, init: RequestInit) => {
init.signal?.addEventListener("abort", () => {
abortSeen = true;
});
return new Promise((res) => {
resolveFetch = res;
});
}),
);
// Kick a reconnect GET (stays pending).
let pending!: Promise<unknown>;
act(() => {
pending = h.state.transport!.fetch!("http://x", { method: "GET" });
});
// Unmount: the cleanup aborts the in-flight attach.
unmount();
expect(abortSeen).toBe(true);
// A late 204 landing after unmount must NOT arm a poll / invalidate the (now
// different) chat.
onResumeFallback.mockClear();
invalidateSpy.mockClear();
await act(async () => {
resolveFetch({ status: 204, ok: false });
await pending;
});
expect(onResumeFallback).not.toHaveBeenCalledWith(true);
expect(invalidateSpy).not.toHaveBeenCalled();
});
it("a resume fetch error clears resumedTurn so the next local turn flushes the queue", async () => {
renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() });
h.state.status = "ready";
vi.stubGlobal(
"fetch",
vi.fn().mockResolvedValue({ status: 500, ok: false }),
);
await act(async () => {
await h.state.transport!.fetch!("http://x", { method: "GET" });
});
// Queue then clean-finish: suppression was cleared, so the queue flushes.
fireEvent.click(screen.getByTestId("queue-btn"));
act(() => {
h.state.onFinish?.({
message: visibleMsg,
isAbort: false,
isDisconnect: false,
isError: false,
});
});
expect(h.state.sendMessage).toHaveBeenCalledWith({ text: "queued text" });
});
it("a resumed turn's onFinish does NOT flush the queue", () => {
renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() });
fireEvent.click(screen.getByTestId("queue-btn"));
act(() => {
h.state.onFinish?.({
message: visibleMsg,
isAbort: false,
isDisconnect: false,
isError: false,
});
});
expect(h.state.sendMessage).not.toHaveBeenCalled();
});
it("a healthy resumed finish (visible content) arms nothing and keeps the store", () => {
h.state.status = "ready";
const { onResumeFallback } = renderThread({
autonomousRunsEnabled: true,
initialRows: streamingTail(),
});
h.state.setMessages.mockClear();
onResumeFallback.mockClear();
act(() => {
h.state.onFinish?.({
message: visibleMsg,
isAbort: false,
isDisconnect: false,
isError: false,
});
});
// No restore (would clobber the fuller streamed message), no poll arm.
expect(h.state.setMessages).not.toHaveBeenCalled();
expect(onResumeFallback).not.toHaveBeenCalledWith(true);
});
it("isDisconnect WITH visible content arms the poll but does NOT restore", () => {
h.state.status = "ready";
const { onResumeFallback, invalidateSpy } = renderThread({
autonomousRunsEnabled: true,
initialRows: streamingTail(),
});
h.state.setMessages.mockClear();
onResumeFallback.mockClear();
invalidateSpy.mockClear();
act(() => {
h.state.onFinish?.({
message: visibleMsg,
isAbort: false,
isDisconnect: true,
isError: false,
});
});
expect(onResumeFallback).toHaveBeenCalledWith(true);
expect(invalidateSpy).toHaveBeenCalledWith({
queryKey: ["ai-chat-messages", "c1"],
});
// Restore forbidden: the on-screen partial must not roll back.
expect(h.state.setMessages).not.toHaveBeenCalled();
});
it("an empty resumed message (starved replay) restores the stripped row AND arms the poll", () => {
h.state.status = "ready";
const { onResumeFallback } = renderThread({
autonomousRunsEnabled: true,
initialRows: streamingTail(),
});
h.state.setMessages.mockClear();
onResumeFallback.mockClear();
act(() => {
h.state.onFinish?.({
message: emptyMsg,
isAbort: false,
isDisconnect: false,
isError: false,
});
});
expect(h.state.setMessages).toHaveBeenCalledTimes(1); // restore
expect(onResumeFallback).toHaveBeenCalledWith(true); // arm
});
it("degraded-merge: merges the tail per initialRows update, and settles disarm the poll", async () => {
h.state.status = "ready";
const { rerender, onResumeFallback } = renderResumable(streamingTail());
// Arm reconcile via a 204.
vi.stubGlobal(
"fetch",
vi.fn().mockResolvedValue({ status: 204, ok: false }),
);
await act(async () => {
await h.state.transport!.fetch!("http://x", { method: "GET" });
});
h.state.setMessages.mockClear();
onResumeFallback.mockClear();
// A streaming-tail update: merge, poll stays armed.
rerender([
row("u1", "user", undefined, "hi"),
row("a1", "assistant", "streaming", "step 1\nstep 2"),
]);
expect(h.state.setMessages).toHaveBeenCalledTimes(1);
expect(onResumeFallback).not.toHaveBeenCalledWith(false);
// A settled-tail update: merge + disarm.
h.state.setMessages.mockClear();
rerender([
row("u1", "user", undefined, "hi"),
row("a1", "assistant", "succeeded", "final"),
]);
expect(h.state.setMessages).toHaveBeenCalledTimes(1);
expect(onResumeFallback).toHaveBeenCalledWith(false);
});
it("a local stream disarms both the merge and the poll", () => {
h.state.status = "streaming";
const { rerender, onResumeFallback } = renderResumable(streamingTail());
onResumeFallback.mockClear();
// A re-render while streaming: the reconciliation effect disarms.
rerender(streamingTail());
expect(onResumeFallback).toHaveBeenCalledWith(false);
});
it("Send now is hidden on a resumed turn but visible on a local stream", () => {
// Resumed turn: hidden.
renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() });
fireEvent.click(screen.getByTestId("queue-btn"));
expect(screen.queryByLabelText("Send now")).toBeNull();
// Local streaming turn (no resume): visible.
cleanup();
resetState();
renderThread({ initialRows: [] });
fireEvent.click(screen.getByTestId("queue-btn"));
expect(screen.getByLabelText("Send now")).toBeTruthy();
});
it("handleStop aborts the attach controller and calls onServerStop", async () => {
const { onServerStop } = renderThread({
autonomousRunsEnabled: true,
initialRows: streamingTail(),
});
// Establish an attach controller via a (pending) reconnect GET.
let abortSeen = false;
vi.stubGlobal(
"fetch",
vi.fn().mockImplementation((_input: unknown, init: RequestInit) => {
init.signal?.addEventListener("abort", () => {
abortSeen = true;
});
return new Promise(() => undefined); // never resolves
}),
);
act(() => {
void h.state.transport!.fetch!("http://x", { method: "GET" });
});
fireEvent.click(screen.getByLabelText("Stop"));
expect(abortSeen).toBe(true);
expect(onServerStop).toHaveBeenCalledWith("c1");
});
});
// Helper: render a resumable thread and expose a rerender that only swaps
// initialRows (the degraded-merge effect depends on it).
function renderResumable(initialRows: IAiChatMessageRow[]) {
const onResumeFallback = vi.fn();
const queryClient = new QueryClient({
defaultOptions: { queries: { retry: false } },
});
const Wrapper = ({ rows }: { rows: IAiChatMessageRow[] }) => (
<QueryClientProvider client={queryClient}>
<MantineProvider>
<ChatThread
chatId="c1"
initialRows={[]}
initialRows={rows}
autonomousRunsEnabled
onTurnFinished={vi.fn()}
observedRow={observedRow as never}
onResumeFallback={onResumeFallback}
/>
</MantineProvider>,
);
}
it("merges the polled run message when this tab is a passive observer", () => {
renderObserver("ready");
expect(h.state.setMessages).toHaveBeenCalledTimes(1);
// The updater replaces/append the observed assistant row by id.
const updater = h.state.setMessages.mock.calls[0][0] as (
prev: { id: string; parts: { text: string }[] }[],
) => { id: string; parts: { text: string }[] }[];
const merged = updater([{ id: "u1", parts: [{ text: "hi" }] }]);
expect(merged).toHaveLength(2);
expect(merged[1].id).toBe("a-run");
expect(merged[1].parts[0].text).toBe("step 1\nstep 2");
});
it("does NOT merge while THIS tab is the streamer (no double-render)", () => {
renderObserver("streaming");
expect(h.state.setMessages).not.toHaveBeenCalled();
});
});
</MantineProvider>
</QueryClientProvider>
);
const view = render(<Wrapper rows={initialRows} />);
const rerender = (rows: IAiChatMessageRow[]) =>
act(() => view.rerender(<Wrapper rows={rows} />));
return { rerender, onResumeFallback };
}
@@ -1,4 +1,5 @@
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { useQueryClient } from "@tanstack/react-query";
import { generateId } from "ai";
import { ActionIcon, Box, Group, Stack, Text, Tooltip } from "@mantine/core";
import {
@@ -24,7 +25,14 @@ import {
} from "@/features/ai-chat/utils/role-launch.ts";
import { describeChatError } from "@/features/ai-chat/utils/error-message.ts";
import { extractServerChatId } from "@/features/ai-chat/utils/adopt-chat-id.ts";
import { mergeObservedMessage } from "@/features/ai-chat/utils/run-polling.ts";
import { assistantMessageHasVisibleContent } from "@/features/ai-chat/utils/message-content.ts";
import {
isStreamingTail,
isSettledAssistantTail,
seedRows,
mergeById,
} from "@/features/ai-chat/utils/resume-helpers.ts";
import { AI_CHAT_MESSAGES_RQ_KEY } from "@/features/ai-chat/queries/ai-chat-query.ts";
import {
dequeue,
enqueueMessage,
@@ -87,19 +95,13 @@ interface ChatThreadProps {
* Copy/export button available mid-stream). Distinct from onTurnFinished,
* which fires only at the terminal outcome. */
onServerChatId?: (serverChatId?: string) => void;
/** #184 reconnect-and-live-follow. When THIS tab reopened a chat whose agent
* run is still going (it is a PASSIVE OBSERVER — it did not start the run here),
* the parent polls the reconnect endpoint and feeds the run's incrementally-
* persisted assistant message here; we merge it into the live list so new
* steps/tool-calls appear as they are persisted. Null when there is nothing to
* observe (no run, feature off, or this tab IS the streamer). The merge is
* ADDITIONALLY guarded by our own `isStreaming`, so a stale value can never
* fight the local stream when we are the streamer. */
observedRow?: IAiChatMessageRow | null;
/** Report this tab's live streaming status up to the parent, so it can stop
* polling the run while WE are the active streamer (the SSE owns the view) and
* resume once we go idle. Called from an effect on every transition. */
onStreamingChange?: (streaming: boolean) => void;
/** #184 phase 1.5: arm/disarm the parent's degraded-poll fallback for THIS
* chat's window. Called `true` when a resume attempt could not attach to the
* live run (attach 204 / starved-or-torn resumed finish), so the window starts
* a dumb timed poll of the message history to follow the detached run to settle;
* called `false` the moment a local stream starts or the terminal settled row is
* merged (invariant 8). The window owns the timer + its 10-min cap. */
onResumeFallback?: (active: boolean) => void;
/** #184: whether detached/autonomous agent runs are enabled for this workspace.
* When true the Stop button must additionally hit the AUTHORITATIVE server stop
* (via onServerStop) — aborting only the local SSE is just a client disconnect,
@@ -155,15 +157,58 @@ export default function ChatThread({
assistantName,
onTurnFinished,
onServerChatId,
observedRow,
onStreamingChange,
onResumeFallback,
autonomousRunsEnabled,
onServerStop,
}: ChatThreadProps) {
const { t } = useTranslation();
const queryClient = useQueryClient();
// resume machinery refs (#184 phase 1.5)
const attachAbortRef = useRef<AbortController | null>(null);
const reconcileTailRef = useRef(false);
const noStreamHandledRef = useRef(false);
const onNoActiveStreamRef = useRef<(() => void) | null>(null);
// Live mount flag. The attach GET and the resumed `onFinish` are async and can
// land AFTER this thread unmounts (the parent remounts per chat via `key`); with
// chatIdRef then pointing at the NEW chat, an ungated late callback would arm a
// spurious poll + foreign invalidation on the newly-opened chat. Every parent-
// facing resume side-effect is gated on this.
const mountedRef = useRef(true);
const [resumedTurn, setResumedTurn] = useState(false);
const resumedTurnRef = useRef(false);
// Identity-stable pair setter (bare useState setter + ref write): it is closed
// over by the transport useMemo([]), so it MUST NOT capture state.
const setResumedTurnPair = useCallback((v: boolean) => {
resumedTurnRef.current = v;
setResumedTurn(v);
}, []);
// Mount-time resume gating (in refs — computed once for this mount; the parent
// remounts per chat via `key`).
//
// Attempt resume for any non-settled tail: a streaming tail (strip + expect
// live replay) or a user tail (the run may exist but its assistant row is not
// seeded yet — attach to the pre-opened registry entry and wait for frames).
// A settled assistant tail must NEVER resume: replaying a finished run into a
// store that already contains its message duplicates parts (SDK text-start
// always pushes a new part).
const stripRef = useRef(chatId !== null && isStreamingTail(initialRows ?? []));
const attemptResumeRef = useRef(
autonomousRunsEnabled === true &&
chatId !== null &&
!isSettledAssistantTail(initialRows ?? []),
);
const strippedRowRef = useRef<IAiChatMessageRow | null>(
stripRef.current ? (initialRows ?? [])[initialRows!.length - 1] : null,
);
const initialMessages = useMemo<UIMessage[]>(
() => (initialRows ?? []).map(rowToUiMessage),
() =>
seedRows(
initialRows ?? [],
attemptResumeRef.current && stripRef.current,
).map(rowToUiMessage),
[initialRows],
);
@@ -261,9 +306,12 @@ export default function ChatThread({
const { head, rest } = dequeue(queuedRef.current);
if (!head) return false;
setQueue(rest);
// Local send: clear any resume-suppression flag so this genuine local turn's
// onFinish flushes normally (invariant 8).
setResumedTurnPair(false);
sendMessageRef.current?.({ text: head.text });
return true;
}, [setQueue]);
}, [setQueue, setResumedTurnPair]);
const enqueue = useCallback(
(text: string) => {
@@ -283,6 +331,47 @@ export default function ChatThread({
new DefaultChatTransport<UIMessage>({
api: "/api/ai-chat/stream",
credentials: "include",
prepareReconnectToStreamRequest: () => ({
// SDK default URL uses the useChat STORE id — always build from the real chat id.
// ?expect=live&anchor=<row id> ONLY when we stripped a streaming tail: expect=live
// is the only case where a finished-retained replay is safe (the row is stripped,
// replay rebuilds it), and the anchor pins the replay to OUR run — a mismatching
// (newer) run must 204 into the restore+poll path instead of replaying a foreign
// transcript into this store.
api: `/api/ai-chat/runs/${chatIdRef.current}/stream${
stripRef.current
? `?expect=live&anchor=${strippedRowRef.current!.id}`
: ""
}`,
}),
fetch: async (input: RequestInfo | URL, init: RequestInit = {}) => {
if ((init.method ?? "GET") !== "GET") return fetch(input, init); // send path untouched
// Reconnect GET: the SDK passes no AbortSignal, so wire our own controller
// for observer Stop / unmount abort.
const controller = new AbortController();
attachAbortRef.current = controller;
try {
const response = await fetch(input, {
...init,
signal: controller.signal,
});
// No onFinish will come for a 204 (silent no-op) OR any non-2xx
// (5xx/502 — a server restart mid-attach). Both run the same
// no-active-stream recovery: restore the stripped row, invalidate, and
// arm the degraded poll (idempotent via noStreamHandledRef; its part-d
// also clears the resumedTurn flag). This is the restart-survival path
// the removed F7 latch used to guard — a transient attach failure must
// NOT drop the in-progress row or stop tracking the durable run.
if (response.status === 204 || !response.ok)
onNoActiveStreamRef.current?.();
return response;
} catch (err) {
// Network throw: same no-onFinish recovery, then rethrow so the SDK
// still surfaces the error to its own machinery.
onNoActiveStreamRef.current?.();
throw err;
}
},
// Inject the chat id and the currently-open page alongside the useChat
// messages so the server can resolve an existing chat (or create one
// when null) and tell the agent which page "this page" refers to. Both
@@ -312,7 +401,15 @@ export default function ChatThread({
[],
);
const { messages, sendMessage, status, stop, error, setMessages } = useChat({
const {
messages,
sendMessage,
status,
stop,
error,
setMessages,
resumeStream,
} = useChat({
// Stable per-mount key. Existing chats use their real id; new chats use a
// generated client id (never `undefined`) so the store is NOT re-created on
// every render mid-stream (see `chatStoreId` above).
@@ -330,6 +427,38 @@ export default function ChatThread({
// would be wrong, so on Stop/disconnect/error the queue is left intact for
// the user to decide.
onFinish: ({ message, isAbort, isDisconnect, isError }) => {
// (1) Capture whether THIS finish belongs to a resumed (attach) turn and
// immediately clear the flag so it can never suppress a LATER local turn.
const wasResumed = resumedTurnRef.current;
setResumedTurnPair(false);
// (2) Recovery after a starved/torn resumed finish (invariant 9). The arm
// and the stripped-row restore are gated DIFFERENTLY. Skip entirely once
// unmounted (an abort-triggered onFinish landing after a chat switch must
// not arm a poll / invalidate on the new chat).
if (wasResumed && mountedRef.current) {
const hasVisibleContent = assistantMessageHasVisibleContent(message);
// ARM the reconcile + degraded poll when the resumed message carries no
// visible content (starved replay) OR the connection dropped mid-run — in
// both cases the poll must drive the row to its real terminal state.
if (isDisconnect || !hasVisibleContent) {
reconcileTailRef.current = true;
queryClient.invalidateQueries({
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current),
});
onResumeFallback?.(true);
}
// RESTORE the stripped streaming row ONLY when the resumed message has no
// visible content. On isDisconnect WITH visible content restore is
// FORBIDDEN: the live stream may have advanced far past the mount-time
// snapshot, so restoring would clobber on-screen content (invariant 9) —
// the arm above suffices, the poll reaches the true terminal.
if (!hasVisibleContent && strippedRowRef.current) {
setMessages((prev) =>
mergeById(prev, rowToUiMessage(strippedRowRef.current!)),
);
}
}
// (3) Standard branches.
// Forward the authoritative server chatId (streamed on the assistant
// message metadata) so the parent adopts the REAL created chat id for a new
// chat — see adopt-chat-id.ts for the full #137 design. `threadKey` lets the
@@ -342,6 +471,10 @@ export default function ChatThread({
else if (isAbort) setStopNotice("manual");
else if (isDisconnect) setStopNotice("disconnect");
else setStopNotice(null);
// A resumed turn NEVER flushes the queue (invariant 7): skip BOTH the
// flush-on-abort branch and the plain flush. The local streamer is the only
// tab that owns the queue.
if (wasResumed) return;
// "Send now": WE triggered this abort to interrupt the current turn and
// immediately send the promoted head. Flush it even though the turn was
// aborted (the normal abort path below keeps the queue intact). The
@@ -423,26 +556,98 @@ export default function ChatThread({
const isStreaming = status === "submitted" || status === "streaming";
// #184: report our live streaming status up so the parent stops polling the run
// while WE are the streamer (the SSE owns the view) and resumes once we go idle.
// Effect (not render) so it never updates parent state during our own render;
// fires on mount with `false`, which also re-syncs the parent after a chat
// switch remounts this thread (a fresh mount is idle until the user sends).
useEffect(() => {
onStreamingChange?.(isStreaming);
}, [isStreaming, onStreamingChange]);
// 204-handler (`onNoActiveStream`): the attach returned 204 — nothing live to
// resume (overflow / begin-failure / after retention / anchor-mismatch). One-
// shot via noStreamHandledRef (we do NOT null onNoActiveStreamRef). Exactly four
// parts. Kept in a ref (read by the transport's fetch closure) and refreshed
// each render below.
const onNoActiveStream = useCallback(() => {
// A late attach outcome after unmount must not arm a poll / invalidate on the
// now-different chat this thread's refs were reused for.
if (!mountedRef.current) return;
if (noStreamHandledRef.current) return;
noStreamHandledRef.current = true;
// (a) Restore the stripped streaming row to the store — ONLY when we actually
// stripped one (a user-tail 204 does NOT reach here with a stripped row, so do
// not dereference null).
if (strippedRowRef.current) {
setMessages((prev) =>
mergeById(prev, rowToUiMessage(strippedRowRef.current!)),
);
}
// (b) Reconcile the tail from the message history + invalidate it so the
// degraded poll starts from a fresh fetch.
reconcileTailRef.current = true;
queryClient.invalidateQueries({
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current),
});
// (c) Arm the degraded poll (a dumb timer with a 10-min cap in the window);
// the thread disarms it via onResumeFallback(false) on settle / local stream.
onResumeFallback?.(true);
// (d) 204 means onFinish will NOT fire — clear the suppression flag so it
// cannot swallow the NEXT local turn's queue flush.
setResumedTurnPair(false);
}, [setMessages, queryClient, onResumeFallback, setResumedTurnPair]);
onNoActiveStreamRef.current = onNoActiveStream;
// #184 passive-observer merge: when the parent feeds a polled run message (we
// reopened a chat whose run is still going and did NOT start it here), merge it
// into the live list so new steps/tool-calls appear as they are persisted. Hard-
// gated by `!isStreaming`: if THIS tab is actually the streamer, the local SSE
// owns the view and a stale observedRow must never overwrite it. `observedRow`
// is a stable per-poll object, so this runs once per poll, not per render.
// Mount effect: kick off the resume attempt for a non-settled tail. Marking the
// turn as resumed BEFORE resumeStream so onFinish (invariant 7/8) sees it.
useEffect(() => {
if (isStreaming || !observedRow) return;
const observed = rowToUiMessage(observedRow);
setMessages((prev) => mergeObservedMessage(prev, observed));
}, [observedRow, isStreaming, setMessages]);
// Re-arm on (re)mount — StrictMode dev-mounts twice, and the cleanup below
// flips this false between the two.
mountedRef.current = true;
if (attemptResumeRef.current) {
setResumedTurnPair(true);
void resumeStream();
}
// Unmount: mark unmounted (gates late attach/onFinish side-effects) and abort
// the in-flight attach GET so its callbacks don't fire against the next chat.
return () => {
mountedRef.current = false;
attachAbortRef.current?.abort();
};
// Mount-only by design; the parent remounts per chat via `key`.
// eslint-disable-next-line react-hooks/exhaustive-deps
}, []);
// Reconciliation + degraded-merge (invariant 8). Deps are EXACTLY
// [initialRows, isStreaming, setMessages].
useEffect(() => {
// A local stream owns the view: disarm BOTH the merge and the window poll.
if (isStreaming) {
reconcileTailRef.current = false;
onResumeFallback?.(false);
return;
}
if (!reconcileTailRef.current) return;
const rows = initialRows ?? [];
const tail = rows[rows.length - 1];
if (!tail || tail.role !== "assistant") return;
// Merge the polled assistant tail on EVERY initialRows update — while the
// degraded poll is active this IS the live per-step progress.
setMessages((prev) => mergeById(prev, rowToUiMessage(tail)));
// Anchor-mismatch coherence: when we restored a stripped streaming row A but a
// DIFFERENT run's row B is now the tail (A finished, B replaced the registry
// entry, so the attach 204'd), A would otherwise linger forever as an orphan
// jumping-dots row over the real run. Settle it from fresh history (where A is
// now persisted) so no phantom row survives. No-op in the common case where A
// IS the tail (id match).
const stripped = strippedRowRef.current;
if (stripped && stripped.id !== tail.id) {
const historical = rows.find((r) => r.id === stripped.id);
if (historical)
setMessages((prev) => mergeById(prev, rowToUiMessage(historical)));
}
// Settled: the terminal merge is done — disarm the flag AND the window poll
// explicitly (the window only has a time cap, it will not disarm itself).
if (tail.status !== "streaming") {
reconcileTailRef.current = false;
onResumeFallback?.(false);
}
// onResumeFallback intentionally omitted (parent-stable callback); deps are
// fixed by the resume design.
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [initialRows, isStreaming, setMessages]);
// "Send now" on a queued message: interrupt the current turn and immediately
// send THIS message, keeping the agent's partial output. Other queued messages
@@ -469,10 +674,12 @@ export default function ChatThread({
const msg = queuedRef.current.find((m) => m.id === id);
if (!msg) return;
setQueue(removeQueuedById(queuedRef.current, id));
// Local send: clear any resume-suppression flag (invariant 8).
setResumedTurnPair(false);
sendMessageRef.current?.({ text: msg.text });
}
},
[setQueue, stop],
[setQueue, stop, setResumedTurnPair],
);
// Stop the current turn. ALWAYS abort the local SSE (`stop()`) so the composer
@@ -485,6 +692,9 @@ export default function ChatThread({
// is not known yet — a brand-new chat in the first moment of its first turn —
// only the local abort happens (there is no server-side run handle to stop yet).
const handleStop = useCallback(() => {
// Abort the resume/attach GET first: the SDK does not pass it a signal, so an
// observer's Stop would otherwise leave the attach fetch running.
attachAbortRef.current?.abort();
stop();
if (!autonomousRunsEnabled) return;
if (chatIdRef.current) {
@@ -617,17 +827,23 @@ export default function ChatThread({
<Text size="xs" lineClamp={2} className={classes.queuedText}>
{m.text}
</Text>
<Tooltip label={t("Interrupt and send now")} withArrow>
<ActionIcon
size="xs"
variant="subtle"
color="blue"
onClick={() => sendNow(m.id)}
aria-label={t("Send now")}
>
<IconPlayerPlayFilled size={12} />
</ActionIcon>
</Tooltip>
{/* "Send now" (interrupt) is hidden on a RESUMED turn: a local
stop() does not abort the resumed attach fetch, so the click
would be swallowed while flushOnAbortRef would fire minutes
later on the natural finish. Only the remove affordance stays. */}
{!resumedTurn && (
<Tooltip label={t("Interrupt and send now")} withArrow>
<ActionIcon
size="xs"
variant="subtle"
color="blue"
onClick={() => sendNow(m.id)}
aria-label={t("Send now")}
>
<IconPlayerPlayFilled size={12} />
</ActionIcon>
</Tooltip>
)}
<ActionIcon
size="xs"
variant="subtle"
@@ -642,7 +858,11 @@ export default function ChatThread({
</Stack>
)}
<ChatInput
onSend={(text) => sendMessage({ text })}
onSend={(text) => {
// Local send: clear any resume-suppression flag (invariant 8).
setResumedTurnPair(false);
sendMessage({ text });
}}
onQueue={enqueue}
onStop={handleStop}
isStreaming={isStreaming}
@@ -0,0 +1,90 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import React from "react";
import { renderHook, waitFor } from "@testing-library/react";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
// react-i18next / notifications are pulled in transitively by ai-chat-query.ts
// (the mutation hooks use them); stub so the module imports cleanly.
vi.mock("react-i18next", () => ({
useTranslation: () => ({ t: (key: string) => key }),
}));
vi.mock("@mantine/notifications", () => ({
notifications: { show: vi.fn() },
}));
// Mock the service module; only getAiChatMessages is exercised, but the other
// named exports must exist so ai-chat-query.ts imports resolve.
vi.mock("@/features/ai-chat/services/ai-chat-service.ts", () => ({
getAiChatMessages: vi.fn(),
getAiChats: vi.fn(),
getAiRoleCatalog: vi.fn(),
getAiRoleCatalogBundle: vi.fn(),
getAiRoles: vi.fn(),
importAiRolesFromCatalog: vi.fn(),
createAiRole: vi.fn(),
deleteAiChat: vi.fn(),
deleteAiRole: vi.fn(),
renameAiChat: vi.fn(),
updateAiRole: vi.fn(),
updateAiRoleFromCatalog: vi.fn(),
}));
import { getAiChatMessages } from "@/features/ai-chat/services/ai-chat-service.ts";
import { useAiChatMessagesQuery } from "@/features/ai-chat/queries/ai-chat-query.ts";
const emptyPage = { items: [], meta: { hasNextPage: false, nextCursor: null } };
function createWrapper() {
const queryClient = new QueryClient({
defaultOptions: { queries: { retry: false } },
});
return function Wrapper({ children }: { children: React.ReactNode }) {
return (
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
);
};
}
// The degraded-poll fallback (#184 phase 1.5) is threaded into this query as a
// `refetchInterval`; AiChatWindow supplies the deliberately-dumb callback. These
// pin the plumbing the window depends on: the interval polls the message history,
// and — critically — fetch ERRORS do NOT stop the tick (TanStack v5 resets the
// failure count each fetch, so the poll must survive a server restart).
describe("useAiChatMessagesQuery — degraded refetchInterval", () => {
beforeEach(() => {
vi.clearAllMocks();
});
it("re-polls at the interval while the callback returns a duration", async () => {
vi.mocked(getAiChatMessages).mockResolvedValue(emptyPage as never);
renderHook(() => useAiChatMessagesQuery("c1", () => 30), {
wrapper: createWrapper(),
});
await waitFor(() =>
expect(vi.mocked(getAiChatMessages).mock.calls.length).toBeGreaterThan(2),
);
});
it("does NOT re-poll when the callback returns false", async () => {
vi.mocked(getAiChatMessages).mockResolvedValue(emptyPage as never);
renderHook(() => useAiChatMessagesQuery("c1", () => false), {
wrapper: createWrapper(),
});
await waitFor(() =>
expect(vi.mocked(getAiChatMessages)).toHaveBeenCalledTimes(1),
);
// Give any errant interval a chance to fire, then assert it did not.
await new Promise((r) => setTimeout(r, 60));
expect(vi.mocked(getAiChatMessages)).toHaveBeenCalledTimes(1);
});
it("keeps ticking through fetch errors (errors do not gate the poll)", async () => {
vi.mocked(getAiChatMessages).mockRejectedValue(new Error("server down"));
renderHook(() => useAiChatMessagesQuery("c1", () => 30), {
wrapper: createWrapper(),
});
await waitFor(() =>
expect(vi.mocked(getAiChatMessages).mock.calls.length).toBeGreaterThan(2),
);
});
});
@@ -13,7 +13,6 @@ import {
deleteAiChat,
deleteAiRole,
getAiChatMessages,
getAiChatRun,
getAiChats,
getAiRoleCatalog,
getAiRoleCatalogBundle,
@@ -26,7 +25,6 @@ import {
import {
IAiChat,
IAiChatMessageRow,
IAiChatRunResponse,
IAiRole,
IAiRoleCatalog,
IAiRoleCatalogBundle,
@@ -37,7 +35,6 @@ import {
IAiRoleUpdateFromCatalogResult,
} from "@/features/ai-chat/types/ai-chat.types.ts";
import { IPagination } from "@/lib/types.ts";
import { runPollInterval } from "@/features/ai-chat/utils/run-polling.ts";
export const AI_CHATS_RQ_KEY = ["ai-chats"];
export const AI_ROLES_RQ_KEY = ["ai-roles"];
@@ -55,7 +52,6 @@ export const AI_CHAT_MESSAGES_RQ_KEY = (chatId: string) => [
"ai-chat-messages",
chatId,
];
export const AI_CHAT_RUN_RQ_KEY = (chatId: string) => ["ai-chat-run", chatId];
/** Paginated list of the current user's chats (auto-loads further pages). */
export function useAiChatsQuery() {
@@ -89,7 +85,15 @@ export function useAiChatsQuery() {
* Load all persisted messages of a chat (oldest first), flattening the
* paginated server response. Used to seed `useChat` initial messages.
*/
export function useAiChatMessagesQuery(chatId: string | undefined) {
export function useAiChatMessagesQuery(
chatId: string | undefined,
// #184 phase 1.5: the degraded-poll fallback. When a tab could not attach to a
// still-running run (the attach returned 204 / the resumed stream ended with no
// terminal row), the window arms a dumb timed poll of the message history to
// follow the detached run to settle. The callback form lives in AiChatWindow;
// threaded here verbatim so this query owns the polling. Undefined => no poll.
refetchInterval?: number | false | (() => number | false),
) {
const query = useInfiniteQuery({
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatId ?? ""),
queryFn: ({ pageParam }) =>
@@ -100,6 +104,7 @@ export function useAiChatMessagesQuery(chatId: string | undefined) {
? (lastPage.meta.nextCursor ?? undefined)
: undefined,
enabled: !!chatId,
refetchInterval,
});
// useInfiniteQuery only fetches the first page on its own. The hook's contract
@@ -139,34 +144,6 @@ export function useAiChatMessagesQuery(chatId: string | undefined) {
};
}
/**
* Reconnect to a chat's latest agent run and LIVE-FOLLOW it (#184). While the run
* is active the query re-polls every {@link runPollInterval} ms (driven off the
* fetched `run.status`, the same status-keyed refetchInterval pattern as the
* embeddings reindex polling); once the run reaches a terminal status — or there
* is no run — the interval returns `false` and polling stops on its own. Polling
* is thus naturally bounded by the run terminating; no separate timeout cap.
*
* `enabled` gates the whole thing: callers pass `false` when the autonomous-runs
* feature is off (the endpoint is NOT flag-gated server-side, but with the feature
* off the chat has no runs, so polling would only ever return `{ run: null }`) OR
* when THIS tab is the one actively streaming the run (the live SSE owns the view,
* so we must not also poll/merge). The global `retry: false` means a failed fetch
* leaves `data` undefined, so refetchInterval(undefined run) returns false — a
* failed fetch can never spin a tight loop.
*/
export function useAiChatRunQuery(
chatId: string | undefined,
enabled: boolean,
) {
return useQuery<IAiChatRunResponse, Error>({
queryKey: AI_CHAT_RUN_RQ_KEY(chatId ?? ""),
queryFn: () => getAiChatRun(chatId as string),
enabled: !!chatId && enabled,
refetchInterval: (query) => runPollInterval(query.state.data?.run),
});
}
export function useRenameAiChatMutation() {
const queryClient = useQueryClient();
const { t } = useTranslation();
@@ -1,92 +0,0 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import React from "react";
import { renderHook, waitFor } from "@testing-library/react";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import type { IAiChatRunResponse } from "@/features/ai-chat/types/ai-chat.types.ts";
// react-i18next is pulled in transitively by ai-chat-query.ts (the mutation hooks
// use it); stub it so the module imports cleanly in this hook test.
vi.mock("react-i18next", () => ({
useTranslation: () => ({ t: (key: string) => key }),
}));
vi.mock("@mantine/notifications", () => ({
notifications: { show: vi.fn() },
}));
// Mock the whole service module; only getAiChatRun is exercised here, but the
// other named exports must exist so ai-chat-query.ts imports resolve.
vi.mock("@/features/ai-chat/services/ai-chat-service.ts", () => ({
getAiChatRun: vi.fn(),
getAiChatMessages: vi.fn(),
getAiChats: vi.fn(),
getAiRoleCatalog: vi.fn(),
getAiRoleCatalogBundle: vi.fn(),
getAiRoles: vi.fn(),
importAiRolesFromCatalog: vi.fn(),
createAiRole: vi.fn(),
deleteAiChat: vi.fn(),
deleteAiRole: vi.fn(),
renameAiChat: vi.fn(),
updateAiRole: vi.fn(),
updateAiRoleFromCatalog: vi.fn(),
}));
import { getAiChatRun } from "@/features/ai-chat/services/ai-chat-service.ts";
import { useAiChatRunQuery } from "@/features/ai-chat/queries/ai-chat-query.ts";
function createWrapper() {
const queryClient = new QueryClient({
defaultOptions: { queries: { retry: false } },
});
return function Wrapper({ children }: { children: React.ReactNode }) {
return (
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
);
};
}
const runningResponse: IAiChatRunResponse = {
run: { id: "run-1", chatId: "c1", status: "running" },
message: {
id: "a1",
role: "assistant",
content: "working...",
createdAt: "2026-01-01T00:00:00Z",
},
};
describe("useAiChatRunQuery — enable gating", () => {
beforeEach(() => {
vi.clearAllMocks();
});
it("fetches the run when enabled (passive observer, feature on)", async () => {
vi.mocked(getAiChatRun).mockResolvedValue(runningResponse);
const { result } = renderHook(() => useAiChatRunQuery("c1", true), {
wrapper: createWrapper(),
});
await waitFor(() => expect(result.current.isSuccess).toBe(true));
expect(getAiChatRun).toHaveBeenCalledWith("c1");
expect(result.current.data?.run?.status).toBe("running");
});
it("does NOT fetch when disabled (this tab is the streamer / feature off)", async () => {
vi.mocked(getAiChatRun).mockResolvedValue(runningResponse);
renderHook(() => useAiChatRunQuery("c1", false), {
wrapper: createWrapper(),
});
// Give any errant fetch a chance to fire, then assert none did.
await new Promise((r) => setTimeout(r, 20));
expect(getAiChatRun).not.toHaveBeenCalled();
});
it("does NOT fetch when there is no chat id", async () => {
vi.mocked(getAiChatRun).mockResolvedValue(runningResponse);
renderHook(() => useAiChatRunQuery(undefined, true), {
wrapper: createWrapper(),
});
await new Promise((r) => setTimeout(r, 20));
expect(getAiChatRun).not.toHaveBeenCalled();
});
});
@@ -5,7 +5,6 @@ import {
IAiChatListParams,
IAiChatMessageRow,
IAiChatMessagesParams,
IAiChatRunResponse,
IAiRole,
IAiRoleCatalog,
IAiRoleCatalogBundle,
@@ -43,23 +42,6 @@ export async function getAiChatMessages(
return req.data;
}
/**
* Reconnect to the latest agent run of a chat (#184). Returns the run's
* persisted lifecycle state and the assistant message it materializes (the
* partial output while the run is in-flight, the final output once it finished).
* The DB is the source of truth, so this works for an in-flight run (the browser
* dropped, the run kept going) and a finished one alike; `{ run: null }` when the
* chat has never had a run. Owner-gated server-side (the requesting user must own
* the chat); it is NOT flag-gated — when the feature is off the chat simply has no
* runs, so the endpoint returns `{ run: null }`.
*/
export async function getAiChatRun(
chatId: string,
): Promise<IAiChatRunResponse> {
const req = await api.post<IAiChatRunResponse>("/ai-chat/run", { chatId });
return req.data;
}
/**
* Explicitly STOP the active agent run of a chat (#184). This is the ONLY thing
* that ends a DETACHED run — a mere browser disconnect (aborting the local SSE)
@@ -210,41 +210,14 @@ export interface IAiChatMessageRow {
// renders a "stopped" marker on interrupted turns.
finishReason?: string;
} | null;
// Persisted lifecycle status of the row's turn, carried on the wire by
// `baseFields`. 'streaming' marks a still-in-progress assistant row (used by
// the resume machinery to decide whether a tail is a live stream to attach to
// or a settled row that must not be replayed).
status?: string;
createdAt: string;
}
/**
* A persisted agent-run row (#184), mirroring the `ai_chat_runs` fields the
* client reads from `POST /ai-chat/run`. Only `status` is load-bearing for the
* reconnect-and-live-update UX (it drives the poll cadence); the rest are carried
* for display/diagnostics. The DB is the source of truth, so this resolves for an
* in-flight run (the browser dropped, the run kept going) and a finished one.
*/
export interface IAiChatRun {
id: string;
chatId: string;
// 'pending' | 'running' | 'succeeded' | 'failed' | 'aborted'. The first two are
// ACTIVE (keep polling); the rest are TERMINAL (stop polling).
status: "pending" | "running" | "succeeded" | "failed" | "aborted" | string;
error?: string | null;
stepCount?: number;
assistantMessageId?: string | null;
startedAt?: string | null;
finishedAt?: string | null;
createdAt?: string;
updatedAt?: string;
}
/**
* Response of `POST /ai-chat/run` (#184): the latest run of a chat and the
* assistant message it materializes (the partial/final output, projected from the
* persisted rows). Both are `null` when the chat has never had a run.
*/
export interface IAiChatRunResponse {
run: IAiChatRun | null;
message: IAiChatMessageRow | null;
}
export interface IAiChatListParams extends QueryParams {}
export interface IAiChatMessagesParams {
@@ -0,0 +1,112 @@
import { describe, it, expect } from "vitest";
import type { UIMessage } from "@ai-sdk/react";
import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.ts";
import {
isStreamingTail,
isSettledAssistantTail,
seedRows,
mergeById,
} from "./resume-helpers.ts";
function row(
id: string,
role: string,
status?: string,
): IAiChatMessageRow {
return { id, role, content: "", status, createdAt: "2026-01-01T00:00:00Z" };
}
function makeMsg(id: string, text: string): UIMessage {
return {
id,
role: "assistant",
parts: [{ type: "text", text }],
} as UIMessage;
}
describe("isStreamingTail", () => {
it("is true when the last row is a streaming assistant row", () => {
expect(
isStreamingTail([row("u1", "user"), row("a1", "assistant", "streaming")]),
).toBe(true);
});
it("is false for a settled assistant tail", () => {
expect(isStreamingTail([row("a1", "assistant", "succeeded")])).toBe(false);
expect(isStreamingTail([row("a1", "assistant")])).toBe(false);
});
it("is false when the tail is a user row or the list is empty", () => {
expect(isStreamingTail([row("u1", "user")])).toBe(false);
expect(isStreamingTail([])).toBe(false);
});
});
describe("isSettledAssistantTail", () => {
it("is true for an assistant tail whose status is not streaming", () => {
expect(isSettledAssistantTail([row("a1", "assistant", "succeeded")])).toBe(
true,
);
expect(isSettledAssistantTail([row("a1", "assistant")])).toBe(true);
expect(isSettledAssistantTail([row("a1", "assistant", "aborted")])).toBe(
true,
);
});
it("is false for a streaming assistant tail", () => {
expect(isSettledAssistantTail([row("a1", "assistant", "streaming")])).toBe(
false,
);
});
it("is false when the tail is a user row or the list is empty", () => {
expect(isSettledAssistantTail([row("u1", "user")])).toBe(false);
expect(isSettledAssistantTail([])).toBe(false);
});
});
describe("seedRows", () => {
const rows = [row("u1", "user"), row("a1", "assistant", "streaming")];
it("returns the rows unchanged when not stripping", () => {
expect(seedRows(rows, false)).toBe(rows);
});
it("drops the last row when stripping", () => {
const seeded = seedRows(rows, true);
expect(seeded).toHaveLength(1);
expect(seeded[0].id).toBe("u1");
});
it("returns an empty list when stripping a single-row list", () => {
expect(seedRows([row("a1", "assistant", "streaming")], true)).toHaveLength(
0,
);
});
});
describe("mergeById", () => {
it("replaces the message with the same id in place (per-step growth)", () => {
const prev = [makeMsg("u1", "hi"), makeMsg("a1", "step 1")];
const incoming = makeMsg("a1", "step 1\nstep 2");
const next = mergeById(prev, incoming);
expect(next).toHaveLength(2);
expect(next[1]).toBe(incoming);
expect(next[0]).toBe(prev[0]); // untouched
expect(next).not.toBe(prev); // new array (never mutates input)
});
it("appends when the incoming message is not yet present", () => {
const prev = [makeMsg("u1", "hi")];
const incoming = makeMsg("a1", "first token");
const next = mergeById(prev, incoming);
expect(next).toHaveLength(2);
expect(next[1]).toBe(incoming);
});
it("returns the original list unchanged when there is nothing to merge", () => {
const prev = [makeMsg("u1", "hi")];
expect(mergeById(prev, null)).toBe(prev);
expect(mergeById(prev, undefined)).toBe(prev);
});
});
@@ -0,0 +1,62 @@
import type { UIMessage } from "@ai-sdk/react";
import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.ts";
/**
* Pure decisions for the resumable-SSE resume machinery (#184 phase 1.5). A tab
* that reopens a chat whose agent run is still going attaches to the server's
* run-stream registry (replay + live tail) instead of polling snapshots; these
* small predicates decide WHICH tail is safe to resume and how to seed the store,
* extracted so they can be unit-tested in isolation.
*/
/**
* A STREAMING tail: the last persisted row is an assistant row still marked
* `status === 'streaming'`. Such a tail is stripped from the seed and rebuilt by
* the replay (`expect=live`), since the SDK's `text-start` always pushes a new
* part and replaying over a seeded in-progress row would duplicate its text.
*/
export function isStreamingTail(rows: IAiChatMessageRow[]): boolean {
const tail = rows[rows.length - 1];
return !!tail && tail.role === "assistant" && tail.status === "streaming";
}
/**
* A SETTLED assistant tail: the last row is an assistant row whose status is
* anything OTHER than 'streaming'. A settled assistant tail must NEVER resume —
* replaying a finished run into a store that already holds its message duplicates
* parts (`text-start` always pushes a new part).
*/
export function isSettledAssistantTail(rows: IAiChatMessageRow[]): boolean {
const tail = rows[rows.length - 1];
return !!tail && tail.role === "assistant" && tail.status !== "streaming";
}
/**
* Seed rows for `useChat`: return the rows unchanged, or without the last row when
* `strip` is set (the streaming tail is stripped so the live replay rebuilds it
* without duplicating parts).
*/
export function seedRows(
rows: IAiChatMessageRow[],
strip: boolean,
): IAiChatMessageRow[] {
return strip ? rows.slice(0, -1) : rows;
}
/**
* Merge an assistant message into the rendered list by id: replace the message
* with the same id in place (the in-progress assistant row is already seeded from
* history, so per-step growth replaces it), or append it when absent. Returns a
* new array; the input is never mutated.
*/
export function mergeById(
messages: UIMessage[],
incoming: UIMessage | null | undefined,
): UIMessage[] {
if (!incoming) return messages;
const idx = messages.findIndex((m) => m.id === incoming.id);
if (idx === -1) return [...messages, incoming];
const next = messages.slice();
next[idx] = incoming;
return next;
}
@@ -1,303 +0,0 @@
import { describe, it, expect } from "vitest";
import type { UIMessage } from "@ai-sdk/react";
import type { IAiChatRun } from "@/features/ai-chat/types/ai-chat.types.ts";
import {
RUN_POLL_INTERVAL_MS,
isRunActive,
runPollInterval,
shouldObserveRun,
shouldClearStoppingLatch,
shouldClearLatchOnQueryError,
mergeObservedMessage,
} from "./run-polling.ts";
function makeRun(status: string): IAiChatRun {
return { id: "run-1", chatId: "c1", status };
}
function makeMsg(id: string, text: string): UIMessage {
return {
id,
role: "assistant",
parts: [{ type: "text", text }],
} as UIMessage;
}
describe("isRunActive", () => {
it("treats pending and running as active", () => {
expect(isRunActive(makeRun("pending"))).toBe(true);
expect(isRunActive(makeRun("running"))).toBe(true);
});
it("treats terminal / unknown / nullish as not active", () => {
expect(isRunActive(makeRun("succeeded"))).toBe(false);
expect(isRunActive(makeRun("failed"))).toBe(false);
expect(isRunActive(makeRun("aborted"))).toBe(false);
expect(isRunActive(makeRun("weird-future-status"))).toBe(false);
expect(isRunActive(null)).toBe(false);
expect(isRunActive(undefined)).toBe(false);
});
});
describe("runPollInterval (the refetchInterval helper)", () => {
it("returns 2000ms while the run is pending/running", () => {
expect(runPollInterval(makeRun("pending"))).toBe(RUN_POLL_INTERVAL_MS);
expect(runPollInterval(makeRun("running"))).toBe(RUN_POLL_INTERVAL_MS);
expect(RUN_POLL_INTERVAL_MS).toBe(2000);
});
it("returns false (stop polling) once the run is terminal", () => {
expect(runPollInterval(makeRun("succeeded"))).toBe(false);
expect(runPollInterval(makeRun("failed"))).toBe(false);
expect(runPollInterval(makeRun("aborted"))).toBe(false);
});
it("returns false (no polling) when there is no run", () => {
expect(runPollInterval(null)).toBe(false);
expect(runPollInterval(undefined)).toBe(false);
});
});
describe("shouldObserveRun (observer-vs-streamer decision)", () => {
it("observes an active run when this tab is NOT the local streamer", () => {
expect(shouldObserveRun(makeRun("running"), false)).toBe(true);
expect(shouldObserveRun(makeRun("pending"), false)).toBe(true);
});
it("observes a terminal run too (so the final output shows on reopen)", () => {
expect(shouldObserveRun(makeRun("succeeded"), false)).toBe(true);
});
it("does NOT observe when this tab IS the streamer (no double-render)", () => {
expect(shouldObserveRun(makeRun("running"), true)).toBe(false);
expect(shouldObserveRun(makeRun("succeeded"), true)).toBe(false);
});
it("does NOT observe when there is no run", () => {
expect(shouldObserveRun(null, false)).toBe(false);
expect(shouldObserveRun(undefined, false)).toBe(false);
});
});
describe("shouldClearStoppingLatch (#234 latch-release decision)", () => {
// The one case the latch SHOULD clear: we requested a stop, we are the passive
// observer (not streaming), and the CURRENT run is terminal.
it("clears only when stopping, observing, and the run is terminal", () => {
expect(
shouldClearStoppingLatch({
stoppingRun: true,
run: makeRun("aborted"),
isLocalStreaming: false,
}),
).toBe(true);
expect(
shouldClearStoppingLatch({
stoppingRun: true,
run: makeRun("succeeded"),
isLocalStreaming: false,
}),
).toBe(true);
expect(
shouldClearStoppingLatch({
stoppingRun: true,
run: makeRun("failed"),
isLocalStreaming: false,
}),
).toBe(true);
});
// Round-3 regression: clearing while THIS tab is still the local streamer would
// re-open the flash for the current turn the moment we switch to observer role.
// A predicate lacking the streaming gate would (wrongly) return true here.
it("does NOT clear while this tab is the local streamer", () => {
expect(
shouldClearStoppingLatch({
stoppingRun: true,
run: makeRun("aborted"),
isLocalStreaming: true,
}),
).toBe(false);
expect(
shouldClearStoppingLatch({
stoppingRun: true,
run: makeRun("succeeded"),
isLocalStreaming: true,
}),
).toBe(false);
});
// The detached run keeps growing after a local abort — while it is still
// active the latch MUST hold so the observer merge stays suppressed.
it("does NOT clear while the run is still active", () => {
expect(
shouldClearStoppingLatch({
stoppingRun: true,
run: makeRun("running"),
isLocalStreaming: false,
}),
).toBe(false);
expect(
shouldClearStoppingLatch({
stoppingRun: true,
run: makeRun("pending"),
isLocalStreaming: false,
}),
).toBe(false);
});
// #234 F4: on Stop the stale PREVIOUS-turn run is removed from the cache, so the
// observed `run` is null until the current turn's run is fetched fresh. A null
// run HOLDS the latch — it can never clear against the just-removed stale run,
// only against the current turn's own terminal run once observed.
it("does NOT clear against a removed/absent run (F4 stale-run guard)", () => {
expect(
shouldClearStoppingLatch({
stoppingRun: true,
run: null,
isLocalStreaming: false,
}),
).toBe(false);
expect(
shouldClearStoppingLatch({
stoppingRun: true,
run: undefined,
isLocalStreaming: false,
}),
).toBe(false);
});
it("does NOT clear when no stop was requested", () => {
expect(
shouldClearStoppingLatch({
stoppingRun: false,
run: makeRun("aborted"),
isLocalStreaming: false,
}),
).toBe(false);
});
});
describe("shouldClearLatchOnQueryError (#234 F7 error-safety-net decision)", () => {
// This guards the REAL anti-flash decision the component's run-query-error
// safety-net effect uses (ai-chat-window.tsx wires the effect to THIS helper,
// not a copy — so the test is non-vacuous vs the live code).
// (b) The F7 hole: a TRANSIENT run-query error while `run` is STILL ACTIVE must
// NOT clear the latch. TanStack Query v5 retains `data` on error, so
// runQueryFailed can be true while the held run is still pending/running.
// Against the PRE-F7 condition (without `!isRunActive(run)`) this would return
// true — so this assertion fails on the buggy code (non-vacuous).
it("does NOT clear on a transient error while the run is still ACTIVE (F7)", () => {
expect(
shouldClearLatchOnQueryError({
stoppingRun: true,
isLocalStreaming: false,
runQueryFailed: true,
run: makeRun("running"),
}),
).toBe(false);
expect(
shouldClearLatchOnQueryError({
stoppingRun: true,
isLocalStreaming: false,
runQueryFailed: true,
run: makeRun("pending"),
}),
).toBe(false);
});
// (a) The genuine permanent-null-freeze: run cache cleared by removeQueries +
// the refetch keeps ERRORING, so `run === null`. This is the ONLY case the
// safety-net exists to cure — it MUST clear so the frozen view resumes.
it("clears on a permanent error when the run is null (permanent-null-freeze)", () => {
expect(
shouldClearLatchOnQueryError({
stoppingRun: true,
isLocalStreaming: false,
runQueryFailed: true,
run: null,
}),
).toBe(true);
expect(
shouldClearLatchOnQueryError({
stoppingRun: true,
isLocalStreaming: false,
runQueryFailed: true,
run: undefined,
}),
).toBe(true);
});
// A TERMINAL run also satisfies `!isRunActive`; clearing then is harmless — the
// terminal effect (shouldClearStoppingLatch) already clears for a terminal run,
// so this only ever agrees with it. Asserted so the (c) reasoning is pinned.
it("clears on an error when the run is terminal (harmless, agrees with terminal effect)", () => {
expect(
shouldClearLatchOnQueryError({
stoppingRun: true,
isLocalStreaming: false,
runQueryFailed: true,
run: makeRun("aborted"),
}),
).toBe(true);
});
it("does NOT clear without an actual query error", () => {
expect(
shouldClearLatchOnQueryError({
stoppingRun: true,
isLocalStreaming: false,
runQueryFailed: false,
run: null,
}),
).toBe(false);
});
it("does NOT clear while this tab is the local streamer", () => {
expect(
shouldClearLatchOnQueryError({
stoppingRun: true,
isLocalStreaming: true,
runQueryFailed: true,
run: null,
}),
).toBe(false);
});
it("does NOT clear when no stop was requested", () => {
expect(
shouldClearLatchOnQueryError({
stoppingRun: false,
isLocalStreaming: false,
runQueryFailed: true,
run: null,
}),
).toBe(false);
});
});
describe("mergeObservedMessage", () => {
it("replaces the message with the same id in place (per-step growth)", () => {
const prev = [makeMsg("u1", "hi"), makeMsg("a1", "step 1")];
const observed = makeMsg("a1", "step 1\nstep 2");
const next = mergeObservedMessage(prev, observed);
expect(next).toHaveLength(2);
expect(next[1]).toBe(observed);
expect(next[0]).toBe(prev[0]); // untouched
expect(next).not.toBe(prev); // new array (never mutates input)
});
it("appends when the observed message is not yet present", () => {
const prev = [makeMsg("u1", "hi")];
const observed = makeMsg("a1", "first token");
const next = mergeObservedMessage(prev, observed);
expect(next).toHaveLength(2);
expect(next[1]).toBe(observed);
});
it("returns the original list unchanged when there is nothing to merge", () => {
const prev = [makeMsg("u1", "hi")];
expect(mergeObservedMessage(prev, null)).toBe(prev);
expect(mergeObservedMessage(prev, undefined)).toBe(prev);
});
});
@@ -1,151 +0,0 @@
import type { UIMessage } from "@ai-sdk/react";
import type { IAiChatRun } from "@/features/ai-chat/types/ai-chat.types.ts";
/**
* Reconnect-and-live-follow helpers (#184). When a chat is reopened while its
* agent run is STILL going, this tab is a PASSIVE OBSERVER: it did not start the
* run here (no local SSE stream), so it catches up by POLLING the reconnect
* endpoint (`POST /ai-chat/run`) and merging the run's incrementally-persisted
* assistant message into the rendered thread. These are the small pure decisions
* that machinery hangs off, extracted so they can be unit-tested in isolation
* (mirrors how reindex polling / editor-sync-state are tested).
*/
/** How often to re-poll the reconnect endpoint while a run is ACTIVE. */
export const RUN_POLL_INTERVAL_MS = 2000;
// 'pending' and 'running' are the two ACTIVE statuses; 'succeeded' | 'failed' |
// 'aborted' are TERMINAL (and any unknown future status is treated as terminal,
// so a stale/odd value never polls forever).
const ACTIVE_STATUSES = new Set(["pending", "running"]);
/** Whether a run is still going (worth polling / merging live updates from). */
export function isRunActive(run: IAiChatRun | null | undefined): boolean {
return !!run && ACTIVE_STATUSES.has(run.status);
}
/**
* The TanStack Query `refetchInterval` value for the run query: poll every
* {@link RUN_POLL_INTERVAL_MS} while the run is active, and `false` (stop) once
* it is terminal or there is no run. Polling is thus naturally bounded by the run
* reaching a terminal status — no separate timeout cap is needed.
*/
export function runPollInterval(
run: IAiChatRun | null | undefined,
): number | false {
return isRunActive(run) ? RUN_POLL_INTERVAL_MS : false;
}
/**
* Observer-vs-streamer decision. We render the polled run message (catch up +
* keep advancing) ONLY when this tab is a passive observer: there IS a run AND
* this tab is NOT the one locally streaming it (we reconnected, we didn't start
* it here). When this tab is the streamer, the live SSE stream owns the view, so
* we neither poll nor merge — avoiding a double-render fight. Terminal runs still
* merge (so the final persisted output is shown on reopen); the poll itself is
* stopped separately by {@link runPollInterval}.
*/
export function shouldObserveRun(
run: IAiChatRun | null | undefined,
localStreaming: boolean,
): boolean {
return !!run && !localStreaming;
}
/**
* Should the "stopping" latch — which suppresses the observer re-stream flash
* after the user pressed Stop — be RELEASED now? All three must hold:
* - `stoppingRun`: we actually requested a stop (otherwise nothing to release);
* - `!isLocalStreaming`: this tab is NOT the local streamer. While we are the
* streamer the run query is disabled, so the observed `run` is not the run we
* are following — releasing the latch then would re-open the flash for the
* current turn the instant we switch to observer role;
* - the observed `run` EXISTS and has reached a TERMINAL status.
*
* The null / still-active `run` case is the #234 F4 invariant. On Stop the stale
* PREVIOUS-turn run is removed from the query cache (`removeQueries`), so `run`
* is null until the CURRENT turn's run is re-fetched fresh; a null or active run
* therefore HOLDS the latch, so it can only ever clear against the current turn's
* OWN terminal run — never a stale cached one. (The cache removal itself is
* integration-level in AiChatWindow; this predicate encodes the decision given
* whatever run is currently observed, and a stale terminal run is
* indistinguishable from a current terminal run at the predicate level — hence
* the cache removal is what guarantees only the current run is ever passed here.)
*/
export function shouldClearStoppingLatch(args: {
stoppingRun: boolean;
run: IAiChatRun | null | undefined;
isLocalStreaming: boolean;
}): boolean {
const { stoppingRun, run, isLocalStreaming } = args;
if (!stoppingRun || isLocalStreaming) return false;
return !!run && !isRunActive(run);
}
/**
* Should the "stopping" latch be RELEASED by the run-query ERROR safety-net?
* (#234 F7 — a NEW path of the same re-stream flash the F4 latch exists to
* prevent.) After Stop, `handleServerStop` clears the run cache; the terminal
* effect then holds the latch via `if (!run) return` until the CURRENT turn's run
* is fetched fresh. If that refetch instead ERRORS permanently, `run` stays null,
* its status-keyed refetchInterval is off, and nothing would ever observe a
* terminal run — freezing the view with the observer merge suppressed. This
* safety-net cures ONLY that genuine permanent-null-freeze.
*
* All four must hold:
* - `stoppingRun`: we actually requested a stop (otherwise nothing to release);
* - `!isLocalStreaming`: this tab is NOT the local streamer (same reason as
* {@link shouldClearStoppingLatch});
* - `runQueryFailed`: the run query is in its error state (TanStack Query v5 with
* retry:false — isError);
* - `!isRunActive(run)`: the observed `run` is NOT an active (pending/running)
* held run. This is the F7 gate. In TanStack Query v5 the query's `data` is
* RETAINED on error, so `runQueryFailed` can be true while `run` is STILL an
* ACTIVE run (a single transient GET-run failure in the window between Stop and
* settle). Without this gate a transient error would release the latch early —
* re-opening the observer merge and flashing the growing detached run over the
* frozen row (exactly the F4 flash). Gating on the run NOT being active means we
* only ever cure the permanent-null-freeze (`run === null`, so
* `isRunActive(null)` is false), never release against an active run.
*
* (A terminal `run` also satisfies `!isRunActive(run)`; clearing then is harmless
* — the terminal effect's {@link shouldClearStoppingLatch} already clears the
* latch for a terminal run, so this only ever agrees with it, never conflicts.)
*
* INVARIANT (do not break): clearing the latch on the `run === null` branch is safe
* ONLY because the run query's `refetchInterval` (see {@link runPollInterval}) stops
* polling when the data is empty — so after we clear on null+error there is no
* subsequent auto-poll that could return a still-active detached run and re-open the
* merge. If `refetchInterval` is ever changed to keep polling on `run === null`/on
* error, this null-branch clear would re-open the F7 flash through the null path.
* Do not change the run query's refetchInterval without re-checking this path.
*/
export function shouldClearLatchOnQueryError(args: {
stoppingRun: boolean;
isLocalStreaming: boolean;
runQueryFailed: boolean;
run: IAiChatRun | null | undefined;
}): boolean {
const { stoppingRun, isLocalStreaming, runQueryFailed, run } = args;
return (
stoppingRun && !isLocalStreaming && runQueryFailed && !isRunActive(run)
);
}
/**
* Merge an observed assistant message into the rendered list: replace the message
* with the same id in place (the in-progress assistant row is already seeded from
* history, so per-step growth replaces it), or append it when absent. Returns a
* new array; the input is never mutated.
*/
export function mergeObservedMessage(
messages: UIMessage[],
observed: UIMessage | null | undefined,
): UIMessage[] {
if (!observed) return messages;
const idx = messages.findIndex((m) => m.id === observed.id);
if (idx === -1) return [...messages, observed];
const next = messages.slice();
next[idx] = observed;
return next;
}