rivetkit 2.0.42 → 2.1.0-rc.1

package/src/actor/schema.ts

@@ -0,0 +1,291 @@
+import type { StandardSchemaV1 } from "@standard-schema/spec";
+import { Unsupported } from "./errors";
+
+export type SchemaHookResult = boolean | Promise<boolean>;
+
+type SchemaHook<TContext = any> = (ctx: TContext) => SchemaHookResult;
+
+export interface EventTypeToken<T, TContext = any> {
+	readonly _eventType?: T;
+	readonly canSubscribe?: SchemaHook<TContext>;
+	readonly schema?: PrimitiveSchema;
+}
+
+export interface QueueTypeToken<
+	TMessage,
+	TComplete = never,
+	TContext = any,
+> {
+	readonly _queueMessage?: TMessage;
+	readonly _queueComplete?: TComplete;
+	readonly canPublish?: SchemaHook<TContext>;
+	readonly message?: PrimitiveSchema;
+	readonly complete?: PrimitiveSchema;
+}
+
+/** @deprecated Use `event<T>()`. */
+export type Type<T> = EventTypeToken<T, any>;
+
+interface EventOptions<TContext = any> {
+	canSubscribe?: SchemaHook<TContext>;
+	schema?: PrimitiveSchema;
+}
+
+interface QueueOptions<TContext = any> {
+	canPublish?: SchemaHook<TContext>;
+	message?: PrimitiveSchema;
+	complete?: PrimitiveSchema;
+}
+
+export function event<T, TContext = any>(
+	options?: EventOptions<TContext>,
+): EventTypeToken<T, TContext> {
+	return (options ?? {}) as EventTypeToken<T, TContext>;
+}
+
+export function queue<TMessage, TComplete = never, TContext = any>(
+	options?: QueueOptions<TContext>,
+): QueueTypeToken<
+	TMessage,
+	TComplete,
+	TContext
+> {
+	return (options ?? {}) as QueueTypeToken<TMessage, TComplete, TContext>;
+}
+
+export type PrimitiveSchema = StandardSchemaV1 | EventTypeToken<unknown, any>;
+
+export interface EventSchemaDefinition<TContext = any> {
+	schema: PrimitiveSchema;
+	canSubscribe?: SchemaHook<TContext>;
+}
+
+export interface QueueSchemaDefinition<TContext = any> {
+	message: PrimitiveSchema;
+	complete?: PrimitiveSchema;
+	canPublish?: SchemaHook<TContext>;
+}
+
+export type EventSchema<TContext = any> =
+	| PrimitiveSchema
+	| EventSchemaDefinition<TContext>;
+export type QueueSchema =
+	| PrimitiveSchema
+	| QueueSchemaDefinition<any>
+	| QueueTypeToken<unknown, unknown, any>;
+export type EventSchemaConfig<TContext = any> = Record<
+	string,
+	EventSchema<TContext>
+>;
+export type QueueSchemaConfig<TContext = any> = Record<string, QueueSchema>;
+export type AnySchemaConfig = EventSchemaConfig | QueueSchemaConfig;
+
+/** @deprecated Use `EventSchema` or `QueueSchema`. */
+export type Schema = QueueSchema;
+/** @deprecated Use `EventSchemaConfig` or `QueueSchemaConfig`. */
+export type SchemaConfig = QueueSchemaConfig;
+
+export type InferSchema<T> =
+	T extends QueueSchemaDefinition<any>
+			? InferSchema<T["message"]>
+			: T extends QueueTypeToken<infer M, unknown, any>
+			? M
+			: T extends EventSchemaDefinition<any>
+			? InferSchema<T["schema"]>
+			: T extends StandardSchemaV1<any, infer O>
+			? O
+			: T extends EventTypeToken<infer R, any>
+				? R
+				: never;
+
+export type InferSchemaMap<T extends Record<string, unknown>> = {
+	[K in keyof T]: InferSchema<T[K]>;
+};
+
+export type InferQueueComplete<T> =
+	T extends QueueTypeToken<unknown, infer C, any>
+			? [C] extends [never]
+				? never
+				: C
+			: T extends QueueSchemaDefinition<any>
+			? T["complete"] extends PrimitiveSchema
+				? InferSchema<T["complete"]>
+				: never
+			: never;
+
+export type InferQueueCompleteMap<T extends QueueSchemaConfig> = {
+	[K in keyof T]: InferQueueComplete<T[K]>;
+};
+
+export type InferEventArgs<T> = T extends readonly unknown[]
+	? number extends T["length"]
+		? [T]
+		: T
+	: [T];
+
+export type ValidationResult<T> =
+	| { success: true; data: T }
+	| { success: false; issues: unknown[] };
+
+export function isStandardSchema(value: unknown): value is StandardSchemaV1 {
+	return typeof value === "object" && value !== null && "~standard" in value;
+}
+
+export function isQueueSchemaDefinition(
+	value: unknown,
+): value is QueueSchemaDefinition<any> {
+	if (isEventSchemaDefinition(value)) {
+		return false;
+	}
+	return (
+		typeof value === "object" &&
+		value !== null &&
+		"message" in value &&
+		(value as { message?: unknown }).message !== undefined
+	);
+}
+
+export function isEventSchemaDefinition(
+	value: unknown,
+): value is EventSchemaDefinition<any> {
+	return (
+		typeof value === "object" &&
+		value !== null &&
+		"schema" in value &&
+		(value as { schema?: unknown }).schema !== undefined
+	);
+}
+
+export function hasSchemaConfigKey<T extends AnySchemaConfig>(
+	schemas: T | undefined,
+	key: string,
+): boolean {
+	if (!schemas) {
+		return false;
+	}
+	return Object.prototype.hasOwnProperty.call(schemas, key);
+}
+
+export function getEventCanSubscribe<TContext = any>(
+	schemas: EventSchemaConfig<TContext> | undefined,
+	key: string,
+): SchemaHook<TContext> | undefined {
+	const schema = schemas?.[key];
+	if (!schema || isStandardSchema(schema)) {
+		return undefined;
+	}
+
+	const maybeCanSubscribe = (schema as { canSubscribe?: unknown })
+		.canSubscribe;
+	return typeof maybeCanSubscribe === "function"
+		? (maybeCanSubscribe as SchemaHook<TContext>)
+		: undefined;
+}
+
+export function getQueueCanPublish<TContext = any>(
+	schemas: QueueSchemaConfig | undefined,
+	key: string,
+): SchemaHook<TContext> | undefined {
+	const schema = schemas?.[key];
+	if (!schema || isStandardSchema(schema)) {
+		return undefined;
+	}
+
+	const maybeCanPublish = (schema as { canPublish?: unknown }).canPublish;
+	return typeof maybeCanPublish === "function"
+		? (maybeCanPublish as SchemaHook<TContext>)
+		: undefined;
+}
+
+function getValidationSchema(
+	schema: QueueSchema | EventSchema | undefined,
+): QueueSchema | EventSchema | undefined {
+	if (!schema) {
+		return undefined;
+	}
+	if (isEventSchemaDefinition(schema)) {
+		return schema.schema;
+	}
+	if (isQueueSchemaDefinition(schema)) {
+		return schema.message;
+	}
+	if (
+		typeof schema === "object" &&
+		schema !== null &&
+		"schema" in schema &&
+		(schema as { schema?: unknown }).schema !== undefined
+	) {
+		return (schema as { schema: QueueSchema | EventSchema }).schema;
+	}
+	if (
+		typeof schema === "object" &&
+		schema !== null &&
+		"message" in schema &&
+		(schema as { message?: unknown }).message !== undefined
+	) {
+		return (schema as { message: QueueSchema | EventSchema }).message;
+	}
+	return schema;
+}
+
+function isPromiseLike(value: unknown): value is PromiseLike<unknown> {
+	return (
+		typeof value === "object" &&
+		value !== null &&
+		"then" in value &&
+		typeof (value as { then?: unknown }).then === "function"
+	);
+}
+
+export async function validateSchema<T extends AnySchemaConfig>(
+	schemas: T | undefined,
+	key: keyof T & string,
+	data: unknown,
+): Promise<ValidationResult<InferSchemaMap<T>[typeof key]>> {
+	const schema = getValidationSchema(schemas?.[key]);
+
+	if (!schema) {
+		return { success: true, data: data as InferSchemaMap<T>[typeof key] };
+	}
+
+	if (isStandardSchema(schema)) {
+		const result = await schema["~standard"].validate(data);
+		if (result.issues) {
+			return { success: false, issues: [...result.issues] };
+		}
+		return {
+			success: true,
+			data: result.value as InferSchemaMap<T>[typeof key],
+		};
+	}
+
+	return { success: true, data: data as InferSchemaMap<T>[typeof key] };
+}
+
+export function validateSchemaSync<T extends AnySchemaConfig>(
+	schemas: T | undefined,
+	key: keyof T & string,
+	data: unknown,
+): ValidationResult<InferSchemaMap<T>[typeof key]> {
+	const schema = getValidationSchema(schemas?.[key]);
+
+	if (!schema) {
+		return { success: true, data: data as InferSchemaMap<T>[typeof key] };
+	}
+
+	if (isStandardSchema(schema)) {
+		const result = schema["~standard"].validate(data);
+		if (isPromiseLike(result)) {
+			throw new Unsupported("async schema validation");
+		}
+		if (result.issues) {
+			return { success: false, issues: [...result.issues] };
+		}
+		return {
+			success: true,
+			data: result.value as InferSchemaMap<T>[typeof key],
+		};
+	}
+
+	return { success: true, data: data as InferSchemaMap<T>[typeof key] };
+}

