@rivet-dev/vercel-world 2.3.6

package/src/actors/hook-token.ts

@@ -0,0 +1,239 @@
+import { actor } from "rivetkit";
+import { hookTokenDb } from "./db.js";
+import { one } from "./shared.js";
+
+export const PENDING_RESERVATION_TTL_MS = 60_000;
+
+type HookTokenContext = {
+	key: readonly unknown[];
+	db: {
+		execute(
+			sql: string,
+			...args: unknown[]
+		): Promise<Record<string, unknown>[]>;
+		transaction<T>(
+			callback: (db: HookTokenContext["db"]) => Promise<T> | T,
+		): Promise<T>;
+	};
+	client<T>(): T;
+};
+
+type HookTokenRow = {
+	generation: number;
+	runId?: string;
+	hookId?: string;
+	status: "empty" | "pending" | "confirmed";
+	createdAt?: number;
+	updatedAt: number;
+	expiresAt?: number;
+};
+
+function assertTokenKey(c: { key: readonly unknown[] }, token: string) {
+	if (String(c.key[0] ?? "") !== token) {
+		throw new Error("hookToken actor key must match token");
+	}
+}
+
+async function readState(c: { db: HookTokenContext["db"] }): Promise<HookTokenRow> {
+	const row = await one(
+		c,
+		`SELECT generation, run_id, hook_id, status, created_at, updated_at, expires_at
+		 FROM hook_token_state WHERE singleton = 1`,
+	);
+	if (!row) throw new Error("hookToken state is missing");
+	return {
+		generation: Number(row.generation),
+		runId: row.run_id == null ? undefined : String(row.run_id),
+		hookId: row.hook_id == null ? undefined : String(row.hook_id),
+		status: String(row.status) as HookTokenRow["status"],
+		createdAt: row.created_at == null ? undefined : Number(row.created_at),
+		updatedAt: Number(row.updated_at),
+		expiresAt: row.expires_at == null ? undefined : Number(row.expires_at),
+	};
+}
+
+async function verifyPendingOwner(
+	c: HookTokenContext,
+	token: string,
+	state: HookTokenRow,
+): Promise<HookTokenRow | null> {
+	if (
+		state.status !== "pending" ||
+		state.runId == null ||
+		state.hookId == null
+	) {
+		return state.status === "confirmed" ? state : null;
+	}
+
+	let ownsHook: boolean;
+	try {
+		ownsHook = await c
+			.client<any>()
+			.workflowRun.getOrCreate([state.runId])
+			.ownsHook(state.hookId, token);
+	} catch {
+		// Canonical ownership could not be checked. A pending reservation must
+		// never become externally visible or be stolen on an inconclusive result.
+		return null;
+	}
+
+	const now = Date.now();
+	if (ownsHook) {
+		await c.db.execute(
+			`UPDATE hook_token_state
+			 SET status = 'confirmed', updated_at = ?, expires_at = NULL
+			 WHERE singleton = 1 AND generation = ? AND run_id = ? AND hook_id = ?
+			   AND status = 'pending'`,
+			now,
+			state.generation,
+			state.runId,
+			state.hookId,
+		);
+		const current = await readState(c);
+		return current.generation === state.generation &&
+			current.status === "confirmed" &&
+			current.runId === state.runId &&
+			current.hookId === state.hookId
+			? current
+			: null;
+	}
+
+	if ((state.expiresAt ?? Number.POSITIVE_INFINITY) > now) return null;
+	await c.db.execute(
+		`UPDATE hook_token_state
+		 SET generation = generation + 1, run_id = NULL, hook_id = NULL,
+		     status = 'empty', created_at = NULL, updated_at = ?, expires_at = NULL
+		 WHERE singleton = 1 AND generation = ? AND run_id = ? AND hook_id = ?
+		   AND status = 'pending'`,
+		now,
+		state.generation,
+		state.runId,
+		state.hookId,
+	);
+	return null;
+}
+
+export const hookToken = actor({
+	db: hookTokenDb,
+	actions: {
+		reserve: async (c, token: string, runId: string, hookId: string) => {
+			const ctx = c as unknown as HookTokenContext;
+			assertTokenKey(ctx, token);
+			let state = await readState(ctx);
+			if (
+				state.status === "pending" &&
+				(state.runId !== runId || state.hookId !== hookId) &&
+				(state.expiresAt ?? Number.POSITIVE_INFINITY) <= Date.now()
+			) {
+				await verifyPendingOwner(ctx, token, state);
+				state = await readState(ctx);
+			}
+			if (
+				state.status !== "empty" &&
+				state.runId === runId &&
+				state.hookId === hookId
+			) {
+				return { ok: true, generation: state.generation };
+			}
+			if (state.status !== "empty") {
+				return {
+					ok: false,
+					generation: state.generation,
+					runId: state.runId,
+					hookId: state.hookId,
+				};
+			}
+
+			const now = Date.now();
+			await ctx.db.execute(
+				`UPDATE hook_token_state
+				 SET generation = generation + 1, run_id = ?, hook_id = ?,
+				     status = 'pending', created_at = ?, updated_at = ?, expires_at = ?
+				 WHERE singleton = 1 AND status = 'empty'`,
+				runId,
+				hookId,
+				now,
+				now,
+				now + PENDING_RESERVATION_TTL_MS,
+			);
+			const reserved = await readState(ctx);
+			return {
+				ok:
+					reserved.status === "pending" &&
+					reserved.runId === runId &&
+					reserved.hookId === hookId,
+				generation: reserved.generation,
+			};
+		},
+		confirm: async (
+			c,
+			token: string,
+			generation: number,
+			runId: string,
+			hookId: string,
+		) => {
+			const ctx = c as unknown as HookTokenContext;
+			assertTokenKey(ctx, token);
+			await ctx.db.execute(
+				`UPDATE hook_token_state
+				 SET status = 'confirmed', updated_at = ?, expires_at = NULL
+				 WHERE singleton = 1 AND generation = ? AND run_id = ? AND hook_id = ?
+				   AND status IN ('pending', 'confirmed')`,
+				Date.now(),
+				generation,
+				runId,
+				hookId,
+			);
+			const state = await readState(ctx);
+			return {
+				ok:
+					state.generation === generation &&
+					state.status === "confirmed" &&
+					state.runId === runId &&
+					state.hookId === hookId,
+			};
+		},
+		get: async (c, token: string) => {
+			const ctx = c as unknown as HookTokenContext;
+			assertTokenKey(ctx, token);
+			let state = await readState(ctx);
+			if (state.status === "pending") {
+				await verifyPendingOwner(ctx, token, state);
+				state = await readState(ctx);
+			}
+			if (state.status !== "confirmed") return null;
+			return {
+				runId: state.runId as string,
+				hookId: state.hookId as string,
+				generation: state.generation,
+			};
+		},
+		release: async (
+			c,
+			token: string,
+			generation: number,
+			runId: string,
+			hookId?: string,
+		) => {
+			const ctx = c as unknown as HookTokenContext;
+			assertTokenKey(ctx, token);
+			await ctx.db.execute(
+				`UPDATE hook_token_state
+				 SET generation = generation + 1, run_id = NULL, hook_id = NULL,
+				     status = 'empty', created_at = NULL, updated_at = ?, expires_at = NULL
+				 WHERE singleton = 1 AND generation = ? AND run_id = ?
+				   AND (? IS NULL OR hook_id = ?) AND status IN ('pending', 'confirmed')`,
+				Date.now(),
+				generation,
+				runId,
+				hookId ?? null,
+				hookId ?? null,
+			);
+			const state = await readState(ctx);
+			return {
+				ok:
+					state.status === "empty" && state.generation >= generation + 1,
+			};
+		},
+	},
+});

