rivetkit 2.0.22-rc.1 → 2.0.22

package/src/actor/router-endpoints.ts

@@ -7,12 +7,12 @@ import { ActionContext } from "@/actor/action";
 import type { AnyConn } from "@/actor/conn";
 import {
 	generateConnId,
-	generateConnSocketId,
+	generateConnRequestId,
 	generateConnToken,
 } from "@/actor/conn";
 import { ConnDriverKind } from "@/actor/conn-drivers";
 import * as errors from "@/actor/errors";
-import type { AnyActorInstance } from "@/actor/instance";
+import { type AnyActorInstance, PERSIST_SYMBOL } from "@/actor/instance";
 import type { InputData } from "@/actor/protocol/serde";
 import { type Encoding, EncodingSchema } from "@/actor/protocol/serde";
 import {
@@ -38,7 +38,11 @@ import {
 	deserializeWithEncoding,
 	serializeWithEncoding,
 } from "@/serde";
-import { bufferToArrayBuffer, promiseWithResolvers } from "@/utils";
+import {
+	arrayBuffersEqual,
+	bufferToArrayBuffer,
+	promiseWithResolvers,
+} from "@/utils";
 import type { ActorDriver } from "./driver";
 import { loggerWithoutContext } from "./log";
 import { parseMessage } from "./protocol/old";
@@ -111,8 +115,9 @@ export async function handleWebSocketConnect(
 	actorId: string,
 	encoding: Encoding,
 	parameters: unknown,
-	connId?: string,
-	connToken?: string,
+	requestId: string,
+	connId: string | undefined,
+	connToken: string | undefined,
 ): Promise<UpgradeWebSocketArgs> {
 	const exposeInternalError = req
 		? getRequestExposeInternalError(req)
@@ -157,7 +162,9 @@ export async function handleWebSocketConnect(
 
 	// Promise used to wait for the websocket close in `disconnect`
 	const closePromise = promiseWithResolvers<void>();
-	const socketId = generateConnSocketId();
+
+	// Track connection outside of scope for cleanup
+	let createdConn: AnyConn | undefined;
 
 	return {
 		onOpen: (_evt: any, ws: WSContext) => {
@@ -179,7 +186,7 @@ export async function handleWebSocketConnect(
 
 					conn = await actor.createConn(
 						{
-							socketId,
+							requestId: requestId,
 							driverState: {
 								[ConnDriverKind.WEBSOCKET]: {
 									encoding,
@@ -194,6 +201,9 @@ export async function handleWebSocketConnect(
 						connToken,
 					);
 
+					// Store connection so we can clean on close
+					createdConn = conn;
+
 					// Unblock other handlers
 					handlersResolve({ conn, actor, connId: conn.id });
 				} catch (error) {
@@ -294,20 +304,13 @@ export async function handleWebSocketConnect(
 			// https://github.com/cloudflare/workerd/issues/2569
 			ws.close(1000, "hack_force_close");
 
-			// Handle cleanup asynchronously
-			handlersPromise
-				.then(({ conn, actor }) => {
+			// Wait for actor.createConn to finish before removing the connection
+			handlersPromise.finally(() => {
+				if (createdConn) {
 					const wasClean = event.wasClean || event.code === 1000;
-					actor.__connDisconnected(conn, wasClean, socketId);
-				})
-				.catch((error) => {
-					deconstructError(
-						error,
-						actor.rLog,
-						{ wsEvent: "close" },
-						exposeInternalError,
-					);
-				});
+					actor.__connDisconnected(createdConn, wasClean, requestId);
+				}
+			});
 		},
 		onError: (_error: unknown) => {
 			try {
@@ -338,7 +341,7 @@ export async function handleSseConnect(
 
 	const encoding = getRequestEncoding(c.req);
 	const parameters = getRequestConnParams(c.req);
-	const socketId = generateConnSocketId();
+	const requestId = generateConnRequestId();
 
 	// Check for reconnection parameters
 	const connId = c.req.header(HEADER_CONN_ID);
@@ -361,7 +364,7 @@ export async function handleSseConnect(
 
 			conn = await actor.createConn(
 				{
-					socketId,
+					requestId: requestId,
 					driverState: {
 						[ConnDriverKind.SSE]: {
 							encoding,
@@ -392,7 +395,7 @@ export async function handleSseConnect(
 
 					// Cleanup
 					if (conn) {
-						actor.__connDisconnected(conn, false, socketId);
+						actor.__connDisconnected(conn, false, requestId);
 					}
 
 					abortResolver.resolve(undefined);
@@ -431,7 +434,7 @@ export async function handleSseConnect(
 
 			// Cleanup on error
 			if (conn && actor !== undefined) {
-				actor.__connDisconnected(conn, false, socketId);
+				actor.__connDisconnected(conn, false, requestId);
 			}
 
 			// Close the stream on error
@@ -461,7 +464,7 @@ export async function handleAction(
 		HTTP_ACTION_REQUEST_VERSIONED,
 	);
 	const actionArgs = cbor.decode(new Uint8Array(request.args));
-	const socketId = generateConnSocketId();
+	const requestId = generateConnRequestId();
 
 	// Invoke the action
 	let actor: AnyActorInstance | undefined;
@@ -475,7 +478,7 @@ export async function handleAction(
 		// Create conn
 		conn = await actor.createConn(
 			{
-				socketId,
+				requestId: requestId,
 				driverState: { [ConnDriverKind.HTTP]: {} },
 			},
 			parameters,
@@ -488,7 +491,7 @@ export async function handleAction(
 	} finally {
 		if (conn) {
 			// HTTP connections don't have persistent sockets, so no socket ID needed
-			actor?.__connDisconnected(conn, true, socketId);
+			actor?.__connDisconnected(conn, true, requestId);
 		}
 	}
 
@@ -595,38 +598,38 @@ export async function handleRawWebSocketHandler(
 
 	// Return WebSocket event handlers
 	return {
-		onOpen: (_evt: any, ws: any) => {
+		onOpen: (evt: any, ws: any) => {
+			// Extract rivetRequestId provided by engine runner
+			const rivetRequestId = evt?.rivetRequestId;
+			const isHibernatable =
+				actor[PERSIST_SYMBOL].hibernatableWebSocket.findIndex((ws) =>
+					arrayBuffersEqual(ws.requestId, rivetRequestId),
+				) !== -1;
+
 			// Wrap the Hono WebSocket in our adapter
-			const adapter = new HonoWebSocketAdapter(ws);
+			const adapter = new HonoWebSocketAdapter(
+				ws,
+				rivetRequestId,
+				isHibernatable,
+			);
 
 			// Store adapter reference on the WebSocket for event handlers
 			(ws as any).__adapter = adapter;
 
-			// Extract the path after prefix and preserve query parameters
-			// Use URL API for cleaner parsing
-			const url = new URL(path, "http://actor");
-			const pathname =
-				url.pathname.replace(/^\/raw\/websocket\/?/, "") || "/";
-			const normalizedPath =
-				(pathname.startsWith("/") ? pathname : "/" + pathname) +
-				url.search;
-
+			const newPath = truncateRawWebSocketPathPrefix(path);
 			let newRequest: Request;
 			if (req) {
-				newRequest = new Request(`http://actor${normalizedPath}`, req);
+				newRequest = new Request(`http://actor${newPath}`, req);
 			} else {
-				newRequest = new Request(`http://actor${normalizedPath}`, {
+				newRequest = new Request(`http://actor${newPath}`, {
 					method: "GET",
 				});
 			}
 
 			actor.rLog.debug({
 				msg: "rewriting websocket url",
-				from: path,
-				to: newRequest.url,
-				pathname: url.pathname,
-				search: url.search,
-				normalizedPath,
+				fromPath: path,
+				toUrl: newRequest.url,
 			});
 
 			// Call the actor's onWebSocket handler with the adapted WebSocket
@@ -711,3 +714,22 @@ export function getRequestConnParams(req: HonoRequest): unknown {
 		);
 	}
 }
+
+/**
+ * Truncase the PATH_RAW_WEBSOCKET_PREFIX path prefix in order to pass a clean
+ * path to the onWebSocket handler.
+ *
+ * Example:
+ * - `/raw/websocket/foo` -> `/foo`
+ * - `/raw/websocket` -> `/`
+ */
+export function truncateRawWebSocketPathPrefix(path: string): string {
+	// Extract the path after prefix and preserve query parameters
+	// Use URL API for cleaner parsing
+	const url = new URL(path, "http://actor");
+	const pathname = url.pathname.replace(/^\/raw\/websocket\/?/, "") || "/";
+	const normalizedPath =
+		(pathname.startsWith("/") ? pathname : "/" + pathname) + url.search;
+
+	return normalizedPath;
+}

package/src/actor/router.ts

@@ -1,5 +1,4 @@
 import { Hono, type Context as HonoContext } from "hono";
-import { cors } from "hono/cors";
 import invariant from "invariant";
 import { EncodingSchema } from "@/actor/protocol/serde";
 import {
@@ -42,6 +41,7 @@ import {
 } from "@/inspector/actor";
 import { isInspectorEnabled, secureInspector } from "@/inspector/utils";
 import type { RunnerConfig } from "@/registry/run-config";
+import { generateConnRequestId } from "./conn";
 import { ConnDriverKind } from "./conn-drivers";
 import type { ActorDriver } from "./driver";
 import { InternalError } from "./errors";
@@ -77,6 +77,17 @@ export function createActorRouter(
 
 	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",
@@ -175,6 +186,7 @@ export function createActorRouter(
 					c.env.actorId,
 					encoding,
 					connParams,
+					generateConnRequestId(),
 					connIdRaw,
 					connTokenRaw,
 				);
@@ -307,22 +319,18 @@ export function createActorRouter(
 			new Hono<
 				ActorInspectorRouterEnv & { Bindings: ActorRouterBindings }
 			>()
-				.use(
-					cors(runConfig.inspector.cors),
-					secureInspector(runConfig),
-					async (c, next) => {
-						const inspector = (
-							await actorDriver.loadActor(c.env.actorId)
-						).inspector;
-						invariant(
-							inspector,
-							"inspector not supported on this platform",
-						);
-
-						c.set("inspector", inspector);
-						return next();
-					},
-				)
+				.use(secureInspector(runConfig), async (c, next) => {
+					const inspector = (
+						await actorDriver.loadActor(c.env.actorId)
+					).inspector;
+					invariant(
+						inspector,
+						"inspector not supported on this platform",
+					);
+
+					c.set("inspector", inspector);
+					return next();
+				})
 				.route("/", createActorInspectorRouter()),
 		);
 	}

package/src/client/actor-conn.ts

@@ -563,17 +563,21 @@ enc
 		// TODO: Handle queue
 		// TODO: Reconnect with backoff
 
-		// Reject open promise
-		if (this.#onOpenPromise) {
-			this.#onOpenPromise.reject(new Error("Closed"));
-		}
-
 		// We can't use `event instanceof CloseEvent` because it's not defined in NodeJS
 		//
 		// These properties will be undefined
 		const closeEvent = event as CloseEvent;
 		const wasClean = closeEvent.wasClean;
 
+		// Reject open promise
+		if (this.#onOpenPromise) {
+			this.#onOpenPromise.reject(
+				new Error(
+					`websocket closed with code ${closeEvent.code}: ${closeEvent.reason}`,
+				),
+			);
+		}
+
 		logger().info({
 			msg: "socket closed",
 			code: closeEvent.code,

package/src/common/cors.ts

@@ -0,0 +1,57 @@
+import type { MiddlewareHandler } from "hono";
+
+/**
+ * Simple CORS middleware that matches the gateway behavior.
+ *
+ * - Echoes back the Origin header from the request
+ * - Echoes back the Access-Control-Request-Headers from preflight
+ * - Supports credentials
+ * - Allows common HTTP methods
+ * - Caches preflight for 24 hours
+ * - Adds Vary header to prevent cache poisoning
+ */
+export const cors = (): MiddlewareHandler => {
+	return async (c, next) => {
+		// Extract origin from request
+		const origin = c.req.header("origin") || "*";
+
+		// Handle preflight OPTIONS request
+		if (c.req.method === "OPTIONS") {
+			const requestHeaders =
+				c.req.header("access-control-request-headers") || "*";
+
+			c.header("access-control-allow-origin", origin);
+			c.header("access-control-allow-credentials", "true");
+			c.header(
+				"access-control-allow-methods",
+				"GET, POST, PUT, DELETE, OPTIONS, PATCH",
+			);
+			c.header("access-control-allow-headers", requestHeaders);
+			c.header("access-control-expose-headers", "*");
+			c.header("access-control-max-age", "86400");
+
+			// Add Vary header to prevent cache poisoning when echoing origin
+			if (origin !== "*") {
+				c.header("vary", "Origin");
+			}
+
+			// Remove content headers from preflight response
+			c.res.headers.delete("content-length");
+			c.res.headers.delete("content-type");
+
+			return c.body(null, 204);
+		}
+
+		await next();
+
+		// Add CORS headers to actual request
+		c.header("access-control-allow-origin", origin);
+		c.header("access-control-allow-credentials", "true");
+		c.header("access-control-expose-headers", "*");
+
+		// Add Vary header to prevent cache poisoning when echoing origin
+		if (origin !== "*") {
+			c.header("vary", "Origin");
+		}
+	};
+};

package/src/common/log.ts

@@ -146,7 +146,7 @@ export async function configureDefaultLogger(
 			},
 		},
 		hooks: {
-			logMethod(inputArgs, _method, level) {
+			logMethod(inputArgs, method, level) {
 				// TODO: This is a hack to not implement our own transport target. We can get better perf if we have our own transport target.
 
 				const levelMap: Record<number, string> = {
@@ -162,20 +162,41 @@ export async function configureDefaultLogger(
 					getEnvUniversal("LOG_TIMESTAMP") === "1"
 						? Date.now()
 						: undefined;
+
+				// Get bindings from the logger instance (child logger fields)
+				const bindings = (this as any).bindings?.() || {};
+
 				// TODO: This can be simplified in logfmt.ts
 				if (inputArgs.length >= 2) {
 					const [objOrMsg, msg] = inputArgs;
 					if (typeof objOrMsg === "object" && objOrMsg !== null) {
-						customWrite(levelName, { ...objOrMsg, msg, time });
+						customWrite(levelName, {
+							...bindings,
+							...objOrMsg,
+							msg,
+							time,
+						});
 					} else {
-						customWrite(levelName, { msg: String(objOrMsg), time });
+						customWrite(levelName, {
+							...bindings,
+							msg: String(objOrMsg),
+							time,
+						});
 					}
 				} else if (inputArgs.length === 1) {
 					const [objOrMsg] = inputArgs;
 					if (typeof objOrMsg === "object" && objOrMsg !== null) {
-						customWrite(levelName, { ...objOrMsg, time });
+						customWrite(levelName, {
+							...bindings,
+							...objOrMsg,
+							time,
+						});
 					} else {
-						customWrite(levelName, { msg: String(objOrMsg), time });
+						customWrite(levelName, {
+							...bindings,
+							msg: String(objOrMsg),
+							time,
+						});
 					}
 				}
 			},

package/src/common/utils.ts

@@ -1,7 +1,7 @@
 import type { Next } from "hono";
 import type { ContentfulStatusCode } from "hono/utils/http-status";
 import * as errors from "@/actor/errors";
-import { getEnvUniversal } from "@/utils";
+import { EXTRA_ERROR_LOG, getEnvUniversal, VERSION } from "@/utils";
 import type { Logger } from "./log";
 
 export function assertUnreachable(x: never): never {
@@ -228,8 +228,7 @@ export function deconstructError(
 			group,
 			code,
 			message,
-			issues: "https://github.com/rivet-dev/rivetkit/issues",
-			support: "https://rivet.dev/discord",
+			...EXTRA_ERROR_LOG,
 			...extraLog,
 		});
 	} else if (exposeInternalError) {
@@ -246,8 +245,7 @@ export function deconstructError(
 				group,
 				code,
 				message,
-				issues: "https://github.com/rivet-dev/rivetkit/issues",
-				support: "https://rivet.dev/discord",
+				...EXTRA_ERROR_LOG,
 				...extraLog,
 			});
 		} else {
@@ -262,8 +260,7 @@ export function deconstructError(
 				group,
 				code,
 				message,
-				issues: "https://github.com/rivet-dev/rivetkit/issues",
-				support: "https://rivet.dev/discord",
+				...EXTRA_ERROR_LOG,
 				...extraLog,
 			});
 		}
@@ -281,8 +278,7 @@ export function deconstructError(
 			msg: "internal error",
 			error: getErrorMessage(error),
 			stack: (error as Error)?.stack,
-			issues: "https://github.com/rivet-dev/rivetkit/issues",
-			support: "https://rivet.dev/discord",
+			...EXTRA_ERROR_LOG,
 			...extraLog,
 		});
 	}

package/src/common/websocket-interface.ts

@@ -3,10 +3,20 @@ export interface RivetEvent {
 	type: string;
 	target?: any;
 	currentTarget?: any;
+	/**
+	 * @experimental
+	 * Request ID for hibernatable websockets (provided by engine runner)
+	 **/
+	rivetRequestId?: ArrayBuffer;
 }
 
 export interface RivetMessageEvent extends RivetEvent {
 	data: any;
+	/**
+	 * @experimental
+	 * Message index for hibernatable websockets (provided by engine runner)
+	 **/
+	rivetMessageIndex?: number;
 }
 
 export interface RivetCloseEvent extends RivetEvent {

package/src/driver-helpers/utils.ts

@@ -15,6 +15,7 @@ export function serializeEmptyPersistData(
 		state: bufferToArrayBuffer(cbor.encode(undefined)),
 		connections: [],
 		scheduledEvents: [],
+		hibernatableWebSocket: [],
 	};
 	return PERSISTED_ACTOR_VERSIONED.serializeWithEmbeddedVersion(persistData);
 }