From ea868e5187369479b5be70f5ba0f8c0f33af2fed Mon Sep 17 00:00:00 2001 From: Robel Estifanos Date: Wed, 12 Aug 2026 14:34:35 -0400 Subject: [PATCH] feat(component): route new producers through Stream V6 Keep stable PTS facade IDs while creating native V6 streams for new data, elect one producer through an OCC claim, and mirror lifecycle, timeout, and deletion transactionally. Preserve legacy chunk rows without migration and bound both durable batching and raw response buffering so slow consumers cannot block persistence. --- src/client/index.test.ts | 275 ++++++++++++++++++++ src/client/index.ts | 283 ++++++++++++++++---- src/component/_generated/component.ts | 7 + src/component/lib.ts | 253 +++++++++++++++--- src/component/producer.test.ts | 354 ++++++++++++++++++++++++++ src/component/schema.ts | 13 +- 6 files changed, 1101 insertions(+), 84 deletions(-) create mode 100644 src/client/index.test.ts create mode 100644 src/component/producer.test.ts diff --git a/src/client/index.test.ts b/src/client/index.test.ts new file mode 100644 index 0000000..5fb478a --- /dev/null +++ b/src/client/index.test.ts @@ -0,0 +1,275 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +import { + PersistentTextStreaming, + type StreamId, + type StreamWriter, +} from "./index.js"; + +const streamId = "stream" as StreamId; + +type Call = + | { kind: "claim" } + | { kind: "append"; text: string; final: boolean } + | { kind: "status"; status: string }; + +function setup(claims: boolean[] = [true]) { + const refs = { + lib: { + addChunk: {}, + claimStream: {}, + createStream: {}, + deleteStream: {}, + getStreamStatus: {}, + getStreamText: {}, + setStreamStatus: {}, + }, + }; + const calls: Call[] = []; + const ctx = { + runMutation: vi.fn(async (ref: object, args: Record) => { + if (ref === refs.lib.claimStream) { + calls.push({ kind: "claim" }); + return claims.shift() ?? false; + } + if (ref === refs.lib.addChunk) { + calls.push({ + kind: "append", + text: args.text as string, + final: args.final as boolean, + }); + return null; + } + if (ref === refs.lib.setStreamStatus) { + calls.push({ kind: "status", status: args.status as string }); + return null; + } + throw new Error("Unexpected mutation."); + }), + }; + return { + calls, + ctx, + refs, + streaming: new PersistentTextStreaming(refs as never), + }; +} + +function request(): Request { + return new Request("https://example.com/stream", { method: "POST" }); +} + +afterEach(() => vi.useRealTimers()); + +describe("PersistentTextStreaming producer", () => { + it("preserves one-winner claim behavior", async () => { + const { calls, ctx, streaming } = setup([true, false]); + const writer = vi.fn>( + async (_ctx, _request, _id, append) => { + await append("winner"); + }, + ); + + const winner = await streaming.stream( + ctx as never, + request(), + streamId, + writer, + ); + const loser = await streaming.stream( + ctx as never, + request(), + streamId, + writer, + ); + + expect(loser.status).toBe(205); + expect(await winner.text()).toBe("winner"); + expect(writer).toHaveBeenCalledTimes(1); + expect(calls.filter((call) => call.kind === "claim")).toHaveLength(2); + }); + + it("flushes on the 100ms timer before final durable completion", async () => { + vi.useFakeTimers(); + const { calls, ctx, streaming } = setup(); + let release: (() => void) | undefined; + const gate = new Promise((resolve) => { + release = resolve; + }); + const writer: StreamWriter = async (_ctx, _request, _id, append) => { + await append("timed"); + await gate; + }; + + const response = await streaming.stream( + ctx as never, + request(), + streamId, + writer, + ); + await vi.advanceTimersByTimeAsync(99); + expect(calls.some((call) => call.kind === "append")).toBe(false); + await vi.advanceTimersByTimeAsync(1); + expect(calls).toContainEqual({ + kind: "append", + text: "timed", + final: false, + }); + + release?.(); + expect(await response.text()).toBe("timed"); + expect(calls.at(-1)).toEqual({ kind: "status", status: "done" }); + }); + + it("splits at the 16KiB byte bound and serializes append order", async () => { + const { calls, ctx, streaming } = setup(); + const first = "a".repeat(12 * 1024); + const second = "πŸ™‚".repeat(2 * 1024); + const writer: StreamWriter = async (_ctx, _request, _id, append) => { + await Promise.all([append(first), append(second)]); + }; + + const response = await streaming.stream( + ctx as never, + request(), + streamId, + writer, + ); + expect(await response.text()).toBe(first + second); + + const appends = calls.filter( + (call): call is Extract => + call.kind === "append", + ); + expect(appends.map((call) => call.text).join("")).toBe(first + second); + expect( + appends.every( + (call) => new TextEncoder().encode(call.text).byteLength <= 16 * 1024, + ), + ).toBe(true); + expect(appends.at(-1)?.final).toBe(true); + }); + + it("persists the failure prefix before recording the producer error", async () => { + const { calls, ctx, streaming } = setup(); + const failure = new Error("producer failed"); + const writer: StreamWriter = async (_ctx, _request, _id, append) => { + await append("durable prefix"); + throw failure; + }; + + const response = await streaming.stream( + ctx as never, + request(), + streamId, + writer, + ); + await expect(response.text()).rejects.toThrow("producer failed"); + expect(calls.slice(-2)).toEqual([ + { kind: "append", text: "durable prefix", final: false }, + { kind: "status", status: "error" }, + ]); + }); + + it("finishes durable writes when the raw body is never consumed", async () => { + const { calls, ctx, refs, streaming } = setup(); + let durableDone: (() => void) | undefined; + const done = new Promise((resolve) => { + durableDone = resolve; + }); + ctx.runMutation.mockImplementation( + async (ref: object, args: Record) => { + if (ref === refs.lib.claimStream) { + calls.push({ kind: "claim" }); + return true; + } + if (ref === refs.lib.addChunk) { + calls.push({ + kind: "append", + text: args.text as string, + final: args.final as boolean, + }); + return null; + } + calls.push({ kind: "status", status: args.status as string }); + if (args.status === "done") durableDone?.(); + return null; + }, + ); + const text = "πŸ™‚".repeat(20 * 1024); + const writer: StreamWriter = async (_ctx, _request, _id, append) => { + await append(text); + }; + + const response = await streaming.stream( + ctx as never, + request(), + streamId, + writer, + ); + + await done; + const appends = calls.filter( + (call): call is Extract => + call.kind === "append", + ); + expect(appends.map((call) => call.text).join("")).toBe(text); + expect(calls.at(-1)).toEqual({ kind: "status", status: "done" }); + await expect(response.text()).rejects.toThrow( + "Raw stream consumer fell behind durable replay.", + ); + }); + + it("continues durable completion after the raw reader cancels", async () => { + const { calls, ctx, refs, streaming } = setup(); + let release: (() => void) | undefined; + let persisted: (() => void) | undefined; + const gate = new Promise((resolve) => { + release = resolve; + }); + const done = new Promise((resolve) => { + persisted = resolve; + }); + ctx.runMutation.mockImplementation( + async (ref: object, args: Record) => { + if (ref === refs.lib.claimStream) { + calls.push({ kind: "claim" }); + return true; + } + if (ref === refs.lib.addChunk) { + calls.push({ + kind: "append", + text: args.text as string, + final: args.final as boolean, + }); + if (args.final === true) persisted?.(); + return null; + } + calls.push({ kind: "status", status: args.status as string }); + return null; + }, + ); + const writer: StreamWriter = async (_ctx, _request, _id, append) => { + await append("first."); + await gate; + await append("second"); + }; + + const response = await streaming.stream( + ctx as never, + request(), + streamId, + writer, + ); + const reader = response.body!.getReader(); + await reader.read(); + await reader.cancel(); + release?.(); + await done; + + expect(calls.filter((call) => call.kind === "append")).toEqual([ + { kind: "append", text: "first.", final: false }, + { kind: "append", text: "second", final: true }, + ]); + }); +}); diff --git a/src/client/index.ts b/src/client/index.ts index fea9d22..74125a5 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -25,11 +25,147 @@ export type StreamWriter> = ( chunkAppender: ChunkAppender, ) => Promise; -// TODO -- make more flexible. # of bytes, etc? -const hasDelimeter = (text: string) => { +const FLUSH_MS = 100; +const MAX_BATCH_BYTES = 16 * 1024; +const MAX_RAW_BUFFER_BYTES = 64 * 1024; +const encoder = new TextEncoder(); + +function shouldFlush(text: string): boolean { return text.includes(".") || text.includes("!") || text.includes("?"); +} + +function splitText(text: string): string[] { + const chunks: string[] = []; + let start = 0; + let index = 0; + let bytes = 0; + + for (const character of text) { + const size = encoder.encode(character).byteLength; + if (bytes > 0 && bytes + size > MAX_BATCH_BYTES) { + chunks.push(text.slice(start, index)); + start = index; + bytes = 0; + } + bytes += size; + index += character.length; + } + if (start < text.length) chunks.push(text.slice(start)); + return chunks; +} + +type Sink = { + write(text: string): void; + close(): void; + error(error: unknown): void; }; +class Batch { + private parts: string[] = []; + private bytes = 0; + private timer: ReturnType | null = null; + private queue = Promise.resolve(); + private ended = false; + + constructor( + private readonly appendChunk: ( + text: string, + final: boolean, + ) => Promise, + private readonly setStatus: (status: StreamStatus) => Promise, + private readonly sink: Sink, + ) {} + + add(text: string): Promise { + if (this.ended) { + return Promise.reject(new Error("Stream is already finalized.")); + } + const next = this.queue.then(() => this.accept(text)); + this.queue = next; + return next; + } + + finish(): Promise { + if (this.ended) return this.queue; + this.ended = true; + this.stopTimer(); + const next = this.queue.then(async () => { + this.stopTimer(); + if (this.parts.length === 0) { + await this.setStatus("done"); + } else { + await this.flush(true); + } + }); + this.queue = next; + return next; + } + + fail(): Promise { + this.ended = true; + this.stopTimer(); + const prior = this.queue; + const next = (async () => { + try { + await prior; + this.stopTimer(); + await this.flush(false); + } finally { + await this.setStatus("error"); + } + })(); + this.queue = next; + return next; + } + + private async accept(text: string): Promise { + if (text.length === 0) return; + const chunks = splitText(text); + for (const chunk of chunks) this.sink.write(chunk); + + for (const chunk of chunks) { + const size = encoder.encode(chunk).byteLength; + if (this.bytes > 0 && this.bytes + size > MAX_BATCH_BYTES) { + await this.flush(false); + } + this.parts.push(chunk); + this.bytes += size; + if (this.bytes >= MAX_BATCH_BYTES) await this.flush(false); + } + + if (shouldFlush(text)) await this.flush(false); + else this.startTimer(); + } + + private async flush(final: boolean): Promise { + this.stopTimer(); + if (this.parts.length === 0) { + if (final) await this.setStatus("done"); + return; + } + const text = this.parts.join(""); + this.parts = []; + this.bytes = 0; + await this.appendChunk(text, final); + } + + private startTimer(): void { + if (this.parts.length === 0 || this.timer !== null) return; + this.timer = setTimeout(() => { + this.timer = null; + const next = this.queue.then(() => this.flush(false)); + this.queue = next; + void next.catch(() => undefined); + }, FLUSH_MS); + } + + private stopTimer(): void { + if (this.timer === null) return; + clearTimeout(this.timer); + this.timer = null; + } +} + // TODO -- some sort of wrapper with easy ergonomics for working with LLMs? export class PersistentTextStreaming { constructor( @@ -111,66 +247,109 @@ export class PersistentTextStreaming { streamId: StreamId, streamWriter: StreamWriter, ) { - const streamState = await ctx.runQuery(this.component.lib.getStreamStatus, { + const claimed = await ctx.runMutation(this.component.lib.claimStream, { streamId, }); - if (streamState !== "pending") { + if (!claimed) { console.log("Stream was already started"); - return new Response("", { + return new Response(null, { status: 205, }); } - // Create a TransformStream to handle streaming data - const { readable, writable } = new TransformStream(); - let writer = - writable.getWriter() as WritableStreamDefaultWriter | null; - const textEncoder = new TextEncoder(); - let pending = ""; - - const doStream = async () => { - const chunkAppender: ChunkAppender = async (text) => { - // write to this handler's response stream on every update - if (writer) { + let connected = true; + const readable = new ReadableStream( + { + start: async (controller) => { + const sink: Sink = { + write(text) { + if (!connected) return; + const chunk = encoder.encode(text); + const capacity = controller.desiredSize; + if (capacity === null || capacity < chunk.byteLength) { + connected = false; + try { + controller.error( + new Error( + "Raw stream consumer fell behind durable replay.", + ), + ); + } catch { + // The response consumer has already disconnected. + } + return; + } + try { + controller.enqueue(chunk); + if ((controller.desiredSize ?? 0) <= 0) { + connected = false; + controller.error( + new Error( + "Raw stream consumer fell behind durable replay.", + ), + ); + } + } catch { + connected = false; + } + }, + close() { + if (!connected) return; + connected = false; + try { + controller.close(); + } catch { + // The response consumer has already disconnected. + } + }, + error(error) { + if (!connected) return; + connected = false; + try { + controller.error(error); + } catch { + // The response consumer has already disconnected. + } + }, + }; + const batch = new Batch( + (text, final) => this.addChunk(ctx, streamId, text, final), + (status) => this.setStreamStatus(ctx, streamId, status), + sink, + ); + try { - await writer.write(textEncoder.encode(text)); - } catch (e) { - console.error("Error writing to stream", e); - console.error( - "Will skip writing to stream but continue database updates", + await streamWriter(ctx, request, streamId, (text) => + batch.add(text), ); - writer = null; + await batch.finish(); + sink.close(); + } catch (error) { + let failure = error; + try { + await batch.fail(); + } catch (persistenceError) { + failure = persistenceError; + } + sink.error(failure); } - } - pending += text; - // write to the database periodically, like at the end of sentences - if (hasDelimeter(text)) { - await this.addChunk(ctx, streamId, pending, false); - pending = ""; - } - }; - try { - await streamWriter(ctx, request, streamId, chunkAppender); - } catch (e) { - await this.setStreamStatus(ctx, streamId, "error"); - if (writer) { - await writer.close(); - } - throw e; - } - - // Success? Flush any last updates - await this.addChunk(ctx, streamId, pending, true); - - if (writer) { - await writer.close(); - } - }; - - // Kick off the streaming, but don't await it. - void doStream(); + }, + cancel() { + connected = false; + }, + }, + { + highWaterMark: MAX_RAW_BUFFER_BYTES, + size: (chunk) => chunk.byteLength, + }, + ); - // Send the readable back to the browser - return new Response(readable); + return new Response(readable, { + headers: { + "Cache-Control": "no-cache, no-transform", + "Content-Type": "text/plain; charset=utf-8", + "X-Accel-Buffering": "no", + }, + }); } /** diff --git a/src/component/_generated/component.ts b/src/component/_generated/component.ts index 3eb969f..4c7fdbd 100644 --- a/src/component/_generated/component.ts +++ b/src/component/_generated/component.ts @@ -31,6 +31,13 @@ export type ComponentApi = any, Name >; + claimStream: FunctionReference< + "mutation", + "internal", + { streamId: string }, + boolean, + Name + >; createStream: FunctionReference<"mutation", "internal", {}, any, Name>; deleteStream: FunctionReference< "mutation", diff --git a/src/component/lib.ts b/src/component/lib.ts index 89931d0..cf5de53 100644 --- a/src/component/lib.ts +++ b/src/component/lib.ts @@ -1,20 +1,97 @@ import { paginator } from "convex-helpers/server/pagination"; -import { v } from "convex/values"; +import { ConvexError, v } from "convex/values"; +import type { StreamReadResult } from "@convex-dev/stream"; import { internal } from "./_generated/api.js"; import { internalMutation, mutation, query } from "./_generated/server.js"; import schema, { streamStatusValidator } from "./schema.js"; +import { textStreams } from "./streams.js"; + +const MAX_PIECE_CODE_UNITS = 128 * 1024; +const MAX_APPEND_MANY_EVENTS = 8; +const FULL_BODY_PAGE_ITEMS = 128; +const MAX_FULL_BODY_PAGES = 32; +const MAX_FULL_BODY_EVENTS = FULL_BODY_PAGE_ITEMS * MAX_FULL_BODY_PAGES; +const MAX_FULL_BODY_BYTES = 4 * 1024 * 1024; +const textEncoder = new TextEncoder(); + +// Split only at code-point boundaries. The code-unit bound leaves ample room +// below Convex's document limit even when every character uses four UTF-8 bytes. +function textPieces(text: string): string[] { + if (text.length === 0) return [""]; + + const pieces: string[] = []; + let start = 0; + while (start < text.length) { + let end = Math.min(start + MAX_PIECE_CODE_UNITS, text.length); + if ( + end < text.length && + text.charCodeAt(end - 1) >= 0xd800 && + text.charCodeAt(end - 1) <= 0xdbff && + text.charCodeAt(end) >= 0xdc00 && + text.charCodeAt(end) <= 0xdfff + ) { + end -= 1; + } + pieces.push(text.slice(start, end)); + start = end; + } + return pieces; +} + +function timeoutError() { + return { code: "timeout", message: "Stream generation timed out." }; +} + +function failureError() { + return { code: "error", message: "Stream generation failed." }; +} + +function fullBodyLimitExceeded(): never { + throw new ConvexError({ + code: "fullBodyLimitExceeded", + message: "Full stream snapshot exceeds its finite compatibility limit.", + }); +} -// Create a new stream with zero chunks. export const createStream = mutation({ args: {}, handler: async (ctx) => { + const core = await textStreams.create(ctx); const streamId = await ctx.db.insert("streams", { status: "pending", + coreId: core.streamId, + engine: "v6", + lastActivityAt: Date.now(), }); return streamId; }, }); +// Claiming serializes producers without changing the externally visible +// lifecycle: the first pending claimant wins, and all later callers lose. +export const claimStream = mutation({ + args: { + streamId: v.id("streams"), + }, + returns: v.boolean(), + handler: async (ctx, args) => { + const stream = await ctx.db.get("streams", args.streamId); + if ( + stream === null || + stream.status !== "pending" || + stream.claimedAt !== undefined + ) { + return false; + } + const claimedAt = Date.now(); + await ctx.db.patch("streams", args.streamId, { + claimedAt, + ...(stream.coreId === undefined ? {} : { lastActivityAt: claimedAt }), + }); + return true; + }, +}); + // Add a chunk to a stream. // If final is true, set the stream to done. // Can only be done on streams which are pending or streaming. @@ -29,20 +106,40 @@ export const addChunk = mutation({ if (!stream) { throw new Error("Stream not found"); } - if (stream.status === "pending") { - await ctx.db.patch("streams", args.streamId, { - status: "streaming", - }); - } else if (stream.status !== "streaming") { + if (stream.status !== "pending" && stream.status !== "streaming") { throw new Error("Stream is not streaming; did it timeout?"); } - await ctx.db.insert("chunks", { - streamId: args.streamId, - text: args.text, - }); - if (args.final) { + + if (stream.coreId === undefined) { + await ctx.db.insert("chunks", { + streamId: args.streamId, + text: args.text, + }); + } else { + const pieces = textPieces(args.text); + for ( + let start = 0; + start < pieces.length; + start += MAX_APPEND_MANY_EVENTS + ) { + const events = pieces.slice(start, start + MAX_APPEND_MANY_EVENTS); + const complete = args.final && start + events.length === pieces.length; + await textStreams.appendMany(ctx, { + streamId: stream.coreId, + events, + ...(complete ? { complete: true as const } : {}), + }); + } + } + + if ( + stream.coreId !== undefined || + stream.status === "pending" || + args.final + ) { await ctx.db.patch("streams", args.streamId, { - status: "done", + status: args.final ? "done" : "streaming", + ...(stream.coreId === undefined ? {} : { lastActivityAt: Date.now() }), }); } }, @@ -73,6 +170,30 @@ export const setStreamStatus = mutation({ ); return; } + if (stream.coreId !== undefined) { + if (args.status === "pending" || args.status === "streaming") { + const core = await textStreams.get(ctx, { streamId: stream.coreId }); + if (core?.status === args.status && stream.status === args.status) { + return; + } + throw new Error( + "Native stream lifecycle transitions require an append or terminal status.", + ); + } + if (args.status === "done") { + await textStreams.complete(ctx, { streamId: stream.coreId }); + } else if (args.status === "error") { + await textStreams.fail(ctx, { + streamId: stream.coreId, + error: failureError(), + }); + } else if (args.status === "timeout") { + await textStreams.fail(ctx, { + streamId: stream.coreId, + error: timeoutError(), + }); + } + } await ctx.db.patch("streams", args.streamId, { status: args.status, }); @@ -92,7 +213,6 @@ export const getStreamStatus = query({ }); // Get the full text of a stream. -// Involves concatenating all the chunks. export const getStreamText = query({ args: { streamId: v.id("streams"), @@ -106,6 +226,40 @@ export const getStreamText = query({ if (!stream) { throw new Error("Stream not found"); } + if (stream.coreId !== undefined) { + const events: string[] = []; + let cursor: string | null = null; + let eventCount = 0; + let bytes = 0; + + // This legacy compatibility query is deliberately finite. The bounded + // reader introduced next is the supported path for unbounded histories. + for (let page = 0; page <= MAX_FULL_BODY_PAGES; page += 1) { + const result: StreamReadResult<"textStreams", string> = + await textStreams.read(ctx, { + streamId: stream.coreId, + cursor, + numItems: FULL_BODY_PAGE_ITEMS, + }); + for (const { event } of result.page) { + eventCount += 1; + bytes += textEncoder.encode(event).byteLength; + if ( + eventCount > MAX_FULL_BODY_EVENTS || + bytes > MAX_FULL_BODY_BYTES + ) { + fullBodyLimitExceeded(); + } + events.push(event); + } + if (result.caughtUp) { + return { text: events.join(""), status: stream.status }; + } + if (page === MAX_FULL_BODY_PAGES) fullBodyLimitExceeded(); + cursor = result.continueCursor; + } + fullBodyLimitExceeded(); + } let text = ""; if (stream.status !== "pending") { const chunks = await ctx.db @@ -125,8 +279,8 @@ const EXPIRATION_TIME = 20 * 60 * 1000; // 20 minutes in milliseconds const BATCH_SIZE = 100; const DELETE_BATCH_SIZE = 64; -// Delete a stream and all its chunks. -// The stream is deleted immediately; chunks are cleaned up asynchronously. +// Delete a stream and all its chunks. The facade read is an OCC fence against +// append: once deletion wins, a retried append observes that the facade is gone. export const deleteStream = mutation({ args: { streamId: v.id("streams"), @@ -137,11 +291,19 @@ export const deleteStream = mutation({ if (!stream) { throw new Error(`Stream ${args.streamId} not found`); } + if (stream.coreId === undefined) { + await ctx.scheduler.runAfter(0, internal.lib._deleteChunksPage, { + streamId: args.streamId, + cursor: null, + }); + } else { + await textStreams.delete( + ctx, + { streamId: stream.coreId }, + { run: internal.streamMaintenance.run }, + ); + } await ctx.db.delete("streams", args.streamId); - await ctx.scheduler.runAfter(0, internal.lib._deleteChunksPage, { - streamId: args.streamId, - cursor: null, - }); }, }); @@ -158,7 +320,9 @@ export const _deleteChunksPage = internalMutation({ .withIndex("byStream", (q) => q.eq("streamId", args.streamId)) .paginate({ cursor: args.cursor, numItems: DELETE_BATCH_SIZE }); - await Promise.all(result.page.map((chunk) => ctx.db.delete("chunks", chunk._id))); + await Promise.all( + result.page.map((chunk) => ctx.db.delete("chunks", chunk._id)), + ); if (!result.isDone) { await ctx.scheduler.runAfter(0, internal.lib._deleteChunksPage, { @@ -169,28 +333,55 @@ export const _deleteChunksPage = internalMutation({ }, }); -// If the last chunk of a stream was added more than 20 minutes ago, -// set the stream to timeout. The action feeding it has to be dead. +// Timeout is a PTS retention policy, not a V6 expiration: it leaves retained +// events readable and marks the native lifecycle failed with a timeout error. export const cleanupExpiredStreams = internalMutation({ args: {}, handler: async (ctx) => { const now = Date.now(); - const pendingStreams = await ctx.db + const legacyPendingStreams = await ctx.db .query("streams") .withIndex("byStatus", (q) => q.eq("status", "pending")) .take(BATCH_SIZE); - const streamingStreams = await ctx.db + const legacyStreamingStreams = await ctx.db .query("streams") .withIndex("byStatus", (q) => q.eq("status", "streaming")) .take(BATCH_SIZE); - for (const stream of [...pendingStreams, ...streamingStreams]) { - if (now - stream._creationTime > EXPIRATION_TIME) { - console.log("Cleaning up expired stream", stream._id); - await ctx.db.patch("streams", stream._id, { - status: "timeout", - }); + const staleNativeBatches = await Promise.all( + (["pending", "streaming"] as const).map((status) => + ctx.db + .query("streams") + .withIndex("byStatusAndEngineAndLastActivityAt", (q) => + q + .eq("status", status) + .eq("engine", "v6") + .lt("lastActivityAt", now - EXPIRATION_TIME), + ) + .take(BATCH_SIZE), + ), + ); + + for (const stream of [...legacyPendingStreams, ...legacyStreamingStreams]) { + if ( + stream.coreId !== undefined || + now - stream._creationTime <= EXPIRATION_TIME + ) { + continue; } + + console.log("Cleaning up expired stream", stream._id); + await ctx.db.patch("streams", stream._id, { status: "timeout" }); + } + + for (const stream of staleNativeBatches.flat()) { + if (stream.coreId === undefined) continue; + console.log("Cleaning up expired stream", stream._id); + await textStreams.fail(ctx, { + streamId: stream.coreId, + error: timeoutError(), + }); + await ctx.db.patch("streams", stream._id, { status: "timeout" }); } }, }); diff --git a/src/component/producer.test.ts b/src/component/producer.test.ts new file mode 100644 index 0000000..8c9658e --- /dev/null +++ b/src/component/producer.test.ts @@ -0,0 +1,354 @@ +/// + +import { convexTest } from "convex-test"; +import { describe, expect, it, vi } from "vitest"; + +import { PersistentTextStreaming } from "../client/index.js"; +import { api, internal } from "./_generated/api.js"; +import schema from "./schema.js"; +import { textStreams } from "./streams.js"; + +const modules = import.meta.glob("./**/*.ts"); + +describe("V6 producer facade", () => { + it("creates and claims a native stream without changing its public lifecycle", async () => { + vi.useFakeTimers(); + try { + vi.setSystemTime(1_000); + const t = convexTest(schema, modules); + const streamId = await t.mutation(api.lib.createStream, {}); + + await expect(t.mutation(api.lib.claimStream, { streamId })).resolves.toBe( + true, + ); + await expect(t.mutation(api.lib.claimStream, { streamId })).resolves.toBe( + false, + ); + + const state = await t.run(async (ctx) => { + const facade = await ctx.db.get("streams", streamId); + const core = facade?.coreId + ? await ctx.db.get("textStreams", facade.coreId) + : null; + return { facade, core }; + }); + expect(state.facade).toMatchObject({ + status: "pending", + claimedAt: 1_000, + lastActivityAt: 1_000, + }); + expect(state.core).toMatchObject({ status: "pending", attempt: 0 }); + } finally { + vi.useRealTimers(); + } + }); + + it("writes Unicode-safe V6 batches while legacy rows keep chunk storage", async () => { + const t = convexTest(schema, modules); + const nativeId = await t.mutation(api.lib.createStream, {}); + const text = "a".repeat(128 * 1024 - 1) + "πŸ˜€tail"; + await t.mutation(api.lib.addChunk, { + streamId: nativeId, + text, + final: false, + }); + await t.mutation(api.lib.addChunk, { + streamId: nativeId, + text: "!", + final: true, + }); + + const legacyId = await t.run((ctx) => + ctx.db.insert("streams", { status: "pending" }), + ); + await t.mutation(api.lib.addChunk, { + streamId: legacyId, + text: "legacy", + final: true, + }); + + const state = await t.run(async (ctx) => { + const native = await ctx.db.get("streams", nativeId); + const core = native?.coreId + ? await ctx.db.get("textStreams", native.coreId) + : null; + const events = await ctx.db.query("textStreamsEvents").collect(); + const chunks = await ctx.db.query("chunks").collect(); + return { native, core, events, chunks }; + }); + expect(state.native).toMatchObject({ status: "done" }); + expect(state.core).toMatchObject({ status: "done" }); + expect(state.events.map((event) => event.event).join("")).toBe(`${text}!`); + expect(state.events).toHaveLength(3); + expect(state.chunks).toMatchObject([ + { streamId: legacyId, text: "legacy" }, + ]); + }); + + it("reconstructs a native body in V6 order across finite compatibility pages", async () => { + const t = convexTest(schema, modules); + const streamId = await t.mutation(api.lib.createStream, {}); + const events = Array.from({ length: 129 }, (_, index) => `${index},`); + for (const [index, text] of events.entries()) { + await t.mutation(api.lib.addChunk, { + streamId, + text, + final: index === events.length - 1, + }); + } + + await expect(t.query(api.lib.getStreamText, { streamId })).resolves.toEqual( + { + text: events.join(""), + status: "done", + }, + ); + + const streaming = new PersistentTextStreaming({ + lib: { getStreamText: api.lib.getStreamText }, + } as never); + await expect( + streaming.getStreamBody( + { + runQuery: ( + _reference: unknown, + args: { streamId: typeof streamId }, + ) => t.query(api.lib.getStreamText, args), + } as never, + streamId as never, + ), + ).resolves.toEqual({ text: events.join(""), status: "done" }); + }); + + it("accepts the exact native body event limit and rejects one more", async () => { + const t = convexTest(schema, modules); + const exact = await t.run(async (ctx) => { + const core = await textStreams.create(ctx); + const streamId = await ctx.db.insert("streams", { + status: "done", + coreId: core.streamId, + engine: "v6", + lastActivityAt: 0, + }); + for (let start = 0; start < 4096; start += 1024) { + await textStreams.appendMany(ctx, { + streamId: core.streamId, + events: Array.from({ length: 1024 }, () => "x"), + }); + } + return streamId; + }); + await expect( + t.query(api.lib.getStreamText, { streamId: exact }), + ).resolves.toEqual({ + text: "x".repeat(4096), + status: "done", + }); + + const over = await t.run(async (ctx) => { + const core = await textStreams.create(ctx); + const streamId = await ctx.db.insert("streams", { + status: "done", + coreId: core.streamId, + engine: "v6", + lastActivityAt: 0, + }); + for (let start = 0; start < 4097; start += 1024) { + await textStreams.appendMany(ctx, { + streamId: core.streamId, + events: Array.from( + { length: Math.min(1024, 4097 - start) }, + () => "x", + ), + }); + } + return streamId; + }); + await expect( + t.query(api.lib.getStreamText, { streamId: over }), + ).rejects.toThrow("fullBodyLimitExceeded"); + }); + + it("rejects a native body beyond the four MiB compatibility cap", async () => { + const t = convexTest(schema, modules); + const streamId = await t.run(async (ctx) => { + const core = await textStreams.create(ctx); + const facade = await ctx.db.insert("streams", { + status: "done", + coreId: core.streamId, + engine: "v6", + lastActivityAt: 0, + }); + for (let index = 0; index < 33; index += 1) { + await textStreams.append(ctx, { + streamId: core.streamId, + event: "x".repeat(128 * 1024), + }); + } + return facade; + }); + await expect(t.query(api.lib.getStreamText, { streamId })).rejects.toThrow( + "fullBodyLimitExceeded", + ); + }); + + it("keeps native and facade lifecycle aligned for nonterminal status calls", async () => { + const t = convexTest(schema, modules); + const pendingId = await t.mutation(api.lib.createStream, {}); + await expect( + t.mutation(api.lib.setStreamStatus, { + streamId: pendingId, + status: "streaming", + }), + ).rejects.toThrow("Native stream lifecycle transitions"); + await expect( + t.query(api.lib.getStreamStatus, { streamId: pendingId }), + ).resolves.toBe("pending"); + + const streamingId = await t.mutation(api.lib.createStream, {}); + await t.mutation(api.lib.addChunk, { + streamId: streamingId, + text: "started", + final: false, + }); + await expect( + t.mutation(api.lib.setStreamStatus, { + streamId: streamingId, + status: "pending", + }), + ).rejects.toThrow("Native stream lifecycle transitions"); + await expect( + t.query(api.lib.getStreamStatus, { streamId: streamingId }), + ).resolves.toBe("streaming"); + }); + + it("mirrors terminal lifecycle and timeout without assigning V6 expiration", async () => { + vi.useFakeTimers(); + try { + vi.setSystemTime(0); + const t = convexTest(schema, modules); + const failedId = await t.mutation(api.lib.createStream, {}); + await t.mutation(api.lib.setStreamStatus, { + streamId: failedId, + status: "error", + }); + const idleId = await t.mutation(api.lib.createStream, {}); + const activeId = await t.mutation(api.lib.createStream, {}); + const legacyId = await t.run((ctx) => + ctx.db.insert("streams", { status: "streaming" }), + ); + + vi.setSystemTime(19 * 60 * 1000); + await t.mutation(api.lib.addChunk, { + streamId: activeId, + text: "active", + final: false, + }); + vi.setSystemTime(20 * 60 * 1000 + 1); + await t.mutation(internal.lib.cleanupExpiredStreams, {}); + + const states = await t.run(async (ctx) => + Promise.all( + [failedId, idleId, activeId, legacyId].map(async (streamId) => { + const facade = await ctx.db.get("streams", streamId); + const core = facade?.coreId + ? await ctx.db.get("textStreams", facade.coreId) + : null; + return { facade, core }; + }), + ), + ); + expect(states).toMatchObject([ + { + facade: { status: "error" }, + core: { status: "failed", error: { code: "error" } }, + }, + { + facade: { status: "timeout" }, + core: { status: "failed", error: { code: "timeout" } }, + }, + { + facade: { status: "streaming" }, + core: { status: "streaming" }, + }, + { facade: { status: "timeout" }, core: null }, + ]); + expect(states[2]?.core).not.toHaveProperty("expiresAt"); + } finally { + vi.useRealTimers(); + } + }); + + it("does not starve stale native streams behind active native rows", async () => { + vi.useFakeTimers(); + try { + vi.setSystemTime(0); + const t = convexTest(schema, modules); + const active = await Promise.all( + Array.from({ length: 101 }, () => t.mutation(api.lib.createStream, {})), + ); + const stale = await t.mutation(api.lib.createStream, {}); + vi.setSystemTime(19 * 60 * 1000); + await Promise.all( + active.map((streamId) => + t.mutation(api.lib.addChunk, { + streamId, + text: "active", + final: false, + }), + ), + ); + vi.setSystemTime(20 * 60 * 1000 + 1); + await t.mutation(internal.lib.cleanupExpiredStreams, {}); + + await expect( + t.query(api.lib.getStreamStatus, { streamId: stale }), + ).resolves.toBe("timeout"); + await expect( + t.query(api.lib.getStreamStatus, { streamId: active[0]! }), + ).resolves.toBe("streaming"); + } finally { + vi.useRealTimers(); + } + }); + + it("deletes native and legacy rows without leaving an appendable facade", async () => { + vi.useFakeTimers(); + try { + const t = convexTest(schema, modules); + const nativeId = await t.mutation(api.lib.createStream, {}); + await t.mutation(api.lib.addChunk, { + streamId: nativeId, + text: "native", + final: true, + }); + const legacyId = await t.run(async (ctx) => { + const streamId = await ctx.db.insert("streams", { status: "done" }); + await ctx.db.insert("chunks", { streamId, text: "legacy" }); + return streamId; + }); + + await t.mutation(api.lib.deleteStream, { streamId: nativeId }); + await t.mutation(api.lib.deleteStream, { streamId: legacyId }); + await expect( + t.mutation(api.lib.addChunk, { + streamId: nativeId, + text: "late", + final: false, + }), + ).rejects.toThrow("Stream not found"); + await t.finishAllScheduledFunctions(vi.runAllTimers); + + await expect( + t.run(async (ctx) => ({ + facades: await ctx.db.query("streams").collect(), + chunks: await ctx.db.query("chunks").collect(), + cores: await ctx.db.query("textStreams").collect(), + events: await ctx.db.query("textStreamsEvents").collect(), + })), + ).resolves.toEqual({ facades: [], chunks: [], cores: [], events: [] }); + } finally { + vi.useRealTimers(); + } + }); +}); diff --git a/src/component/schema.ts b/src/component/schema.ts index 9d88a10..64c6924 100644 --- a/src/component/schema.ts +++ b/src/component/schema.ts @@ -14,7 +14,18 @@ export type StreamStatus = Infer; export default defineSchema({ streams: defineTable({ status: streamStatusValidator, - }).index("byStatus", ["status"]), + // These fields are optional so existing facade rows remain legacy streams. + coreId: v.optional(v.id("textStreams")), + engine: v.optional(v.literal("v6")), + claimedAt: v.optional(v.number()), + lastActivityAt: v.optional(v.number()), + }) + .index("byStatus", ["status"]) + .index("byStatusAndEngineAndLastActivityAt", [ + "status", + "engine", + "lastActivityAt", + ]), chunks: defineTable({ streamId: v.id("streams"), text: v.string(),