rivetkit 2.0.6 → 2.0.7

package/src/actor/instance.ts

@@ -1,7 +1,9 @@
 import * as cbor from "cbor-x";
+import type { SSEStreamingApi } from "hono/streaming";
+import type { WSContext } from "hono/ws";
 import invariant from "invariant";
 import onChange from "on-change";
-import type { ActorKey } from "@/actor/mod";
+import type { ActorKey, Encoding } from "@/actor/mod";
 import type { Client } from "@/client/client";
 import { getBaseLogger, getIncludeTarget, type Logger } from "@/common/log";
 import { isCborSerializable, stringifyError } from "@/common/utils";
@@ -15,22 +17,24 @@ import { TO_CLIENT_VERSIONED } from "@/schemas/client-protocol/versioned";
 import {
 	bufferToArrayBuffer,
 	getEnvUniversal,
+	promiseWithResolvers,
 	SinglePromiseQueue,
 } from "@/utils";
 import type { ActionContext } from "./action";
 import type { ActorConfig, OnConnectOptions } from "./config";
+import { Conn, type ConnId, generateConnId, generateConnToken } from "./conn";
 import {
-	CONNECTION_CHECK_LIVENESS_SYMBOL,
-	Conn,
-	type ConnectionDriver,
-	type ConnId,
-} from "./connection";
+	CONN_DRIVERS,
+	type ConnDriver,
+	type ConnDriverState,
+	getConnDriverKindFromState,
+} from "./conn-drivers";
+import type { ConnSocket } from "./conn-socket";
 import { ActorContext } from "./context";
 import type { AnyDatabaseProvider, InferDatabaseClient } from "./database";
-import type { ActorDriver, ConnDriver, ConnectionDriversMap } from "./driver";
+import type { ActorDriver } from "./driver";
 import * as errors from "./errors";
 import { serializeActorKey } from "./keys";
-import { loggerWithoutContext } from "./log";
 import type {
 	PersistedActor,
 	PersistedConn,
@@ -161,7 +165,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	#backgroundPromises: Promise<void>[] = [];
 	#abortController = new AbortController();
 	#config: ActorConfig<S, CP, CS, V, I, DB>;
-	#connectionDrivers!: ConnectionDriversMap;
 	#actorDriver!: ActorDriver;
 	#inlineClient!: Client<Registry<any>>;
 	#actorId!: string;
@@ -206,9 +209,9 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			getConnections: async () => {
 				return Array.from(this.#connections.entries()).map(([id, conn]) => ({
 					id,
-					stateEnabled: conn._stateEnabled,
-					params: conn.params as {},
-					state: conn._stateEnabled ? conn.state : undefined,
+					stateEnabled: conn.__stateEnabled,
+					params: conn.params as any,
+					state: conn.__stateEnabled ? conn.state : undefined,
 				}));
 			},
 			setState: async (state: unknown) => {
@@ -254,7 +257,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	}
 
 	async start(
-		connectionDrivers: ConnectionDriversMap,
 		actorDriver: ActorDriver,
 		inlineClient: Client<Registry<any>>,
 		actorId: string,
@@ -278,7 +280,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			),
 		);
 
-		this.#connectionDrivers = connectionDrivers;
 		this.#actorDriver = actorDriver;
 		this.#inlineClient = inlineClient;
 		this.#actorId = actorId;
@@ -528,7 +529,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		}
 	}
 
-	get #connStateEnabled() {
+	get connStateEnabled() {
 		return "createConnState" in this.#config || "connState" in this.#config;
 	}
 
@@ -543,7 +544,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	}
 
 	/** Promise used to wait for a save to complete. This is required since you cannot await `#saveStateThrottled`. */
-	#onPersistSavedPromise?: PromiseWithResolvers<void>;
+	#onPersistSavedPromise?: ReturnType<typeof promiseWithResolvers<void>>;
 
 	/** Throttled save state method. Used to write to KV at a reasonable cadence. */
 	#savePersistThrottled() {