package/src/actor/utils.ts

@@ -84,6 +84,46 @@ export class Lock<T> {
 	}
 }
 
+export interface JoinedAbortSignal {
+	signal?: AbortSignal;
+	cleanup: () => void;
+}
+
+export function joinAbortSignals(
+	...signals: Array<AbortSignal | undefined>
+): JoinedAbortSignal {
+	const activeSignals = signals.filter(
+		(signal): signal is AbortSignal => signal !== undefined,
+	);
+	if (activeSignals.length === 0) {
+		return { signal: undefined, cleanup: () => {} };
+	}
+	if (activeSignals.length === 1) {
+		return { signal: activeSignals[0], cleanup: () => {} };
+	}
+
+	const controller = new AbortController();
+	if (activeSignals.some((signal) => signal.aborted)) {
+		controller.abort();
+		return { signal: controller.signal, cleanup: () => {} };
+	}
+
+	const cleanup = () => {
+		for (const signal of activeSignals) {
+			signal.removeEventListener("abort", onAbort);
+		}
+	};
+	const onAbort = () => {
+		controller.abort();
+		cleanup();
+	};
+	for (const signal of activeSignals) {
+		signal.addEventListener("abort", onAbort, { once: true });
+	}
+
+	return { signal: controller.signal, cleanup };
+}
+
 export function generateSecureToken(length = 32) {
 	const array = new Uint8Array(length);
 	crypto.getRandomValues(array);

package/src/client/actor-common.ts

@@ -21,7 +21,7 @@ export type ActorActionFunction<
  */
 export type ActorDefinitionActions<AD extends AnyActorDefinition> =
 	// biome-ignore lint/suspicious/noExplicitAny: safe to use any here
-	AD extends ActorDefinition<any, any, any, any, any, any, infer R>
+	AD extends ActorDefinition<any, any, any, any, any, any, any, any, infer R>
 		? {
 				[K in keyof R]: R[K] extends (
 					...args: infer Args

package/src/client/actor-conn.ts

@@ -30,6 +30,13 @@ import { checkForSchedulingError, queryActor } from "./actor-query";
 import { ACTOR_CONNS_SYMBOL, type ClientRaw } from "./client";
 import * as errors from "./errors";
 import { logger } from "./log";
+import {
+	createQueueSender,
+	type QueueSendNoWaitOptions,
+	type QueueSendOptions,
+	type QueueSendResult,
+	type QueueSendWaitOptions,
+} from "./queue";
 import {
 	type WebSocketMessage as ConnMessage,
 	messageLength,
@@ -131,10 +138,12 @@ export class ActorConnRaw {
 
 	#errorHandlers = new Set<ActorErrorCallback>();
 	#openHandlers = new Set<ConnectionStateCallback>();
+	#openScheduled = false;
 	#closeHandlers = new Set<ConnectionStateCallback>();
 	#statusChangeHandlers = new Set<StatusChangeCallback>();
 
 	#actionIdCounter = 0;
+	#queueSender: ReturnType<typeof createQueueSender>;
 
 	/**
 	 * Interval that keeps the NodeJS process alive if this is the only thing running.
@@ -175,10 +184,43 @@ export class ActorConnRaw {
 		this.#params = params;
 		this.#encoding = encoding;
 		this.#actorQuery = actorQuery;
+		this.#queueSender = createQueueSender({
+			encoding: this.#encoding,
+			params: this.#params,
+			customFetch: async (request: Request) => {
+				if (!this.#actorId) {
+					const { actorId } = await queryActor(
+						undefined,
+						this.#actorQuery,
+						this.#driver,
+					);
+					this.#actorId = actorId;
+				}
+				return this.#driver.sendRequest(this.#actorId, request);
+			},
+		});
 
 		this.#keepNodeAliveInterval = setInterval(() => 60_000);
 	}
 
+	send(
+		name: string,
+		body: unknown,
+		options: QueueSendWaitOptions,
+	): Promise<QueueSendResult>;
+	send(
+		name: string,
+		body: unknown,
+		options?: QueueSendNoWaitOptions,
+	): Promise<void>;
+	send(
+		name: string,
+		body: unknown,
+		options?: QueueSendOptions,
+	): Promise<QueueSendResult | void> {
+		return this.#queueSender.send(name, body, options as any);
+	}
+
 	/**
 	 * Call a raw action connection. See {@link ActorConn} for type-safe action calls.
 	 *
@@ -206,7 +248,7 @@ export class ActorConnRaw {
 		const { promise, resolve, reject } = promiseWithResolvers<{
 			id: bigint;
 			output: unknown;
-		}>();
+		}>((reason) => logger().warn({ msg: "unhandled action promise rejection", reason }));
 		this.#actionsInFlight.set(actionId, {
 			name: opts.name,
 			resolve,
@@ -243,7 +285,6 @@ export class ActorConnRaw {
 
 	/**
 	 * Do not call this directly.
-enc
 	 * Establishes a connection to the server using the specified endpoint & encoding & driver.
 	 *
 	 * @protected
@@ -338,7 +379,7 @@ enc
 			// Create promise for open
 			if (this.#onOpenPromise)
 				throw new Error("#onOpenPromise already defined");
-			this.#onOpenPromise = promiseWithResolvers();
+			this.#onOpenPromise = promiseWithResolvers((reason) => logger().warn({ msg: "unhandled open promise rejection", reason }));
 
 			await this.#connectWebSocket();
 
@@ -424,39 +465,58 @@ enc
 			return;
 		}
 
-		logger().debug({
-			msg: "socket open",
-			messageQueueLength: this.#messageQueue.length,
-			connId: this.#connId,
-		});
+		if (this.#connStatus === "connected" || this.#openScheduled) {
+			return;
+		}
+		this.#openScheduled = true;
+
+		queueMicrotask(() => {
+			this.#openScheduled = false;
+			if (this.#disposed) {
+				logger().debug({
+					msg: "handleOnOpen scheduled after dispose, closing websocket",
+				});
+				if (this.#websocket) {
+					this.#websocket.close(1000, "Disposed");
+					this.#websocket = undefined;
+				}
+				return;
+			}
 
-		// Update connection state (this also notifies handlers)
-		this.#setConnStatus("connected");
+			logger().debug({
+				msg: "socket open",
+				messageQueueLength: this.#messageQueue.length,
+				connId: this.#connId,
+			});
 
-		// Resolve open promise
-		if (this.#onOpenPromise) {
-			this.#onOpenPromise.resolve(undefined);
-		} else {
-			logger().warn({ msg: "#onOpenPromise is undefined" });
-		}
+			// Update connection state (this also notifies handlers)
+			this.#setConnStatus("connected");
 
-		// Resubscribe to all active events
-		for (const eventName of this.#eventSubscriptions.keys()) {
-			this.#sendSubscription(eventName, true);
-		}
+			// Resolve open promise
+			if (this.#onOpenPromise) {
+				this.#onOpenPromise.resolve(undefined);
+			} else {
+				logger().warn({ msg: "#onOpenPromise is undefined" });
+			}
 
-		// Flush queue
-		//
-		// If the message fails to send, the message will be re-queued
-		const queue = this.#messageQueue;
-		this.#messageQueue = [];
-		logger().debug({
-			msg: "flushing message queue",
-			queueLength: queue.length,
+			// Resubscribe to all active events
+			for (const eventName of this.#eventSubscriptions.keys()) {
+				this.#sendSubscription(eventName, true);
+			}
+
+			// Flush queue
+			//
+			// If the message fails to send, the message will be re-queued
+			const queue = this.#messageQueue;
+			this.#messageQueue = [];
+			logger().debug({
+				msg: "flushing message queue",
+				queueLength: queue.length,
+			});
+			for (const msg of queue) {
+				this.#sendMessage(msg);
+			}
 		});
-		for (const msg of queue) {
-			this.#sendMessage(msg);
-		}
 	}
 
 	/** Called by the onmessage event from drivers. */
@@ -939,7 +999,14 @@ enc
 				messageType: (message.body as any).tag,
 				actionName: (message.body as any).val?.name,
 			});
-			if (readyState === 1) {
+			if (this.#connStatus !== "connected") {
+				logger().debug({
+					msg: "websocket init pending, queueing message",
+					connStatus: this.#connStatus,
+					messageType: (message.body as any).tag,
+				});
+				queueMessage = true;
+			} else if (readyState === 1) {
 				try {
 					const messageSerialized = serializeWithEncoding(
 						this.#encoding,
@@ -1149,7 +1216,7 @@ enc
 				ws.readyState !== 2 /* CLOSING */ &&
 				ws.readyState !== 3 /* CLOSED */
 			) {
-				const { promise, resolve } = promiseWithResolvers();
+				const { promise, resolve } = promiseWithResolvers((reason) => logger().warn({ msg: "unhandled websocket close promise rejection", reason }));
 				ws.addEventListener("close", () => resolve(undefined));
 				ws.close(1000, "Disposed");
 				await promise;

package/src/client/actor-handle.ts

@@ -29,6 +29,13 @@ import { checkForSchedulingError, queryActor } from "./actor-query";
 import { type ClientRaw, CREATE_ACTOR_CONN_PROXY } from "./client";
 import { ActorError, isSchedulingError } from "./errors";
 import { logger } from "./log";
+import {
+	createQueueSender,
+	type QueueSendNoWaitOptions,
+	type QueueSendOptions,
+	type QueueSendResult,
+	type QueueSendWaitOptions,
+} from "./queue";
 import { rawHttpFetch, rawWebSocket } from "./raw-utils";
 import { sendHttpRequest } from "./utils";
 
@@ -44,6 +51,7 @@ export class ActorHandleRaw {
 	#encoding: Encoding;
 	#actorQuery: ActorQuery;
 	#params: unknown;
+	#queueSender: ReturnType<typeof createQueueSender>;
 
 	/**
 	 * Do not call this directly.
@@ -64,6 +72,36 @@ export class ActorHandleRaw {
 		this.#encoding = encoding;
 		this.#actorQuery = actorQuery;
 		this.#params = params;
+		this.#queueSender = createQueueSender({
+			encoding: this.#encoding,
+			params: this.#params,
+			customFetch: async (request: Request) => {
+				const { actorId } = await queryActor(
+					undefined,
+					this.#actorQuery,
+					this.#driver,
+				);
+				return this.#driver.sendRequest(actorId, request);
+			},
+		});
+	}
+
+	send(
+		name: string,
+		body: unknown,
+		options: QueueSendWaitOptions,
+	): Promise<QueueSendResult>;
+	send(
+		name: string,
+		body: unknown,
+		options?: QueueSendNoWaitOptions,
+	): Promise<void>;
+	send(
+		name: string,
+		body: unknown,
+		options?: QueueSendOptions,
+	): Promise<QueueSendResult | void> {
+		return this.#queueSender.send(name, body, options as any);
 	}
 
 	/**
@@ -199,16 +237,10 @@ export class ActorHandleRaw {
 	}
 
 	/**
-	 * Makes a raw HTTP request to the actor.
-	 *
-	 * @param input - The URL, path, or Request object
-	 * @param init - Standard fetch RequestInit options
-	 * @returns Promise<Response> - The raw HTTP response
+	 * Fetches a resource from this actor via the /request endpoint. This is a
+	 * convenience wrapper around the raw HTTP API.
 	 */
-	async fetch(
-		input: string | URL | Request,
-		init?: RequestInit,
-	): Promise<Response> {
+	fetch(input: string | URL | Request, init?: RequestInit) {
 		return rawHttpFetch(
 			this.#driver,
 			this.#actorQuery,
@@ -219,16 +251,9 @@ export class ActorHandleRaw {
 	}
 
 	/**
-	 * Creates a raw WebSocket connection to the actor.
-	 *
-	 * @param path - The path for the WebSocket connection (e.g., "stream")
-	 * @param protocols - Optional WebSocket subprotocols
-	 * @returns WebSocket - A raw WebSocket connection
+	 * Opens a raw WebSocket connection to this actor.
 	 */
-	async websocket(
-		path?: string,
-		protocols?: string | string[],
-	): Promise<WebSocket> {
+	webSocket(path?: string, protocols?: string | string[]) {
 		return rawWebSocket(
 			this.#driver,
 			this.#actorQuery,
@@ -241,21 +266,24 @@ export class ActorHandleRaw {
 	/**
 	 * Resolves the actor to get its unique actor ID.
 	 */
-	async resolve({ signal }: { signal?: AbortSignal } = {}): Promise<string> {
-		if (
-			"getForKey" in this.#actorQuery ||
-			"getOrCreateForKey" in this.#actorQuery
-		) {
-			// TODO:
-			let name: string;
-			if ("getForKey" in this.#actorQuery) {
-				name = this.#actorQuery.getForKey.name;
-			} else if ("getOrCreateForKey" in this.#actorQuery) {
-				name = this.#actorQuery.getOrCreateForKey.name;
-			} else {
-				assertUnreachable(this.#actorQuery);
-			}
+	async resolve(): Promise<string> {
+		if ("getForKey" in this.#actorQuery) {
+			const name = this.#actorQuery.getForKey.name;
+
+			// Query the actor to get the id
+			const { actorId } = await queryActor(
+				undefined,
+				this.#actorQuery,
+				this.#driver,
+			);
+
+			this.#actorQuery = { getForId: { actorId, name } };
+
+			return actorId;
+		} else if ("getOrCreateForKey" in this.#actorQuery) {
+			const name = this.#actorQuery.getOrCreateForKey.name;
 
+			// Query the actor to get or create and get the id
 			const { actorId } = await queryActor(
 				undefined,
 				this.#actorQuery,
@@ -266,7 +294,7 @@ export class ActorHandleRaw {
 
 			return actorId;
 		} else if ("getForId" in this.#actorQuery) {
-			// SKip since it's already resolved
+			// Skip since it's already resolved
 			return this.#actorQuery.getForId.actorId;
 		} else if ("create" in this.#actorQuery) {
 			// Cannot create a handle with this query