rivetkit 2.0.5 → 2.0.7-rc.1

package/src/common/actor-router-consts.ts

@@ -0,0 +1,59 @@
+// NOTE: This is in a separate file from the router since it needs to be shared between the client & the server. If this was in the router file, the client would end up importing the *entire* actor router and tree shaking would not work.
+
+// MARK: Paths
+export const PATH_CONNECT_WEBSOCKET = "/connect/websocket";
+export const PATH_RAW_WEBSOCKET_PREFIX = "/raw/websocket/";
+
+// MARK: Headers
+export const HEADER_ACTOR_QUERY = "x-rivet-query";
+
+export const HEADER_ENCODING = "x-rivet-encoding";
+
+// IMPORTANT: Params must be in headers or in an E2EE part of the request (i.e. NOT the URL or query string) in order to ensure that tokens can be securely passed in params.
+export const HEADER_CONN_PARAMS = "x-rivet-conn-params";
+
+export const HEADER_ACTOR_ID = "x-rivet-actor";
+
+export const HEADER_CONN_ID = "x-rivet-conn";
+
+export const HEADER_CONN_TOKEN = "x-rivet-conn-token";
+
+export const HEADER_RIVET_TOKEN = "x-rivet-token";
+
+// MARK: Manager Gateway Headers
+export const HEADER_RIVET_TARGET = "x-rivet-target";
+export const HEADER_RIVET_ACTOR = "x-rivet-actor";
+
+// MARK: WebSocket Protocol Prefixes
+/** Some servers (such as node-ws & Cloudflare) require explicitly match a certain WebSocket protocol. This gives us a static protocol to match against. */
+export const WS_PROTOCOL_STANDARD = "rivet";
+export const WS_PROTOCOL_TARGET = "rivet_target.";
+export const WS_PROTOCOL_ACTOR = "rivet_actor.";
+export const WS_PROTOCOL_ENCODING = "rivet_encoding.";
+export const WS_PROTOCOL_CONN_PARAMS = "rivet_conn_params.";
+export const WS_PROTOCOL_CONN_ID = "rivet_conn.";
+export const WS_PROTOCOL_CONN_TOKEN = "rivet_conn_token.";
+export const WS_PROTOCOL_TOKEN = "rivet_token.";
+
+// MARK: WebSocket Inline Test Protocol Prefixes
+export const WS_PROTOCOL_TRANSPORT = "test_transport.";
+export const WS_PROTOCOL_PATH = "test_path.";
+
+/**
+ * Headers that publics can send from public clients.
+ *
+ * Used for CORS.
+ **/
+export const ALLOWED_PUBLIC_HEADERS = [
+	"Content-Type",
+	"User-Agent",
+	HEADER_ACTOR_QUERY,
+	HEADER_ENCODING,
+	HEADER_CONN_PARAMS,
+	HEADER_ACTOR_ID,
+	HEADER_CONN_ID,
+	HEADER_CONN_TOKEN,
+	HEADER_RIVET_TARGET,
+	HEADER_RIVET_ACTOR,
+	HEADER_RIVET_TOKEN,
+];

package/src/common/router.ts

@@ -76,5 +76,6 @@ export function handleRouteError(error: unknown, c: HonoContext) {
 		HTTP_RESPONSE_ERROR_VERSIONED,
 	);
 
-	return c.body(output, { status: statusCode });
+	// TODO: Remove any
+	return c.body(output as any, { status: statusCode });
 }

package/src/common/versioned-data.ts

