vvzvlad/gitmost
1
0
Fork 0
Code Issues 15 Pull Requests 5 Packages Projects Releases Activity

refactor(ai-chat): registry — step-aligned retention + tail-only attach (#491)

Browse Source 
Реестр ран-стримов больше не буферизует до 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) <noreply@anthropic.com>
This commit is contained in:
agent_coder
2026-07-11 18:34:33 +03:00
parent a919c79cb2
commit daeeb1f3f2
8 changed files with 881 additions and 312 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts
+280 -68
View File
@@ -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<Subscriber>;
@@ -80,6 +180,10 @@ interface Entry {
export class AiChatStreamRegistryService implements OnModuleDestroy {
private readonly logger = new Logger(AiChatStreamRegistryService.name);
private readonly entries = new Map<string, Entry>(); // 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<Subscriber>(),
@@ -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<RunStreamAttachment | null> {
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);
apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts
+326 -157
View File
@@ -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, unknown>): 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<string>;
@@ -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);
apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts
+12 -10
View File
@@ -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,
apps/server/src/core/ai-chat/ai-chat.controller.ts
+90 -38
View File
@@ -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<void> {
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<void>((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<void> {
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();
apps/server/src/core/ai-chat/ai-chat.service.ts
+39 -13
View File
@@ -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<void> => {
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<number | null> => {
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);