vvzvlad/gitmost
1
0
Fork 0
Code Issues 27 Pull Requests 6 Packages Projects Releases Activity

Merge remote-tracking branch 'gitea/develop' into feat/501-mcp-apikey-auth

Browse Source
This commit is contained in:
agent_coder
2026-07-12 04:33:17 +03:00
parent d42ca8dc57 ab133fa0d0
commit 55e8f61b3c
172 changed files with 10805 additions and 1560 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
+302 -69
View File
@@ -1,42 +1,122 @@
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
* AI_CHAT_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 AI_CHAT_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.
export const SUBSCRIBER_MAX_BUFFERED_BYTES = 2 * RUN_STREAM_MAX_BUFFER_BYTES;
// 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 * AI_CHAT_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.AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES;
if (!raw) return AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES;
const parsed = Number(raw);
return Number.isFinite(parsed) && parsed > 0
? Math.floor(parsed)
: AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES;
}
export interface RunStreamCallbacks {
onFrame: (frame: string) => void;
@@ -44,6 +124,9 @@ export interface RunStreamCallbacks {
}
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,77 @@ 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,
// The client's persisted step frontier. `null` = a NOT-tail-aware client (no
// `n` query param) — a legacy/parameterless tab that expects the old
// "finished -> 204 -> poll" contract; distinct from `0` (a tail-aware client
// with nothing persisted yet).
n: number | null,
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;
// #491 regression guard (#137/#161 dup): a NOT-tail-aware client (no `n`)
// resuming a FINISHED run must 204 and poll — the old `finished && !expectLive`
// gate. Without this, a missing `n` collapsing to frontier 0 would serve the
// WHOLE tail of a finished, NON-rotated run (coverageFloor 0), and a
// parameterless client that never stripped its transcript would APPEND that
// full replay onto the steps it already shows -> duplicated text. A tail-aware
// client (n present, incl. n=0) still gets the tail past its frontier.
if (entry.finished && n === null) 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;
// A LIVE run with no `n` (legacy parameterless) replays from step 0 (the old
// behavior); a tail-aware client resumes from its frontier.
const frontier = n ?? 0;
const floor = this.coverageFloor(entry);
if (floor > frontier) {
this.logger.warn(
`run-stream attach gap for run=${entry.runId}: coverageFloor=${floor} ` +
`> client frontier=${frontier} -> 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] >= frontier) 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 +383,12 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
pendingBytes: 0,
overflowed: false,
pendingEnd: false,
minStamp: frontier,
};
// 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 +437,83 @@ 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 {
// Empty ring: only the live tail is coming. The floor is the current step,
// but never below persistedFloor — a confirmed persist can rotate the ring
// empty while currentStamp still lags a beat behind on another connection, so
// max() keeps the invariant STRUCTURAL (a client with n = persistedFloor is
// always covered) rather than timing-dependent.
if (entry.frames.length === 0)
return Math.max(entry.currentStamp, entry.persistedFloor);
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 +522,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
+350 -158
View File
@@ -1,19 +1,27 @@
import {
AiChatStreamRegistryService,
RUN_STREAM_MAX_BUFFER_BYTES,
AI_CHAT_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,276 @@ 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(AI_CHAT_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(AI_CHAT_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('#491 regression (#137/#161 dup): a PARAMETERLESS attach (n=null) to a finished NON-rotated run -> 204, but n=0 still gets the tail', async () => {
// A finished, non-rotated run: frames present, coverageFloor 0. A missing `n`
// (null — a legacy/parameterless tab that never stripped its transcript) must
// 204 -> poll, NOT receive the whole tail it would append (duplicate). A
// tail-aware client (n=0 present) still resumes.
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
src.close();
await flush();
// NOT rotated (no confirmPersistedStep) -> stamps[0]=0, coverageFloor=0.
// MUTATION-VERIFY: revert the `finished && n === null -> null` gate (default n
// to 0) and the parameterless attach below serves the full tail instead of 204.
expect(await registry.attach(CHAT, 'assist-1', null, collector().cb)).toBeNull();
// A tail-aware client at frontier 0 IS served (the distinction: null != 0).
const tailAware = await registry.attach(CHAT, 'assist-1', 0, collector().cb);
expect(tailAware).not.toBeNull();
expect(tailAware!.finished).toBe(true);
});
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 +555,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 +563,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
+44 -11
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,44 @@ 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('#491: an ABSENT/invalid n passes null (not 0) so a finished run 204s (not-tail-aware)', async () => {
// Distinguishing a MISSING `n` from `n=0` is the #137/#161 dup guard: a
// parameterless/legacy tab must be handed null (-> the registry 204s a finished
// run) rather than frontier 0 (which would serve a finished non-rotated run's
// whole tail). MUTATION-VERIFY: revert to `Number(n) || 0` and this asserts 0.
const { controller, streamRegistry } = makeController({
chat: owned,
attachment: null,
});
for (const bad of [undefined, '', 'abc']) {
streamRegistry.attach.mockClear();
const { res } = makeRawRes();
const { req } = makeReq();
await controller.attachRunStream(
'c1',
undefined,
bad,
req,
res,
user,
workspace,
);
expect(streamRegistry.attach).toHaveBeenCalledWith(
'c1',
undefined,
null,
expect.anything(),
);
}
});
it('#491: a PRESENT n=0 passes 0 (tail-aware, distinct from absent)', async () => {
const { controller, streamRegistry } = makeController({
chat: owned,
attachment: null,
@@ -190,7 +223,7 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
await controller.attachRunStream(
'c1',
undefined,
undefined,
'0',
req,
res,
user,
@@ -198,8 +231,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
);
expect(streamRegistry.attach).toHaveBeenCalledWith(
'c1',
false,
undefined,
0,
expect.anything(),
);
});
@@ -245,8 +278,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.delta.spec.ts
+108
View File
@@ -0,0 +1,108 @@
import { ForbiddenException } from '@nestjs/common';
import { AiChatController } from './ai-chat.controller';
import type { User, Workspace } from '@docmost/db/types/entity.types';
/**
* Wiring spec for the #491 delta-poll endpoint (`POST /ai-chat/messages/delta`).
* Owner-gated via assertOwnedChat (same gate as the other reads), NOT flag-gated.
* The run fact rides IN the delta response (no separate /run poll). Hand-rolled
* mocks — no Nest graph, no DB. Constructor order: (aiChatService,
* aiChatRunService, aiChatRepo, aiChatMessageRepo, aiTranscription, pageRepo).
*/
describe('AiChatController POST /ai-chat/messages/delta (#491)', () => {
const user = { id: 'u1' } as User;
const workspace = { id: 'ws1' } as Workspace;
function makeController(opts: {
chat?: unknown;
delta?: { rows: unknown[]; cursor: string };
run?: unknown;
}) {
const aiChatRunService = {
getLatestForChat: jest.fn().mockResolvedValue(opts.run),
};
const aiChatRepo = {
findById: jest.fn().mockResolvedValue(opts.chat),
};
const aiChatMessageRepo = {
findByChatUpdatedAfter: jest
.fn()
.mockResolvedValue(opts.delta ?? { rows: [], cursor: 'C1' }),
};
const controller = new AiChatController(
{} as never,
aiChatRunService as never,
aiChatRepo as never,
aiChatMessageRepo as never,
{} as never,
{} as never,
);
return { controller, aiChatRunService, aiChatRepo, aiChatMessageRepo };
}
it('owner-gates: a chat the user does not own throws, never reaching the repo', async () => {
const { controller, aiChatMessageRepo, aiChatRunService } = makeController({
chat: { id: 'c1', creatorId: 'someone-else' },
});
await expect(
controller.getMessagesDelta({ chatId: 'c1' }, user, workspace),
).rejects.toBeInstanceOf(ForbiddenException);
expect(aiChatMessageRepo.findByChatUpdatedAfter).not.toHaveBeenCalled();
expect(aiChatRunService.getLatestForChat).not.toHaveBeenCalled();
});
it('returns { rows, cursor, run:{id,status} } with the run fact inlined', async () => {
const rows = [{ id: 'm1' }];
const { controller } = makeController({
chat: { id: 'c1', creatorId: 'u1' },
delta: { rows, cursor: 'C2' },
run: { id: 'r1', status: 'running', error: 'ignored', stepCount: 3 },
});
const res = await controller.getMessagesDelta(
{ chatId: 'c1', cursor: 'C1' },
user,
workspace,
);
expect(res).toEqual({
rows,
cursor: 'C2',
// ONLY id + status — never the whole run row.
run: { id: 'r1', status: 'running' },
});
});
it('run is null when the chat has never had a run', async () => {
const { controller } = makeController({
chat: { id: 'c1', creatorId: 'u1' },
run: undefined,
});
const res = await controller.getMessagesDelta(
{ chatId: 'c1' },
user,
workspace,
);
expect(res.run).toBeNull();
});
it('passes cursor through, defaulting a missing cursor to null (first poll)', async () => {
const { controller, aiChatMessageRepo } = makeController({
chat: { id: 'c1', creatorId: 'u1' },
});
await controller.getMessagesDelta({ chatId: 'c1' }, user, workspace);
expect(aiChatMessageRepo.findByChatUpdatedAfter).toHaveBeenCalledWith(
'c1',
'ws1',
null,
);
await controller.getMessagesDelta(
{ chatId: 'c1', cursor: 'CX' },
user,
workspace,
);
expect(aiChatMessageRepo.findByChatUpdatedAfter).toHaveBeenLastCalledWith(
'c1',
'ws1',
'CX',
);
});
});
apps/server/src/core/ai-chat/ai-chat.controller.ts
+139 -38
View File
@@ -51,6 +51,7 @@ import {
ChatIdDto,
ExportChatDto,
GeneratePageTitleDto,
GetChatDeltaDto,
GetChatMessagesDto,
GetRunDto,
RenameChatDto,
@@ -63,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';
/**
@@ -149,6 +191,46 @@ export class AiChatController {
);
}
/**
* Delta poll (#491) — the degraded-poll fallback's payload. Returns the chat's
* message rows changed since `cursor` (a DB-clock timestamp from the previous
* poll), a FRESH cursor, AND the current run fact `{ id, status } | null`. This
* replaces the old degraded poll that refetched ALL infinite-query pages (full
* parts) every 2.5s: the client seeds once and thereafter merges only the
* deltas by id (the overlap window guarantees repeats — the merge is idempotent,
* see mergeById). The run fact rides IN the delta (a separate /run poll would
* double the poll QPS), so the client FSM gets the run's status on the same tick.
* Owner-gated via assertOwnedChat (same gate as the other read endpoints).
*/
@HttpCode(HttpStatus.OK)
@Post('messages/delta')
async getMessagesDelta(
@Body() dto: GetChatDeltaDto,
@AuthUser() user: User,
@AuthWorkspace() workspace: Workspace,
): Promise<{
rows: AiChatMessage[];
cursor: string;
run: { id: string; status: string } | null;
}> {
await this.assertOwnedChat(dto.chatId, user, workspace);
const { rows, cursor } =
await this.aiChatMessageRepo.findByChatUpdatedAfter(
dto.chatId,
workspace.id,
dto.cursor ?? null,
);
const run = await this.aiChatRunService.getLatestForChat(
dto.chatId,
workspace.id,
);
return {
rows,
cursor,
run: run ? { id: run.id, status: run.status } : null,
};
}
/**
* Export a chat to Markdown (#183). The DB is the single source of truth: the
* whole transcript is loaded (oldest -> newest) and rendered server-side. Now
@@ -249,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)
@@ -269,39 +357,49 @@ 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. #491: distinguish a MISSING/invalid `n`
// (null — a NOT-tail-aware, legacy/parameterless tab expecting the old
// "finished -> 204 -> poll" contract) from `n=0` (a tail-aware client with
// nothing persisted yet). Passing 0 for a missing `n` would serve a finished,
// non-rotated run's WHOLE tail and a parameterless client would append it onto
// the steps it already shows -> #137/#161 duplicate. null makes the registry
// 204 such a finished run (see attach); a tail-aware n=0 still resumes.
const frontier: number | null =
n === undefined || n === '' || !Number.isFinite(Number(n))
? null
: Math.max(0, Number(n));
// 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;
@@ -330,13 +428,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.run-race.spec.ts
+48 -3
View File
@@ -181,7 +181,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
{} as never, // pageAccess
{ isAiChatDeferredToolsEnabled: () => false, isAiChatFinalStepLockdownEnabled: () => false } as never, // environment
);
return { svc };
return { svc, aiChatMessageRepo };
}
const body = {
@@ -287,7 +287,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
// Drive stream() to the point streamText is called, capturing the options object
// (which carries onStepFinish/onFinish/onError/onAbort) and the run hooks.
async function captureStreamCallbacks() {
const { svc } = makeService();
const { svc, aiChatMessageRepo } = makeService();
let capturedOpts: any;
streamTextMock.mockImplementation((opts: any) => {
capturedOpts = opts;
@@ -314,7 +314,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
runHooks: runHooks as never,
});
expect(capturedOpts).toBeDefined();
return { capturedOpts, runHooks };
return { capturedOpts, runHooks, aiChatMessageRepo };
}
it('F9: onStepFinish bumps the run step count, onFinish settles the run "completed" (the dominant autonomous-run path)', async () => {
@@ -369,6 +369,51 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
expect.stringContaining('provider exploded'),
);
});
// #490 reactive branch: a provider CONTEXT-OVERFLOW 400 in onError is classified,
// records a distinguishable cause, and stamps metadata.replayOverflow so the NEXT
// turn's budgeter trims aggressively (the recovery that un-bricks the chat).
it('#490: a context-overflow 400 stamps replayOverflow on the finalized row', async () => {
jest
.spyOn(Logger.prototype, 'error')
.mockImplementation(() => undefined as never);
jest
.spyOn(Logger.prototype, 'warn')
.mockImplementation(() => undefined as never);
const { capturedOpts, aiChatMessageRepo } = await captureStreamCallbacks();
const overflow = Object.assign(new Error('too large'), {
statusCode: 400,
message:
"This model's maximum context length is 128000 tokens. However, your messages resulted in 214000 tokens. Please reduce the length.",
});
await capturedOpts.onError({ error: overflow });
// The seed row exists (finalizeOwner is the owner-write path).
expect(aiChatMessageRepo.finalizeOwner).toHaveBeenCalled();
const calls = aiChatMessageRepo.finalizeOwner.mock.calls as any[][];
const patch = calls[calls.length - 1][2] as {
status: string;
metadata: Record<string, unknown>;
};
expect(patch.status).toBe('error');
expect(patch.metadata.replayOverflow).toBe(true);
expect(patch.metadata.error).toContain('контекстное окно');
});
it('#490: a non-overflow error does NOT stamp replayOverflow', async () => {
jest
.spyOn(Logger.prototype, 'error')
.mockImplementation(() => undefined as never);
const { capturedOpts, aiChatMessageRepo } = await captureStreamCallbacks();
await capturedOpts.onError({ error: new Error('network reset') });
const calls = aiChatMessageRepo.finalizeOwner.mock.calls as any[][];
const patch = calls[calls.length - 1][2] as {
status: string;
metadata: Record<string, unknown>;
};
expect('replayOverflow' in patch.metadata).toBe(false);
});
});
/**
apps/server/src/core/ai-chat/ai-chat.service.spec.ts
+342 -13
View File
@@ -13,6 +13,7 @@ import {
compactToolOutput,
assistantParts,
serializeSteps,
type StepPartsCache,
rowToUiMessage,
prepareAgentStep,
stepBudgetWarning,
@@ -28,10 +29,14 @@ import {
FINAL_STEP_NUDGE,
STEP_LIMIT_NO_ANSWER_MARKER,
OUTPUT_DEGENERATION_ERROR,
lastAssistantContextTokens,
lastAssistantReplayOverflow,
seedActivatedTools,
} from './ai-chat.service';
import type { AiChatMessage, Workspace } from '@docmost/db/types/entity.types';
import { buildSystemPrompt } from './ai-chat.prompt';
import type { McpClientsService } from './external-mcp/mcp-clients.service';
import { resolveEffectiveReplayThreshold } from './history-budget';
/**
* Unit tests for compactToolOutput: the pure helper that shrinks tool outputs
@@ -114,6 +119,54 @@ describe('compactToolOutput', () => {
describe('assistantParts', () => {
type AnyPart = Record<string, unknown>;
// #490 memoization: assistantParts builds each step's parts once and caches
// them by the step OBJECT's identity, so a mid-stream flush does not
// re-stringify every prior step's (large) output. Observable property: with a
// shared cache, the second call over the SAME step object returns the cached
// (identical) part array even if the step's underlying output was swapped —
// proving the work was memoized, not redone.
it('memoizes a step by identity (shared cache => one build per step)', () => {
const cache: StepPartsCache = new WeakMap();
const step = {
text: 'x',
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
toolResults: [{ toolCallId: 'c1', toolName: 'getPage', output: { v: 1 } }],
};
const first = assistantParts([step], '', cache) as AnyPart[];
expect((first.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
1,
);
// Swap the output for a NEW value; a re-build would pick it up, a cache hit
// keeps the first result.
step.toolResults[0] = {
toolCallId: 'c1',
toolName: 'getPage',
output: { v: 2 },
};
const second = assistantParts([step], '', cache) as AnyPart[];
expect((second.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
1,
);
// Same cached part objects are reused.
expect(second.find((p) => p.type === 'tool-getPage')).toBe(
first.find((p) => p.type === 'tool-getPage'),
);
});
it('without a cache, each call rebuilds (no stale memo)', () => {
const step = {
text: 'x',
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
toolResults: [{ toolCallId: 'c1', toolName: 'getPage', output: { v: 1 } }],
};
const first = assistantParts([step], '') as AnyPart[];
step.toolResults[0].output = { v: 2 };
const second = assistantParts([step], '') as AnyPart[];
expect((second.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
2,
);
});
it('emits output-available for a tool-call WITH a paired result', () => {
const steps = [
{
@@ -231,61 +284,320 @@ describe('assistantParts', () => {
});
});
describe('serializeSteps', () => {
// #490 trace format v2: per call the trace stores { input } for the call and an
// OUTCOME element — { ok: true } on success, { error, kind: 'thrown' } on a
// thrown tool-error, { error, kind: 'interrupted' } on a mid-step abort. The tool
// OUTPUT is no longer duplicated here (it lives once in metadata.parts).
describe('serializeSteps (trace v2)', () => {
it('returns null when there are no calls or results', () => {
expect(serializeSteps([])).toBeNull();
});
it('flattens calls and results into a compact trace', () => {
it('pairs a successful call with an { ok: true } outcome and NO output', () => {
const trace = serializeSteps([
{
toolCalls: [{ toolName: 'getPage', input: { id: 'p1' } }],
toolResults: [{ toolName: 'getPage', output: { title: 'T' } }],
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: { id: 'p1' } }],
toolResults: [{ toolCallId: 'c1', toolName: 'getPage' }],
},
]) as Array<Record<string, unknown>>;
expect(trace).toHaveLength(2);
expect(trace[0]).toEqual({ toolName: 'getPage', input: { id: 'p1' } });
expect(trace[1]).toEqual({ toolName: 'getPage', output: { title: 'T' } });
expect(trace[1]).toEqual({ toolName: 'getPage', ok: true });
// The output is NOT stored in the trace any more (dedup: it lives in parts).
expect(trace.some((e) => 'output' in e)).toBe(false);
});
it('records a THROWN tool failure (tool-error part) with its error message', () => {
it('records a THROWN failure with { error, kind: "thrown" }', () => {
const trace = serializeSteps([
{
toolCalls: [{ toolName: 'editPageText', input: { id: 'p1' } }],
toolCalls: [
{ toolCallId: 'c1', toolName: 'editPageText', input: { id: 'p1' } },
],
toolResults: [],
content: [
{
type: 'tool-error',
toolCallId: 'c1',
toolName: 'editPageText',
error: new Error('page is locked'),
},
],
},
]) as Array<Record<string, unknown>>;
// The call element is followed by a paired error element (mirroring how a
// successful result is appended), so the failure survives in the trace.
expect(trace).toHaveLength(2);
expect(trace[0]).toEqual({ toolName: 'editPageText', input: { id: 'p1' } });
expect(trace[1]).toEqual({
toolName: 'editPageText',
error: 'page is locked',
kind: 'thrown',
});
});
it('truncates a very long tool-error message to the tool-output limit', () => {
it('marks an interrupted call (no result, no throw) with kind "interrupted"', () => {
const trace = serializeSteps([
{
toolCalls: [
{ toolCallId: 'c1', toolName: 'createComment', input: { x: 1 } },
],
toolResults: [],
content: [],
},
]) as Array<Record<string, unknown>>;
expect(trace).toHaveLength(2);
expect(trace[1]).toEqual({
toolName: 'createComment',
error: 'Tool call did not complete.',
kind: 'interrupted',
});
// Structurally distinct from a thrown hard-fail so it never inflates an
// error-rate scan.
expect((trace[1] as { kind: string }).kind).not.toBe('thrown');
});
it('truncates a very long thrown-error message to the tool-output limit', () => {
const long = 'x'.repeat(5000);
const trace = serializeSteps([
{
toolCalls: [{ toolName: 'editPageText', input: {} }],
toolCalls: [{ toolCallId: 'c1', toolName: 'editPageText', input: {} }],
toolResults: [],
content: [{ type: 'tool-error', toolName: 'editPageText', error: long }],
content: [
{
type: 'tool-error',
toolCallId: 'c1',
toolName: 'editPageText',
error: long,
},
],
},
]) as Array<Record<string, unknown>>;
const errorText = trace[1].error as string;
// Truncated (not the full 5000 chars) and carries the omission marker.
expect(errorText.length).toBeLessThan(long.length);
expect(errorText).toContain('chars omitted');
});
it('pairs parallel calls in one step with their outcomes by id', () => {
const trace = serializeSteps([
{
toolCalls: [
{ toolCallId: 'a', toolName: 'getPage', input: {} },
{ toolCallId: 'b', toolName: 'searchPages', input: {} },
],
toolResults: [{ toolCallId: 'b', toolName: 'searchPages' }],
content: [
{ type: 'tool-error', toolCallId: 'a', toolName: 'getPage', error: 'nope' },
],
},
]) as Array<Record<string, unknown>>;
// call a, outcome a (thrown), call b, outcome b (ok)
expect(trace).toHaveLength(4);
expect(trace[1]).toEqual({ toolName: 'getPage', error: 'nope', kind: 'thrown' });
expect(trace[3]).toEqual({ toolName: 'searchPages', ok: true });
});
});
// #490: every assistant row flushAssistant writes carries the v2 era marker so a
// dual-shape diagnostic query can branch on the trace shape without inspecting it.
describe('toolTraceVersion era marker (#490)', () => {
it('stamps metadata.toolTraceVersion = 2 on every flushed row', () => {
const seed = flushAssistant([], '', 'streaming');
expect(seed.metadata.toolTraceVersion).toBe(2);
const done = flushAssistant(
[
{
text: 'ok',
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
toolResults: [{ toolCallId: 'c1', toolName: 'getPage' }],
},
],
'',
'completed',
{ finishReason: 'stop' },
);
expect(done.metadata.toolTraceVersion).toBe(2);
});
});
// #490 replay-budget signal helpers over persisted history.
describe('lastAssistantContextTokens', () => {
const row = (
role: string,
metadata: Record<string, unknown> | null,
): AiChatMessage => ({ role, metadata }) as unknown as AiChatMessage;
it('reads the most recent assistant turn contextTokens (provider fact)', () => {
const hist = [
row('user', null),
row('assistant', { contextTokens: 12000 }),
row('user', null),
row('assistant', { contextTokens: 41000 }),
];
expect(lastAssistantContextTokens(hist)).toBe(41000);
});
it('returns undefined when the last assistant turn recorded no usage', () => {
const hist = [row('assistant', { error: 'boom' }), row('user', null)];
expect(lastAssistantContextTokens(hist)).toBeUndefined();
expect(lastAssistantContextTokens([])).toBeUndefined();
});
});
// #490 snapshotOpenPage fast-path: skip the full Markdown export + upsert when a
// snapshot already exists at the page's CURRENT version (same updated_at instant).
describe('snapshotOpenPage fast-path (#490)', () => {
function makeSvc(existingSnapshot: unknown, pageUpdatedAt: Date) {
const exportPageMarkdown = jest.fn(async () => '# md');
const upsert = jest.fn(async () => undefined);
const findByChatPage = jest.fn(async () => existingSnapshot);
const pageRepo = {
findById: jest.fn(async () => ({
id: 'p1',
workspaceId: 'ws1',
updatedAt: pageUpdatedAt,
})),
};
const svc = new AiChatService(
{} as never, // ai
{} as never, // aiChatRepo
{} as never, // aiChatMessageRepo
{ findByChatPage, upsert } as never, // aiChatPageSnapshotRepo
{} as never, // aiSettings
{ exportPageMarkdown } as never, // tools
{} as never, // mcpClients
{} as never, // aiAgentRoleRepo
pageRepo as never, // pageRepo
{} as never, // pageAccess
{} as never, // environment
);
return { svc, exportPageMarkdown, upsert, findByChatPage };
}
const args = () =>
[
'chat1',
'p1',
{ id: 'ws1' } as never,
{ id: 'u1' } as never,
'sess',
] as const;
it('skips export + upsert when the snapshot is already at this page version', async () => {
const t = new Date('2026-07-07T10:00:00Z');
const { svc, exportPageMarkdown, upsert } = makeSvc(
{ pageUpdatedAt: t, contentMd: '# md' },
t,
);
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
.snapshotOpenPage(...args());
expect(exportPageMarkdown).not.toHaveBeenCalled();
expect(upsert).not.toHaveBeenCalled();
});
it('exports + upserts when the page advanced since the snapshot', async () => {
const { svc, exportPageMarkdown, upsert } = makeSvc(
{ pageUpdatedAt: new Date('2026-07-07T10:00:00Z'), contentMd: 'old' },
new Date('2026-07-07T11:00:00Z'),
);
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
.snapshotOpenPage(...args());
expect(exportPageMarkdown).toHaveBeenCalledTimes(1);
expect(upsert).toHaveBeenCalledTimes(1);
});
it('seeds (exports + upserts) on the first turn (no snapshot yet)', async () => {
const { svc, exportPageMarkdown, upsert } = makeSvc(
undefined,
new Date('2026-07-07T10:00:00Z'),
);
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
.snapshotOpenPage(...args());
expect(exportPageMarkdown).toHaveBeenCalledTimes(1);
expect(upsert).toHaveBeenCalledTimes(1);
});
});
// #490 deferred-tool activation persisted across turns.
describe('seedActivatedTools', () => {
const valid = new Set(['Search_web', 'getPageJson', 'diffPageVersions']);
it('seeds from persisted metadata, intersected with current valid names', () => {
expect(
seedActivatedTools(
{ activatedTools: ['Search_web', 'getPageJson'] },
valid,
),
).toEqual(['Search_web', 'getPageJson']);
});
it('drops a stored tool that is no longer valid (allowlist/role changed)', () => {
// 'Habr_publish' was activated before but is not in the current allowlist.
expect(
seedActivatedTools({ activatedTools: ['Search_web', 'Habr_publish'] }, valid),
).toEqual(['Search_web']);
});
it('is empty/robust for missing, non-array, or unknown-shaped metadata', () => {
expect(seedActivatedTools(undefined, valid)).toEqual([]);
expect(seedActivatedTools({}, valid)).toEqual([]);
expect(seedActivatedTools({ activatedTools: 'nope' }, valid)).toEqual([]);
expect(
seedActivatedTools({ activatedTools: [1, 'getPageJson', null] }, valid),
).toEqual(['getPageJson']);
});
it('de-duplicates stored names', () => {
expect(
seedActivatedTools(
{ activatedTools: ['getPageJson', 'getPageJson'] },
valid,
),
).toEqual(['getPageJson']);
});
});
describe('lastAssistantReplayOverflow', () => {
const row = (
role: string,
metadata: Record<string, unknown> | null,
): AiChatMessage => ({ role, metadata }) as unknown as AiChatMessage;
it('is true only when the LAST assistant turn overflowed', () => {
expect(
lastAssistantReplayOverflow([
row('assistant', { replayOverflow: true }),
row('user', null),
]),
).toBe(true);
// A recovered (later, non-overflow) assistant turn clears it.
expect(
lastAssistantReplayOverflow([
row('assistant', { replayOverflow: true }),
row('user', null),
row('assistant', { contextTokens: 5 }),
]),
).toBe(false);
expect(lastAssistantReplayOverflow([])).toBe(false);
});
// #490 reactive recovery: a prior turn stamped `replayOverflow` must make the
// NEXT turn's effective budget the AGGRESSIVE 0.5x cut — that harder trim is
// what un-bricks a chat that just 400'd on the context window. This exercises
// the exact wiring the service uses: read the stamp, then scale the threshold.
it('#490: a prior replayOverflow drives the next turn to the 0.5x aggressive budget', () => {
const history = [
row('assistant', { replayOverflow: true }),
row('user', null),
];
const priorOverflowed = lastAssistantReplayOverflow(history);
expect(priorOverflowed).toBe(true);
// Base budget 100k -> aggressive recovery halves it to 50k this turn.
expect(resolveEffectiveReplayThreshold(100_000, priorOverflowed)).toBe(50_000);
// Odd base floors, not rounds.
expect(resolveEffectiveReplayThreshold(99_999, true)).toBe(49_999);
// No prior overflow -> the base budget is used verbatim (no aggressive cut).
expect(resolveEffectiveReplayThreshold(100_000, false)).toBe(100_000);
// An explicit off-switch (null) is never overridden, even on recovery.
expect(resolveEffectiveReplayThreshold(null, true)).toBeNull();
});
});
describe('rowToUiMessage', () => {
@@ -618,6 +930,23 @@ describe('flushAssistant', () => {
expect(flushed.metadata.error).toBe('boom');
});
// #490 observability: the replay budgeter's decision is stamped on the turn.
it('records replayTrimmedToTokens + replayOverflow when provided', () => {
const f = flushAssistant([], '', 'error', {
error: 'ctx',
replayTrimmedToTokens: 42_000,
replayOverflow: true,
});
expect(f.metadata.replayTrimmedToTokens).toBe(42_000);
expect(f.metadata.replayOverflow).toBe(true);
});
it('omits the replay metadata when not provided', () => {
const f = flushAssistant([], '', 'completed', { finishReason: 'stop' });
expect('replayTrimmedToTokens' in f.metadata).toBe(false);
expect('replayOverflow' in f.metadata).toBe(false);
});
// #274 observability: the page-change diff the agent saw this turn is persisted
// to metadata.pageChanged when a non-empty diff was injected, and omitted when
// the diff is empty/whitespace or the arg is not supplied.
apps/server/src/core/ai-chat/ai-chat.service.ts
+470 -97
View File
@@ -55,6 +55,12 @@ import {
type SelectionContext,
} from './tools/current-page.util';
import { roleModelOverride } from './roles/role-model-config';
import {
resolveReplayBudget,
resolveEffectiveReplayThreshold,
isContextOverflowError,
trimHistoryForReplay,
} from './history-budget';
import {
startSseHeartbeat,
stripStreamingHopByHopHeaders,
@@ -117,9 +123,14 @@ const FINAL_STEP_NUDGE =
// NO text at all (#444, mitigates the "empty turn" the lockdown used to prevent
// when the toggle is OFF). Makes the exhausted-without-answer state explicit to
// the user and, on replay, to the model on the next turn.
// The persisted content is the app's base locale (en-US) — which is ALSO the
// i18n key the client localizes through `t()` — instead of a hardcoded Russian
// string (it used to render Russian for every locale, and fed Russian back to
// the model on replay). Keep it a plain, model-readable English sentence so the
// next turn's replay reads cleanly; the client resolves the locale.
const STEP_LIMIT_NO_ANSWER_MARKER =
'(Достигнут лимит шагов — итоговый ответ не сформулирован; работа могла ' +
'остаться незавершённой. Напишите «продолжай», чтобы агент продолжил.)';
'(Step limit reached — no final answer was produced; the work may be ' +
'unfinished. Reply "continue" to let the agent carry on.)';
// Reason recorded in ai_chat_runs.error / the assistant row when the token-
// degeneration detector (#444) aborts a run. Distinct from a user Stop (no error)
@@ -127,6 +138,15 @@ const STEP_LIMIT_NO_ANSWER_MARKER =
const OUTPUT_DEGENERATION_ERROR =
'Output degeneration detected (repeated token loop)';
// Prefix recorded on the assistant row when the provider rejected the turn for
// CONTEXT OVERFLOW (#490): the replayed history exceeded the model's window. The
// row is ALSO stamped `metadata.replayOverflow` so the NEXT turn's budgeter trims
// aggressively (the reactive recovery — the overflowing turn had no usage signal
// to trigger preventive trimming, so the classified 400 is what un-bricks it).
export const CONTEXT_OVERFLOW_ERROR_PREFIX =
'Диалог превысил контекстное окно модели; история будет агрессивно ' +
'сокращена на следующем ходу.';
/**
* Compute the step-budget warning text (#444), or '' when this step is outside
* the warning band. The warning fires on steps
@@ -882,6 +902,21 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
const freshPage = await this.pageRepo.findById(pageId);
// Page deleted during the turn (or somehow foreign) => don't write.
if (!freshPage || freshPage.workspaceId !== workspace.id) return;
// Fast-path (#490): if a snapshot already exists at THIS page version
// (same updated_at instant), its content is already current — skip the full
// Markdown export + upsert entirely. A turn that did NOT touch the open page
// (the common case) thus does no snapshot work. This mirrors the read-side
// fast path in detectPageChange (sameInstant): both trust that a page edit
// bumps updated_at. When the agent (or a human) DID edit the page this turn,
// updated_at advanced, so this does not match and we re-export as before.
const existing = await this.aiChatPageSnapshotRepo.findByChatPage(
chatId,
pageId,
workspace.id,
);
if (existing && sameInstant(existing.pageUpdatedAt, freshPage.updatedAt)) {
return;
}
const currentMd = await this.tools.exportPageMarkdown(
user,
sessionId,
@@ -921,10 +956,17 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// supplied or the supplied one does not belong to this workspace.
let isNewChat = false;
let chatId = body.chatId;
// Persisted chat-level metadata bag (#490): read once here so the deferred-tool
// activation set can be seeded from the previous turn. Undefined for a new chat.
let chatMetadata: Record<string, unknown> | undefined;
if (chatId) {
const existing = await this.aiChatRepo.findById(chatId, workspace.id);
if (!existing) {
chatId = undefined;
} else {
chatMetadata = (existing.metadata ?? undefined) as
| Record<string, unknown>
| undefined;
}
}
// The open page the client sent is attacker-controllable — BOTH its id and
@@ -1086,7 +1128,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// per-row conversion and degraded to plain text with a "[tool context
// omitted]" marker rather than 500-ing the whole turn (silent loss of tool
// context is not acceptable — the model must see the truncation).
const messages = await convertHistoryResilient(uiMessages, (index, err) =>
let messages = await convertHistoryResilient(uiMessages, (index, err) =>
this.logger.warn(
`Degraded unconvertible history row ${index} on chat ${chatId} to text: ${
err instanceof Error ? err.message : 'unknown error'
@@ -1135,6 +1177,56 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// Here we only need the admin-configured system prompt.
const resolved = await this.aiSettings.resolve(workspace.id);
// History-replay token budget (#490). The full conversation is replayed to
// the provider every turn, so a long chat eventually 400s on the context
// window — forever. Bound the REPLAYED history (never the persisted rows).
// PRIMARY signal is the provider's own fact: the last turn's contextTokens.
const replayBudget = resolveReplayBudget(resolved?.chatContextWindowRaw);
if (replayBudget.usedDefault) {
// The default fires precisely for installs with NO configured window —
// the ones that hit terminal overflow. Warn so it is observable.
this.logger.warn(
`AI chat (chat ${chatId}): no chatContextWindow configured; ` +
`applying the default replay budget (${replayBudget.thresholdTokens} tokens).`,
);
}
// Last turn's provider-reported context size (authoritative when present).
const priorContextTokens = lastAssistantContextTokens(oldHistory);
// Reactive recovery (#490): if the LAST turn was rejected for context
// overflow (stamped by onError), trim AGGRESSIVELY this turn — the
// overflowing turn produced no usage signal, so a normal-threshold trim may
// not shrink enough to fit. This is what un-bricks a chat that just 400'd.
const priorOverflowed = lastAssistantReplayOverflow(oldHistory);
const effectiveThreshold = resolveEffectiveReplayThreshold(
replayBudget.thresholdTokens,
priorOverflowed,
);
if (priorOverflowed) {
this.logger.warn(
`AI chat (chat ${chatId}): previous turn hit context overflow; ` +
`applying aggressive replay budget (${effectiveThreshold} tokens).`,
);
}
const preTrim = trimHistoryForReplay(
messages,
effectiveThreshold,
// A prior OVERFLOW means the provider count is stale/absent — force the
// char-estimate path by ignoring priorContextTokens on recovery.
priorOverflowed ? undefined : priorContextTokens,
);
messages = preTrim.messages;
// Observability (#490): record the budgeter's decision on the turn so the UI
// can surface "replay truncated at N tokens". Threaded into flushAssistant.
let replayTrimmedToTokens: number | undefined = preTrim.trimmed
? preTrim.estimatedTokens
: undefined;
if (preTrim.trimmed) {
this.logger.log(
`AI chat (chat ${chatId}): replay history trimmed to ~${preTrim.estimatedTokens} ` +
`tokens (budget ${replayBudget.thresholdTokens}).`,
);
}
// Build the external MCP toolset FIRST so the system prompt can carry each
// connected server's admin-authored guidance (#180). Merge in admin-
// configured external MCP tools (web search, etc.; §6.8). A down/slow
@@ -1326,10 +1418,19 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// tools + ALL external MCP tools), computed from the ACTUAL toolset so an
// external tool is loadable by its namespaced name. loadTools rejects any
// name outside this set.
const activatedTools = new Set<string>();
const validDeferredNames = new Set<string>(
Object.keys(baseTools).filter((k) => !CORE_TOOL_SET.has(k)),
);
// #490: seed the activation set from the chat's PERSISTED set so the model
// does not re-run loadTools every turn to re-activate the same tools. Only
// when deferred loading is enabled, and ALWAYS intersected with the CURRENT
// valid deferred names — an allowlist/role change must never resurrect a tool
// that no longer exists (prepareAgentStep would get a phantom active name).
const activatedTools = new Set<string>(
deferredEnabled
? seedActivatedTools(chatMetadata, validDeferredNames)
: [],
);
// Add the loadTools meta-tool ONLY when the feature is enabled; when off the
// toolset and behavior are exactly as before.
const tools = deferredEnabled
@@ -1339,6 +1440,39 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
}
: baseTools;
// #490: persist the (deterministically ordered) activation set back onto the
// chat metadata at turn end, so the NEXT turn seeds from it. Once-guarded and
// skipped when nothing new was activated (the set equals its seed) so an
// ordinary turn adds no extra write. Preserves other metadata keys.
let activatedToolsPersisted = false;
const persistActivatedTools = async (): Promise<void> => {
if (!deferredEnabled || activatedToolsPersisted || !chatId) return;
activatedToolsPersisted = true;
const current = [...activatedTools].sort();
const seeded = seedActivatedTools(chatMetadata, validDeferredNames).sort();
if (current.length === 0 || current.join('') === seeded.join('')) {
return; // nothing new activated -> no write
}
try {
await this.aiChatRepo.update(
chatId,
{
metadata: {
...(chatMetadata ?? {}),
activatedTools: current,
},
} as never,
workspace.id,
);
} catch (err) {
this.logger.warn(
`Failed to persist activated tools (chat ${chatId}): ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
}
};
// Accumulate the turn's streamed output so a provider error / disconnect can
// persist the PARTIAL answer the user already saw — the SDK's onError/onAbort
// callbacks don't hand us the in-progress text. `capturedSteps` holds finished
@@ -1347,6 +1481,11 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
const capturedSteps: StepLike[] = [];
let inProgressText = '';
// Per-turn step->parts memo (#490): shared across every flushAssistant call
// this turn so each finished step's (large) output is JSON-stringified ONCE,
// not re-stringified on every subsequent onStepFinish flush (was O(N²)).
const partsCache: StepPartsCache = new WeakMap();
// Token-degeneration guard (#444). When the final-step lockdown is OFF, a
// runaway repetition loop (the 255KB "loadTools." incident) is aborted via
// this internal controller, unioned with the run/socket signal below. The
@@ -1404,27 +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 }),
{ 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;
}
};
@@ -1514,6 +1665,13 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
system,
messages,
tools,
// Pin the AI SDK per-request retry budget explicitly instead of relying
// on its default (which is also 2). Connection arithmetic per turn:
// (1 + maxRetries=2) × (1 + AI_STREAM_PRE_RESPONSE_RETRIES) network
// connects worst-case — the two retry layers compose, so making the SDK
// side explicit keeps that ceiling visible and pinned against SDK-default
// drift.
maxRetries: 2,
// No maxOutputTokens cap on the agent: tool-call arguments (e.g. a full
// page body for the write tools) are emitted as OUTPUT tokens, so a fixed
// cap would truncate complex tool calls mid-argument. Let the model use its
@@ -1593,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);
@@ -1649,6 +1824,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// closure scope here). Omitted/0 = no limit.
maxContextTokens: resolved?.chatContextWindow,
pageChanged,
partsCache,
replayTrimmedToTokens,
}),
);
// #184/#487: the RUN is finalized ALWAYS (never gated on the message).
@@ -1673,6 +1850,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// own edits are baked in — and this also SEEDS the snapshot on the first
// turn. Runs once across every terminal path (see snapshotTurnEnd).
await snapshotTurnEnd();
// #490: persist the deferred-tool activation set for the next turn.
await persistActivatedTools();
// Generate the chat title for a freshly created chat AFTER the stream's
// provider call has completed — NOT concurrently with it. The z.ai coding
@@ -1696,7 +1875,16 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// object, so the actual provider cause is clearly logged. Reuse the
// shared formatter so provider error formatting stays unified.
const e = error as { stack?: string };
const errorText = describeProviderError(error, String(error));
// #490 reactive branch: classify a CONTEXT-OVERFLOW rejection (the
// replayed history exceeded the model window). The overflowing turn had
// no prior usage to trigger preventive trimming, so we record a clear,
// distinguishable cause AND stamp the row so the NEXT turn's budgeter
// trims aggressively — the reactive recovery that un-bricks the chat.
const overflow = isContextOverflowError(error);
const providerError = describeProviderError(error, String(error));
const errorText = overflow
? `${CONTEXT_OVERFLOW_ERROR_PREFIX} (${providerError})`
: providerError;
this.logger.error(`AI chat stream error: ${errorText}`, e?.stack);
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: timing of
// an error-terminated stream.
@@ -1714,6 +1902,9 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
flushAssistant(capturedSteps, inProgressText, 'error', {
error: errorText,
pageChanged,
partsCache,
replayTrimmedToTokens,
replayOverflow: overflow || undefined,
}),
);
// #184: settle the RUN as failed, carrying the provider/transport cause.
@@ -1723,6 +1914,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// committed before the error must be baked into the snapshot, or the
// next turn would mis-report it as a user edit.
await snapshotTurnEnd();
// #490: persist the deferred-tool activation set for the next turn.
await persistActivatedTools();
},
onAbort: async ({ steps }) => {
// #444: distinguish a degeneration abort (our internal controller) from
@@ -1737,6 +1930,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
flushAssistant(capturedSteps, truncated, 'error', {
error: OUTPUT_DEGENERATION_ERROR,
pageChanged,
partsCache,
}),
);
if (runId)
@@ -1747,6 +1941,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
);
await closeExternalClients();
await snapshotTurnEnd();
// #490: persist the deferred-tool activation set for the next turn.
await persistActivatedTools();
return;
}
const partialChars =
@@ -1771,6 +1967,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
await finalizeAssistant(
flushAssistant(capturedSteps, inProgressText, 'aborted', {
pageChanged,
partsCache,
}),
);
// #184: settle the RUN as aborted (an explicit user stop reached the
@@ -1781,6 +1978,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// committed before the client disconnect / stop() must be baked into the
// snapshot, or the next turn would mis-report it as a user edit.
await snapshotTurnEnd();
// #490: persist the deferred-tool activation set for the next turn.
await persistActivatedTools();
},
});
@@ -2091,6 +2290,70 @@ export function chatStreamMetadata(
return undefined;
}
/**
* The provider-reported context size of the most recent assistant turn, read from
* its persisted `metadata.contextTokens` (#490 replay budgeter's PRIMARY signal —
* the provider's own fact, not an estimate). Returns undefined for a chat with no
* assistant turn yet, or one whose last turn recorded no usage (e.g. it errored),
* in which case the budgeter falls back to the char-estimate.
*/
export function lastAssistantContextTokens(
history: ReadonlyArray<AiChatMessage>,
): number | undefined {
for (let i = history.length - 1; i >= 0; i--) {
const row = history[i];
if (row.role !== 'assistant') continue;
const meta = (row.metadata ?? {}) as { contextTokens?: unknown };
const n = meta.contextTokens;
return typeof n === 'number' && Number.isFinite(n) && n > 0 ? n : undefined;
}
return undefined;
}
/**
* Seed the per-turn deferred-tool activation set from a chat's persisted metadata
* (#490), INTERSECTED with the current valid deferred names. Persisting the set
* across turns saves the model re-running loadTools every turn to re-activate the
* same tools; intersecting on load means a changed allowlist / role can never
* resurrect a tool that no longer exists (which would hand prepareAgentStep a
* phantom active name). Tolerant of any stored shape — a non-array is ignored.
*/
export function seedActivatedTools(
metadata: Record<string, unknown> | undefined,
validDeferredNames: ReadonlySet<string>,
): string[] {
const stored = metadata?.activatedTools;
if (!Array.isArray(stored)) return [];
const seen = new Set<string>();
const out: string[] = [];
for (const name of stored) {
if (typeof name === 'string' && validDeferredNames.has(name) && !seen.has(name)) {
seen.add(name);
out.push(name);
}
}
return out;
}
/**
* Whether the most recent assistant turn was rejected for CONTEXT OVERFLOW
* (#490): its row carries `metadata.replayOverflow` (stamped by the stream's
* onError). The next turn's budgeter reads this to trim aggressively — the
* reactive recovery. Only the LAST assistant turn matters (an older overflow was
* already recovered), so we stop at the first assistant row scanning backwards.
*/
export function lastAssistantReplayOverflow(
history: ReadonlyArray<AiChatMessage>,
): boolean {
for (let i = history.length - 1; i >= 0; i--) {
const row = history[i];
if (row.role !== 'assistant') continue;
const meta = (row.metadata ?? {}) as { replayOverflow?: unknown };
return meta.replayOverflow === true;
}
return false;
}
/** The last message with role 'user' from a useChat payload, if any. */
function lastUserMessage(
messages: UIMessage[] | undefined,
@@ -2152,6 +2415,15 @@ export function sanitizeUserParts(
/** Marker for a history row whose tool parts could not be replayed (#489). */
export const TOOL_CONTEXT_OMITTED_MARKER = '[tool context omitted]';
/**
* Synthetic error text for a tool call that neither returned a result nor threw
* a `tool-error` — i.e. it was interrupted mid-step (an abort / server restart).
* Shared by `assistantParts` (the replayed `output-error` part) and
* `serializeSteps` (the `{ kind: 'interrupted' }` trace element) so the replay
* text and the trace stay in lockstep (#490).
*/
export const TOOL_CALL_INCOMPLETE_TEXT = 'Tool call did not complete.';
/**
* Convert persisted UI history to model messages, tolerating a single poisoned
* row (#489). `convertToModelMessages` over the WHOLE array throws if ANY row is
@@ -2359,71 +2631,97 @@ function normalizeToolError(error: unknown): string {
*/
// Exported only so the unit tests can import these pure helpers; exporting
// them does not change runtime behavior.
/**
* Per-turn memo for {@link assistantParts}: a step's rebuilt parts keyed by the
* step OBJECT's identity (#490). A finished step in `capturedSteps` keeps a stable
* reference across every mid-stream flush, and `compactToolOutput` inside it does a
* `JSON.stringify` of the whole (often 50–200 KB) output — so without a memo each
* `onStepFinish` re-stringifies EVERY prior step's output (O(N²) stringify over a
* turn). Keyed by step identity => one stringify per step per turn. WeakMap so a
* turn's steps are GC'd with the turn.
*/
export type StepPartsCache = WeakMap<object, Array<Record<string, unknown>>>;
/** Build the parts for ONE step (text + a part per tool call). Pure. */
function buildStepParts(step: StepLike): Array<Record<string, unknown>> {
const parts: Array<Record<string, unknown>> = [];
if (step.text) {
parts.push({ type: 'text', text: step.text });
}
// Index this step's results by tool call id to pair calls with outputs.
const resultsById = new Map<string, unknown>();
for (const r of step.toolResults ?? []) {
if (r.toolCallId) resultsById.set(r.toolCallId, r.output);
}
// Index this step's THROWN tool failures (ai@6 `tool-error` content parts)
// by tool call id, so a call that failed replays with its real error text.
const errorsById = new Map<string, unknown>();
for (const part of step.content ?? []) {
if (part.type === 'tool-error' && part.toolCallId) {
errorsById.set(part.toolCallId, part.error);
}
}
for (const call of step.toolCalls ?? []) {
if (!call.toolName || !call.toolCallId) continue;
const hasResult = resultsById.has(call.toolCallId);
if (hasResult) {
// output-available: the tool returned; the next turn replays its result.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-available',
input: call.input,
output: compactToolOutput(resultsById.get(call.toolCallId)),
});
} else if (errorsById.has(call.toolCallId)) {
// The tool THREW: replay the REAL error so the model on the next turn
// knows WHY the call failed (and does not blindly repeat it). An
// output-error round-trips through convertToModelMessages as a balanced
// tool-call + tool-result, keeping the rebuilt history valid.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-error',
input: call.input,
errorText: normalizeToolError(errorsById.get(call.toolCallId)),
});
} else {
// No paired result AND no tool-error (e.g. aborted mid-step). Persisting
// a bare tool-call (input-available) would replay as an unpaired call and
// throw MissingToolResultsError on the next turn (convertToModelMessages
// emits no tool-result for it). Emit a SYNTHETIC paired result instead:
// an output-error round-trips through convertToModelMessages as a
// balanced tool-call + tool-result, keeping the rebuilt history valid.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-error',
input: call.input,
errorText: TOOL_CALL_INCOMPLETE_TEXT,
});
}
}
return parts;
}
export function assistantParts(
steps: ReadonlyArray<StepLike> | undefined,
fallbackText: string,
cache?: StepPartsCache,
): UIMessage['parts'] {
const parts: Array<Record<string, unknown>> = [];
let sawText = false;
for (const step of steps ?? []) {
if (step.text) {
parts.push({ type: 'text', text: step.text });
sawText = true;
}
// Index this step's results by tool call id to pair calls with outputs.
const resultsById = new Map<string, unknown>();
for (const r of step.toolResults ?? []) {
if (r.toolCallId) resultsById.set(r.toolCallId, r.output);
}
// Index this step's THROWN tool failures (ai@6 `tool-error` content parts)
// by tool call id, so a call that failed replays with its real error text.
const errorsById = new Map<string, unknown>();
for (const part of step.content ?? []) {
if (part.type === 'tool-error' && part.toolCallId) {
errorsById.set(part.toolCallId, part.error);
}
}
for (const call of step.toolCalls ?? []) {
if (!call.toolName || !call.toolCallId) continue;
const hasResult = resultsById.has(call.toolCallId);
if (hasResult) {
// output-available: the tool returned; the next turn replays its result.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-available',
input: call.input,
output: compactToolOutput(resultsById.get(call.toolCallId)),
});
} else if (errorsById.has(call.toolCallId)) {
// The tool THREW: replay the REAL error so the model on the next turn
// knows WHY the call failed (and does not blindly repeat it). An
// output-error round-trips through convertToModelMessages as a balanced
// tool-call + tool-result, keeping the rebuilt history valid.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-error',
input: call.input,
errorText: normalizeToolError(errorsById.get(call.toolCallId)),
});
} else {
// No paired result AND no tool-error (e.g. aborted mid-step). Persisting
// a bare tool-call (input-available) would replay as an unpaired call and
// throw MissingToolResultsError on the next turn (convertToModelMessages
// emits no tool-result for it). Emit a SYNTHETIC paired result instead:
// an output-error round-trips through convertToModelMessages as a
// balanced tool-call + tool-result, keeping the rebuilt history valid.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-error',
input: call.input,
errorText: 'Tool call did not complete.',
});
}
// Memoize per step object (#490): a finished step is immutable and keeps its
// reference across flushes, so its parts (and the costly output stringify) are
// built exactly once per turn. A cache miss (or no cache) just rebuilds.
let stepParts = cache?.get(step as object);
if (!stepParts) {
stepParts = buildStepParts(step);
cache?.set(step as object, stepParts);
}
parts.push(...stepParts);
}
const sawText = parts.some((p) => p.type === 'text');
if (!sawText && fallbackText) {
// No per-step text (e.g. a single final block): append the final text after
// any tool parts so the natural call -> result -> answer order is preserved.
@@ -2586,6 +2884,16 @@ export function flushAssistant(
maxContextTokens?: number;
error?: string;
pageChanged?: { title: string; diff: string } | null;
// Per-turn step->parts memo (#490): pass the SAME cache on every flush of a
// turn so each finished step's output is stringified once, not once per flush.
partsCache?: StepPartsCache;
// #490 observability: when the replay budgeter trimmed this turn's history,
// the (estimated) token size it trimmed to — the UI can show "replay truncated
// at N tokens". Omitted when nothing was trimmed.
replayTrimmedToTokens?: number;
// #490 reactive branch: set when the provider rejected this turn for context
// overflow. Stamped into metadata so the NEXT turn's budgeter trims aggressively.
replayOverflow?: boolean;
},
): AssistantFlush {
const finished = capturedSteps ?? [];
@@ -2595,13 +2903,32 @@ export function flushAssistant(
// in-progress step's text (the partial answer cut off by an error/abort, or
// simply not yet flushed mid-stream) as the last text part so the persisted
// parts match what streamed to the client.
const parts = assistantParts(finished, '') as unknown as Array<
Record<string, unknown>
>;
const parts = assistantParts(
finished,
'',
extra?.partsCache,
) as unknown as Array<Record<string, unknown>>;
if (trailing) parts.push({ type: 'text', text: trailing });
const metadata: Record<string, unknown> = {
parts: parts as unknown as UIMessage['parts'],
// Era marker for the `tool_calls` trace shape (#490): v2 stores outcome flags
// ({ ok } / { error, kind }) and NO tool output (the output lives once in
// `parts`). Old rows have no marker and the legacy { output } shape; a
// dual-shape query branches on this. Old rows are deliberately NOT migrated.
toolTraceVersion: 2,
// #491 STEP MARKER: the number of FINISHED steps whose parts are in THIS row,
// written by the SAME flush that builds `parts` (atomically — they are both
// derived from `finished`, so the marker can NEVER disagree with the persisted
// parts). This is the step-alignment anchor the resume stack builds on:
// - the registry rotates its retention ring only on a CONFIRMED persist of
// step N (commit 3);
// - attach slices the tail at "step > N" from the client's persisted seed.
// It is NOT `run.stepCount`: recordStep is fire-and-forget and NOT atomic with
// the parts write, so stepCount could race ahead of the persisted parts
// (seed↔marker drift). The in-progress trailing text (an error/abort partial,
// or a mid-stream flush) is NOT a finished step and is excluded from the count.
stepsPersisted: finished.length,
};
// finishReason: prefer an explicit one; else derive a sensible value from the
// terminal status (so onError/onAbort records keep their historical reason).
@@ -2617,6 +2944,9 @@ export function flushAssistant(
if (extra?.contextTokens) metadata.contextTokens = extra.contextTokens;
if (extra?.maxContextTokens)
metadata.maxContextTokens = extra.maxContextTokens;
if (extra?.replayTrimmedToTokens)
metadata.replayTrimmedToTokens = extra.replayTrimmedToTokens;
if (extra?.replayOverflow) metadata.replayOverflow = true;
if (extra?.error) metadata.error = extra.error;
// Persist the page-change diff the agent saw this turn (#274 observability),
// so history / the Markdown export can show what the user changed. Only when
@@ -2642,42 +2972,85 @@ export function flushAssistant(
/**
* Reduce SDK step objects to a compact, JSON-serializable trace for the
* `tool_calls` column. Stores only what the UI action-log and history need —
* never raw provider payloads or keys.
* `tool_calls` column — trace format **v2** (#490).
*
* v2 stores, per call, ONLY the metadata a queryable trace needs — never the
* tool OUTPUT. Before #490 each output was persisted TWICE: once here (compacted)
* and once in `metadata.parts` (via `assistantParts`), so a 50-step run with
* 50–200 KB outputs wrote hundreds of MB per turn (each `onStepFinish` rewrote
* the whole row). The parts copy is the one the model replays and the UI/Markdown
* export render, so the trace copy of the output was pure duplication. v2 keeps
* the output ONLY in parts and reduces the trace to outcome flags.
*
* Element shapes (paired per call, in order):
* - `{ toolName, input }` — the call
* - `{ toolName, ok: true }` — it returned a result (success)
* - `{ toolName, error, kind: 'thrown' }` — it threw a `tool-error`
* - `{ toolName, error, kind: 'interrupted' }` — no result and no throw (an
* abort / server restart mid-step). `kind` is MANDATORY: without it a
* synthetic "Tool call did not complete." is indistinguishable from a real
* hard-fail and pollutes any error-rate scan. The distinction is STRUCTURAL
* (an `errorsById` hit vs the synthetic fallback branch), NOT a per-tool
* classifier — soft failures stay OUT of the trace (they live in
* `metadata.parts` outputs; a per-tool mirror would persist its own bugs).
*
* Rows carry `metadata.toolTraceVersion: 2` (set by {@link flushAssistant}) so a
* dual-shape query can branch on the era. Old rows are NOT migrated (rewriting
* giant jsonb is the very WAL churn this removes); see docs/reading-ai-logs.md.
*/
export function serializeSteps(
steps: ReadonlyArray<{
toolCalls?: ReadonlyArray<{ toolName?: string; input?: unknown }>;
toolResults?: ReadonlyArray<{ toolName?: string; output?: unknown }>;
toolCalls?: ReadonlyArray<{
toolCallId?: string;
toolName?: string;
input?: unknown;
}>;
toolResults?: ReadonlyArray<{ toolCallId?: string; toolName?: string }>;
content?: ReadonlyArray<{
type?: string;
toolCallId?: string;
toolName?: string;
error?: unknown;
}>;
}>,
): unknown {
const calls: Array<{
toolName?: string;
input?: unknown;
output?: unknown;
error?: string;
}> = [];
const calls: Array<
| { toolName?: string; input?: unknown }
| { toolName?: string; ok: true }
| { toolName?: string; error: string; kind: 'thrown' | 'interrupted' }
> = [];
for (const step of steps ?? []) {
// Index this step's results + thrown errors by tool call id, so each call is
// paired with its outcome (mirrors assistantParts' pairing exactly).
const resultIds = new Set<string>();
for (const r of step.toolResults ?? []) {
if (r.toolCallId) resultIds.add(r.toolCallId);
}
const errorsById = new Map<string, unknown>();
for (const part of step.content ?? []) {
if (part.type === 'tool-error' && part.toolCallId) {
errorsById.set(part.toolCallId, part.error);
}
}
for (const call of step.toolCalls ?? []) {
calls.push({ toolName: call.toolName, input: call.input });
}
for (const r of step.toolResults ?? []) {
calls.push({ toolName: r.toolName, output: compactToolOutput(r.output) });
}
// ai@6 surfaces a THROWN tool failure as a `tool-error` content part, NOT as
// a `toolResults` entry. Record it as its own paired element (mirroring how a
// successful result is appended) so the failure and its reason survive in the
// trace instead of leaving an orphaned call with no result.
for (const part of step.content ?? []) {
if (part.type === 'tool-error') {
if (call.toolCallId && resultIds.has(call.toolCallId)) {
// Success: the output itself lives in metadata.parts, not here.
calls.push({ toolName: call.toolName, ok: true });
} else if (call.toolCallId && errorsById.has(call.toolCallId)) {
// Hard fail: the tool threw. Persist the real (bounded) reason.
calls.push({
toolName: part.toolName,
error: normalizeToolError(part.error),
toolName: call.toolName,
error: normalizeToolError(errorsById.get(call.toolCallId)),
kind: 'thrown',
});
} else {
// Neither a result nor a throw: interrupted mid-step (abort/restart).
// Marked structurally so it never inflates a thrown-error count.
calls.push({
toolName: call.toolName,
error: TOOL_CALL_INCOMPLETE_TEXT,
kind: 'interrupted',
});
}
}
apps/server/src/core/ai-chat/ai-chat.step-marker.spec.ts
+65
View File
@@ -0,0 +1,65 @@
import { flushAssistant } from './ai-chat.service';
/**
* #491 STEP MARKER — `metadata.stepsPersisted` is written by the SAME flush that
* builds `metadata.parts`, so the marker can never disagree with the persisted
* parts (the step-alignment anchor the resume stack builds on). These are
* PROPERTY tests: they assert the marker tracks the number of FINISHED steps for
* every flush shape.
*/
// A finished step carrying one line of text and one tool call/result.
function step(i: number) {
return {
text: `step ${i}`,
toolCalls: [
{ toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } },
],
toolResults: [
{ toolCallId: `c${i}`, toolName: 'getPage', output: { title: `T${i}` } },
],
};
}
describe('flushAssistant step marker (#491)', () => {
it('seed (no steps) → stepsPersisted 0', () => {
const f = flushAssistant([], '', 'streaming');
expect(f.metadata.stepsPersisted).toBe(0);
});
it('PROPERTY: stepsPersisted equals the number of FINISHED steps, for any N', () => {
for (let n = 0; n <= 6; n++) {
const steps = Array.from({ length: n }, (_, i) => step(i));
const f = flushAssistant(steps, '', 'streaming');
expect(f.metadata.stepsPersisted).toBe(n);
// ...and the parts actually contain those N steps' text (marker agrees with
// the persisted parts — the atomicity the whole design relies on).
const parts = f.metadata.parts as Array<Record<string, unknown>>;
const textParts = parts.filter((p) => p.type === 'text');
expect(textParts).toHaveLength(n);
}
});
it('an in-progress trailing partial does NOT increment the marker', () => {
// 2 finished steps + a partial (not-yet-finished) trailing text: the marker
// counts only the CONFIRMED step boundaries, not the partial.
const f = flushAssistant([step(0), step(1)], 'partial third step', 'error', {
error: 'boom',
});
expect(f.metadata.stepsPersisted).toBe(2);
// The partial text IS persisted in parts (so the user sees it), but it is not a
// counted step.
const parts = f.metadata.parts as Array<Record<string, unknown>>;
expect(parts[parts.length - 1]).toEqual({
type: 'text',
text: 'partial third step',
});
});
it('terminal completed flush counts all finished steps', () => {
const f = flushAssistant([step(0), step(1), step(2)], '', 'completed', {
finishReason: 'stop',
});
expect(f.metadata.stepsPersisted).toBe(3);
});
});
apps/server/src/core/ai-chat/ai-chat.write-volume.int-spec.ts
+209
View File
@@ -0,0 +1,209 @@
import { randomBytes } from 'crypto';
import { Client } from 'pg';
import { flushAssistant, serializeSteps } from './ai-chat.service';
/**
* #490 write-volume regression — an OBSERVABLE-PROPERTY test on a LIVE Postgres,
* not "bytes through a mock repo" (a mock measures exactly the thing that does not
* hurt). It drives a realistic 50-step run where each step returns a ~100 KB tool
* output and, at every `onStepFinish`, UPDATEs the assistant row the way the
* service does — then reads the REAL write volume via the `pg_current_wal_lsn()`
* delta around the run.
*
* The property proven: v2 stores each tool OUTPUT only in `metadata.parts`, no
* longer ALSO in the `tool_calls` trace. So:
* 1. the trace (`tool_calls`) column's write volume is now O(Σ steps) — tiny,
* linear outcome flags — vs the pre-#490 O(N²) that re-persisted every prior
* output on every step; and
* 2. the FULL-row write volume drops sharply (the duplicated output copy is gone).
*
* Connects to the local gitmost test Postgres (docker `gitmost-test-pg` on :5432);
* SKIPS cleanly when that DB is not reachable so it never breaks a DB-less CI.
*/
const CONN =
process.env.WAL_TEST_DATABASE_URL ??
'postgresql://docmost:docmost_dev_pw@localhost:5432/docmost';
// A step whose tool output is ~100 KB (a page read), in the SDK StepLike shape.
// The body is INCOMPRESSIBLE random text — a `'x'.repeat()` filler would TOAST-
// compress to nothing and hide the real write volume (a page body does not).
function makeStep(i: number, outputBytes = 100_000) {
const body = randomBytes(Math.ceil(outputBytes * 0.75)).toString('base64');
return {
text: `step ${i} reasoning`,
toolCalls: [{ toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } }],
toolResults: [
{
toolCallId: `c${i}`,
toolName: 'getPage',
output: { id: `p${i}`, title: `Page ${i}`, body },
},
],
};
}
// The pre-#490 (v1) trace: outputs stored a SECOND time in `tool_calls`
// (the duplication #490 removed). Mirrors the OLD serializeSteps shape.
function v1Trace(steps: ReturnType<typeof makeStep>[]): unknown {
const calls: unknown[] = [];
for (const s of steps) {
for (const c of s.toolCalls) calls.push({ toolName: c.toolName, input: c.input });
for (const r of s.toolResults)
calls.push({ toolName: r.toolName, output: r.output });
}
return calls;
}
async function walDelta(
client: Client,
fn: () => Promise<void>,
): Promise<number> {
const before = (await client.query('SELECT pg_current_wal_lsn() AS l')).rows[0]
.l as string;
await fn();
// NOTE: do NOT pg_switch_wal() here — a segment switch pads the LSN to the next
// 16 MB boundary and would swamp the actual write delta. The raw LSN advances by
// the bytes of WAL emitted, which is exactly what we want to measure.
const after = (await client.query('SELECT pg_current_wal_lsn() AS l')).rows[0]
.l as string;
return Number(
(await client.query('SELECT pg_wal_lsn_diff($1,$2) AS d', [after, before]))
.rows[0].d,
);
}
describe('#490 write-volume on a live Postgres (pg_current_wal_lsn delta)', () => {
let client: Client | undefined;
let available = false;
beforeAll(async () => {
try {
client = new Client(CONN);
await client.connect();
await client.query('SELECT pg_current_wal_lsn()');
available = true;
} catch {
available = false;
client = undefined;
}
});
afterAll(async () => {
await client?.end().catch(() => undefined);
});
const STEPS = 50;
it('v2 trace write volume is O(Σ steps) — a tiny fraction of the v1 duplicate', async () => {
if (!available || !client) {
console.warn('SKIP: gitmost-test-pg not reachable; skipping WAL test.');
return;
}
const c = client;
// Isolated table so we measure only the tool_calls (trace) column's writes.
await c.query('DROP TABLE IF EXISTS _wal_trace');
await c.query('CREATE TABLE _wal_trace(id int primary key, tool_calls jsonb)');
await c.query("INSERT INTO _wal_trace VALUES (1, '[]'::jsonb)");
const steps: ReturnType<typeof makeStep>[] = [];
// v1: each step re-persists ALL prior outputs into the trace (the O(N²) churn).
const v1 = await walDelta(c, async () => {
const acc: ReturnType<typeof makeStep>[] = [];
for (let i = 0; i < STEPS; i++) {
acc.push(makeStep(i));
await c.query('UPDATE _wal_trace SET tool_calls=$1 WHERE id=1', [
JSON.stringify(v1Trace(acc)),
]);
}
steps.push(...acc);
});
await c.query("UPDATE _wal_trace SET tool_calls='[]'::jsonb WHERE id=1");
// v2: the REAL serializeSteps — outcome flags only, NO outputs.
const v2 = await walDelta(c, async () => {
const acc: ReturnType<typeof makeStep>[] = [];
for (let i = 0; i < STEPS; i++) {
acc.push(makeStep(i));
await c.query('UPDATE _wal_trace SET tool_calls=$1 WHERE id=1', [
JSON.stringify(serializeSteps(acc)),
]);
}
});
await c.query('DROP TABLE IF EXISTS _wal_trace');
// eslint-disable-next-line no-console
console.log(
`[#490 WAL] trace column over ${STEPS} steps: v1=${(v1 / 1e6).toFixed(1)}MB ` +
`v2=${(v2 / 1e6).toFixed(2)}MB (${(v1 / v2).toFixed(0)}x smaller)`,
);
// The trace no longer carries outputs: v2 is a tiny fraction of v1's WAL.
expect(v2).toBeLessThan(v1 * 0.1);
// And v2's trace WAL is small in absolute terms — O(Σ steps) of flags, not
// O(N² × output). 50 steps of ~40-byte flags is well under a few MB of WAL.
expect(v2).toBeLessThan(5_000_000);
// v1's duplicate alone is huge (≈ the 100 KB output re-written N² times).
expect(v1).toBeGreaterThan(50_000_000);
}, 120_000);
it('the full assistant row write drops sharply once the duplicate is gone', async () => {
if (!available || !client) return;
const c = client;
await c.query('DROP TABLE IF EXISTS _wal_full');
await c.query(
'CREATE TABLE _wal_full(id int primary key, content text, tool_calls jsonb, metadata jsonb, status text)',
);
await c.query("INSERT INTO _wal_full VALUES (1, '', '[]'::jsonb, '{}'::jsonb, 'streaming')");
const writeRow = async (patch: {
content: string;
toolCalls: unknown;
metadata: unknown;
status: string;
}) =>
c.query(
'UPDATE _wal_full SET content=$1, tool_calls=$2, metadata=$3, status=$4 WHERE id=1',
[
patch.content,
JSON.stringify(patch.toolCalls ?? null),
JSON.stringify(patch.metadata),
patch.status,
],
);
// v2 (real flushAssistant): outputs live once, in metadata.parts.
const v2 = await walDelta(c, async () => {
const acc: ReturnType<typeof makeStep>[] = [];
for (let i = 0; i < STEPS; i++) {
acc.push(makeStep(i));
await writeRow(flushAssistant(acc as never, '', 'streaming'));
}
});
await c.query("UPDATE _wal_full SET content='', tool_calls='[]'::jsonb, metadata='{}'::jsonb WHERE id=1");
// v1: same row PLUS the duplicated outputs in the trace column.
const v1 = await walDelta(c, async () => {
const acc: ReturnType<typeof makeStep>[] = [];
for (let i = 0; i < STEPS; i++) {
acc.push(makeStep(i));
const f = flushAssistant(acc as never, '', 'streaming');
await writeRow({ ...f, toolCalls: v1Trace(acc) });
}
});
await c.query('DROP TABLE IF EXISTS _wal_full');
// eslint-disable-next-line no-console
console.log(
`[#490 WAL] full row over ${STEPS} steps: v1=${(v1 / 1e6).toFixed(1)}MB ` +
`v2=${(v2 / 1e6).toFixed(1)}MB (saved ${((1 - v2 / v1) * 100).toFixed(0)}%)`,
);
// Removing the duplicated trace copy is a large, real write-volume reduction.
expect(v2).toBeLessThan(v1 * 0.75);
}, 120_000);
});
apps/server/src/core/ai-chat/chat-markdown.util.spec.ts
+46 -1
View File
@@ -1,5 +1,10 @@
import { buildChatMarkdown, normalizeLang } from './chat-markdown.util';
import {
buildChatMarkdown,
normalizeLang,
labelledToolNames,
} from './chat-markdown.util';
import type { AiChatMessage } from '@docmost/db/types/entity.types';
import { SHARED_TOOL_SPECS } from '../../../../../packages/mcp/src/tool-specs';
/**
* normalizeLang: the client sends `i18n.language` — a FULL locale tag like
@@ -455,3 +460,43 @@ describe('buildChatMarkdown (server) — structure', () => {
expect(md).toContain('````');
});
});
/**
* #494 — REVERSE drift-guard for the export's friendly tool labels. A label keyed
* by a tool name that no longer exists silently degrades to the generic
* "Ran tool <name>" line; nothing reddened before. This asserts every labelled
* name is a real in-app tool and that both languages label the same set.
*/
describe('tool-label parity (#494)', () => {
// In-app tool names come from the shared registry (inAppKey, excluding
// mcpOnly specs) PLUS the inline in-app-only tools that carry a friendly label.
// The only labelled inline tool is the hybrid semantic search.
const INLINE_INAPP_LABELLED = new Set(['searchPages']);
function validInAppToolNames(): Set<string> {
const names = new Set<string>(INLINE_INAPP_LABELLED);
for (const spec of Object.values(SHARED_TOOL_SPECS)) {
if ((spec as { mcpOnly?: boolean }).mcpOnly) continue;
names.add((spec as { inAppKey: string }).inAppKey);
}
return names;
}
it('en and ru label the SAME set of tools', () => {
expect(labelledToolNames('en').sort()).toEqual(
labelledToolNames('ru').sort(),
);
});
it('every labelled tool name is a real in-app tool', () => {
const valid = validInAppToolNames();
const dead = labelledToolNames('en').filter((n) => !valid.has(n));
expect(dead).toEqual([]);
});
it('the guard REDDENS for an unknown label key (mutation check)', () => {
const valid = validInAppToolNames();
// A hypothetical renamed-away label must be caught.
expect(valid.has('getPageRenamedAway')).toBe(false);
});
});
apps/server/src/core/ai-chat/chat-markdown.util.ts
+11
View File
@@ -154,6 +154,17 @@ function toolLabel(name: string, lang: ExportLang): string {
return LABELS[lang].tools[name] ?? LABELS[lang].ranTool(name);
}
/**
* The tool names that carry a hand-written friendly export label, per language.
* Exported for the drift-guard (#494): a label keyed by a tool name that no
* longer exists is a DEAD entry (the tool was renamed and now silently falls back
* to the generic `ranTool(name)` line). The guard asserts every key here is a
* real in-app tool AND that the two languages label the SAME set of tools.
*/
export function labelledToolNames(lang: ExportLang): string[] {
return Object.keys(LABELS[lang].tools);
}
/**
* Stringify an arbitrary tool input/output value for a fenced block. Strings
* pass through as-is; everything else is pretty-printed JSON, falling back to
apps/server/src/core/ai-chat/dto/ai-chat.dto.ts
+25 -1
View File
@@ -1,4 +1,10 @@
import { IsOptional, IsString, MaxLength, MinLength } from 'class-validator';
import {
IsISO8601,
IsOptional,
IsString,
MaxLength,
MinLength,
} from 'class-validator';
/** Identify a chat by id (workspace-scoped on the server). */
export class ChatIdDto {
@@ -37,6 +43,24 @@ export class GetChatMessagesDto {
cursor?: string;
}
/**
* Delta poll (#491): pull the chat's rows changed since `cursor` (a DB-clock
* timestamp from the previous poll) plus the current run fact — the degraded-poll
* fallback's payload, replacing the full infinite-query refetch. Omit `cursor` on
* the first poll (returns just a fresh cursor to start the chain).
*/
export class GetChatDeltaDto {
@IsString()
chatId: string;
// ISO-8601 timestamp echoed from the previous poll's response. Validated as
// ISO-8601 (not a bare string): a malformed cursor would otherwise reach the
// `::timestamptz` cast in findByChatUpdatedAfter and 500 instead of a clean 400.
@IsOptional()
@IsISO8601()
cursor?: string;
}
/** Resolve the chat bound to a document (the page's most-recent owned chat). */
export class BoundChatDto {
@IsString()
apps/server/src/core/ai-chat/external-mcp/dto/create-mcp-server.dto.ts
+4 -1
View File
@@ -37,10 +37,13 @@ export class CreateMcpServerDto {
@IsObject()
headers?: Record<string, string>;
// Omit/null => no restriction; `[]` is persisted verbatim and means deny-all
// (zero tools) since #476. @IsOptional() skips validation for null as well,
// so an explicit null is accepted.
@IsOptional()
@IsArray()
@IsString({ each: true })
toolAllowlist?: string[];
toolAllowlist?: string[] | null;
// Admin-authored guidance ("how/when to use this server's tools") injected
// into the agent system prompt next to the tool descriptions (#180). Trusted,
apps/server/src/core/ai-chat/external-mcp/dto/update-mcp-server.dto.ts
+4 -1
View File
@@ -38,10 +38,13 @@ export class UpdateMcpServerDto {
@IsObject()
headers?: Record<string, string>;
// Absent => unchanged; null => no restriction; `[]` is persisted verbatim
// and means deny-all (zero tools) since #476. @IsOptional() skips validation
// for null as well, so an explicit null is accepted.
@IsOptional()
@IsArray()
@IsString({ each: true })
toolAllowlist?: string[];
toolAllowlist?: string[] | null;
// Admin-authored prompt guidance (#180). Absent => unchanged; blank => cleared
// (stored as null by the repo). Capped to bound prompt/token size.
apps/server/src/core/ai-chat/external-mcp/mcp-allowlist-filter.spec.ts
+169
View File
@@ -0,0 +1,169 @@
import { type Tool } from 'ai';
import { McpClientsService } from './mcp-clients.service';
/**
* Tool-allowlist filtering semantics on the merged external toolset (#476).
*
* COVERAGE CHOICE (documented per issue #476): the full corrupt-row chain
* (DB value -> repo normalizeRow -> toolsFor filter) is covered on TWO levels
* instead of one live-stub-MCP-server integration test:
* (a) apps/server/test/integration/ai-mcp-server-repo.int-spec.ts pins the
* repo read/write semantics against a real Postgres — `[]` round-trips
* as jsonb `[]`, a present-but-corrupt value fails CLOSED to `[]` with
* an error log;
* (b) THIS spec pins what the toolset builder does with the repo's output —
* null = unrestricted, `['alpha']` = only alpha, `[]` (including the
* corrupt-row fallback) = ZERO tools.
* Together they prove the end-to-end property "corrupt/empty allowlist can
* never widen to all tools" without a live stub HTTP MCP server.
*
* The drive path mirrors mcp-namespacing.spec.ts: stub the repo's listEnabled,
* spy the private `connect` to return a fake client, inspect the merged keys.
*/
function fakeTool(): Tool {
return { description: 'x', inputSchema: undefined } as unknown as Tool;
}
interface FakeServer {
id: string;
name: string;
transport: string;
url: string;
headersEnc: string | null;
toolAllowlist: string[] | null;
}
function server(
over: Partial<FakeServer> & { id: string; name: string },
): FakeServer {
return {
transport: 'http',
url: 'https://example.com/mcp',
headersEnc: null,
toolAllowlist: null,
...over,
};
}
/**
* Build a service whose repo returns `servers` and whose fake clients expose
* `rawTools` from tools(). Returns the merged tool keys produced by toolsFor.
*/
async function mergedKeysFor(
servers: FakeServer[],
rawTools: Record<string, Tool>,
): Promise<string[]> {
const repoStub = {
listEnabled: jest.fn().mockResolvedValue(servers),
};
const service = new McpClientsService(repoStub as never, {} as never);
jest
.spyOn(
service as unknown as { connect: (s: FakeServer) => unknown },
'connect',
)
.mockImplementation(() =>
Promise.resolve({
tools: () => Promise.resolve(rawTools),
close: () => Promise.resolve(),
}),
);
const toolset = await service.toolsFor('ws-1');
// Release the lease so the service does not hold the fake clients open.
await Promise.all(toolset.clients.map((c) => c.close()));
return Object.keys(toolset.tools);
}
describe('external MCP tool-allowlist filtering (via toolsFor, #476)', () => {
afterEach(() => jest.restoreAllMocks());
const RAW = () => ({
alpha: fakeTool(),
beta: fakeTool(),
gamma: fakeTool(),
});
it("['alpha'] lets ONLY alpha through", async () => {
const keys = await mergedKeysFor(
[server({ id: 'id-1', name: 'srv', toolAllowlist: ['alpha'] })],
RAW(),
);
expect(keys).toEqual(['srv_alpha']);
});
it('null (no restriction) lets every tool through', async () => {
const keys = await mergedKeysFor(
[server({ id: 'id-1', name: 'srv', toolAllowlist: null })],
RAW(),
);
expect(keys.sort()).toEqual(['srv_alpha', 'srv_beta', 'srv_gamma']);
});
it('[] (deny-all) yields ZERO tools — an empty array is authoritative, not falsy (#476)', async () => {
// This is the regression the #476 change guards: `[]` used to fall through
// the old `allow.length > 0` check and expose ALL tools. It must expose NONE.
const keys = await mergedKeysFor(
[server({ id: 'id-1', name: 'srv', toolAllowlist: [] })],
RAW(),
);
expect(keys).toEqual([]);
});
it('the corrupt-row fallback ([] from the repo) also yields ZERO tools (#476)', async () => {
// The repo turns a present-but-corrupt tool_allowlist into `[]` (fail-closed,
// see normalizeRow in ai-mcp-server.repo.ts + the int-spec); this pins that
// the toolset builder honours that fallback as deny-all rather than allow-all.
const corruptFallback: string[] = [];
const keys = await mergedKeysFor(
[server({ id: 'id-1', name: 'srv', toolAllowlist: corruptFallback })],
RAW(),
);
expect(keys).toEqual([]);
});
it('allowlisted names not exposed by the server are ignored (no phantom tools)', async () => {
const keys = await mergedKeysFor(
[
server({
id: 'id-1',
name: 'srv',
toolAllowlist: ['alpha', 'does-not-exist'],
}),
],
RAW(),
);
expect(keys).toEqual(['srv_alpha']);
});
it('a deny-all server contributes no prompt instructions (0 tools merged)', async () => {
const repoStub = {
listEnabled: jest.fn().mockResolvedValue([
{
...server({ id: 'id-1', name: 'srv', toolAllowlist: [] }),
instructions: 'use the tools wisely',
},
]),
};
const service = new McpClientsService(repoStub as never, {} as never);
jest
.spyOn(
service as unknown as { connect: (s: FakeServer) => unknown },
'connect',
)
.mockImplementation(() =>
Promise.resolve({
tools: () => Promise.resolve(RAW()),
close: () => Promise.resolve(),
}),
);
const toolset = await service.toolsFor('ws-1');
await Promise.all(toolset.clients.map((c) => c.close()));
expect(Object.keys(toolset.tools)).toEqual([]);
// mergeNamespaced reported 0 contributed tools, so no guidance is attached.
expect(toolset.instructions).toEqual([]);
});
});
apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.ts
+6 -2
View File
@@ -444,9 +444,13 @@ export class McpClientsService {
try {
client = await this.connectWithTimeout(server, CONNECT_TIMEOUT_MS);
const raw = await withTimeout(client.tools(), CONNECT_TIMEOUT_MS);
// Allowlist semantics (#476): null/absent = no restriction (all tools);
// ANY array — including `[]` — is authoritative, so an EMPTY allowlist
// yields ZERO tools (deny-all). Do NOT add a `.length > 0` escape here:
// that read `[]` as falsy and silently widened deny-all to allow-all
// (the repo also fails corrupt rows closed to `[]` for the same reason).
const allow = server.toolAllowlist;
const picked =
Array.isArray(allow) && allow.length > 0 ? pick(raw, allow) : raw;
const picked = Array.isArray(allow) ? pick(raw, allow) : raw;
// Bound each tool's execute with a per-call total-timeout guard before
// merging, so a single chatty-but-stuck call is aborted after the cap.
const guarded = wrapToolsWithCallTimeout(picked, callTimeoutMs);
apps/server/src/core/ai-chat/external-mcp/mcp-servers.service.ts
+2 -1
View File
@@ -100,7 +100,8 @@ export class McpServersService {
transport: dto.transport,
url: dto.url,
headersEnc,
// undefined => unchanged; [] / value handled by repo (empty => null).
// undefined => unchanged; null => no restriction; `[]` is persisted
// verbatim and means deny-all (#476).
toolAllowlist: dto.toolAllowlist,
// undefined => unchanged; blank => cleared (null) by the repo.
instructions: dto.instructions,
apps/server/src/core/ai-chat/history-budget.spec.ts
+266
View File
@@ -0,0 +1,266 @@
import type { ModelMessage } from 'ai';
import {
resolveReplayBudget,
isContextOverflowError,
estimateMessagesTokens,
trimHistoryForReplay,
REPLAY_BUDGET_DEFAULT_TOKENS,
REPLAY_TRUNCATION_MARKER,
REPLAY_TURN_COLLAPSED_MARKER,
} from './history-budget';
describe('resolveReplayBudget', () => {
it('uses floor(0.7 x window) for a configured window (no cap)', () => {
// 0.7 x 60k = 42k
expect(resolveReplayBudget(60_000)).toEqual({
thresholdTokens: 42_000,
usedDefault: false,
});
// 0.7 x 1M = 700k — NOT capped (anti-brick vs the window, not a cost limiter).
expect(resolveReplayBudget(1_000_000)).toEqual({
thresholdTokens: 700_000,
usedDefault: false,
});
});
it('accepts the raw ::text stored form', () => {
expect(resolveReplayBudget('60000').thresholdTokens).toBe(42_000);
});
// The crux (#490): a chat with NO context window configured must STILL be
// budgeted — those are exactly the installs that hit terminal overflow.
it('applies the flat default when the window is unset/empty', () => {
expect(resolveReplayBudget(undefined)).toEqual({
thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS,
usedDefault: true,
});
expect(resolveReplayBudget('')).toEqual({
thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS,
usedDefault: true,
});
expect(resolveReplayBudget(' ')).toEqual({
thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS,
usedDefault: true,
});
});
it('treats an explicit 0 as the off-switch (distinct from unset)', () => {
expect(resolveReplayBudget(0)).toEqual({
thresholdTokens: null,
usedDefault: false,
});
expect(resolveReplayBudget('0')).toEqual({
thresholdTokens: null,
usedDefault: false,
});
});
it('falls back to the default on a negative/garbage value', () => {
expect(resolveReplayBudget(-5).usedDefault).toBe(true);
expect(resolveReplayBudget('abc').usedDefault).toBe(true);
});
});
describe('isContextOverflowError', () => {
it('classifies a real provider 400 context-overflow shape', () => {
// OpenAI-compatible shape.
expect(
isContextOverflowError({
statusCode: 400,
message:
"This model's maximum context length is 128000 tokens. However, your messages resulted in 214000 tokens. Please reduce the length of the messages.",
}),
).toBe(true);
// Anthropic-style wording.
expect(
isContextOverflowError({
status: 400,
message: 'prompt is too long: 250000 tokens > 200000 maximum',
}),
).toBe(true);
// Nested body + string status.
expect(
isContextOverflowError({
response: { status: '400' },
message: 'input is too long for the requested model',
}),
).toBe(true);
// Error instance with the cause carrying the body.
const e = new Error('Bad request');
(e as any).statusCode = 400;
(e as any).cause = new Error('maximum context window exceeded');
expect(isContextOverflowError(e)).toBe(true);
});
it('does NOT classify unrelated 400s or auth/rate-limit errors', () => {
expect(
isContextOverflowError({ statusCode: 400, message: 'invalid tool schema' }),
).toBe(false);
expect(
isContextOverflowError({
statusCode: 429,
message: 'context length exceeded but rate limited',
}),
).toBe(false);
expect(isContextOverflowError({ statusCode: 500, message: 'server error' })).toBe(
false,
);
expect(isContextOverflowError(undefined)).toBe(false);
expect(isContextOverflowError('some random string')).toBe(false);
});
});
// Helpers to build ModelMessage fixtures in the ai@6 shape.
const userMsg = (text: string): ModelMessage =>
({ role: 'user', content: [{ type: 'text', text }] }) as ModelMessage;
const assistantMsg = (
text: string,
toolCallId?: string,
toolName?: string,
): ModelMessage =>
({
role: 'assistant',
content: [
{ type: 'text', text },
...(toolCallId
? [{ type: 'tool-call', toolCallId, toolName, input: {} }]
: []),
],
}) as ModelMessage;
const toolMsg = (
toolCallId: string,
toolName: string,
value: unknown,
): ModelMessage =>
({
role: 'tool',
content: [
{ type: 'tool-result', toolCallId, toolName, output: { type: 'json', value } },
],
}) as ModelMessage;
describe('trimHistoryForReplay', () => {
it('null budget disables trimming (returns the same reference)', () => {
const msgs = [userMsg('hi'), assistantMsg('yo')];
const r = trimHistoryForReplay(msgs, null);
expect(r.trimmed).toBe(false);
expect(r.messages).toBe(msgs);
});
it('leaves history under budget untouched (same reference)', () => {
const msgs = [userMsg('hi'), assistantMsg('a short answer')];
const r = trimHistoryForReplay(msgs, 100_000);
expect(r.trimmed).toBe(false);
expect(r.messages).toBe(msgs);
});
it('truncates OLD tool outputs but keeps recent turns full', () => {
const big = 'X'.repeat(40_000); // ~16k tokens on its own
const msgs: ModelMessage[] = [];
// 6 OLD turns (indices 0..5), each with a huge tool output.
for (let i = 0; i < 6; i++) {
msgs.push(userMsg(`old q${i}`));
msgs.push(assistantMsg('looking', `c${i}`, 'getPage'));
msgs.push(toolMsg(`c${i}`, 'getPage', { body: big }));
msgs.push(assistantMsg(`old a${i}`));
}
// 3 small recent turns, then the CURRENT turn with its own huge tool output.
// With REPLAY_KEEP_RECENT_TURNS=4 the last 4 user-turns stay full, so only
// these small recent turns + the current big one are kept full; the 6 old
// turns above fall in the trim region.
for (let i = 0; i < 3; i++) {
msgs.push(userMsg(`recent q${i}`));
msgs.push(assistantMsg(`recent a${i}`));
}
msgs.push(userMsg('current q'));
msgs.push(assistantMsg('looking', 'cR', 'getPage'));
msgs.push(toolMsg('cR', 'getPage', { body: big }));
msgs.push(assistantMsg('current a'));
// Budget large enough that phase-1 tool truncation alone brings it under.
const r = trimHistoryForReplay(msgs, 30_000);
expect(r.trimmed).toBe(true);
const flat = JSON.stringify(r.messages);
// The CURRENT turn's tool output survives in full.
expect(flat).toContain(big);
// Old outputs were truncated with the marker.
expect(flat).toContain(REPLAY_TRUNCATION_MARKER);
// Phase 1 sufficed: the oldest turns were NOT collapsed.
expect(flat).not.toContain(REPLAY_TURN_COLLAPSED_MARKER);
expect(estimateMessagesTokens(r.messages)).toBeLessThan(
estimateMessagesTokens(msgs),
);
});
it('collapses the oldest turns when tool truncation is not enough', () => {
// Many turns with LARGE assistant TEXT (not tool output) so phase 1 can't help.
const bigText = 'слово '.repeat(8_000); // large Cyrillic text per turn
const msgs: ModelMessage[] = [];
for (let i = 0; i < 12; i++) {
msgs.push(userMsg(`q${i}`));
msgs.push(assistantMsg(bigText));
}
const r = trimHistoryForReplay(msgs, 30_000);
expect(r.trimmed).toBe(true);
// Oldest turns collapsed; result fits (best-effort) and is much smaller.
expect(estimateMessagesTokens(r.messages)).toBeLessThan(
estimateMessagesTokens(msgs),
);
// The LAST turn's text is preserved in full (recent turns stay full).
expect(JSON.stringify(r.messages[r.messages.length - 1])).toContain(bigText);
});
it('is deterministic / byte-stable for identical inputs', () => {
const big = 'Y'.repeat(30_000);
const build = (): ModelMessage[] => {
const m: ModelMessage[] = [];
for (let i = 0; i < 10; i++) {
m.push(userMsg(`q${i}`));
m.push(assistantMsg('t', `c${i}`, 'getPage'));
m.push(toolMsg(`c${i}`, 'getPage', { body: big }));
}
return m;
};
const a = trimHistoryForReplay(build(), 15_000);
const b = trimHistoryForReplay(build(), 15_000);
expect(JSON.stringify(a.messages)).toBe(JSON.stringify(b.messages));
});
it('never leaves an unpaired tool-call after collapsing (balanced history)', () => {
const big = 'Z'.repeat(40_000);
const msgs: ModelMessage[] = [];
for (let i = 0; i < 10; i++) {
msgs.push(userMsg(`q${i}`));
msgs.push(assistantMsg('t', `c${i}`, 'getPage'));
msgs.push(toolMsg(`c${i}`, 'getPage', { body: big }));
}
const r = trimHistoryForReplay(msgs, 8_000);
// Count tool-call vs tool-result parts in the trimmed output.
let calls = 0;
let results = 0;
for (const m of r.messages) {
if (!Array.isArray(m.content)) continue;
for (const p of m.content as Array<{ type?: string }>) {
if (p.type === 'tool-call') calls++;
if (p.type === 'tool-result' || p.type === 'tool-error') results++;
}
}
// Every surviving tool-call has a surviving result (collapsing drops BOTH).
expect(calls).toBe(results);
// Collapsed turns carry the marker.
expect(JSON.stringify(r.messages)).toContain(REPLAY_TURN_COLLAPSED_MARKER);
});
it('respects the provider fact: under-budget contextTokens skips trimming', () => {
const big = 'W'.repeat(60_000);
const msgs = [
userMsg('q'),
assistantMsg('t', 'c1', 'getPage'),
toolMsg('c1', 'getPage', { body: big }),
];
// char-estimate is high, but the provider says we are well under budget.
const r = trimHistoryForReplay(msgs, 100_000, 5_000);
expect(r.trimmed).toBe(false);
expect(r.messages).toBe(msgs);
});
});
apps/server/src/core/ai-chat/history-budget.ts
+375
View File
@@ -0,0 +1,375 @@
/**
* History-replay token budget (#490).
*
* The whole persisted conversation is replayed to the provider on EVERY turn, so
* a long chat eventually exceeds the model's context window and the provider 400s
* on every turn — terminally (the chat "bricks"). This module bounds the replayed
* history at REPLAY TIME only: it never mutates what is persisted (the DB stays
* the full record), and its output is a deterministic, byte-stable function of its
* input so the trimmed prefix is identical turn to turn (provider prompt-cache
* friendliness — real money on long chats).
*
* The PRIMARY signal is the provider's own fact: `metadata.contextTokens` from the
* last turn. The chars-based {@link estimateTokens} (shared with the client) is
* used only for the DELTA of not-yet-sent messages, to decide WHAT to trim, and as
* the fallback for chats with no usage yet.
*/
import type { ModelMessage } from 'ai';
import { estimateTokens } from '@docmost/token-estimate';
/** Flat default budget when no context window is configured (tokens). */
export const REPLAY_BUDGET_DEFAULT_TOKENS = 100_000;
/** Fraction of a configured context window used as the budget. */
export const REPLAY_BUDGET_WINDOW_FRACTION = 0.7;
/**
* Fraction of the normal budget used for the REACTIVE re-trim after a provider
* context-overflow 400 — the preventive estimate under-counted, so cut harder.
*/
export const REPLAY_AGGRESSIVE_FRACTION = 0.5;
/**
* Turns (a user message + its assistant/tool replies) kept FULL at the tail,
* including the current one — never trimmed. Older turns are compacted first.
*/
export const REPLAY_KEEP_RECENT_TURNS = 4;
/** Leading chars kept from a truncated old tool output. */
export const REPLAY_TOOL_OUTPUT_HEAD = 800;
/** Trailing chars kept from a truncated old tool output. */
export const REPLAY_TOOL_OUTPUT_TAIL = 300;
/** Marker inserted where an old tool output was truncated for replay. */
export const REPLAY_TRUNCATION_MARKER =
'[…truncated for replay; call the tool again to read the full output]';
/** Marker for a whole old turn collapsed to its text. */
export const REPLAY_TURN_COLLAPSED_MARKER =
'[earlier tool activity omitted for replay]';
export interface ReplayBudget {
/** Token threshold above which replay history is trimmed; `null` = OFF. */
thresholdTokens: number | null;
/** True when the flat default was used (no context window configured). */
usedDefault: boolean;
}
/**
* Resolve the replay budget from the RAW stored `chatContextWindow` (text/number).
* - a positive value -> `floor(fraction × window)` (NO cap — the budgeter is
* anti-brick protection against the window itself, not a cost/economy limiter,
* exactly as the codebase already treats maxOutputTokens; the reactive branch
* still guarantees anti-brick regardless of how high this budget is)
* - explicit `0` -> OFF (admin opt-out; `null` threshold)
* - unset/empty/invalid-> the flat default (still protects — the installations
* that hit terminal overflow are exactly the ones that never set a window)
*
* Note the raw value is needed because the parsed `chatContextWindow` collapses
* both `0` and unset to `undefined`, which would erase the explicit off-switch.
*/
export function resolveReplayBudget(rawContextWindow: unknown): ReplayBudget {
let n: number | undefined;
if (typeof rawContextWindow === 'number') {
n = rawContextWindow;
} else if (typeof rawContextWindow === 'string') {
const t = rawContextWindow.trim();
n = t === '' ? undefined : Number(t);
}
// Unset / empty / non-numeric / negative -> flat default (the protective case).
if (n === undefined || !Number.isFinite(n) || n < 0) {
return { thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS, usedDefault: true };
}
// Explicit 0 -> off-switch.
if (n === 0) {
return { thresholdTokens: null, usedDefault: false };
}
return {
thresholdTokens: Math.floor(REPLAY_BUDGET_WINDOW_FRACTION * n),
usedDefault: false,
};
}
/**
* The effective replay threshold for THIS turn, given the base budget and whether
* the PREVIOUS turn hit a context-overflow 400 (the reactive-recovery signal,
* `metadata.replayOverflow`). On recovery the base budget is scaled down by
* {@link REPLAY_AGGRESSIVE_FRACTION}: the overflowing turn produced no usage
* signal, so the preventive estimate under-counted and a normal-threshold trim may
* not shrink enough to fit — this harder cut is what un-bricks the chat.
*
* A `null` base budget (trimming OFF) is passed through unchanged: an explicit
* off-switch is never overridden by the recovery path.
*/
export function resolveEffectiveReplayThreshold(
thresholdTokens: number | null,
priorOverflowed: boolean,
): number | null {
if (!priorOverflowed || thresholdTokens == null) return thresholdTokens;
return Math.floor(thresholdTokens * REPLAY_AGGRESSIVE_FRACTION);
}
/**
* True when a provider error is a CONTEXT-OVERFLOW rejection (the prompt exceeds
* the model's window). Providers surface this as an HTTP 400 with a recognizable
* message; match both the status and the message patterns robustly across
* OpenAI-compatible / Anthropic / Gemini wordings, since the exact shape varies.
*/
export function isContextOverflowError(error: unknown): boolean {
const status = extractStatus(error);
const msg = extractMessage(error).toLowerCase();
// Message patterns seen across providers for "prompt too long".
const overflowPattern =
/context (?:length|window)|maximum context|too many tokens|too large for|reduce the length|prompt is too long|input (?:is )?too long|exceeds? the (?:maximum )?(?:context|token)|maximum.*tokens|string too long/;
if (!overflowPattern.test(msg)) return false;
// A 400/413 with an overflow-shaped message is an overflow. Some providers
// omit/rewrite the status, so accept the message match when the status is
// unknown, but reject it for auth/rate-limit statuses that never mean overflow.
if (status === 400 || status === 413) return true;
if (status === 401 || status === 403 || status === 429) return false;
return true;
}
function extractStatus(error: unknown): number | undefined {
if (!error || typeof error !== 'object') return undefined;
const e = error as Record<string, unknown>;
for (const k of ['statusCode', 'status']) {
const v = e[k];
if (typeof v === 'number') return v;
if (typeof v === 'string' && /^\d+$/.test(v)) return Number(v);
}
// Nested (e.g. { response: { status } } / { cause: { statusCode } }).
for (const k of ['response', 'cause', 'data']) {
const nested = e[k];
if (nested && typeof nested === 'object') {
const s = extractStatus(nested);
if (s !== undefined) return s;
}
}
return undefined;
}
function extractMessage(error: unknown): string {
if (error == null) return '';
if (typeof error === 'string') return error;
if (error instanceof Error) {
// Include nested causes (provider libs wrap the real body in `cause`).
const cause = (error as { cause?: unknown }).cause;
return `${error.message} ${cause ? extractMessage(cause) : ''}`;
}
if (typeof error === 'object') {
const e = error as Record<string, unknown>;
const parts: string[] = [];
for (const k of ['message', 'error', 'body', 'responseBody', 'data']) {
const v = e[k];
if (typeof v === 'string') parts.push(v);
else if (v && typeof v === 'object') parts.push(extractMessage(v));
}
return parts.join(' ');
}
return String(error);
}
/** Rough token size of a ModelMessage array via the shared chars estimator. */
export function estimateMessagesTokens(
messages: ReadonlyArray<ModelMessage>,
): number {
let total = 0;
for (const m of messages) {
total += estimateTokens(serializeContent(m.content));
}
return total;
}
function serializeContent(content: unknown): string {
if (typeof content === 'string') return content;
try {
return JSON.stringify(content) ?? '';
} catch {
return '';
}
}
/** Deep JSON string of an arbitrary value, bounded so estimation never throws. */
function stringifyValue(value: unknown): string {
if (typeof value === 'string') return value;
try {
return JSON.stringify(value) ?? String(value);
} catch {
return String(value);
}
}
export interface TrimResult {
messages: ModelMessage[];
/** Whether any trimming was applied. */
trimmed: boolean;
/** Estimated tokens of the returned messages (chars-based). */
estimatedTokens: number;
}
/**
* Bound the replayed history to `budgetTokens`, deterministically. Returns the
* SAME array reference (no copy) when nothing needs trimming, so the common case
* is free and byte-identical. Trimming order (spec #490):
* 1. truncate OLD turns' tool outputs (head+tail + marker) — the bulk of the size
* 2. mechanically collapse the OLDEST turns to their text (concatenation, no LLM)
* 3. the current + last {@link REPLAY_KEEP_RECENT_TURNS} turns stay FULL
*
* `budgetTokens === null` disables trimming. `priorContextTokens` (the provider's
* fact from last turn) short-circuits the decision: when it is known and already
* under budget we skip trimming even if the char-estimate is higher (the provider
* count is authoritative). The char-estimate drives WHAT to cut.
*/
export function trimHistoryForReplay(
messages: ModelMessage[],
budgetTokens: number | null,
priorContextTokens?: number,
): TrimResult {
if (budgetTokens == null) {
return { messages, trimmed: false, estimatedTokens: 0 };
}
const estimated = estimateMessagesTokens(messages);
// Decision signal: prefer the provider's fact (last turn's contextTokens) plus
// the estimated delta of the messages appended since; fall back to the pure
// char-estimate for a chat with no usage yet.
const projected =
priorContextTokens != null
? Math.max(priorContextTokens, estimated)
: estimated;
if (projected <= budgetTokens) {
return { messages, trimmed: false, estimatedTokens: estimated };
}
// The tail we always keep full: from the Nth-from-last user message onward.
const boundary = recentBoundaryIndex(messages, REPLAY_KEEP_RECENT_TURNS);
const tail = messages.slice(boundary);
let head = messages.slice(0, boundary).map(cloneMessage);
// Phase 1: truncate old tool outputs.
for (const m of head) {
if (m.role === 'tool') truncateToolMessage(m);
}
let out = [...head, ...tail];
let est = estimateMessagesTokens(out);
if (est <= budgetTokens) {
return { messages: out, trimmed: true, estimatedTokens: est };
}
// Phase 2: collapse the oldest turns (in `head`) to their text, one at a time,
// from the oldest, until we fit or the whole head is collapsed.
const turns = splitTurns(head);
const collapsed: ModelMessage[] = [];
let i = 0;
for (; i < turns.length; i++) {
if (est <= budgetTokens) break;
collapsed.push(...collapseTurn(turns[i]));
// Re-estimate the whole prospective output.
const remaining = turns.slice(i + 1).flat();
out = [...collapsed, ...remaining, ...tail];
est = estimateMessagesTokens(out);
}
// Include any turns we didn't need to collapse.
const remaining = turns.slice(i).flat();
out = [...collapsed, ...remaining, ...tail];
est = estimateMessagesTokens(out);
return { messages: out, trimmed: true, estimatedTokens: est };
}
/** Index of the first message of the Nth-from-last user turn (0 if fewer). */
function recentBoundaryIndex(
messages: ReadonlyArray<ModelMessage>,
keepTurns: number,
): number {
const userIdx: number[] = [];
for (let i = 0; i < messages.length; i++) {
if (messages[i].role === 'user') userIdx.push(i);
}
if (userIdx.length <= keepTurns) return 0;
return userIdx[userIdx.length - keepTurns];
}
/** Split a message list into turns; each turn starts at a `user` message. */
function splitTurns(messages: ModelMessage[]): ModelMessage[][] {
const turns: ModelMessage[][] = [];
for (const m of messages) {
if (m.role === 'user' || turns.length === 0) turns.push([m]);
else turns[turns.length - 1].push(m);
}
return turns;
}
/**
* Collapse a whole turn to its plain text (mechanical concatenation, not an LLM
* summary). Keeps the user message; replaces the assistant/tool messages with a
* single assistant text message = the assistant's concatenated text + a marker
* when tool activity was dropped. Dropping BOTH the tool-call and tool-result
* parts together keeps the rebuilt history balanced (no unpaired calls).
*/
function collapseTurn(turn: ModelMessage[]): ModelMessage[] {
const out: ModelMessage[] = [];
let assistantText = '';
let hadTools = false;
for (const m of turn) {
if (m.role === 'user') {
out.push(m);
} else if (m.role === 'assistant') {
const { text, tools } = extractAssistantText(m.content);
assistantText += text;
hadTools = hadTools || tools;
} else if (m.role === 'tool') {
hadTools = true;
} else {
out.push(m);
}
}
const note =
(assistantText ? assistantText.trimEnd() : '') +
(hadTools
? `${assistantText ? '\n\n' : ''}${REPLAY_TURN_COLLAPSED_MARKER}`
: '');
if (note) out.push({ role: 'assistant', content: note } as ModelMessage);
return out;
}
function extractAssistantText(content: unknown): {
text: string;
tools: boolean;
} {
if (typeof content === 'string') return { text: content, tools: false };
if (!Array.isArray(content)) return { text: '', tools: false };
let text = '';
let tools = false;
for (const part of content) {
const type = (part as { type?: string })?.type;
if (type === 'text') text += (part as { text?: string }).text ?? '';
else if (type === 'tool-call') tools = true;
}
return { text, tools };
}
/** Truncate every tool-result output in a `tool` message to head+tail+marker. */
function truncateToolMessage(message: ModelMessage): void {
const content = message.content;
if (!Array.isArray(content)) return;
for (const part of content) {
const p = part as { type?: string; output?: { type?: string; value?: unknown } };
if (p.type !== 'tool-result' && p.type !== 'tool-error') continue;
if (!p.output) continue;
const raw = stringifyValue(p.output.value);
const budget = REPLAY_TOOL_OUTPUT_HEAD + REPLAY_TOOL_OUTPUT_TAIL;
if (raw.length <= budget + REPLAY_TRUNCATION_MARKER.length) continue;
const truncated =
raw.slice(0, REPLAY_TOOL_OUTPUT_HEAD) +
`\n${REPLAY_TRUNCATION_MARKER}\n` +
raw.slice(raw.length - REPLAY_TOOL_OUTPUT_TAIL);
// Represent the shrunk output as a text output (a valid tool-result output).
p.output = { type: 'text', value: truncated };
}
}
/** Shallow-ish clone so trimming never mutates the caller's (persisted-derived)
* message objects — only the OLD region is cloned before it is edited. */
function cloneMessage(m: ModelMessage): ModelMessage {
if (typeof m.content === 'string') return { ...m };
return {
...m,
content: (m.content as unknown[]).map((p) =>
p && typeof p === 'object' ? { ...(p as object) } : p,
),
} as ModelMessage;
}
apps/server/src/core/ai-chat/output-degeneration.spec.ts
+64
View File
@@ -9,9 +9,73 @@ import {
DEGENERATION_CHECK_STEP,
REPEATED_LINES_THRESHOLD,
MIN_PERIOD_REPEATS,
degenerationThresholds,
} from './output-degeneration';
import { AiChatService } from './ai-chat.service';
// Part A (#495 iter10): the detector thresholds are env-tunable. These drive the
// resolver against real repeat-count shapes and mutation-verify that the env
// override actually changes the trigger point (not a vacuous read).
describe('degeneration thresholds are env-configurable', () => {
const VARS = [
'AI_CHAT_DEGENERATION_REPEATED_LINES',
'AI_CHAT_DEGENERATION_PERIOD_MAX_LEN',
'AI_CHAT_DEGENERATION_PERIOD_MIN_REPEATS',
'AI_CHAT_DEGENERATION_CHECK_STEP',
];
const saved: Record<string, string | undefined> = {};
beforeEach(() => {
for (const v of VARS) saved[v] = process.env[v];
});
afterEach(() => {
for (const v of VARS) {
if (saved[v] === undefined) delete process.env[v];
else process.env[v] = saved[v];
}
});
it('defaults to the compiled constants when unset', () => {
for (const v of VARS) delete process.env[v];
expect(degenerationThresholds()).toEqual({
repeatedLines: REPEATED_LINES_THRESHOLD,
maxPeriodLen: 150,
minPeriodRepeats: MIN_PERIOD_REPEATS,
checkStep: DEGENERATION_CHECK_STEP,
});
});
it('falls back to the default on blank / invalid / non-positive values', () => {
for (const bad of ['', ' ', 'abc', '0', '-3', '1.5']) {
process.env.AI_CHAT_DEGENERATION_REPEATED_LINES = bad;
// '1.5' floors to 1 (still ≥1, valid); every other bad value → default.
const expected = bad === '1.5' ? 1 : REPEATED_LINES_THRESHOLD;
expect(degenerationThresholds().repeatedLines).toBe(expected);
}
});
it('a RAISED check-step suppresses a burst the default would have flagged', () => {
// A ~3.3KB periodic burst is periodic-degenerate, but shouldCheckDegeneration
// is the throttle gate. Default checkStep=2000 arms on it; raising the step
// above the burst size means the throttle never re-fires for it.
const burstLen = 'loadTools.\n'.repeat(300).length; // ~3300
delete process.env.AI_CHAT_DEGENERATION_CHECK_STEP;
expect(shouldCheckDegeneration(burstLen, 0)).toBe(true); // default 2000
process.env.AI_CHAT_DEGENERATION_CHECK_STEP = String(burstLen + 1);
expect(shouldCheckDegeneration(burstLen, 0)).toBe(false); // raised gate
});
it('a LOWERED repeated-lines threshold trips on a shorter identical-line run', () => {
// 8 identical lines: below the default 25 (rule 1) and below the periodic
// rule's 20 repeats — so isDegenerateOutput is false by default.
const shortRun = 'x\n'.repeat(8);
delete process.env.AI_CHAT_DEGENERATION_REPEATED_LINES;
expect(isDegenerateOutput(shortRun)).toBe(false);
// Lower rule 1 to 5 → the 8-line run now trips.
process.env.AI_CHAT_DEGENERATION_REPEATED_LINES = '5';
expect(isDegenerateOutput(shortRun)).toBe(true);
});
});
// Mock ONLY streamText so we can capture the onChunk/onStepFinish callbacks the
// service registers and drive them by hand; every other `ai` export the service
// uses (convertToModelMessages, stepCountIs, …) stays real.
apps/server/src/core/ai-chat/output-degeneration.ts
+54 -2
View File
@@ -23,6 +23,54 @@ export const MAX_PERIOD_LEN = 150;
/** Rule 2: minimum number of consecutive block repeats to trigger. */
export const MIN_PERIOD_REPEATS = 20;
/**
* Read a positive-integer threshold from an env var, falling back to `fallback`
* on unset/blank/invalid/non-positive. Mirrors the `AI_STREAM_PRE_RESPONSE_RETRIES`
* resolver in `ai-streaming-fetch.ts`: read the RAW string first so a blank value
* is treated as "unset" (→ fallback) rather than coercing to 0. Thresholds must
* stay ≥ 1 — a 0/negative would make the detector fire on any text (or never), so
* a bad value degrades to the safe compiled default instead. Env-tunable so an
* operator can retune the anti-babble guard (#444) without a redeploy, following
* the `AI_CHAT_FINAL_STEP_LOCKDOWN` toggle convention.
*/
function envThreshold(name: string, fallback: number): number {
const rawStr = process.env[name];
if (rawStr === undefined || rawStr.trim() === '') return fallback;
const raw = Number(rawStr);
return Number.isFinite(raw) && raw >= 1 ? Math.floor(raw) : fallback;
}
/**
* Resolve the degeneration-detector thresholds from the environment, each
* defaulting to the compiled constant above. Read fresh per call (not cached at
* import) so a test — or a runtime env change — takes effect deterministically.
*/
export function degenerationThresholds(): {
repeatedLines: number;
maxPeriodLen: number;
minPeriodRepeats: number;
checkStep: number;
} {
return {
repeatedLines: envThreshold(
'AI_CHAT_DEGENERATION_REPEATED_LINES',
REPEATED_LINES_THRESHOLD,
),
maxPeriodLen: envThreshold(
'AI_CHAT_DEGENERATION_PERIOD_MAX_LEN',
MAX_PERIOD_LEN,
),
minPeriodRepeats: envThreshold(
'AI_CHAT_DEGENERATION_PERIOD_MIN_REPEATS',
MIN_PERIOD_REPEATS,
),
checkStep: envThreshold(
'AI_CHAT_DEGENERATION_CHECK_STEP',
DEGENERATION_CHECK_STEP,
),
};
}
/**
* Rule 1 — ≥`REPEATED_LINES_THRESHOLD` consecutive IDENTICAL non-empty lines at
* the tail. Catches the classic newline-delimited loop ("loadTools.\n" ×N).
@@ -128,7 +176,11 @@ export function hasPeriodicTail(
* Pure — the caller owns the abort side effect.
*/
export function isDegenerateOutput(text: string): boolean {
return hasRepeatedLineRun(text) || hasPeriodicTail(text);
const cfg = degenerationThresholds();
return (
hasRepeatedLineRun(text, cfg.repeatedLines) ||
hasPeriodicTail(text, cfg.maxPeriodLen, cfg.minPeriodRepeats)
);
}
/**
@@ -154,7 +206,7 @@ export function shouldCheckDegeneration(
textLen: number,
lastCheckLen: number,
): boolean {
return textLen - lastCheckLen >= DEGENERATION_CHECK_STEP;
return textLen - lastCheckLen >= degenerationThresholds().checkStep;
}
/**
apps/server/src/core/ai-chat/public-share-chat.service.ts
+4
View File
@@ -307,6 +307,10 @@ export class PublicShareChatService {
system,
messages: modelMessages,
tools,
// Pin the AI SDK per-request retry budget explicitly (matches the SDK
// default of 2). Connection arithmetic: (1 + maxRetries) × (1 +
// AI_STREAM_PRE_RESPONSE_RETRIES) worst-case connects per turn.
maxRetries: 2,
// Bound the agent loop for anonymous callers.
stopWhen: stepCountIs(5),
// Cap per-request output so one anonymous call cannot run up the provider
apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts
+17 -26
View File
@@ -426,6 +426,7 @@ export class AiChatToolsService {
const {
sharedToolSpecs,
createCommentSignalTracker,
createListCommentsProbe,
searchShapes,
getGuideSection,
} = await loadDocmostMcp();
@@ -718,7 +719,11 @@ export class AiChatToolsService {
if (spec.mcpOnly) continue;
if (spec.inlineBothHosts) continue;
const run = spec.inAppExecute ?? spec.execute;
if (!run) continue; // defensive: a shared spec always carries one of them.
// Guaranteed present by assertEverySpecIsRegisterable() (#494), which runs
// at tool-specs module load and throws if a non-inline spec the in-app host
// registers carries neither inAppExecute nor execute — so this can no longer
// silently drop a mis-declared tool. Kept as a type-narrowing guard.
if (!run) continue;
tools[spec.inAppKey] = sharedTool(
spec,
(async (args) =>
@@ -764,35 +769,21 @@ export class AiChatToolsService {
// wrapper below) so the race governs the whole call. The client carries the
// per-call composite signal via setToolAbortSignal.
const capMs = inAppToolCallCapMs();
if (!createCommentSignalTracker) {
// The signal needs BOTH the tracker factory AND the shared count-source probe
// factory (#494). Either being absent (a stale @docmost/mcp build or a mocked
// loader) => signal disabled, tool results byte-identical.
if (!createCommentSignalTracker || !createListCommentsProbe) {
return wrapInAppToolsWithCap(tools, client, capMs);
}
// Shared probe (#494): the SAME factory the standalone MCP host uses, so the
// in-app probe body is no longer a hand-mirror that could drift (counting the
// full feed newer than the watermark, labelling a hit with the light page
// title). `client` supplies the loopback listComments/getPageRaw reads.
const tracker = createCommentSignalTracker({
probe: async (pageId: string, sinceMs: number) => {
const { items } = await client.listComments(pageId, true);
const count = (items as Array<{ createdAt?: string }>).filter((c) => {
const created = c?.createdAt ? new Date(c.createdAt).getTime() : NaN;
return Number.isFinite(created) && created > sinceMs;
}).length;
let title: string | undefined;
if (count > 0) {
// Title labels the signal; untrusted, defanged by the shared builder.
// Fetched only on a hit so the no-signal path never pays for it. Uses
// the LIGHT raw page info (title only) — mirroring the standalone MCP
// probe's getPageRaw — instead of the heavy getPage (which also renders
// Markdown + subpages) just to read one field.
try {
const res = (await client.getPageRaw(pageId)) as {
title?: string;
} | null;
title = res?.title ?? undefined;
} catch {
// Title is optional — omit it when the page can't be fetched.
}
}
return { count, title };
},
probe: createListCommentsProbe(
client as unknown as Parameters<typeof createListCommentsProbe>[0],
),
});
return wrapInAppToolsWithCap(
apps/server/src/core/ai-chat/tools/comment-signal-inapp.spec.ts
+9 -2
View File
@@ -21,7 +21,10 @@ import { SHARED_TOOL_SPECS } from '../../../../../../packages/mcp/src/tool-specs
// The REAL shared tracker factory, imported from source (same cross-boundary
// approach the tool-specs spec uses) so the in-app wiring is exercised against
// exactly the watermark/debounce/injection-safe logic the package ships.
import { createCommentSignalTracker } from '../../../../../../packages/mcp/src/comment-signal';
import {
createCommentSignalTracker,
createListCommentsProbe,
} from '../../../../../../packages/mcp/src/comment-signal';
// The REAL client-side citation extractor: proves that the passive signal does
// NOT strip a tool's citations (the #417 in-app regression this spec guards).
import { toolCitations } from '../../../../../../apps/client/src/features/ai-chat/utils/tool-parts';
@@ -284,9 +287,13 @@ describe('AiChatToolsService forUser + comment signal (real tracker)', () => {
return fakeClient as DocmostClientLike;
} as unknown as loader.DocmostClientCtor,
sharedToolSpecs: SHARED_TOOL_SPECS as unknown as Record<string, loader.SharedToolSpec>,
// Wire the REAL factory so the in-app path is exercised end to end.
// Wire the REAL factories so the in-app path is exercised end to end —
// including the shared count-source probe (#494) the service now builds the
// tracker's `probe` from.
createCommentSignalTracker:
createCommentSignalTracker as unknown as loader.CommentSignalTrackerFactory,
createListCommentsProbe:
createListCommentsProbe as unknown as loader.CreateListCommentsProbeFn,
// Pure no-network draw.io helpers (#424) — required on the loader return;
// this comment-signal test doesn't exercise them, so no-op stubs suffice.
searchShapes: (() => []) as unknown as loader.SearchShapesFn,
apps/server/src/core/ai-chat/tools/docmost-client.loader.ts
+29
View File
@@ -150,6 +150,27 @@ export type CommentSignalTrackerFactory = (options: {
debounceMs?: number;
}) => CommentSignalTrackerLike;
/**
* Local mirror of `@docmost/mcp`'s `createListCommentsProbe` (#494): the SHARED
* count-source probe both hosts use, so the in-app probe body is no longer a
* hand-copy of the standalone MCP one. Given a client with the light comment feed
* + raw-page-title reads, it returns the tracker's `probe` (count comments newer
* than the watermark, label a hit with the page title). Loosely typed at this
* cross-package boundary, like the rest of this loader.
*/
export type CreateListCommentsProbeFn = (client: {
listComments(
pageId: string,
includeResolved: boolean,
): Promise<{ items: Array<{ createdAt?: string | null }> }>;
getPageRaw(
pageId: string,
): Promise<{ title?: string | null } | null | undefined>;
}) => (
pageId: string,
sinceMs: number,
) => Promise<CommentSignalProbeResultLike>;
// Pure, no-network draw.io helpers (#424). These are plain functions on the
// module (NOT DocmostClient methods) — the in-app AI-SDK service calls them
// directly to wire drawioShapes / drawioGuide, mirroring the MCP server.
@@ -170,6 +191,10 @@ interface DocmostMcpModule {
// loader in unit tests. The in-app layer treats an absent factory as "signal
// disabled" — a pure no-op that leaves tool results byte-identical.
createCommentSignalTracker?: CommentSignalTrackerFactory;
// Optional (#494): the shared count-source probe factory. Absent on a pre-#494
// build or a mocked loader; the in-app layer only builds a probe when the
// signal factory above is also present.
createListCommentsProbe?: CreateListCommentsProbeFn;
// Optional (#447): a deterministic hash of the tool-specs registry content,
// generated into build/ by the package's build. Absent on a pre-#447 build (or
// the mocked loader in unit tests) — the stale-check below is a NO-OP when it
@@ -284,6 +309,7 @@ export async function loadDocmostMcp(): Promise<{
DocmostClient: DocmostClientCtor;
sharedToolSpecs: Record<string, SharedToolSpec>;
createCommentSignalTracker?: CommentSignalTrackerFactory;
createListCommentsProbe?: CreateListCommentsProbeFn;
searchShapes: SearchShapesFn;
getGuideSection: GetGuideSectionFn;
}> {
@@ -331,6 +357,9 @@ export async function loadDocmostMcp(): Promise<{
// Optional: forwarded when present so the in-app layer can build the passive
// comment signal (#417); undefined on a stale build => signal disabled.
createCommentSignalTracker: mod.createCommentSignalTracker,
// Optional (#494): the shared count-source probe factory; undefined on a
// stale build => the in-app layer falls back to no signal.
createListCommentsProbe: mod.createListCommentsProbe,
// Pure no-network draw.io helpers (#424); not client methods.
searchShapes: mod.searchShapes,
getGuideSection: mod.getGuideSection,
apps/server/src/core/comment/comment.controller.ts
+34
View File
@@ -16,6 +16,7 @@ import { UpdateCommentDto } from './dto/update-comment.dto';
import { ResolveCommentDto } from './dto/resolve-comment.dto';
import { ApplySuggestionDto } from './dto/apply-suggestion.dto';
import { DismissSuggestionDto } from './dto/dismiss-suggestion.dto';
import { ResyncSuggestionAnchorDto } from './dto/resync-suggestion-anchor.dto';
import { PageIdDto, CommentIdDto } from './dto/comments.input';
import { AuthUser } from '../../common/decorators/auth-user.decorator';
import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
@@ -235,6 +236,39 @@ export class CommentController {
return this.commentService.applySuggestion(comment, user, provenance);
}
@HttpCode(HttpStatus.OK)
@Post('resync-suggestion-anchor')
async resyncSuggestionAnchor(
@Body() dto: ResyncSuggestionAnchorDto,
@AuthUser() user: User,
@AuthWorkspace() workspace: Workspace,
) {
const comment = await this.commentRepo.findById(dto.commentId, {
includeCreator: true,
includeResolvedBy: true,
});
if (!comment) {
throw new NotFoundException('Comment not found');
}
const page = await this.pageRepo.findById(comment.pageId);
if (!page || page.deletedAt) {
throw new NotFoundException('Page not found');
}
// Authorize BEFORE revealing structural detail (mirrors apply/dismiss).
// Re-anchoring does NOT change the page text — it only corrects the stored
// selection metadata — so the page-level gate is comment access. The service
// further restricts it to the suggestion's own author.
await this.pageAccessService.validateCanComment(page, user, workspace.id);
return this.commentService.resyncSuggestionAnchor(
comment,
dto.selection,
user,
);
}
@HttpCode(HttpStatus.OK)
@Post('dismiss-suggestion')
async dismissSuggestion(
apps/server/src/core/comment/comment.service.apply-suggestion.spec.ts
+48 -2
View File
@@ -146,11 +146,19 @@ describe('CommentService — applySuggestion', () => {
'page-1',
expect.objectContaining({ operation: 'commentDeleted', commentId: 'c-1' }),
);
// #496: hard-deleted row → the audit payload is the only surviving record.
expect(auditService.log).toHaveBeenCalledWith(
expect.objectContaining({
event: AuditEvent.COMMENT_SUGGESTION_APPLIED,
resourceType: AuditResource.COMMENT,
resourceId: 'c-1',
metadata: expect.objectContaining({
pageId: 'page-1',
suggestedText: 'new text',
selection: 'old text',
commentAuthor: 'user-1',
decidedBy: 'user-1',
}),
}),
);
expect(result.outcome).toBe('deleted');
@@ -189,17 +197,25 @@ describe('CommentService — applySuggestion', () => {
expect(resolvePatch.resolvedAt).toBeInstanceOf(Date);
expect(resolvePatch.resolvedById).toBe('user-1');
// NOT deleted; broadcast an update, not a deletion.
// NOT deleted.
expect(commentRepo.deleteComment).not.toHaveBeenCalled();
expect(collaborationGateway.handleYjsEvent).not.toHaveBeenCalledWith(
'deleteCommentMark',
expect.anything(),
expect.anything(),
);
// #496 dedup: resolveComment broadcasts `commentResolved` with the enriched
// row; finalize must NOT ALSO emit a redundant `commentUpdated`. So the
// thread receives exactly ONE resolve broadcast and no update broadcast.
expect(wsService.emitCommentEvent).toHaveBeenCalledWith(
'space-1',
'page-1',
expect.objectContaining({ operation: 'commentUpdated', comment: UPDATED }),
expect.objectContaining({ operation: 'commentResolved', comment: UPDATED }),
);
expect(wsService.emitCommentEvent).not.toHaveBeenCalledWith(
'space-1',
'page-1',
expect.objectContaining({ operation: 'commentUpdated' }),
);
expect(auditService.log).toHaveBeenCalledWith(
@@ -211,6 +227,36 @@ describe('CommentService — applySuggestion', () => {
expect(result.outcome).toBe('resolved');
});
it('re-entry: already applied+resolved WITH replies → emits commentUpdated (dedup does not over-suppress)', async () => {
// suggestionAppliedAt set → idempotent finalize; resolvedAt set → resolveComment
// is skipped, so there is NO commentResolved broadcast. The applied-stamp state
// must still reach clients via a single commentUpdated.
const { service, wsService } = makeService(
{ applied: false, currentText: 'new text' },
true,
);
await service.applySuggestion(
suggestionComment({
suggestionAppliedAt: new Date(),
resolvedAt: new Date(),
}),
user(),
);
expect(wsService.emitCommentEvent).toHaveBeenCalledWith(
'space-1',
'page-1',
expect.objectContaining({ operation: 'commentUpdated', comment: UPDATED }),
);
// Nothing resolved this time (already resolved) → no resolve broadcast.
expect(wsService.emitCommentEvent).not.toHaveBeenCalledWith(
'space-1',
'page-1',
expect.objectContaining({ operation: 'commentResolved' }),
);
});
// --- error / rejection branches -----------------------------------------
it('applied=false and currentText differs → ConflictException with currentText in payload', async () => {
apps/server/src/core/comment/comment.service.dismiss-suggestion.spec.ts
+10
View File
@@ -107,11 +107,21 @@ describe('CommentService — dismissSuggestion', () => {
'page-1',
expect.objectContaining({ operation: 'commentDeleted', commentId: 'c-1' }),
);
// #496: the row is hard-deleted, so the audit payload must carry the
// decision's substance (what was suggested, the anchored text, who authored
// it, who decided) — it is the only surviving record.
expect(auditService.log).toHaveBeenCalledWith(
expect.objectContaining({
event: AuditEvent.COMMENT_SUGGESTION_DISMISSED,
resourceType: AuditResource.COMMENT,
resourceId: 'c-1',
metadata: expect.objectContaining({
pageId: 'page-1',
suggestedText: 'new text',
selection: 'old text',
commentAuthor: 'user-1',
decidedBy: 'user-1',
}),
}),
);
expect(result.outcome).toBe('deleted');
apps/server/src/core/comment/comment.service.resync-anchor.spec.ts
+126
View File
@@ -0,0 +1,126 @@
import { BadRequestException, ForbiddenException } from '@nestjs/common';
import { CommentService } from './comment.service';
/**
* Coverage for CommentService.resyncSuggestionAnchor (#496): re-anchoring a
* suggestion's stored selection (== apply-time expectedText) to the live-doc
* substring. The service is built directly with jest-mocked deps (the
* @InjectQueue tokens can't be resolved by Test.createTestingModule — see the
* sibling specs).
*/
describe('CommentService — resyncSuggestionAnchor', () => {
const UPDATED = { id: 'c-1', selection: 'new anchor', __updated: true } as any;
function makeService() {
const commentRepo: any = {
updateComment: jest.fn(async () => undefined),
findById: jest.fn(async () => UPDATED),
};
const service = new CommentService(
commentRepo,
{} as any,
{ emitCommentEvent: jest.fn() } as any,
{} as any,
{ add: jest.fn() } as any,
{ add: jest.fn() } as any,
{ log: jest.fn() } as any,
);
return { service, commentRepo };
}
const suggestion = (over?: Partial<any>): any => ({
id: 'c-1',
creatorId: 'user-1',
parentCommentId: null,
selection: 'old anchor',
suggestedText: 'new text',
suggestionAppliedAt: null,
resolvedAt: null,
...over,
});
const user = (over?: Partial<any>): any => ({ id: 'user-1', ...over });
it('persists the new selection and returns the enriched comment', async () => {
const { service, commentRepo } = makeService();
const out = await service.resyncSuggestionAnchor(
suggestion(),
'new anchor',
user(),
);
expect(commentRepo.updateComment).toHaveBeenCalledWith(
{ selection: 'new anchor' },
'c-1',
);
expect(out).toBe(UPDATED);
});
it('is idempotent: no write when the anchor already matches', async () => {
const { service, commentRepo } = makeService();
const out = await service.resyncSuggestionAnchor(
suggestion({ selection: 'same' }),
'same',
user(),
);
expect(commentRepo.updateComment).not.toHaveBeenCalled();
expect(out).toEqual(suggestion({ selection: 'same' }));
});
it('rejects a non-author (only the suggestion owner may re-anchor)', async () => {
const { service, commentRepo } = makeService();
await expect(
service.resyncSuggestionAnchor(suggestion(), 'new anchor', user({ id: 'other' })),
).rejects.toBeInstanceOf(ForbiddenException);
expect(commentRepo.updateComment).not.toHaveBeenCalled();
});
it('rejects a reply / a comment with no suggestion', async () => {
const { service } = makeService();
await expect(
service.resyncSuggestionAnchor(
suggestion({ parentCommentId: 'p-1' }),
'x',
user(),
),
).rejects.toBeInstanceOf(BadRequestException);
await expect(
service.resyncSuggestionAnchor(
suggestion({ suggestedText: null }),
'x',
user(),
),
).rejects.toBeInstanceOf(BadRequestException);
});
it('rejects re-anchoring an already applied or resolved suggestion', async () => {
const { service } = makeService();
await expect(
service.resyncSuggestionAnchor(
suggestion({ suggestionAppliedAt: new Date() }),
'x',
user(),
),
).rejects.toBeInstanceOf(BadRequestException);
await expect(
service.resyncSuggestionAnchor(
suggestion({ resolvedAt: new Date() }),
'x',
user(),
),
).rejects.toBeInstanceOf(BadRequestException);
});
it('rejects a no-op selection equal to the suggested text', async () => {
const { service } = makeService();
await expect(
service.resyncSuggestionAnchor(
suggestion({ suggestedText: 'new text' }),
'new text',
user(),
),
).rejects.toBeInstanceOf(BadRequestException);
});
});
apps/server/src/core/comment/comment.service.ts
+122 -17
View File
@@ -370,6 +370,76 @@ export class CommentService {
return updatedComment;
}
/**
* Re-sync a suggestion's stored `selection` (== apply-time expectedText) to the
* RAW substring the inline mark actually covers in the LIVE document (#496).
*
* The MCP client creates the comment from a DEBOUNCED REST snapshot, then
* anchors the mark in the live collab doc. When the two disagree (the doc moved
* on in the debounce window) the stored selection no longer equals the marked
* text, so EVERY apply 409s ("the commented text changed"). After anchoring the
* client re-reads the exact marked substring and calls this to store it, making
* apply's strict equality hold.
*
* Only meaningful for an un-settled top-level suggestion authored by the
* caller: applying/resolving freezes the anchor, and a reply-carrying thread is
* preserved rather than mutated. The new text must still differ from the
* suggestion (else "apply" would be a no-op), preserving create()'s invariant.
*/
async resyncSuggestionAnchor(
comment: Comment,
selection: string,
user: User,
): Promise<Comment> {
if (comment.creatorId !== user.id) {
throw new ForbiddenException(
'You can only re-anchor your own suggestion',
);
}
if (comment.parentCommentId) {
throw new BadRequestException(
'Only a top-level comment can carry a suggested edit',
);
}
if (!comment.suggestedText) {
throw new BadRequestException('This comment has no suggested edit');
}
// A settled suggestion's anchor is frozen: re-anchoring an applied/resolved
// thread is meaningless and could resurrect a stale expectedText.
if (comment.suggestionAppliedAt || comment.resolvedAt) {
throw new BadRequestException(
'Cannot re-anchor a suggestion that was already applied or resolved',
);
}
const trimmed = selection.trim();
if (trimmed.length === 0) {
throw new BadRequestException('The re-anchored selection cannot be empty');
}
// Same no-op guard as create(): the suggestion must differ from the text it
// replaces, or apply becomes indistinguishable from already-applied.
if (trimmed === comment.suggestedText.trim()) {
throw new BadRequestException(
'A suggested edit must differ from the selected text',
);
}
// Idempotent: nothing to persist when the anchor already matches.
if (comment.selection === selection) {
return comment;
}
await this.commentRepo.updateComment({ selection }, comment.id);
const updatedComment = await this.commentRepo.findById(comment.id, {
includeCreator: true,
includeResolvedBy: true,
});
// Re-anchoring only corrects stored metadata; it does not change the page
// text or the comment body, so no ws broadcast / notification is warranted.
return updatedComment;
}
/**
* Apply the suggested edit carried by a top-level inline comment: atomically
* replace the text under the comment mark in the collaborative document with
@@ -524,7 +594,7 @@ export class CommentService {
resourceType: AuditResource.COMMENT,
resourceId: comment.id,
spaceId: comment.spaceId,
metadata: { pageId: comment.pageId },
metadata: this.suggestionAuditMetadata(comment, user),
});
return { ...updatedComment, outcome: 'resolved' };
}
@@ -538,7 +608,7 @@ export class CommentService {
resourceType: AuditResource.COMMENT,
resourceId: comment.id,
spaceId: comment.spaceId,
metadata: { pageId: comment.pageId },
metadata: this.suggestionAuditMetadata(comment, user),
});
return settled;
}
@@ -577,8 +647,10 @@ export class CommentService {
// Auto-resolve the thread. resolveComment handles the resolve mark, its ws
// broadcast and the resolve notification. Stay defensive on re-entry.
let didResolveBroadcast = false;
if (!comment.resolvedAt) {
await this.resolveComment(comment, true, user, provenance);
didResolveBroadcast = true;
}
const updatedComment = await this.commentRepo.findById(comment.id, {
@@ -586,18 +658,27 @@ export class CommentService {
includeResolvedBy: true,
});
this.wsService.emitCommentEvent(comment.spaceId, comment.pageId, {
operation: 'commentUpdated',
pageId: comment.pageId,
comment: updatedComment,
});
// #496 dedup: resolveComment already broadcast `commentResolved` carrying
// the fully-enriched row (the applied stamps were persisted above, before
// that call, so its re-read reflects them). Emitting `commentUpdated` here
// too made the client receive TWO events for one apply. Broadcast the
// update ONLY when we did NOT resolve — i.e. the rare re-entry on an
// already-resolved thread, where the applied-stamp change still needs a
// broadcast and resolveComment did not run.
if (!didResolveBroadcast) {
this.wsService.emitCommentEvent(comment.spaceId, comment.pageId, {
operation: 'commentUpdated',
pageId: comment.pageId,
comment: updatedComment,
});
}
this.auditService.log({
event: AuditEvent.COMMENT_SUGGESTION_APPLIED,
resourceType: AuditResource.COMMENT,
resourceId: comment.id,
spaceId: comment.spaceId,
metadata: { pageId: comment.pageId },
metadata: this.suggestionAuditMetadata(comment, user),
});
return { ...updatedComment, outcome: 'resolved' };
@@ -616,7 +697,7 @@ export class CommentService {
resourceType: AuditResource.COMMENT,
resourceId: comment.id,
spaceId: comment.spaceId,
metadata: { pageId: comment.pageId },
metadata: this.suggestionAuditMetadata(comment, user),
});
return settled;
@@ -627,14 +708,17 @@ export class CommentService {
* inline `comment` anchor mark, then ATOMICALLY hard-delete the row only if it
* is still childless. Shared by the apply/dismiss no-replies branches (#329).
*
* ORDER MATTERS: the anchor mark is removed FIRST and FATALLY (mirrors
* applySuggestion, which mutates the doc before writing the DB). The row
* delete is irreversible, so if the mark removal fails — including the
* COLLAB_DISABLE_REDIS "no live instance" hard-error — we must NOT delete the
* row and report success, or the document is left with a permanent orphan
* anchor pointing at a comment that no longer exists (the exact data-integrity
* bug #329 targets). Let the exception propagate (→ 5xx); the operation is
* then repeatable with row + mark still consistent.
* ORDER MATTERS (updated #399 → #496): what runs FIRST and FATALLY here is the
* mark-removal ENQUEUE (a fast, durable Redis add), NOT the mark op itself.
* deleteCommentMark awaits only the enqueue, so a failed add throws BEFORE the
* irreversible row delete — the row + mark stay consistent and the operation is
* repeatable. The actual anchor strip then runs off the HTTP path in the worker
* (idempotent, 3 retries). Only an EXHAUSTED-retries job could leave the doc
* with an orphan anchor pointing at a hard-deleted comment (the data-integrity
* bug #329 targets); that residual divergence is now self-healed by the
* resolve/unresolve mark worker, which strips an orphan mark whenever its
* comment row is gone (#496), and it is meanwhile VISIBLE via BullMQ failed-job
* metrics rather than a silently-swallowed warn.
*
* RACE (#338 F4): the caller read `hasChildren` BEFORE the (slow) mark
* removal, so a reply can land in that window. `comments.parent_comment_id` is
@@ -732,6 +816,27 @@ export class CommentService {
return this.generalQueue.add(QueueJob.COMMENT_MARK_UPDATE, jobData);
}
/**
* Build the audit metadata for a suggestion apply/dismiss decision (#496).
* The subject comment is HARD-DELETED on the childless path, so the audit row
* is the only surviving record — capture the decision's substance (what was
* suggested, the anchored text it replaced, who authored it, who decided)
* before the row can vanish. `decidedBy` is the acting user; `commentAuthor`
* is the suggestion's creator.
*/
private suggestionAuditMetadata(
comment: Comment,
user: User,
): Record<string, any> {
return {
pageId: comment.pageId,
suggestedText: comment.suggestedText ?? null,
selection: comment.selection ?? null,
commentAuthor: comment.creatorId ?? null,
decidedBy: user.id,
};
}
private async queueCommentNotification(
content: any,
oldMentionIds: string[],
apps/server/src/core/comment/dto/resync-suggestion-anchor.dto.ts
+20
View File
@@ -0,0 +1,20 @@
import { IsString, IsUUID, MaxLength, MinLength } from 'class-validator';
/**
* #496: after the MCP client anchors a suggestion in the LIVE collab doc, it
* re-reads the exact substring under the new mark and syncs it here as the
* comment's stored `selection` (== apply-time expectedText). Fixes the perpetual
* 409 where expectedText came from a debounced REST snapshot while the mark sat
* in the live doc.
*/
export class ResyncSuggestionAnchorDto {
@IsUUID()
commentId: string;
// The raw substring the mark now covers in the live document. Bounded like the
// create-time selection (2000) so a legitimate anchored span is never cut.
@IsString()
@MinLength(1)
@MaxLength(2000)
selection: string;
}
apps/server/src/core/page/dto/move-page.dto.spec.ts
+33 -23
View File
@@ -17,9 +17,10 @@ import { MovePageDto } from './move-page.dto';
// a valid ordering key the server itself generated would be refused on move.
//
// The tests below assert the CORRECT contract: any key the generator can produce
// must satisfy the DTO. The genuinely-failing case is marked `test.failing` so the
// suite stays green while locking the bug; it flips red (alerting us) once the DTO
// bounds are widened to cover the generator's real range.
// must satisfy the DTO. FIXED (#495 item 9): the DTO now validates `position` by
// CHARSET ([0-9A-Za-z], the generator's base-62 alphabet) instead of the wrong
// @MaxLength(12) length bound, so dense between-inserts are accepted; the former
// `test.failing` bug-lock is now a passing assertion.
function constraintErrors(position: unknown) {
const dto = plainToInstance(MovePageDto, {
@@ -47,24 +48,33 @@ describe('MovePageDto.position vs generateJitteredKeyBetween parity', () => {
expect(hasError(errors, 'position')).toBe(false);
});
// BUG LOCK: dense between-inserts produce keys longer than 12 chars, which
// MaxLength(12) rejects even though they are valid ordering keys. This SHOULD
// pass; it currently fails. Flips green when the DTO bound is fixed.
test.failing(
'accepts dense between-inserted keys (currently rejected by MaxLength(12))',
async () => {
let lo = generateJitteredKeyBetween(null, null);
let hi = generateJitteredKeyBetween(lo, null);
// Repeatedly insert just above `lo`, shrinking the gap so the key grows.
let longest = lo;
for (let i = 0; i < 40; i++) {
const mid = generateJitteredKeyBetween(lo, hi);
if (mid.length > longest.length) longest = mid;
hi = mid;
}
expect(longest.length).toBeGreaterThan(12); // sanity: we produced a long key
const errors = await constraintErrors(longest);
expect(hasError(errors, 'position')).toBe(false);
},
);
// FIXED: dense between-inserts produce keys longer than 12 chars, which the old
// MaxLength(12) rejected even though they are valid ordering keys. Now accepted.
it('accepts dense between-inserted keys longer than 12 chars', async () => {
let lo = generateJitteredKeyBetween(null, null);
let hi = generateJitteredKeyBetween(lo, null);
// Repeatedly insert just above `lo`, shrinking the gap so the key grows. The
// generator is JITTERED (random), so use enough iterations that the longest
// key reliably clears the old 12-char bound: measured min-over-50-trials is
// ~11 at 40 iterations (flaky) but ~36 at 200 (robust margin).
let longest = lo;
for (let i = 0; i < 200; i++) {
const mid = generateJitteredKeyBetween(lo, hi);
if (mid.length > longest.length) longest = mid;
hi = mid;
}
expect(longest.length).toBeGreaterThan(12); // sanity: we produced a long key
const errors = await constraintErrors(longest);
expect(hasError(errors, 'position')).toBe(false);
});
// The charset guard replaces the length bound: reject anything outside the
// generator's [0-9A-Za-z] alphabet (control chars, separators, injection) and
// the empty string, while still accepting every real key.
it('rejects a position with characters outside the fractional-index alphabet', async () => {
for (const bad of ['a0/b', 'a b', 'a\n0', 'a.b', '', "a';--"]) {
const errors = await constraintErrors(bad);
expect(hasError(errors, 'position')).toBe(true);
}
});
});
apps/server/src/core/page/dto/move-page.dto.ts
+13 -3
View File
@@ -1,8 +1,8 @@
import {
IsString,
IsOptional,
MinLength,
MaxLength,
Matches,
IsNotEmpty,
} from 'class-validator';
@@ -10,9 +10,19 @@ export class MovePageDto {
@IsString()
pageId: string;
// `position` is a fractional-indexing key from `generateJitteredKeyBetween`
// (the SAME generator page.service uses). Validate by CHARSET, not length: the
// generator's default base-62 alphabet is [0-9A-Za-z], and DENSE between-inserts
// legitimately grow a key well past a dozen chars (measured >40), so the old
// @MinLength(5)/@MaxLength(12) bounds wrongly 400'd valid ordering keys the
// server itself produced (Gitea #139 item 6). The charset regex rejects control
// chars / separators / injection, and a generous MaxLength stays only as a
// DoS guard — far above any realistic key, so it never rejects a real move.
@IsString()
@MinLength(5)
@MaxLength(12)
@Matches(/^[0-9A-Za-z]+$/, {
message: 'position must be a fractional-index key ([0-9A-Za-z])',
})
@MaxLength(256)
position: string;
@IsOptional()
apps/server/src/core/page/dto/page-identity.dto.spec.ts
+49
View File
@@ -0,0 +1,49 @@
import 'reflect-metadata';
import { plainToInstance } from 'class-transformer';
import { validate } from 'class-validator';
import { PageIdDto } from './page.dto';
// #435: PageIdDto.pageId carries a page's DOUBLE identity (internal UUID OR
// public 10-char slugId), both as bare strings. The DTO must accept exactly
// those two FORMATS and reject a malformed / swapped identity at the boundary.
async function pageIdErrors(pageId: unknown) {
const dto = plainToInstance(PageIdDto, { pageId });
const errors = await validate(dto as object);
return errors.some((e) => e.property === 'pageId');
}
const UUID = '019f499a-9f8c-7d68-b7be-ce100d7c6c56';
const SLUG = 'aB3xQ7kR2p';
describe('PageIdDto pageId format validation', () => {
it('accepts a canonical page UUID', async () => {
expect(await pageIdErrors(UUID)).toBe(false);
});
it('accepts a 10-char slugId', async () => {
expect(await pageIdErrors(SLUG)).toBe(false);
});
it('rejects a truncated / wrong-length slug', async () => {
expect(await pageIdErrors('aB3xQ7kR2')).toBe(true); // 9 chars
expect(await pageIdErrors('aB3xQ7kR2pX')).toBe(true); // 11 chars
});
it('rejects a slug with an illegal character', async () => {
expect(await pageIdErrors('aB3xQ7kR2!')).toBe(true);
});
it('rejects a full URL / path-shaped identity (not the bare id)', async () => {
expect(await pageIdErrors(`my-page-title-${SLUG}`)).toBe(true);
expect(await pageIdErrors(`https://x/p/${SLUG}`)).toBe(true);
});
it('rejects a malformed UUID', async () => {
expect(await pageIdErrors('019f499a-9f8c-7d68-b7be')).toBe(true);
expect(await pageIdErrors('not-a-uuid-at-all-really')).toBe(true);
});
it('rejects an empty string', async () => {
expect(await pageIdErrors('')).toBe(true);
});
});
apps/server/src/core/page/dto/page-identity.validator.ts
+39
View File
@@ -0,0 +1,39 @@
import { applyDecorators } from '@nestjs/common';
import { Matches, ValidationOptions } from 'class-validator';
/**
* A page identity at the API boundary is EITHER the internal page UUID or the
* public 10-char slugId (page.repo.findById matches a non-UUID input as a
* slugId). Both arrive as bare strings, which is exactly how the two got swapped
* silently (incident family #435). This regex pins the two accepted FORMATS so a
* malformed / cross-wired identity (a truncated slug, a full URL, an email, an
* id from another entity kind that isn't even shaped like either) is rejected at
* the boundary instead of falling through to a confusing 404.
*
* - UUID: canonical 8-4-4-4-12 hex (version-agnostic — page ids are UUIDv7,
* so only the shape/length is enforced, matching the MCP's UUID_RE
* and the server's isValidUUID acceptance).
* - slugId: exactly 10 chars over [0-9A-Za-z] (nanoid `generateSlugId`).
*
* The two are disjoint (a UUID is 36 chars WITH dashes, a slugId 10 chars
* WITHOUT), so a value can only satisfy one branch.
*/
export const PAGE_ID_OR_SLUG_ID_REGEX =
/^(?:[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}|[0-9A-Za-z]{10})$/i;
/**
* Validate that a DTO string field is a well-formed page identity (page UUID OR
* 10-char slugId). Composed decorator so the same format rule is applied
* consistently wherever a DTO accepts a `pageId` that may be either form.
*/
export function IsPageIdOrSlugId(
validationOptions?: ValidationOptions,
): PropertyDecorator {
return applyDecorators(
Matches(PAGE_ID_OR_SLUG_ID_REGEX, {
message:
'pageId must be a page UUID or a 10-character slugId',
...validationOptions,
}),
);
}
apps/server/src/core/page/dto/page.dto.ts
+6
View File
@@ -9,10 +9,16 @@ import {
import { Transform } from 'class-transformer';
import { ContentFormat } from './create-page.dto';
import { IsPageIdOrSlugId } from './page-identity.validator';
export class PageIdDto {
@IsString()
@IsNotEmpty()
// Format-validate the double identity (#435): accept only a page UUID or a
// 10-char slugId so a malformed / swapped identity is rejected at the boundary
// rather than passed to the repo as a bare string. Base for PageInfoDto,
// DeletePageDto, BacklinksListDto, AddLabelsDto/RemoveLabelDto, etc.
@IsPageIdOrSlugId()
pageId: string;
}
apps/server/src/core/page/services/page.service.ts
+4 -2
View File
@@ -53,8 +53,10 @@ import {
extractPageSlugId,
} from '../../../integrations/export/utils';
import { canonicalizeFootnotes } from '@docmost/editor-ext';
import { markdownToProseMirror } from '@docmost/prosemirror-markdown';
import { normalizeForeignMarkdown } from '../../../integrations/import/utils/foreign-markdown';
import {
markdownToProseMirror,
normalizeForeignMarkdown,
} from '@docmost/prosemirror-markdown';
import { WatcherService } from '../../watcher/watcher.service';
import { sql } from 'kysely';
import { TransclusionService } from '../transclusion/transclusion.service';
apps/server/src/core/page/services/spec/temporary-note-cleanup.service.spec.ts
+96 -63
View File
@@ -1,19 +1,38 @@
import { TemporaryNoteCleanupService } from '../temporary-note-cleanup.service';
/**
* Chainable Kysely stub that records every `.where(...)` call so the test can
* assert the sweep only selects armed, expired, not-yet-trashed notes. The
* terminal `.execute()` resolves the configured expired rows (the batch SELECT);
* `.executeTakeFirst()` resolves the per-row deadline re-read done just before
* each `removePage`. By default the re-read reports the note as still armed and
* still expired (epoch deadline < now), so the sweep proceeds to delete it;
* tests override `reReadFirst` to simulate a concurrent "Make permanent".
* Chainable Kysely stub for the temporary-note sweep.
*
* `this.db` serves the non-locking candidate SELECT (selectFrom/select/where/
* limit/execute -> the configured expired rows) AND `.transaction().execute(cb)`,
* which runs `cb` with a separate `trx` builder. The `trx` builder serves the
* per-row LOCKED re-check (selectFrom/select/where/forUpdate/skipLocked/
* executeTakeFirst). `lockedRows` drives what that locked re-check returns per
* candidate — an id/creator/workspace row means "still expired, delete it";
* `undefined` means the predicate no longer matched (made permanent / re-armed /
* already trashed) or the row was SKIP-LOCKED by another worker, so it is skipped.
*/
function makeDbStub(expiredRows: any[]) {
function makeDbStub(expiredRows: any[], lockedRows?: any[]) {
const whereCalls: any[][] = [];
const reReadFirst = jest
.fn()
.mockResolvedValue({ temporaryExpiresAt: new Date(0), deletedAt: null });
const locked = [
...(lockedRows ??
expiredRows.map((r) => ({
id: r.id,
creatorId: r.creatorId,
workspaceId: r.workspaceId,
}))),
];
const lockedTakeFirst = jest.fn(() => Promise.resolve(locked.shift()));
const forUpdate = jest.fn(() => trxBuilder);
const skipLocked = jest.fn(() => trxBuilder);
const trxBuilder: any = {
selectFrom: jest.fn(() => trxBuilder),
select: jest.fn(() => trxBuilder),
where: jest.fn(() => trxBuilder),
forUpdate,
skipLocked,
executeTakeFirst: lockedTakeFirst,
};
const builder: any = {
selectFrom: jest.fn(() => builder),
select: jest.fn(() => builder),
@@ -23,9 +42,11 @@ function makeDbStub(expiredRows: any[]) {
}),
limit: jest.fn(() => builder),
execute: jest.fn().mockResolvedValue(expiredRows),
executeTakeFirst: reReadFirst,
transaction: jest.fn(() => ({
execute: (cb: (trx: any) => Promise<any>) => Promise.resolve(cb(trxBuilder)),
})),
};
return { builder, whereCalls, reReadFirst };
return { builder, whereCalls, lockedTakeFirst, forUpdate, skipLocked };
}
describe('TemporaryNoteCleanupService.sweepExpiredTemporaryNotes', () => {
@@ -52,20 +73,36 @@ describe('TemporaryNoteCleanupService.sweepExpiredTemporaryNotes', () => {
expect(builder.limit.mock.calls[0][0]).toBeGreaterThan(0);
});
it('soft-deletes each expired note via removePage, attributed to its creator', async () => {
it('soft-deletes each expired note via removePage under a row lock, attributed to its creator', async () => {
const expired = [
{ id: 'p1', creatorId: 'u1', workspaceId: 'w1' },
{ id: 'p2', creatorId: 'u2', workspaceId: 'w1' },
];
const { builder } = makeDbStub(expired);
const { builder, forUpdate, skipLocked } = makeDbStub(expired);
const pageRepo = { removePage: jest.fn().mockResolvedValue(undefined) } as any;
const service = new TemporaryNoteCleanupService(builder, pageRepo);
await service.sweepExpiredTemporaryNotes();
expect(pageRepo.removePage).toHaveBeenCalledTimes(2);
expect(pageRepo.removePage).toHaveBeenNthCalledWith(1, 'p1', 'u1', 'w1');
expect(pageRepo.removePage).toHaveBeenNthCalledWith(2, 'p2', 'u2', 'w1');
// The 4th arg is the locking transaction — the delete runs inside it.
expect(pageRepo.removePage).toHaveBeenNthCalledWith(
1,
'p1',
'u1',
'w1',
expect.anything(),
);
expect(pageRepo.removePage).toHaveBeenNthCalledWith(
2,
'p2',
'u2',
'w1',
expect.anything(),
);
// The re-check acquired a FOR UPDATE SKIP LOCKED lock (once per candidate).
expect(forUpdate).toHaveBeenCalledTimes(2);
expect(skipLocked).toHaveBeenCalledTimes(2);
});
it('continues past a failing note (one bad removePage does not abort the sweep)', async () => {
@@ -86,60 +123,29 @@ describe('TemporaryNoteCleanupService.sweepExpiredTemporaryNotes', () => {
service.sweepExpiredTemporaryNotes(),
).resolves.toBeUndefined();
expect(pageRepo.removePage).toHaveBeenCalledTimes(2);
expect(pageRepo.removePage).toHaveBeenNthCalledWith(2, 'good', 'u2', 'w1');
expect(pageRepo.removePage).toHaveBeenNthCalledWith(
2,
'good',
'u2',
'w1',
expect.anything(),
);
});
it('does NOT trash a note made permanent in the race window', async () => {
// The batch SELECT saw the note as expired, but before its turn in the loop
// the user clicked "Make permanent" (temporary_expires_at -> null). The
// deadline re-read must catch this and skip the delete so the keep wins.
it('does NOT trash a note made permanent / re-armed / already trashed (locked re-check returns nothing)', async () => {
// The batch SELECT saw the note as expired, but by the time the LOCKED
// re-check runs the row no longer matches the still-armed+expired+not-trashed
// predicate (make-permanent, re-arm to a future deadline, or already trashed),
// OR another worker holds the row (SKIP LOCKED). In every case the locked
// SELECT returns nothing and the delete is skipped so the keep/other worker wins.
const expired = [{ id: 'p1', creatorId: 'u1', workspaceId: 'w1' }];
const { builder, reReadFirst } = makeDbStub(expired);
reReadFirst.mockResolvedValueOnce({
temporaryExpiresAt: null,
deletedAt: null,
});
const { builder, lockedTakeFirst } = makeDbStub(expired, [undefined]);
const pageRepo = { removePage: jest.fn() } as any;
const service = new TemporaryNoteCleanupService(builder, pageRepo);
await service.sweepExpiredTemporaryNotes();
expect(reReadFirst).toHaveBeenCalledTimes(1);
expect(pageRepo.removePage).not.toHaveBeenCalled();
});
it('skips a note already trashed since the batch SELECT', async () => {
const expired = [{ id: 'p1', creatorId: 'u1', workspaceId: 'w1' }];
const { builder, reReadFirst } = makeDbStub(expired);
reReadFirst.mockResolvedValueOnce({
temporaryExpiresAt: new Date(0),
deletedAt: new Date(),
});
const pageRepo = { removePage: jest.fn() } as any;
const service = new TemporaryNoteCleanupService(builder, pageRepo);
await service.sweepExpiredTemporaryNotes();
expect(pageRepo.removePage).not.toHaveBeenCalled();
});
it('does NOT trash a note re-armed to a future deadline in the race window', async () => {
// The batch SELECT saw the note as expired, but before its turn in the loop
// the user disarmed it and re-armed it to a fresh, still-future deadline
// (temporary_expires_at -> now + 1h). The deadline re-read must catch that
// the note is no longer expired and skip the delete so the keep wins.
const expired = [{ id: 'p1', creatorId: 'u1', workspaceId: 'w1' }];
const { builder, reReadFirst } = makeDbStub(expired);
reReadFirst.mockResolvedValueOnce({
temporaryExpiresAt: new Date(Date.now() + 60 * 60 * 1000),
deletedAt: null,
});
const pageRepo = { removePage: jest.fn() } as any;
const service = new TemporaryNoteCleanupService(builder, pageRepo);
await service.sweepExpiredTemporaryNotes();
expect(reReadFirst).toHaveBeenCalledTimes(1);
expect(lockedTakeFirst).toHaveBeenCalledTimes(1);
expect(pageRepo.removePage).not.toHaveBeenCalled();
});
@@ -151,4 +157,31 @@ describe('TemporaryNoteCleanupService.sweepExpiredTemporaryNotes', () => {
await service.sweepExpiredTemporaryNotes();
expect(pageRepo.removePage).not.toHaveBeenCalled();
});
it('sweeps once on application bootstrap (catches notes expired during downtime)', async () => {
const expired = [{ id: 'p1', creatorId: 'u1', workspaceId: 'w1' }];
const { builder } = makeDbStub(expired);
const pageRepo = { removePage: jest.fn().mockResolvedValue(undefined) } as any;
const service = new TemporaryNoteCleanupService(builder, pageRepo);
await service.onApplicationBootstrap();
expect(pageRepo.removePage).toHaveBeenCalledTimes(1);
expect(pageRepo.removePage).toHaveBeenCalledWith(
'p1',
'u1',
'w1',
expect.anything(),
);
});
it('a startup-sweep failure never blocks application boot', async () => {
const { builder } = makeDbStub([]);
// Make the candidate SELECT throw to simulate a boot-time DB hiccup.
builder.execute.mockRejectedValueOnce(new Error('db not ready'));
const pageRepo = { removePage: jest.fn() } as any;
const service = new TemporaryNoteCleanupService(builder, pageRepo);
await expect(service.onApplicationBootstrap()).resolves.toBeUndefined();
});
});
apps/server/src/core/page/services/temporary-note-cleanup.service.ts
+66 -42
View File
@@ -1,8 +1,13 @@
import { Injectable, Logger } from '@nestjs/common';
import {
Injectable,
Logger,
OnApplicationBootstrap,
} from '@nestjs/common';
import { Interval } from '@nestjs/schedule';
import { InjectKysely } from 'nestjs-kysely';
import { KyselyDB } from '@docmost/db/types/kysely.types';
import { PageRepo } from '@docmost/db/repos/page/page.repo';
import { executeTx } from '@docmost/db/utils';
/**
* Background sweeper for temporary notes ("structure or die"). A note whose
@@ -11,7 +16,7 @@ import { PageRepo } from '@docmost/db/repos/page/page.repo';
* TrashCleanupService; `@nestjs/schedule` is already enabled globally.
*/
@Injectable()
export class TemporaryNoteCleanupService {
export class TemporaryNoteCleanupService implements OnApplicationBootstrap {
private readonly logger = new Logger(TemporaryNoteCleanupService.name);
// Cap a single sweep so a large backlog (e.g. many notes created during
@@ -24,6 +29,20 @@ export class TemporaryNoteCleanupService {
private readonly pageRepo: PageRepo,
) {}
// Sweep once at startup so notes that expired during downtime are trashed
// right away instead of waiting up to an hour for the first @Interval tick.
// Best-effort: never let a startup-sweep failure block application boot.
async onApplicationBootstrap() {
try {
await this.sweepExpiredTemporaryNotes();
} catch (error) {
this.logger.error(
'Temporary-note startup sweep failed',
error instanceof Error ? error.stack : undefined,
);
}
}
// Hourly granularity: lifetimes are configured in hours, so a sub-hour
// overshoot past the deadline is acceptable.
@Interval('temporary-note-cleanup', 60 * 60 * 1000)
@@ -31,9 +50,11 @@ export class TemporaryNoteCleanupService {
try {
const now = new Date();
// Candidate ids (non-locking). The authoritative re-check happens per row
// under a row lock below, so this cheap pass just bounds the batch.
const expired = await this.db
.selectFrom('pages')
.select(['id', 'creatorId', 'workspaceId'])
.select(['id'])
.where('temporaryExpiresAt', 'is not', null)
.where('temporaryExpiresAt', '<', now)
.where('deletedAt', 'is', null) // not already in trash
@@ -41,50 +62,53 @@ export class TemporaryNoteCleanupService {
.execute();
let trashed = 0;
for (const page of expired) {
for (const candidate of expired) {
try {
// Re-check the deadline at deletion time. The SELECT above is not
// transactional, so a user may click "Make permanent"
// (toggleTemporary sets temporary_expires_at = null) in the window
// between the SELECT and this per-row removePage. removePage deletes
// by id with only a `deletedAt IS NULL` filter and never re-reads the
// deadline, so without this guard a concurrently-kept note would
// still be trashed. Re-read the row and skip it unless it is still
// armed AND still expired, so a concurrent make-permanent wins.
const current = await this.db
.selectFrom('pages')
.select(['temporaryExpiresAt', 'deletedAt'])
.where('id', '=', page.id)
.executeTakeFirst();
const didTrash = await executeTx(this.db, async (trx) => {
// Re-check the row UNDER A LOCK inside the transaction. `FOR UPDATE
// SKIP LOCKED`:
// - serialises against a concurrent "Make permanent"
// (toggleTemporary UPDATE takes the same row lock): if it commits
// first, the deadline predicate below no longer matches and we
// skip; if we lock first, it waits until this delete commits.
// - SKIP LOCKED lets a second worker/instance skip a row another
// sweeper already claimed instead of blocking on it (no double
// processing, no thundering herd).
// The predicate re-asserts still-armed AND still-expired AND
// not-already-trashed, so a make-permanent / prior sweep drops the row.
const locked = await trx
.selectFrom('pages')
.select(['id', 'creatorId', 'workspaceId'])
.where('id', '=', candidate.id)
.where('temporaryExpiresAt', 'is not', null)
.where('temporaryExpiresAt', '<', now)
.where('deletedAt', 'is', null)
.forUpdate()
.skipLocked()
.executeTakeFirst();
if (
!current ||
current.deletedAt !== null ||
current.temporaryExpiresAt === null ||
new Date(current.temporaryExpiresAt) >= now
) {
// Made permanent, already trashed, or no longer expired since the
// SELECT — leave it alone.
continue;
}
if (!locked) return false;
// Reuse the exact soft-delete path: recursive over children, removes
// shares in a transaction, and emits PAGE_SOFT_DELETED (tree
// invalidation + watcher notifications). Attribute the automatic
// deletion to the note's creator (no schema change). Both the SELECT
// above and removePage filter `deletedAt IS NULL`, so a double sweep
// is idempotent.
await this.pageRepo.removePage(
page.id,
// creatorId is set on every created page; a temporary note always
// has one. Cast to satisfy the non-null deletedById parameter.
page.creatorId as string,
page.workspaceId,
);
trashed++;
// Reuse the exact soft-delete path (recursive children + share
// removal + PAGE_SOFT_DELETED broadcast), running IN this locked
// transaction so the delete is atomic with the re-check and cannot
// deadlock on a nested independent transaction. The broadcast is
// deferred by removePage to this transaction's commit. Attribute the
// automatic deletion to the note's creator (no schema change).
await this.pageRepo.removePage(
locked.id,
// creatorId is set on every created page; a temporary note always
// has one. Cast to satisfy the non-null deletedById parameter.
locked.creatorId as string,
locked.workspaceId,
trx,
);
return true;
});
if (didTrash) trashed++;
} catch (error) {
this.logger.error(
`Failed to trash expired temporary note ${page.id}`,
`Failed to trash expired temporary note ${candidate.id}`,
error instanceof Error ? error.stack : undefined,
);
}
apps/server/src/core/share/share-alias.controller.spec.ts
+3
View File
@@ -133,6 +133,9 @@ describe('ShareAliasController authz gates', () => {
creatorId: 'u-1',
alias: 'promo',
confirmReassign: true,
// The requesting user is forwarded so setAlias can gate the reassign
// 409 title disclosure on target-page view permission (#495).
user,
});
expect(result).toEqual({ id: 'alias-1' });
});
apps/server/src/core/share/share-alias.controller.ts
+3
View File
@@ -79,6 +79,9 @@ export class ShareAliasController {
creatorId: user.id,
alias: dto.alias,
confirmReassign: dto.confirmReassign,
// Gates whether the reassign 409 may reveal the current target's title
// (view-permission check on that page) — see setAlias (#495).
user,
});
}
apps/server/src/core/share/share-alias.service.spec.ts
+84 -9
View File
@@ -1,4 +1,8 @@
import { BadRequestException, ConflictException } from '@nestjs/common';
import {
BadRequestException,
ConflictException,
ForbiddenException,
} from '@nestjs/common';
import { NoResultError } from 'kysely';
import { ShareAliasService } from './share-alias.service';
@@ -7,6 +11,8 @@ import { ShareAliasService } from './share-alias.service';
* 409 reassign guard, uniqueness-race handling, availability probe, and the
* request-time readable-target resolution (which re-runs the share boundary).
*/
const USER = { id: 'u-1' } as any;
describe('ShareAliasService', () => {
// Sentinel handed to repo calls so tests can assert they ran inside the tx.
const trx = { __trx: true };
@@ -27,6 +33,10 @@ describe('ShareAliasService', () => {
resolveReadableSharePage: jest.fn(),
isSharingAllowed: jest.fn(),
};
// Default: the requester CAN view the target page (validateCanView resolves),
// so the reassign 409 may disclose its title. Tests override to reject to
// assert the no-leak path.
const pageAccessService = { validateCanView: jest.fn().mockResolvedValue(undefined) };
// Fake kysely db: only .transaction().execute(cb) is used by setAlias.
const db = {
transaction: jest.fn(() => ({
@@ -37,9 +47,10 @@ describe('ShareAliasService', () => {
shareAliasRepo as any,
pageRepo as any,
shareService as any,
pageAccessService as any,
db as any,
);
return { service, shareAliasRepo, pageRepo, shareService, db };
return { service, shareAliasRepo, pageRepo, shareService, pageAccessService, db };
}
describe('setAlias', () => {
@@ -50,6 +61,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'A', // too short + uppercase
}),
).rejects.toBeInstanceOf(BadRequestException);
@@ -66,6 +78,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: ' My Page ',
});
@@ -114,6 +127,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'ted',
});
@@ -144,6 +158,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'foo',
});
@@ -179,6 +194,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'new',
});
@@ -190,30 +206,77 @@ describe('ShareAliasService', () => {
);
});
it('throws 409 with current target when name is taken and not confirmed', async () => {
const { service, shareAliasRepo, pageRepo } = makeService();
it('throws 409 with the target TITLE (never its id) when the requester CAN view it', async () => {
const { service, shareAliasRepo, pageRepo, pageAccessService } =
makeService();
shareAliasRepo.findByAliasAndWorkspace.mockResolvedValue({
id: 'a-1',
alias: 'foo',
pageId: 'p-other',
});
pageRepo.findById.mockResolvedValue({ id: 'p-other', title: 'Other' });
pageAccessService.validateCanView.mockResolvedValue(undefined); // can view
try {
await service.setAlias({
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'foo',
});
fail('expected ConflictException');
} catch (err) {
expect(err).toBeInstanceOf(ConflictException);
expect((err as ConflictException).getResponse()).toMatchObject({
const body = (err as ConflictException).getResponse();
expect(body).toMatchObject({
code: 'ALIAS_REASSIGN_REQUIRED',
currentPageId: 'p-other',
currentPageTitle: 'Other',
});
// SECURITY (#495): the page id is NEVER disclosed, even to a viewer.
expect(body).not.toHaveProperty('currentPageId');
expect(pageAccessService.validateCanView).toHaveBeenCalledWith(
expect.objectContaining({ id: 'p-other' }),
USER,
);
}
expect(shareAliasRepo.updatePageId).not.toHaveBeenCalled();
});
it('throws 409 WITHOUT the title or id when the requester CANNOT view the target (#495)', async () => {
const { service, shareAliasRepo, pageRepo, pageAccessService } =
makeService();
shareAliasRepo.findByAliasAndWorkspace.mockResolvedValue({
id: 'a-1',
alias: 'foo',
pageId: 'p-secret',
});
pageRepo.findById.mockResolvedValue({ id: 'p-secret', title: 'Secret' });
// No view permission on the target page -> validateCanView throws.
pageAccessService.validateCanView.mockRejectedValue(
new ForbiddenException(),
);
try {
await service.setAlias({
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'foo',
});
fail('expected ConflictException');
} catch (err) {
expect(err).toBeInstanceOf(ConflictException);
const body = (err as ConflictException).getResponse() as Record<
string,
unknown
>;
expect(body).toMatchObject({ code: 'ALIAS_REASSIGN_REQUIRED' });
// The enumeration hole: neither the id nor the title of a page the
// requester cannot see may leak.
expect(body).not.toHaveProperty('currentPageId');
expect(body.currentPageTitle ?? null).toBeNull();
}
expect(shareAliasRepo.updatePageId).not.toHaveBeenCalled();
});
@@ -231,6 +294,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'foo',
confirmReassign: true,
});
@@ -269,6 +333,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'foo',
});
fail('expected ConflictException');
@@ -294,6 +359,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'foo',
});
fail('expected ConflictException');
@@ -317,6 +383,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'foo',
});
fail('expected ConflictException');
@@ -346,6 +413,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'foo',
});
fail('expected ConflictException');
@@ -375,6 +443,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'foo',
confirmReassign: true,
});
@@ -406,6 +475,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'ted',
});
fail('expected ConflictException');
@@ -428,6 +498,7 @@ describe('ShareAliasService', () => {
workspaceId: 'ws-1',
pageId: 'p-1',
creatorId: 'u-1',
user: USER,
alias: 'foo',
}),
).rejects.toBeInstanceOf(BadRequestException);
@@ -450,18 +521,22 @@ describe('ShareAliasService', () => {
alias: 'free-name',
valid: true,
available: true,
currentPageId: null,
});
// SECURITY (#495): the availability probe must NOT leak any page id.
expect(res).not.toHaveProperty('currentPageId');
});
it('reports taken with the current target page', async () => {
it('reports taken WITHOUT leaking the current target page id (#495)', async () => {
const { service, shareAliasRepo } = makeService();
shareAliasRepo.findByAliasAndWorkspace.mockResolvedValue({
id: 'a-1',
pageId: 'p-9',
});
const res = await service.checkAvailability('taken', 'ws-1');
expect(res).toMatchObject({ available: false, currentPageId: 'p-9' });
expect(res).toMatchObject({ available: false });
// The row exists (available:false) but its pageId is never returned — an
// authenticated member cannot map an alias name to a page id it can't view.
expect(res).not.toHaveProperty('currentPageId');
});
});
apps/server/src/core/share/share-alias.service.ts
+42 -10
View File
@@ -7,7 +7,8 @@ import {
import { ShareAliasRepo } from '@docmost/db/repos/share-alias/share-alias.repo';
import { PageRepo } from '@docmost/db/repos/page/page.repo';
import { ShareService } from './share.service';
import { Page, ShareAlias } from '@docmost/db/types/entity.types';
import { PageAccessService } from '../page/page-access/page-access.service';
import { Page, ShareAlias, User } from '@docmost/db/types/entity.types';
import { isValidShareAlias, normalizeShareAlias } from './share-alias.util';
import { InjectKysely } from 'nestjs-kysely';
import { KyselyDB } from '@docmost/db/types/kysely.types';
@@ -43,6 +44,7 @@ export class ShareAliasService {
private readonly shareAliasRepo: ShareAliasRepo,
private readonly pageRepo: PageRepo,
private readonly shareService: ShareService,
private readonly pageAccessService: PageAccessService,
@InjectKysely() private readonly db: KyselyDB,
) {}
@@ -55,9 +57,13 @@ export class ShareAliasService {
* `/l/<old>` link survives
* - name already points at pageId -> no-op (idempotent)
* - name points at ANOTHER page -> the "swap". Without confirmReassign
* we throw 409 carrying the current target so the client can confirm;
* with it we UPDATE the single row's page_id (every /l/<alias> link
* follows the 302 to the new page instantly — no stale cache).
* we throw 409 so the client can confirm. SECURITY (#495): the 409 reveals
* the current target's title ONLY when `user` may VIEW that page, and never
* its id — otherwise any member with one editable+shared page could iterate
* alias names with confirmReassign=false and map them to (id, title) of
* pages they cannot see. With confirmReassign we UPDATE the single row's
* page_id (every /l/<alias> link follows the 302 to the new page instantly
* — no stale cache).
*
* To keep the invariant self-healing we DELETE every other alias row still
* pointing at this page (a legacy duplicate, or the target page's own former
@@ -77,8 +83,12 @@ export class ShareAliasService {
creatorId: string;
alias: string;
confirmReassign?: boolean;
// The requesting user — used ONLY to gate whether the reassign 409 may reveal
// the current target page's title (view-permission check). Not an authz gate
// for the write itself (the controller already validated edit on `pageId`).
user: User;
}): Promise<ShareAlias> {
const { workspaceId, pageId, creatorId, confirmReassign } = opts;
const { workspaceId, pageId, creatorId, confirmReassign, user } = opts;
const alias = normalizeShareAlias(opts.alias);
if (!isValidShareAlias(alias)) {
throw new BadRequestException(
@@ -97,14 +107,30 @@ export class ShareAliasService {
// The name is occupied by a DIFFERENT (or dangling) target page.
if (byName && byName.pageId !== pageId) {
if (!confirmReassign) {
// SECURITY (#495): only disclose the current target's TITLE, and only
// when the requester may VIEW that page. Never disclose its id (the
// client's confirm-reassign UX doesn't use it, and it is an enumerable
// identity). A member with one editable+shared page must NOT be able to
// iterate alias names and map them to (id, title) of pages they cannot
// see. When view is denied (or the alias is dangling) the 409 is the
// bare "occupied" fact — the client still shows a generic confirm modal.
const currentPage = byName.pageId
? await this.pageRepo.findById(byName.pageId)
: null;
let currentPageTitle: string | null = null;
if (currentPage) {
try {
await this.pageAccessService.validateCanView(currentPage, user);
currentPageTitle = currentPage.title ?? null;
} catch {
// No view permission on the target -> do not reveal its title.
currentPageTitle = null;
}
}
throw new ConflictException({
message: 'Alias already in use',
code: 'ALIAS_REASSIGN_REQUIRED',
currentPageId: byName.pageId,
currentPageTitle: currentPage?.title ?? null,
currentPageTitle,
});
}
// Confirmed swap. ORDER MATTERS: the partial unique index on
@@ -223,21 +249,27 @@ export class ShareAliasService {
alias: string;
valid: boolean;
available: boolean;
currentPageId: string | null;
}> {
const alias = normalizeShareAlias(rawAlias);
if (!isValidShareAlias(alias)) {
return { alias, valid: false, available: false, currentPageId: null };
return { alias, valid: false, available: false };
}
const existing = await this.shareAliasRepo.findByAliasAndWorkspace(
alias,
workspaceId,
);
// SECURITY (#495): return ONLY the boolean availability. The previous shape
// leaked `currentPageId` — the id of whatever page the alias already targets —
// to ANY authenticated workspace member, with no view-permission check on that
// page. An attacker could enumerate alias names and map them to page ids they
// have no access to. The taken/free bit is all the "is this address free"
// probe needs. The reassign flow (setAlias 409) may surface the target's
// TITLE, but only behind a `validateCanView` on that page (see setAlias); it
// never returns the page id.
return {
alias,
valid: true,
available: !existing,
currentPageId: existing?.pageId ?? null,
};
}
apps/server/src/core/telemetry/client-metrics.constants.ts
+56 -2
View File
@@ -24,6 +24,54 @@ export const ALLOWED_RATINGS = new Set<string>([
'poor',
]);
// The ONLY route labels accepted. The endpoint is anonymous, so an un-checked
// `route` is a free-text write surface (arbitrary high-cardinality strings /
// injected text into the metrics table). The client only ever sends a label from
// a finite template dictionary (`templateRoute`), so we drop anything not in it.
//
// PARITY: this MUST mirror `KNOWN_ROUTE_TEMPLATES` in the client's
// `apps/client/src/lib/telemetry/route-template.ts` (the canonical source). A
// drift means legit client routes get dropped — keep the two in lockstep; the
// client self-consistency test asserts `templateRoute` only emits these values.
export const ALLOWED_ROUTE_TEMPLATES = new Set<string>([
'/',
'other',
// Static routes.
'/home',
'/spaces',
'/favorites',
'/login',
'/forgot-password',
'/password-reset',
'/setup/register',
'/settings/account/profile',
'/settings/account/preferences',
'/settings/workspace',
'/settings/ai',
'/settings/members',
'/settings/groups',
'/settings/spaces',
'/settings/sharing',
// Dynamic templates (slugs/ids are already collapsed to `:param`).
'/share/:shareId/p/:slug',
'/share/p/:slug',
'/share/:shareId',
'/p/:slug',
'/s/:space/p/:slug',
'/s/:space/trash',
'/s/:space',
'/labels/:label',
'/invites/:invitationId',
'/settings/groups/:groupId',
]);
// `attr` is a web-vitals attribution TARGET: a CSS-selector-ish string (an
// element path like `html>body>div#app>button.cta`), never free prose. Constrain
// it to a conservative CSS-selector charset so the anonymous endpoint cannot be
// used to write arbitrary text / PII / markup into the metrics table. A value
// containing anything outside this set is DROPPED (-> null); the event is kept.
export const ATTR_ALLOWED_CHARSET = /^[A-Za-z0-9#.\-_> :()[\]="'*+~,]+$/;
// Max events accepted per batch; the rest are ignored.
export const MAX_EVENTS_PER_BATCH = 50;
@@ -77,14 +125,20 @@ export function sanitizeVitalEvent(
? e.rating
: null;
// route: accept ONLY a known template label (dictionary check), else drop to
// null. The length cap stays as a cheap pre-guard before the Set lookup.
let route: string | null = null;
if (typeof e.route === 'string' && e.route.length > 0) {
route = e.route.slice(0, MAX_ROUTE_LENGTH);
const candidate = e.route.slice(0, MAX_ROUTE_LENGTH);
route = ALLOWED_ROUTE_TEMPLATES.has(candidate) ? candidate : null;
}
// attr: truncate, then accept ONLY if it is a CSS-selector-shaped string
// (charset whitelist); anything with characters outside the set is dropped.
let attr: string | null = null;
if (typeof e.attr === 'string' && e.attr.length > 0) {
attr = e.attr.slice(0, MAX_ATTR_LENGTH);
const candidate = e.attr.slice(0, MAX_ATTR_LENGTH);
attr = ATTR_ALLOWED_CHARSET.test(candidate) ? candidate : null;
}
let docSize: number | null = null;
apps/server/src/core/telemetry/vitals.service.spec.ts
+38
View File
@@ -90,6 +90,44 @@ describe('VitalsService.buildRows', () => {
expect(rows[0].attr).toHaveLength(MAX_ATTR_LENGTH);
});
it('keeps a known route template but DROPS an unknown/free-text route (#495)', () => {
const rows = svc.buildRows(
{
events: [
{ name: 'INP', value: 1, route: '/s/:space/p/:slug' }, // known
{ name: 'INP', value: 2, route: '/s/acme-corp/p/secret-slug' }, // raw path (slugs) — not a template
{ name: 'INP', value: 3, route: 'DROP TABLE client_metrics;--' }, // injected free text
{ name: 'INP', value: 4, route: '/home' }, // known static
],
},
WS,
);
expect(rows.map((r) => r.route)).toEqual([
'/s/:space/p/:slug',
null, // raw path dropped
null, // free text dropped
'/home',
]);
});
it('DROPS an attr that is not a CSS-selector-shaped string (#495)', () => {
const rows = svc.buildRows(
{
events: [
{ name: 'INP', value: 1, attr: 'div#app>button.cta' }, // valid selector
{ name: 'INP', value: 2, attr: 'user@example.com wrote a note' }, // free text / PII
{ name: 'INP', value: 3, attr: '<script>alert(1)</script>' }, // markup
],
},
WS,
);
expect(rows.map((r) => r.attr)).toEqual([
'div#app>button.cta',
null,
null,
]);
});
it('caps the batch at 50 events', () => {
const events = Array.from({ length: 200 }, () => ({ name: 'CLS', value: 1 }));
const rows = svc.buildRows({ events }, WS);