@@ -63,8 +63,8 @@ export class VersionedDataHandler<T> {
 	}
 
 	private embedVersion(data: VersionedData<Uint8Array>): Uint8Array {
-		const versionBytes = new Uint8Array(4);
-		new DataView(versionBytes.buffer).setUint32(0, data.version, true);
+		const versionBytes = new Uint8Array(2);
+		new DataView(versionBytes.buffer).setUint16(0, data.version, true);
 
 		const result = new Uint8Array(versionBytes.length + data.data.length);
 		result.set(versionBytes);
@@ -74,15 +74,15 @@ export class VersionedDataHandler<T> {
 	}
 
 	private extractVersion(bytes: Uint8Array): VersionedData<Uint8Array> {
-		if (bytes.length < 4) {
+		if (bytes.length < 2) {
 			throw new Error("Invalid versioned data: too short");
 		}
 
-		const version = new DataView(bytes.buffer, bytes.byteOffset).getUint32(
+		const version = new DataView(bytes.buffer, bytes.byteOffset).getUint16(
 			0,
 			true,
 		);
-		const data = bytes.slice(4);
+		const data = bytes.slice(2);
 
 		return { version, data };
 	}

package/src/driver-helpers/mod.ts

@@ -1,14 +1,27 @@
 export type { ActorDriver } from "@/actor/driver";
 export type { ActorInstance, AnyActorInstance } from "@/actor/instance";
 export {
+	ALLOWED_PUBLIC_HEADERS,
 	HEADER_ACTOR_ID,
 	HEADER_ACTOR_QUERY,
-	HEADER_AUTH_DATA,
 	HEADER_CONN_ID,
 	HEADER_CONN_PARAMS,
 	HEADER_CONN_TOKEN,
 	HEADER_ENCODING,
-} from "@/actor/router-endpoints";
+	HEADER_RIVET_ACTOR,
+	HEADER_RIVET_TARGET,
+	PATH_CONNECT_WEBSOCKET,
+	PATH_RAW_WEBSOCKET_PREFIX,
+	WS_PROTOCOL_ACTOR,
+	WS_PROTOCOL_CONN_ID,
+	WS_PROTOCOL_CONN_PARAMS,
+	WS_PROTOCOL_CONN_TOKEN,
+	WS_PROTOCOL_ENCODING,
+	WS_PROTOCOL_PATH,
+	WS_PROTOCOL_STANDARD,
+	WS_PROTOCOL_TARGET,
+	WS_PROTOCOL_TRANSPORT,
+} from "@/common/actor-router-consts";
 export type {
 	ActorOutput,
 	CreateInput,

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

@@ -4,6 +4,7 @@ import { bundleRequire } from "bundle-require";
 import invariant from "invariant";
 import { describe } from "vitest";
 import type { Transport } from "@/client/mod";
+import { configureInspectorAccessToken } from "@/inspector/utils";
 import { createManagerRouter } from "@/manager/router";
 import type { DriverConfig, Registry, RunConfig } from "@/mod";
 import { RunConfigSchema } from "@/registry/run-config";
@@ -22,6 +23,7 @@ import { runActorInlineClientTests } from "./tests/actor-inline-client";
 import { runActorInspectorTests } from "./tests/actor-inspector";
 import { runActorMetadataTests } from "./tests/actor-metadata";
 import { runActorOnStateChangeTests } from "./tests/actor-onstatechange";
+import { runActorReconnectTests } from "./tests/actor-reconnect";
 import { runActorVarsTests } from "./tests/actor-vars";
 import { runManagerDriverTests } from "./tests/manager-driver";
 import { runRawHttpTests } from "./tests/raw-http";
@@ -32,6 +34,7 @@ import { runRequestAccessTests } from "./tests/request-access";
 export interface SkipTests {
 	schedule?: boolean;
 	sleep?: boolean;
+	sse?: boolean;
 }
 
 export interface DriverTestConfig {
@@ -86,7 +89,10 @@ export function runDriverTests(
 			runActorDriverTests(driverTestConfig);
 			runManagerDriverTests(driverTestConfig);
 
-			for (const transport of ["websocket", "sse"] as Transport[]) {
+			const transports: Transport[] = driverTestConfig.skip?.sse
+				? ["websocket"]
+				: ["websocket", "sse"];
+			for (const transport of transports) {
 				describe(`transport (${transport})`, () => {
 					runActorConnTests({
 						...driverTestConfig,
@@ -95,6 +101,8 @@ export function runDriverTests(
 
 					runActorConnStateTests({ ...driverTestConfig, transport });
 
+					runActorReconnectTests({ ...driverTestConfig, transport });
+
 					runRequestAccessTests({ ...driverTestConfig, transport });
 
 					runActorDriverTestsWithTransport({ ...driverTestConfig, transport });
@@ -193,11 +201,12 @@ export async function createTestRuntime(
 
 		// Create router
 		const managerDriver = driver.manager(registry.config, config);
+		configureInspectorAccessToken(config, managerDriver);
 		const { router } = createManagerRouter(
 			registry.config,
 			config,
 			managerDriver,
-			false,
+			undefined,
 		);
 
 		// Inject WebSocket

package/src/driver-test-suite/test-inline-client-driver.ts

@@ -3,14 +3,20 @@ import type { Context as HonoContext } from "hono";
 import invariant from "invariant";
 import type { WebSocket } from "ws";
 import type { Encoding } from "@/actor/protocol/serde";
+import { assertUnreachable } from "@/actor/utils";
+import { ActorError as ClientActorError } from "@/client/errors";
+import type { Transport } from "@/client/mod";
 import {
 	HEADER_ACTOR_QUERY,
 	HEADER_CONN_PARAMS,
 	HEADER_ENCODING,
-} from "@/actor/router-endpoints";
-import { assertUnreachable } from "@/actor/utils";
-import { ActorError as ClientActorError } from "@/client/errors";
-import type { Transport } from "@/client/mod";
+	WS_PROTOCOL_ACTOR,
+	WS_PROTOCOL_CONN_PARAMS,
+	WS_PROTOCOL_ENCODING,
+	WS_PROTOCOL_PATH,
+	WS_PROTOCOL_TARGET,
+	WS_PROTOCOL_TRANSPORT,
+} from "@/common/actor-router-consts";
 import type { UniversalEventSource } from "@/common/eventsource-interface";
 import type { DeconstructedError } from "@/common/utils";
 import { importWebSocket } from "@/common/websocket";
@@ -111,7 +117,8 @@ export function createTestInlineClientDriver(
 					headers,
 					body: actorRequest.body,
 					signal: actorRequest.signal,
-				}),
+					duplex: "half",
+				} as RequestInit),
 			);
 
 			// Check if it's an error response from our handler
@@ -153,39 +160,48 @@ export function createTestInlineClientDriver(
 			actorId: string,
 			encoding: Encoding,
 			params: unknown,
+			connId?: string,
+			connToken?: string,
 		): Promise<UniversalWebSocket> {
 			const WebSocket = await importWebSocket();
 
 			// Normalize path to match other drivers
 			const normalizedPath = path.startsWith("/") ? path.slice(1) : path;
 
-			logger().debug({
-				msg: "creating websocket connection via test inline driver",
-			});
-
 			// Create WebSocket connection to the test endpoint
-			// Use a placeholder path and pass the actual path as a query param to avoid mixing user query params with internal ones
 			const wsUrl = new URL(
 				`${endpoint}/.test/inline-driver/connect-websocket/ws`,
 			);
-			wsUrl.searchParams.set("path", normalizedPath);
-			wsUrl.searchParams.set("actorId", actorId);
-			if (params !== undefined)
-				wsUrl.searchParams.set("params", JSON.stringify(params));
-			wsUrl.searchParams.set("encodingKind", encoding);
+
+			logger().debug({
+				msg: "creating websocket connection via test inline driver",
+				url: wsUrl.toString(),
+			});
 
 			// Convert http/https to ws/wss
 			const wsProtocol = wsUrl.protocol === "https:" ? "wss:" : "ws:";
-			const finalWsUrl = `${wsProtocol}//${wsUrl.host}${wsUrl.pathname}${wsUrl.search}`;
+			const finalWsUrl = `${wsProtocol}//${wsUrl.host}${wsUrl.pathname}`;
 
 			logger().debug({ msg: "connecting to websocket", url: finalWsUrl });
 
+			// Build protocols for the connection
+			const protocols: string[] = [];
+			protocols.push(`${WS_PROTOCOL_TARGET}actor`);
+			protocols.push(`${WS_PROTOCOL_ACTOR}${actorId}`);
+			protocols.push(`${WS_PROTOCOL_ENCODING}${encoding}`);
+			protocols.push(`${WS_PROTOCOL_TRANSPORT}${transport}`);
+			protocols.push(
+				`${WS_PROTOCOL_PATH}${encodeURIComponent(normalizedPath)}`,
+			);
+			if (params !== undefined) {
+				protocols.push(
+					`${WS_PROTOCOL_CONN_PARAMS}${encodeURIComponent(JSON.stringify(params))}`,
+				);
+			}
+
 			// Create and return the WebSocket
 			// Node & browser WebSocket types are incompatible
-			const ws = new WebSocket(finalWsUrl, [
-				// HACK: See packages/drivers/cloudflare-workers/src/websocket.ts
-				"rivetkit",
-			]) as any;
+			const ws = new WebSocket(finalWsUrl, protocols) as any;
 
 			return ws;
 		},
@@ -202,7 +218,6 @@ export function createTestInlineClientDriver(
 			_actorId: string,
 			_encoding: Encoding,
 			_params: unknown,
-			_authData: unknown,
 		): Promise<Response> {
 			throw "UNIMPLEMENTED";
 			// const upgradeWebSocket = this.#runConfig.getUpgradeWebSocket?.();
@@ -214,6 +229,8 @@ export function createTestInlineClientDriver(
 		displayInformation(): ManagerDisplayInformation {
 			return { name: "Test Inline", properties: {} };
 		},
+		// TODO:
+		getOrCreateInspectorAccessToken: () => "",
 
 		// action: async <Args extends Array<unknown> = unknown[], Response = unknown>(
 		// 	_c: HonoContext | undefined,
@@ -560,7 +577,8 @@ async function makeInlineRequest<T>(
 			method,
 			args,
 		} satisfies TestInlineDriverCallRequest),
-	});
+		duplex: "half",
+	} as RequestInit);
 
 	if (!response.ok) {
 		throw new Error(`Failed to call inline ${method}: ${response.statusText}`);

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();