package/src/actors/shared.ts

@@ -0,0 +1,153 @@
+import { WorkflowRunNotFoundError } from "@workflow/errors";
+import {
+	EventSchema,
+	HookSchema,
+	StepSchema,
+	WorkflowRunSchema,
+	type Hook,
+	type Wait,
+} from "@workflow/world";
+import { compact, decodeValue } from "../codec.js";
+
+export type Db = {
+	execute(sql: string, ...args: unknown[]): Promise<Record<string, unknown>[]>;
+	transaction<T>(callback: (tx: Db) => Promise<T> | T): Promise<T>;
+};
+
+export type Ctx = { db: Db };
+
+export async function withActorTransaction<T>(
+	c: Ctx,
+	fn: (tx: Ctx) => Promise<T>,
+): Promise<T> {
+	return c.db.transaction((db) => fn({ db }));
+}
+
+export function toDate(value: unknown): Date | undefined {
+	return value == null ? undefined : new Date(Number(value));
+}
+
+export function toMs(value: unknown): number | null {
+	if (value == null) return null;
+	if (value instanceof Date) return value.getTime();
+	return new Date(value as string | number).getTime();
+}
+
+export function rowToRun(row: Record<string, unknown>) {
+	return WorkflowRunSchema.parse(
+		compact({
+			runId: row.run_id,
+			status: row.status,
+			deploymentId: row.deployment_id,
+			workflowName: row.workflow_name,
+			specVersion:
+				row.spec_version == null ? undefined : Number(row.spec_version),
+			executionContext: decodeValue(row.execution_context),
+			attributes: decodeValue(row.attributes) ?? {},
+			input: decodeValue(row.input),
+			output: decodeValue(row.output),
+			error: decodeValue(row.error),
+			expiredAt: toDate(row.expired_at),
+			startedAt: toDate(row.started_at),
+			completedAt: toDate(row.completed_at),
+			createdAt: toDate(row.created_at),
+			updatedAt: toDate(row.updated_at),
+		}),
+	);
+}
+
+export function rowToEvent(row: Record<string, unknown>) {
+	return EventSchema.parse(
+		compact({
+			runId: row.run_id,
+			eventId: row.event_id,
+			eventType: row.event_type,
+			correlationId: row.correlation_id,
+			eventData: decodeValue(row.event_data),
+			specVersion:
+				row.spec_version == null ? undefined : Number(row.spec_version),
+			createdAt: toDate(row.created_at),
+		}),
+	);
+}
+
+export function rowToStep(row: Record<string, unknown>) {
+	return StepSchema.parse(
+		compact({
+			runId: row.run_id,
+			stepId: row.step_id,
+			stepName: row.step_name,
+			status: row.status,
+			input: decodeValue(row.input),
+			output: decodeValue(row.output),
+			error: decodeValue(row.error),
+			attempt: Number(row.attempt),
+			startedAt: toDate(row.started_at),
+			completedAt: toDate(row.completed_at),
+			createdAt: toDate(row.created_at),
+			updatedAt: toDate(row.updated_at),
+			retryAfter: toDate(row.retry_after),
+			specVersion:
+				row.spec_version == null ? undefined : Number(row.spec_version),
+		}),
+	);
+}
+
+export function rowToHook(row: Record<string, unknown>) {
+	const parsed = HookSchema.parse(
+		compact({
+			runId: row.run_id,
+			hookId: row.hook_id,
+			token: row.token,
+			ownerId: row.owner_id,
+			projectId: row.project_id,
+			environment: row.environment,
+			metadata: decodeValue(row.metadata),
+			createdAt: toDate(row.created_at),
+			specVersion:
+				row.spec_version == null ? undefined : Number(row.spec_version),
+			isWebhook: row.is_webhook == null ? undefined : Boolean(row.is_webhook),
+		}),
+	);
+	parsed.isWebhook ??= true;
+	return parsed;
+}
+
+export function rowToWait(row: Record<string, unknown>): Wait {
+	return compact({
+		waitId: row.wait_id,
+		runId: row.run_id,
+		status: row.status,
+		resumeAt: toDate(row.resume_at),
+		completedAt: toDate(row.completed_at),
+		createdAt: toDate(row.created_at),
+		updatedAt: toDate(row.updated_at),
+		specVersion: row.spec_version == null ? undefined : Number(row.spec_version),
+	}) as Wait;
+}
+
+export function filterData<T extends { input?: unknown; output?: unknown }>(
+	value: T,
+	resolveData?: string,
+) {
+	if (resolveData !== "none") return value;
+	return { ...value, input: undefined, output: undefined };
+}
+
+export function filterHookData<T extends { metadata?: unknown }>(
+	value: T,
+	resolveData?: string,
+) {
+	if (resolveData !== "none") return value;
+	return { ...value, metadata: undefined };
+}
+
+export async function one(c: Ctx, sql: string, ...args: unknown[]) {
+	return (await c.db.execute(sql, ...args))[0];
+}
+
+export async function requireRun(c: Ctx, runId: string) {
+	const row = await one(c, "SELECT * FROM workflow_runs WHERE run_id = ?", runId);
+	if (!row) throw new WorkflowRunNotFoundError(runId);
+	return rowToRun(row);
+}

