rivetkit 3.0.0-preview.3 → 3.0.0-preview.5

package/src/client/actor-conn.ts

@@ -50,9 +50,7 @@ import {
 import { ACTOR_CONNS_SYMBOL, type ClientRaw } from "./client";
 import * as errors from "./errors";
 import { isRetryableLifecycleReconnectSignal } from "./lifecycle-errors";
-import { retryAfterMs } from "./retry-budget";
 import { logger } from "./log";
-import { isInvalidToken } from "./token-provider";
 import { outboundTelemetryHeaders } from "./outbound-telemetry";
 import {
 	createQueueSender,
@@ -61,6 +59,8 @@ import {
 	type QueueSendResult,
 	type QueueSendWaitOptions,
 } from "./queue";
+import { retryAfterMs } from "./retry-budget";
+import { isInvalidToken } from "./token-provider";
 import {
 	type WebSocketMessage as ConnMessage,
 	messageLength,
@@ -160,6 +160,8 @@ export class ActorConnRaw {
 	#abortController = new AbortController();
 
 	#connStatus: ActorConnStatus = "idle";
+	/** The opening failure that stopped retries. Later actions must fail instead of queueing. */
+	#connectionOpenError?: errors.ActorError;
 
 	#actorId?: string;
 	#connId?: string;
@@ -356,6 +358,8 @@ export class ActorConnRaw {
 				`Invalid action call: expected an options object { name, args }, got ${typeof opts}. Use conn.actionName(...args) for the shorthand API.`,
 			);
 		}
+		if (this.#disposed) throw new errors.ActorConnDisposed();
+		if (this.#connectionOpenError) throw this.#connectionOpenError;
 		logger().debug({ msg: "action", name: opts.name, args: opts.args });
 
 		// If we have an active connection, use the websockactionId
@@ -466,6 +470,7 @@ export class ActorConnRaw {
 	}
 
 	#connectWithRetry() {
+		this.#connectionOpenError = undefined;
 		this.#setConnStatus("connecting");
 
 		// Attempt to reconnect indefinitely
@@ -523,6 +528,11 @@ export class ActorConnRaw {
 			// Wait for result
 			await this.#waitForOpen(this.#onOpenPromise.promise);
 		} catch (error) {
+			this.#abandonCurrentSocket(
+				isConnectionOpenTimeout(error)
+					? "Connect timeout"
+					: "Connection opening failed",
+			);
 			if (this.#shouldRetryConnectionOpenError(error)) {
 				throw error;
 			}
@@ -537,7 +547,7 @@ export class ActorConnRaw {
 							{ error: stringifyError(error) },
 						);
 
-			this.#clearQueuedMessages();
+			this.#connectionOpenError = actorError;
 			this.#rejectPendingPromises(actorError, false);
 			this.#dispatchActorError(actorError);
 			this.#setConnStatus("idle");
@@ -573,29 +583,24 @@ export class ActorConnRaw {
 		});
 		try {
 			await Promise.race([open, timedOut]);
-		} catch (error) {
-			if (isConnectionOpenTimeout(error)) {
-				this.#abandonCurrentSocket();
-			}
-			throw error;
 		} finally {
 			clearTimeout(timer);
 		}
 	}
 
-	#abandonCurrentSocket() {
+	#abandonCurrentSocket(reason: string) {
 		const ws = this.#websocket;
 		if (!ws) return;
 		this.#websocket = undefined;
 		this.#abandonedSockets.add(ws);
-		logger().warn({
-			msg: "connection did not open in time, retrying",
+		logger().debug({
+			msg: "closing failed connection attempt",
 			connId: this.#connId,
 			readyState: ws.readyState,
-			timeoutMs: CONNECT_TIMEOUT_MS,
+			reason,
 		});
 		try {
-			ws.close(1000, "Connect timeout");
+			ws.close(1000, reason);
 		} catch (err) {
 			logger().debug({
 				msg: "failed to close abandoned websocket",
@@ -611,6 +616,13 @@ export class ActorConnRaw {
 		if (isConnectionOpenTimeout(error)) {
 			return true;
 		}
+		if (
+			error instanceof errors.ActorError &&
+			error.group === "client" &&
+			error.code === "connection_open_interrupted"
+		) {
+			return true;
+		}
 		if (
 			error instanceof errors.ActorError &&
 			isInvalidToken(error.group, error.code) &&
@@ -692,16 +704,18 @@ export class ActorConnRaw {
 		});
 		this.#websocket = ws;
 		ws.addEventListener("open", () => {
-			if (this.#abandonedSockets.has(ws)) return;
+			if (this.#websocket !== ws || this.#abandonedSockets.has(ws))
+				return;
 			logger().debug({
 				msg: "client websocket open",
 				connId: this.#connId,
 			});
 		});
 		ws.addEventListener("message", async (ev: { data: ConnMessage }) => {
-			if (this.#abandonedSockets.has(ws)) return;
+			if (this.#websocket !== ws || this.#abandonedSockets.has(ws))
+				return;
 			try {
-				await this.#handleOnMessage(ev.data);
+				await this.#handleOnMessage(ev.data, ws);
 			} catch (err) {
 				logger().error({
 					msg: "error in websocket message handler",
@@ -710,7 +724,8 @@ export class ActorConnRaw {
 			}
 		});
 		ws.addEventListener("close", async (ev: Event | CloseEventLike) => {
-			if (this.#abandonedSockets.has(ws)) return;
+			if (this.#websocket !== ws || this.#abandonedSockets.has(ws))
+				return;
 			try {
 				await this.#handleOnClose(ev);
 			} catch (err) {
@@ -721,7 +736,8 @@ export class ActorConnRaw {
 			}
 		});
 		ws.addEventListener("error", () => {
-			if (this.#abandonedSockets.has(ws)) return;
+			if (this.#websocket !== ws || this.#abandonedSockets.has(ws))
+				return;
 			try {
 				this.#handleOnError();
 			} catch (err) {
@@ -803,7 +819,7 @@ export class ActorConnRaw {
 	}
 
 	/** Called by the onmessage event from drivers. */
-	async #handleOnMessage(data: any) {
+	async #handleOnMessage(data: any, socket: UniversalWebSocket) {
 		logger().trace({
 			msg: "received message",
 			dataType: typeof data,
@@ -812,6 +828,10 @@ export class ActorConnRaw {
 		});
 
 		const response = await this.#parseMessage(data as ConnMessage);
+		// Decoding can yield while this socket closes and another connects. Apply
+		// the message only while its socket still owns the connection.
+		if (this.#websocket !== socket || this.#abandonedSockets.has(socket))
+			return;
 		logger().trace(
 			getLogMessage()
 				? {
@@ -897,23 +917,19 @@ export class ActorConnRaw {
 						this.#actorResolutionState,
 						this.#driver,
 					);
+					if (
+						this.#websocket !== socket ||
+						this.#abandonedSockets.has(socket)
+					) {
+						return;
+					}
 					if (schedulingError) {
 						errorToThrow = schedulingError;
 					}
 				}
 
-				// If we have an onOpenPromise, reject it with the error
-				if (this.#onOpenPromise) {
-					this.#onOpenPromise.reject(errorToThrow);
-				}
-
 				this.#invalidateActorIfStale(group, code);
-
-				// Reject any in-flight requests
-				for (const [id, inFlight] of this.#actionsInFlight.entries()) {
-					inFlight.reject(errorToThrow);
-					this.#actionsInFlight.delete(id);
-				}
+				this.#rejectPendingPromises(errorToThrow, false);
 
 				this.#dispatchActorError(errorToThrow);
 			}
@@ -1047,9 +1063,17 @@ export class ActorConnRaw {
 				this.#invalidateActorIfStale(group, code);
 			} else {
 				// Default error for non-structured close reasons
-				error = new Error(
-					`${wasClean ? "Connection closed" : "Connection lost"} (code: ${closeEvent.code}, reason: ${reason})`,
-				);
+				const message = `${wasClean ? "Connection closed" : "Connection lost"} (code: ${closeEvent.code}, reason: ${reason})`;
+				// A transport loss during the handshake can outlast an Engine
+				// restart. Retry opening without replaying actions with unknown outcomes.
+				error =
+					!wasConnected && closeEvent.code === 1006
+						? new errors.ActorError(
+								"client",
+								"connection_open_interrupted",
+								message,
+							)
+						: new Error(message);
 			}
 
 			this.#rejectPendingPromises(error, false);
@@ -1071,6 +1095,8 @@ export class ActorConnRaw {
 	}
 
 	#rejectPendingPromises(error: Error, suppressUnhandled: boolean) {
+		// Rejected actions must not survive in the queue and execute after a reconnect.
+		this.#clearQueuedMessages();
 		if (this.#onOpenPromise) {
 			if (suppressUnhandled) {
 				this.#onOpenPromise.promise.catch(() => {});

package/src/client/lifecycle-errors.ts

@@ -77,6 +77,19 @@ function classifyActorError(
 		);
 	}
 
+	if (
+		(error.group === "ups" && error.code === "consumer_closed") ||
+		(error.group === "depot" && error.code === "request_deadline_exceeded")
+	) {
+		// The connection can recover, but an interrupted action may already have run.
+		return buildLifecycleBoundaryInfo(
+			"reconnect_only",
+			"actor_error",
+			error.message,
+			{ group: error.group, code: error.code },
+		);
+	}
+
 	if (
 		error.group === "guard" &&
 		isRetryableGuardGatewayHttpError(error.code)

package/src/registry/config/index.ts

@@ -291,7 +291,8 @@ export const RegistryConfigSchema = z
 				 * after calling `CoreRegistry::shutdown()`. Defaults to the
 				 * engine-provided actor stop threshold once the node connects. A
 				 * signal-driven shutdown of a managed local development engine waits
-				 * one second by default so a file watcher can restart promptly.
+				 * four seconds by default so actors can acknowledge stopping before
+				 * a file watcher restarts the process.
 				 *
 				 * Set a longer period when a complete graceful node drain is required.
 				 */

package/src/registry/index.ts

@@ -29,6 +29,13 @@ import type {
 type ShutdownSignal = "SIGINT" | "SIGTERM";
 
 const REGISTRY_READY_TIMEOUT_MS = 30_000;
+// A managed Engine restart can wait for the old process to release storage and
+// for its worker lease to expire before the replacement can register.
+const LOCAL_ENGINE_READY_TIMEOUT_MS = 60_000;
+// Core gives local actors three seconds to drain and acknowledge their stopped
+// reports. Leave time to close the node connection after that, while staying
+// below the five-second force-kill ceiling used by tsx watch.
+const LOCAL_DEV_SHUTDOWN_GRACE_PERIOD_MS = 4_000;
 
 // Serverless is disabled for v3; keep the implementation for a later release.
 function rejectServerless(): void {
@@ -789,7 +796,7 @@ export class Registry<A extends RegistryActors> {
 		let gracePeriodMs = config.shutdown?.gracePeriodMs;
 		if (gracePeriodMs === undefined) {
 			gracePeriodMs = localDevSignal
-				? 1_000
+				? LOCAL_DEV_SHUTDOWN_GRACE_PERIOD_MS
 				: ((await this.#actorStopThresholdMs(
 						modeAPromise ?? modeBPromise,
 					)) ?? 30 * 60 * 1000);
@@ -942,16 +949,22 @@ export class Registry<A extends RegistryActors> {
 			throw new Error("registry envoy startup did not initialize");
 		}
 
+		const readyTimeoutMs =
+			config.engineSpawn &&
+			config.endpoint !== undefined &&
+			isLocalEngineEndpoint(config.endpoint)
+				? LOCAL_ENGINE_READY_TIMEOUT_MS
+				: REGISTRY_READY_TIMEOUT_MS;
 		let timeout: ReturnType<typeof setTimeout> | undefined;
 		const timeoutPromise = new Promise<never>((_resolve, reject) => {
 			timeout = setTimeout(
 				() =>
 					reject(
 						new Error(
-							`RivetKit registry did not register with the Engine within ${REGISTRY_READY_TIMEOUT_MS}ms`,
+							`RivetKit registry did not register with the Engine within ${readyTimeoutMs}ms`,
 						),
 					),
-				REGISTRY_READY_TIMEOUT_MS,
+				readyTimeoutMs,
 			);
 			timeout.unref?.();
 		});