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..525fba9f24 --- /dev/null +++ b/apps/web/src/lib/ai/core/__tests__/frame-log-cursor.test.ts @@ -0,0 +1,369 @@ +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; +/** + * 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() })); + +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 kind = isSeek ? 'seek' as const : isPayload ? 'payload' as const : 'index' as const; + + 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 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(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) }) }; + }, + }, +})); + +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 = []; + 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)]; + + 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 }, + }); + }); + + // ── 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 + // 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..d7d7ab9353 --- /dev/null +++ b/apps/web/src/lib/ai/core/__tests__/remote-frame-follower.test.ts @@ -0,0 +1,548 @@ +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}` }); + +/** + * 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, + 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 AND terminated, serves the rest then ends cleanly', async () => { + mockReadFramesFrom + .mockResolvedValueOnce(read({ nextSeq: 0 })) + .mockResolvedValueOnce(read({ frames: [frame(0), FINISH], nextSeq: 2 })); + mockStatusRow.mockResolvedValue(COMPLETE); + + const { channel } = acquireRemoteChannel('msg-1'); + const seen = watch(channel); + await settle(); + + assert({ + 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), 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, + }); + }); + + 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, empty: true })); + 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..0513e6ca2a --- /dev/null +++ b/apps/web/src/lib/ai/core/frame-log-cursor.ts @@ -0,0 +1,317 @@ +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. + * + * ── 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 + * 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; + +/** + * 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[]; + /** 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); + } + } + + // 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 { + index = await db + .select({ + fromSeq: aiStreamFrames.fromSeq, + frameCount: aiStreamFrames.frameCount, + byteSize: aiStreamFrames.byteSize, + }) + .from(aiStreamFrames) + .where(and( + eq(aiStreamFrames.messageId, messageId), + gte(aiStreamFrames.fromSeq, containing), + )) + .orderBy(asc(aiStreamFrames.fromSeq)) + .limit(MAX_TICK_ROWS); + } catch (error) { + loggers.ai.warn('frame-log-cursor: index read failed', { + messageId, + fromSeq, + error: error instanceof Error ? error.message : 'unknown', + }); + return nothing(fromSeq, false); + } + + // 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 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. + 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 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; + 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, 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 new file mode 100644 index 0000000000..6bf3f53112 --- /dev/null +++ b/apps/web/src/lib/ai/core/remote-frame-follower.ts @@ -0,0 +1,356 @@ +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'; +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; + +/** + * 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. + * + * 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; + /** 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; + 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 }; + } +}; + +/** 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; + 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) { + 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 + // 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. 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) { + 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, + truncated: tail.truncated || tail.empty || !follower.sawTerminator, + }); + 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, + sawTerminator: false, + 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) => {