rivetkit 2.0.6 → 2.0.7

package/src/driver-test-suite/tests/actor-conn-state.ts

@@ -145,38 +145,79 @@ export function runActorConnStateTests(driverTestConfig: DriverTestConfig) {
 		});
 
 		describe("Connection Lifecycle", () => {
-			test("should track connection and disconnection events", async (c) => {
+			test.skipIf(
+				driverTestConfig.transport === "sse" &&
+					driverTestConfig.clientType === "inline",
+			)("should track connection and disconnection events", async (c) => {
 				const { client } = await setupDriverTest(c, driverTestConfig);
 
+				const debugHandle = client.connStateActor.getOrCreate(undefined, {
+					params: { noCount: true },
+				});
+
 				// Create a connection
-				const handle = client.connStateActor.getOrCreate();
-				const conn = handle.connect();
+				const conn = client.connStateActor.getOrCreate().connect();
 
 				// Get the connection state
 				const connState = await conn.getConnectionState();
 
 				// Verify the connection is tracked
-				const connectionIds = await conn.getConnectionIds();
-				expect(connectionIds).toContain(connState.id);
+				await vi.waitFor(async () => {
+					const connectionIds = await debugHandle.getConnectionIds();
+					expect(connectionIds).toContain(connState.id);
+				});
 
 				// Initial disconnection count
-				const initialDisconnections = await conn.getDisconnectionCount();
+				await vi.waitFor(async () => {
+					const disconnects = await debugHandle.getDisconnectionCount();
+					expect(disconnects).toBe(0);
+				});
 
 				// Dispose the connection
 				await conn.dispose();
 
+				// Validate conn count
+				await vi.waitFor(
+					async () => {
+						console.log("disconnects before");
+						const disconnects = await debugHandle.getDisconnectionCount();
+						console.log("disconnects", disconnects);
+						expect(disconnects).toBe(1);
+					},
+					// SSE takes a long time to disconnect on CF Workers
+					{
+						timeout: 10_000,
+						interval: 100,
+					},
+				);
+
 				// Create a new connection to check the disconnection count
-				const newConn = handle.connect();
+				const newConn = client.connStateActor.getOrCreate().connect();
 
-				// Verify disconnection was tracked
+				// Verify the connection is tracked
 				await vi.waitFor(async () => {
-					const newDisconnections = await newConn.getDisconnectionCount();
-
-					expect(newDisconnections).toBeGreaterThan(initialDisconnections);
+					const connectionIds = await debugHandle.getConnectionIds();
+					console.log("conn ids", connectionIds);
+					expect(connectionIds.length).toBe(1);
 				});
 
 				// Clean up
 				await newConn.dispose();
+
+				// Verify disconnection was tracked
+				await vi.waitFor(
+					async () => {
+						console.log("A");
+						const disconnects = await debugHandle.getDisconnectionCount();
+						console.log(`B ${disconnects}`);
+						expect(disconnects).toBe(2);
+					},
+					// SSE takes a long time to disconnect on CF Workers
+					{
+						timeout: 10_000,
+						interval: 100,
+					},
+				);
 			});
 
 			test("should update connection state", async (c) => {
@@ -228,17 +269,20 @@ export function runActorConnStateTests(driverTestConfig: DriverTestConfig) {
 					receivedMessages.push(data);
 				});
 
-				// Send message from first connection to second
-				const success = await conn1.sendToConnection(
-					state2.id,
-					"Hello from conn1",
-				);
-				expect(success).toBe(true);
-
-				// Verify message was received
-				expect(receivedMessages.length).toBe(1);
-				expect(receivedMessages[0].from).toBe(state1.id);
-				expect(receivedMessages[0].message).toBe("Hello from conn1");
+				// TODO: SSE has race condition between subscrib eand publish message
+				await vi.waitFor(async () => {
+					// Send message from first connection to second
+					const success = await conn1.sendToConnection(
+						state2.id,
+						"Hello from conn1",
+					);
+					expect(success).toBe(true);
+
+					// Verify message was received
+					expect(receivedMessages.length).toBe(1);
+					expect(receivedMessages[0].from).toBe(state1.id);
+					expect(receivedMessages[0].message).toBe("Hello from conn1");
+				});
 
 				// Clean up
 				await conn1.dispose();

package/src/driver-test-suite/tests/actor-conn.ts

@@ -1,4 +1,5 @@
-import { describe, expect, test } from "vitest";
+import { describe, expect, test, vi } from "vitest";
+import { SSE_PING_INTERVAL } from "@/actor/router-endpoints";
 import type { DriverTestConfig } from "../mod";
 import { FAKE_TIME, setupDriverTest, waitFor } from "../utils";
 
@@ -84,28 +85,27 @@ export function runActorConnTests(driverTestConfig: DriverTestConfig) {
 
 				// Set up event listener
 				const receivedEvents: number[] = [];
-				const receivedEventsPromise = new Promise((resolve) => {
-					connection.on("newCount", (count: number) => {
-						receivedEvents.push(count);
-						if (
-							receivedEvents.includes(1) &&
-							receivedEvents.includes(6) &&
-							receivedEvents.includes(9)
-						)
-							resolve(undefined);
-					});
+				connection.on("newCount", (count: number) => {
+					receivedEvents.push(count);
 				});
 
-				// Send one RPC call over the connection to ensure it's open
-				await connection.increment(1);
+				// TODO: There is a race condition with opening subscription and sending events on SSE, so we need to wait for a successful round trip on the event
+				await vi.waitFor(async () => {
+					// Send one RPC call over the connection to ensure it's open
+					await connection.setCount(1);
+					expect(receivedEvents).includes(1);
+				});
 
 				// Now use stateless RPC calls through the handle (no connection)
 				// These should still trigger events that the connection receives
-				await handle.increment(5);
-				await handle.increment(3);
+				await handle.setCount(2);
+				await handle.setCount(3);
 
 				// Wait for all events to be received
-				await receivedEventsPromise;
+				await vi.waitFor(() => {
+					expect(receivedEvents).includes(2);
+					expect(receivedEvents).includes(3);
+				});
 
 				// Clean up
 				await connection.dispose();
@@ -124,13 +124,17 @@ export function runActorConnTests(driverTestConfig: DriverTestConfig) {
 					receivedEvents.push(count);
 				});
 
-				// Trigger broadcast events
-				await connection.increment(5);
-				await connection.increment(3);
-
+				// HACK: Race condition between subscribing & sending events in SSE
 				// Verify events were received
-				expect(receivedEvents).toContain(5);
-				expect(receivedEvents).toContain(8);
+				await vi.waitFor(
+					async () => {
+						await connection.setCount(5);
+						await connection.setCount(8);
+						expect(receivedEvents).toContain(5);
+						expect(receivedEvents).toContain(8);
+					},
+					{ timeout: 10_000 },
+				);
 
 				// Clean up
 				await connection.dispose();
@@ -154,8 +158,10 @@ export function runActorConnTests(driverTestConfig: DriverTestConfig) {
 				await connection.increment(3);
 
 				// Verify only the first event was received
-				expect(receivedEvents).toEqual([5]);
-				expect(receivedEvents).not.toContain(8);
+				await vi.waitFor(() => {
+					expect(receivedEvents).toEqual([5]);
+					expect(receivedEvents).not.toContain(8);
+				});
 
 				// Clean up
 				await connection.dispose();
@@ -174,17 +180,20 @@ export function runActorConnTests(driverTestConfig: DriverTestConfig) {
 					receivedEvents.push(count);
 				});
 
+				// TODO: SSE has race condition with subscriptions & publishing messages
 				// Trigger first event
-				await connection.increment(5);
+				await vi.waitFor(async () => {
+					await connection.setCount(5);
+					expect(receivedEvents).toEqual([5]);
+				});
 
 				// Unsubscribe
 				unsubscribe();
 
 				// Trigger second event, should not be received
-				await connection.increment(3);
+				await connection.setCount(8);
 
 				// Verify only the first event was received
-				expect(receivedEvents).toEqual([5]);
 				expect(receivedEvents).not.toContain(8);
 
 				// Clean up
@@ -225,7 +234,10 @@ export function runActorConnTests(driverTestConfig: DriverTestConfig) {
 		});
 
 		describe("Lifecycle Hooks", () => {
-			test("should trigger lifecycle hooks", async (c) => {
+			test.skipIf(
+				driverTestConfig.transport === "sse" &&
+					driverTestConfig.clientType === "inline",
+			)("should trigger lifecycle hooks", async (c) => {
 				const { client } = await setupDriverTest(c, driverTestConfig);
 
 				// Create and connect
@@ -244,104 +256,31 @@ export function runActorConnTests(driverTestConfig: DriverTestConfig) {
 				// Disconnect should trigger onDisconnect
 				await connection.dispose();
 
-				// Reconnect to check if onDisconnect was called
-				const handle = client.counterWithLifecycle.getOrCreate([
-					"test-lifecycle",
-				]);
-				const finalEvents = await handle.getEvents();
-				expect(finalEvents).toBeOneOf([
-					// Still active
-					["onStart", "onBeforeConnect", "onConnect", "onDisconnect"],
-					// Went to sleep and woke back up
-					[
-						"onStart",
-						"onBeforeConnect",
-						"onConnect",
-						"onDisconnect",
-						"onStart",
-					],
-				]);
-			});
-		});
-
-		describe("Connection Liveness", () => {
-			// TODO: KIT-242
-			test.skip("should return correct liveness status for connections", async (c) => {
-				const { client } = await setupDriverTest(c, driverTestConfig);
-
-				// Create actor and connection
-				const handle = client.connLivenessActor.getOrCreate([
-					"test-liveness-status",
-				]);
-				const connA = handle.connect();
-				const connB = handle.connect();
-
-				const connAId = await connA.getConnectionId();
-				const connBId = await connB.getConnectionId();
-
-				// Verify connection works initially
-				await connA.increment(5);
-				await connB.increment(5);
-
-				const counter = await handle.getCounter();
-				expect(counter).toBe(10);
-
-				const connectionsStatusBeforeKill =
-					await handle.getWsConnectionsLiveness();
-				expect(connectionsStatusBeforeKill).toHaveLength(2);
-				expect(connectionsStatusBeforeKill).toContainEqual(
-					expect.objectContaining({
-						id: connAId,
-						status: "connected",
-						lastSeen: FAKE_TIME.getTime(),
-					}),
-				);
-				expect(connectionsStatusBeforeKill).toContainEqual(
-					expect.objectContaining({
-						id: connBId,
-						status: "connected",
-						lastSeen: FAKE_TIME.getTime(),
-					}),
-				);
-
-				// Kill one connection
-				await handle.kill(connAId); // instead of dispose, we use kill to simulate a disconnection (e.g. network failure)
-				// connA.dispose();
-				// we killed the connection, but the actor instance does not know about it yet
-				// it should still be in the list of connections, but with a status of "reconnecting"
-				const connectionsStatusAfterKill =
-					await handle.getWsConnectionsLiveness();
-				expect(connectionsStatusAfterKill).toEqual(
-					expect.arrayContaining([
-						expect.objectContaining({
-							id: connAId,
-							status: "reconnecting",
-							lastSeen: FAKE_TIME.getTime(),
-						}),
-						expect.objectContaining({
-							id: connBId,
-							status: "connected",
-							lastSeen: FAKE_TIME.getTime(),
-						}),
-					]),
-				);
-
-				// default time to wait for cleanup is 5 seconds
-				// check actor options
-				await waitFor(driverTestConfig, 5_000);
-
-				// After timeout, the killed connection should be unavailable, since the manager has cleaned it up
-				const connectionsStatusAfterCleanup =
-					await handle.getWsConnectionsLiveness();
-				expect(connectionsStatusAfterCleanup).not.toContainEqual(
-					expect.objectContaining({
-						id: connAId,
-					}),
-				);
-				expect(connectionsStatusAfterCleanup).toContainEqual(
-					expect.objectContaining({
-						id: connBId,
-					}),
+				await vi.waitFor(
+					async () => {
+						// Reconnect to check if onDisconnect was called
+						const handle = client.counterWithLifecycle.getOrCreate([
+							"test-lifecycle",
+						]);
+						const finalEvents = await handle.getEvents();
+						expect(finalEvents).toBeOneOf([
+							// Still active
+							["onStart", "onBeforeConnect", "onConnect", "onDisconnect"],
+							// Went to sleep and woke back up
+							[
+								"onStart",
+								"onBeforeConnect",
+								"onConnect",
+								"onDisconnect",
+								"onStart",
+							],
+						]);
+					},
+					// NOTE: High timeout required for Cloudflare Workers
+					{
+						timeout: 10_000,
+						interval: 100,
+					},
 				);
 			});
 		});

package/src/driver-test-suite/tests/actor-reconnect.ts

@@ -0,0 +1,160 @@
+import { describe, expect, test, vi } from "vitest";
+import type { ActorConnRaw } from "@/client/actor-conn";
+import type { DriverTestConfig } from "../mod";
+import { setupDriverTest } from "../utils";
+
+export function runActorReconnectTests(driverTestConfig: DriverTestConfig) {
+	describe("Actor Reconnection Tests", () => {
+		test("should reconnect and preserve connection state after non-clean disconnect", async (c) => {
+			const { client, endpoint } = await setupDriverTest(c, driverTestConfig);
+
+			// Create actor and connect
+			const handle = client.counterConn.getOrCreate(["test-reconnect"]);
+			const connection = handle.connect();
+
+			// Set an initial count on the connection
+			await connection.increment(5);
+
+			// Verify connection count is 1
+			const connCount1 = await connection.getConnectionCount();
+			expect(connCount1).toBe(1);
+
+			// Force disconnect (non-clean) - simulates network failure
+			const connRaw = connection as unknown as ActorConnRaw;
+			await forceUncleanDisconnect(
+				endpoint,
+				connRaw.actorId!,
+				connRaw.connectionId!,
+			);
+
+			// Wait a bit for the disconnection to be processed
+			await vi.waitFor(
+				async () => {
+					const countAfterReconnect = await connection.getCount();
+					expect(countAfterReconnect).toBe(5); // Should preserve the count
+				},
+				{ timeout: 5000, interval: 100 },
+			);
+
+			// Verify connection count is still 1 (same connection reconnected)
+			const connCount2 = await connection.getConnectionCount();
+			expect(connCount2).toBe(1);
+
+			// Verify we can still increment the counter
+			const newCount = await connection.getCount();
+			expect(newCount).toBe(5);
+
+			// Clean up
+			await connection.dispose();
+		});
+
+		test("should not preserve connection state after clean disconnect", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+
+			// Create actor and connect
+			const handle = client.counterConn.getOrCreate(["test-clean-disconnect"]);
+			const connection = handle.connect();
+
+			// Set an initial count on the connection
+			await connection.increment(10);
+
+			// Verify connection count is 1
+			const connCount1 = await connection.getConnectionCount();
+			expect(connCount1).toBe(1);
+
+			// Clean disconnect
+			await connection.dispose();
+
+			// Wait a bit to ensure disconnection is processed
+			await vi.waitFor(
+				async () => {
+					// Check that connection count is now 0
+					const handle2 = client.counterConn.get(["test-clean-disconnect"]);
+					const connCount = await handle2.getConnectionCount();
+					// This counts the current action caller
+					expect(connCount).toBe(1);
+				},
+				{ timeout: 5000 },
+			);
+
+			// Create a new connection
+			const connection2 = handle.connect();
+
+			// The count should be reset since it's a new connection
+			const countNewConnection = await connection2.getCount();
+			expect(countNewConnection).toBe(0); // Should be reset
+
+			// Verify connection count is 1 again (new connection)
+			const connCount3 = await connection2.getConnectionCount();
+			expect(connCount3).toBe(1);
+
+			// Clean up
+			await connection2.dispose();
+		});
+
+		test("should handle multiple non-clean disconnects and reconnects", async (c) => {
+			const { client, endpoint } = await setupDriverTest(c, driverTestConfig);
+
+			// Create actor and connect
+			const handle = client.counterConn.getOrCreate([
+				"test-multiple-reconnect",
+			]);
+			const connection = handle.connect();
+
+			// Set an initial count
+			await connection.setCount(100);
+
+			// Perform multiple disconnect-reconnect cycles
+			for (let i = 0; i < 3; i++) {
+				// Increment before disconnect
+				await connection.increment(1);
+
+				// Force disconnect
+				const connRaw = connection as unknown as ActorConnRaw;
+				await forceUncleanDisconnect(
+					endpoint,
+					connRaw.actorId!,
+					connRaw.connectionId!,
+				);
+
+				// Wait for reconnection and verify state is preserved
+				await vi.waitFor(
+					async () => {
+						const countAfter = await connection.getCount();
+						expect(countAfter).toBe(101 + i);
+					},
+					{ timeout: 5000 },
+				);
+
+				// Verify connection count remains 1
+				const connCount = await connection.getConnectionCount();
+				expect(connCount).toBe(1);
+			}
+
+			// Final verification
+			const finalCount = await connection.getCount();
+			expect(finalCount).toBe(103);
+
+			// Clean up
+			await connection.dispose();
+		});
+	});
+}
+
+async function forceUncleanDisconnect(
+	endpoint: string,
+	actorId: string,
+	connId: string,
+): Promise<void> {
+	const response = await fetch(
+		`${endpoint}/.test/force-disconnect?actor=${actorId}&conn=${connId}`,
+		{
+			method: "POST",
+		},
+	);
+
+	if (!response.ok) {
+		const text = await response.text();
+		throw new Error(`Failed to force disconnect: ${text}`);
+	}
+}

package/src/driver-test-suite/tests/actor-sleep.ts

@@ -318,7 +318,6 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 
 			// Close WebSocket
 			ws.close();
-			await new Promise((resolve) => setTimeout(resolve, 100));
 
 			// Wait for sleep timeout after WebSocket closed
 			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);

package/src/driver-test-suite/tests/raw-websocket.ts

@@ -277,41 +277,6 @@ export function runRawWebSocketTests(driverTestConfig: DriverTestConfig) {
 			ws.close();
 		});
 
-		test("should pass connection parameters through subprotocols", async (c) => {
-			const { client } = await setupDriverTest(c, driverTestConfig);
-
-			// Create actor with connection parameters
-			const testParams = { userId: "test123", role: "admin" };
-			const actor = client.rawWebSocketActor.getOrCreate(["params"], {
-				params: testParams,
-			});
-
-			const ws = await actor.websocket();
-
-			await new Promise<void>((resolve) => {
-				ws.addEventListener("open", () => resolve(), { once: true });
-			});
-
-			// Send a request to echo the auth data (which should include conn params from auth)
-			ws.send(JSON.stringify({ type: "getAuthData" }));
-
-			const response = await new Promise<any>((resolve, reject) => {
-				ws.addEventListener("message", (event: any) => {
-					const data = JSON.parse(event.data as string);
-					if (data.type === "authData") {
-						resolve(data);
-					}
-				});
-				ws.addEventListener("close", reject);
-			});
-
-			// For now, just verify we get a response
-			// The actual connection params handling needs to be implemented
-			expect(response).toBeDefined();
-
-			ws.close();
-		});
-
 		test("should handle connection close properly", async (c) => {
 			const { client } = await setupDriverTest(c, driverTestConfig);
 			const actor = client.rawWebSocketActor.getOrCreate(["close-test"]);

package/src/driver-test-suite/utils.ts

@@ -39,14 +39,14 @@ export async function setupDriverTest(
 		});
 	} else if (driverTestConfig.clientType === "inline") {
 		// Use inline client from driver
+		const transport = driverTestConfig.transport ?? "websocket";
 		const managerDriver = createTestInlineClientDriver(
 			endpoint,
 			"bare",
-			driverTestConfig.transport ?? "websocket",
+			transport,
 		);
-		invariant(driverTestConfig.transport, "missing transport");
 		const runConfig = RunConfigSchema.parse({
-			transport: driverTestConfig.transport,
+			transport: transport,
 		});
 		client = createClientWithDriver(managerDriver, runConfig);
 	} else {

package/src/drivers/default.ts

@@ -3,15 +3,12 @@ import { loggerWithoutContext } from "@/actor/log";
 import { createEngineDriver } from "@/drivers/engine/mod";
 import { createFileSystemOrMemoryDriver } from "@/drivers/file-system/mod";
 import type { DriverConfig, RunConfig } from "@/registry/run-config";
-import { getEnvUniversal } from "@/utils";
 
 /**
  * Chooses the appropriate driver based on the run configuration.
  */
 export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig {
-	const engineEndpoint = runConfig.endpoint ?? getEnvUniversal("RIVET_ENGINE");
-
-	if (engineEndpoint && runConfig.driver) {
+	if (runConfig.endpoint && runConfig.driver) {
 		throw new UserError(
 			"Cannot specify both 'engine' and 'driver' in configuration",
 		);
@@ -21,12 +18,16 @@ export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig {
 		return runConfig.driver;
 	}
 
-	if (engineEndpoint) {
+	if (runConfig.endpoint) {
 		loggerWithoutContext().debug({
 			msg: "using rivet engine driver",
-			endpoint: engineEndpoint,
+			endpoint: runConfig.endpoint,
+		});
+		// TODO: Add all properties from config
+		return createEngineDriver({
+			endpoint: runConfig.endpoint,
+			token: runConfig.token,
 		});
-		return createEngineDriver({ endpoint: engineEndpoint });
 	}
 
 	loggerWithoutContext().debug({ msg: "using default file system driver" });