rivetkit 2.0.2 → 2.0.4-rc.1

package/src/actor/router.ts

@@ -0,0 +1,380 @@
+import { Hono } from "hono";
+import {
+	type ActionOpts,
+	type ActionOutput,
+	type ConnsMessageOpts,
+	handleAction,
+	handleQueueSend,
+	handleRawRequest,
+} from "@/actor/router-endpoints";
+
+import {
+	PATH_CONNECT,
+	PATH_INSPECTOR_CONNECT,
+	PATH_WEBSOCKET_PREFIX,
+} from "@/common/actor-router-consts";
+import {
+	handleRouteError,
+	handleRouteNotFound,
+	loggerMiddleware,
+} from "@/common/router";
+import { noopNext } from "@/common/utils";
+import { inspectorLogger } from "@/inspector/log";
+import type { RegistryConfig } from "@/registry/config";
+import { type GetUpgradeWebSocket, VERSION } from "@/utils";
+import { timingSafeEqual } from "@/utils/crypto";
+import { isDev } from "@/utils/env-vars";
+import { CONN_DRIVER_SYMBOL } from "./conn/mod";
+import type { ActorDriver } from "./driver";
+import { loggerWithoutContext } from "./log";
+import {
+	parseWebSocketProtocols,
+	routeWebSocket,
+} from "./router-websocket-endpoints";
+
+export type { ActionOpts, ActionOutput, ConnsMessageOpts };
+
+interface ActorRouterBindings {
+	actorId: string;
+}
+
+export type ActorRouter = Hono<{ Bindings: ActorRouterBindings }>;
+
+export interface MetadataResponse {
+	runtime: string;
+	version: string;
+}
+
+/**
+ * Creates a router that runs on the partitioned instance.
+ *
+ * You only need to pass `getUpgradeWebSocket` if this router is exposed
+ * directly publicly. Usually WebSockets are routed manually in the
+ * ManagerDriver instead of via Hono. The only platform that uses this
+ * currently is Cloudflare Workers.
+ */
+export function createActorRouter(
+	config: RegistryConfig,
+	actorDriver: ActorDriver,
+	getUpgradeWebSocket: GetUpgradeWebSocket | undefined,
+	isTest: boolean,
+): ActorRouter {
+	const router = new Hono<{ Bindings: ActorRouterBindings }>({
+		strict: false,
+	});
+
+	router.use("*", loggerMiddleware(loggerWithoutContext()));
+
+	// Track all HTTP requests to prevent actor from sleeping during active requests
+	router.use("*", async (c, next) => {
+		const actor = await actorDriver.loadActor(c.env.actorId);
+		actor.beginHonoHttpRequest();
+		try {
+			await next();
+		} finally {
+			actor.endHonoHttpRequest();
+		}
+	});
+
+	router.get("/", (c) => {
+		return c.text(
+			"This is an RivetKit actor.\n\nLearn more at https://rivetkit.org",
+		);
+	});
+
+	router.get("/health", (c) => {
+		return c.text("ok");
+	});
+
+	router.get("/metadata", async (c) => {
+		return c.json({
+			runtime: "rivetkit",
+			version: VERSION,
+		} satisfies MetadataResponse);
+	});
+
+	if (isTest) {
+		// Test endpoint to force disconnect a connection non-cleanly
+		router.post("/.test/force-disconnect", async (c) => {
+			const connId = c.req.query("conn");
+
+			if (!connId) {
+				return c.text("Missing conn query parameter", 400);
+			}
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			const conn = actor.connectionManager.getConnForId(connId);
+
+			if (!conn) {
+				return c.text(`Connection not found: ${connId}`, 404);
+			}
+
+			// Force close the connection without clean shutdown
+			if (conn[CONN_DRIVER_SYMBOL]?.terminate) {
+				conn[CONN_DRIVER_SYMBOL].terminate(actor, conn);
+			}
+
+			return c.json({ success: true });
+		});
+	}
+
+	// Route all WebSocket paths using the same handler
+	//
+	// All WebSockets use a separate underlying router in routeWebSocket since
+	// WebSockets also need to be routed from ManagerDriver.proxyWebSocket and
+	// ManagerDriver.openWebSocket.
+	if (getUpgradeWebSocket) {
+		router.on(
+			"GET",
+			[PATH_CONNECT, `${PATH_WEBSOCKET_PREFIX}*`, PATH_INSPECTOR_CONNECT],
+			async (c) => {
+				const upgradeWebSocket = getUpgradeWebSocket();
+				if (upgradeWebSocket) {
+					return upgradeWebSocket(async (c) => {
+						const protocols = c.req.header(
+							"sec-websocket-protocol",
+						);
+						const { encoding, connParams } =
+							parseWebSocketProtocols(protocols);
+
+						return await routeWebSocket(
+							c.req.raw,
+							c.req.path,
+							c.req.header(),
+							config,
+							actorDriver,
+							c.env.actorId,
+							encoding,
+							connParams,
+							undefined,
+							undefined,
+							false,
+							false,
+						);
+					})(c, noopNext());
+				} else {
+					return c.text(
+						"WebSockets are not enabled for this driver.",
+						400,
+					);
+				}
+			},
+		);
+	}
+
+	// Inspector HTTP endpoints for agent-based debugging
+	if (config.inspector.enabled) {
+		// Auth middleware for inspector routes
+		const inspectorAuth = async (c: any): Promise<Response | undefined> => {
+			if (isDev() && !config.inspector.token()) {
+				inspectorLogger().warn({
+					msg: "RIVET_INSPECTOR_TOKEN is not set, skipping inspector auth in development mode",
+				});
+				return undefined;
+			}
+
+			const userToken = c.req
+				.header("Authorization")
+				?.replace("Bearer ", "");
+			if (!userToken) {
+				return c.text("Unauthorized", 401);
+			}
+
+			const inspectorToken = config.inspector.token();
+			if (!inspectorToken) {
+				return c.text("Unauthorized", 401);
+			}
+
+			if (!timingSafeEqual(userToken, inspectorToken)) {
+				return c.text("Unauthorized", 401);
+			}
+
+			return undefined;
+		};
+
+		router.get("/inspector/state", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			const isStateEnabled = actor.inspector.isStateEnabled();
+			const state = isStateEnabled
+				? actor.inspector.getStateJson()
+				: undefined;
+			return c.json({ state, isStateEnabled });
+		});
+
+		router.patch("/inspector/state", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			const body = await c.req.json<{ state: unknown }>();
+			await actor.inspector.setStateJson(body.state);
+			return c.json({ ok: true });
+		});
+
+		router.get("/inspector/connections", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			const connections = actor.inspector.getConnectionsJson();
+			return c.json({ connections });
+		});
+
+		router.get("/inspector/rpcs", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			const rpcs = actor.inspector.getRpcs();
+			return c.json({ rpcs });
+		});
+
+		router.post("/inspector/action/:name", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			const name = c.req.param("name");
+			const body = await c.req.json<{ args: unknown[] }>();
+			const output = await actor.inspector.executeActionJson(
+				name,
+				body.args ?? [],
+			);
+			return c.json({ output });
+		});
+
+		router.get("/inspector/queue", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			const limit = parseInt(c.req.query("limit") ?? "50", 10);
+			const status = await actor.inspector.getQueueStatusJson(limit);
+			return c.json(status);
+		});
+
+		router.get("/inspector/traces", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			const startMs = parseInt(c.req.query("startMs") ?? "0", 10);
+			const endMs = parseInt(
+				c.req.query("endMs") ?? String(Date.now()),
+				10,
+			);
+			const limit = parseInt(c.req.query("limit") ?? "1000", 10);
+
+			await actor.traces.flush();
+			const result = await actor.inspector.getTracesJson({
+				startMs,
+				endMs,
+				limit,
+			});
+			return c.json(result);
+		});
+
+		router.get("/inspector/workflow-history", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			const result = actor.inspector.getWorkflowHistoryJson();
+			return c.json(result);
+		});
+
+		router.get("/inspector/summary", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+
+			const isStateEnabled = actor.inspector.isStateEnabled();
+			const isDatabaseEnabled = actor.inspector.isDatabaseEnabled();
+			const isWorkflowEnabled = actor.inspector.isWorkflowEnabled();
+
+			const state = isStateEnabled
+				? actor.inspector.getStateJson()
+				: undefined;
+			const connections = actor.inspector.getConnectionsJson();
+			const rpcs = actor.inspector.getRpcs();
+			const queueSize = actor.inspector.getQueueSize();
+			const workflowHistory = actor.inspector.getWorkflowHistory();
+
+			// Convert BigInt values in workflow history to numbers for JSON serialization.
+			const bigIntReplacer = (_key: string, value: unknown) =>
+				typeof value === "bigint" ? Number(value) : value;
+			const safeWorkflowHistory = workflowHistory
+				? JSON.parse(JSON.stringify(workflowHistory, bigIntReplacer))
+				: null;
+
+			return c.json({
+				state,
+				connections,
+				rpcs,
+				queueSize,
+				isStateEnabled,
+				isDatabaseEnabled,
+				isWorkflowEnabled,
+				workflowHistory: safeWorkflowHistory,
+			});
+		});
+	}
+
+	router.post("/action/:action", async (c) => {
+		const actionName = c.req.param("action");
+
+		return handleAction(c, config, actorDriver, actionName, c.env.actorId);
+	});
+
+	router.post("/queue", async (c) => {
+		return handleQueueSend(c, config, actorDriver, c.env.actorId);
+	});
+
+	router.post("/queue/:name", async (c) => {
+		return handleQueueSend(
+			c,
+			config,
+			actorDriver,
+			c.env.actorId,
+			c.req.param("name"),
+		);
+	});
+
+	router.all("/request/*", async (c) => {
+		// TODO: This is not a clean way of doing this since `/http/` might exist mid-path
+		// Strip the /http prefix from the URL to get the original path
+		const url = new URL(c.req.url);
+		const originalPath = url.pathname.replace(/^\/request/, "") || "/";
+
+		// Create a new request with the corrected URL
+		const correctedUrl = new URL(originalPath + url.search, url.origin);
+		const correctedRequest = new Request(correctedUrl, {
+			method: c.req.method,
+			headers: c.req.raw.headers,
+			body: c.req.raw.body,
+			duplex: "half",
+		} as RequestInit);
+
+		loggerWithoutContext().debug({
+			msg: "rewriting http url",
+			from: c.req.url,
+			to: correctedRequest.url,
+		});
+
+		return await handleRawRequest(
+			c,
+			correctedRequest,
+			actorDriver,
+			c.env.actorId,
+		);
+	});
+
+	router.notFound(handleRouteNotFound);
+	router.onError(handleRouteError);
+
+	return router;
+}

package/src/actor/schedule.ts

@@ -0,0 +1,17 @@
+import type { AnyActorInstance } from "./instance/mod";
+
+export class Schedule {
+	#actor: AnyActorInstance;
+
+	constructor(actor: AnyActorInstance) {
+		this.#actor = actor;
+	}
+
+	async after(duration: number, fn: string, ...args: unknown[]) {
+		await this.#actor.scheduleEvent(Date.now() + duration, fn, args);
+	}
+
+	async at(timestamp: number, fn: string, ...args: unknown[]) {
+		await this.#actor.scheduleEvent(timestamp, fn, args);
+	}
+}

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] };
+}