package/src/actors/streams.ts

@@ -0,0 +1,215 @@
+import type { StreamChunksResponse, StreamInfoResponse } from "@workflow/world";
+import { event, type Type } from "rivetkit";
+import { decodeValue, encodeValue } from "../codec.js";
+import { withActorTransaction, type Ctx } from "./shared.js";
+
+function envInteger(name: string, fallback: number, min: number) {
+	const value = Number(process.env[name]);
+	return Number.isFinite(value) ? Math.max(min, Math.floor(value)) : fallback;
+}
+
+export const STREAM_TTL_MS = envInteger(
+	"RIVET_WORLD_RIVET_STREAM_TTL_MS",
+	86_400_000,
+	0,
+);
+
+export const streamAppended: Type<{
+	streamId: string;
+	runId: string;
+	chunkId: string;
+	eof: boolean;
+}> = event();
+
+type StreamContext = Ctx & {
+	key: readonly string[];
+	broadcast(name: "streamAppended", value: unknown): void;
+};
+
+function assertRunKey(c: StreamContext, runId: string) {
+	if (c.key[0] !== runId) throw new Error("workflowRun actor key must match runId");
+}
+
+function decodeBase64Chunk(value: string) {
+	return new Uint8Array(Buffer.from(value, "base64"));
+}
+
+let broadcastCounter = 0;
+function shouldDropBroadcast(eof: boolean) {
+	if (eof) return false;
+	const rate = Number(process.env.RIVET_WORLD_RIVET_INJECT_BROADCAST_DROP_RATE);
+	if (!Number.isFinite(rate) || rate <= 0) return false;
+	const r = Math.min(rate, 1);
+	const n = broadcastCounter++;
+	return Math.floor((n + 1) * r) > Math.floor(n * r);
+}
+
+async function injectReadDelay() {
+	const ms = Number(process.env.RIVET_WORLD_RIVET_INJECT_STREAM_READ_DELAY_MS);
+	if (Number.isFinite(ms) && ms > 0) {
+		await new Promise((resolve) => setTimeout(resolve, ms));
+	}
+}
+
+async function cleanupExpired(c: Ctx, streamId: string) {
+	if (STREAM_TTL_MS === 0) return 0;
+	const expired = await c.db.execute(
+		`SELECT 1 FROM stream_chunks WHERE stream_id = ? AND eof = 1
+		 AND created_at <= ? LIMIT 1`,
+		streamId,
+		Date.now() - STREAM_TTL_MS,
+	);
+	if (expired.length === 0) return 0;
+	return (
+		await c.db.execute(
+			"DELETE FROM stream_chunks WHERE stream_id = ? RETURNING sequence",
+			streamId,
+		)
+	).length;
+}
+
+export async function writeStream(
+	c: StreamContext,
+	streamId: string,
+	runId: string,
+	chunkBase64: string,
+	eof: boolean,
+) {
+	assertRunKey(c, runId);
+	await cleanupExpired(c, streamId);
+	const result = await withActorTransaction(c, async (tx) => {
+		const closed = await tx.db.execute(
+			"SELECT sequence FROM stream_chunks WHERE stream_id = ? AND eof = 1 LIMIT 1",
+			streamId,
+		);
+		if (closed[0]) {
+			if (eof) return { sequence: Number(closed[0].sequence), inserted: false };
+			throw new Error(`stream "${streamId}" is already closed`);
+		}
+		const next = await tx.db.execute(
+			"SELECT COALESCE(MAX(sequence), -1) + 1 AS sequence FROM stream_chunks WHERE stream_id = ?",
+			streamId,
+		);
+		const value = Number(next[0]?.sequence ?? 0);
+		await tx.db.execute(
+			`INSERT INTO stream_chunks (stream_id, sequence, chunk_data, eof, created_at)
+			 VALUES (?, ?, ?, ?, ?)`,
+			streamId,
+			value,
+			encodeValue(decodeBase64Chunk(chunkBase64)),
+			eof ? 1 : 0,
+			Date.now(),
+		);
+		await tx.db.execute(
+			"INSERT OR IGNORE INTO run_streams (stream_id) VALUES (?)",
+			streamId,
+		);
+		return { sequence: value, inserted: true };
+	});
+	const chunkId = String(result.sequence);
+	if (result.inserted && !shouldDropBroadcast(eof)) {
+		c.broadcast("streamAppended", { streamId, runId, chunkId, eof });
+	}
+	return { chunkId };
+}
+
+export async function getStreamChunks(
+	c: StreamContext,
+	streamId: string,
+	runId: string,
+	startIndex = 0,
+	limit = 100,
+): Promise<StreamChunksResponse> {
+	assertRunKey(c, runId);
+	await cleanupExpired(c, streamId);
+	const rows = await c.db.execute(
+		`SELECT sequence, chunk_data FROM stream_chunks
+		 WHERE stream_id = ? AND eof = 0 AND sequence >= ?
+		 ORDER BY sequence ASC LIMIT ?`,
+		streamId,
+		startIndex,
+		limit + 1,
+	);
+	const hasMore = rows.length > limit;
+	const values = rows.slice(0, limit);
+	const data = values.map((row) => ({
+		index: Number(row.sequence),
+		data: decodeValue<Uint8Array>(row.chunk_data) ?? new Uint8Array(),
+	}));
+	const closed = await c.db.execute(
+		"SELECT 1 FROM stream_chunks WHERE stream_id = ? AND eof = 1 LIMIT 1",
+		streamId,
+	);
+	await injectReadDelay();
+	const nextSequence = values.length
+		? Number(values.at(-1)?.sequence) + 1
+		: startIndex;
+	return {
+		data,
+		cursor: hasMore
+			? Buffer.from(JSON.stringify({ i: nextSequence })).toString("base64")
+			: null,
+		hasMore,
+		done: !hasMore && closed.length > 0,
+	};
+}
+
+export async function getStreamInfo(
+	c: StreamContext,
+	streamId: string,
+	runId: string,
+): Promise<StreamInfoResponse> {
+	assertRunKey(c, runId);
+	await cleanupExpired(c, streamId);
+	const row = (
+		await c.db.execute(
+			`SELECT SUM(CASE WHEN eof = 0 THEN 1 ELSE 0 END) AS chunk_count,
+			 MAX(eof) AS done FROM stream_chunks WHERE stream_id = ?`,
+			streamId,
+		)
+	)[0];
+	const count = Number(row?.chunk_count ?? 0);
+	return { tailIndex: count - 1, done: Number(row?.done ?? 0) === 1 };
+}
+
+export async function listStreams(c: StreamContext, runId: string) {
+	assertRunKey(c, runId);
+	return (
+		await c.db.execute(`SELECT r.stream_id FROM run_streams r
+			WHERE EXISTS (SELECT 1 FROM stream_chunks c WHERE c.stream_id = r.stream_id)
+			ORDER BY r.stream_id`)
+	).map((row) => String(row.stream_id));
+}
+
+export async function cleanupExpiredStreams(c: Ctx) {
+	if (STREAM_TTL_MS === 0) return;
+	const cutoff = Date.now() - STREAM_TTL_MS;
+	const streams = await c.db.execute(
+		`SELECT DISTINCT stream_id FROM stream_chunks
+		 WHERE eof = 1 AND created_at <= ?
+		 ORDER BY stream_id LIMIT 32`,
+		cutoff,
+	);
+	for (const row of streams) {
+		const streamId = String(row.stream_id);
+		await c.db.execute("DELETE FROM stream_chunks WHERE stream_id = ?", streamId);
+		await c.db.execute("DELETE FROM run_streams WHERE stream_id = ?", streamId);
+	}
+}
+
+export async function expireClosedStreamForTesting(
+	c: StreamContext,
+	streamId: string,
+	runId: string,
+) {
+	assertRunKey(c, runId);
+	if (process.env.RIVET_WORLD_RIVET_TESTING !== "1") {
+		throw new Error("expireClosedForTesting requires test mode");
+	}
+	await c.db.execute(
+		"UPDATE stream_chunks SET created_at = ? WHERE stream_id = ? AND eof = 1",
+		Date.now() - STREAM_TTL_MS - 1_000,
+		streamId,
+	);
+	return { deletedChunks: await cleanupExpired(c, streamId) };
+}