From daeeb1f3f2ce2f9807694532fe9dd9f00dd5f198 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 18:34:33 +0300 Subject: [PATCH] =?UTF-8?q?refactor(ai-chat):=20registry=20=E2=80=94=20ste?= =?UTF-8?q?p-aligned=20retention=20+=20tail-only=20attach=20(#491)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Реестр ран-стримов больше не буферизует до 32МБ сырых SSE-кадров на активный ран и не выливает весь буфер в сокет синхронно при attach — это давало OOM на 1ГБ-контейнере при нескольких марафонских ранах. Теперь кольцо ограничено (env-настраиваемо, по умолчанию 4МБ) и держится в границах за счёт ротации по шагам. Серверная часть (суть коммита): - Штамповка кадров по шагам в ingestFrame. Штамп кадра = число `finish-step` кадров ДО него (с 0); сам finish-step несёт текущее значение, затем счётчик инкрементится. Так штамп совпадает с `metadata.stepsPersisted`: клиент с N персистнутыми шагами имеет 0..N-1 в сиде и просит хвост `stamp >= N`. Границу ловим дешёвым startsWith по `data: {"type":"finish-step"` — форма кадра проверена эмпирически против ai@6.0.207 (одна часть на кадр, type всегда первый ключ; кавычки в text-delta экранированы, ложных срабатываний нет). - Кольцо ротируется ТОЛЬКО на подтверждённом персисте шага N (`confirmPersistedStep`), сбрасывая кадры `stamp < N` (эти шаги уже на диске и придут в свежем сиде). `updateStreaming` теперь СИГНАЛИЗИРУЕТ исход (число персистнутых шагов или null), и ротация вызывается лишь при не-null возврате — провал персиста ничего не ротирует, кольцо покрывает БОЛЬШЕ (анти-инверсия: наивная ротация в .then() после НЕзаписанного шага дырявила бы гарантию). - Переполнение кольца сверх байтового капа вытесняет старейшие кадры; вытеснение ещё-не-персистнутого кадра открывает GAP. Гэп НЕ липкий: floor покрытия считается из кольца, поздний персист, проротировав дырявые шаги, его чистит. - attach(chatId, anchor, n): маркер шага N приходит ТОЛЬКО от клиента (сервер не читает строку — N из устаревшего сида дал бы тихую дыру в один шаг). Покрытие ОК ⟺ coverageFloor <= n; иначе 204 → клиент рефетчит (больший N) и переподключается. Хвост = синтетический `start`-кадр (ран-факт runId/chatId) + кадры `stamp >= n`. Инвариант 6 (нет кросс-ран реплея) сохранён через anchor; инвариант 4 (снапшот+регистрация в один синхронный тик) сохранён. N-срез применяется во ВСЕХ ветках, включая finished-retained: finished + N=N_final → пустой хвост + finish-кадр, клиент закрывает стрим. - Контроллер пишет хвост чанками с учётом drain (writeTailRespectingDrain), а не синхронным залпом (вторая половина OOM). Кап подписчика — производное 2× кап кольца, обе величины env-резолвятся на инстансе. Клиент: в тип строки добавлен `metadata.stepsPersisted` (источник N). PIN-SPEC трип-вайр на ai@6.0.207: `readUIMessageStream({ message })` продолжает последнее сообщение, `start`-кадр не сбрасывает parts, текст не пересекает finish-step — на этом держится продолжение при attach; апгрейд ai теперь падает громко. Тесты (observable-property против РЕАЛЬНОГО реестра/БД): детектор границы на реальной форме кадра, N-срез (в т.ч. посреди шага), ротация только на подтверждённом персисте, «персист провалился но кольцо влезло → attach успешен / провал + переполнение → 204», «устаревший N → 204 → после рефетча успех», очистка гэпа поздним персистом, finished-retained + N_final, memory-bound (5 параллельных марафонов сверх 32МБ, каждое кольцо ≤ кап). Обновлены registry/controller specs и DB-backed интеграционный attach-spec под новую сигнатуру/семантику. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../features/ai-chat/types/ai-chat.types.ts | 6 + .../utils/sdk-continuation-tripwire.test.ts | 83 +++ .../ai-chat-stream-registry.service.ts | 348 ++++++++++--- .../ai-chat/ai-chat-stream-registry.spec.ts | 483 ++++++++++++------ .../ai-chat/ai-chat.controller.attach.spec.ts | 22 +- .../src/core/ai-chat/ai-chat.controller.ts | 128 +++-- .../src/core/ai-chat/ai-chat.service.ts | 52 +- .../integration/ai-chat-attach.int-spec.ts | 71 ++- 8 files changed, 881 insertions(+), 312 deletions(-) create mode 100644 apps/client/src/features/ai-chat/utils/sdk-continuation-tripwire.test.ts diff --git a/apps/client/src/features/ai-chat/types/ai-chat.types.ts b/apps/client/src/features/ai-chat/types/ai-chat.types.ts index 2c987263..acb84de5 100644 --- a/apps/client/src/features/ai-chat/types/ai-chat.types.ts +++ b/apps/client/src/features/ai-chat/types/ai-chat.types.ts @@ -181,6 +181,12 @@ export interface IAiChatMessageRow { toolCalls?: unknown; metadata?: { parts?: UIMessage["parts"]; + // #491 step-alignment anchor: the count of FINISHED steps whose parts are in + // THIS row, written atomically with `parts` server-side (flushAssistant). The + // resume client reads it as its persisted step frontier N — the tail-only + // attach asks the run-stream registry for the frames of step N onward (the + // seed already carries steps 0..N-1). Absent on pre-#491 rows -> read as 0. + stepsPersisted?: number; // AI SDK v6 `totalUsage` persisted on assistant rows. Legacy cumulative // figure (sum of every step's usage for the turn); kept for back-compat and // as the fallback for older rows that have no `contextTokens`. diff --git a/apps/client/src/features/ai-chat/utils/sdk-continuation-tripwire.test.ts b/apps/client/src/features/ai-chat/utils/sdk-continuation-tripwire.test.ts new file mode 100644 index 00000000..c9fc19dc --- /dev/null +++ b/apps/client/src/features/ai-chat/utils/sdk-continuation-tripwire.test.ts @@ -0,0 +1,83 @@ +import { describe, it, expect } from "vitest"; +import { readUIMessageStream, type UIMessage } from "ai"; +import pkg from "../../../../package.json"; + +/** + * PIN-SPEC TRIP-WIRE (#491). The tail-only attach continuation relies on THREE + * behaviors of `ai@6.0.207`, verified line-by-line in the issue. Without this + * test, an `ai` bump could silently break attach (the client would append the + * live tail to the wrong message, or duplicate a step): + * + * 1. `readUIMessageStream({ message })` CONTINUES the passed message — it does + * not start a fresh one — so the tail streamed after a re-seed is appended to + * the seeded assistant row (the same DB id). + * 2. A `start` frame does NOT reset the existing message's parts (so the seeded + * steps 0..N-1 survive; the synthetic `start` the registry prepends only + * carries the run-fact metadata). + * 3. Text parts do NOT cross a `finish-step` boundary — a new `text-start` after + * `finish-step` is a NEW part — so the reconstructed steps stay separated and + * the step frontier stays meaningful. + * + * If an `ai` upgrade changes any of these, this test fails LOUD instead of the + * resume path silently corrupting. + */ +describe("ai SDK continuation trip-wire (#491, tail-only attach)", () => { + it("is pinned to the exact ai version the continuation was verified against", () => { + // A caret/range bump is exactly what would silently break attach — require an + // exact pin. Bumping ai MUST re-verify the behavior asserted below, then this. + expect((pkg as { dependencies: Record }).dependencies.ai).toBe( + "6.0.207", + ); + }); + + it("continues the seeded message: start does not reset parts, the tail appends as new parts", async () => { + // A seeded assistant row with ONE finished step already reconstructed. + const seeded: UIMessage = { + id: "assistant-1", + role: "assistant", + parts: [ + { type: "step-start" }, + { type: "text", text: "STEP0", state: "done" }, + ], + } as UIMessage; + + // The tail the registry delivers on re-attach: a synthetic start (run-fact), + // then step 1's frames, then finish. As UI-message chunks (what the SSE frames + // decode to). + const chunks = [ + { type: "start", messageMetadata: { runId: "r1", chatId: "c1" } }, + { type: "start-step" }, + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: "STEP1" }, + { type: "text-end", id: "t1" }, + { type: "finish-step" }, + { type: "finish" }, + ]; + const stream = new ReadableStream({ + start(c) { + for (const ch of chunks) c.enqueue(ch); + c.close(); + }, + }); + + let last: UIMessage | undefined; + for await (const msg of readUIMessageStream({ message: seeded, stream })) { + last = msg; + } + + expect(last).toBeDefined(); + // Same message id (continuation, not a fresh message). + expect(last!.id).toBe("assistant-1"); + // The seeded step-0 parts SURVIVED the `start` frame, and step 1 was appended + // as SEPARATE parts (text did not cross the finish-step boundary). + const shape = last!.parts.map((p) => `${p.type}:${(p as { text?: string }).text ?? ""}`); + expect(shape).toEqual([ + "step-start:", + "text:STEP0", + "step-start:", + "text:STEP1", + ]); + // The run-fact metadata from the synthetic start frame is applied. + expect(last!.metadata).toMatchObject({ runId: "r1", chatId: "c1" }); + }); +}); diff --git a/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts b/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts index 1ef476f5..5fb3c15a 100644 --- a/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts @@ -1,49 +1,132 @@ 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. + * In-memory run-stream registry (#184 phase 1.5, step-aligned retention #491). 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`, be handed the TAIL past the step it already + * has persisted, 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. + * + * ── #491 step-aligned retention (the OOM fix) ──────────────────────────────── + * The old registry buffered up to 32MB of raw SSE frames PER active run (V8 ~2× + * in memory) and, on attach, blasted the WHOLE buffer to the socket synchronously + * with no drain — a handful of marathon runs on a 1GB container OOM'd. #491 caps + * the ring at a few MB (env-tunable, default 4MB) and keeps it there by ROTATING: + * + * - Every buffered frame is STAMPED with a step number at tee (see ingestFrame). + * Convention: the stamp of a frame is the number of `finish-step` parts seen + * BEFORE it (starting at 0). The finish-step frame itself carries the current + * value, THEN the counter increments. So a frame stamped `s` is the content of + * the (s+1)-th step — 0-based step index `s` — and the stamp aligns EXACTLY + * with `metadata.stepsPersisted`: a client whose persisted `stepsPersisted` is + * N has steps 0..N-1 on disk (and in its seed) and needs the tail `stamp >= N`. + * + * - The ring rotates ONLY on a CONFIRMED persist of step N + * (`confirmPersistedStep`), dropping frames with `stamp < N` (those steps are + * now on disk and a fresh client seed carries them). A NON-confirmed step is + * never rotated away, so a persist FAILURE just makes the ring cover MORE + * (auto-safe). This is the anti-inversion rule: a naive "rotate in .then()" + * that rotated after an UNwritten step would drop a step nobody has → silent + * hole. Rotation is gated on a real, successful persist. + * + * - If the ring still exceeds its byte cap after rotation (a single fat step, or + * a lagging persist), the OLDEST frames are evicted to stay bounded. Evicting a + * not-yet-persisted frame opens a GAP: an attach whose N falls at or below an + * evicted step answers 204 and the client degrades to restore+poll. The gap is + * NOT sticky — the coverage floor is recomputed from the ring, so a later + * persist that rotates past the holey steps clears it. + * + * ── attach numbering / coverage (the wire convention) ──────────────────────── + * The step marker N comes ONLY FROM THE CLIENT (a query param). The server never + * reads the row to derive N — a server-side N from a stale seed would open a + * silent one-step hole. N is the client's persisted `stepsPersisted` (a COUNT): + * - the tail it needs = frames with `stamp >= N`; + * - coverage is OK ⟺ `coverageFloor(entry) <= N`, where coverageFloor is the + * smallest step FULLY present in the ring (its smallest retained stamp, bumped + * by one when that leading step was only partially evicted by overflow). If + * `coverageFloor > N` the ring starts AFTER the client's frontier (a hole, or + * the client's seed simply lagged behind a rotation) → 204 → the client + * refetches (a larger N) and re-attaches. + * The N cutoff is applied in ALL branches, INCLUDING the finished-retained replay. + * + * ── same-tick invariants (unchanged, still load-bearing) ───────────────────── + * invariant 1: only the matching run may mutate/observe an entry (runId check). + * invariant 2: retention deletes ONLY its own entry (a replacement may own the key). + * invariant 3: open() over a live entry mirrors the done-path (subscribers released). + * invariant 4: the tail SLICE + subscriber registration happen in ONE synchronous + * tick inside attach() — no await between them — so a concurrently + * ingested frame is EITHER in the snapshot (buffered before the sync + * block, and the just-added subscriber never sees it) OR fanned out to + * the paused subscriber's `pending` (ingested after) — never both and + * never neither: no loss, no duplication. NOTE (#491): the controller + * now AWAITS the drain-respecting tail write BEFORE calling start(), so + * frames ingested during that await accumulate in `pending`; this is + * bounded by the subscriber cap (an overflow degrades start() to an + * end(), a 204-equivalent). It is the SYNCHRONOUS snapshot+registration + * — not a same-tick start() — that makes this correct. + * invariant 5: the controller wires close-cleanup BEFORE any write. + * invariant 6: no cross-run replay — the `anchor` (the client's assistant row id) + * must match this run's assistant id, or a foreign run's transcript + * would be appended to the client's message. */ /** 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, and - * the client falls back to its restore + degraded-poll path, #430). - * - * Raised from 4MB to 32MB (#430): marathon autonomous runs (11-25 min observed) - * stream far more than 4MB of SSE frames, so a live disconnect mid-run would find - * an already-overflowed buffer and could only degrade-poll instead of re-attaching - * to the live tail. 32MB comfortably covers those runs while staying bounded. - * - * Memory cost: this is the WORST-CASE retained size PER ACTIVE run (the buffer is - * freed on finish + retention, or dropped immediately on overflow). With the small - * number of concurrent autonomous runs a single workspace realistically has, 32MB - * each is an acceptable ceiling; the overflow->204->degraded-poll fallback remains - * the backstop for anything larger, so correctness never depends on this bound. + * DEFAULT per-run replay ring cap (#491, down from 32MB). SSE frames carry + * UNcompacted tool outputs + framing overhead (×1.5–2 vs the persisted parts), so + * a "2–3 large reads + reasoning" step routinely blows past 2MB; 4MB comfortably + * holds a step or two of TAIL, which is all a resuming client needs (steps below + * its persisted frontier come from the seed, not the ring). The ring stays bounded + * because it rotates on every confirmed persist; this cap is only the ceiling for + * the un-persisted tail between rotations. Env-tunable via + * RUN_STREAM_MAX_BUFFER_BYTES (bytes); a 0/invalid value falls back to this. */ -export const RUN_STREAM_MAX_BUFFER_BYTES = 32 * 1024 * 1024; +export const RUN_STREAM_MAX_BUFFER_BYTES = 4 * 1024 * 1024; -// 2x the replay cap: a just-written full-replay burst alone can never trip the -// per-subscriber cap (see controller); only a genuinely stalled socket can. +// 2× the ring cap: a just-written full-tail burst alone can never trip the +// per-subscriber cap (see controller); only a genuinely stalled socket can. This +// derivative relationship is preserved even when the ring cap is env-overridden. export const SUBSCRIBER_MAX_BUFFERED_BYTES = 2 * RUN_STREAM_MAX_BUFFER_BYTES; +/** + * A finish-step boundary frame is exactly `data: {"type":"finish-step"...}\n\n` + * (verified empirically against ai@6.0.207 — each UI-message-stream part is a + * single `data: {json}\n\n` event, never split across `data:` lines, and `type` + * is always the first key). A prefix match is cheaper than JSON.parse-per-frame + * and has no false positives: a literal `"type":"finish-step"` inside a text + * delta is JSON-escaped (`\"type\":...`), and the frame would start with + * `data: {"type":"text-delta"` anyway. + */ +const FINISH_STEP_FRAME_PREFIX = 'data: {"type":"finish-step"'; + +/** Resolve the ring cap from the environment, falling back to the default. */ +function resolveMaxBufferBytes(): number { + const raw = process.env.RUN_STREAM_MAX_BUFFER_BYTES; + if (!raw) return RUN_STREAM_MAX_BUFFER_BYTES; + const parsed = Number(raw); + return Number.isFinite(parsed) && parsed > 0 + ? Math.floor(parsed) + : RUN_STREAM_MAX_BUFFER_BYTES; +} + export interface RunStreamCallbacks { onFrame: (frame: string) => void; onEnd: () => void; } export interface RunStreamAttachment { + // The synthetic `start` frame (carrying { runId, chatId }) followed by the + // buffered TAIL filtered to `stamp >= N`. The controller writes these to the + // socket in chunks respecting drain, then calls start(). replay: string[]; finished: boolean; start(): void; // drain pending frames (order preserved) and go live @@ -53,14 +136,19 @@ export interface RunStreamAttachment { 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. + // Byte size of `pending`, capped at the subscriber cap. `start()` is called in + // the SAME tick as `attach()` today, so `pending` never holds more than one + // microtask of frames — but the controller writes the (potentially large) tail + // respecting drain BEFORE start(), so a stalled socket can accumulate here; the + // cap is the structural backstop (an overflow degrades start() to an end()). pendingBytes: number; overflowed: boolean; pendingEnd: boolean; + // The client's step frontier N: this subscriber only receives frames with + // `stamp >= minStamp` (the tail past what it already persisted). Live frames + // always satisfy this (their stamp is the current, highest step), so it only + // filters the rare out-of-order below-frontier frame. + minStamp: number; } interface Entry { @@ -68,8 +156,20 @@ interface Entry { // 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; + // Parallel arrays: frames[i] is the SSE string, stamps[i] its step number. frames: string[]; + stamps: number[]; bytes: number; + // The running step counter used to stamp the NEXT frame (number of finish-step + // frames seen so far). + currentStamp: number; + // The highest confirmed `stepsPersisted`: frames with stamp < persistedFloor are + // on disk (safe to drop, never re-buffered). Monotonic (confirmPersistedStep). + persistedFloor: number; + // The highest stamp EVICTED by an overflow (unsafe) drop, -1 if none. Used to + // detect a partially-evicted leading step when computing the coverage floor. + overflowThroughStamp: number; + // Sticky-for-logging only: at least one unsafe (overflow) eviction happened. overflowed: boolean; finished: boolean; subscribers: Set; @@ -80,6 +180,10 @@ interface Entry { export class AiChatStreamRegistryService implements OnModuleDestroy { private readonly logger = new Logger(AiChatStreamRegistryService.name); private readonly entries = new Map(); // key: chatId + // Env-resolved caps (per instance) so a deployment can tune the ceiling without + // a code change. The subscriber cap keeps the documented 2× relationship. + readonly maxBufferBytes = resolveMaxBufferBytes(); + readonly subscriberMaxBufferedBytes = 2 * this.maxBufferBytes; /** * Register a fresh entry at the START of a run (before any frame), so a tab @@ -105,7 +209,11 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { this.entries.set(chatId, { runId, frames: [], + stamps: [], bytes: 0, + currentStamp: 0, + persistedFloor: 0, + overflowThroughStamp: -1, overflowed: false, finished: false, subscribers: new Set(), @@ -150,6 +258,34 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { void pump(); } + /** + * Confirm that step `stepsPersisted` (a COUNT: steps 0..stepsPersisted-1) is on + * disk for this run, and ROTATE the ring: drop the buffered frames of those + * now-persisted steps (stamp < stepsPersisted). This is the ONLY thing that + * rotates the ring, and it is called ONLY after a genuinely SUCCESSFUL per-step + * persist (see ai-chat.service updateStreaming). A failed persist never calls + * it, so the ring covers more (auto-safe). Identity-checked (invariant 1) and + * monotonic (a stale lower count is ignored). + */ + confirmPersistedStep( + chatId: string, + runId: string, + stepsPersisted: number, + ): void { + const entry = this.entries.get(chatId); + if (!entry || entry.runId !== runId) return; + if (!Number.isFinite(stepsPersisted) || stepsPersisted <= entry.persistedFloor) + return; + entry.persistedFloor = stepsPersisted; + // Clean rotation: drop the persisted steps from the head. These frames are on + // disk + carried by a fresh client seed, so this NEVER opens a gap. + while (entry.frames.length > 0 && entry.stamps[0] < stepsPersisted) { + entry.bytes -= Buffer.byteLength(entry.frames[0]); + entry.frames.shift(); + entry.stamps.shift(); + } + } + /** * 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 @@ -162,36 +298,62 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { } /** - * 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. + * Attach to a run's stream from the client's step frontier `n` (its persisted + * `stepsPersisted`). Async only for the phase-2 Redis seam — the body runs + * synchronously so the tail SLICE and the subscriber registration happen in ONE + * tick with no await between them (invariant 4). * * 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. + * - there is no entry; + * - the `anchor` does not match this run's assistant id (invariant 6); + * - the ring does not cover the client's frontier (coverageFloor > n): a hole + * from overflow, or the client's seed simply lagged behind a rotation. The + * client then refetches (a larger n) and re-attaches. + * + * Otherwise the attachment's `replay` is a synthetic `start` frame (the run-fact + * on re-attach) followed by the buffered tail filtered to `stamp >= n`. For a + * FINISHED run this is replay-only (no subscriber) and ends after the replay — + * with n = N_final that tail is just the run's `finish` frame, so the client + * closes the stream. For a LIVE run a paused subscriber is registered; the + * caller writes the replay (respecting drain) then calls start() to drain the + * pending frames and go live. */ async attach( chatId: string, - expectLive: boolean, anchor: string | undefined, + n: number, cb: RunStreamCallbacks, ): Promise { const entry = this.entries.get(chatId); - if (!entry || entry.overflowed) return null; + if (!entry) 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) { + if (anchor && entry.assistantMessageId !== anchor) return null; + // A finished entry with NOTHING in the ring (aborted before the first frame, + // or fully overflowed) has no tail to deliver -> 204 -> the client polls. + if (entry.finished && entry.frames.length === 0) return null; + const floor = this.coverageFloor(entry); + if (floor > n) { + this.logger.warn( + `run-stream attach gap for run=${entry.runId}: coverageFloor=${floor} ` + + `> client n=${n} -> 204 (client refetches + re-attaches)`, + ); + return null; + } + + const startFrame = this.buildStartFrame(chatId, entry.runId); + const sliceTail = (): string[] => { + const out: string[] = [startFrame]; + for (let i = 0; i < entry.frames.length; i++) { + if (entry.stamps[i] >= n) out.push(entry.frames[i]); + } + return out; + }; + + if (entry.finished) { // Replay-only: the run is done, no subscriber is registered. return { - replay: entry.frames.slice(), + replay: sliceTail(), finished: true, start: () => undefined, unsubscribe: () => undefined, @@ -206,15 +368,12 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { pendingBytes: 0, overflowed: false, pendingEnd: false, + minStamp: n, }; + // Register + snapshot in the SAME synchronous block (invariant 4). No await + // separates them, so a concurrently ingested frame cannot be lost/duplicated. 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). + const replay = sliceTail(); return { replay, finished: false, @@ -263,24 +422,77 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { this.entries.clear(); } - /** Buffer + fan-out a single frame. See invariant/overflow semantics inline. */ + /** The synthetic `start` frame the tail is prefixed with — the source of the + * run-fact (runId/chatId) on re-attach. A `start` frame does NOT reset the + * client's message parts (ai@6.0.207 createStreamingUIMessageState), so it is + * safe to prepend even when the sliced tail begins mid-message. */ + private buildStartFrame(chatId: string, runId: string): string { + return `data: ${JSON.stringify({ + type: 'start', + messageMetadata: { runId, chatId }, + })}\n\n`; + } + + /** + * The smallest step FULLY present in the ring: its smallest retained stamp, or + * (when the leading step was only partially evicted by an overflow) one past it. + * When the ring is empty it is the current step (only the live tail is coming). + * An attach at frontier `n` is covered ⟺ coverageFloor <= n. + */ + private coverageFloor(entry: Entry): number { + if (entry.frames.length === 0) return entry.currentStamp; + const min = entry.stamps[0]; + return entry.overflowThroughStamp >= min ? min + 1 : min; + } + + /** + * Buffer (step-stamped) + fan-out a single frame. The stamp is the number of + * finish-step frames seen BEFORE this one; a finish-step frame carries the + * current value and THEN increments the counter (so its stamp equals the 0-based + * index of the step it closes). Only frames at/above persistedFloor are buffered + * (already-persisted steps are on disk); the ring is then trimmed to the byte + * cap, an unsafe eviction opening a gap. Fan-out is always live (filtered per + * subscriber by its frontier). + */ private ingestFrame(entry: Entry, frame: string): void { - entry.bytes += Buffer.byteLength(frame); - if (!entry.overflowed) { + const size = Buffer.byteLength(frame); + const stamp = entry.currentStamp; + if (frame.startsWith(FINISH_STEP_FRAME_PREFIX)) { + entry.currentStamp = stamp + 1; + } + + // Buffer for replay only if this step is not already persisted+rotated away. + if (stamp >= entry.persistedFloor) { 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`, - ); + entry.stamps.push(stamp); + entry.bytes += size; + // Enforce the ring cap. Evicting a not-yet-persisted frame (stamp >= + // persistedFloor) opens a GAP; a leftover persisted frame (< floor) is a + // safe drop. Keep evicting until the ring is back under the cap. + while (entry.bytes > this.maxBufferBytes && entry.frames.length > 0) { + const evStamp = entry.stamps[0]; + entry.bytes -= Buffer.byteLength(entry.frames[0]); + entry.frames.shift(); + entry.stamps.shift(); + if (evStamp >= entry.persistedFloor) { + if (evStamp > entry.overflowThroughStamp) + entry.overflowThroughStamp = evStamp; + if (!entry.overflowed) { + entry.overflowed = true; + this.logger.warn( + `run-stream ring overflow for run=${entry.runId}: an un-persisted ` + + `step was evicted to stay under ${this.maxBufferBytes}B; a late ` + + `attach at an evicted step will 204 until a later persist confirms`, + ); + } + } } } + + // Fan out live, filtered to each subscriber's frontier (a subscriber only + // wants the tail past the step it already persisted). for (const sub of entry.subscribers) { + if (stamp < sub.minStamp) continue; if (sub.started) { try { sub.onFrame(frame); @@ -289,12 +501,12 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { } } else { sub.pending.push(frame); - sub.pendingBytes += Buffer.byteLength(frame); - if (sub.pendingBytes > SUBSCRIBER_MAX_BUFFERED_BYTES) { + sub.pendingBytes += size; + if (sub.pendingBytes > this.subscriberMaxBufferedBytes) { // 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. + // was delayed (the controller's drain-respecting tail write, or the + // phase-2 await seam). Drop it rather than buffer the whole run; on + // start() it degrades to an immediate end (a 204-equivalent). sub.overflowed = true; sub.pending = []; entry.subscribers.delete(sub); diff --git a/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts b/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts index 193b76b1..06df6d75 100644 --- a/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts @@ -2,18 +2,26 @@ import { AiChatStreamRegistryService, RUN_STREAM_MAX_BUFFER_BYTES, RUN_STREAM_RETAIN_FINISHED_MS, - SUBSCRIBER_MAX_BUFFERED_BYTES, 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. + * Unit tests for the in-memory run-stream registry (#184 phase 1.5, step-aligned + * retention #491). The registry is the whole of the resumable-transport contract: + * step-stamped retention, tail-only attach at the client's frontier N, the + * confirmed-persist ring rotation (and the anti-inversion rule), the memory bound, + * the overflow gap, paused -> live hand-off, retention, the anchor check + * (invariant 6), and the mirror-the-done-path replace semantics (invariant 3). */ +// Real ai@6 UI-message-stream SSE frames are `data: {json}\n\n`, one part each. +const sse = (part: Record): string => + `data: ${JSON.stringify(part)}\n\n`; +const finishStep = (): string => sse({ type: 'finish-step' }); +const textDelta = (id: string, delta: string): string => + sse({ type: 'text-delta', id, delta }); +const finish = (): string => sse({ type: 'finish' }); + // A ReadableStream whose frames the test pushes explicitly, plus close/error. function makePushStream(): { stream: ReadableStream; @@ -58,6 +66,9 @@ function collector(): { }; } +// The tail past the synthetic start frame (replay[0] is always the start frame). +const tail = (replay: string[]): string[] => replay.slice(1); + describe('AiChatStreamRegistryService', () => { const CHAT = 'chat-1'; let registry: AiChatStreamRegistryService; @@ -71,7 +82,21 @@ describe('AiChatStreamRegistryService', () => { registry.onModuleDestroy(); }); - it('replays frames in arrival order (live attach)', async () => { + it('prepends a synthetic start frame carrying { runId, chatId }', 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, 'assist-1', 0, c.cb))!; + const start = JSON.parse(att.replay[0].replace(/^data: /, '').trim()); + expect(start.type).toBe('start'); + expect(start.messageMetadata).toEqual({ runId: 'run-1', chatId: CHAT }); + }); + + it('replays the buffered tail (from frontier 0) in arrival order (live attach)', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); @@ -81,13 +106,13 @@ describe('AiChatStreamRegistryService', () => { await flush(); const c = collector(); - const att = await registry.attach(CHAT, false, undefined, c.cb); + const att = await registry.attach(CHAT, 'assist-1', 0, c.cb); expect(att).not.toBeNull(); - expect(att!.replay).toEqual(['a', 'b', 'c']); + expect(tail(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 () => { + it('late attach gets the buffered prefix as tail plus the live tail', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); @@ -96,17 +121,16 @@ describe('AiChatStreamRegistryService', () => { await flush(); const c = collector(); - const att = (await registry.attach(CHAT, false, undefined, c.cb))!; - expect(att.replay).toEqual(['a', 'b']); + const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!; + expect(tail(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 () => { + it('a paused subscriber receives frames buffered during pause in order, then live', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); @@ -114,81 +138,45 @@ describe('AiChatStreamRegistryService', () => { 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']); + const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!; + expect(tail(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 + att.start(); expect(c.frames).toEqual(['b', 'c']); - src.push('d'); // now live + src.push('d'); 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'); + registry.bind(CHAT, 'run-1', 'assist-1', makePushStream().stream); const c = collector(); - const att = (await registry.attach(CHAT, false, undefined, c.cb))!; - // Terminate the run while the subscriber is still paused. + const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!; 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 () => { + it('anchor mismatch 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(); + expect(await registry.attach(CHAT, 'assist-1', 0, 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(); + expect(await registry.attach(CHAT, 'other-id', 0, c.cb)).toBeNull(); }); - it('matching anchor with expect=live attaches', async () => { + it('matching anchor attaches', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); @@ -196,97 +184,60 @@ describe('AiChatStreamRegistryService', () => { await flush(); const c = collector(); - const att = await registry.attach(CHAT, true, 'assist-1', c.cb); + const att = await registry.attach(CHAT, 'assist-1', 0, c.cb); expect(att).not.toBeNull(); - expect(att!.replay).toEqual(['a']); + expect(tail(att!.replay)).toEqual(['a']); }); - it('overflow: attach returns null, but the LIVE subscriber keeps receiving (incl. the crossing frame)', async () => { + 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); - // A live (started) subscriber attached before the flood. + const bad = collector(); + const badAtt = (await registry.attach(CHAT, 'assist-1', 0, { + onFrame: () => { + throw new Error('boom'); + }, + onEnd: bad.cb.onEnd, + }))!; + badAtt.start(); + + const good = collector(); + const goodAtt = (await registry.attach(CHAT, 'assist-1', 0, 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); + expect(good.frames).toEqual(['a', 'b']); + }); + + it('open() over a LIVE entry ends started subscribers once; a late done never touches 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))!; + const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!; att.start(); - // Cap-relative so it survives a buffer-cap change (#430): a quarter-cap frame - // means 5 frames comfortably exceed the replay cap; the last one crosses. - const chunk = 'x'.repeat(Math.floor(RUN_STREAM_MAX_BUFFER_BYTES / 4)); - for (let i = 0; i < 5; i++) src.push(chunk + 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(chunk + 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(); - - // Cap-relative so it survives a buffer-cap change (#430): a quarter-of-the- - // per-subscriber-cap frame means 5 frames exceed A's paused-pending cap while - // B streams every frame live. - const chunk = 'x'.repeat(Math.floor(SUBSCRIBER_MAX_BUFFERED_BYTES / 4)); - for (let i = 0; i < 5; i++) src.push(chunk + 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(5); - - // 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 + expect(c.ended()).toBe(1); 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 + expect(c.ended()).toBe(1); const still = (registry as any).entries.get(CHAT); expect(still).toBe(newEntry); expect(still.runId).toBe('run-2'); @@ -299,7 +250,6 @@ describe('AiChatStreamRegistryService', () => { 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(); }); @@ -310,32 +260,253 @@ describe('AiChatStreamRegistryService', () => { 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 () => { +/** + * #491 step-stamped retention: the boundary detector, tail-only slicing at the + * client's frontier N, the confirmed-persist rotation (+ anti-inversion), the + * overflow gap, the memory bound, and the finished-retained tail. All observable + * against the REAL registry driven through open/bind/ingest. + */ +describe('AiChatStreamRegistryService step-aligned retention (#491)', () => { + const CHAT = 'chat-s'; + let registry: AiChatStreamRegistryService; + + beforeEach(() => { + registry = new AiChatStreamRegistryService(); + jest.spyOn((registry as any).logger, 'warn').mockImplementation(() => {}); + }); + afterEach(() => registry.onModuleDestroy()); + + const entryOf = () => (registry as any).entries.get(CHAT); + + it('stamps frames by finish-step count, aligned with stepsPersisted', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + // step 0 content, its finish-step, step 1 content, its finish-step, finish. + src.push(textDelta('t0', 'a')); // stamp 0 + src.push(finishStep()); // stamp 0 (the finish-step frame carries the pre value) + src.push(textDelta('t1', 'b')); // stamp 1 + src.push(finishStep()); // stamp 1 + src.push(finish()); // stamp 2 + await flush(); + const e = entryOf(); + expect(e.stamps).toEqual([0, 0, 1, 1, 2]); + expect(e.currentStamp).toBe(2); + }); - const bad = collector(); - const badAtt = (await registry.attach(CHAT, false, undefined, { - onFrame: () => { - throw new Error('boom'); - }, - onEnd: bad.cb.onEnd, - }))!; - badAtt.start(); + it('does NOT treat a text delta that merely quotes "finish-step" as a boundary', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + // A model that literally types "type":"finish-step" — JSON-escaped in the frame. + src.push(textDelta('t0', '"type":"finish-step"')); + await flush(); + expect(entryOf().currentStamp).toBe(0); // no false boundary + }); - 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 + it('tail-only: attach at N slices frames with stamp >= N', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(textDelta('t1', 'b')); // 1 + src.push(finishStep()); // 1 + src.push(textDelta('t2', 'c')); // 2 (in-progress) 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']); + const c = collector(); + // Client persisted 2 steps -> wants the tail from step 2. + const att = (await registry.attach(CHAT, 'assist-1', 2, c.cb))!; + expect(tail(att.replay)).toEqual([textDelta('t2', 'c')]); + }); + + it('attach in the MIDDLE of a step (N between finish-steps) slices from that step', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(textDelta('t1', 'b1')); // 1 + src.push(textDelta('t1', 'b2')); // 1 (still step 1, no finish-step yet) + await flush(); + + const c = collector(); + const att = (await registry.attach(CHAT, 'assist-1', 1, c.cb))!; + // Step 0's frames are dropped from the tail; the whole in-progress step 1 is kept. + expect(tail(att.replay)).toEqual([textDelta('t1', 'b1'), textDelta('t1', 'b2')]); + }); + + it('rotates the ring ONLY on a confirmed persist (drops stamp < N)', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(textDelta('t1', 'b')); // 1 + await flush(); + expect(entryOf().stamps).toEqual([0, 0, 1]); + + // Confirm step 0 persisted (stepsPersisted = 1) -> drop stamp < 1. + registry.confirmPersistedStep(CHAT, 'run-1', 1); + expect(entryOf().stamps).toEqual([1]); + expect(entryOf().persistedFloor).toBe(1); + }); + + it('persist FAILED but the ring still fits -> attach SUCCEEDS and the tail includes step N', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(textDelta('t1', 'b')); // 1 (step 1's persist FAILED -> no confirm) + await flush(); + // No confirmPersistedStep for step 1: the ring still holds step 1. + + const c = collector(); + // Client's last successful persist was step 0 -> stepsPersisted = 1. + const att = await registry.attach(CHAT, 'assist-1', 1, c.cb); + expect(att).not.toBeNull(); + expect(tail(att!.replay)).toEqual([textDelta('t1', 'b')]); // includes step 1 + }); + + it('persist failed AND the ring overflowed past N -> 204 (coverage gap)', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + // Step 0: a fat step that blows past the cap with NO persist confirmation. + const big = 'x'.repeat(Math.floor(RUN_STREAM_MAX_BUFFER_BYTES / 2)); + src.push(textDelta('t0', big)); // 0 + src.push(textDelta('t0', big)); // 0 + src.push(textDelta('t0', big)); // 0 -> overflow evicts stamp-0 frames + await flush(); + const e = entryOf(); + expect(e.overflowed).toBe(true); + expect(e.bytes).toBeLessThanOrEqual(registry.maxBufferBytes); + + // A client at frontier 0 falls at/below an evicted step -> gap -> null. + const c = collector(); + expect(await registry.attach(CHAT, 'assist-1', 0, c.cb)).toBeNull(); + }); + + it('stale N (client seed lagged behind a rotation) -> 204; after a refetch (larger N) -> success', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(textDelta('t1', 'b')); // 1 + src.push(finishStep()); // 1 + src.push(textDelta('t2', 'c')); // 2 + await flush(); + // Server confirmed steps 0 and 1 -> rotate away stamp < 2. + registry.confirmPersistedStep(CHAT, 'run-1', 2); + expect(entryOf().stamps).toEqual([2]); + + // A client whose seed still says stepsPersisted = 1 -> below minStamp -> 204. + const stale = collector(); + expect(await registry.attach(CHAT, 'assist-1', 1, stale.cb)).toBeNull(); + + // It refetches (now stepsPersisted = 2) and re-attaches -> success. + const fresh = collector(); + const att = await registry.attach(CHAT, 'assist-1', 2, fresh.cb); + expect(att).not.toBeNull(); + expect(tail(att!.replay)).toEqual([textDelta('t2', 'c')]); + }); + + it('overflow gap CLEARS once a later persist rotates out the holey steps', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + const big = 'x'.repeat(Math.floor(RUN_STREAM_MAX_BUFFER_BYTES / 2)); + src.push(textDelta('t0', big)); // 0 + src.push(textDelta('t0', big)); // 0 + src.push(finishStep()); // 0 (still stamp 0) + src.push(textDelta('t1', 'small')); // 1 + src.push(finishStep()); // 1 + src.push(textDelta('t2', 'c')); // 2 + await flush(); + expect(entryOf().overflowed).toBe(true); + + // Late persist confirms steps 0..1 -> rotates out the holey step-0 frames. + registry.confirmPersistedStep(CHAT, 'run-1', 2); + // A client at frontier 2 is now cleanly covered (the hole was below it). + const c = collector(); + const att = await registry.attach(CHAT, 'assist-1', 2, c.cb); + expect(att).not.toBeNull(); + expect(tail(att!.replay)).toEqual([textDelta('t2', 'c')]); + }); + + it('finished-retained + N = N_final -> empty tail plus the finish frame', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(finish()); // 1 (N_final = 1) + src.close(); + await flush(); + // The last step's per-step persist confirmed stepsPersisted = 1. + registry.confirmPersistedStep(CHAT, 'run-1', 1); + + const c = collector(); + const att = (await registry.attach(CHAT, 'assist-1', 1, c.cb))!; + expect(att.finished).toBe(true); + // Empty step tail; just the finish frame so the client's SDK closes the stream. + expect(tail(att.replay)).toEqual([finish()]); + // No subscriber registered for a finished run. + expect(entryOf().subscribers.size).toBe(0); + }); + + it('confirmPersistedStep is monotonic and identity-checked', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); + src.push(finishStep()); + src.push(textDelta('t1', 'b')); + await flush(); + registry.confirmPersistedStep(CHAT, 'run-1', 1); + expect(entryOf().persistedFloor).toBe(1); + // A stale lower count is ignored. + registry.confirmPersistedStep(CHAT, 'run-1', 0); + expect(entryOf().persistedFloor).toBe(1); + // A foreign runId is ignored. + registry.confirmPersistedStep(CHAT, 'WRONG', 5); + expect(entryOf().persistedFloor).toBe(1); + }); + + it('MEMORY BOUND: 5 parallel marathon runs each stream well past 32MB; each ring stays <= the cap', async () => { + const cap = registry.maxBufferBytes; + const chats = ['m0', 'm1', 'm2', 'm3', 'm4']; + const srcs = chats.map((chat) => { + registry.open(chat, `run-${chat}`); + const s = makePushStream(); + registry.bind(chat, `run-${chat}`, `assist-${chat}`, s.stream); + return s; + }); + // ~256KB frames; 160 per chat = 40MB streamed each, well past the old 32MB. + // Interleave a finish-step every 8 frames so steps advance realistically. No + // persist confirmation -> the ONLY thing keeping memory bounded is the cap. + const frame = 'y'.repeat(256 * 1024); + for (let batch = 0; batch < 20; batch++) { + for (let i = 0; i < 8; i++) { + for (const s of srcs) s.push(textDelta('t', frame)); + } + for (const s of srcs) s.push(finishStep()); + await flush(); // drain the pump so queues never hold a whole run + } + let total = 0; + for (const chat of chats) { + const e = (registry as any).entries.get(chat); + expect(e.bytes).toBeLessThanOrEqual(cap); + total += e.bytes; + } + // Total retained across all 5 runs is bounded by 5x the per-run cap — the old + // registry would have retained ~5x40MB = 200MB here. + expect(total).toBeLessThanOrEqual(cap * chats.length); }); }); @@ -361,7 +532,7 @@ describe('AiChatStreamRegistryService retention timers', () => { it('a finished entry is removed after the retention window', () => { registry.open(CHAT, 'run-1'); - registry.abortEntry(CHAT, 'run-1'); // finalize -> retention armed + registry.abortEntry(CHAT, 'run-1'); expect((registry as any).entries.get(CHAT)).toBeDefined(); jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1); expect((registry as any).entries.get(CHAT)).toBeUndefined(); @@ -369,20 +540,18 @@ describe('AiChatStreamRegistryService retention timers', () => { 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. + registry.abortEntry(CHAT, 'run-1'); 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 + registry.abortEntry(CHAT, 'run-1'); const clearSpy = jest.spyOn(global, 'clearTimeout'); - registry.open(CHAT, 'run-2'); // must clear run-1's retain timer + registry.open(CHAT, 'run-2'); expect(clearSpy).toHaveBeenCalled(); jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1); const entry = (registry as any).entries.get(CHAT); diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts b/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts index 67aa7cbe..05bc19b8 100644 --- a/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts @@ -8,10 +8,12 @@ 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 + * Wiring spec for the #184 phase 1.5 attach endpoint (tail-only #491) * (`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, + * registry is mocked so this exercises ONLY the controller's tail-write/live/204/ + * cleanup wiring against a fake raw socket. The attach signature is now + * `(chatId, anchor, n, cb)` — the client hands its persisted step frontier `n` + * and its assistant row id `anchor`. Constructor order is (aiChatService, * aiChatRunService, aiChatRepo, aiChatMessageRepo, aiTranscription, pageRepo, * streamRegistry, environment). */ @@ -86,8 +88,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { attach: jest.fn( ( _chatId: string, - _live: boolean, _anchor: string | undefined, + _n: number, cb: RunStreamCallbacks, ) => { capturedCb = cb; @@ -156,7 +158,7 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { expect(res.hijack).not.toHaveBeenCalled(); }); - it('threads expect=live and anchor through to the registry', async () => { + it('threads anchor and the numeric frontier n through to the registry', async () => { const { controller, streamRegistry } = makeController({ chat: owned, attachment: null, @@ -165,8 +167,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { const { req } = makeReq(); await controller.attachRunStream( 'c1', - 'live', 'anchor-1', + '2', req, res, user, @@ -174,13 +176,13 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { ); expect(streamRegistry.attach).toHaveBeenCalledWith( 'c1', - true, 'anchor-1', + 2, // parsed to a number expect.anything(), ); }); - it('passes expect=false when the query is absent', async () => { + it('floors n to 0 when the query is absent/invalid', async () => { const { controller, streamRegistry } = makeController({ chat: owned, attachment: null, @@ -198,8 +200,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { ); expect(streamRegistry.attach).toHaveBeenCalledWith( 'c1', - false, undefined, + 0, expect.anything(), ); }); @@ -245,8 +247,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { const { req } = makeReq(); await controller.attachRunStream( 'c1', - 'live', 'a1', + '1', req, res, user, diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.ts b/apps/server/src/core/ai-chat/ai-chat.controller.ts index 49267968..b97ad6ff 100644 --- a/apps/server/src/core/ai-chat/ai-chat.controller.ts +++ b/apps/server/src/core/ai-chat/ai-chat.controller.ts @@ -64,6 +64,47 @@ import { SUBSCRIBER_MAX_BUFFERED_BYTES, } from './ai-chat-stream-registry.service'; import { startSseHeartbeat } from './sse-resilience'; + +/** + * Write the attach TAIL to the hijacked socket in chunks that RESPECT drain + * (#491): each `write()` that returns false (the kernel buffer is full) is awaited + * on the next 'drain' before continuing. The old code wrote the whole buffer + * synchronously, which — with the pre-#491 32MB ring — spiked memory (half the + * OOM). Bails immediately if the socket ended/errored mid-write. Frames that the + * paused registry subscriber buffers while this awaits are delivered by start(). + */ +async function writeTailRespectingDrain( + raw: { + write(chunk: string): boolean; + writableEnded?: boolean; + destroyed?: boolean; + once(event: string, cb: () => void): unknown; + removeListener?(event: string, cb: () => void): unknown; + }, + frames: string[], +): Promise { + for (const frame of frames) { + if (raw.writableEnded || raw.destroyed) return; + const ok = raw.write(frame); + if (!ok) { + // Kernel buffer full — wait for drain (or an early close/error) before the + // next chunk, so a slow reader never forces the whole tail into memory. + // Remove ALL three listeners once any fires, so a many-chunk tail with + // repeated backpressure never leaks (MaxListenersExceededWarning). + await new Promise((resolve) => { + const finish = (): void => { + raw.removeListener?.('drain', finish); + raw.removeListener?.('close', finish); + raw.removeListener?.('error', finish); + resolve(); + }; + raw.once('drain', finish); + raw.once('close', finish); + raw.once('error', finish); + }); + } + } +} import { EnvironmentService } from '../../integrations/environment/environment.service'; /** @@ -290,19 +331,25 @@ export class AiChatController { } /** - * 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. + * Attach to a chat's live run stream from the client's step frontier (#184 phase + * 1.5, tail-only #491). A late/reloaded tab hands the server the step count it + * has PERSISTED (`n` = the seeded row's `metadata.stepsPersisted`) and its + * assistant row id (`anchor`); the registry answers with the TAIL past step `n` + * (a synthetic `start` frame + the buffered frames stamped >= n) and then the + * live tail. Owner-gated via assertOwnedChat (same gate as getRun). When there + * is nothing to resume — no entry, a ring that does not cover the client's + * frontier (overflow gap, or the client's seed lagged a rotation), or an anchor + * that pins a DIFFERENT run (invariant 6) — the endpoint answers 204, the ONLY + * "nothing to resume" signal the AI SDK's reconnect accepts (it maps 204 to a + * silent no-op); the client then refetches (a larger n) and re-attaches. 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. + * The step marker `n` comes ONLY from the client — the server never reads the + * row to derive it, because a server-side n from a stale seed would open a + * silent one-step hole. The tail is written to the socket in CHUNKS respecting + * drain (writeTailRespectingDrain): the old code synchronously blasted the whole + * buffer, which — with the old 32MB cap — was half the OOM. */ @SkipTransform() @UseGuards(JwtAuthGuard, UserThrottlerGuard) @@ -310,39 +357,41 @@ export class AiChatController { @Get('runs/:chatId/stream') async attachRunStream( @Param('chatId', new ParseUUIDPipe()) chatId: string, - @Query('expect') expect: string | undefined, @Query('anchor') anchor: string | undefined, + @Query('n') n: string | undefined, @Req() req: FastifyRequest, @Res() res: FastifyReply, @AuthUser() user: User, @AuthWorkspace() workspace: Workspace, ): Promise { await this.assertOwnedChat(chatId, user, workspace); // same gate as getRun + // The client's persisted step frontier. A missing/invalid value floors to 0 + // ("give me everything") which, past any rotation, safely 204s. + const frontier = Number.isFinite(Number(n)) ? Math.max(0, Number(n)) : 0; + // The per-subscriber backpressure cap tracks the (env-tunable) ring cap. + const subscriberCap = + this.streamRegistry?.subscriberMaxBufferedBytes ?? + SUBSCRIBER_MAX_BUFFERED_BYTES; 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(); + const attachment = await this.streamRegistry?.attach(chatId, anchor, frontier, { + onFrame: (frame) => { + // Backpressure guard: 2x the ring cap, so the initial tail burst alone + // can never trip it; only a genuinely stalled socket can. + try { + if (res.raw.writableLength > subscriberCap) { + res.raw.destroy(); // 'close' fires -> unsubscribe below + return; } - }, - onEnd: () => { - stopHeartbeat(); - if (!res.raw.writableEnded) res.raw.end(); - }, + 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; @@ -371,13 +420,16 @@ export class AiChatController { // deliberately NO Connection/Keep-Alive (hop-by-hop; Safari/HTTP2) }); res.raw.flushHeaders?.(); - for (const frame of attachment.replay) res.raw.write(frame); + // Write the tail in chunks respecting drain (not a synchronous blast, which + // was half the OOM). Frames the paused subscriber buffers meanwhile are + // drained by start() below; its cap is the backstop for a stalled socket. + await writeTailRespectingDrain(res.raw, attachment.replay); if (attachment.finished) { - res.raw.end(); + if (!res.raw.writableEnded) res.raw.end(); return; } stopHeartbeat = startSseHeartbeat(res.raw, 15_000); - attachment.start(); // drain pending accumulated during replay, go live + attachment.start(); // drain pending accumulated during the tail write, go live } catch { attachment.unsubscribe(); stopHeartbeat(); diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index deb5b127..e58788b0 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -1543,30 +1543,39 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // Per-step (non-terminal) update: persist the finished steps the moment a // step ends. Tolerant — a failed update is logged and swallowed so it never // throws into the stream. Keeps status 'streaming'. - const updateStreaming = async (): Promise => { - if (!assistantId) return; + // + // #491: it now SIGNALS its outcome — the persisted `stepsPersisted` count on + // a CONFIRMED write, or null when it was skipped/failed. The caller rotates + // the run-stream registry ring ONLY on a non-null return (a confirmed + // persist), so a failed persist never rotates away a step nobody has (the + // classic inversion bug); a failure just makes the ring cover more. + const updateStreaming = async (): Promise => { + if (!assistantId) return null; // Cheap short-circuit once the turn is finalized (see `finalized` below). // The AUTHORITATIVE guard is `onlyIfStreaming` on the UPDATE: a late // fire-and-forget step update could still be in flight on another pool // connection when finalize runs, so the SQL `WHERE status='streaming'` // (not this flag) is what prevents it clobbering the terminal row. - if (finalized) return; + if (finalized) return null; + // Build the flush ONCE so the returned count is EXACTLY the persisted + // `stepsPersisted` (both derive from capturedSteps.length at this instant). + const flushed = flushAssistant(capturedSteps, '', 'streaming', { + pageChanged, + partsCache, + }); + const stepsPersisted = flushed.metadata.stepsPersisted as number; try { - await this.aiChatMessageRepo.update( - assistantId, - workspace.id, - flushAssistant(capturedSteps, '', 'streaming', { - pageChanged, - partsCache, - }), - { onlyIfStreaming: true }, - ); + await this.aiChatMessageRepo.update(assistantId, workspace.id, flushed, { + onlyIfStreaming: true, + }); + return stepsPersisted; } catch (err) { this.logger.warn( `Failed to update streaming assistant row: ${ err instanceof Error ? err.message : 'unknown error' }`, ); + return null; } }; @@ -1742,7 +1751,24 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // this point still recovers the step. Not awaited here (never block the // stream), but SERIALIZED via stepUpdateChain so the writes commit in // step order; updateStreaming is error-tolerant (logs + swallows). - stepUpdateChain = stepUpdateChain.then(() => updateStreaming()); + // #491: on a CONFIRMED persist, rotate the run-stream registry ring to + // drop the now-on-disk steps (stamp < stepsPersisted). Gated on the + // resumable flag (same as open/bind) and identity-checked in the + // registry; a null return (skipped/failed) rotates NOTHING (auto-safe). + stepUpdateChain = stepUpdateChain.then(async () => { + const persisted = await updateStreaming(); + if ( + persisted != null && + runId && + this.environment?.isAiChatResumableStreamEnabled?.() + ) { + this.streamRegistry?.confirmPersistedStep( + chatId, + runId, + persisted, + ); + } + }); // #184: persist the run's progress (finished-step count). Fire-and- // forget; the hook swallows its own errors. if (runId) runHooks?.onStep?.(runId, capturedSteps.length); diff --git a/apps/server/test/integration/ai-chat-attach.int-spec.ts b/apps/server/test/integration/ai-chat-attach.int-spec.ts index 70da3cbd..6e57199e 100644 --- a/apps/server/test/integration/ai-chat-attach.int-spec.ts +++ b/apps/server/test/integration/ai-chat-attach.int-spec.ts @@ -30,11 +30,13 @@ import { * 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. + * Proven here (tail-only #491): a finished run attached at its persisted frontier + * N_final delivers only the TAIL past N (a synthetic `start` carrying the run-fact + * + the terminal `finish`/`[DONE]`) — the step content below N lives in the seeded + * DB row, NOT the ring; 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)); @@ -133,14 +135,16 @@ function liveSink(): { }; } -// The SSE `start` frame carries the message id; pull it out of a `data: {...}`. -function parseStartMessageId(frames: string[]): string | undefined { +// Parse the first `start` frame's JSON out of a `data: {...}` sequence. +function parseStartFrame( + frames: string[], +): { messageId?: string; messageMetadata?: any } | 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; + if (json.type === 'start') return json; } catch { /* not this frame */ } @@ -270,7 +274,7 @@ describe('AiChatService run-stream attach [integration]', () => { await destroyTestDb(); }); - it('run-wrapped: replay is the full frame sequence incl [DONE], start.messageId == the seeded DB row id', async () => { + it('run-wrapped, tail-only: a finished run at N_final delivers the run-fact start + finish/[DONE]; the step content lives in the seeded row', async () => { const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; const registry = new AiChatStreamRegistryService(); const runService = new AiChatRunService(runRepo, { @@ -300,27 +304,37 @@ describe('AiChatService run-stream attach [integration]', () => { ); }); const rowId = await assistantRowId(chatId); + // The client reads its persisted step frontier N from the seeded row. + const row: any = await msgRepo.findById(rowId, workspaceId); + const nFinal = row.metadata.stepsPersisted as number; + expect(nFinal).toBe(1); // a single finished step + // The step content is in the SEEDED row (parts/content), not the ring. + expect(JSON.stringify(row.metadata.parts)).toContain('Hello'); - // Finished-run replay with expect=live + the correct anchor. + // Attach at N_final with the correct anchor: the tail past step 1 is just + // the terminal frames; step 0's 'Hello' is BELOW the frontier (seeded). const sink = liveSink(); - const att = await registry.attach(chatId, true, rowId, sink.cb); + const att = await registry.attach(chatId, rowId, nFinal, 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'); + // The synthetic start frame carries the run-fact (runId/chatId), the source + // of the run-fact on re-attach. + const start = parseStartFrame(att!.replay); + expect(start?.messageMetadata).toMatchObject({ + runId: box.runId, + chatId, + }); + // The terminal marker is delivered so the client's SDK closes the stream. expect(att!.replay.some((f) => f.includes('[DONE]'))).toBe(true); + // 'Hello' (step 0, below the frontier) is NOT re-streamed — it is seeded. + expect(att!.replay.some((f) => f.includes('Hello'))).toBe(false); } finally { registry.onModuleDestroy(); await cleanup(); } }); - it('anchor mismatch with expect=live returns null (invariant 6)', async () => { + it('anchor mismatch returns null (invariant 6)', async () => { const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; const registry = new AiChatStreamRegistryService(); const runService = new AiChatRunService(runRepo, { @@ -347,7 +361,7 @@ describe('AiChatService run-stream attach [integration]', () => { 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), + await registry.attach(chatId, 'a-different-run-row', 1, sink.cb), ).toBeNull(); } finally { registry.onModuleDestroy(); @@ -376,8 +390,11 @@ describe('AiChatService run-stream attach [integration]', () => { 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 + const att = (await registry.attach(chatId, undefined, 0, sink.cb))!; + // Nothing streamed yet -> the tail is just the synthetic start frame; the + // whole live stream (start..DONE) follows via onFrame after start(). + expect(att.replay).toHaveLength(1); + expect(att.replay[0]).toContain('"type":"start"'); att.start(); // go live (drains nothing, then follows) // Now emit the whole turn. @@ -448,7 +465,7 @@ describe('AiChatService run-stream attach [integration]', () => { }); try { const sink = liveSink(); - const att = (await registry.attach(chatId, false, undefined, sink.cb))!; + const att = (await registry.attach(chatId, undefined, 0, sink.cb))!; att.start(); // Give streamText a beat to begin consuming the partial output. @@ -526,7 +543,9 @@ describe('AiChatService run-stream attach [integration]', () => { expect(entry).toBeDefined(); expect(entry.finished).toBe(true); const sink = liveSink(); - expect(await registry.attach(chatId, false, undefined, sink.cb)).toBeNull(); + // Finished with an EMPTY ring (aborted before any frame) -> null -> the + // client degrades to poll instead of hanging on an empty stream. + expect(await registry.attach(chatId, undefined, 0, sink.cb)).toBeNull(); } finally { registry.onModuleDestroy(); await cleanup(); @@ -556,8 +575,8 @@ describe('AiChatService run-stream attach [integration]', () => { }); 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(); + expect(await registry.attach(chatId, undefined, 0, sink.cb)).toBeNull(); + expect(await registry.attach(chatId, 'anything', 1, sink.cb)).toBeNull(); } finally { registry.onModuleDestroy(); await cleanup();