@@ -719,13 +720,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			// Load connections
 			for (const connPersist of this.#persist.connections) {
 				// Create connections
-				const driver = this.__getConnDriver(connPersist.connDriver);
-				const conn = new Conn<S, CP, CS, V, I, DB>(
-					this,
-					connPersist,
-					driver,
-					this.#connStateEnabled,
-				);
+				const conn = new Conn<S, CP, CS, V, I, DB>(this, connPersist);
 				this.#connections.set(conn.id, conn);
 
 				// Register event subscriptions
@@ -791,14 +786,52 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	}
 
 	/**
-	 * Removes a connection and cleans up its resources.
+	 * Connection disconnected.
+	 *
+	 * If a clean diconnect, will be removed immediately.
+	 *
+	 * If not a clean disconnect, will keep the connection alive for a given interval to wait for reconnect.
 	 */
-	__removeConn(conn: Conn<S, CP, CS, V, I, DB> | undefined) {
-		if (!conn) {
-			this.#rLog.warn({ msg: "`conn` does not exist" });
+	__connDisconnected(
+		conn: Conn<S, CP, CS, V, I, DB>,
+		wasClean: boolean,
+		socketId: string,
+	) {
+		// If socket ID is provided, check if it matches the current socket ID
+		// If it doesn't match, this is a stale disconnect event from an old socket
+		if (socketId && conn.__socket && socketId !== conn.__socket.socketId) {
+			this.rLog.debug({
+				msg: "ignoring stale disconnect event",
+				connId: conn.id,
+				eventSocketId: socketId,
+				currentSocketId: conn.__socket.socketId,
+			});
 			return;
 		}
 
+		if (wasClean) {
+			// Disconnected cleanly, remove the conn
+
+			this.#removeConn(conn);
+		} else {
+			// Disconnected uncleanly, allow reconnection
+
+			if (!conn.__driverState) {
+				this.rLog.warn("called conn disconnected without driver state");
+			}
+
+			// Update last seen so we know when to clean it up
+			conn.__persist.lastSeen = Date.now();
+
+			// Remove socket
+			conn.__socket = undefined;
+		}
+	}
+
+	/**
+	 * Removes a connection and cleans up its resources.
+	 */
+	#removeConn(conn: Conn<S, CP, CS, V, I, DB>) {
 		// Remove from persist & save immediately
 		const connIdx = this.#persist.connections.findIndex(
 			(c) => c.connId === conn.id,
@@ -846,12 +879,96 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		this.#resetSleepTimer();
 	}
 
-	async prepareConn(
+	/**
+	 * Called to create a new connection or reconnect an existing one.
+	 */
+	async createConn(
+		socket: ConnSocket,
 		// biome-ignore lint/suspicious/noExplicitAny: TypeScript bug with ExtractActorConnParams<this>,
 		params: any,
 		request?: Request,
-	): Promise<CS> {
-		// Authenticate connection
+		connectionId?: string,
+		connectionToken?: string,
+	): Promise<Conn<S, CP, CS, V, I, DB>> {
+		this.#assertReady();
+
+		// If connection ID and token are provided, try to reconnect
+		if (connectionId && connectionToken) {
+			this.rLog.debug({
+				msg: "checking for existing connection",
+				connectionId,
+			});
+			const existingConn = this.#connections.get(connectionId);
+			if (existingConn && existingConn._token === connectionToken) {
+				// This is a valid reconnection
+				this.rLog.debug({
+					msg: "reconnecting existing connection",
+					connectionId,
+				});
+
+				// If there's an existing driver state, clean it up without marking as clean disconnect
+				if (existingConn.__driverState) {
+					const driverKind = getConnDriverKindFromState(
+						existingConn.__driverState,
+					);
+					const driver = CONN_DRIVERS[driverKind];
+					if (driver.disconnect) {
+						// Call driver disconnect to clean up directly. Don't use Conn.disconnect since that will remove the connection entirely.
+						driver.disconnect(
+							this,
+							existingConn,
+							(existingConn.__driverState as any)[driverKind],
+							"Reconnecting with new driver state",
+						);
+					}
+				}
+
+				// Update with new driver state
+				existingConn.__socket = socket;
+				existingConn.__persist.lastSeen = Date.now();
+
+				// Update sleep timer since connection is now active
+				this.#resetSleepTimer();
+
+				this.inspector.emitter.emit("connectionUpdated");
+
+				// Send init message for reconnection
+				existingConn._sendMessage(
+					new CachedSerializer<protocol.ToClient>(
+						{
+							body: {
+								tag: "Init",
+								val: {
+									actorId: this.id,
+									connectionId: existingConn.id,
+									connectionToken: existingConn._token,
+								},
+							},
+						},
+						TO_CLIENT_VERSIONED,
+					),
+				);
+
+				return existingConn;
+			}
+
+			// If we get here, either connection doesn't exist or token doesn't match
+			// Fall through to create new connection with new IDs
+			this.rLog.debug({
+				msg: "connection not found or token mismatch, creating new connection",
+				connectionId,
+			});
+		}
+
+		// Generate new connection ID and token if not provided or if reconnection failed
+		const newConnId = generateConnId();
+		const newConnToken = generateConnToken();
+
+		if (this.#connections.has(newConnId)) {
+			throw new Error(`Connection already exists: ${newConnId}`);
+		}
+
+		// Prepare connection state
 		let connState: CS | undefined;
 
 		const onBeforeConnectOpts = {
@@ -866,7 +983,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			);
 		}
 
-		if (this.#connStateEnabled) {
+		if (this.connStateEnabled) {
 			if ("createConnState" in this.#config) {
 				const dataOrPromise = this.#config.createConnState(
 					this.actorContext as unknown as ActorContext<
@@ -897,53 +1014,17 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			}
 		}
 
-		return connState as CS;
-	}
-
-	__getConnDriver(driverId: ConnectionDriver): ConnDriver {
-		// Get driver
-		const driver = this.#connectionDrivers[driverId];
-		if (!driver) throw new Error(`No connection driver: ${driverId}`);
-		return driver;
-	}
-
-	/**
-	 * Called after establishing a connection handshake.
-	 */
-	async createConn(
-		connectionId: string,
-		connectionToken: string,
-		params: CP,
-		state: CS,
-		driverId: ConnectionDriver,
-		driverState: unknown,
-		authData: unknown,
-	): Promise<Conn<S, CP, CS, V, I, DB>> {
-		this.#assertReady();
-
-		if (this.#connections.has(connectionId)) {
-			throw new Error(`Connection already exists: ${connectionId}`);
-		}
-
 		// Create connection
-		const driver = this.__getConnDriver(driverId);
 		const persist: PersistedConn<CP, CS> = {
-			connId: connectionId,
-			token: connectionToken,
-			connDriver: driverId,
-			connDriverState: driverState,
+			connId: newConnId,
+			token: newConnToken,
 			params: params,
-			state: state,
-			authData: authData,
+			state: connState as CS,
 			lastSeen: Date.now(),
 			subscriptions: [],
 		};
-		const conn = new Conn<S, CP, CS, V, I, DB>(
-			this,
-			persist,
-			driver,
-			this.#connStateEnabled,
-		);
+		const conn = new Conn<S, CP, CS, V, I, DB>(this, persist);
+		conn.__socket = socket;
 		this.#connections.set(conn.id, conn);
 
 		// Update sleep
@@ -1129,11 +1210,10 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		this.#rLog.debug({ msg: "checking connections liveness" });
 
 		for (const conn of this.#connections.values()) {
-			const liveness = conn[CONNECTION_CHECK_LIVENESS_SYMBOL]();
-			if (liveness.status === "connected") {
+			if (conn.__status === "connected") {
 				this.#rLog.debug({ msg: "connection is alive", connId: conn.id });
 			} else {
-				const lastSeen = liveness.lastSeen;
+				const lastSeen = conn.__persist.lastSeen;
 				const sinceLastSeen = Date.now() - lastSeen;
 				if (sinceLastSeen < this.#config.options.connectionLivenessTimeout) {
 					this.#rLog.debug({
@@ -1152,9 +1232,8 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 					lastSeen,
 				});
 
-				// TODO: Do we need to force disconnect the connection here?
-
-				this.__removeConn(conn);
+				// Assume that the connection is dead here, no need to disconnect anything
+				this.#removeConn(conn);
 			}
 		}
 	}
@@ -1309,7 +1388,10 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	/**
 	 * Handles raw HTTP requests to the actor.
 	 */
-	async handleFetch(request: Request, opts: {}): Promise<Response> {
+	async handleFetch(
+		request: Request,
+		opts: Record<never, never>,
+	): Promise<Response> {
 		this.#assertReady();
 
 		if (!this.#config.onFetch) {
@@ -1562,7 +1644,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			} else {
 				// Create callback
 				if (!this.#onPersistSavedPromise) {
-					this.#onPersistSavedPromise = Promise.withResolvers();
+					this.#onPersistSavedPromise = promiseWithResolvers();
 				}
 
 				// Save state throttled
@@ -1783,10 +1865,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			connections: persist.connections.map((conn) => ({
 				id: conn.connId,
 				token: conn.token,
-				driver: conn.connDriver as string,
-				driverState: bufferToArrayBuffer(
-					cbor.encode(conn.connDriverState || {}),
-				),
 				parameters: bufferToArrayBuffer(cbor.encode(conn.params || {})),
 				state: bufferToArrayBuffer(cbor.encode(conn.state || {})),
 				subscriptions: conn.subscriptions.map((sub) => ({
@@ -1820,8 +1898,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			connections: bareData.connections.map((conn) => ({
 				connId: conn.id,
 				token: conn.token,
-				connDriver: conn.driver as ConnectionDriver,
-				connDriverState: cbor.decode(new Uint8Array(conn.driverState)),
 				params: cbor.decode(new Uint8Array(conn.parameters)),
 				state: cbor.decode(new Uint8Array(conn.state)),
 				subscriptions: conn.subscriptions.map((sub) => ({

package/src/actor/mod.ts

@@ -76,16 +76,10 @@ export type { ActionContext } from "./action";
 export type * from "./config";
 export type {
 	Conn,
-	ConnectionDriver,
 	ConnectionStatus,
 	generateConnId,
 	generateConnToken,
-} from "./connection";
-export {
-	CONNECTION_DRIVER_HTTP,
-	CONNECTION_DRIVER_SSE,
-	CONNECTION_DRIVER_WEBSOCKET,
-} from "./connection";
+} from "./conn";
 export type { ActorContext } from "./context";
 export type {
 	ActionContextOf,
@@ -95,10 +89,6 @@ export type {
 } from "./definition";
 export { lookupInRegistry } from "./definition";
 export { UserError, type UserErrorOptions } from "./errors";
-export {
-	createGenericConnDrivers,
-	GenericConnGlobalState,
-} from "./generic-conn-driver";
 export type { AnyActorInstance } from "./instance";
 export {
 	type ActorRouter,

package/src/actor/persisted.ts

@@ -1,5 +1,3 @@
-import type { ConnectionDriver } from "./connection";
-
 /** State object that gets automatically persisted to storage. */
 export interface PersistedActor<S, CP, CS, I> {
 	input?: I;
@@ -13,12 +11,11 @@ export interface PersistedActor<S, CP, CS, I> {
 export interface PersistedConn<CP, CS> {
 	connId: string;
 	token: string;
-	connDriver: ConnectionDriver;
-	connDriverState: unknown;
 	params: CP;
 	state: CS;
-	authData?: unknown;
 	subscriptions: PersistedSubscription[];
+
+	/** Last time the socket was seen. This is set when disconencted so we can determine when we need to clean this up. */
 	lastSeen: number;
 }
 

package/src/actor/protocol/old.ts

@@ -16,7 +16,7 @@ import {
 import { deserializeWithEncoding } from "@/serde";
 import { assertUnreachable, bufferToArrayBuffer } from "../../utils";
 import { ActionContext } from "../action";
-import type { Conn } from "../connection";
+import type { Conn } from "../conn";
 import type { ActorInstance } from "../instance";
 
 export const TransportSchema = z.enum(["websocket", "sse"]);