Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 53d367662b | |||
| f0334d7f9b | |||
| 10d5220f5e | |||
| 52ee3c1f3e |
@@ -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;
|
||||
}
|
||||
@@ -0,0 +1,329 @@
|
||||
import { Injectable, Logger, OnModuleDestroy } from '@nestjs/common';
|
||||
|
||||
/**
|
||||
* In-memory run-stream registry (#184 phase 1.5). A durable agent run tees its
|
||||
* SSE frames here (via `pipeUIMessageStreamToResponse({ consumeSseStream })`)
|
||||
* so a LATE tab — one that reloaded, or opened after the starter dropped — can
|
||||
* attach through `GET /ai-chat/runs/:chatId/stream`, replay the frames buffered
|
||||
* so far, and then follow the live tail as a normal streamer.
|
||||
*
|
||||
* This is deliberately single-process and best-effort: it holds nothing the DB
|
||||
* does not (the run + assistant row are the source of truth), so a process
|
||||
* restart simply drops in-flight entries and the client falls back to its
|
||||
* restore + degraded-poll path. The async `attach` return type is the seam for a
|
||||
* future phase-2 cross-process backend (Redis) — the interface does not change.
|
||||
*/
|
||||
|
||||
/** How long a finished entry is retained for late attach (replay + immediate end). */
|
||||
export const RUN_STREAM_RETAIN_FINISHED_MS = 30_000;
|
||||
|
||||
/** Per-run replay buffer cap. Past this the buffer is dropped (attach -> 204). */
|
||||
export const RUN_STREAM_MAX_BUFFER_BYTES = 4 * 1024 * 1024;
|
||||
|
||||
// 2x the replay cap: a just-written 4MB replay burst alone can never trip the
|
||||
// per-subscriber cap (see controller); only a genuinely stalled socket can.
|
||||
export const SUBSCRIBER_MAX_BUFFERED_BYTES = 2 * RUN_STREAM_MAX_BUFFER_BYTES;
|
||||
|
||||
export interface RunStreamCallbacks {
|
||||
onFrame: (frame: string) => void;
|
||||
onEnd: () => void;
|
||||
}
|
||||
|
||||
export interface RunStreamAttachment {
|
||||
replay: string[];
|
||||
finished: boolean;
|
||||
start(): void; // drain pending frames (order preserved) and go live
|
||||
unsubscribe(): void; // safe to call at any point, idempotent
|
||||
}
|
||||
|
||||
interface Subscriber extends RunStreamCallbacks {
|
||||
started: boolean;
|
||||
pending: string[];
|
||||
// Byte size of `pending`, capped at SUBSCRIBER_MAX_BUFFERED_BYTES. `start()` is
|
||||
// called in the SAME tick as `attach()` today (see attach), so `pending` never
|
||||
// holds more than one microtask of frames — but the async `attach` signature is
|
||||
// a phase-2 seam: an await between attach and start would let a stalled paused
|
||||
// subscriber buffer the WHOLE run here. The cap is the structural backstop.
|
||||
pendingBytes: number;
|
||||
overflowed: boolean;
|
||||
pendingEnd: boolean;
|
||||
}
|
||||
|
||||
interface Entry {
|
||||
runId: string;
|
||||
// The persisted assistant row id of this run (set at bind; undefined if the
|
||||
// seed failed). Used by the attach anchor check (invariant 6).
|
||||
assistantMessageId?: string;
|
||||
frames: string[];
|
||||
bytes: number;
|
||||
overflowed: boolean;
|
||||
finished: boolean;
|
||||
subscribers: Set<Subscriber>;
|
||||
retainTimer?: NodeJS.Timeout;
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class AiChatStreamRegistryService implements OnModuleDestroy {
|
||||
private readonly logger = new Logger(AiChatStreamRegistryService.name);
|
||||
private readonly entries = new Map<string, Entry>(); // key: chatId
|
||||
|
||||
/**
|
||||
* Register a fresh entry at the START of a run (before any frame), so a tab
|
||||
* that attaches in the begin->seed window finds an entry to wait on. If an
|
||||
* entry already exists for this chat (a previous, possibly still-live run whose
|
||||
* tee loop is draining), it is terminated MIRRORING the done-path (invariant 3)
|
||||
* so its subscribers are released and its retention timer is cleared; a late
|
||||
* `done` from that old tee then fires against the closed-over old reference and,
|
||||
* thanks to identity checks, never touches this new entry.
|
||||
*/
|
||||
open(chatId: string, runId: string): void {
|
||||
const existing = this.entries.get(chatId);
|
||||
if (existing) {
|
||||
if (existing.retainTimer) {
|
||||
clearTimeout(existing.retainTimer);
|
||||
existing.retainTimer = undefined;
|
||||
}
|
||||
// Started subscribers get exactly one onEnd() and are removed; paused ones
|
||||
// are marked pendingEnd (their start() will end them). finished=true guards
|
||||
// any later done from the old tee loop from double-notifying.
|
||||
this.terminateSubscribers(existing);
|
||||
}
|
||||
this.entries.set(chatId, {
|
||||
runId,
|
||||
frames: [],
|
||||
bytes: 0,
|
||||
overflowed: false,
|
||||
finished: false,
|
||||
subscribers: new Set<Subscriber>(),
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Tee a run's SSE frame stream into its entry (called from consumeSseStream).
|
||||
* No-op with a warning when there is no entry or the entry belongs to a
|
||||
* different run (invariant 1). The reader loop is fire-and-forget: the tee
|
||||
* branch outlives the client socket by design.
|
||||
*/
|
||||
bind(
|
||||
chatId: string,
|
||||
runId: string,
|
||||
assistantMessageId: string | undefined,
|
||||
stream: ReadableStream<string>,
|
||||
): void {
|
||||
const entry = this.entries.get(chatId);
|
||||
if (!entry || entry.runId !== runId) {
|
||||
// Invariant 1: only the matching run may mutate the entry.
|
||||
this.logger.warn(
|
||||
`bind: no matching run-stream entry for chat=${chatId} run=${runId}`,
|
||||
);
|
||||
return;
|
||||
}
|
||||
entry.assistantMessageId = assistantMessageId;
|
||||
const reader = stream.getReader();
|
||||
const pump = async (): Promise<void> => {
|
||||
try {
|
||||
for (;;) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
this.ingestFrame(entry, value);
|
||||
}
|
||||
this.finalizeEntry(chatId, entry);
|
||||
} catch {
|
||||
// A read error is a terminal event too — release subscribers.
|
||||
this.finalizeEntry(chatId, entry);
|
||||
}
|
||||
};
|
||||
void pump();
|
||||
}
|
||||
|
||||
/**
|
||||
* Terminate a run's entry from the OUTER catch of the stream method (a failure
|
||||
* before/while wiring the pipe, so `done` will never arrive). Identity-checked
|
||||
* on runId (invariant 1); the shared terminal path is idempotent.
|
||||
*/
|
||||
abortEntry(chatId: string, runId: string): void {
|
||||
const entry = this.entries.get(chatId);
|
||||
if (!entry || entry.runId !== runId) return;
|
||||
this.finalizeEntry(chatId, entry);
|
||||
}
|
||||
|
||||
/**
|
||||
* Attach to a run's stream. Async only for the phase-2 Redis seam — the body
|
||||
* runs synchronously so the replay snapshot and the subscriber registration
|
||||
* happen in ONE tick with no await between them (invariant 4): a frame ingested
|
||||
* concurrently cannot slip into the gap and be lost or duplicated.
|
||||
*
|
||||
* Returns null (-> the caller answers 204) when:
|
||||
* - there is no entry, or it overflowed (replay is gone);
|
||||
* - expect=live with an anchor that does not match this run's assistant id
|
||||
* (invariant 6: a stripped tab must never replay a FOREIGN run's transcript);
|
||||
* - the run finished and the caller did not expect a live tail.
|
||||
* A finished run with expect=live yields a replay-only attachment (no
|
||||
* subscriber registered). Otherwise a paused subscriber is registered and the
|
||||
* caller replays `replay`, then calls start() to drain and go live.
|
||||
*/
|
||||
async attach(
|
||||
chatId: string,
|
||||
expectLive: boolean,
|
||||
anchor: string | undefined,
|
||||
cb: RunStreamCallbacks,
|
||||
): Promise<RunStreamAttachment | null> {
|
||||
const entry = this.entries.get(chatId);
|
||||
if (!entry || entry.overflowed) return null;
|
||||
// Invariant 6: cross-run replay is forbidden. Before bind, assistantMessageId
|
||||
// is undefined and mismatches any anchor -> 204 -> client restore+poll path.
|
||||
if (expectLive && anchor && entry.assistantMessageId !== anchor) return null;
|
||||
if (entry.finished && !expectLive) return null;
|
||||
if (entry.finished && expectLive) {
|
||||
// Replay-only: the run is done, no subscriber is registered.
|
||||
return {
|
||||
replay: entry.frames.slice(),
|
||||
finished: true,
|
||||
start: () => undefined,
|
||||
unsubscribe: () => undefined,
|
||||
};
|
||||
}
|
||||
|
||||
const sub: Subscriber = {
|
||||
onFrame: cb.onFrame,
|
||||
onEnd: cb.onEnd,
|
||||
started: false,
|
||||
pending: [],
|
||||
pendingBytes: 0,
|
||||
overflowed: false,
|
||||
pendingEnd: false,
|
||||
};
|
||||
entry.subscribers.add(sub);
|
||||
// Snapshot in the SAME synchronous block as the registration (invariant 4).
|
||||
const replay = entry.frames.slice();
|
||||
// CONTRACT: the caller MUST call start() in the SAME tick as this attach()
|
||||
// returns — no await between them. While a subscriber is paused, every frame
|
||||
// is buffered in sub.pending; a delayed start() lets a whole run accumulate
|
||||
// there. The pendingBytes cap (see ingestFrame) is the structural backstop if
|
||||
// that contract is ever broken (e.g. the phase-2 Redis await seam).
|
||||
return {
|
||||
replay,
|
||||
finished: false,
|
||||
start: () => {
|
||||
if (sub.overflowed) {
|
||||
// The pending buffer overflowed while paused: end the stream instead of
|
||||
// replaying a partial (a 204-equivalent post-attach degrade).
|
||||
try {
|
||||
sub.onEnd();
|
||||
} catch {
|
||||
// The socket is gone; nothing to end.
|
||||
}
|
||||
entry.subscribers.delete(sub);
|
||||
return;
|
||||
}
|
||||
// Deliver frames buffered while paused, in order, then go live.
|
||||
for (const frame of sub.pending) {
|
||||
try {
|
||||
sub.onFrame(frame);
|
||||
} catch {
|
||||
entry.subscribers.delete(sub);
|
||||
return;
|
||||
}
|
||||
}
|
||||
sub.pending = [];
|
||||
sub.started = true;
|
||||
if (sub.pendingEnd) {
|
||||
try {
|
||||
sub.onEnd();
|
||||
} catch {
|
||||
// The socket is gone; nothing to end.
|
||||
}
|
||||
entry.subscribers.delete(sub);
|
||||
}
|
||||
},
|
||||
unsubscribe: () => {
|
||||
entry.subscribers.delete(sub);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
onModuleDestroy(): void {
|
||||
for (const entry of this.entries.values()) {
|
||||
if (entry.retainTimer) clearTimeout(entry.retainTimer);
|
||||
}
|
||||
this.entries.clear();
|
||||
}
|
||||
|
||||
/** Buffer + fan-out a single frame. See invariant/overflow semantics inline. */
|
||||
private ingestFrame(entry: Entry, frame: string): void {
|
||||
entry.bytes += Buffer.byteLength(frame);
|
||||
if (!entry.overflowed) {
|
||||
entry.frames.push(frame);
|
||||
if (entry.bytes > RUN_STREAM_MAX_BUFFER_BYTES) {
|
||||
// The crossing frame was already counted AND (below) fanned out; only the
|
||||
// replay buffer is dropped. After overflow no more frames are buffered,
|
||||
// but live fan-out continues.
|
||||
entry.overflowed = true;
|
||||
entry.frames = [];
|
||||
this.logger.warn(
|
||||
`run-stream buffer overflow for run=${entry.runId}; ` +
|
||||
`late attach will 204 until the run ends`,
|
||||
);
|
||||
}
|
||||
}
|
||||
for (const sub of entry.subscribers) {
|
||||
if (sub.started) {
|
||||
try {
|
||||
sub.onFrame(frame);
|
||||
} catch {
|
||||
entry.subscribers.delete(sub);
|
||||
}
|
||||
} else {
|
||||
sub.pending.push(frame);
|
||||
sub.pendingBytes += Buffer.byteLength(frame);
|
||||
if (sub.pendingBytes > SUBSCRIBER_MAX_BUFFERED_BYTES) {
|
||||
// The paused subscriber's buffer overflowed — only possible if start()
|
||||
// was delayed past the same-tick contract (the phase-2 await seam).
|
||||
// Drop it rather than buffer the whole run; on start() it degrades to an
|
||||
// immediate end (a 204-equivalent) instead of replaying a partial.
|
||||
sub.overflowed = true;
|
||||
sub.pending = [];
|
||||
entry.subscribers.delete(sub);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Shared terminal path for done / read-error / external-abort. Idempotent: a
|
||||
* second call (already finished) is a no-op, so an open()-replaced or
|
||||
* abort-then-done entry is never double-armed or double-ended.
|
||||
*/
|
||||
private finalizeEntry(chatId: string, entry: Entry): void {
|
||||
if (entry.finished) return;
|
||||
this.terminateSubscribers(entry);
|
||||
const timer = setTimeout(() => {
|
||||
// Invariant 2: only delete OUR entry (a replacement may already own the key).
|
||||
if (this.entries.get(chatId) === entry) this.entries.delete(chatId);
|
||||
}, RUN_STREAM_RETAIN_FINISHED_MS);
|
||||
timer.unref?.();
|
||||
entry.retainTimer = timer;
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark the entry finished and release its subscribers, mirroring the done-path:
|
||||
* started subscribers get exactly one onEnd() and are removed; paused ones are
|
||||
* flagged pendingEnd so their start() ends them. Deleting the current element
|
||||
* during Set iteration is safe.
|
||||
*/
|
||||
private terminateSubscribers(entry: Entry): void {
|
||||
entry.finished = true;
|
||||
for (const sub of entry.subscribers) {
|
||||
if (sub.started) {
|
||||
try {
|
||||
sub.onEnd();
|
||||
} catch {
|
||||
// The socket is gone; nothing to end.
|
||||
}
|
||||
entry.subscribers.delete(sub);
|
||||
} else {
|
||||
sub.pendingEnd = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,388 @@
|
||||
import {
|
||||
AiChatStreamRegistryService,
|
||||
RUN_STREAM_MAX_BUFFER_BYTES,
|
||||
RUN_STREAM_RETAIN_FINISHED_MS,
|
||||
RunStreamCallbacks,
|
||||
} from './ai-chat-stream-registry.service';
|
||||
|
||||
/**
|
||||
* Unit tests for the in-memory run-stream registry (#184 phase 1.5). The registry
|
||||
* is the whole of the resumable-transport contract: replay ordering, paused ->
|
||||
* live hand-off, overflow, retention, the anchor check (invariant 6), and the
|
||||
* mirror-the-done-path replace semantics (invariant 3). Every enumerated case in
|
||||
* the issue's task 1.5 has a test here.
|
||||
*/
|
||||
|
||||
// A ReadableStream whose frames the test pushes explicitly, plus close/error.
|
||||
function makePushStream(): {
|
||||
stream: ReadableStream<string>;
|
||||
push: (f: string) => void;
|
||||
close: () => void;
|
||||
error: (e?: unknown) => void;
|
||||
} {
|
||||
let controller!: ReadableStreamDefaultController<string>;
|
||||
const stream = new ReadableStream<string>({
|
||||
start(c) {
|
||||
controller = c;
|
||||
},
|
||||
});
|
||||
return {
|
||||
stream,
|
||||
push: (f) => controller.enqueue(f),
|
||||
close: () => controller.close(),
|
||||
error: (e) => controller.error(e ?? new Error('read error')),
|
||||
};
|
||||
}
|
||||
|
||||
// Let the fire-and-forget pump drain queued frames (reader.read() resolves on a
|
||||
// macrotask boundary for an already-enqueued value).
|
||||
const flush = (): Promise<void> => new Promise((r) => setTimeout(r, 0));
|
||||
|
||||
function collector(): {
|
||||
cb: RunStreamCallbacks;
|
||||
frames: string[];
|
||||
ended: () => number;
|
||||
} {
|
||||
const frames: string[] = [];
|
||||
let ends = 0;
|
||||
return {
|
||||
frames,
|
||||
ended: () => ends,
|
||||
cb: {
|
||||
onFrame: (f) => frames.push(f),
|
||||
onEnd: () => {
|
||||
ends += 1;
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
describe('AiChatStreamRegistryService', () => {
|
||||
const CHAT = 'chat-1';
|
||||
let registry: AiChatStreamRegistryService;
|
||||
|
||||
beforeEach(() => {
|
||||
registry = new AiChatStreamRegistryService();
|
||||
jest.spyOn((registry as any).logger, 'warn').mockImplementation(() => {});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
registry.onModuleDestroy();
|
||||
});
|
||||
|
||||
it('replays frames in arrival order (live attach)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
src.push('b');
|
||||
src.push('c');
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = await registry.attach(CHAT, false, undefined, c.cb);
|
||||
expect(att).not.toBeNull();
|
||||
expect(att!.replay).toEqual(['a', 'b', 'c']);
|
||||
expect(att!.finished).toBe(false);
|
||||
});
|
||||
|
||||
it('late attach gets the full prefix as replay plus the live tail', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
src.push('b');
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
|
||||
expect(att.replay).toEqual(['a', 'b']);
|
||||
att.start();
|
||||
// Live tail arrives after start().
|
||||
src.push('c');
|
||||
src.push('d');
|
||||
await flush();
|
||||
expect(c.frames).toEqual(['c', 'd']);
|
||||
});
|
||||
|
||||
it('a paused subscriber receives frames buffered during pause in order, then live (no loss/reorder)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
// Attach (paused). Frames that arrive BEFORE start() must queue, not drop.
|
||||
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
|
||||
expect(att.replay).toEqual(['a']);
|
||||
src.push('b'); // arrives while paused -> pending
|
||||
src.push('c');
|
||||
await flush();
|
||||
expect(c.frames).toEqual([]); // nothing delivered yet (paused)
|
||||
att.start(); // drains pending in order
|
||||
expect(c.frames).toEqual(['b', 'c']);
|
||||
src.push('d'); // now live
|
||||
await flush();
|
||||
expect(c.frames).toEqual(['b', 'c', 'd']);
|
||||
});
|
||||
|
||||
it('a run that finishes while a subscriber is paused ends it on start()', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
|
||||
// Terminate the run while the subscriber is still paused.
|
||||
registry.abortEntry(CHAT, 'run-1');
|
||||
expect(c.ended()).toBe(0); // paused: not ended yet
|
||||
att.start();
|
||||
expect(c.ended()).toBe(1); // start() drains + ends
|
||||
});
|
||||
|
||||
it('finished + expect=live returns a replay WITHOUT registering a subscriber', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
src.push('b');
|
||||
src.close();
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, true, undefined, c.cb))!;
|
||||
expect(att.finished).toBe(true);
|
||||
expect(att.replay).toEqual(['a', 'b']);
|
||||
// No subscriber registered: start()/unsubscribe are no-ops and the entry has
|
||||
// zero subscribers.
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
expect(entry.subscribers.size).toBe(0);
|
||||
att.start();
|
||||
expect(c.frames).toEqual([]);
|
||||
});
|
||||
|
||||
it('finished WITHOUT expect=live returns null', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
src.close();
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
expect(await registry.attach(CHAT, false, undefined, c.cb)).toBeNull();
|
||||
});
|
||||
|
||||
it('anchor mismatch with expect=live returns null (and null before bind sets assistantMessageId)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const c = collector();
|
||||
// Before bind: assistantMessageId is undefined -> mismatches any anchor.
|
||||
expect(
|
||||
await registry.attach(CHAT, true, 'assist-1', c.cb),
|
||||
).toBeNull();
|
||||
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
await flush();
|
||||
// Wrong anchor -> null (cross-run replay forbidden, invariant 6).
|
||||
expect(await registry.attach(CHAT, true, 'other-id', c.cb)).toBeNull();
|
||||
});
|
||||
|
||||
it('matching anchor with expect=live attaches', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = await registry.attach(CHAT, true, 'assist-1', c.cb);
|
||||
expect(att).not.toBeNull();
|
||||
expect(att!.replay).toEqual(['a']);
|
||||
});
|
||||
|
||||
it('overflow: attach returns null, but the LIVE subscriber keeps receiving (incl. the crossing frame)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
|
||||
// A live (started) subscriber attached before the flood.
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
|
||||
att.start();
|
||||
|
||||
const oneMb = 'x'.repeat(1024 * 1024);
|
||||
// 5 x 1MB = 5MB > 4MB cap; the 5th frame is the one that crosses.
|
||||
for (let i = 0; i < 5; i++) src.push(oneMb + i);
|
||||
await flush();
|
||||
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
expect(entry.overflowed).toBe(true);
|
||||
expect(entry.bytes).toBeGreaterThan(RUN_STREAM_MAX_BUFFER_BYTES);
|
||||
// The live subscriber received ALL 5 frames, including the crossing one.
|
||||
expect(c.frames).toHaveLength(5);
|
||||
expect(c.frames[4]).toBe(oneMb + 4);
|
||||
|
||||
// A NEW attach after overflow gets null (replay buffer is gone).
|
||||
const c2 = collector();
|
||||
expect(await registry.attach(CHAT, false, undefined, c2.cb)).toBeNull();
|
||||
});
|
||||
|
||||
it('a paused subscriber whose pending buffer overflows is dropped and ends on start(); other subscribers keep receiving', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
|
||||
// A: paused (start() deliberately delayed to simulate the phase-2 await seam).
|
||||
const a = collector();
|
||||
const attA = (await registry.attach(CHAT, false, undefined, a.cb))!;
|
||||
// B: live (started) — its delivery must be unaffected by A's overflow.
|
||||
const b = collector();
|
||||
const attB = (await registry.attach(CHAT, false, undefined, b.cb))!;
|
||||
attB.start();
|
||||
|
||||
const oneMb = 'x'.repeat(1024 * 1024);
|
||||
// 9 x 1MB = 9MB > 8MB per-subscriber cap; A's pending overflows, B streams live.
|
||||
for (let i = 0; i < 9; i++) src.push(oneMb + i);
|
||||
await flush();
|
||||
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
// A was dropped from the subscriber set on overflow; B (started) remains.
|
||||
expect(entry.subscribers.size).toBe(1);
|
||||
expect(a.frames).toEqual([]); // paused + overflowed: nothing was delivered
|
||||
// B received every frame live (delivery unaffected by A's overflow).
|
||||
expect(b.frames).toHaveLength(9);
|
||||
|
||||
// A's start() (arriving late) degrades to an immediate end, not a partial replay.
|
||||
attA.start();
|
||||
expect(a.frames).toEqual([]);
|
||||
expect(a.ended()).toBe(1);
|
||||
});
|
||||
|
||||
it('open() over a LIVE entry ends started subscribers exactly once and a late done does not touch the new entry (invariant 3)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
|
||||
att.start(); // started subscriber on run-1
|
||||
|
||||
// run-2 starts on the same chat while run-1's tee is still reading.
|
||||
registry.open(CHAT, 'run-2');
|
||||
expect(c.ended()).toBe(1); // exactly one onEnd from the replace
|
||||
|
||||
const newEntry = (registry as any).entries.get(CHAT);
|
||||
expect(newEntry.runId).toBe('run-2');
|
||||
expect(newEntry.finished).toBe(false);
|
||||
|
||||
// The old tee now completes: its late done must NOT double-end nor delete the
|
||||
// new entry.
|
||||
src.push('b');
|
||||
src.close();
|
||||
await flush();
|
||||
expect(c.ended()).toBe(1); // still exactly one
|
||||
const still = (registry as any).entries.get(CHAT);
|
||||
expect(still).toBe(newEntry);
|
||||
expect(still.runId).toBe('run-2');
|
||||
});
|
||||
|
||||
it('bind with a foreign runId is a no-op (invariant 1)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'WRONG-run', 'assist-x', src.stream);
|
||||
src.push('a');
|
||||
await flush();
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
// Frames were NOT ingested (bind bailed), assistantMessageId untouched.
|
||||
expect(entry.frames).toEqual([]);
|
||||
expect(entry.assistantMessageId).toBeUndefined();
|
||||
});
|
||||
|
||||
it('abortEntry with a foreign runId is a no-op (invariant 1)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
registry.abortEntry(CHAT, 'WRONG-run');
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
expect(entry.finished).toBe(false);
|
||||
});
|
||||
|
||||
it('a throwing onFrame ejects only that subscriber; the ingest loop stays alive', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
|
||||
const bad = collector();
|
||||
const badAtt = (await registry.attach(CHAT, false, undefined, {
|
||||
onFrame: () => {
|
||||
throw new Error('boom');
|
||||
},
|
||||
onEnd: bad.cb.onEnd,
|
||||
}))!;
|
||||
badAtt.start();
|
||||
|
||||
const good = collector();
|
||||
const goodAtt = (await registry.attach(CHAT, false, undefined, good.cb))!;
|
||||
goodAtt.start();
|
||||
|
||||
src.push('a'); // bad throws on this frame -> ejected
|
||||
src.push('b'); // good still receives both
|
||||
await flush();
|
||||
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
expect(entry.subscribers.size).toBe(1); // bad ejected, good remains
|
||||
expect(good.frames).toEqual(['a', 'b']);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* Retention + replace timer behavior. Fake timers, and entries are finalized via
|
||||
* the synchronous abortEntry() path so no stream pump / microtask juggling is
|
||||
* needed.
|
||||
*/
|
||||
describe('AiChatStreamRegistryService retention timers', () => {
|
||||
const CHAT = 'chat-r';
|
||||
let registry: AiChatStreamRegistryService;
|
||||
|
||||
beforeEach(() => {
|
||||
jest.useFakeTimers();
|
||||
registry = new AiChatStreamRegistryService();
|
||||
jest.spyOn((registry as any).logger, 'warn').mockImplementation(() => {});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
registry.onModuleDestroy();
|
||||
jest.useRealTimers();
|
||||
});
|
||||
|
||||
it('a finished entry is removed after the retention window', () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
registry.abortEntry(CHAT, 'run-1'); // finalize -> retention armed
|
||||
expect((registry as any).entries.get(CHAT)).toBeDefined();
|
||||
jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1);
|
||||
expect((registry as any).entries.get(CHAT)).toBeUndefined();
|
||||
});
|
||||
|
||||
it('retention deletes ONLY its own entry (invariant 2)', () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
registry.abortEntry(CHAT, 'run-1'); // arm retention for entry A
|
||||
// Simulate the race where the key was replaced without clearing A's timer.
|
||||
const sentinel = { marker: true };
|
||||
(registry as any).entries.set(CHAT, sentinel);
|
||||
jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1);
|
||||
// A's timer saw entries.get(CHAT) !== A, so it did NOT delete the successor.
|
||||
expect((registry as any).entries.get(CHAT)).toBe(sentinel);
|
||||
});
|
||||
|
||||
it('open() over a retained entry clears its timer and the successor survives', () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
registry.abortEntry(CHAT, 'run-1'); // retained, timer armed
|
||||
const clearSpy = jest.spyOn(global, 'clearTimeout');
|
||||
registry.open(CHAT, 'run-2'); // must clear run-1's retain timer
|
||||
expect(clearSpy).toHaveBeenCalled();
|
||||
jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1);
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
expect(entry).toBeDefined();
|
||||
expect(entry.runId).toBe('run-2');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,423 @@
|
||||
import { ForbiddenException } from '@nestjs/common';
|
||||
import { AiChatController } from './ai-chat.controller';
|
||||
import type {
|
||||
RunStreamAttachment,
|
||||
RunStreamCallbacks,
|
||||
} from './ai-chat-stream-registry.service';
|
||||
import { SUBSCRIBER_MAX_BUFFERED_BYTES } from './ai-chat-stream-registry.service';
|
||||
import type { User, Workspace } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
* Wiring spec for the #184 phase 1.5 attach endpoint
|
||||
* (`GET /ai-chat/runs/:chatId/stream`). Owner-gated via assertOwnedChat; the
|
||||
* registry is mocked so this exercises ONLY the controller's replay/live/204/
|
||||
* cleanup wiring against a fake raw socket. Constructor order is (aiChatService,
|
||||
* aiChatRunService, aiChatRepo, aiChatMessageRepo, aiTranscription, pageRepo,
|
||||
* streamRegistry, environment).
|
||||
*/
|
||||
describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
|
||||
const user = { id: 'u1' } as User;
|
||||
const workspace = { id: 'ws1' } as Workspace;
|
||||
|
||||
function makeRawRes() {
|
||||
const raw: any = {
|
||||
writableEnded: false,
|
||||
writableLength: 0,
|
||||
destroyed: false,
|
||||
written: [] as string[],
|
||||
head: null as any,
|
||||
write: jest.fn((f: string) => {
|
||||
raw.written.push(f);
|
||||
return true;
|
||||
}),
|
||||
writeHead: jest.fn((code: number, headers: any) => {
|
||||
raw.head = { code, headers };
|
||||
return raw;
|
||||
}),
|
||||
flushHeaders: jest.fn(),
|
||||
end: jest.fn(() => {
|
||||
raw.writableEnded = true;
|
||||
}),
|
||||
destroy: jest.fn(() => {
|
||||
raw.destroyed = true;
|
||||
}),
|
||||
on: jest.fn(),
|
||||
once: jest.fn(),
|
||||
};
|
||||
const res: any = {
|
||||
raw,
|
||||
status: jest.fn(() => res),
|
||||
send: jest.fn(),
|
||||
hijack: jest.fn(),
|
||||
};
|
||||
return { res, raw };
|
||||
}
|
||||
|
||||
function makeReq(destroyed = false) {
|
||||
const handlers: Record<string, () => void> = {};
|
||||
const raw: any = {
|
||||
destroyed,
|
||||
once: jest.fn((ev: string, fn: () => void) => {
|
||||
handlers[ev] = fn;
|
||||
}),
|
||||
};
|
||||
return { req: { raw } as any, raw, fireClose: () => handlers['close']?.() };
|
||||
}
|
||||
|
||||
function makeAttachment(
|
||||
over: Partial<RunStreamAttachment> = {},
|
||||
): RunStreamAttachment {
|
||||
return {
|
||||
replay: [],
|
||||
finished: false,
|
||||
start: jest.fn(),
|
||||
unsubscribe: jest.fn(),
|
||||
...over,
|
||||
};
|
||||
}
|
||||
|
||||
function makeController(opts: {
|
||||
chat?: unknown;
|
||||
attachment?: RunStreamAttachment | null;
|
||||
}) {
|
||||
const aiChatRepo = { findById: jest.fn().mockResolvedValue(opts.chat) };
|
||||
let capturedCb: RunStreamCallbacks | undefined;
|
||||
const streamRegistry = {
|
||||
attach: jest.fn(
|
||||
(
|
||||
_chatId: string,
|
||||
_live: boolean,
|
||||
_anchor: string | undefined,
|
||||
cb: RunStreamCallbacks,
|
||||
) => {
|
||||
capturedCb = cb;
|
||||
return Promise.resolve(
|
||||
opts.attachment === undefined ? makeAttachment() : opts.attachment,
|
||||
);
|
||||
},
|
||||
),
|
||||
};
|
||||
const environment = { isAiChatResumableStreamEnabled: () => true };
|
||||
const controller = new AiChatController(
|
||||
{} as never, // aiChatService
|
||||
{} as never, // aiChatRunService
|
||||
aiChatRepo as never,
|
||||
{} as never, // aiChatMessageRepo
|
||||
{} as never, // aiTranscription
|
||||
{} as never, // pageRepo
|
||||
streamRegistry as never,
|
||||
environment as never,
|
||||
);
|
||||
return {
|
||||
controller,
|
||||
aiChatRepo,
|
||||
streamRegistry,
|
||||
getCb: () => capturedCb!,
|
||||
};
|
||||
}
|
||||
|
||||
const owned = { id: 'c1', creatorId: 'u1' };
|
||||
|
||||
it('owner-gates: a foreign chat throws ForbiddenException and never attaches', async () => {
|
||||
const { controller, streamRegistry } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'someone-else' },
|
||||
});
|
||||
const { res } = makeRawRes();
|
||||
const { req } = makeReq();
|
||||
await expect(
|
||||
controller.attachRunStream(
|
||||
'c1',
|
||||
undefined,
|
||||
undefined,
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
workspace,
|
||||
),
|
||||
).rejects.toBeInstanceOf(ForbiddenException);
|
||||
expect(streamRegistry.attach).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('answers 204 when the registry has nothing to resume (no entry / finished / anchor-mismatch)', async () => {
|
||||
const { controller } = makeController({ chat: owned, attachment: null });
|
||||
const { res } = makeRawRes();
|
||||
const { req } = makeReq();
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
undefined,
|
||||
undefined,
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
expect(res.status).toHaveBeenCalledWith(204);
|
||||
expect(res.send).toHaveBeenCalled();
|
||||
expect(res.hijack).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('threads expect=live and anchor through to the registry', async () => {
|
||||
const { controller, streamRegistry } = makeController({
|
||||
chat: owned,
|
||||
attachment: null,
|
||||
});
|
||||
const { res } = makeRawRes();
|
||||
const { req } = makeReq();
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
'live',
|
||||
'anchor-1',
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
expect(streamRegistry.attach).toHaveBeenCalledWith(
|
||||
'c1',
|
||||
true,
|
||||
'anchor-1',
|
||||
expect.anything(),
|
||||
);
|
||||
});
|
||||
|
||||
it('passes expect=false when the query is absent', async () => {
|
||||
const { controller, streamRegistry } = makeController({
|
||||
chat: owned,
|
||||
attachment: null,
|
||||
});
|
||||
const { res } = makeRawRes();
|
||||
const { req } = makeReq();
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
undefined,
|
||||
undefined,
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
expect(streamRegistry.attach).toHaveBeenCalledWith(
|
||||
'c1',
|
||||
false,
|
||||
undefined,
|
||||
expect.anything(),
|
||||
);
|
||||
});
|
||||
|
||||
it('hijacks, writes headers + replay, registers a close cleanup, then goes live', async () => {
|
||||
const start = jest.fn();
|
||||
const attachment = makeAttachment({
|
||||
replay: ['f1', 'f2'],
|
||||
finished: false,
|
||||
start,
|
||||
});
|
||||
const { controller } = makeController({ chat: owned, attachment });
|
||||
const { res, raw } = makeRawRes();
|
||||
const { req } = makeReq();
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
undefined,
|
||||
undefined,
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
expect(res.hijack).toHaveBeenCalled();
|
||||
expect(raw.writeHead).toHaveBeenCalledWith(
|
||||
200,
|
||||
expect.objectContaining({ 'content-type': 'text/event-stream' }),
|
||||
);
|
||||
expect(raw.written).toEqual(['f1', 'f2']); // replay
|
||||
expect(start).toHaveBeenCalled(); // go live after replay
|
||||
expect(req.raw.once).toHaveBeenCalledWith('close', expect.any(Function));
|
||||
});
|
||||
|
||||
it('finished replay ends the response immediately without going live', async () => {
|
||||
const start = jest.fn();
|
||||
const attachment = makeAttachment({
|
||||
replay: ['f1'],
|
||||
finished: true,
|
||||
start,
|
||||
});
|
||||
const { controller } = makeController({ chat: owned, attachment });
|
||||
const { res, raw } = makeRawRes();
|
||||
const { req } = makeReq();
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
'live',
|
||||
'a1',
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
expect(raw.written).toEqual(['f1']);
|
||||
expect(raw.end).toHaveBeenCalled();
|
||||
expect(start).not.toHaveBeenCalled(); // finished -> returns before start()
|
||||
});
|
||||
|
||||
it('a close during the awaits (req.raw.destroyed) unsubscribes and writes nothing', async () => {
|
||||
const attachment = makeAttachment({ replay: ['f1'] });
|
||||
const { controller } = makeController({ chat: owned, attachment });
|
||||
const { res, raw } = makeRawRes();
|
||||
const { req } = makeReq(true); // destroyed already at registration time
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
undefined,
|
||||
undefined,
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
expect(attachment.unsubscribe).toHaveBeenCalled();
|
||||
expect(raw.writeHead).not.toHaveBeenCalled();
|
||||
expect(raw.written).toEqual([]);
|
||||
});
|
||||
|
||||
it('the registered close handler unsubscribes the attachment', async () => {
|
||||
const attachment = makeAttachment({ replay: [] });
|
||||
const { controller } = makeController({ chat: owned, attachment });
|
||||
const { res } = makeRawRes();
|
||||
const { req, fireClose } = makeReq();
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
undefined,
|
||||
undefined,
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
expect(attachment.unsubscribe).not.toHaveBeenCalled();
|
||||
fireClose(); // socket closed
|
||||
expect(attachment.unsubscribe).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('onFrame destroys the socket when the buffered length exceeds the cap', async () => {
|
||||
const attachment = makeAttachment({ replay: [] });
|
||||
const { controller, getCb } = makeController({ chat: owned, attachment });
|
||||
const { res, raw } = makeRawRes();
|
||||
const { req } = makeReq();
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
undefined,
|
||||
undefined,
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
const cb = getCb();
|
||||
// Normal live frame writes.
|
||||
raw.writableLength = 0;
|
||||
cb.onFrame('live-1');
|
||||
expect(raw.written).toContain('live-1');
|
||||
// A stalled socket over the cap is destroyed instead of buffering.
|
||||
raw.writableLength = SUBSCRIBER_MAX_BUFFERED_BYTES + 1;
|
||||
raw.write.mockClear();
|
||||
cb.onFrame('too-much');
|
||||
expect(raw.destroy).toHaveBeenCalled();
|
||||
expect(raw.write).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* The begin-hook `open()` flag gate (#184 phase 1.5). `open()` lives ONLY in the
|
||||
* stream() begin-hook, gated on the resumable flag. If it regressed, a flag-off
|
||||
* turn would create an EMPTY registry entry (never bound, never finished) and a
|
||||
* later attach would find a non-null paused attachment -> a hung SSE that never
|
||||
* gets a frame and never ends, instead of a clean 204. These drive stream() only
|
||||
* far enough to capture the runHooks it hands to the service, then invoke the
|
||||
* begin-hook and assert whether the registry was opened.
|
||||
*/
|
||||
describe('AiChatController begin-hook open() flag gate (#184 phase 1.5)', () => {
|
||||
const user = { id: 'u1' } as User;
|
||||
const workspace = {
|
||||
id: 'ws1',
|
||||
settings: { ai: { chat: true, autonomousRuns: true } },
|
||||
} as unknown as Workspace;
|
||||
|
||||
function makeController(opts: { resumable: boolean }) {
|
||||
let capturedArgs: any;
|
||||
const aiChatService = {
|
||||
resolveRoleForRequest: jest.fn(async () => null),
|
||||
getChatModel: jest.fn(async () => ({})),
|
||||
stream: jest.fn(async (args: any) => {
|
||||
capturedArgs = args;
|
||||
}),
|
||||
};
|
||||
const aiChatRunService = {
|
||||
getActiveForChat: jest.fn(async () => undefined),
|
||||
beginRun: jest.fn(async () => ({
|
||||
runId: 'run-1',
|
||||
signal: new AbortController().signal,
|
||||
})),
|
||||
};
|
||||
const streamRegistry = { open: jest.fn(), attach: jest.fn() };
|
||||
const environment = {
|
||||
isAiChatResumableStreamEnabled: () => opts.resumable,
|
||||
};
|
||||
const controller = new AiChatController(
|
||||
aiChatService as never,
|
||||
aiChatRunService as never,
|
||||
{} as never, // aiChatRepo
|
||||
{} as never, // aiChatMessageRepo
|
||||
{} as never, // aiTranscription
|
||||
{} as never, // pageRepo
|
||||
streamRegistry as never,
|
||||
environment as never,
|
||||
);
|
||||
return {
|
||||
controller,
|
||||
streamRegistry,
|
||||
aiChatRunService,
|
||||
getRunHooks: () => capturedArgs?.runHooks,
|
||||
};
|
||||
}
|
||||
|
||||
function makeReqRes() {
|
||||
const req: any = {
|
||||
raw: { sessionId: 'sess-1', once: jest.fn() },
|
||||
body: {
|
||||
messages: [
|
||||
{ id: 'm1', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
|
||||
],
|
||||
},
|
||||
};
|
||||
const res: any = {
|
||||
raw: { once: jest.fn(), on: jest.fn(), headersSent: false, writableEnded: false },
|
||||
hijack: jest.fn(),
|
||||
};
|
||||
return { req, res };
|
||||
}
|
||||
|
||||
it('flag OFF: the begin-hook does NOT open a registry entry (no hung empty entry)', async () => {
|
||||
const { controller, streamRegistry, aiChatRunService, getRunHooks } =
|
||||
makeController({ resumable: false });
|
||||
const { req, res } = makeReqRes();
|
||||
await controller.stream(req, res, user, workspace);
|
||||
|
||||
const runHooks = getRunHooks();
|
||||
expect(runHooks).toBeDefined();
|
||||
const handle = await runHooks.begin('chat-1');
|
||||
// The run still begins (the durable-run feature is independent of resume)...
|
||||
expect(aiChatRunService.beginRun).toHaveBeenCalled();
|
||||
expect(handle).toEqual({ runId: 'run-1', signal: expect.anything() });
|
||||
// ...but with the flag off the registry entry is NEVER opened.
|
||||
expect(streamRegistry.open).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('flag ON: the begin-hook opens the registry entry with (chatId, runId)', async () => {
|
||||
const { controller, streamRegistry, getRunHooks } = makeController({
|
||||
resumable: true,
|
||||
});
|
||||
const { req, res } = makeReqRes();
|
||||
await controller.stream(req, res, user, workspace);
|
||||
|
||||
const runHooks = getRunHooks();
|
||||
await runHooks.begin('chat-1');
|
||||
expect(streamRegistry.open).toHaveBeenCalledWith('chat-1', 'run-1');
|
||||
});
|
||||
});
|
||||
@@ -4,11 +4,15 @@ import {
|
||||
ConflictException,
|
||||
Controller,
|
||||
ForbiddenException,
|
||||
Get,
|
||||
HttpCode,
|
||||
HttpException,
|
||||
HttpStatus,
|
||||
Logger,
|
||||
Param,
|
||||
ParseUUIDPipe,
|
||||
Post,
|
||||
Query,
|
||||
Req,
|
||||
Res,
|
||||
ServiceUnavailableException,
|
||||
@@ -54,6 +58,12 @@ import {
|
||||
} from './dto/ai-chat.dto';
|
||||
import { describeProviderError } from '../../integrations/ai/ai-error.util';
|
||||
import { buildChatMarkdown } from './chat-markdown.util';
|
||||
import {
|
||||
AiChatStreamRegistryService,
|
||||
SUBSCRIBER_MAX_BUFFERED_BYTES,
|
||||
} from './ai-chat-stream-registry.service';
|
||||
import { startSseHeartbeat } from './sse-resilience';
|
||||
import { EnvironmentService } from '../../integrations/environment/environment.service';
|
||||
|
||||
/**
|
||||
* Per-user AI chat API (§6.1). Routes are POST to match this codebase's
|
||||
@@ -72,6 +82,11 @@ export class AiChatController {
|
||||
private readonly aiChatMessageRepo: AiChatMessageRepo,
|
||||
private readonly aiTranscription: AiTranscriptionService,
|
||||
private readonly pageRepo: PageRepo,
|
||||
// #184 phase 1.5. OPTIONAL so existing positional constructions (controller
|
||||
// specs) compile unchanged; Nest always injects the real providers in
|
||||
// production. Only touched on the resumable-stream (flag-on) path.
|
||||
private readonly streamRegistry?: AiChatStreamRegistryService,
|
||||
private readonly environment?: EnvironmentService,
|
||||
) {}
|
||||
|
||||
/** List the requesting user's chats in this workspace (paginated). */
|
||||
@@ -233,6 +248,102 @@ export class AiChatController {
|
||||
return { stopped };
|
||||
}
|
||||
|
||||
/**
|
||||
* Attach to a chat's live run stream (#184 phase 1.5). A late/reloaded tab
|
||||
* replays the frames buffered so far and then follows the live tail as a normal
|
||||
* streamer. Owner-gated via assertOwnedChat (same gate as getRun). When there is
|
||||
* nothing to resume — no entry, a finished run without expect=live, an
|
||||
* overflowed buffer, or an anchor that pins a DIFFERENT run — the endpoint
|
||||
* answers 204, the ONLY "nothing to resume" signal the AI SDK's reconnect
|
||||
* accepts (it maps 204 to a silent no-op). With AI_CHAT_RESUMABLE_STREAM off the
|
||||
* registry is never populated, so attach always 204s.
|
||||
*
|
||||
* `expect=live` opts into replaying a finished-but-retained run (safe only when
|
||||
* the client stripped the streaming tail); `anchor` is the client's assistant
|
||||
* row id, which must match this run's (invariant 6) or a foreign run's
|
||||
* transcript would be replayed into the store.
|
||||
*/
|
||||
@SkipTransform()
|
||||
@UseGuards(JwtAuthGuard, UserThrottlerGuard)
|
||||
@Throttle({ [AI_CHAT_THROTTLER]: { limit: 60, ttl: 60000 } })
|
||||
@Get('runs/:chatId/stream')
|
||||
async attachRunStream(
|
||||
@Param('chatId', new ParseUUIDPipe()) chatId: string,
|
||||
@Query('expect') expect: string | undefined,
|
||||
@Query('anchor') anchor: string | undefined,
|
||||
@Req() req: FastifyRequest,
|
||||
@Res() res: FastifyReply,
|
||||
@AuthUser() user: User,
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
): Promise<void> {
|
||||
await this.assertOwnedChat(chatId, user, workspace); // same gate as getRun
|
||||
let stopHeartbeat: () => void = () => undefined;
|
||||
const attachment = await this.streamRegistry?.attach(
|
||||
chatId,
|
||||
expect === 'live',
|
||||
anchor,
|
||||
{
|
||||
onFrame: (frame) => {
|
||||
// Backpressure guard: 2x the replay cap, so the initial replay burst
|
||||
// alone can never trip it; only a genuinely stalled socket can.
|
||||
try {
|
||||
if (res.raw.writableLength > SUBSCRIBER_MAX_BUFFERED_BYTES) {
|
||||
res.raw.destroy(); // 'close' fires -> unsubscribe below
|
||||
return;
|
||||
}
|
||||
if (!res.raw.writableEnded) res.raw.write(frame);
|
||||
} catch {
|
||||
res.raw.destroy();
|
||||
}
|
||||
},
|
||||
onEnd: () => {
|
||||
stopHeartbeat();
|
||||
if (!res.raw.writableEnded) res.raw.end();
|
||||
},
|
||||
},
|
||||
);
|
||||
if (!attachment) {
|
||||
res.status(204).send(); // the ONLY "nothing to resume" signal the SDK accepts
|
||||
return;
|
||||
}
|
||||
res.hijack();
|
||||
// Cleanup BEFORE any write (invariant 5): a torn-down socket must not orphan
|
||||
// a paused subscriber whose pending queue would buffer the whole run.
|
||||
req.raw.once('close', () => {
|
||||
attachment.unsubscribe();
|
||||
stopHeartbeat();
|
||||
});
|
||||
// A close emitted DURING the awaits above was missed by the listener — check.
|
||||
// (Healthy pending GETs have req.raw.destroyed === false, so no false
|
||||
// positives; returning without end() is fine — the socket is gone.)
|
||||
if (req.raw.destroyed) {
|
||||
attachment.unsubscribe();
|
||||
return;
|
||||
}
|
||||
res.raw.on('error', () => undefined);
|
||||
try {
|
||||
res.raw.writeHead(200, {
|
||||
'content-type': 'text/event-stream',
|
||||
'cache-control': 'no-cache',
|
||||
'x-vercel-ai-ui-message-stream': 'v1',
|
||||
'x-accel-buffering': 'no',
|
||||
// deliberately NO Connection/Keep-Alive (hop-by-hop; Safari/HTTP2)
|
||||
});
|
||||
res.raw.flushHeaders?.();
|
||||
for (const frame of attachment.replay) res.raw.write(frame);
|
||||
if (attachment.finished) {
|
||||
res.raw.end();
|
||||
return;
|
||||
}
|
||||
stopHeartbeat = startSseHeartbeat(res.raw, 15_000);
|
||||
attachment.start(); // drain pending accumulated during replay, go live
|
||||
} catch {
|
||||
attachment.unsubscribe();
|
||||
stopHeartbeat();
|
||||
res.raw.destroy();
|
||||
}
|
||||
}
|
||||
|
||||
/** Rename a chat. */
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@Post('rename')
|
||||
@@ -344,13 +455,25 @@ export class AiChatController {
|
||||
// its progress, and settle its terminal status — see AiChatRunService.
|
||||
const runHooks: AiChatRunHooks | undefined = autonomousRuns
|
||||
? {
|
||||
begin: (chatId) =>
|
||||
this.aiChatRunService.beginRun({
|
||||
begin: async (chatId) => {
|
||||
const handle = await this.aiChatRunService.beginRun({
|
||||
chatId,
|
||||
workspaceId: workspace.id,
|
||||
userId: user.id,
|
||||
trigger: 'user',
|
||||
}),
|
||||
});
|
||||
// #184 phase 1.5: register the run-stream entry at BEGIN (before any
|
||||
// frame) so a tab that attaches in the begin->seed window finds an
|
||||
// entry to wait on. Gated on AI_CHAT_RESUMABLE_STREAM: with the flag
|
||||
// off nothing is registered and attach always 204s.
|
||||
if (
|
||||
handle?.runId &&
|
||||
this.environment?.isAiChatResumableStreamEnabled?.()
|
||||
) {
|
||||
this.streamRegistry?.open(chatId, handle.runId);
|
||||
}
|
||||
return handle;
|
||||
},
|
||||
onAssistantSeeded: (runId, messageId) =>
|
||||
this.aiChatRunService.linkAssistantMessage(
|
||||
runId,
|
||||
|
||||
@@ -4,6 +4,7 @@ import { TokenModule } from '../auth/token.module';
|
||||
import { AiChatController } from './ai-chat.controller';
|
||||
import { AiChatService } from './ai-chat.service';
|
||||
import { AiChatRunService } from './ai-chat-run.service';
|
||||
import { AiChatStreamRegistryService } from './ai-chat-stream-registry.service';
|
||||
import { AiTranscriptionService } from './ai-transcription.service';
|
||||
import { AiChatToolsService } from './tools/ai-chat-tools.service';
|
||||
import { EmbeddingModule } from './embedding/embedding.module';
|
||||
@@ -44,6 +45,7 @@ import { PublicShareChatToolsService } from './tools/public-share-chat-tools.ser
|
||||
providers: [
|
||||
AiChatService,
|
||||
AiChatRunService,
|
||||
AiChatStreamRegistryService,
|
||||
AiTranscriptionService,
|
||||
AiChatToolsService,
|
||||
PublicShareChatService,
|
||||
|
||||
@@ -1,4 +1,13 @@
|
||||
import { ForbiddenException } from '@nestjs/common';
|
||||
import { ForbiddenException, Logger } from '@nestjs/common';
|
||||
// Mock ONLY streamText so a driven stream() call can capture the pipe-options
|
||||
// object (consumeSseStream / generateMessageId). Everything else in the AI SDK
|
||||
// stays REAL (requireActual), so the pure-helper suites in this file are
|
||||
// unaffected — none of them call stream()/streamText.
|
||||
jest.mock('ai', () => ({
|
||||
...jest.requireActual('ai'),
|
||||
streamText: jest.fn(),
|
||||
}));
|
||||
import { streamText } from 'ai';
|
||||
import {
|
||||
AiChatService,
|
||||
compactToolOutput,
|
||||
@@ -1059,3 +1068,181 @@ describe('isInterruptResume', () => {
|
||||
expect(isInterruptResume(withPrev(null), true)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* #184 phase 1.5 — the run-wrapped pipe options (unit). Drives stream() to the
|
||||
* pipe call with streamText mocked, capturing the options object, and asserts:
|
||||
* - flag OFF while a runId IS present -> the LEGACY option shape (no
|
||||
* consumeSseStream, no generateMessageId), and the registry is never touched.
|
||||
* This is the exact dormancy guarantee this PR rests on.
|
||||
* - flag ON + runId -> consumeSseStream tees into the registry and
|
||||
* generateMessageId returns the seeded assistant DB row id.
|
||||
* - flag ON but no runHooks (runId undefined) -> legacy (the runId gate).
|
||||
* - flag ON + runId -> the outer catch releases the entry via abortEntry.
|
||||
*/
|
||||
describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', () => {
|
||||
const streamTextMock = streamText as unknown as jest.Mock;
|
||||
let pipeMock: jest.Mock;
|
||||
|
||||
beforeEach(() => {
|
||||
streamTextMock.mockReset();
|
||||
pipeMock = jest.fn();
|
||||
streamTextMock.mockReturnValue({
|
||||
consumeStream: jest.fn(),
|
||||
pipeUIMessageStreamToResponse: pipeMock,
|
||||
});
|
||||
// Silence the service's diagnostic logging for a clean test run.
|
||||
jest.spyOn(Logger.prototype, 'log').mockImplementation(() => undefined as never);
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined as never);
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined as never);
|
||||
});
|
||||
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
// A raw-response stub sufficient for the post-streamText wiring.
|
||||
function makeRes() {
|
||||
return {
|
||||
raw: {
|
||||
writeHead: jest.fn(),
|
||||
write: jest.fn(),
|
||||
once: jest.fn(),
|
||||
on: jest.fn(),
|
||||
flushHeaders: jest.fn(),
|
||||
writableEnded: false,
|
||||
destroyed: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
// Wire only the deps reached on the way to the pipe call, plus a spy registry.
|
||||
function makeService(opts: { resumable: boolean }) {
|
||||
const aiChatRepo = {
|
||||
findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })),
|
||||
insert: jest.fn(),
|
||||
};
|
||||
const aiChatMessageRepo = {
|
||||
// Both the user insert and the assistant seed return the same row id.
|
||||
insert: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findAllByChat: jest.fn(async () => []),
|
||||
update: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
};
|
||||
const aiSettings = { resolve: jest.fn(async () => ({})) };
|
||||
const tools = { forUser: jest.fn(async () => ({})) };
|
||||
const mcpClients = {
|
||||
toolsFor: jest.fn(async () => ({
|
||||
tools: {},
|
||||
clients: [],
|
||||
outcomes: [],
|
||||
instructions: [],
|
||||
})),
|
||||
};
|
||||
const streamRegistry = {
|
||||
open: jest.fn(),
|
||||
bind: jest.fn(),
|
||||
abortEntry: jest.fn(),
|
||||
};
|
||||
const svc = new AiChatService(
|
||||
{} as never, // ai (model is injected)
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiChatPageSnapshotRepo
|
||||
aiSettings as never,
|
||||
tools as never,
|
||||
mcpClients as never,
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo (openPage undefined -> never touched)
|
||||
{} as never, // pageAccess
|
||||
{
|
||||
isAiChatDeferredToolsEnabled: () => false,
|
||||
isAiChatResumableStreamEnabled: () => opts.resumable,
|
||||
} as never,
|
||||
streamRegistry as never,
|
||||
);
|
||||
return { svc, streamRegistry };
|
||||
}
|
||||
|
||||
const body = {
|
||||
chatId: 'chat-1',
|
||||
messages: [
|
||||
{ id: 'm1', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
|
||||
],
|
||||
};
|
||||
|
||||
const makeRunHooks = () => ({
|
||||
begin: jest.fn(async () => ({
|
||||
runId: 'run-1',
|
||||
signal: new AbortController().signal,
|
||||
})),
|
||||
onAssistantSeeded: jest.fn(),
|
||||
onStep: jest.fn(),
|
||||
onSettled: jest.fn(),
|
||||
});
|
||||
|
||||
async function drive(svc: AiChatService, hooks: unknown): Promise<void> {
|
||||
await svc.stream({
|
||||
user: { id: 'u1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 's1',
|
||||
body: body as never,
|
||||
res: makeRes() as never,
|
||||
signal: new AbortController().signal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks: hooks as never,
|
||||
});
|
||||
}
|
||||
|
||||
it('flag OFF + runId present: LEGACY option shape (no consumeSseStream / generateMessageId); registry untouched', async () => {
|
||||
const { svc, streamRegistry } = makeService({ resumable: false });
|
||||
await drive(svc, makeRunHooks());
|
||||
expect(pipeMock).toHaveBeenCalledTimes(1);
|
||||
const options = pipeMock.mock.calls[0][1];
|
||||
// The dormancy guarantee: a live run with the flag off tees NOTHING and does
|
||||
// not stamp a message id — byte-for-byte the pre-1.5 wire.
|
||||
expect(options.consumeSseStream).toBeUndefined();
|
||||
expect(options.generateMessageId).toBeUndefined();
|
||||
expect(streamRegistry.bind).not.toHaveBeenCalled();
|
||||
expect(streamRegistry.abortEntry).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('flag ON + runId: consumeSseStream tees into the registry; generateMessageId returns the seeded row id', async () => {
|
||||
const { svc, streamRegistry } = makeService({ resumable: true });
|
||||
await drive(svc, makeRunHooks());
|
||||
const options = pipeMock.mock.calls[0][1];
|
||||
expect(typeof options.consumeSseStream).toBe('function');
|
||||
expect(typeof options.generateMessageId).toBe('function');
|
||||
// generateMessageId stamps the seeded assistant DB row id.
|
||||
expect(options.generateMessageId()).toBe('msg-1');
|
||||
// consumeSseStream binds the tee: (chatId, runId, assistantId, stream).
|
||||
const fakeStream = {} as ReadableStream<string>;
|
||||
options.consumeSseStream({ stream: fakeStream });
|
||||
expect(streamRegistry.bind).toHaveBeenCalledWith(
|
||||
'chat-1',
|
||||
'run-1',
|
||||
'msg-1',
|
||||
fakeStream,
|
||||
);
|
||||
});
|
||||
|
||||
it('flag ON but NO runHooks (runId undefined): pipe options stay legacy (the runId gate)', async () => {
|
||||
const { svc, streamRegistry } = makeService({ resumable: true });
|
||||
await drive(svc, undefined);
|
||||
const options = pipeMock.mock.calls[0][1];
|
||||
expect(options.consumeSseStream).toBeUndefined();
|
||||
expect(options.generateMessageId).toBeUndefined();
|
||||
expect(streamRegistry.bind).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('flag ON + runId: the outer catch calls abortEntry when the stream throws', async () => {
|
||||
const { svc, streamRegistry } = makeService({ resumable: true });
|
||||
streamTextMock.mockImplementation(() => {
|
||||
throw new Error('boom');
|
||||
});
|
||||
await expect(drive(svc, makeRunHooks())).rejects.toThrow('boom');
|
||||
expect(streamRegistry.abortEntry).toHaveBeenCalledWith('chat-1', 'run-1');
|
||||
});
|
||||
});
|
||||
|
||||
@@ -32,6 +32,7 @@ import {
|
||||
import { AiChatToolsService } from './tools/ai-chat-tools.service';
|
||||
import { McpClientsService } from './external-mcp/mcp-clients.service';
|
||||
import { EnvironmentService } from '../../integrations/environment/environment.service';
|
||||
import { AiChatStreamRegistryService } from './ai-chat-stream-registry.service';
|
||||
import { buildSystemPrompt } from './ai-chat.prompt';
|
||||
import {
|
||||
CORE_TOOL_KEYS,
|
||||
@@ -276,6 +277,10 @@ export class AiChatService implements OnModuleInit {
|
||||
// Reads the AI_CHAT_DEFERRED_TOOLS toggle (#332). Injected last so existing
|
||||
// positional constructor callers (tests) only append one stub.
|
||||
private readonly environment: EnvironmentService,
|
||||
// #184 phase 1.5 run-stream registry. OPTIONAL so existing positional
|
||||
// constructions (int-specs) compile unchanged; Nest always injects the real
|
||||
// provider in production. Only ever touched on the run-wrapped + flag-on path.
|
||||
private readonly streamRegistry?: AiChatStreamRegistryService,
|
||||
) {}
|
||||
|
||||
/**
|
||||
@@ -1174,6 +1179,35 @@ export class AiChatService implements OnModuleInit {
|
||||
// as the cumulative authoritative usage so the client never jumps DOWN.
|
||||
let cumulativeStepUsage: ChatStreamUsage | undefined;
|
||||
result.pipeUIMessageStreamToResponse(res.raw, {
|
||||
// #184 phase 1.5: run-wrapped mode only — the legacy path (flag off) stays
|
||||
// byte-for-byte identical, including the absence of start.messageId. Both
|
||||
// fields are gated on `runId` (present only for a durable run) AND the
|
||||
// AI_CHAT_RESUMABLE_STREAM flag; the seed `assistantId` is unconditional,
|
||||
// so gating on `assistantId` alone would change the legacy wire.
|
||||
...(runId && this.environment?.isAiChatResumableStreamEnabled?.()
|
||||
? {
|
||||
// Tee the SSE frames into the run-stream registry so late tabs can
|
||||
// attach (replay + live tail).
|
||||
consumeSseStream: ({
|
||||
stream,
|
||||
}: {
|
||||
stream: ReadableStream<string>;
|
||||
}) =>
|
||||
this.streamRegistry?.bind(
|
||||
chatId,
|
||||
runId!,
|
||||
assistantId,
|
||||
stream,
|
||||
),
|
||||
// Stamp the persisted assistant row's DB id onto the streamed
|
||||
// message so every tab renders the SAME id as the DB row (id-based
|
||||
// reconciliation). Seeding is best-effort: when it failed, let the
|
||||
// client generate the id.
|
||||
...(assistantId
|
||||
? { generateMessageId: () => assistantId }
|
||||
: {}),
|
||||
}
|
||||
: {}),
|
||||
headers: { 'X-Accel-Buffering': 'no' },
|
||||
// Surface the authoritative chatId on the streamed assistant UI message so
|
||||
// the client adopts the REAL id of the row we created, instead of guessing
|
||||
@@ -1239,6 +1273,12 @@ export class AiChatService implements OnModuleInit {
|
||||
// finalizeRun (onSettled) is idempotent — a settle here and a settle from a
|
||||
// streamText callback collapse to a single terminal write.
|
||||
if (runId) {
|
||||
// #184 phase 1.5: a failure here means the tee `done` will never arrive,
|
||||
// so release the registry entry's subscribers explicitly — otherwise an
|
||||
// attached tab hangs forever. Same flag gate as the tee wiring above.
|
||||
if (this.environment?.isAiChatResumableStreamEnabled?.()) {
|
||||
this.streamRegistry?.abortEntry(chatId, runId);
|
||||
}
|
||||
await runHooks?.onSettled?.(
|
||||
runId,
|
||||
'error',
|
||||
|
||||
@@ -292,6 +292,23 @@ export class EnvironmentService {
|
||||
return enabled === 'true';
|
||||
}
|
||||
|
||||
/**
|
||||
* Resumable SSE transport for durable agent runs (#184 phase 1.5). When
|
||||
* enabled, a run tees its SSE frames into the in-memory run-stream registry so
|
||||
* a late/reloaded tab can attach (replay + live tail) via
|
||||
* `GET /ai-chat/runs/:chatId/stream`. Defaults to DISABLED: PR 1 ships the
|
||||
* server code dormant — with the flag off, `open`/`bind`/`generateMessageId`
|
||||
* are never called and attach always answers 204, so the legacy and #184
|
||||
* phase-1 wire paths stay byte-for-byte identical. Set
|
||||
* AI_CHAT_RESUMABLE_STREAM=true to activate it (paired with the PR 2 client).
|
||||
*/
|
||||
isAiChatResumableStreamEnabled(): boolean {
|
||||
const enabled = this.configService
|
||||
.get<string>('AI_CHAT_RESUMABLE_STREAM', 'false')
|
||||
.toLowerCase();
|
||||
return enabled === 'true';
|
||||
}
|
||||
|
||||
getPostHogHost(): string {
|
||||
return this.configService.get<string>('POSTHOG_HOST');
|
||||
}
|
||||
|
||||
@@ -0,0 +1,564 @@
|
||||
import * as http from 'node:http';
|
||||
import { Kysely } from 'kysely';
|
||||
import {
|
||||
MockLanguageModelV3,
|
||||
convertArrayToReadableStream,
|
||||
} from 'ai/test';
|
||||
import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
|
||||
import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
|
||||
import { AiChatRunRepo } from '@docmost/db/repos/ai-chat/ai-chat-run.repo';
|
||||
import { AiChatService } from 'src/core/ai-chat/ai-chat.service';
|
||||
import { AiChatRunService } from 'src/core/ai-chat/ai-chat-run.service';
|
||||
import {
|
||||
AiChatStreamRegistryService,
|
||||
RunStreamCallbacks,
|
||||
} from 'src/core/ai-chat/ai-chat-stream-registry.service';
|
||||
import {
|
||||
getTestDb,
|
||||
destroyTestDb,
|
||||
createWorkspace,
|
||||
createUser,
|
||||
createChat,
|
||||
} from './db';
|
||||
|
||||
/**
|
||||
* #184 phase 1.5 — the resumable transport end to end against REAL Postgres,
|
||||
* the REAL `streamText` (seeded via MockLanguageModelV3) and a REAL Node
|
||||
* ServerResponse, driving the REAL `AiChatService.stream` run-wrapped path with a
|
||||
* REAL `AiChatStreamRegistryService`. The run-hooks mirror the controller: they
|
||||
* begin a durable run and `open()` the registry entry at begin, and the service
|
||||
* tees the SSE frames into it via `consumeSseStream` while stamping the DB row id
|
||||
* via `generateMessageId` (both gated on runId + the resumable flag).
|
||||
*
|
||||
* Proven here: a finished run's replay is the full frame sequence incl `[DONE]`
|
||||
* with `start.messageId` == the seeded DB row id; the anchor check (invariant 6);
|
||||
* an attach opened BEFORE the first frame follows the live stream from frame 0; an
|
||||
* explicit stop surfaces `{"type":"abort"}` + `[DONE]` + end to the subscriber;
|
||||
* and the legacy (non-run) path tees nothing.
|
||||
*/
|
||||
|
||||
const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
|
||||
|
||||
async function waitFor(
|
||||
cond: () => Promise<boolean> | boolean,
|
||||
{ timeoutMs = 15_000, stepMs = 25 } = {},
|
||||
): Promise<void> {
|
||||
const start = Date.now();
|
||||
while (Date.now() - start < timeoutMs) {
|
||||
if (await cond()) return;
|
||||
await sleep(stepMs);
|
||||
}
|
||||
throw new Error('waitFor: condition not met within timeout');
|
||||
}
|
||||
|
||||
// A real Node ServerResponse wired to a live socket (as in the stream int-spec).
|
||||
function makeRealResponse(): Promise<{
|
||||
res: http.ServerResponse;
|
||||
cleanup: () => Promise<void>;
|
||||
}> {
|
||||
return new Promise((resolve) => {
|
||||
const server = http.createServer((_req, res) => {
|
||||
resolve({
|
||||
res,
|
||||
cleanup: () =>
|
||||
new Promise<void>((done) => {
|
||||
try {
|
||||
if (!res.writableEnded) res.end();
|
||||
} catch {
|
||||
/* socket already gone */
|
||||
}
|
||||
server.close(() => done());
|
||||
}),
|
||||
});
|
||||
});
|
||||
server.listen(0, () => {
|
||||
const port = (server.address() as any).port;
|
||||
const creq = http.request({ port, method: 'GET' }, (cres) => {
|
||||
cres.resume();
|
||||
});
|
||||
creq.on('error', () => undefined);
|
||||
creq.end();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
// A full, successful single-step turn.
|
||||
function successStream() {
|
||||
return convertArrayToReadableStream([
|
||||
{ type: 'stream-start', warnings: [] },
|
||||
{ type: 'text-start', id: 't1' },
|
||||
{ type: 'text-delta', id: 't1', delta: 'Hello' },
|
||||
{ type: 'text-delta', id: 't1', delta: ' there' },
|
||||
{ type: 'text-end', id: 't1' },
|
||||
{
|
||||
type: 'finish',
|
||||
finishReason: 'stop',
|
||||
usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
|
||||
},
|
||||
] as any);
|
||||
}
|
||||
|
||||
// A stream the test feeds chunk by chunk (to attach mid-flight / drive a stop).
|
||||
function makeControlledStream() {
|
||||
let controller!: ReadableStreamDefaultController<any>;
|
||||
const stream = new ReadableStream<any>({
|
||||
start(c) {
|
||||
controller = c;
|
||||
},
|
||||
});
|
||||
return {
|
||||
stream,
|
||||
emit: (chunk: any) => controller.enqueue(chunk),
|
||||
close: () => controller.close(),
|
||||
};
|
||||
}
|
||||
|
||||
// Collect replay + live frames from an attachment.
|
||||
function liveSink(): {
|
||||
cb: RunStreamCallbacks;
|
||||
frames: string[];
|
||||
ended: () => boolean;
|
||||
} {
|
||||
const frames: string[] = [];
|
||||
let ended = false;
|
||||
return {
|
||||
frames,
|
||||
ended: () => ended,
|
||||
cb: {
|
||||
onFrame: (f) => frames.push(f),
|
||||
onEnd: () => {
|
||||
ended = true;
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
// The SSE `start` frame carries the message id; pull it out of a `data: {...}`.
|
||||
function parseStartMessageId(frames: string[]): string | undefined {
|
||||
for (const f of frames) {
|
||||
const m = /^data: (\{.*\})\s*$/m.exec(f.trim());
|
||||
if (!m) continue;
|
||||
try {
|
||||
const json = JSON.parse(m[1]);
|
||||
if (json.type === 'start') return json.messageId;
|
||||
} catch {
|
||||
/* not this frame */
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
describe('AiChatService run-stream attach [integration]', () => {
|
||||
let db: Kysely<any>;
|
||||
let aiChatRepo: AiChatRepo;
|
||||
let msgRepo: AiChatMessageRepo;
|
||||
let runRepo: AiChatRunRepo;
|
||||
let workspaceId: string;
|
||||
let userId: string;
|
||||
|
||||
const mcpClients = {
|
||||
toolsFor: async () => ({
|
||||
tools: {},
|
||||
clients: [],
|
||||
outcomes: [],
|
||||
instructions: [],
|
||||
}),
|
||||
};
|
||||
|
||||
// Build the service with the run-stream registry wired and the resumable flag
|
||||
// ON (the property under test). Deferred tools OFF (irrelevant here).
|
||||
function buildService(registry: AiChatStreamRegistryService): AiChatService {
|
||||
return new AiChatService(
|
||||
{ getChatModel: async () => null } as any,
|
||||
aiChatRepo,
|
||||
msgRepo,
|
||||
{} as any,
|
||||
{ resolve: async () => null } as any,
|
||||
{ forUser: async () => ({}) } as any,
|
||||
mcpClients as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
{
|
||||
isAiChatDeferredToolsEnabled: () => false,
|
||||
isAiChatResumableStreamEnabled: () => true,
|
||||
} as any,
|
||||
registry,
|
||||
);
|
||||
}
|
||||
|
||||
// Run-hooks mirroring the controller: begin the durable run AND open() the
|
||||
// registry entry at begin. Captures the runId so a test can stop it.
|
||||
function makeRunHooks(
|
||||
runService: AiChatRunService,
|
||||
registry: AiChatStreamRegistryService,
|
||||
box: { runId?: string },
|
||||
) {
|
||||
return {
|
||||
begin: async (chatId: string) => {
|
||||
const handle = await runService.beginRun({
|
||||
chatId,
|
||||
workspaceId,
|
||||
userId,
|
||||
trigger: 'user',
|
||||
});
|
||||
box.runId = handle.runId;
|
||||
registry.open(chatId, handle.runId);
|
||||
return handle;
|
||||
},
|
||||
onAssistantSeeded: (runId: string, messageId: string) =>
|
||||
runService.linkAssistantMessage(runId, workspaceId, messageId),
|
||||
onStep: (runId: string, n: number) =>
|
||||
void runService.recordStep(runId, workspaceId, n),
|
||||
onSettled: (runId: string, status: any, error?: string) =>
|
||||
runService.finalizeRun(runId, workspaceId, status, error),
|
||||
};
|
||||
}
|
||||
|
||||
function userUiMessage(text: string) {
|
||||
return {
|
||||
id: `u-${Math.random()}`,
|
||||
role: 'user',
|
||||
parts: [{ type: 'text', text }],
|
||||
};
|
||||
}
|
||||
|
||||
async function startRun(opts: {
|
||||
registry: AiChatStreamRegistryService;
|
||||
runService?: AiChatRunService;
|
||||
model: MockLanguageModelV3;
|
||||
chatId: string;
|
||||
body: any;
|
||||
box?: { runId?: string };
|
||||
}): Promise<{ res: http.ServerResponse; cleanup: () => Promise<void> }> {
|
||||
const service = buildService(opts.registry);
|
||||
const { res, cleanup } = await makeRealResponse();
|
||||
const runHooks = opts.runService
|
||||
? makeRunHooks(opts.runService, opts.registry, opts.box ?? {})
|
||||
: undefined;
|
||||
await service.stream({
|
||||
user: { id: userId, workspaceId } as any,
|
||||
workspace: { id: workspaceId, name: 'WS' } as any,
|
||||
sessionId: 'sess-1',
|
||||
body: opts.body,
|
||||
res: { raw: res } as any,
|
||||
signal: new AbortController().signal,
|
||||
model: opts.model as any,
|
||||
role: null,
|
||||
runHooks,
|
||||
} as any);
|
||||
return { res, cleanup };
|
||||
}
|
||||
|
||||
async function assistantRowId(chatId: string): Promise<string> {
|
||||
const rows = await msgRepo.findAllByChat(chatId, workspaceId);
|
||||
const row = rows.find((r: any) => r.role === 'assistant');
|
||||
return row!.id as string;
|
||||
}
|
||||
|
||||
beforeAll(async () => {
|
||||
db = getTestDb();
|
||||
aiChatRepo = new AiChatRepo(db as any);
|
||||
msgRepo = new AiChatMessageRepo(db as any);
|
||||
runRepo = new AiChatRunRepo(db as any);
|
||||
workspaceId = (await createWorkspace(db)).id;
|
||||
userId = (await createUser(db, workspaceId)).id;
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await destroyTestDb();
|
||||
});
|
||||
|
||||
it('run-wrapped: replay is the full frame sequence incl [DONE], start.messageId == the seeded DB row id', async () => {
|
||||
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
const registry = new AiChatStreamRegistryService();
|
||||
const runService = new AiChatRunService(runRepo, {
|
||||
isCloud: () => false,
|
||||
} as never);
|
||||
const model = new MockLanguageModelV3({
|
||||
doStream: async () => ({ stream: successStream() }),
|
||||
} as any);
|
||||
|
||||
const box: { runId?: string } = {};
|
||||
const { cleanup } = await startRun({
|
||||
registry,
|
||||
runService,
|
||||
model,
|
||||
chatId,
|
||||
body: { chatId, messages: [userUiMessage('Hi')] },
|
||||
box,
|
||||
});
|
||||
try {
|
||||
// Wait for the assistant row to settle (terminal callbacks run async).
|
||||
await waitFor(async () => {
|
||||
const rows = await msgRepo.findAllByChat(chatId, workspaceId);
|
||||
return rows.some(
|
||||
(r: any) =>
|
||||
r.role === 'assistant' &&
|
||||
['completed', 'error', 'aborted'].includes(r.status),
|
||||
);
|
||||
});
|
||||
const rowId = await assistantRowId(chatId);
|
||||
|
||||
// Finished-run replay with expect=live + the correct anchor.
|
||||
const sink = liveSink();
|
||||
const att = await registry.attach(chatId, true, rowId, sink.cb);
|
||||
expect(att).not.toBeNull();
|
||||
expect(att!.finished).toBe(true);
|
||||
// The tee captured frames (consumeSseStream was wired).
|
||||
expect(att!.replay.length).toBeGreaterThan(0);
|
||||
// generateMessageId stamped the DB row id onto the streamed start frame.
|
||||
expect(parseStartMessageId(att!.replay)).toBe(rowId);
|
||||
// The full sequence includes the streamed text and the terminal marker.
|
||||
const joined = att!.replay.join('');
|
||||
expect(joined).toContain('Hello');
|
||||
expect(att!.replay.some((f) => f.includes('[DONE]'))).toBe(true);
|
||||
} finally {
|
||||
registry.onModuleDestroy();
|
||||
await cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it('anchor mismatch with expect=live returns null (invariant 6)', async () => {
|
||||
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
const registry = new AiChatStreamRegistryService();
|
||||
const runService = new AiChatRunService(runRepo, {
|
||||
isCloud: () => false,
|
||||
} as never);
|
||||
const model = new MockLanguageModelV3({
|
||||
doStream: async () => ({ stream: successStream() }),
|
||||
} as any);
|
||||
|
||||
const { cleanup } = await startRun({
|
||||
registry,
|
||||
runService,
|
||||
model,
|
||||
chatId,
|
||||
body: { chatId, messages: [userUiMessage('Hi')] },
|
||||
});
|
||||
try {
|
||||
await waitFor(async () => {
|
||||
const rows = await msgRepo.findAllByChat(chatId, workspaceId);
|
||||
return rows.some(
|
||||
(r: any) => r.role === 'assistant' && r.status === 'completed',
|
||||
);
|
||||
});
|
||||
const sink = liveSink();
|
||||
// A foreign anchor must NOT replay this run's transcript.
|
||||
expect(
|
||||
await registry.attach(chatId, true, 'a-different-run-row', sink.cb),
|
||||
).toBeNull();
|
||||
} finally {
|
||||
registry.onModuleDestroy();
|
||||
await cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it('an attach opened BEFORE the first frame follows the live stream from frame 0', async () => {
|
||||
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
const registry = new AiChatStreamRegistryService();
|
||||
const runService = new AiChatRunService(runRepo, {
|
||||
isCloud: () => false,
|
||||
} as never);
|
||||
const controlled = makeControlledStream();
|
||||
const model = new MockLanguageModelV3({
|
||||
doStream: async () => ({ stream: controlled.stream }),
|
||||
} as any);
|
||||
|
||||
const { cleanup } = await startRun({
|
||||
registry,
|
||||
runService,
|
||||
model,
|
||||
chatId,
|
||||
body: { chatId, messages: [userUiMessage('Slow please')] },
|
||||
});
|
||||
try {
|
||||
// Attach while the entry exists (opened at begin) but before any frame.
|
||||
const sink = liveSink();
|
||||
const att = (await registry.attach(chatId, false, undefined, sink.cb))!;
|
||||
expect(att.replay).toEqual([]); // nothing streamed yet -> replay from 0
|
||||
att.start(); // go live (drains nothing, then follows)
|
||||
|
||||
// Now emit the whole turn.
|
||||
controlled.emit({ type: 'stream-start', warnings: [] });
|
||||
controlled.emit({ type: 'text-start', id: 't1' });
|
||||
controlled.emit({ type: 'text-delta', id: 't1', delta: 'Zero' });
|
||||
controlled.emit({ type: 'text-end', id: 't1' });
|
||||
controlled.emit({
|
||||
type: 'finish',
|
||||
finishReason: 'stop',
|
||||
usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 },
|
||||
});
|
||||
controlled.close();
|
||||
|
||||
await waitFor(() => sink.frames.some((f) => f.includes('[DONE]')));
|
||||
// The subscriber saw the stream from the very first frame (`start`) through
|
||||
// the terminal marker, with the streamed text present.
|
||||
expect(sink.frames.some((f) => f.includes('"type":"start"'))).toBe(true);
|
||||
expect(sink.frames.join('')).toContain('Zero');
|
||||
expect(sink.frames[sink.frames.length - 1]).toContain('[DONE]');
|
||||
expect(sink.ended()).toBe(true);
|
||||
} finally {
|
||||
registry.onModuleDestroy();
|
||||
await cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it('requestStop surfaces {"type":"abort"} + [DONE] + end to the attached subscriber', async () => {
|
||||
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
const registry = new AiChatStreamRegistryService();
|
||||
const runService = new AiChatRunService(runRepo, {
|
||||
isCloud: () => false,
|
||||
} as never);
|
||||
// An abort-AWARE model: it streams some partial output, then errors the model
|
||||
// stream with an AbortError when the run signal aborts — exactly as a real
|
||||
// provider network stream is torn down on abort (a plain in-memory stream
|
||||
// would just stall, so streamText would never observe the stop).
|
||||
const model = new MockLanguageModelV3({
|
||||
doStream: async ({ abortSignal }: any) => {
|
||||
const stream = new ReadableStream<any>({
|
||||
start(controller) {
|
||||
controller.enqueue({ type: 'stream-start', warnings: [] });
|
||||
controller.enqueue({ type: 'text-start', id: 't1' });
|
||||
controller.enqueue({ type: 'text-delta', id: 't1', delta: 'partial' });
|
||||
abortSignal?.addEventListener('abort', () => {
|
||||
try {
|
||||
controller.error(
|
||||
new DOMException('Aborted', 'AbortError'),
|
||||
);
|
||||
} catch {
|
||||
/* already errored/closed */
|
||||
}
|
||||
});
|
||||
},
|
||||
});
|
||||
return { stream };
|
||||
},
|
||||
} as any);
|
||||
|
||||
const box: { runId?: string } = {};
|
||||
const { cleanup } = await startRun({
|
||||
registry,
|
||||
runService,
|
||||
model,
|
||||
chatId,
|
||||
body: { chatId, messages: [userUiMessage('Start then stop')] },
|
||||
box,
|
||||
});
|
||||
try {
|
||||
const sink = liveSink();
|
||||
const att = (await registry.attach(chatId, false, undefined, sink.cb))!;
|
||||
att.start();
|
||||
|
||||
// Give streamText a beat to begin consuming the partial output.
|
||||
await sleep(250);
|
||||
// User presses Stop -> the run signal aborts -> the SDK emits an abort chunk.
|
||||
await runService.requestStop(box.runId!, workspaceId);
|
||||
|
||||
await waitFor(() => sink.ended());
|
||||
expect(sink.frames.some((f) => f.includes('"type":"abort"'))).toBe(true);
|
||||
expect(sink.frames.some((f) => f.includes('[DONE]'))).toBe(true);
|
||||
expect(sink.ended()).toBe(true);
|
||||
} finally {
|
||||
registry.onModuleDestroy();
|
||||
await cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it('the outer catch calls abortEntry so an open entry is released (finished)', async () => {
|
||||
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
const registry = new AiChatStreamRegistryService();
|
||||
const runService = new AiChatRunService(runRepo, {
|
||||
isCloud: () => false,
|
||||
} as never);
|
||||
const model = new MockLanguageModelV3({
|
||||
doStream: async () => ({ stream: successStream() }),
|
||||
} as any);
|
||||
|
||||
// A msgRepo whose user-row insert throws: the turn fails AFTER begin (the
|
||||
// entry is already open) but BEFORE the pipe, exercising the outer catch.
|
||||
const throwingMsgRepo = {
|
||||
insert: async () => {
|
||||
throw new Error('db boom');
|
||||
},
|
||||
findAllByChat: (...a: any[]) => (msgRepo as any).findAllByChat(...a),
|
||||
update: (...a: any[]) => (msgRepo as any).update(...a),
|
||||
findById: (...a: any[]) => (msgRepo as any).findById(...a),
|
||||
};
|
||||
const service = new AiChatService(
|
||||
{ getChatModel: async () => null } as any,
|
||||
aiChatRepo,
|
||||
throwingMsgRepo as any,
|
||||
{} as any,
|
||||
{ resolve: async () => null } as any,
|
||||
{ forUser: async () => ({}) } as any,
|
||||
mcpClients as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
{
|
||||
isAiChatDeferredToolsEnabled: () => false,
|
||||
isAiChatResumableStreamEnabled: () => true,
|
||||
} as any,
|
||||
registry,
|
||||
);
|
||||
const { res, cleanup } = await makeRealResponse();
|
||||
const box: { runId?: string } = {};
|
||||
try {
|
||||
await expect(
|
||||
service.stream({
|
||||
user: { id: userId, workspaceId } as any,
|
||||
workspace: { id: workspaceId, name: 'WS' } as any,
|
||||
sessionId: 'sess-1',
|
||||
body: { chatId, messages: [userUiMessage('will throw')] },
|
||||
res: { raw: res } as any,
|
||||
signal: new AbortController().signal,
|
||||
model: model as any,
|
||||
role: null,
|
||||
runHooks: makeRunHooks(runService, registry, box),
|
||||
} as any),
|
||||
).rejects.toThrow();
|
||||
// The entry opened at begin was terminated by abortEntry (from the catch),
|
||||
// so it is finished and a plain attach returns null instead of hanging.
|
||||
const entry = (registry as any).entries.get(chatId);
|
||||
expect(entry).toBeDefined();
|
||||
expect(entry.finished).toBe(true);
|
||||
const sink = liveSink();
|
||||
expect(await registry.attach(chatId, false, undefined, sink.cb)).toBeNull();
|
||||
} finally {
|
||||
registry.onModuleDestroy();
|
||||
await cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it('legacy (no run-hooks): the registry is never populated', async () => {
|
||||
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
const registry = new AiChatStreamRegistryService();
|
||||
const model = new MockLanguageModelV3({
|
||||
doStream: async () => ({ stream: successStream() }),
|
||||
} as any);
|
||||
|
||||
// No runHooks -> runId undefined -> the run-wrapped tee is never wired.
|
||||
const { cleanup } = await startRun({
|
||||
registry,
|
||||
model,
|
||||
chatId,
|
||||
body: { chatId, messages: [userUiMessage('Legacy hi')] },
|
||||
});
|
||||
try {
|
||||
await waitFor(async () => {
|
||||
const rows = await msgRepo.findAllByChat(chatId, workspaceId);
|
||||
return rows.some(
|
||||
(r: any) => r.role === 'assistant' && r.status === 'completed',
|
||||
);
|
||||
});
|
||||
const sink = liveSink();
|
||||
// No entry was ever opened; attach always yields null.
|
||||
expect(await registry.attach(chatId, false, undefined, sink.cb)).toBeNull();
|
||||
expect(await registry.attach(chatId, true, 'anything', sink.cb)).toBeNull();
|
||||
} finally {
|
||||
registry.onModuleDestroy();
|
||||
await cleanup();
|
||||
}
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user