Compare commits

...

7 Commits

Author SHA1 Message Date
vvzvlad 40227bbf51 Merge pull request 'feat(ai-chat): resumable SSE — клиент (#381 PR 2, переоткрыт на develop после стек-мёржа)' (#389) from feat/381-resumable-sse-pr2 into develop
Reviewed-on: #389
2026-07-06 16:29:04 +03:00
agent_vscode a26803a1bc docs(env): document AI_CHAT_RESUMABLE_STREAM staged-rollout flag (#381)
The resumable-SSE run-stream registry (PR #386/#387) ships behind the
server-side AI_CHAT_RESUMABLE_STREAM env flag, OFF by default: with the
flag off attach always answers 204 and reopened tabs of an active run
fall back to degraded history polling. Document the flag, its default,
its relation to the per-workspace autonomousRuns setting, and the
single-instance constraint in .env.example next to the autonomous-runs
section.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-06 16:27:06 +03:00
agent_coder 18eee98b7f fix(ai-chat): #381 PR2 review round 1 — F7 restart-survival + unmount abort + anchor-orphan
Do 1 [F7 regression]: транзиентный сбой attach больше не роняет строку и не
теряет ран. В transport fetch-wrapper `204 || !response.ok` оба зовут
onNoActiveStream (восстановить stripped-строку + invalidate + арм poll), и catch
зовёт его перед rethrow — раньше !ok/throw только сбрасывали флаг, и на 5xx/502/
network-blip in-progress ассистент-турн исчезал, durable-ран не отслеживался.
onNoActiveStream — суперсет (его часть-г всё ещё чистит флаг), идемпотентен.
Расширяет литеральный block-3 спеки (там был только сброс флага) — по ревью и
в согласии с интенцией окна «poll must survive a server restart».

Do 2 [stability]: attach-GET абортится при unmount + mount-гейтинг сайд-эффектов.
mountedRef: mount-эффект ре-армит true и в cleanup ставит false + abort
attachAbortRef; onNoActiveStream рано выходит на !mounted, onFinish-recovery
гейтится `wasResumed && mountedRef.current`. Снимает до-10-мин спурьёзный поллинг
+ чужую invalidateQueries + утёкший fetch на новооткрытом чате (и StrictMode
double-resume).

Do 3 [coherence]: anchor-mismatch не оставляет вечную dots-строку. Reconcile
после мержа хвоста мержит fresh-history версию stripped-строки, если её id !=
id хвоста — settl'ит осиротевшую streaming-A над раном B.

Тесты: F7 500 → restore+арм; F7 network-throw → restore+арм; unmount при pending
attach → abort + поздние колбэки не летят. vitest src/features/ai-chat 304
зелёных, grep-guard пуст.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-06 16:26:27 +03:00
agent_coder 067fc46170 feat(ai-chat): единый resumable SSE-транспорт — клиент + удаление поллинга/латчей (#381 PR 2)
PR 2 из 2 (#381, фаза 1.5 #184). Все вкладки теперь на ОДНОМ транспорте:
любая подключается к рану через GET-attach (реплей кадров + живой хвост,
реестр из PR 1). Наблюдатель — обычный стример; Stop, точки, инвалидации
работают штатно. Двухпутёвый поллинг снапшотов и вся latch-механика F4/F5/F7
удалены.

- utils/resume-helpers.ts (новый): isStreamingTail / isSettledAssistantTail /
  seedRows / mergeById (дословный перенос mergeObservedMessage из run-polling).
- components/chat-thread.tsx: resume-машинерия (гейтинг по не-settled хвосту,
  strip streaming-хвоста + attach ?expect=live&anchor=<row id>, транспорт
  prepareReconnectToStreamRequest+fetch, 204-обработчик из 4 частей,
  reconcile+degraded-merge, recovery с АСИММЕТРИЕЙ arm-vs-restore — при
  isDisconnect с видимым контентом только arm, без restore-клоббера живого
  стрима (инв. 9), строгий порядок onFinish с ранним return до обеих веток
  отправки (инв. 7), «Send now» скрыт на resumed-ходе, Stop абортит attach).
- components/ai-chat-window.tsx: degraded-poll фолбэк вместо латчей — тупой
  таймер (2500ms, 10-мин кап, без проверок ошибок/хвоста; переживает рестарт
  сервера), гасится тредом через onResumeFallback(false).
- Удалено: run-polling.ts(+test), useAiChatRunQuery/AI_CHAT_RUN_RQ_KEY,
  getAiChatRun (stopRun оставлен), IAiChatRun/IAiChatRunResponse, латчи
  stoppingRun/localStreaming/observedRow/onStreamingChange, F7-эффект,
  observer-merge. Серверный POST /ai-chat/run не тронут.

Проверка: tsc (мои файлы чисты), vitest src/features/ai-chat 34 файла/301 тест
зелёные, grep-guard по удалённым символам пуст. Отдельное внутреннее ревью на
инварианты 7/8/9 + 204-null-safety — чисто.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-06 16:26:27 +03:00
vvzvlad ab1da408e5 Merge pull request 'feat(ai-chat): resumable SSE run-stream registry — сервер, спящий (#381 PR 1)' (#386) from feat/381-resumable-sse-pr1 into develop
Reviewed-on: #386
2026-07-06 16:24:38 +03:00
agent_coder 10d5220f5e fix(ai-chat): #381 PR1 review round 1 — open()-gate test + paused-pending byte-cap
Do 1 [test-coverage]: контроллерный тест на флаг-гейт begin-hook open() — при
OFF beginRun зовётся (durable-ран независим), а streamRegistry.open НЕ зовётся
(закрывает регресс: пустая entry → non-null paused attach → зависший SSE вместо
204); при ON — open зовётся с (chatId, runId).

Do 2 [stability]: байт-кап очереди pending paused-подписчика. pendingBytes +
overflowed на Subscriber; в paused-ветке ingestFrame при превышении
SUBSCRIBER_MAX_BUFFERED_BYTES (8MB) подписчик помечается overflowed, pending
чистится, он выбрасывается из entry.subscribers (как overflowed-entry). start()
на overflowed → onEnd (чистый 204-эквивалент, без частичного реплея). Контракт
«start() в том же тике, что attach()» задокументирован в коде — кап это
структурный бэкстоп для phase-2 Redis-await шва. Юнит-тест: paused A + live B,
9×1MB > cap → A выброшен (0 доставок), B получает все 9 живьём, поздний start(A)
→ один onEnd без реплея.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-06 07:42:36 +03:00
agent_coder 52ee3c1f3e feat(ai-chat): resumable SSE run-stream registry — server, dormant (#381 PR 1)
PR 1 of 2 (#381, фаза 1.5 #184): серверный реестр SSE-стримов агентских ранов,
чтобы любая вкладка могла подключиться к живому рану с реплеем кадров + живым
хвостом. «Спящий» — весь провод за флагом AI_CHAT_RESUMABLE_STREAM (off по
умолчанию); клиент (PR 2) ещё не написан.

- ai-chat-stream-registry.service.ts: in-memory реестр (open/bind/abortEntry/
  attach). attach — снапшот+подписка в ОДНОМ синхронном блоке (инвариант 4: нет
  await между `subscribers.add` и `frames.slice()`), paused-подписчик, overflow,
  retention с identity-guard (инвариант 2), open поверх live entry даёт ровно
  один onEnd (инвариант 3), anchor против кросс-ранового реплея (инвариант 6).
- ai-chat.controller.ts: begin-хук open(chatId, runId) + GET-attach эндпоинт
  (403 чужой чат; 204 нет-entry/finished/anchor-мисматч; cleanup до первой
  записи + recheck req.raw.destroyed; cap→destroy).
- ai-chat.service.ts: tee SSE-кадров в реестр (consumeSseStream + generateMessageId,
  гейт на runId && flag) + abortEntry из внешнего catch.
- environment.service.ts: флаг isAiChatResumableStreamEnabled().

Флаг OFF ⇒ байт-в-байт legacy И #184-фаза-1 (нет start.messageId, нет tee).
Инжектируемые провайдеры НЕ @Optional() → поломка вайринга роняет старт, а не
тихо выключает фичу.

Тесты: registry unit (16), controller.attach (9), service pipe-options (4, вкл.
flag-off-with-runId негатив), int-spec ai-chat-attach (6, реальный MockLanguageModelV3).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-06 07:22:20 +03:00
22 changed files with 3183 additions and 933 deletions
+12
View File
@@ -222,6 +222,18 @@ MCP_DOCMOST_PASSWORD=
# CLOUD=true) — run a single instance instead. The server logs a startup WARNING # CLOUD=true) — run a single instance instead. The server logs a startup WARNING
# when it detects a multi-instance deployment (CLOUD=true) so the constraint is # when it detects a multi-instance deployment (CLOUD=true) so the constraint is
# visible, and a startup sweep settles any run left dangling by a restart. # visible, and a startup sweep settles any run left dangling by a restart.
#
# Resumable run streams (#184 phase 1.5, #381). With the flag ON, an active
# durable run tees its SSE frames into an in-memory registry, and a
# reloaded/second tab attaches via GET /ai-chat/runs/:chatId/stream to follow the
# run LIVE (replay of the buffered frames + the live tail). With the flag OFF
# (default) the registry is never populated and attach always answers 204, so a
# reopened tab of an active run silently falls back to degraded 2.5s history
# polling — every wire path stays byte-for-byte identical to a build without the
# feature. Staged-rollout switch: only meaningful when autonomousRuns (above) is
# enabled for a workspace, and the same single-instance constraint applies (the
# registry is process-local).
# AI_CHAT_RESUMABLE_STREAM=false
# --- Anonymous public-share AI assistant --- # --- Anonymous public-share AI assistant ---
# Opt-in per workspace (AI settings -> "public share assistant"; off by default). # Opt-in per workspace (AI settings -> "public share assistant"; off by default).
@@ -41,17 +41,10 @@ import { extractPageSlugId } from "@/lib";
import { import {
AI_CHATS_RQ_KEY, AI_CHATS_RQ_KEY,
AI_CHAT_MESSAGES_RQ_KEY, AI_CHAT_MESSAGES_RQ_KEY,
AI_CHAT_RUN_RQ_KEY,
useAiChatMessagesQuery, useAiChatMessagesQuery,
useAiChatRunQuery,
useAiChatsQuery, useAiChatsQuery,
useAiRolesQuery, useAiRolesQuery,
} from "@/features/ai-chat/queries/ai-chat-query.ts"; } 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 { workspaceAtom } from "@/features/user/atoms/current-user-atom";
import ConversationList from "@/features/ai-chat/components/conversation-list.tsx"; import ConversationList from "@/features/ai-chat/components/conversation-list.tsx";
import ChatThread from "@/features/ai-chat/components/chat-thread.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. // Margin kept between the window and the viewport edges while dragging.
const EDGE_MARGIN = 8; 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. */ /** Compact token formatter: 1.2M / 3.4k / 950. */
function formatTokens(n: number): string { function formatTokens(n: number): string {
if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`; if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
@@ -242,150 +241,62 @@ export default function AiChatWindow() {
[roles], [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 } = 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 // #184 reconnect-and-live-follow. Whether detached agent runs are enabled for
// this workspace. The reconnect endpoint itself is NOT flag-gated server-side // this workspace. When the feature is off no runs are ever created, so the
// (it is only owner-gated and returns `{ run: null }` when the chat has no // resume attempt would only ever 204; gating ChatThread's resume on it avoids a
// run); but when the feature is off no runs are ever created, so polling it // pointless attach round-trip.
// would always come back empty — we gate it off here to avoid pointless polls.
const workspace = useAtomValue(workspaceAtom); const workspace = useAtomValue(workspaceAtom);
const autonomousRunsEnabled = const autonomousRunsEnabled =
workspace?.settings?.ai?.autonomousRuns === true; 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 // Authoritative stop of the open chat's detached run (the Stop button in
// autonomous mode). Latch "stopping" first (suppresses the re-stream flash), // autonomous mode). Request the server stop — the ONLY thing that ends a
// then request the server stop — the ONLY thing that ends a detached run; a mere // detached run; a mere local SSE abort is a client disconnect the server
// local SSE abort is a client disconnect the server ignores. On failure we // ignores. On failure surface the error.
// release the latch so the observer resumes (better to show the live run than to
// freeze the view) and surface the error.
const handleServerStop = useCallback( const handleServerStop = useCallback(
(chatId: string): void => { (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(() => { void stopRun(chatId).catch(() => {
setStoppingRun(false);
notifications.show({ notifications.show({
message: t("Failed to stop the run"), message: t("Failed to stop the run"),
color: "red", 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 // The page the user is currently viewing. AiChatWindow lives in a pathless
// parent layout route, so useParams() can't see :pageSlug. Match the full // parent layout route, so useParams() can't see :pageSlug. Match the full
// pathname against the authenticated page route instead so "the current page" // pathname against the authenticated page route instead so "the current page"
@@ -1034,16 +945,13 @@ export default function AiChatWindow() {
assistantName={currentRole?.name} assistantName={currentRole?.name}
onTurnFinished={onTurnFinished} onTurnFinished={onTurnFinished}
onServerChatId={onServerChatId} onServerChatId={onServerChatId}
// #184: live-follow a still-running run when we reopened the chat as // #184 phase 1.5: arm/disarm the degraded-poll fallback when a
// a passive observer; null when there is nothing to observe or this // resume attempt could not attach to the live run; the thread
// tab is the streamer. onStreamingChange lets the window stop polling // disarms it on settle / local stream.
// while we are the streamer. onResumeFallback={onResumeFallback}
observedRow={observedRow}
onStreamingChange={onStreamingChange}
// #184: in autonomous mode the Stop button must hit the authoritative // #184: in autonomous mode the Stop button must hit the authoritative
// server stop (a local SSE abort is a client disconnect the server // server stop (a local SSE abort is a client disconnect the server
// ignores). onServerStop also arms the "stopping" latch above so the // ignores).
// stopped run's output does not re-stream via the observer merge.
autonomousRunsEnabled={autonomousRunsEnabled} autonomousRunsEnabled={autonomousRunsEnabled}
onServerStop={handleServerStop} onServerStop={handleServerStop}
/> />
@@ -1,6 +1,13 @@
import { describe, it, expect, beforeEach, vi } from "vitest"; import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { render, screen, fireEvent, act, cleanup } from "@testing-library/react"; import {
render,
screen,
fireEvent,
act,
cleanup,
} from "@testing-library/react";
import { MantineProvider } from "@mantine/core"; 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 // 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 // above the imports) can expose the captured useChat callbacks / transport and
@@ -12,50 +19,61 @@ const h = vi.hoisted(() => ({
sendMessage: vi.fn(), sendMessage: vi.fn(),
stop: vi.fn(), stop: vi.fn(),
setMessages: 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 | { transport: null as null | {
prepareSendMessagesRequest: (arg: { prepareSendMessagesRequest?: (arg: {
messages: unknown[]; messages: unknown[];
body: Record<string, unknown>; body: Record<string, 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", () => ({ 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.onFinish = opts.onFinish ?? null;
h.state.seededMessages = opts.messages ?? null;
return { return {
messages: [], messages: [],
sendMessage: h.state.sendMessage, sendMessage: h.state.sendMessage,
status: h.state.status, status: h.state.status,
stop: h.state.stop, stop: h.state.stop,
error: null, error: null,
// #184: ChatThread reads setMessages to merge a polled observer run.
setMessages: h.state.setMessages, setMessages: h.state.setMessages,
resumeStream: h.state.resumeStream,
}; };
}, },
})); }));
// Mock "ai": deterministic ids + a transport that records its options so the test // 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", () => { vi.mock("ai", () => {
let counter = 0; let counter = 0;
return { return {
generateId: () => `gid-${counter++}`, generateId: () => `gid-${counter++}`,
DefaultChatTransport: class { DefaultChatTransport: class {
constructor(opts: { constructor(opts: Record<string, unknown>) {
prepareSendMessagesRequest: (arg: { h.state.transport = opts as never;
messages: unknown[];
body: Record<string, unknown>;
}) => { body: Record<string, unknown> };
}) {
h.state.transport = opts;
} }
}, },
}; };
}); });
// 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 // Stub the heavy children: MessageList (markdown/render) and ChatInput (the
// composer). The ChatInput stub exposes a button that queues a message, the only // composer). The ChatInput stub exposes a button that queues a message, the only
// interaction this test needs to populate the queue while "streaming". // 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" />, default: () => <div data-testid="message-list" />,
})); }));
vi.mock("@/features/ai-chat/components/chat-input.tsx", () => ({ vi.mock("@/features/ai-chat/components/chat-input.tsx", () => ({
default: ({ onQueue }: { onQueue: (text: string) => void }) => ( default: ({
onQueue,
onStop,
}: {
onQueue: (text: string) => void;
onStop: () => void;
}) => (
<>
<button data-testid="queue-btn" onClick={() => onQueue("queued text")}> <button data-testid="queue-btn" onClick={() => onQueue("queued text")}>
queue queue
</button> </button>
<button aria-label="Stop" onClick={() => onStop()}>
stop
</button>
</>
), ),
})); }));
import ChatThread from "./chat-thread"; 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(); const onTurnFinished = vi.fn();
render( 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> <MantineProvider>
<ChatThread chatId="c1" initialRows={[]} onTurnFinished={onTurnFinished} /> <ChatThread
</MantineProvider>, 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)", () => { describe("ChatThread — send now (#198)", () => {
beforeEach(() => { beforeEach(resetState);
h.state.status = "streaming";
h.state.onFinish = null;
h.state.sendMessage.mockClear();
h.state.stop.mockClear();
h.state.transport = null;
});
it("aborts the current turn and resends the queued message on the abort", () => { it("aborts the current turn and resends the queued message on the abort", () => {
renderThread(); renderThread();
// Queue a message while the turn is streaming.
fireEvent.click(screen.getByTestId("queue-btn")); fireEvent.click(screen.getByTestId("queue-btn"));
const sendNowBtn = screen.getByLabelText("Send now"); const sendNowBtn = screen.getByLabelText("Send now");
expect(sendNowBtn).toBeTruthy(); 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); fireEvent.click(sendNowBtn);
expect(h.state.stop).toHaveBeenCalledTimes(1); expect(h.state.stop).toHaveBeenCalledTimes(1);
expect(h.state.sendMessage).not.toHaveBeenCalled(); expect(h.state.sendMessage).not.toHaveBeenCalled();
// The abort we triggered reaches onFinish: the promoted head is flushed.
act(() => { act(() => {
h.state.onFinish?.({ h.state.onFinish?.({
message: { id: "a", role: "assistant", parts: [] }, message: { id: "a", role: "assistant", parts: [] },
@@ -122,10 +181,8 @@ describe("ChatThread — send now (#198)", () => {
fireEvent.click(screen.getByTestId("queue-btn")); fireEvent.click(screen.getByTestId("queue-btn"));
fireEvent.click(screen.getByLabelText("Send now")); fireEvent.click(screen.getByLabelText("Send now"));
const prep = h.state.transport!.prepareSendMessagesRequest; const prep = h.state.transport!.prepareSendMessagesRequest!;
// The send right after "send now" carries interrupted: true...
expect(prep({ messages: [], body: {} }).body.interrupted).toBe(true); 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); 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.getByTestId("queue-btn"));
fireEvent.click(screen.getByLabelText("Send now")); 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.stop).not.toHaveBeenCalled();
expect(h.state.sendMessage).toHaveBeenCalledWith({ text: "queued text" }); 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); 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)", () => { describe("ChatThread — turn-end decision (onFinish)", () => {
beforeEach(() => { beforeEach(resetState);
h.state.status = "streaming";
h.state.onFinish = null;
h.state.sendMessage.mockClear();
h.state.stop.mockClear();
h.state.transport = null;
});
// 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: { function finishWith(flags: {
isAbort?: boolean; isAbort?: boolean;
isDisconnect?: boolean; isDisconnect?: boolean;
isError?: 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(); cleanup();
h.state.sendMessage.mockClear(); h.state.sendMessage.mockClear();
const { onTurnFinished } = renderThread(); const { onTurnFinished } = renderThread();
// Populate the queue while the turn is streaming.
fireEvent.click(screen.getByTestId("queue-btn")); fireEvent.click(screen.getByTestId("queue-btn"));
act(() => { act(() => {
h.state.onFinish?.({ h.state.onFinish?.({
@@ -187,16 +226,12 @@ describe("ChatThread — turn-end decision (onFinish)", () => {
it("CONTINUES — flushes the next queued message on a clean finish", () => { it("CONTINUES — flushes the next queued message on a clean finish", () => {
finishWith({}); finishWith({});
// Clean finish (no terminal flag): the queued message is auto-sent.
expect(h.state.sendMessage).toHaveBeenCalledWith({ text: "queued text" }); expect(h.state.sendMessage).toHaveBeenCalledWith({ text: "queued text" });
// A clean finish shows no stop notice.
expect(screen.queryByText("Response stopped.")).toBeNull(); expect(screen.queryByText("Response stopped.")).toBeNull();
}); });
it("ENDS — keeps the queue intact on a user abort and shows the stopped notice", () => { it("ENDS — keeps the queue intact on a user abort and shows the stopped notice", () => {
finishWith({ isAbort: true }); 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(h.state.sendMessage).not.toHaveBeenCalled();
expect(screen.getByText("Response stopped.")).toBeTruthy(); 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)", () => { it("ENDS — keeps the queue intact on a stream error (no auto-retry, no stopped notice)", () => {
finishWith({ isError: true }); finishWith({ isError: true });
// Blindly retrying after a failure would be wrong; the queue is left alone.
expect(h.state.sendMessage).not.toHaveBeenCalled(); expect(h.state.sendMessage).not.toHaveBeenCalled();
// isError clears the neutral notice (the error banner covers this case).
expect(screen.queryByText("Response stopped.")).toBeNull(); expect(screen.queryByText("Response stopped.")).toBeNull();
}); });
it("notifies the parent on EVERY terminal outcome", () => { 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 [ for (const flags of [
{}, {},
{ isAbort: true }, { isAbort: true },
@@ -232,55 +263,411 @@ describe("ChatThread — turn-end decision (onFinish)", () => {
}); });
}); });
// #184 passive-observer merge: when reconnecting to a still-running run, the // #184 phase 1.5: the resumable-SSE client. A reopened tab resumes the live run
// parent feeds the polled run message via `observedRow`; ChatThread merges it via // via the SDK's reconnect transport (attach: replay + tail) instead of polling.
// setMessages — but ONLY when this tab is NOT itself streaming (the streamer's describe("ChatThread — resume (attach) machinery (#184)", () => {
// SSE owns the view, so a stale observedRow must never overwrite it). const streamingTail = () => [
describe("ChatThread — observer run merge (#184)", () => { row("u1", "user", undefined, "hi"),
beforeEach(() => { row("a1", "assistant", "streaming", "partial"),
h.state.onFinish = null; ];
h.state.setMessages.mockReset(); 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 = { it("strips the streaming tail from the seed, but keeps a user tail whole", () => {
id: "a-run", renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() });
role: "assistant", // 2 rows in, streaming tail stripped -> 1 seeded message.
content: "step 1\nstep 2", expect(h.state.seededMessages).toHaveLength(1);
metadata: {
parts: [{ type: "text", text: "step 1\nstep 2" }],
},
createdAt: "2026-01-01T00:00:00Z",
} as const;
function renderObserver(status: string) { cleanup();
h.state.status = status; renderThread({ autonomousRunsEnabled: true, initialRows: userTail() });
render( // 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> <MantineProvider>
<ChatThread <ChatThread
chatId="c1" chatId="c1"
initialRows={[]} initialRows={rows}
autonomousRunsEnabled
onTurnFinished={vi.fn()} onTurnFinished={vi.fn()}
observedRow={observedRow as never} onResumeFallback={onResumeFallback}
/> />
</MantineProvider>, </MantineProvider>
</QueryClientProvider>
); );
} const view = render(<Wrapper rows={initialRows} />);
const rerender = (rows: IAiChatMessageRow[]) =>
it("merges the polled run message when this tab is a passive observer", () => { act(() => view.rerender(<Wrapper rows={rows} />));
renderObserver("ready"); return { rerender, onResumeFallback };
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();
});
});
@@ -1,4 +1,5 @@
import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { useQueryClient } from "@tanstack/react-query";
import { generateId } from "ai"; import { generateId } from "ai";
import { ActionIcon, Box, Group, Stack, Text, Tooltip } from "@mantine/core"; import { ActionIcon, Box, Group, Stack, Text, Tooltip } from "@mantine/core";
import { import {
@@ -24,7 +25,14 @@ import {
} from "@/features/ai-chat/utils/role-launch.ts"; } from "@/features/ai-chat/utils/role-launch.ts";
import { describeChatError } from "@/features/ai-chat/utils/error-message.ts"; import { describeChatError } from "@/features/ai-chat/utils/error-message.ts";
import { extractServerChatId } from "@/features/ai-chat/utils/adopt-chat-id.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 { import {
dequeue, dequeue,
enqueueMessage, enqueueMessage,
@@ -87,19 +95,13 @@ interface ChatThreadProps {
* Copy/export button available mid-stream). Distinct from onTurnFinished, * Copy/export button available mid-stream). Distinct from onTurnFinished,
* which fires only at the terminal outcome. */ * which fires only at the terminal outcome. */
onServerChatId?: (serverChatId?: string) => void; onServerChatId?: (serverChatId?: string) => void;
/** #184 reconnect-and-live-follow. When THIS tab reopened a chat whose agent /** #184 phase 1.5: arm/disarm the parent's degraded-poll fallback for THIS
* run is still going (it is a PASSIVE OBSERVER it did not start the run here), * chat's window. Called `true` when a resume attempt could not attach to the
* the parent polls the reconnect endpoint and feeds the run's incrementally- * live run (attach 204 / starved-or-torn resumed finish), so the window starts
* persisted assistant message here; we merge it into the live list so new * a dumb timed poll of the message history to follow the detached run to settle;
* steps/tool-calls appear as they are persisted. Null when there is nothing to * called `false` the moment a local stream starts or the terminal settled row is
* observe (no run, feature off, or this tab IS the streamer). The merge is * merged (invariant 8). The window owns the timer + its 10-min cap. */
* ADDITIONALLY guarded by our own `isStreaming`, so a stale value can never onResumeFallback?: (active: boolean) => void;
* 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: whether detached/autonomous agent runs are enabled for this workspace. /** #184: whether detached/autonomous agent runs are enabled for this workspace.
* When true the Stop button must additionally hit the AUTHORITATIVE server stop * When true the Stop button must additionally hit the AUTHORITATIVE server stop
* (via onServerStop) aborting only the local SSE is just a client disconnect, * (via onServerStop) aborting only the local SSE is just a client disconnect,
@@ -155,15 +157,58 @@ export default function ChatThread({
assistantName, assistantName,
onTurnFinished, onTurnFinished,
onServerChatId, onServerChatId,
observedRow, onResumeFallback,
onStreamingChange,
autonomousRunsEnabled, autonomousRunsEnabled,
onServerStop, onServerStop,
}: ChatThreadProps) { }: ChatThreadProps) {
const { t } = useTranslation(); 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[]>( const initialMessages = useMemo<UIMessage[]>(
() => (initialRows ?? []).map(rowToUiMessage), () =>
seedRows(
initialRows ?? [],
attemptResumeRef.current && stripRef.current,
).map(rowToUiMessage),
[initialRows], [initialRows],
); );
@@ -261,9 +306,12 @@ export default function ChatThread({
const { head, rest } = dequeue(queuedRef.current); const { head, rest } = dequeue(queuedRef.current);
if (!head) return false; if (!head) return false;
setQueue(rest); 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 }); sendMessageRef.current?.({ text: head.text });
return true; return true;
}, [setQueue]); }, [setQueue, setResumedTurnPair]);
const enqueue = useCallback( const enqueue = useCallback(
(text: string) => { (text: string) => {
@@ -283,6 +331,47 @@ export default function ChatThread({
new DefaultChatTransport<UIMessage>({ new DefaultChatTransport<UIMessage>({
api: "/api/ai-chat/stream", api: "/api/ai-chat/stream",
credentials: "include", 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 // Inject the chat id and the currently-open page alongside the useChat
// messages so the server can resolve an existing chat (or create one // 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 // 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 // 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 // generated client id (never `undefined`) so the store is NOT re-created on
// every render mid-stream (see `chatStoreId` above). // 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 // would be wrong, so on Stop/disconnect/error the queue is left intact for
// the user to decide. // the user to decide.
onFinish: ({ message, isAbort, isDisconnect, isError }) => { 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 // Forward the authoritative server chatId (streamed on the assistant
// message metadata) so the parent adopts the REAL created chat id for a new // 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 // 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 (isAbort) setStopNotice("manual");
else if (isDisconnect) setStopNotice("disconnect"); else if (isDisconnect) setStopNotice("disconnect");
else setStopNotice(null); 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 // "Send now": WE triggered this abort to interrupt the current turn and
// immediately send the promoted head. Flush it even though the turn was // immediately send the promoted head. Flush it even though the turn was
// aborted (the normal abort path below keeps the queue intact). The // 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"; const isStreaming = status === "submitted" || status === "streaming";
// #184: report our live streaming status up so the parent stops polling the run // 204-handler (`onNoActiveStream`): the attach returned 204 — nothing live to
// while WE are the streamer (the SSE owns the view) and resumes once we go idle. // resume (overflow / begin-failure / after retention / anchor-mismatch). One-
// Effect (not render) so it never updates parent state during our own render; // shot via noStreamHandledRef (we do NOT null onNoActiveStreamRef). Exactly four
// fires on mount with `false`, which also re-syncs the parent after a chat // parts. Kept in a ref (read by the transport's fetch closure) and refreshed
// switch remounts this thread (a fresh mount is idle until the user sends). // each render below.
useEffect(() => { const onNoActiveStream = useCallback(() => {
onStreamingChange?.(isStreaming); // A late attach outcome after unmount must not arm a poll / invalidate on the
}, [isStreaming, onStreamingChange]); // 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 // Mount effect: kick off the resume attempt for a non-settled tail. Marking the
// reopened a chat whose run is still going and did NOT start it here), merge it // turn as resumed BEFORE resumeStream so onFinish (invariant 7/8) sees 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.
useEffect(() => { useEffect(() => {
if (isStreaming || !observedRow) return; // Re-arm on (re)mount — StrictMode dev-mounts twice, and the cleanup below
const observed = rowToUiMessage(observedRow); // flips this false between the two.
setMessages((prev) => mergeObservedMessage(prev, observed)); mountedRef.current = true;
}, [observedRow, isStreaming, setMessages]); 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 now" on a queued message: interrupt the current turn and immediately
// send THIS message, keeping the agent's partial output. Other queued messages // 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); const msg = queuedRef.current.find((m) => m.id === id);
if (!msg) return; if (!msg) return;
setQueue(removeQueuedById(queuedRef.current, id)); setQueue(removeQueuedById(queuedRef.current, id));
// Local send: clear any resume-suppression flag (invariant 8).
setResumedTurnPair(false);
sendMessageRef.current?.({ text: msg.text }); sendMessageRef.current?.({ text: msg.text });
} }
}, },
[setQueue, stop], [setQueue, stop, setResumedTurnPair],
); );
// Stop the current turn. ALWAYS abort the local SSE (`stop()`) so the composer // 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 — // 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). // only the local abort happens (there is no server-side run handle to stop yet).
const handleStop = useCallback(() => { 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(); stop();
if (!autonomousRunsEnabled) return; if (!autonomousRunsEnabled) return;
if (chatIdRef.current) { if (chatIdRef.current) {
@@ -617,6 +827,11 @@ export default function ChatThread({
<Text size="xs" lineClamp={2} className={classes.queuedText}> <Text size="xs" lineClamp={2} className={classes.queuedText}>
{m.text} {m.text}
</Text> </Text>
{/* "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> <Tooltip label={t("Interrupt and send now")} withArrow>
<ActionIcon <ActionIcon
size="xs" size="xs"
@@ -628,6 +843,7 @@ export default function ChatThread({
<IconPlayerPlayFilled size={12} /> <IconPlayerPlayFilled size={12} />
</ActionIcon> </ActionIcon>
</Tooltip> </Tooltip>
)}
<ActionIcon <ActionIcon
size="xs" size="xs"
variant="subtle" variant="subtle"
@@ -642,7 +858,11 @@ export default function ChatThread({
</Stack> </Stack>
)} )}
<ChatInput <ChatInput
onSend={(text) => sendMessage({ text })} onSend={(text) => {
// Local send: clear any resume-suppression flag (invariant 8).
setResumedTurnPair(false);
sendMessage({ text });
}}
onQueue={enqueue} onQueue={enqueue}
onStop={handleStop} onStop={handleStop}
isStreaming={isStreaming} 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, deleteAiChat,
deleteAiRole, deleteAiRole,
getAiChatMessages, getAiChatMessages,
getAiChatRun,
getAiChats, getAiChats,
getAiRoleCatalog, getAiRoleCatalog,
getAiRoleCatalogBundle, getAiRoleCatalogBundle,
@@ -26,7 +25,6 @@ import {
import { import {
IAiChat, IAiChat,
IAiChatMessageRow, IAiChatMessageRow,
IAiChatRunResponse,
IAiRole, IAiRole,
IAiRoleCatalog, IAiRoleCatalog,
IAiRoleCatalogBundle, IAiRoleCatalogBundle,
@@ -37,7 +35,6 @@ import {
IAiRoleUpdateFromCatalogResult, IAiRoleUpdateFromCatalogResult,
} from "@/features/ai-chat/types/ai-chat.types.ts"; } from "@/features/ai-chat/types/ai-chat.types.ts";
import { IPagination } from "@/lib/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_CHATS_RQ_KEY = ["ai-chats"];
export const AI_ROLES_RQ_KEY = ["ai-roles"]; export const AI_ROLES_RQ_KEY = ["ai-roles"];
@@ -55,7 +52,6 @@ export const AI_CHAT_MESSAGES_RQ_KEY = (chatId: string) => [
"ai-chat-messages", "ai-chat-messages",
chatId, 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). */ /** Paginated list of the current user's chats (auto-loads further pages). */
export function useAiChatsQuery() { export function useAiChatsQuery() {
@@ -89,7 +85,15 @@ export function useAiChatsQuery() {
* Load all persisted messages of a chat (oldest first), flattening the * Load all persisted messages of a chat (oldest first), flattening the
* paginated server response. Used to seed `useChat` initial messages. * 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({ const query = useInfiniteQuery({
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatId ?? ""), queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatId ?? ""),
queryFn: ({ pageParam }) => queryFn: ({ pageParam }) =>
@@ -100,6 +104,7 @@ export function useAiChatMessagesQuery(chatId: string | undefined) {
? (lastPage.meta.nextCursor ?? undefined) ? (lastPage.meta.nextCursor ?? undefined)
: undefined, : undefined,
enabled: !!chatId, enabled: !!chatId,
refetchInterval,
}); });
// useInfiniteQuery only fetches the first page on its own. The hook's contract // 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() { export function useRenameAiChatMutation() {
const queryClient = useQueryClient(); const queryClient = useQueryClient();
const { t } = useTranslation(); 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, IAiChatListParams,
IAiChatMessageRow, IAiChatMessageRow,
IAiChatMessagesParams, IAiChatMessagesParams,
IAiChatRunResponse,
IAiRole, IAiRole,
IAiRoleCatalog, IAiRoleCatalog,
IAiRoleCatalogBundle, IAiRoleCatalogBundle,
@@ -43,23 +42,6 @@ export async function getAiChatMessages(
return req.data; 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 * 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) * 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. // renders a "stopped" marker on interrupted turns.
finishReason?: string; finishReason?: string;
} | null; } | 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; 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 IAiChatListParams extends QueryParams {}
export interface IAiChatMessagesParams { 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, ConflictException,
Controller, Controller,
ForbiddenException, ForbiddenException,
Get,
HttpCode, HttpCode,
HttpException, HttpException,
HttpStatus, HttpStatus,
Logger, Logger,
Param,
ParseUUIDPipe,
Post, Post,
Query,
Req, Req,
Res, Res,
ServiceUnavailableException, ServiceUnavailableException,
@@ -54,6 +58,12 @@ import {
} from './dto/ai-chat.dto'; } from './dto/ai-chat.dto';
import { describeProviderError } from '../../integrations/ai/ai-error.util'; import { describeProviderError } from '../../integrations/ai/ai-error.util';
import { buildChatMarkdown } from './chat-markdown.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 * 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 aiChatMessageRepo: AiChatMessageRepo,
private readonly aiTranscription: AiTranscriptionService, private readonly aiTranscription: AiTranscriptionService,
private readonly pageRepo: PageRepo, 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). */ /** List the requesting user's chats in this workspace (paginated). */
@@ -233,6 +248,102 @@ export class AiChatController {
return { stopped }; 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. */ /** Rename a chat. */
@HttpCode(HttpStatus.OK) @HttpCode(HttpStatus.OK)
@Post('rename') @Post('rename')
@@ -344,13 +455,25 @@ export class AiChatController {
// its progress, and settle its terminal status — see AiChatRunService. // its progress, and settle its terminal status — see AiChatRunService.
const runHooks: AiChatRunHooks | undefined = autonomousRuns const runHooks: AiChatRunHooks | undefined = autonomousRuns
? { ? {
begin: (chatId) => begin: async (chatId) => {
this.aiChatRunService.beginRun({ const handle = await this.aiChatRunService.beginRun({
chatId, chatId,
workspaceId: workspace.id, workspaceId: workspace.id,
userId: user.id, userId: user.id,
trigger: 'user', 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) => onAssistantSeeded: (runId, messageId) =>
this.aiChatRunService.linkAssistantMessage( this.aiChatRunService.linkAssistantMessage(
runId, runId,
@@ -4,6 +4,7 @@ import { TokenModule } from '../auth/token.module';
import { AiChatController } from './ai-chat.controller'; import { AiChatController } from './ai-chat.controller';
import { AiChatService } from './ai-chat.service'; import { AiChatService } from './ai-chat.service';
import { AiChatRunService } from './ai-chat-run.service'; import { AiChatRunService } from './ai-chat-run.service';
import { AiChatStreamRegistryService } from './ai-chat-stream-registry.service';
import { AiTranscriptionService } from './ai-transcription.service'; import { AiTranscriptionService } from './ai-transcription.service';
import { AiChatToolsService } from './tools/ai-chat-tools.service'; import { AiChatToolsService } from './tools/ai-chat-tools.service';
import { EmbeddingModule } from './embedding/embedding.module'; import { EmbeddingModule } from './embedding/embedding.module';
@@ -44,6 +45,7 @@ import { PublicShareChatToolsService } from './tools/public-share-chat-tools.ser
providers: [ providers: [
AiChatService, AiChatService,
AiChatRunService, AiChatRunService,
AiChatStreamRegistryService,
AiTranscriptionService, AiTranscriptionService,
AiChatToolsService, AiChatToolsService,
PublicShareChatService, 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 { import {
AiChatService, AiChatService,
compactToolOutput, compactToolOutput,
@@ -1059,3 +1068,181 @@ describe('isInterruptResume', () => {
expect(isInterruptResume(withPrev(null), true)).toBe(false); 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 { AiChatToolsService } from './tools/ai-chat-tools.service';
import { McpClientsService } from './external-mcp/mcp-clients.service'; import { McpClientsService } from './external-mcp/mcp-clients.service';
import { EnvironmentService } from '../../integrations/environment/environment.service'; import { EnvironmentService } from '../../integrations/environment/environment.service';
import { AiChatStreamRegistryService } from './ai-chat-stream-registry.service';
import { buildSystemPrompt } from './ai-chat.prompt'; import { buildSystemPrompt } from './ai-chat.prompt';
import { import {
CORE_TOOL_KEYS, CORE_TOOL_KEYS,
@@ -276,6 +277,10 @@ export class AiChatService implements OnModuleInit {
// Reads the AI_CHAT_DEFERRED_TOOLS toggle (#332). Injected last so existing // Reads the AI_CHAT_DEFERRED_TOOLS toggle (#332). Injected last so existing
// positional constructor callers (tests) only append one stub. // positional constructor callers (tests) only append one stub.
private readonly environment: EnvironmentService, 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. // as the cumulative authoritative usage so the client never jumps DOWN.
let cumulativeStepUsage: ChatStreamUsage | undefined; let cumulativeStepUsage: ChatStreamUsage | undefined;
result.pipeUIMessageStreamToResponse(res.raw, { 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' }, headers: { 'X-Accel-Buffering': 'no' },
// Surface the authoritative chatId on the streamed assistant UI message so // 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 // 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 // finalizeRun (onSettled) is idempotent — a settle here and a settle from a
// streamText callback collapse to a single terminal write. // streamText callback collapse to a single terminal write.
if (runId) { 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?.( await runHooks?.onSettled?.(
runId, runId,
'error', 'error',
@@ -292,6 +292,23 @@ export class EnvironmentService {
return enabled === 'true'; 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 { getPostHogHost(): string {
return this.configService.get<string>('POSTHOG_HOST'); 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();
}
});
});