From 12976b6f93286be72a37a1376aaae986e0a39a45 Mon Sep 17 00:00:00 2001 From: 2witstudios <2witstudios@gmail.com> Date: Sat, 15 Aug 2026 16:48:35 -0500 Subject: [PATCH 1/2] feat(streams): a reply generated on one machine now streams to every machine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The join could classify a cross-instance stream (previous leaf) but had no frames to serve it. Cross-instance delivery was `stream-join-poll-fallback` — whole-array `parts` snapshot replacements at roughly 1200ms. This tails the durable frame log instead: incremental frames at ~300ms. THE ROUTE BODY IS NOT FORKED, and that is the design. A follower produces a real `StreamChannel`, so `subscribe` / `onFrame` / `onEnd`, the SSE encoding, the 20s ping, the 5s permission recheck, teardown, and the `overflow` / `resumeFromSeq` semantics are literally the same code for both sources. The alternative — an `if (remote)` branch in the route — meant two copies of five behaviours, one of which nothing local exercises. - `frame-log-cursor.ts` — `readFramesFrom({messageId, fromSeq})`. TWO CURSORS, not the schema docblock's `from_seq + frame_count > $X`: that predicate is an expression over two columns, so it is not sargable against the `(message_id, from_seq)` PK and would scan the whole log on every tick, per follower. Instead: `max(from_seq) WHERE from_seq <= $X` seeks the containing row, then `from_seq >= that` range-scans forward, and the overshoot is sliced. Contiguity is enforced exactly as `readFrames` does — a row whose `from_seq` is not where the walk expected it is a HOLE, and the walk STOPS and reports `truncated`. It never skips. Folding across a hole streams a confidently-wrong message to a live user who cannot tell. - `remote-frame-follower.ts` — refcounted per messageId, so co-located tabs share one poller rather than making Postgres load scale with viewers. Starts at seq 0 so any subscriber cursor is servable (a channel starting at one reader's cursor answers `overflow` to every reader below it, and `overflow` is a reseed, not a resume). Polls at 250ms — matched to the writer's 64-frame / 200ms / 256KB flush, since polling faster cannot see fresher data — backing off to 1s after 4 empty ticks and resetting on any frame. The session status is read ONLY on ticks that found nothing. Lingers ~5s past the last reader so a reconnect does not cold-start, plus a `STREAM_MAX_LIFETIME_MS` eviction. - THREE HONEST ANSWERS AT TERMINAL, because frames are deleted on the terminal write: log present → serve the rest, then a clean `done`; log released → `{done, aborted, reload}`; row gone → the same, never a clean end. A short bubble that looks whole is the failure being avoided. - `reload` is a new wire field, distinct from `resumeFromSeq` and not to be collapsed into it: `resumeFromSeq` says "ask again from here", `reload` says there is nowhere to resume from. It plumbs `ChannelEnd.truncated` → `consumeStreamJoin` → `endSession(joinFailed: true)`, which already triggers the durable-message reload. - `streamSessionRegistry`'s success path ended `joinFailed: false` unconditionally. Now `!session.delivered || result.reload`. A join that resolved having delivered NOTHING — it opened, the stream was already over — is not authoritative whatever its terminal frame said, and consumers were keeping a debounced seed snapshot that can be shorter than the finished reply. - `X-Stream-Join-Source: local | remote | terminal`. An N=2 smoke test has no other way to PROVE it exercised the follower rather than getting lucky with the load balancer; in production the remote share should sit near (N-1)/N, which makes it a free load-balancer monitor. Out of scope and untouched, per the brief: `stream-join-poll-fallback` stays (it remains a fallback until N=2 has baked), the `parts` / `rawPartsCount` columns stay, and the cross-instance abort path is unchanged. Twelve seeded mutations verified red, including skipping a hole instead of stopping, dropping the containing-row slice, reporting a read failure as an empty log, ending a released log cleanly, treating an unreadable status as terminal, never resetting the poll backoff, a non-idempotent release, removing the linger, and both `joinFailed` regressions. Two follower tests were strengthened after their first versions failed to discriminate. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SQwbEoZweBEBgsGq9wnqXj --- .../[messageId]/__tests__/route.test.ts | 153 ++++++ .../ai/chat/stream-join/[messageId]/route.ts | 66 ++- .../core/__tests__/frame-log-cursor.test.ts | 275 +++++++++++ .../__tests__/remote-frame-follower.test.ts | 450 ++++++++++++++++++ .../ai/core/__tests__/stream-channel.test.ts | 4 +- .../core/__tests__/stream-join-client.test.ts | 2 +- apps/web/src/lib/ai/core/frame-log-cursor.ts | 221 +++++++++ .../src/lib/ai/core/remote-frame-follower.ts | 309 ++++++++++++ apps/web/src/lib/ai/core/stream-channel.ts | 32 +- .../web/src/lib/ai/core/stream-join-client.ts | 14 +- .../__tests__/streamSessionRegistry.test.ts | 45 ++ .../lib/ai/streams/streamSessionRegistry.ts | 27 +- 12 files changed, 1571 insertions(+), 27 deletions(-) create mode 100644 apps/web/src/lib/ai/core/__tests__/frame-log-cursor.test.ts create mode 100644 apps/web/src/lib/ai/core/__tests__/remote-frame-follower.test.ts create mode 100644 apps/web/src/lib/ai/core/frame-log-cursor.ts create mode 100644 apps/web/src/lib/ai/core/remote-frame-follower.ts diff --git a/apps/web/src/app/api/ai/chat/stream-join/[messageId]/__tests__/route.test.ts b/apps/web/src/app/api/ai/chat/stream-join/[messageId]/__tests__/route.test.ts index af81fb6d8a..5a18f174ec 100644 --- a/apps/web/src/app/api/ai/chat/stream-join/[messageId]/__tests__/route.test.ts +++ b/apps/web/src/app/api/ai/chat/stream-join/[messageId]/__tests__/route.test.ts @@ -1,6 +1,7 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { NextResponse } from 'next/server'; import { StreamChannelRegistry } from '@/lib/ai/core/stream-channel-registry'; +import { openStreamChannel } from '@/lib/ai/core/stream-channel'; import type { SessionAuthResult, AuthError } from '@/lib/auth'; // Fresh registry per test — module-level let, updated in beforeEach @@ -38,6 +39,19 @@ vi.mock('@pagespace/lib/audit/audit-log', () => ({ * the registry answers, or nothing does. The cross-instance cases opt in. */ const mockSessionRow = vi.fn<() => Promise>(); + +/** + * The remote follower, mocked at its module boundary. + * + * The route's job is to acquire one, hand its channel to the SAME subscribe/SSE code a local + * channel goes through, and release it on every teardown path. The follower's own polling and + * terminal classification live in `remote-frame-follower.test.ts`. + */ +const mockRemoteRelease = vi.fn(); +const mockAcquireRemoteChannel = vi.fn(); +vi.mock('@/lib/ai/core/remote-frame-follower', () => ({ + acquireRemoteChannel: (messageId: string) => mockAcquireRemoteChannel(messageId), +})); vi.mock('@pagespace/db/db', () => ({ db: { select: () => ({ @@ -127,6 +141,13 @@ describe('GET /api/ai/chat/stream-join/[messageId]', () => { mockCanSubscribeToStream.mockResolvedValue(true); vi.mocked(canUserViewPage).mockResolvedValue(true); mockSessionRow.mockResolvedValue([]); + // Default: a follower that yields an already-finished, empty channel, so a case that does + // not care about the remote path still terminates promptly. + mockAcquireRemoteChannel.mockImplementation((id: string) => { + const channel = openStreamChannel({ messageId: id }); + channel.finish(false); + return { channel, release: mockRemoteRelease }; + }); }); describe('authentication', () => { @@ -260,6 +281,138 @@ describe('GET /api/ai/chat/stream-join/[messageId]', () => { }); }); + /** + * SERVING a remote stream. + * + * The route body is deliberately NOT forked here: a follower tails the durable frame log and + * presents it as an ordinary `StreamChannel`, so the SSE framing, the ping, the recheck, the + * teardown and the overflow semantics below are one code path for both sources. These cases + * pin that the follower is wired in and released, not the follower's own behaviour + * (`remote-frame-follower.test.ts` covers that). + */ + describe('serving a remote stream through the follower', () => { + const remoteRow = () => [{ + channelId: mockPageId, + userId: mockUserId, + displayName: mockDisplayName, + conversationId: mockConversationId, + browserSessionId: mockBrowserSessionId, + status: 'streaming', + }]; + + it('given a remote stream, serves its frames instead of 404ing', async () => { + mockSessionRow.mockResolvedValue(remoteRow()); + mockAcquireRemoteChannel.mockImplementation((id: string) => { + const channel = openStreamChannel({ messageId: id }); + channel.append(textChunk('from the durable log') as never); + channel.finish(false); + return { channel, release: mockRemoteRelease }; + }); + + const response = await GET(makeRequest(), makeContext(mockMessageId)); + const body = await readSSEBody(response); + + expect(response.status).toBe(200); + expect(body).toContain(`data: ${JSON.stringify({ seq: 0, chunk: textChunk('from the durable log') })}\n\n`); + }); + + it('labels the response with where the frames came from', async () => { + mockSessionRow.mockResolvedValue(remoteRow()); + + const response = await GET(makeRequest(), makeContext(mockMessageId)); + + // An N=2 smoke test has no other way to PROVE it exercised the follower rather than + // getting lucky with the load balancer; in production the remote share should sit near + // (N-1)/N. + expect(response.headers.get('X-Stream-Join-Source')).toBe('remote'); + }); + + it('labels a locally-owned join as local', async () => { + testRegistry.open(mockMessageId, mockMeta); + + const response = await GET(makeRequest(), makeContext(mockMessageId)); + + expect(response.headers.get('X-Stream-Join-Source')).toBe('local'); + }); + + it('labels a terminal join as terminal, and follows it rather than 404ing', async () => { + // Frames are deleted on the terminal write, so a terminal row is exactly the case where + // the follower's honest-answer logic is needed — not a case to short-circuit. + mockSessionRow.mockResolvedValue([{ ...remoteRow()[0], status: 'complete' }]); + + const response = await GET(makeRequest(), makeContext(mockMessageId)); + + expect(response.status).toBe(200); + expect(response.headers.get('X-Stream-Join-Source')).toBe('terminal'); + }); + + it('given a truncated end, tells the client to RELOAD rather than sending a bare done', async () => { + mockSessionRow.mockResolvedValue(remoteRow()); + mockAcquireRemoteChannel.mockImplementation((id: string) => { + const channel = openStreamChannel({ messageId: id }); + channel.append(textChunk('a prefix') as never); + channel.finish(false, { truncated: true }); + return { channel, release: mockRemoteRelease }; + }); + + const body = await readSSEBody(await GET(makeRequest(), makeContext(mockMessageId))); + + // A bare `done` would leave a short reply on screen looking whole. `reload` is a different + // answer from `resumeFromSeq`: there is no seq to resume from, only a durable message to + // re-read. + expect(body).toContain('data: {"done":true,"aborted":false,"reload":true}\n\n'); + }); + + it('releases the follower reference when the stream ends', async () => { + mockSessionRow.mockResolvedValue(remoteRow()); + mockAcquireRemoteChannel.mockImplementation((id: string) => { + const channel = openStreamChannel({ messageId: id }); + channel.finish(false); + return { channel, release: mockRemoteRelease }; + }); + + await readSSEBody(await GET(makeRequest(), makeContext(mockMessageId))); + + // A leaked reference keeps a poller hitting Postgres for a reader that has already gone. + expect(mockRemoteRelease).toHaveBeenCalled(); + }); + + it('releases the follower reference when the client disconnects', async () => { + mockSessionRow.mockResolvedValue(remoteRow()); + const controller = new AbortController(); + mockAcquireRemoteChannel.mockImplementation((id: string) => ({ + channel: openStreamChannel({ messageId: id }), + release: mockRemoteRelease, + })); + + const response = await GET(makeRequest(controller.signal), makeContext(mockMessageId)); + const reader = response.body!.getReader(); + controller.abort(); + await reader.read().catch(() => undefined); + + expect(mockRemoteRelease).toHaveBeenCalled(); + }); + + it('never acquires a follower for a locally-owned stream', async () => { + testRegistry.open(mockMessageId, mockMeta); + + await GET(makeRequest(), makeContext(mockMessageId)); + + expect(mockAcquireRemoteChannel).not.toHaveBeenCalled(); + }); + + it('never acquires a follower for a caller who may not subscribe', async () => { + mockSessionRow.mockResolvedValue(remoteRow()); + mockCanSubscribeToStream.mockResolvedValue(false); + + await GET(makeRequest(), makeContext(mockMessageId)); + + // Acquiring first would start a poller — and a DB read loop over another member's private + // conversation — for a request that is about to 404. + expect(mockAcquireRemoteChannel).not.toHaveBeenCalled(); + }); + }); + describe('authorization', () => { it('given a user without view access, should return 403', async () => { testRegistry.open(mockMessageId, mockMeta); diff --git a/apps/web/src/app/api/ai/chat/stream-join/[messageId]/route.ts b/apps/web/src/app/api/ai/chat/stream-join/[messageId]/route.ts index 934b99ebab..bc041743e9 100644 --- a/apps/web/src/app/api/ai/chat/stream-join/[messageId]/route.ts +++ b/apps/web/src/app/api/ai/chat/stream-join/[messageId]/route.ts @@ -5,6 +5,7 @@ import { auditRequest } from '@pagespace/lib/audit/audit-log'; import { parseGlobalChannelId } from '@pagespace/lib/ai/global-channel-id'; import { canSubscribeToStream } from '@/lib/ai/core/stream-subscription-authz'; import { resolveStreamJoinContext } from '@/lib/ai/core/stream-join-context'; +import { acquireRemoteChannel } from '@/lib/ai/core/remote-frame-follower'; export const dynamic = 'force-dynamic'; @@ -131,22 +132,42 @@ export async function GET( // frames to discard — under-skipping duplicated visible text, over-skipping left a silent // permanent gap, and neither was detectable from either side. // - // A `remote` or `terminal` context has no channel to read from YET — serving one is the next - // leaf (the durable log's follower). Until then this is the same 404 the route already gave - // for a registry miss, reached AFTER authorization rather than instead of it. The client's - // poll fallback (`stream-join-poll-fallback.ts`) handles a 404 exactly as it does today, so - // cross-instance delivery is unchanged by this leaf — what changes is that the answer is now - // classified, and that a caller who may not subscribe is refused for the right reason. - if (joinContext.kind !== 'local') { - return NextResponse.json({ error: 'Stream not found' }, { status: 404 }); - } - const channel = joinContext.channel; + // A REMOTE or TERMINAL context is served by a follower that tails the durable frame log and + // presents it as an ordinary `StreamChannel` — so everything below this line (the SSE framing, + // the ping, the recheck, the teardown, the overflow/resumeFromSeq semantics) is ONE code path + // for both sources. Forking the route body here is what this deliberately does not do: it + // would mean two copies of five behaviours, one of which nothing local exercises. + // + // `terminal` is followed too, not short-circuited. The frames are deleted on the terminal + // write, so a terminal row is exactly the case where the follower's honest-answer logic is + // needed: serve whatever the log still holds, then end — with `truncated` when the log was + // already released, which tells the client to reload rather than trust a short reply. + const remote = joinContext.kind === 'local' ? null : acquireRemoteChannel(messageId); + const channel = joinContext.kind === 'local' ? joinContext.channel : remote!.channel; + const joinSource = joinContext.kind === 'local' ? 'local' : joinContext.kind; + // Dropped in EVERY exit below, not only the happy one — a follower reference leaked by an + // early return keeps a poller alive for a reader that never arrived. + const releaseRemote = () => remote?.release(); const requestedFromSeq = Number(new URL(request.url).searchParams.get('fromSeq') ?? '0'); const fromSeq = Number.isFinite(requestedFromSeq) && requestedFromSeq >= 0 ? Math.floor(requestedFromSeq) : 0; + // ONE teardown for both halves of this joiner's hold: the channel subscription, and — when + // the channel is a follower's — this reader's reference to it. Splitting them was the obvious + // shape and the wrong one: every path that unsubscribed would have had to remember the + // release too, and the one that forgot would keep a poller running against Postgres for a + // reader that had already gone. + let unsubscribeChannel: (() => void) | null = null; + let detached = false; + const detach = (): void => { + if (detached) return; + detached = true; + unsubscribeChannel?.(); + releaseRemote(); + }; + const unsubscribe = channel.subscribe({ fromSeq, onFrame: ({ seq, chunk }) => { @@ -163,9 +184,18 @@ export async function GET( // An `overflow` end means the cursor names a frame the ring no longer holds. Say so with // the seq that IS available rather than quietly serving a later prefix — the client can // then reseed deliberately instead of rendering a gap it cannot see. + // + // A `truncated` end is the OTHER thing, and they must not be conflated: it says there is + // no resume point at all — the durable log was released or holds a hole — so the only + // correct answer is `reload`, which the client turns into a read of the durably-persisted + // message. Sending a bare `done` there would leave a short reply on screen looking whole. const done = end.reason === 'overflow' ? encoder.encode(`data: ${JSON.stringify({ done: true, resumeFromSeq: end.resumeFromSeq })}\n\n`) - : encodeDoneFrame(end.aborted); + : encoder.encode(`data: ${JSON.stringify( + end.truncated === true + ? { done: true, aborted: end.aborted, reload: true } + : { done: true, aborted: end.aborted }, + )}\n\n`); if (streamController) { streamController.enqueue(done); streamController.close(); @@ -173,12 +203,15 @@ export async function GET( preBuffer.push(done); } streamClosed = true; + detach(); }, }); + unsubscribeChannel = unsubscribe; // finish() deletes entries before notifying subscribers, so subscribe() returns // null for both unknown and already-finished streams. if (unsubscribe === null) { + detach(); return NextResponse.json({ error: 'Stream not found' }, { status: 404 }); } @@ -202,7 +235,7 @@ export async function GET( } if (request.signal.aborted) { streamClosed = true; - unsubscribe(); + detach(); controller.close(); return; } @@ -211,14 +244,14 @@ export async function GET( streamClosed = true; clearRecheckTimeout(); clearPingInterval(); - unsubscribe(); + detach(); controller.close(); }, { once: true }); const closeStreamAsDenied = (reason: string) => { streamClosed = true; clearPingInterval(); - unsubscribe(); + detach(); auditRequest(request, { eventType: 'authz.access.denied', resourceType: 'ai_stream', @@ -283,6 +316,11 @@ export async function GET( 'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no', + // WHERE THE FRAMES CAME FROM. An N=2 smoke test has no other way to PROVE it exercised + // the follower rather than getting lucky with the load balancer, and in production the + // remote share should sit near (N-1)/N — a free check that traffic is actually spread and + // that cross-instance joins are being served rather than silently falling back to polling. + 'X-Stream-Join-Source': joinSource, }, }); } diff --git a/apps/web/src/lib/ai/core/__tests__/frame-log-cursor.test.ts b/apps/web/src/lib/ai/core/__tests__/frame-log-cursor.test.ts new file mode 100644 index 0000000000..0a9bd23d61 --- /dev/null +++ b/apps/web/src/lib/ai/core/__tests__/frame-log-cursor.test.ts @@ -0,0 +1,275 @@ +import { describe, it, beforeEach, vi } from 'vitest'; +import { assert } from './riteway'; + +/** + * A REAL in-memory `ai_stream_frames`, driven through the module's own two queries. + * + * The whole point of this leaf is WHICH rows it asks for and how it walks them, so a mock that + * replayed one canned array to both passes would make every contiguity and cursor case vacuous. + * The store below models the table; the db mock routes each query shape to it. + */ +interface Row { messageId: string; fromSeq: number; frameCount: number; frames: unknown[]; byteSize: number } +let table: Row[] = []; +let seekError: Error | null = null; +let rangeError: Error | null = null; + +const { mockLoggerWarn } = vi.hoisted(() => ({ mockLoggerWarn: vi.fn() })); + +vi.mock('@pagespace/db/db', () => ({ + db: { + select: (projection: Record) => { + // `max(from_seq)` — the containing-row seek. Distinguished by the marker the `max` mock + // below returns, so the two selects cannot be confused for one another. + const isSeek = (projection.fromSeq as { max?: unknown } | undefined)?.max !== undefined; + const isPayload = 'frames' in projection; + + const run = (cond: Cond): Row[] | { fromSeq: number | null }[] => { + const rows = table + .filter((r) => r.messageId === cond.messageId) + .filter((r) => (cond.lte === undefined ? true : r.fromSeq <= cond.lte)) + .filter((r) => (cond.gte === undefined ? true : r.fromSeq >= cond.gte)) + .sort((a, b) => a.fromSeq - b.fromSeq); + + if (isSeek) { + if (seekError) throw seekError; + const highest = rows.length === 0 ? null : rows[rows.length - 1].fromSeq; + return [{ fromSeq: highest }]; + } + if (rangeError && isPayload) throw rangeError; + if (seekError && !isPayload) throw seekError; + return rows; + }; + + const chain = (cond: Cond) => ({ + orderBy: () => Object.assign(Promise.resolve(run(cond)), { + limit: (n: number) => Promise.resolve((run(cond) as Row[]).slice(0, n)), + }), + then: (resolve: (rows: unknown) => unknown) => Promise.resolve(run(cond)).then(resolve), + }); + + return { from: () => ({ where: (cond: Cond) => chain(cond) }) }; + }, + }, +})); + +interface Cond { messageId: string; lte?: number; gte?: number } + +/** + * Operators collapsed into the ONE thing the queries express: a messageId plus optional seq + * bounds. Keeps the store above readable without pretending to be a SQL engine. + */ +vi.mock('@pagespace/db/operators', () => ({ + and: (...args: Partial[]) => Object.assign({}, ...args), + eq: (_f: unknown, v: string) => ({ messageId: v }), + lte: (_f: unknown, v: number) => ({ lte: v }), + gte: (_f: unknown, v: number) => ({ gte: v }), + asc: (f: unknown) => f, + max: (f: unknown) => ({ max: f }), +})); + +vi.mock('@pagespace/db/schema/ai-streams', () => ({ + aiStreamFrames: { + messageId: 'message_id', + fromSeq: 'from_seq', + frameCount: 'frame_count', + frames: 'frames', + byteSize: 'byte_size', + }, +})); + +vi.mock('@pagespace/lib/logging/logger-config', () => ({ + loggers: { ai: { info: vi.fn(), warn: mockLoggerWarn, error: vi.fn(), debug: vi.fn() } }, +})); + +import { readFramesFrom } from '../frame-log-cursor'; + +const frame = (n: number) => ({ type: 'text-delta', id: 't1', delta: `f${n}` }); + +const row = (fromSeq: number, count: number, byteSize = 100): Row => ({ + messageId: 'msg-1', + fromSeq, + frameCount: count, + frames: Array.from({ length: count }, (_, i) => frame(fromSeq + i)), + byteSize, +}); + +const deltas = (frames: unknown[]) => frames.map((f) => (f as { delta: string }).delta); + +beforeEach(() => { + vi.clearAllMocks(); + table = []; + seekError = null; + rangeError = null; +}); + +describe('readFramesFrom — the cursor', () => { + it('given a cursor of 0 and a whole log, returns everything in seq order', async () => { + table = [row(0, 3), row(3, 2)]; + + const read = await readFramesFrom({ messageId: 'msg-1', fromSeq: 0 }); + + assert({ + given: 'a fresh follower starting at seq 0', + should: 'return every frame, contiguous, with the next cursor past the end', + actual: { frames: deltas(read.frames), nextSeq: read.nextSeq, truncated: read.truncated }, + expected: { frames: ['f0', 'f1', 'f2', 'f3', 'f4'], nextSeq: 5, truncated: false }, + }); + }); + + it('given a cursor INSIDE a row, slices that row\'s earlier frames off', async () => { + table = [row(0, 4)]; + + const read = await readFramesFrom({ messageId: 'msg-1', fromSeq: 2 }); + + // The containing-row seek deliberately starts the walk BEFORE the cursor — that is what + // makes the query use the `(message_id, from_seq)` primary key instead of scanning on + // `from_seq + frame_count > $X`, which is not sargable. The overshoot is sliced here. + assert({ + given: 'a cursor part-way through a batch', + should: 'return only the frames after it', + actual: { frames: deltas(read.frames), nextSeq: read.nextSeq }, + expected: { frames: ['f2', 'f3'], nextSeq: 4 }, + }); + }); + + it('given a cursor exactly at a row boundary, returns that row whole', async () => { + table = [row(0, 2), row(2, 2)]; + + const read = await readFramesFrom({ messageId: 'msg-1', fromSeq: 2 }); + + assert({ + given: 'a cursor on a batch boundary', + should: 'return that batch and nothing before it', + actual: deltas(read.frames), + expected: ['f2', 'f3'], + }); + }); + + it('given a cursor at the end of the log, returns nothing and holds the cursor', async () => { + table = [row(0, 3)]; + + const read = await readFramesFrom({ messageId: 'msg-1', fromSeq: 3 }); + + assert({ + given: 'a follower that is fully caught up', + should: 'return no frames and leave the cursor where it was', + actual: { frames: read.frames.length, nextSeq: read.nextSeq, truncated: read.truncated, empty: read.empty }, + expected: { frames: 0, nextSeq: 3, truncated: false, empty: false }, + }); + }); + + // ── CONTIGUITY ────────────────────────────────────────────────────────────────────────────── + // + // Folding across a hole does not produce a slightly-wrong message; it produces a confidently + // wrong one — a `tool-output-available` whose `tool-input-start` fell in the gap attaches to + // nothing, and text after the gap concatenates as though the missing tokens were never spoken. + // Here it is being streamed to a LIVE reader, who has no way to tell. + + it('given a HOLE, stops at it and reports truncated — never skips the gap', async () => { + table = [row(0, 2), row(5, 2)]; + + const read = await readFramesFrom({ messageId: 'msg-1', fromSeq: 0 }); + + assert({ + given: 'a log missing seqs 2-4', + should: 'return the contiguous prefix only, and say it is truncated', + actual: { frames: deltas(read.frames), nextSeq: read.nextSeq, truncated: read.truncated }, + expected: { frames: ['f0', 'f1'], nextSeq: 2, truncated: true }, + }); + }); + + it('given a log that BEGINS after the cursor, serves nothing and reports truncated', async () => { + // Not "empty" — the rows exist, they just start past this reader. Serving from the first + // surviving row would hand it a gap it cannot see. + table = [row(4, 2)]; + + const read = await readFramesFrom({ messageId: 'msg-1', fromSeq: 0 }); + + assert({ + given: 'a reader whose cursor predates every surviving row', + should: 'refuse to serve rather than start mid-message', + actual: { frames: read.frames.length, truncated: read.truncated, empty: read.empty }, + expected: { frames: 0, truncated: true, empty: false }, + }); + }); + + it('given no rows at all, reports empty', async () => { + const read = await readFramesFrom({ messageId: 'msg-1', fromSeq: 0 }); + + // `empty` is what lets a follower tell a RELEASED log (the stream ended and retention + // deleted it) from a stream that simply has not flushed since the last tick. + assert({ + given: 'a messageId the log holds nothing for', + should: 'report empty, not truncated', + actual: { empty: read.empty, truncated: read.truncated }, + expected: { empty: true, truncated: false }, + }); + }); + + it('bounds one tick, resuming from where it stopped', async () => { + // The FIRST tick starts at seq 0 and can face the whole log of a long reply. Bounded so + // that read is spread across ticks rather than pulling tens of megabytes into memory at + // once — the difference between a follower and an OOM after a fleet-wide reconnect. + const big = 2 * 1024 * 1024; + table = [row(0, 1, big), row(1, 1, big), row(2, 1, big)]; + + const first = await readFramesFrom({ messageId: 'msg-1', fromSeq: 0 }); + const second = await readFramesFrom({ messageId: 'msg-1', fromSeq: first.nextSeq }); + + assert({ + given: 'a log far larger than one tick\'s budget', + should: 'stop inside the budget without reporting truncation, and continue on the next tick', + actual: { + first: deltas(first.frames), + firstTruncated: first.truncated, + second: deltas(second.frames), + }, + expected: { first: ['f0'], firstTruncated: false, second: ['f1'] }, + }); + }); + + // ── FAILURE MUST NOT LOOK LIKE A RELEASED LOG ─────────────────────────────────────────────── + + it('given the cursor seek fails, reports neither frames nor emptiness', async () => { + table = [row(0, 2)]; + seekError = new Error('db down'); + + const read = await readFramesFrom({ messageId: 'msg-1', fromSeq: 0 }); + + // `empty: false` on failure is load-bearing: a follower reads `empty` as "the log was + // released", and a DB blip that said so would end every viewer's stream early. + assert({ + given: 'a read that threw', + should: 'answer a quiet tick, never "the log is gone"', + actual: { frames: read.frames.length, nextSeq: read.nextSeq, empty: read.empty, truncated: read.truncated }, + expected: { frames: 0, nextSeq: 0, empty: false, truncated: false }, + }); + }); + + it('given the range read fails, degrades the same way', async () => { + table = [row(0, 2)]; + rangeError = new Error('db down'); + + const read = await readFramesFrom({ messageId: 'msg-1', fromSeq: 0 }); + + assert({ + given: 'a payload read that threw', + should: 'answer a quiet tick', + actual: { frames: read.frames.length, empty: read.empty }, + expected: { frames: 0, empty: false }, + }); + }); + + it('scopes every read to the messageId asked for', async () => { + table = [row(0, 2), { ...row(0, 2), messageId: 'msg-other' }]; + + const read = await readFramesFrom({ messageId: 'msg-other', fromSeq: 0 }); + + assert({ + given: 'two messages\' logs in the table', + should: 'read only the one named', + actual: read.frames.length, + expected: 2, + }); + }); +}); diff --git a/apps/web/src/lib/ai/core/__tests__/remote-frame-follower.test.ts b/apps/web/src/lib/ai/core/__tests__/remote-frame-follower.test.ts new file mode 100644 index 0000000000..324f797afb --- /dev/null +++ b/apps/web/src/lib/ai/core/__tests__/remote-frame-follower.test.ts @@ -0,0 +1,450 @@ +import { describe, it, beforeEach, afterEach, vi } from 'vitest'; +import { assert } from './riteway'; +import type { UIMessageChunk } from 'ai'; +import type { ChannelEnd } from '../stream-channel'; +import type { FrameCursorRead } from '../frame-log-cursor'; + +const { mockReadFramesFrom, mockStatusRow, mockLoggerWarn } = vi.hoisted(() => ({ + mockReadFramesFrom: vi.fn(), + mockStatusRow: vi.fn<() => Promise>(), + mockLoggerWarn: vi.fn(), +})); + +/** + * The cursor read is mocked at its module boundary; the CHANNEL is not. + * + * These cases are about the poller's behaviour — cadence, terminal classification, refcounting — + * and the channel it drives is the real one, because "a follower is indistinguishable from a + * local channel" is the property under test rather than an implementation note. + * `frame-log-cursor.test.ts` covers the SQL side. + */ +vi.mock('../frame-log-cursor', () => ({ readFramesFrom: mockReadFramesFrom })); + +vi.mock('@pagespace/db/db', () => ({ + db: { + select: () => ({ + from: () => ({ where: () => ({ limit: () => mockStatusRow() }) }), + }), + }, +})); + +vi.mock('@pagespace/db/operators', () => ({ eq: vi.fn() })); +vi.mock('@pagespace/db/schema/ai-streams', () => ({ + aiStreamSessions: { messageId: 'message_id', status: 'status' }, +})); +vi.mock('@pagespace/lib/logging/logger-config', () => ({ + loggers: { ai: { info: vi.fn(), warn: mockLoggerWarn, error: vi.fn(), debug: vi.fn() } }, +})); + +import { acquireRemoteChannel, resetRemoteFrameFollowers } from '../remote-frame-follower'; + +const frame = (n: number): UIMessageChunk => ({ type: 'text-delta', id: 't1', delta: `f${n}` }); + +const read = (over: Partial = {}): FrameCursorRead => ({ + frames: [], + nextSeq: 0, + truncated: false, + empty: false, + ...over, +}); + +const STREAMING = [{ status: 'streaming' }]; +const COMPLETE = [{ status: 'complete' }]; +const ABORTED = [{ status: 'aborted' }]; + +/** Let the poller's awaits settle without advancing timers. */ +const settle = async () => { + for (let i = 0; i < 30; i += 1) await Promise.resolve(); +}; + +/** Run the next scheduled tick and let it settle. */ +const nextTick = async () => { + await vi.advanceTimersByTimeAsync(1000); + await settle(); +}; + +interface Watcher { frames: unknown[]; end: ChannelEnd | null } + +const watch = (channel: { subscribe: (o: never) => () => void }): Watcher => { + const w: Watcher = { frames: [], end: null }; + channel.subscribe({ + fromSeq: 0, + onFrame: ({ chunk }: { chunk: unknown }) => w.frames.push(chunk), + onEnd: (end: ChannelEnd) => { w.end = end; }, + } as never); + return w; +}; + +beforeEach(() => { + vi.useFakeTimers(); + vi.clearAllMocks(); + resetRemoteFrameFollowers(); + mockReadFramesFrom.mockResolvedValue(read({ empty: true })); + mockStatusRow.mockResolvedValue(STREAMING); +}); + +afterEach(() => { + resetRemoteFrameFollowers(); + vi.useRealTimers(); +}); + +describe('acquireRemoteChannel — following another instance\'s log', () => { + it('appends what the log holds onto a channel indistinguishable from a local one', async () => { + mockReadFramesFrom.mockResolvedValueOnce(read({ frames: [frame(0), frame(1)], nextSeq: 2 })); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + + assert({ + given: 'a remote generation with two durable frames', + should: 'fan them out through the ordinary channel, seq-addressed', + actual: { frames: seen.frames, nextSeq: channel.nextSeq }, + expected: { frames: [frame(0), frame(1)], nextSeq: 2 }, + }); + }); + + it('always starts at seq 0, whatever cursor a subscriber holds', async () => { + acquireRemoteChannel('msg-1'); + await settle(); + + // A channel that began at a subscriber's cursor would answer `overflow` to every OTHER + // subscriber below it — including the common case of a second tab on the same reply — and + // `overflow` is a reseed, not a resume. + assert({ + given: 'a newly acquired follower', + should: 'read from seq 0', + actual: mockReadFramesFrom.mock.calls[0][0], + expected: { messageId: 'msg-1', fromSeq: 0 }, + }); + }); + + it('advances its cursor rather than re-reading what it already served', async () => { + mockReadFramesFrom + .mockResolvedValueOnce(read({ frames: [frame(0)], nextSeq: 1 })) + .mockResolvedValue(read({ nextSeq: 1 })); + + acquireRemoteChannel('msg-1'); + await settle(); + await nextTick(); + + assert({ + given: 'a tick that delivered one frame', + should: 'ask for what comes after it next time', + actual: mockReadFramesFrom.mock.calls[1][0], + expected: { messageId: 'msg-1', fromSeq: 1 }, + }); + }); + + it('reads the session status ONLY on ticks that found nothing', async () => { + mockReadFramesFrom.mockResolvedValue(read({ frames: [frame(0)], nextSeq: 1 })); + + acquireRemoteChannel('msg-1'); + await settle(); + await nextTick(); + + // A tick that found frames has already proven the generation was alive. A status query per + // frame batch would double this follower's query rate for no information at all. + assert({ + given: 'ticks that keep finding frames', + should: 'never query the session row', + actual: mockStatusRow.mock.calls.length, + expected: 0, + }); + }); + + // ── THE THREE HONEST ANSWERS AT TERMINAL ──────────────────────────────────────────────────── + // + // `stream-lifecycle` deletes the frames on its terminal write, so "the row is over" and "the + // frames are still there" are independent facts. Collapsing them is what would make a followed + // reply silently shorter than the one the originating tab saw. + + it('given a terminal row whose log is still present, serves the rest then ends cleanly', async () => { + mockReadFramesFrom + .mockResolvedValueOnce(read({ nextSeq: 0 })) + .mockResolvedValueOnce(read({ frames: [frame(0)], nextSeq: 1 })); + mockStatusRow.mockResolvedValue(COMPLETE); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + + assert({ + given: 'a finished generation whose frames survive', + should: 'deliver the tail and end WITHOUT asking for a reload', + actual: { frames: seen.frames, end: seen.end }, + expected: { frames: [frame(0)], end: { reason: 'finished', aborted: false, truncated: false } }, + }); + }); + + it('given a terminal row whose log was RELEASED, ends asking for a reload', async () => { + mockReadFramesFrom.mockResolvedValue(read({ empty: true })); + mockStatusRow.mockResolvedValue(COMPLETE); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + + // Retention deleted the frames once the terminal `messages` row was confirmed. The reply is + // durably saved and this follower cannot prove it delivered all of it — so it says so, + // rather than handing over a clean end for a possibly-short bubble. + assert({ + given: 'a finished generation whose frames retention already reclaimed', + should: 'end with truncated, which the route turns into `reload`', + actual: seen.end, + expected: { reason: 'finished', aborted: false, truncated: true }, + }); + }); + + it('given an aborted row, carries aborted through', async () => { + mockReadFramesFrom + .mockResolvedValueOnce(read({ nextSeq: 0 })) + .mockResolvedValueOnce(read({ nextSeq: 0 })); + mockStatusRow.mockResolvedValue(ABORTED); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + + assert({ + given: 'a generation that was Stopped', + should: 'report aborted, the same distinction a local end makes', + actual: seen.end?.aborted, + expected: true, + }); + }); + + it('given the row is GONE, ends as truncated rather than as a clean finish', async () => { + mockReadFramesFrom.mockResolvedValue(read({ empty: true })); + mockStatusRow.mockResolvedValue([]); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + + // A conversation hard-deleted mid-follow, or a retention sweep. Not "finished", and nothing + // durable remains to reload from either — but a clean end would leave a partial bubble + // looking whole. + assert({ + given: 'a session row that vanished mid-follow', + should: 'end truncated', + actual: { aborted: seen.end?.aborted, truncated: seen.end?.truncated }, + expected: { aborted: true, truncated: true }, + }); + }); + + it('given an UNREADABLE status, keeps following rather than declaring the stream over', async () => { + mockReadFramesFrom.mockResolvedValue(read({ nextSeq: 0 })); + mockStatusRow.mockRejectedValue(new Error('db down')); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + + // Ending on a DB blip would tell every viewer the reply had finished while it was still + // being generated — the exact failure this workstream exists to end. + assert({ + given: 'a status read that threw', + should: 'leave the channel open', + actual: { end: seen.end, finished: channel.finished }, + expected: { end: null, finished: false }, + }); + }); + + it('given a HOLE in the log, ends truncated and keeps the valid prefix', async () => { + mockReadFramesFrom.mockResolvedValueOnce(read({ frames: [frame(0)], nextSeq: 1, truncated: true })); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + + assert({ + given: 'a log the reader cannot fold past', + should: 'keep what came before the hole and end truncated — never fold across it', + actual: { frames: seen.frames, truncated: seen.end?.truncated }, + expected: { frames: [frame(0)], truncated: true }, + }); + }); + + // ── CADENCE ───────────────────────────────────────────────────────────────────────────────── + + it('backs off after consecutive empty ticks', async () => { + mockReadFramesFrom.mockResolvedValue(read({ nextSeq: 0 })); + + acquireRemoteChannel('msg-1'); + await settle(); + + // Four empty ticks at the fast cadence, then the fifth is scheduled slowly. + for (let i = 0; i < 4; i += 1) { + await vi.advanceTimersByTimeAsync(250); + await settle(); + } + const beforeIdleWindow = mockReadFramesFrom.mock.calls.length; + await vi.advanceTimersByTimeAsync(250); + await settle(); + + assert({ + given: 'a stream sitting in a long tool call', + should: 'stop polling four times a second', + actual: mockReadFramesFrom.mock.calls.length, + expected: beforeIdleWindow, + }); + }); + + it('returns to the fast cadence the moment a frame arrives', async () => { + mockReadFramesFrom.mockResolvedValue(read({ nextSeq: 0 })); + + acquireRemoteChannel('msg-1'); + await settle(); + // Sink well into the idle cadence. + for (let i = 0; i < 6; i += 1) { + await vi.advanceTimersByTimeAsync(1000); + await settle(); + } + + // One frame — the stream is chatty again. + mockReadFramesFrom.mockResolvedValueOnce(read({ frames: [frame(0)], nextSeq: 1 })); + await vi.advanceTimersByTimeAsync(1000); + await settle(); + + // …then one quiet tick. THIS is where the reset shows: without it the empty counter is + // still above the threshold and the follower would schedule the NEXT read a second out + // instead of 250ms out, staying four times slower for the rest of a chatty reply. + await vi.advanceTimersByTimeAsync(250); + await settle(); + const afterFirstQuietTick = mockReadFramesFrom.mock.calls.length; + + await vi.advanceTimersByTimeAsync(250); + await settle(); + + assert({ + given: 'a stream that went quiet, spoke again, then paused for one tick', + should: 'still be polling at the fast cadence — the backoff resets on any frame', + actual: mockReadFramesFrom.mock.calls.length > afterFirstQuietTick, + expected: true, + }); + }); + + // ── REFCOUNTING ───────────────────────────────────────────────────────────────────────────── + + it('given co-located readers, shares ONE poller and one channel', async () => { + mockReadFramesFrom.mockResolvedValue(read({ nextSeq: 0 })); + + const a = acquireRemoteChannel('msg-1'); + const b = acquireRemoteChannel('msg-1'); + await settle(); + + // Without this, N tabs on one instance is N pollers against the same rows, and a follower's + // load on Postgres scales with viewers rather than with streams. + assert({ + given: 'two tabs on this instance watching one remote reply', + should: 'hand both the same channel and start one poller', + actual: { sameChannel: a.channel === b.channel, initialReads: mockReadFramesFrom.mock.calls.length }, + expected: { sameChannel: true, initialReads: 1 }, + }); + }); + + it('keeps polling while any reader remains', async () => { + mockReadFramesFrom.mockResolvedValue(read({ nextSeq: 0 })); + + const a = acquireRemoteChannel('msg-1'); + acquireRemoteChannel('msg-1'); + await settle(); + a.release(); + + // WELL PAST the linger window, twice — a follower that ignored the refcount would have + // started lingering on this release and stopped, so the second window is what discriminates. + // Checking only the first would pass on the broken version, since the linger itself keeps + // polling for five seconds. + await vi.advanceTimersByTimeAsync(20_000); + await settle(); + const afterFirstWindow = mockReadFramesFrom.mock.calls.length; + await vi.advanceTimersByTimeAsync(20_000); + await settle(); + + assert({ + given: 'one of two readers leaving', + should: 'go on following indefinitely for the other, not linger and stop', + actual: mockReadFramesFrom.mock.calls.length > afterFirstWindow, + expected: true, + }); + }); + + it('lingers after the last reader leaves, then stops', async () => { + mockReadFramesFrom.mockResolvedValue(read({ nextSeq: 0 })); + + const a = acquireRemoteChannel('msg-1'); + await settle(); + a.release(); + + // A reconnect — a network blip, a reseed, a pane remounting — would otherwise pay for a cold + // read from seq 0 for a stream this instance was following a moment ago. + await vi.advanceTimersByTimeAsync(1000); + await settle(); + const duringLinger = mockReadFramesFrom.mock.calls.length; + + await vi.advanceTimersByTimeAsync(30_000); + await settle(); + const afterLinger = mockReadFramesFrom.mock.calls.length; + await vi.advanceTimersByTimeAsync(30_000); + await settle(); + + assert({ + given: 'the last reader leaving', + should: 'keep polling briefly, then stop', + actual: { + polledDuringLinger: duringLinger > 1, + stoppedAfter: mockReadFramesFrom.mock.calls.length === afterLinger, + }, + expected: { polledDuringLinger: true, stoppedAfter: true }, + }); + }); + + it('given a reader returning during the linger, keeps the warm follower', async () => { + mockReadFramesFrom.mockResolvedValue(read({ nextSeq: 0 })); + + const a = acquireRemoteChannel('msg-1'); + await settle(); + a.release(); + const b = acquireRemoteChannel('msg-1'); + + await vi.advanceTimersByTimeAsync(30_000); + await settle(); + const before = mockReadFramesFrom.mock.calls.length; + await vi.advanceTimersByTimeAsync(2000); + await settle(); + + assert({ + given: 'a reconnect inside the linger window', + should: 'cancel the linger and go on polling, rather than cold-starting from seq 0', + actual: { sameChannel: a.channel === b.channel, stillPolling: mockReadFramesFrom.mock.calls.length > before }, + expected: { sameChannel: true, stillPolling: true }, + }); + }); + + it('release is idempotent, so a double teardown cannot orphan another reader\'s poller', async () => { + mockReadFramesFrom.mockResolvedValue(read({ nextSeq: 0 })); + + const a = acquireRemoteChannel('msg-1'); + const b = acquireRemoteChannel('msg-1'); + await settle(); + a.release(); + // A double release that decremented twice would take the refcount to zero and strand the + // reader that never let go. + a.release(); + + await vi.advanceTimersByTimeAsync(20_000); + await settle(); + const afterFirstWindow = mockReadFramesFrom.mock.calls.length; + await vi.advanceTimersByTimeAsync(20_000); + await settle(); + + assert({ + given: 'one reader releasing twice while another still holds the follower', + should: 'keep following for the reader that remains', + actual: mockReadFramesFrom.mock.calls.length > afterFirstWindow, + expected: true, + }); + b.release(); + }); +}); diff --git a/apps/web/src/lib/ai/core/__tests__/stream-channel.test.ts b/apps/web/src/lib/ai/core/__tests__/stream-channel.test.ts index 7163068321..66571dc057 100644 --- a/apps/web/src/lib/ai/core/__tests__/stream-channel.test.ts +++ b/apps/web/src/lib/ai/core/__tests__/stream-channel.test.ts @@ -125,7 +125,7 @@ describe('stream-channel — seq and fan-out', () => { const sub = collect(channel); expect(sub.frames).toHaveLength(2); - expect(sub.end).toEqual({ reason: 'finished', aborted: false }); + expect(sub.end).toEqual({ reason: 'finished', aborted: false, truncated: false }); }); it('given finish(true), should report aborted to subscribers', () => { @@ -133,7 +133,7 @@ describe('stream-channel — seq and fan-out', () => { const sub = collect(channel); channel.finish(true); - expect(sub.end).toEqual({ reason: 'finished', aborted: true }); + expect(sub.end).toEqual({ reason: 'finished', aborted: true, truncated: false }); }); it('given an append after finish, should refuse it', () => { diff --git a/apps/web/src/lib/ai/core/__tests__/stream-join-client.test.ts b/apps/web/src/lib/ai/core/__tests__/stream-join-client.test.ts index 5d7aabf6df..33ad9687f7 100644 --- a/apps/web/src/lib/ai/core/__tests__/stream-join-client.test.ts +++ b/apps/web/src/lib/ai/core/__tests__/stream-join-client.test.ts @@ -50,7 +50,7 @@ describe('consumeStreamJoin', () => { const result = await consumeStreamJoin('m1', new AbortController().signal, (parts, seq) => seen.push({ parts, seq })); - expect(result).toEqual({ aborted: false, resumeFromSeq: undefined }); + expect(result).toEqual({ aborted: false, resumeFromSeq: undefined, reload: false }); // The caller gets the FULL folded array each time, not a delta — so it writes with // replace semantics and there is no skip count to get wrong. expect(seen.at(-1)).toEqual({ parts: [textPart('hello world')], seq: 2 }); diff --git a/apps/web/src/lib/ai/core/frame-log-cursor.ts b/apps/web/src/lib/ai/core/frame-log-cursor.ts new file mode 100644 index 0000000000..f9c714e85e --- /dev/null +++ b/apps/web/src/lib/ai/core/frame-log-cursor.ts @@ -0,0 +1,221 @@ +import type { UIMessageChunk } from 'ai'; +import { db } from '@pagespace/db/db'; +import { and, asc, eq, gte, lte, max } from '@pagespace/db/operators'; +import { aiStreamFrames } from '@pagespace/db/schema/ai-streams'; +import { loggers } from '@pagespace/lib/logging/logger-config'; + +/** + * INCREMENTAL reads of the durable frame log, from a cursor. + * + * `readFrames` (frame-log.ts) answers "the whole message, once, for a recovery". This answers + * "what is new since seq N", called every few hundred milliseconds while a stream on ANOTHER + * web instance is still generating. Same table, same contiguity rule, different shape of + * question — and keeping them separate is deliberate: the recovery read is allowed to be + * expensive and exhaustive, this one has to be cheap enough to run on a timer. + * + * ── WHY TWO CURSORS AND NOT THE SCHEMA DOCBLOCK'S PREDICATE ───────────────────────────────── + * + * `ai_stream_frames`'s own docblock suggests + * + * WHERE message_id = $1 AND from_seq + frame_count > $X ORDER BY from_seq + * + * which is exact but NOT SARGABLE: `from_seq + frame_count` is an expression over two columns, + * so the `(message_id, from_seq)` primary key cannot be used to seek and Postgres falls back to + * scanning every row of the message's log — on every poll tick, for every follower. + * + * The same answer is available in two indexed steps. `max(from_seq) WHERE from_seq <= $X` finds + * the row that CONTAINS the cursor (a backward index seek, one row), and `from_seq >= that` + * range-scans forward from it. Both use the PK. The first row read may start before the cursor, + * so its leading frames are sliced off — which is the same arithmetic the docblock describes, + * just after an indexable seek rather than instead of one. + * + * ── CONTIGUITY IS ENFORCED EXACTLY AS `readFrames` ENFORCES IT ────────────────────────────── + * + * A row whose `from_seq` is not where the walk expected it is a HOLE, and the walk STOPS there + * and reports `truncated`. It never skips to the next row. + * + * This is not tidiness. Folding across a hole does not produce a slightly-wrong message; it + * produces a confidently-wrong one — a `tool-output-available` whose `tool-input-start` fell in + * the gap attaches to nothing, and text after the gap concatenates as though the missing tokens + * were never spoken. Here that goes further than it does in a recovery: this content is streamed + * to a LIVE user, one frame at a time, and they have no way to tell. Stopping and saying so lets + * the reader reload the durable message instead. + */ + +/** + * Ceiling on what one tick materializes. + * + * A poll tick normally reads one or two rows — whatever the writer flushed since the last one. + * The FIRST tick is different: it starts at seq 0 and can face the whole log of a long reply. + * Bounded so that read is spread across several ticks rather than pulling tens of megabytes into + * one instance's memory at once, which is the difference between a follower and an OOM after a + * fleet-wide reconnect. + * + * Deliberately smaller than `readFrames`'s MAX_READ_BYTES (24 MB): that one had to reconstruct a + * whole message in a single call, and this one simply continues on the next tick 250ms later. + */ +const MAX_TICK_BYTES = 2 * 1024 * 1024; + +export interface FrameCursorRead { + /** Frames from `fromSeq` onward, contiguous, in seq order. */ + frames: UIMessageChunk[]; + /** The cursor to pass next. Equals `fromSeq` when nothing new was found. */ + nextSeq: number; + /** + * The walk stopped at a HOLE. Whatever is in `frames` is still a valid contiguous prefix, but + * the log cannot serve past it and the reader must not wait for more — it must tell the client + * to reload the durable message. + */ + truncated: boolean; + /** + * The log holds NO rows for this messageId at all. + * + * Distinct from "no new rows", and the distinction is what lets a follower tell a released log + * (the stream ended and retention deleted it) from a stream that simply has not flushed since + * the last tick. Reported as `false` on a read FAILURE, so a DB blip is never mistaken for a + * released log. + */ + empty: boolean; +} + +const nothing = (fromSeq: number, empty: boolean): FrameCursorRead => ({ + frames: [], + nextSeq: fromSeq, + truncated: false, + empty, +}); + +/** + * Read the durable log from `fromSeq`. + * + * NEVER THROWS. A follower runs this on a timer against a stream it does not own; a read failure + * is a tick that found nothing, and the next one tries again. It reports `empty: false` in that + * case so the caller cannot read a transient failure as "the log was released". + */ +export const readFramesFrom = async ({ + messageId, + fromSeq, +}: { + messageId: string; + fromSeq: number; +}): Promise => { + // CURSOR 1 — the row containing `fromSeq`. A backward seek on the PK, one row. + let containing: number | null; + try { + const [head] = await db + .select({ fromSeq: max(aiStreamFrames.fromSeq) }) + .from(aiStreamFrames) + .where(and( + eq(aiStreamFrames.messageId, messageId), + lte(aiStreamFrames.fromSeq, fromSeq), + )); + containing = head?.fromSeq ?? null; + } catch (error) { + loggers.ai.warn('frame-log-cursor: cursor seek failed', { + messageId, + fromSeq, + error: error instanceof Error ? error.message : 'unknown', + }); + return nothing(fromSeq, false); + } + + if (containing === null) { + // Nothing at or before the cursor. Two very different situations, and the follower's + // behaviour turns on which: an empty log (never written, or released after the terminal + // message write) versus a log whose earliest row starts AFTER the cursor — a hole at the + // front, which is unservable from here. + try { + const [first] = await db + .select({ fromSeq: aiStreamFrames.fromSeq }) + .from(aiStreamFrames) + .where(eq(aiStreamFrames.messageId, messageId)) + .orderBy(asc(aiStreamFrames.fromSeq)) + .limit(1); + + if (!first) return nothing(fromSeq, true); + + loggers.ai.warn('frame-log-cursor: log begins after the cursor — cannot serve this reader', { + messageId, + fromSeq, + firstAvailableSeq: first.fromSeq, + }); + return { frames: [], nextSeq: fromSeq, truncated: true, empty: false }; + } catch (error) { + loggers.ai.warn('frame-log-cursor: emptiness check failed', { + messageId, + fromSeq, + error: error instanceof Error ? error.message : 'unknown', + }); + return nothing(fromSeq, false); + } + } + + // CURSOR 2 — forward range scan from the containing row. Also a PK seek. + let rows: { fromSeq: number; frameCount: number; frames: unknown[]; byteSize: number }[]; + try { + rows = await db + .select({ + fromSeq: aiStreamFrames.fromSeq, + frameCount: aiStreamFrames.frameCount, + frames: aiStreamFrames.frames, + byteSize: aiStreamFrames.byteSize, + }) + .from(aiStreamFrames) + .where(and( + eq(aiStreamFrames.messageId, messageId), + gte(aiStreamFrames.fromSeq, containing), + )) + .orderBy(asc(aiStreamFrames.fromSeq)); + } catch (error) { + loggers.ai.warn('frame-log-cursor: range read failed', { + messageId, + fromSeq, + error: error instanceof Error ? error.message : 'unknown', + }); + return nothing(fromSeq, false); + } + + const collected: UIMessageChunk[] = []; + // The walk starts at the CONTAINING row's seq, not at `fromSeq` — the leading frames it + // contributes are sliced off at the end. Starting the contiguity check at `fromSeq` instead + // would report a hole for every ordinary mid-row cursor. + let expectedSeq = containing; + let truncated = false; + let readBytes = 0; + + for (const row of rows) { + if (row.fromSeq !== expectedSeq) { + // A HOLE. Stop — never skip. See the module docblock: past this point the fold would be + // confidently wrong rather than merely short, and it is being streamed to a live reader. + loggers.ai.warn('frame-log-cursor: gap in durable frames — stopping at the hole', { + messageId, + expectedSeq, + foundSeq: row.fromSeq, + }); + truncated = true; + break; + } + // Checked AFTER the contiguity test and BEFORE taking the row, so the budget bounds what is + // materialized. Not `truncated`: the walk stopped because it had read enough, not because + // the log is unservable, and the next tick resumes exactly here. + if (readBytes >= MAX_TICK_BYTES && collected.length > 0) break; + + readBytes += row.byteSize; + collected.push(...(row.frames as UIMessageChunk[])); + // Advance by the RECORDED count, exactly as `readFrames` does. Trusting the column is what + // makes a disagreement between it and the payload surface as a gap (and stop the walk) + // rather than silently shifting every subsequent row's seq. + expectedSeq = row.fromSeq + row.frameCount; + } + + // Drop what the containing row contributed before the cursor. + const skip = fromSeq - containing; + const frames = skip > 0 ? collected.slice(skip) : collected; + + return { + frames, + nextSeq: Math.max(fromSeq, expectedSeq), + truncated, + empty: false, + }; +}; diff --git a/apps/web/src/lib/ai/core/remote-frame-follower.ts b/apps/web/src/lib/ai/core/remote-frame-follower.ts new file mode 100644 index 0000000000..608777f249 --- /dev/null +++ b/apps/web/src/lib/ai/core/remote-frame-follower.ts @@ -0,0 +1,309 @@ +import { db } from '@pagespace/db/db'; +import { eq } from '@pagespace/db/operators'; +import { aiStreamSessions } from '@pagespace/db/schema/ai-streams'; +import { loggers } from '@pagespace/lib/logging/logger-config'; +import { openStreamChannel, type StreamChannel } from '@/lib/ai/core/stream-channel'; +import { readFramesFrom } from '@/lib/ai/core/frame-log-cursor'; +import { STREAM_MAX_LIFETIME_MS } from '@/lib/ai/core/stream-horizons'; + +/** + * Follow a generation owned by ANOTHER web instance, by tailing its durable frame log — and + * present it as an ordinary `StreamChannel`. + * + * ── THE SHAPE IS THE WHOLE DESIGN ─────────────────────────────────────────────────────────── + * + * The obvious implementation is a second branch in the SSE route: if remote, poll and write + * frames. That would fork the route body, and the fork would then have to re-implement — and + * keep in step with — the SSE framing, the 20s ping, the 5s permission recheck, the teardown on + * client abort, and the `overflow` / `resumeFromSeq` semantics. Five behaviours, two copies, + * one of which nobody exercises locally. + * + * So this produces a real `StreamChannel` instead: an ordinary `openStreamChannel`, with a + * poller as its only writer where a generation would have had a pump. `subscribe`, `onFrame`, + * `onEnd`, the ring, eviction, backpressure and the end reasons are then LITERALLY the same + * code, and the route cannot tell a followed stream from a local one. + * + * ── IT ALWAYS STARTS AT SEQ 0 ─────────────────────────────────────────────────────────────── + * + * Not at the first subscriber's cursor. A channel that began at seq 40 would answer `overflow` + * to every OTHER subscriber that arrives with a lower cursor — including the common case of a + * second tab on the same reply — and `overflow` is a reseed, not a resume. Starting at 0 makes + * the channel servable from any cursor, at the cost of one initial read that the tick budget + * (`frame-log-cursor.ts`) already bounds. + * + * ── REFCOUNTED PER messageId ──────────────────────────────────────────────────────────────── + * + * Co-located tabs watching one remote reply share ONE poller. Without this, N tabs on one + * instance is N pollers against the same rows, and the load a follower puts on Postgres scales + * with viewers rather than with streams. + */ + +/** + * Poll cadence while frames are arriving. + * + * Matched to the WRITER, not to a latency target. `frame-log-writer` flushes at 64 frames / + * 200ms / 256 KB, so polling faster than its flush interval cannot see fresher data — it only + * spends round trips discovering that nothing changed. 250ms sits just past a flush period, so + * a follower is typically one flush behind, which is where the ~300ms incremental figure comes + * from (against ~1200ms for the whole-array snapshot polling this replaces). + */ +const POLL_INTERVAL_MS = 250; + +/** + * Cadence after the stream goes quiet. + * + * A generation inside a long tool call emits nothing for minutes. At the fast cadence that is + * four wasted queries a second, per followed stream, for the length of the tool call. + */ +const IDLE_POLL_INTERVAL_MS = 1000; + +/** Empty ticks before backing off. Small, so a brief pause between flushes does not slow the + * common case, and the backoff resets on the very next frame. */ +const EMPTY_TICKS_BEFORE_BACKOFF = 4; + +/** + * How long a channel outlives its last subscriber. + * + * A tab reconnecting (a network blip, a `resumeFromSeq` reseed, a pane remounting) would + * otherwise pay for a cold start — a fresh read from seq 0 — for a stream this instance was + * following moments ago. Short enough that an abandoned follower stops polling promptly. + */ +const LINGER_MS = 5000; + +export interface RemoteChannelHandle { + /** Indistinguishable from a locally-owned channel, by construction. */ + channel: StreamChannel; + /** Drop this holder's reference. Idempotent. */ + release: () => void; +} + +interface Follower { + channel: StreamChannel; + refs: number; + cursor: number; + emptyTicks: number; + pollTimer: ReturnType | null; + lingerTimer: ReturnType | null; + evictionTimer: ReturnType; + stopped: boolean; +} + +const followers = new Map(); + +/** The owning instance's own verdict on whether the generation is over. */ +const readTerminalState = async ( + messageId: string, +): Promise<{ present: boolean; terminal: boolean; aborted: boolean }> => { + try { + const [row] = await db + .select({ status: aiStreamSessions.status }) + .from(aiStreamSessions) + .where(eq(aiStreamSessions.messageId, messageId)) + .limit(1); + + if (!row) return { present: false, terminal: true, aborted: true }; + if (row.status === 'streaming') return { present: true, terminal: false, aborted: false }; + return { present: true, terminal: true, aborted: row.status === 'aborted' }; + } catch (error) { + // Unreadable is NOT terminal. Ending the channel on a DB blip would tell every viewer the + // reply had finished while it was still being generated — the exact failure this whole + // workstream exists to end. + loggers.ai.warn('remote-frame-follower: could not read terminal state', { + messageId, + error: error instanceof Error ? error.message : 'unknown', + }); + return { present: true, terminal: false, aborted: false }; + } +}; + +const stop = (messageId: string, follower: Follower): void => { + if (follower.stopped) return; + follower.stopped = true; + if (follower.pollTimer !== null) clearTimeout(follower.pollTimer); + if (follower.lingerTimer !== null) clearTimeout(follower.lingerTimer); + clearTimeout(follower.evictionTimer); + if (followers.get(messageId) === follower) followers.delete(messageId); +}; + +/** + * End the followed channel, and say honestly WHY. + * + * Three answers, and they are genuinely different to the reader. `stream-lifecycle` deletes the + * frames on its terminal write, so "the row is over" and "the frames are still there" are + * independent facts, and collapsing them is what would make a followed reply silently shorter + * than the one the originating tab saw. + */ +const finishFollower = ( + messageId: string, + follower: Follower, + { aborted, truncated }: { aborted: boolean; truncated: boolean }, +): void => { + stop(messageId, follower); + try { + follower.channel.finish(aborted, { truncated }); + } catch (error) { + loggers.ai.warn('remote-frame-follower: channel finish threw', { + messageId, + error: error instanceof Error ? error.message : 'unknown', + }); + } +}; + +const schedule = (messageId: string, follower: Follower, delayMs: number): void => { + if (follower.stopped) return; + follower.pollTimer = setTimeout(() => { + void tick(messageId, follower); + }, delayMs); + follower.pollTimer.unref?.(); +}; + +const tick = async (messageId: string, follower: Follower): Promise => { + if (follower.stopped) return; + + const read = await readFramesFrom({ messageId, fromSeq: follower.cursor }); + if (follower.stopped) return; + + if (read.frames.length > 0) { + for (const chunk of read.frames) follower.channel.append(chunk); + follower.cursor = read.nextSeq; + // Reset on ANY new frame — the backoff is about quiet streams, and a stream that spoke is + // not quiet. Without this a reader that fell into the idle cadence during one long tool call + // would stay four times slower for the rest of a chatty reply. + follower.emptyTicks = 0; + } + + if (read.truncated) { + // A hole, or a log that begins after this channel's cursor. Nothing further can be served + // from here and there is no seq to resume from — the reader must reload the durable message. + // Whatever was appended before the hole stays: it is a valid prefix, and it is what a client + // that disconnected at that seq would have. + finishFollower(messageId, follower, { aborted: false, truncated: true }); + return; + } + + if (read.frames.length > 0) { + // Frames arrived, so the generation was alive as of this read. Deliberately NOT paying for + // the status query — that is the whole reason the status read is on the empty branch only. + schedule(messageId, follower, POLL_INTERVAL_MS); + return; + } + + follower.emptyTicks += 1; + + // Only a tick that found NOTHING asks whether the stream is over. A tick that found frames has + // already answered it, and a status read per frame batch would double this follower's query + // rate for no information at all. + const state = await readTerminalState(messageId); + if (follower.stopped) return; + + if (state.terminal) { + if (!state.present) { + // The row is GONE — a conversation hard-deleted mid-follow, or a retention sweep. Not + // "finished", and not resumable: nothing durable remains to reload from either, so the + // reader is told its copy is incomplete rather than being handed a clean end. + finishFollower(messageId, follower, { aborted: true, truncated: true }); + return; + } + + // The row is terminal. Which of the two honest answers depends on whether the log is STILL + // THERE: `stream-lifecycle` releases the frames once the terminal `messages` row is + // confirmed, so an empty log means the reply is durably saved and this follower cannot + // prove it delivered all of it. + const tail = await readFramesFrom({ messageId, fromSeq: follower.cursor }); + if (follower.stopped) return; + + if (tail.frames.length > 0) { + for (const chunk of tail.frames) follower.channel.append(chunk); + follower.cursor = tail.nextSeq; + } + + finishFollower(messageId, follower, { + aborted: state.aborted, + // `empty` is reported only for a log with no rows AT ALL — never for a failed read (see + // `FrameCursorRead.empty`), so a DB blip cannot masquerade as a released log and send + // every viewer into a needless reload. + truncated: tail.truncated || tail.empty, + }); + return; + } + + schedule( + messageId, + follower, + follower.emptyTicks >= EMPTY_TICKS_BEFORE_BACKOFF ? IDLE_POLL_INTERVAL_MS : POLL_INTERVAL_MS, + ); +}; + +/** + * Take a reference to this messageId's follower, starting one if none is running. + * + * The caller MUST `release()` — the SSE route does so in the same teardown that unsubscribes. + */ +export const acquireRemoteChannel = (messageId: string): RemoteChannelHandle => { + const existing = followers.get(messageId); + if (existing) { + existing.refs += 1; + // Cancel a linger in progress: this stream has a viewer again, and the poller it already + // has is warmer than a cold restart from seq 0 would be. + if (existing.lingerTimer !== null) { + clearTimeout(existing.lingerTimer); + existing.lingerTimer = null; + } + return { channel: existing.channel, release: releaser(messageId, existing) }; + } + + const channel = openStreamChannel({ messageId }); + const follower: Follower = { + channel, + refs: 1, + // ALWAYS 0. See the module docblock: a channel that started at a subscriber's cursor would + // answer `overflow` to every other subscriber below it. + cursor: 0, + emptyTicks: 0, + pollTimer: null, + lingerTimer: null, + // Leak backstop, sharing STREAM_MAX_LIFETIME_MS with the channel registry, the abort + // registry and the heartbeat cap — see stream-horizons.ts for why those must not drift. + // Past this horizon nothing else in the system will serve this stream either, so a follower + // still polling for it is holding a timer and a ring for a stream nobody can act on. + evictionTimer: setTimeout(() => { + const current = followers.get(messageId); + if (current) finishFollower(messageId, current, { aborted: false, truncated: true }); + }, STREAM_MAX_LIFETIME_MS), + stopped: false, + }; + follower.evictionTimer.unref?.(); + followers.set(messageId, follower); + + // First read immediately rather than after a tick: the subscriber that just arrived is + // waiting on the reply's existing content, and 250ms of blank bubble is the one latency a + // reader actually notices. + void tick(messageId, follower); + + return { channel, release: releaser(messageId, follower) }; +}; + +const releaser = (messageId: string, follower: Follower): (() => void) => { + let released = false; + return () => { + if (released) return; + released = true; + follower.refs -= 1; + if (follower.refs > 0 || follower.stopped) return; + + // LINGER rather than stop. A reconnect — a network blip, a reseed, a pane remounting — + // would otherwise pay for a cold read from seq 0 for a stream this instance was following a + // moment ago. + follower.lingerTimer = setTimeout(() => { + const current = followers.get(messageId); + if (current === follower && follower.refs <= 0) stop(messageId, follower); + }, LINGER_MS); + follower.lingerTimer.unref?.(); + }; +}; + +/** Test-only. Module state otherwise leaks across cases. */ +export const resetRemoteFrameFollowers = (): void => { + for (const [messageId, follower] of [...followers.entries()]) stop(messageId, follower); + followers.clear(); +}; diff --git a/apps/web/src/lib/ai/core/stream-channel.ts b/apps/web/src/lib/ai/core/stream-channel.ts index f19711b5f9..dc7a7812ed 100644 --- a/apps/web/src/lib/ai/core/stream-channel.ts +++ b/apps/web/src/lib/ai/core/stream-channel.ts @@ -70,6 +70,20 @@ export interface ChannelEnd { aborted: boolean; /** Set when `reason === 'overflow'`. */ resumeFromSeq?: number; + /** + * What this subscriber received is NOT the whole message, and no resume point can fix it. + * + * Distinct from `overflow`, which names a seq the reader can restart from. This says the + * SOURCE cannot serve the rest at all — the durable log was released while a follower was + * reading it, or it holds a hole. There is nothing to resume from; the only correct answer is + * to reload the durably-persisted message. + * + * Never set by a locally-owned channel, whose ring is the live record by construction. It + * exists for the remote follower (`remote-frame-follower.ts`), and it rides `ChannelEnd` — + * rather than being a special case in the SSE route — precisely so the route stays one code + * path for both sources. + */ + truncated?: boolean; } export interface SubscribeOptions { @@ -133,7 +147,12 @@ export interface StreamChannel { readonly aborted: boolean; readonly subscriberCount: number; append(chunk: UIMessageChunk): void; - finish(aborted: boolean): void; + /** + * End the channel. `truncated` marks an end where what was delivered is provably incomplete + * and unresumable — see `ChannelEnd.truncated`. A generation never passes it; a follower + * reading a released or holed durable log does. + */ + finish(aborted: boolean, options?: { truncated?: boolean }): void; /** Frames still in memory from `fromSeq` onward. For the checkpoint / terminal fold. */ getFrames(fromSeq?: number): UIMessageChunk[]; subscribe(options: SubscribeOptions): () => void; @@ -222,6 +241,7 @@ export const openStreamChannel = (options: StreamChannelOptions): StreamChannel let firstAvailableSeq = 0; let finished = false; let aborted = false; + let truncated = false; const subscribers = new Set(); @@ -270,12 +290,16 @@ export const openStreamChannel = (options: StreamChannelOptions): StreamChannel } }; - const finish = (nextAborted: boolean): void => { + const finish = (nextAborted: boolean, options?: { truncated?: boolean }): void => { if (finished) return; finished = true; aborted = nextAborted; + // Recorded, so a subscriber that joins AFTER the end is told the same thing as one that was + // already attached. A late joiner replaying a truncated log and being handed a clean `done` + // would render the short reply as if it were the whole one. + truncated = options?.truncated === true; for (const subscriber of [...subscribers]) { - endSubscriber(subscriber, { reason: 'finished', aborted: nextAborted }); + endSubscriber(subscriber, { reason: 'finished', aborted: nextAborted, truncated }); } }; @@ -318,7 +342,7 @@ export const openStreamChannel = (options: StreamChannelOptions): StreamChannel // A stream that already ended still replays in full above, then ends — so a client // that joins a beat late gets the whole reply rather than an empty stream. if (!subscriber.ended && finished) { - endSubscriber(subscriber, { reason: 'finished', aborted }); + endSubscriber(subscriber, { reason: 'finished', aborted, truncated }); } return () => { diff --git a/apps/web/src/lib/ai/core/stream-join-client.ts b/apps/web/src/lib/ai/core/stream-join-client.ts index a47ee7fb75..06781b740c 100644 --- a/apps/web/src/lib/ai/core/stream-join-client.ts +++ b/apps/web/src/lib/ai/core/stream-join-client.ts @@ -22,6 +22,17 @@ export class StreamJoinError extends Error { * data: {"seq": , "chunk": } — one raw SDK frame, in seq order * data: {"done": true, "aborted": } — end sentinel * data: {"done": true, "resumeFromSeq": } — cursor too old; reseed from n + * data: {"done": true, "aborted": , "reload": true} + * — what you have is INCOMPLETE and there is + * no seq to resume from. Reload the durable + * message. + * + * `reload` and `resumeFromSeq` are deliberately different answers and must not be collapsed. + * `resumeFromSeq` says "ask again from here" — the content exists, this connection just cannot + * serve it from the cursor given. `reload` says the SOURCE cannot serve the rest at all: a + * cross-instance follower found the durable frame log released (the stream ended and retention + * deleted it) or holed. There is nowhere to resume from, and the durably-persisted message is + * the only complete copy. * * The frames are the SDK's own `UIMessageChunk`s, and they are folded here by the SAME * reduction the server uses, so a joiner's view is the SDK's view rather than a re-derivation. @@ -38,7 +49,7 @@ export async function consumeStreamJoin( signal: AbortSignal, onParts: (parts: UIMessagePart[], seq: number) => void, fromSeq = 0, -): Promise<{ aborted: boolean; resumeFromSeq?: number }> { +): Promise<{ aborted: boolean; resumeFromSeq?: number; reload?: boolean }> { let response: Response; try { response = await fetch( @@ -99,6 +110,7 @@ export async function consumeStreamJoin( return { aborted: (parsed.aborted as boolean | undefined) ?? false, resumeFromSeq: parsed.resumeFromSeq as number | undefined, + reload: parsed.reload === true, }; } const chunk = parsed.chunk as UIMessageChunk | undefined; diff --git a/apps/web/src/lib/ai/streams/__tests__/streamSessionRegistry.test.ts b/apps/web/src/lib/ai/streams/__tests__/streamSessionRegistry.test.ts index bec49a68d0..3ef1907e53 100644 --- a/apps/web/src/lib/ai/streams/__tests__/streamSessionRegistry.test.ts +++ b/apps/web/src/lib/ai/streams/__tests__/streamSessionRegistry.test.ts @@ -169,9 +169,54 @@ describe('endings', () => { await settle(); expect(ends[0]?.channelId).toBe('page-1'); + }); + + it('given a join that RESOLVED but delivered nothing, still reports joinFailed', async () => { + // A join can resolve cleanly having delivered NOTHING: it opened, the stream was already + // over, and the end sentinel was the first thing on the wire. What the store holds then is + // the seeded snapshot — debounced, and possibly SHORTER than the finished reply. Reporting + // `joinFailed: false` told consumers to keep it, and they kept a truncated bubble. + consumeStreamJoin.mockResolvedValue({ aborted: false }); + const ends: StreamSessionEnd[] = []; + onStreamSessionEnd((end) => ends.push(end)); + + openStreamSession(descriptor()); + await settle(); + + expect(ends[0]?.joinFailed).toBe(true); + }); + + it('given a join that delivered frames, reports joinFailed false', async () => { + consumeStreamJoin.mockImplementation(async (_id, _signal, onParts) => { + onParts([{ type: 'text', text: 'hello' }], 0); + return { aborted: false }; + }); + const ends: StreamSessionEnd[] = []; + onStreamSessionEnd((end) => ends.push(end)); + + openStreamSession(descriptor()); + await settle(); + expect(ends[0]?.joinFailed).toBe(false); }); + it('given the server ended the join with reload, reports joinFailed even though frames arrived', async () => { + // The cross-instance case: a follower delivered a real prefix and then found the durable + // frame log released or holed. The content on screen is genuine but incomplete, and there is + // no seq to resume from — only the durably-persisted message is whole. + consumeStreamJoin.mockImplementation(async (_id, _signal, onParts) => { + onParts([{ type: 'text', text: 'partial' }], 0); + return { aborted: false, reload: true }; + }); + const ends: StreamSessionEnd[] = []; + onStreamSessionEnd((end) => ends.push(end)); + + openStreamSession(descriptor()); + await settle(); + + expect(ends[0]?.joinFailed).toBe(true); + }); + it('given a completion for a stream we never watched, reports joinFailed so the consumer reloads', () => { const ends: StreamSessionEnd[] = []; onStreamSessionEnd((end) => ends.push(end)); diff --git a/apps/web/src/lib/ai/streams/streamSessionRegistry.ts b/apps/web/src/lib/ai/streams/streamSessionRegistry.ts index a73a0d8189..76bea05821 100644 --- a/apps/web/src/lib/ai/streams/streamSessionRegistry.ts +++ b/apps/web/src/lib/ai/streams/streamSessionRegistry.ts @@ -148,10 +148,14 @@ export interface StreamSessionEnd { */ aborted?: boolean; /** - * The join never delivered anything authoritative — usually because the stream lives on - * another web instance whose in-process channel registry this one cannot reach. The - * generation was fine; we just could not watch it. The consumer must reload the durably - * persisted message rather than trust whatever partial content the store held. + * The join never delivered a complete, authoritative copy. The generation was fine; we just + * could not watch all of it. The consumer must reload the durably persisted message rather + * than trust whatever partial content the store held. + * + * Three ways to get here, and they are all the same instruction to a consumer: the join + * errored; it resolved having delivered nothing; or the server ended it with `reload`, meaning + * a cross-instance follower found the durable frame log released or holed and could not serve + * the rest. */ joinFailed: boolean; } @@ -451,7 +455,20 @@ export const openStreamSession = (descriptor: StreamSessionDescriptor): void => endSession(descriptor.messageId, { conversationId: descriptor.conversationId, channelId: descriptor.channelId, - joinFailed: false, + // `!delivered`, NOT a flat `false` — and the difference is the whole point of this line. + // + // A join can resolve cleanly having delivered NOTHING: it opened, the stream was already + // over, and the end sentinel was the first and only thing on the wire. What the store + // holds in that case is the seeded snapshot, which is debounced and can be SHORTER than + // the finished reply. Reporting `joinFailed: false` there told every consumer "keep what + // you have, it is authoritative", and they kept a truncated bubble. + // + // `reload` is the server saying the same thing explicitly: a cross-instance follower + // could not serve the whole message (the durable log was released or holed), so whatever + // arrived is a prefix however much of it there was. Neither is a failure — the + // generation was fine — but in both cases the durably-persisted message is the only + // complete copy, and `joinFailed` is exactly the flag that sends consumers to it. + joinFailed: !session.delivered || result.reload === true, }); }) .catch((error: unknown) => { From ff6f77084bf2d7514a9e782cdd4133a7b1201fa0 Mon Sep 17 00:00:00 2001 From: 2witstudios <2witstudios@gmail.com> Date: Sat, 15 Aug 2026 17:26:38 -0500 Subject: [PATCH 2/2] fix(streams): bound the follower's read in SQL, and require proof before ending a stream cleanly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two P1 review findings on the follower. 1. THE TICK BUDGET BOUNDED NOTHING. `readFramesFrom` selected every row from the cursor onward and applied `MAX_TICK_BYTES` while walking the result — so the driver had already materialized and parsed the entire remaining log, up to the writer's 64 MB per-stream ceiling, before any JavaScript ceiling could look at it. A reconnect against a long stream, and especially a fleet-wide reconnect across many streams, allocated that per follower. `readFrames` documents avoiding exactly this; this module reintroduced it and then claimed in a comment that it did not. Now two passes, the same shape `readFrames` uses: metadata only (three integers per row, `LIMIT`ed) decides the contiguous, budgeted prefix, then the payload query fetches `frames` for exactly those rows, bounded by `from_seq <= $lastWanted`. It is also CHEAPER on the common path than the version it replaces: a tick that finds nothing new — most of them, on a stream inside a tool call — now stops after the metadata pass and never touches the `frames` jsonb at all. 2. A SURVIVING LOG IS NOT A COMPLETE ONE. `frame-log-writer` stops early in three documented ways (its pre-write delete failing, a batch insert failing, exhausting the durable budget), and each leaves a valid, contiguous, HOLE-FREE prefix. A follower that delivered such a prefix and then found the row terminal saw `empty: false` and `truncated: false`, sent a clean `done`, and — because frames HAD been delivered — left `joinFailed` false too. The user kept a truncated reply that nothing ever reloaded, with no signal anywhere that it was short. There is no length marker to check, and `raw_parts_count` (the comparator `materializeInterruptedStream` uses for this exact question) is zeroed by the terminal write. But the stream carries its own terminator: the pump appends every SDK frame verbatim, so a log containing `finish`, `abort`, or the pump's synthetic `error` contains everything the generation produced. The follower now requires one before ending cleanly, and asks for a reload otherwise. Tracked as "have we ever seen one" rather than "is the last frame one" — a trailing `message-metadata` would defeat the stricter test, and the two directions are not symmetric: failing to recognise completeness costs one needless reload, falsely claiming it is the silent truncation above. The cursor's existing cases passed against the unbounded query, so four new ones assert on what was FETCHED rather than what was returned. Seven seeded mutations verified red across both fixes, including removing the `<= lastWanted` bound, dropping the metadata LIMIT, removing the completeness proof, and narrowing the terminator set. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SQwbEoZweBEBgsGq9wnqXj --- .../core/__tests__/frame-log-cursor.test.ts | 102 +++++++++++++- .../__tests__/remote-frame-follower.test.ts | 108 ++++++++++++++- apps/web/src/lib/ai/core/frame-log-cursor.ts | 124 ++++++++++++++++-- .../src/lib/ai/core/remote-frame-follower.ts | 67 ++++++++-- 4 files changed, 368 insertions(+), 33 deletions(-) diff --git a/apps/web/src/lib/ai/core/__tests__/frame-log-cursor.test.ts b/apps/web/src/lib/ai/core/__tests__/frame-log-cursor.test.ts index 0a9bd23d61..525fba9f24 100644 --- a/apps/web/src/lib/ai/core/__tests__/frame-log-cursor.test.ts +++ b/apps/web/src/lib/ai/core/__tests__/frame-log-cursor.test.ts @@ -12,6 +12,13 @@ interface Row { messageId: string; fromSeq: number; frameCount: number; frames: let table: Row[] = []; let seekError: Error | null = null; let rangeError: Error | null = null; +/** + * Every query the module issued, in order — so a case can assert on WHAT WAS FETCHED rather + * than only on what was returned. That distinction is the entire point of the byte budget: an + * earlier version applied it while walking a result the driver had already materialized in + * full, which bounded nothing. + */ +let issued: { kind: 'seek' | 'index' | 'payload'; cond: Cond; limit?: number; rows: number }[] = []; const { mockLoggerWarn } = vi.hoisted(() => ({ mockLoggerWarn: vi.fn() })); @@ -22,6 +29,7 @@ vi.mock('@pagespace/db/db', () => ({ // below returns, so the two selects cannot be confused for one another. const isSeek = (projection.fromSeq as { max?: unknown } | undefined)?.max !== undefined; const isPayload = 'frames' in projection; + const kind = isSeek ? 'seek' as const : isPayload ? 'payload' as const : 'index' as const; const run = (cond: Cond): Row[] | { fromSeq: number | null }[] => { const rows = table @@ -40,11 +48,24 @@ vi.mock('@pagespace/db/db', () => ({ return rows; }; + const record = (cond: Cond, limit: number | undefined, rows: unknown[]) => { + issued.push({ kind, cond, limit, rows: rows.length }); + return rows; + }; + const chain = (cond: Cond) => ({ - orderBy: () => Object.assign(Promise.resolve(run(cond)), { - limit: (n: number) => Promise.resolve((run(cond) as Row[]).slice(0, n)), - }), - then: (resolve: (rows: unknown) => unknown) => Promise.resolve(run(cond)).then(resolve), + orderBy: () => Object.assign( + Promise.resolve(null).then(() => record(cond, undefined, run(cond) as unknown[])), + { limit: (n: number) => Promise.resolve(record(cond, n, (run(cond) as Row[]).slice(0, n))) }, + ), + // BOTH handlers forwarded. `await` calls `then(resolve, reject)`, so a thenable that + // drops the second argument swallows every rejection and the await hangs forever — + // which is what a failed-read case looks like from the outside: a 5s timeout, not a + // failure. + then: (resolve: (rows: unknown) => unknown, reject?: (err: unknown) => unknown) => + Promise.resolve(null) + .then(() => record(cond, undefined, run(cond) as unknown[])) + .then(resolve, reject), }); return { from: () => ({ where: (cond: Cond) => chain(cond) }) }; @@ -98,10 +119,15 @@ const deltas = (frames: unknown[]) => frames.map((f) => (f as { delta: string }) beforeEach(() => { vi.clearAllMocks(); table = []; + issued = []; seekError = null; rangeError = null; }); +/** Rows the PAYLOAD query actually pulled out of the database this tick. */ +const payloadRowsFetched = (): number => + issued.filter((q) => q.kind === 'payload').reduce((n, q) => n + q.rows, 0); + describe('readFramesFrom — the cursor', () => { it('given a cursor of 0 and a whole log, returns everything in seq order', async () => { table = [row(0, 3), row(3, 2)]; @@ -206,6 +232,74 @@ describe('readFramesFrom — the cursor', () => { }); }); + // ── THE BUDGET MUST BOUND THE QUERY, NOT THE LOOP ─────────────────────────────────────────── + // + // An earlier version selected every row from the cursor onward and applied MAX_TICK_BYTES + // while walking the result — which bounded nothing, because the driver had already + // materialized and parsed the entire remaining log (up to the writer's 64 MB per-stream + // ceiling) before any JavaScript ceiling could look at it. These cases assert on what was + // FETCHED; the one below only asserts on what was returned, and passed against that version. + + it('fetches payload rows only up to the tick budget, not the whole remaining log', async () => { + const big = 2 * 1024 * 1024; + table = [row(0, 1, big), row(1, 1, big), row(2, 1, big), row(3, 1, big), row(4, 1, big)]; + + await readFramesFrom({ messageId: 'msg-1', fromSeq: 0 }); + + assert({ + given: 'a log far larger than one tick\'s budget', + should: 'pull only the budgeted rows out of the database', + actual: payloadRowsFetched(), + expected: 1, + }); + }); + + it('caps the metadata read in SQL rather than in the walk', async () => { + table = [row(0, 1), row(1, 1)]; + + await readFramesFrom({ messageId: 'msg-1', fromSeq: 0 }); + + // Three integers per row, so this is a formality against a pathological log — but it has to + // be expressed as a LIMIT, in the statement, for the same reason the byte budget does. + assert({ + given: 'the metadata pass', + should: 'carry a row LIMIT', + actual: issued.find((q) => q.kind === 'index')?.limit, + expected: 512, + }); + }); + + it('given a tick with nothing new, never issues the payload query at all', async () => { + table = [row(0, 3)]; + + await readFramesFrom({ messageId: 'msg-1', fromSeq: 3 }); + + // The common case on a stream inside a long tool call. Stopping after the metadata pass is + // what makes the two-pass split cheaper than the single over-fetching query it replaced, + // rather than merely safer. + assert({ + given: 'a follower that is fully caught up', + should: 'stop after the metadata pass', + actual: issued.filter((q) => q.kind === 'payload').length, + expected: 0, + }); + }); + + it('bounds the payload query by an upper seq, not just a lower one', async () => { + table = [row(0, 1, 2 * 1024 * 1024), row(1, 1), row(2, 1)]; + + await readFramesFrom({ messageId: 'msg-1', fromSeq: 0 }); + + // Without the `<= lastWanted` half, the payload query is open-ended again and the budget + // decided nothing. + assert({ + given: 'a budget that admitted only the first row', + should: 'ask the database for exactly that range', + actual: issued.find((q) => q.kind === 'payload')?.cond.lte, + expected: 0, + }); + }); + it('bounds one tick, resuming from where it stopped', async () => { // The FIRST tick starts at seq 0 and can face the whole log of a long reply. Bounded so // that read is spread across ticks rather than pulling tens of megabytes into memory at diff --git a/apps/web/src/lib/ai/core/__tests__/remote-frame-follower.test.ts b/apps/web/src/lib/ai/core/__tests__/remote-frame-follower.test.ts index 324f797afb..d7d7ab9353 100644 --- a/apps/web/src/lib/ai/core/__tests__/remote-frame-follower.test.ts +++ b/apps/web/src/lib/ai/core/__tests__/remote-frame-follower.test.ts @@ -40,6 +40,14 @@ import { acquireRemoteChannel, resetRemoteFrameFollowers } from '../remote-frame const frame = (n: number): UIMessageChunk => ({ type: 'text-delta', id: 't1', delta: `f${n}` }); +/** + * The SDK's own stream terminators, which are what prove a durable log holds a WHOLE + * generation rather than a prefix the writer gave up on. See STREAM_TERMINATOR_FRAMES. + */ +const FINISH = { type: 'finish' } as UIMessageChunk; +const ABORT = { type: 'abort' } as UIMessageChunk; +const ERROR_FRAME = { type: 'error', errorText: 'boom' } as UIMessageChunk; + const read = (over: Partial = {}): FrameCursorRead => ({ frames: [], nextSeq: 0, @@ -159,10 +167,10 @@ describe('acquireRemoteChannel — following another instance\'s log', () => { // frames are still there" are independent facts. Collapsing them is what would make a followed // reply silently shorter than the one the originating tab saw. - it('given a terminal row whose log is still present, serves the rest then ends cleanly', async () => { + it('given a terminal row whose log is still present AND terminated, serves the rest then ends cleanly', async () => { mockReadFramesFrom .mockResolvedValueOnce(read({ nextSeq: 0 })) - .mockResolvedValueOnce(read({ frames: [frame(0)], nextSeq: 1 })); + .mockResolvedValueOnce(read({ frames: [frame(0), FINISH], nextSeq: 2 })); mockStatusRow.mockResolvedValue(COMPLETE); const { channel } = acquireRemoteChannel('msg-1'); @@ -170,10 +178,100 @@ describe('acquireRemoteChannel — following another instance\'s log', () => { await settle(); assert({ - given: 'a finished generation whose frames survive', + given: 'a finished generation whose log survives and carries the stream terminator', should: 'deliver the tail and end WITHOUT asking for a reload', actual: { frames: seen.frames, end: seen.end }, - expected: { frames: [frame(0)], end: { reason: 'finished', aborted: false, truncated: false } }, + expected: { + frames: [frame(0), FINISH], + end: { reason: 'finished', aborted: false, truncated: false }, + }, + }); + }); + + // ── A NON-EMPTY LOG IS NOT A COMPLETE ONE ─────────────────────────────────────────────────── + // + // `frame-log-writer` stops early in three documented ways (pre-write delete failed, batch + // insert failed, durable budget exhausted), and every one leaves a valid, contiguous, + // HOLE-FREE prefix. Without a completeness proof that reads as a clean end: `empty` false, + // `truncated` false, and — because frames were delivered — `joinFailed` false too. The user + // keeps a truncated reply nothing ever reloads. + + it('given a surviving log that never terminated, asks for a reload rather than ending cleanly', async () => { + mockReadFramesFrom + .mockResolvedValueOnce(read({ frames: [frame(0), frame(1)], nextSeq: 2 })) + .mockResolvedValue(read({ nextSeq: 2 })); + mockStatusRow.mockResolvedValue(COMPLETE); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + await nextTick(); + + assert({ + given: 'a writer that gave up partway, leaving a hole-free PREFIX', + should: 'keep the frames but end truncated — a short log must not read as a complete one', + actual: { frames: seen.frames.length, truncated: seen.end?.truncated }, + expected: { frames: 2, truncated: true }, + }); + }); + + it('accepts an abort frame as proof of completeness', async () => { + mockReadFramesFrom + .mockResolvedValueOnce(read({ nextSeq: 0 })) + .mockResolvedValueOnce(read({ frames: [frame(0), ABORT], nextSeq: 2 })); + mockStatusRow.mockResolvedValue(ABORTED); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + + // A stopped generation is still a COMPLETE record of what it produced — the SDK's `abort` + // frame is its terminator, and the client's bubble should not be replaced by a reload. + assert({ + given: 'a Stopped generation whose log carries the abort terminator', + should: 'end aborted but not truncated', + actual: { aborted: seen.end?.aborted, truncated: seen.end?.truncated }, + expected: { aborted: true, truncated: false }, + }); + }); + + it('accepts the pump\'s synthetic error frame as proof of completeness', async () => { + mockReadFramesFrom + .mockResolvedValueOnce(read({ nextSeq: 0 })) + .mockResolvedValueOnce(read({ frames: [frame(0), ERROR_FRAME], nextSeq: 2 })); + mockStatusRow.mockResolvedValue(COMPLETE); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + + // The pump appends this and then stops reading, so it is the last frame there will ever be. + assert({ + given: 'a generation whose SDK stream threw', + should: 'treat the recorded error as the end of the record, not as a truncation', + actual: seen.end?.truncated, + expected: false, + }); + }); + + it('remembers a terminator seen on an EARLIER tick', async () => { + mockReadFramesFrom + .mockResolvedValueOnce(read({ frames: [frame(0), FINISH], nextSeq: 2 })) + .mockResolvedValue(read({ nextSeq: 2 })); + mockStatusRow.mockResolvedValue(COMPLETE); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + await nextTick(); + + // The terminal branch runs on a LATER tick than the one that delivered the terminator, so + // the proof has to be remembered rather than re-derived from the final read. + assert({ + given: 'a terminator delivered before the row went terminal', + should: 'still count as proof when the follower ends', + actual: seen.end?.truncated, + expected: false, }); }); @@ -199,7 +297,7 @@ describe('acquireRemoteChannel — following another instance\'s log', () => { it('given an aborted row, carries aborted through', async () => { mockReadFramesFrom .mockResolvedValueOnce(read({ nextSeq: 0 })) - .mockResolvedValueOnce(read({ nextSeq: 0 })); + .mockResolvedValueOnce(read({ nextSeq: 0, empty: true })); mockStatusRow.mockResolvedValue(ABORTED); const { channel } = acquireRemoteChannel('msg-1'); diff --git a/apps/web/src/lib/ai/core/frame-log-cursor.ts b/apps/web/src/lib/ai/core/frame-log-cursor.ts index f9c714e85e..0513e6ca2a 100644 --- a/apps/web/src/lib/ai/core/frame-log-cursor.ts +++ b/apps/web/src/lib/ai/core/frame-log-cursor.ts @@ -29,6 +29,25 @@ import { loggers } from '@pagespace/lib/logging/logger-config'; * so its leading frames are sliced off — which is the same arithmetic the docblock describes, * just after an indexable seek rather than instead of one. * + * ── THE BUDGET BOUNDS THE QUERY, NOT THE LOOP ─────────────────────────────────────────────── + * + * An earlier version selected every row from the cursor onward and applied `MAX_TICK_BYTES` + * while walking the result. That bounded nothing: the driver had already materialized and + * parsed the entire remaining log — up to the writer's 64 MB per-stream durable ceiling — + * before any JavaScript ceiling could look at it. A reconnect against a long stream, and + * especially a fleet-wide reconnect across many streams, would allocate that per follower + * (review finding — chatgpt-codex-connector, PR #2421). `readFrames` documents avoiding + * exactly this and then this module reintroduced it. + * + * So the payload is never over-fetched. The METADATA pass reads three integers per row — + * negligible whatever the log holds, and `LIMIT`ed besides — and the contiguous prefix and + * the byte budget are both decided from it. The PAYLOAD pass then fetches `frames` for exactly + * the rows that survived that decision, bounded by `from_seq <= $lastWanted`. + * + * It also costs LESS than the version it replaces on the common path. A tick that finds nothing + * new — most of them, on a stream inside a tool call — now stops after the metadata query and + * never issues the payload query at all. + * * ── CONTIGUITY IS ENFORCED EXACTLY AS `readFrames` ENFORCES IT ────────────────────────────── * * A row whose `from_seq` is not where the walk expected it is a HOLE, and the walk STOPS there @@ -56,6 +75,19 @@ import { loggers } from '@pagespace/lib/logging/logger-config'; */ const MAX_TICK_BYTES = 2 * 1024 * 1024; +/** + * Ceiling on how many row HEADERS one tick inspects. + * + * The metadata pass is three integers per row, so this is a formality against a pathological + * log rather than a real constraint — but it makes the first query bounded in the same + * statement rather than in the loop that reads it, which is the whole lesson of `MAX_TICK_BYTES` + * above. Comfortably more rows than `MAX_TICK_BYTES` can admit at any realistic batch size, so + * the byte budget is what actually decides the prefix and this never truncates first in + * practice. Reaching it behaves exactly like reaching the byte budget: the tick stops, and the + * next one resumes from the cursor it left. + */ +const MAX_TICK_ROWS = 512; + export interface FrameCursorRead { /** Frames from `fromSeq` onward, contiguous, in seq order. */ frames: UIMessageChunk[]; @@ -150,14 +182,15 @@ export const readFramesFrom = async ({ } } - // CURSOR 2 — forward range scan from the containing row. Also a PK seek. - let rows: { fromSeq: number; frameCount: number; frames: unknown[]; byteSize: number }[]; + // PASS 1 — METADATA ONLY, forward from the containing row. Three integers per row, `LIMIT`ed: + // this is what makes the budget below a real bound rather than a ceiling applied to something + // the driver has already parsed into memory. + let index: { fromSeq: number; frameCount: number; byteSize: number }[]; try { - rows = await db + index = await db .select({ fromSeq: aiStreamFrames.fromSeq, frameCount: aiStreamFrames.frameCount, - frames: aiStreamFrames.frames, byteSize: aiStreamFrames.byteSize, }) .from(aiStreamFrames) @@ -165,9 +198,10 @@ export const readFramesFrom = async ({ eq(aiStreamFrames.messageId, messageId), gte(aiStreamFrames.fromSeq, containing), )) - .orderBy(asc(aiStreamFrames.fromSeq)); + .orderBy(asc(aiStreamFrames.fromSeq)) + .limit(MAX_TICK_ROWS); } catch (error) { - loggers.ai.warn('frame-log-cursor: range read failed', { + loggers.ai.warn('frame-log-cursor: index read failed', { messageId, fromSeq, error: error instanceof Error ? error.message : 'unknown', @@ -175,15 +209,17 @@ export const readFramesFrom = async ({ return nothing(fromSeq, false); } - const collected: UIMessageChunk[] = []; + // Decide the contiguous, budgeted prefix from the metadata alone. + // // The walk starts at the CONTAINING row's seq, not at `fromSeq` — the leading frames it // contributes are sliced off at the end. Starting the contiguity check at `fromSeq` instead // would report a hole for every ordinary mid-row cursor. let expectedSeq = containing; + let lastWantedSeq = -1; let truncated = false; let readBytes = 0; - for (const row of rows) { + for (const row of index) { if (row.fromSeq !== expectedSeq) { // A HOLE. Stop — never skip. See the module docblock: past this point the fold would be // confidently wrong rather than merely short, and it is being streamed to a live reader. @@ -195,26 +231,86 @@ export const readFramesFrom = async ({ truncated = true; break; } - // Checked AFTER the contiguity test and BEFORE taking the row, so the budget bounds what is - // materialized. Not `truncated`: the walk stopped because it had read enough, not because - // the log is unservable, and the next tick resumes exactly here. - if (readBytes >= MAX_TICK_BYTES && collected.length > 0) break; + // Checked AFTER the contiguity test and BEFORE the row is taken, so the budget decides what + // the payload query will FETCH. Deliberately not `truncated`: the walk stopped because it + // had taken enough for one tick, not because the log is unservable, and the next tick + // resumes exactly here. The first row is always taken, so a single oversized row still makes + // progress rather than stalling the follower forever. + if (readBytes >= MAX_TICK_BYTES && lastWantedSeq >= 0) break; readBytes += row.byteSize; - collected.push(...(row.frames as UIMessageChunk[])); + lastWantedSeq = row.fromSeq; // Advance by the RECORDED count, exactly as `readFrames` does. Trusting the column is what // makes a disagreement between it and the payload surface as a gap (and stop the walk) // rather than silently shifting every subsequent row's seq. expectedSeq = row.fromSeq + row.frameCount; } + // NOTHING NEW TO FETCH. Two ways to land here, and both are ordinary rather than exceptional: + // the very first row was a hole (`lastWantedSeq < 0`), or the admitted rows reach no further + // than the cursor already does (`expectedSeq <= fromSeq`) — which is what EVERY tick on a + // stream sitting in a long tool call looks like. + // + // Returning here is most of what makes the two-pass split cheaper than the single + // over-fetching query it replaced rather than merely safer: the common tick now costs two + // integer reads and never touches the `frames` jsonb at all. + if (lastWantedSeq < 0 || expectedSeq <= fromSeq) { + return { frames: [], nextSeq: fromSeq, truncated, empty: false }; + } + + // PASS 2 — the payload for exactly those rows, and no more. + let rows: { fromSeq: number; frameCount: number; frames: unknown[] }[]; + try { + rows = await db + .select({ + fromSeq: aiStreamFrames.fromSeq, + frameCount: aiStreamFrames.frameCount, + frames: aiStreamFrames.frames, + }) + .from(aiStreamFrames) + .where(and( + eq(aiStreamFrames.messageId, messageId), + gte(aiStreamFrames.fromSeq, containing), + lte(aiStreamFrames.fromSeq, lastWantedSeq), + )) + .orderBy(asc(aiStreamFrames.fromSeq)); + } catch (error) { + loggers.ai.warn('frame-log-cursor: payload read failed', { + messageId, + fromSeq, + error: error instanceof Error ? error.message : 'unknown', + }); + return nothing(fromSeq, false); + } + + // PASS 2 RE-WALKS CONTIGUITY rather than trusting pass 1 to still describe the table — + // the same rule `readFrames` follows, and for the same reason: between the two queries a + // release or the retention backstop can delete this message's rows wholesale. Re-deriving here + // makes pass 1 purely a decision about HOW MUCH to fetch, and leaves this walk the single + // authority on what is contiguous. + const collected: UIMessageChunk[] = []; + let seq = containing; + for (const row of rows) { + if (row.fromSeq !== seq) { + loggers.ai.warn('frame-log-cursor: log changed between the index and payload reads — stopping', { + messageId, + expectedSeq: seq, + foundSeq: row.fromSeq, + }); + truncated = true; + break; + } + collected.push(...(row.frames as UIMessageChunk[])); + seq = row.fromSeq + row.frameCount; + } + // Drop what the containing row contributed before the cursor. const skip = fromSeq - containing; const frames = skip > 0 ? collected.slice(skip) : collected; return { frames, - nextSeq: Math.max(fromSeq, expectedSeq), + nextSeq: Math.max(fromSeq, seq), truncated, empty: false, }; diff --git a/apps/web/src/lib/ai/core/remote-frame-follower.ts b/apps/web/src/lib/ai/core/remote-frame-follower.ts index 608777f249..6bf3f53112 100644 --- a/apps/web/src/lib/ai/core/remote-frame-follower.ts +++ b/apps/web/src/lib/ai/core/remote-frame-follower.ts @@ -1,3 +1,4 @@ +import type { UIMessageChunk } from 'ai'; import { db } from '@pagespace/db/db'; import { eq } from '@pagespace/db/operators'; import { aiStreamSessions } from '@pagespace/db/schema/ai-streams'; @@ -61,6 +62,33 @@ const IDLE_POLL_INTERVAL_MS = 1000; * common case, and the backoff resets on the very next frame. */ const EMPTY_TICKS_BEFORE_BACKOFF = 4; +/** + * The frames that PROVE the log holds a whole generation. + * + * A NON-EMPTY LOG IS NOT A COMPLETE ONE, and conflating them was a real bug. `frame-log-writer` + * has three documented ways to stop early — its pre-write delete failing, a batch insert + * failing, and exhausting the per-stream durable budget — and every one of them leaves a valid, + * contiguous, HOLE-FREE prefix behind. A follower that delivered such a prefix and then found + * the row terminal would see `empty: false` and `truncated: false`, hand the client a clean + * `done`, and — because frames HAD been delivered — leave `joinFailed` false too. The user would + * be left looking at a truncated reply that nothing ever reloaded, with no signal anywhere that + * it was short (review finding — chatgpt-codex-connector, PR #2421). + * + * The log has no length marker to check against, and `raw_parts_count` — the comparator + * `materializeInterruptedStream` uses for exactly this question — is zeroed by the terminal + * write, so it is gone by the time a follower needs it. But the STREAM carries its own + * terminator: `pumpSdkStreamToChannel` appends every SDK frame verbatim, so a log that contains + * one contains everything the generation produced. `finish` ends a normal turn, `abort` a + * stopped one, and `error` is the synthetic frame the pump appends when the SDK stream throws — + * after which it stops reading. Any of the three is proof that capture ran to the end. + * + * Deliberately "have we EVER seen one" rather than "is the last frame one": a trailing + * `message-metadata` frame would defeat the stricter test, and the cost of being wrong differs + * enormously by direction. Failing to recognise completeness costs one needless reload; falsely + * claiming it is the silent truncation above. + */ +const STREAM_TERMINATOR_FRAMES: ReadonlySet = new Set(['finish', 'abort', 'error']); + /** * How long a channel outlives its last subscriber. * @@ -82,6 +110,9 @@ interface Follower { refs: number; cursor: number; emptyTicks: number; + /** A stream terminator has reached this channel — see STREAM_TERMINATOR_FRAMES. Until it has, + * a non-empty log is not evidence of a complete one. */ + sawTerminator: boolean; pollTimer: ReturnType | null; lingerTimer: ReturnType | null; evictionTimer: ReturnType; @@ -116,6 +147,14 @@ const readTerminalState = async ( } }; +/** Append frames and record whether any of them ended the stream. */ +const deliver = (follower: Follower, frames: readonly UIMessageChunk[]): void => { + for (const chunk of frames) { + follower.channel.append(chunk); + if (STREAM_TERMINATOR_FRAMES.has(chunk.type)) follower.sawTerminator = true; + } +}; + const stop = (messageId: string, follower: Follower): void => { if (follower.stopped) return; follower.stopped = true; @@ -164,7 +203,7 @@ const tick = async (messageId: string, follower: Follower): Promise => { if (follower.stopped) return; if (read.frames.length > 0) { - for (const chunk of read.frames) follower.channel.append(chunk); + deliver(follower, read.frames); follower.cursor = read.nextSeq; // Reset on ANY new frame — the backoff is about quiet streams, and a stream that spoke is // not quiet. Without this a reader that fell into the idle cadence during one long tool call @@ -205,24 +244,31 @@ const tick = async (messageId: string, follower: Follower): Promise => { return; } - // The row is terminal. Which of the two honest answers depends on whether the log is STILL - // THERE: `stream-lifecycle` releases the frames once the terminal `messages` row is - // confirmed, so an empty log means the reply is durably saved and this follower cannot - // prove it delivered all of it. + // The row is terminal. One last read, because frames can land between the tick that found + // nothing and the status query that answered. const tail = await readFramesFrom({ messageId, fromSeq: follower.cursor }); if (follower.stopped) return; if (tail.frames.length > 0) { - for (const chunk of tail.frames) follower.channel.append(chunk); + deliver(follower, tail.frames); follower.cursor = tail.nextSeq; } + // A CLEAN END REQUIRES PROOF, not merely the absence of evidence to the contrary. + // + // - `tail.empty` — the log was RELEASED. `stream-lifecycle` deletes the frames once the + // terminal `messages` row is confirmed, so the reply is durably saved and this follower + // cannot show it delivered all of it. (Reported only for a log with no rows AT ALL, + // never for a failed read — see `FrameCursorRead.empty` — so a DB blip cannot + // masquerade as a released log and send every viewer into a needless reload.) + // - `tail.truncated` — a hole, or a log that begins after this cursor. + // - `!sawTerminator` — the log survives and is hole-free, but no `finish`/`abort`/`error` + // ever arrived, so the writer stopped early and this is a PREFIX. See + // STREAM_TERMINATOR_FRAMES: without this clause a short log reads exactly like a + // complete one. finishFollower(messageId, follower, { aborted: state.aborted, - // `empty` is reported only for a log with no rows AT ALL — never for a failed read (see - // `FrameCursorRead.empty`), so a DB blip cannot masquerade as a released log and send - // every viewer into a needless reload. - truncated: tail.truncated || tail.empty, + truncated: tail.truncated || tail.empty || !follower.sawTerminator, }); return; } @@ -260,6 +306,7 @@ export const acquireRemoteChannel = (messageId: string): RemoteChannelHandle => // answer `overflow` to every other subscriber below it. cursor: 0, emptyTicks: 0, + sawTerminator: false, pollTimer: null, lingerTimer: null, // Leak backstop, sharing STREAM_MAX_LIFETIME_MS with the channel registry, the abort