rivetkit 2.0.4 → 2.0.6

package/src/drivers/engine/actor-driver.ts

@@ -6,10 +6,26 @@ import { Runner } from "@rivetkit/engine-runner";
 import * as cbor from "cbor-x";
 import { WSContext } from "hono/ws";
 import invariant from "invariant";
+import { lookupInRegistry } from "@/actor/definition";
+import {
+	createGenericConnDrivers,
+	GenericConnGlobalState,
+} from "@/actor/generic-conn-driver";
 import { deserializeActorKey } from "@/actor/keys";
 import { EncodingSchema } from "@/actor/protocol/serde";
+import { type ActorRouter, createActorRouter } from "@/actor/router";
+import {
+	handleRawWebSocketHandler,
+	handleWebSocketConnect,
+} from "@/actor/router-endpoints";
 import type { Client } from "@/client/client";
+import {
+	PATH_CONNECT_WEBSOCKET,
+	PATH_RAW_WEBSOCKET_PREFIX,
+} from "@/common/actor-router-consts";
+import type { UpgradeWebSocketArgs } from "@/common/inline-websocket-adapter2";
 import { getLogger } from "@/common/log";
+import type { UniversalWebSocket } from "@/common/websocket-interface";
 import {
 	type ActorDriver,
 	type AnyActorInstance,
@@ -19,23 +35,8 @@ import {
 	type ManagerDriver,
 	serializeEmptyPersistData,
 } from "@/driver-helpers/mod";
-import type {
-	ActorRouter,
-	RegistryConfig,
-	RunConfig,
-	UniversalWebSocket,
-	UpgradeWebSocketArgs,
-} from "@/mod";
-import {
-	createActorRouter,
-	createGenericConnDrivers,
-	GenericConnGlobalState,
-	handleRawWebSocketHandler,
-	handleWebSocketConnect,
-	lookupInRegistry,
-	PATH_CONNECT_WEBSOCKET,
-	PATH_RAW_WEBSOCKET_PREFIX,
-} from "@/mod";
+import type { RegistryConfig } from "@/registry/config";
+import type { RunConfig } from "@/registry/run-config";
 import type { Config } from "./config";
 import { KEYS } from "./kv";
 import { logger } from "./log";

package/src/drivers/file-system/global-state.ts

@@ -451,7 +451,9 @@ export class FileSystemGlobalState {
 		// Get the actor metadata
 		const entry = await this.loadActor(actorId);
 		if (!entry.state) {
-			throw new Error(`Actor does exist and cannot be started: ${actorId}`);
+			throw new Error(
+				`Actor does not exist and cannot be started: "${actorId}"`,
+			);
 		}
 
 		// Actor already starting

package/src/drivers/file-system/manager.ts

@@ -142,10 +142,10 @@ export class FileSystemManagerDriver implements ManagerDriver {
 		encoding: Encoding,
 		params: unknown,
 	): Promise<UniversalWebSocket> {
-		// TODO:
-
 		// Handle raw WebSocket paths
-		if (path === PATH_CONNECT_WEBSOCKET) {
+		const pathOnly = path.split("?")[0];
+		const normalizedPath = pathOnly.startsWith("/") ? pathOnly : `/${pathOnly}`;
+		if (normalizedPath === PATH_CONNECT_WEBSOCKET) {
 			// Handle standard connect
 			const wsHandler = await handleWebSocketConnect(
 				undefined,
@@ -158,16 +158,14 @@ export class FileSystemManagerDriver implements ManagerDriver {
 			);
 			return new InlineWebSocketAdapter2(wsHandler);
 		} else if (
-			path.startsWith(PATH_RAW_WEBSOCKET_PREFIX) ||
-			path === "/raw/websocket"
+			normalizedPath.startsWith(PATH_RAW_WEBSOCKET_PREFIX) ||
+			normalizedPath === "/raw/websocket"
 		) {
 			// Handle websocket proxy
-			// Normalize path to include trailing slash if missing
-			const normalizedPath =
-				path === "/raw/websocket" ? "/raw/websocket/" : path;
+			// Use the full path with query parameters
 			const wsHandler = await handleRawWebSocketHandler(
 				undefined,
-				normalizedPath,
+				path,
 				this.#actorDriver,
 				actorId,
 				undefined,
@@ -194,13 +192,14 @@ export class FileSystemManagerDriver implements ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		connParams: unknown,
-		authData: unknown,
 	): Promise<Response> {
 		const upgradeWebSocket = this.#runConfig.getUpgradeWebSocket?.();
 		invariant(upgradeWebSocket, "missing getUpgradeWebSocket");
 
 		// Handle raw WebSocket paths
-		if (path === PATH_CONNECT_WEBSOCKET) {
+		const pathOnly = path.split("?")[0];
+		const normalizedPath = pathOnly.startsWith("/") ? pathOnly : `/${pathOnly}`;
+		if (normalizedPath === PATH_CONNECT_WEBSOCKET) {
 			// Handle standard connect
 			const wsHandler = await handleWebSocketConnect(
 				c.req.raw,
@@ -209,26 +208,22 @@ export class FileSystemManagerDriver implements ManagerDriver {
 				actorId,
 				encoding,
 				connParams,
-				authData,
+				undefined,
 			);
-
 			return upgradeWebSocket(() => wsHandler)(c, noopNext());
 		} else if (
-			path.startsWith(PATH_RAW_WEBSOCKET_PREFIX) ||
-			path === "/raw/websocket"
+			normalizedPath.startsWith(PATH_RAW_WEBSOCKET_PREFIX) ||
+			normalizedPath === "/raw/websocket"
 		) {
 			// Handle websocket proxy
-			// Normalize path to include trailing slash if missing
-			const normalizedPath =
-				path === "/raw/websocket" ? "/raw/websocket/" : path;
+			// Use the full path with query parameters
 			const wsHandler = await handleRawWebSocketHandler(
 				c.req.raw,
-				normalizedPath,
+				path,
 				this.#actorDriver,
 				actorId,
-				authData,
+				undefined,
 			);
-
 			return upgradeWebSocket(() => wsHandler)(c, noopNext());
 		} else {
 			throw new Error(`Unreachable path: ${path}`);

package/src/manager/driver.ts

@@ -1,8 +1,8 @@
 import type { Env, Hono, Context as HonoContext } from "hono";
 import type { ActorKey, Encoding, UniversalWebSocket } from "@/actor/mod";
 import type { ManagerInspector } from "@/inspector/manager";
-import type { RunConfig } from "@/mod";
 import type { RegistryConfig } from "@/registry/config";
+import type { RunConfig } from "@/registry/run-config";
 
 export type ManagerDriverBuilder = (
 	registryConfig: RegistryConfig,

package/src/manager/protocol/query.ts

@@ -7,7 +7,7 @@ import {
 	HEADER_CONN_PARAMS,
 	HEADER_CONN_TOKEN,
 	HEADER_ENCODING,
-} from "@/actor/router-endpoints";
+} from "@/common/actor-router-consts";
 
 // Maximum size of a key component in bytes
 // Set to 128 bytes to allow for separators and escape characters in the full key

package/src/manager/router.ts

@@ -3,6 +3,8 @@ import * as cbor from "cbor-x";
 import { Hono } from "hono";
 import { cors as corsMiddleware } from "hono/cors";
 import { createMiddleware } from "hono/factory";
+import type { WSContext } from "hono/ws";
+import invariant from "invariant";
 import { z } from "zod";
 import {
 	ActorNotFound,
@@ -11,11 +13,18 @@ import {
 	Unsupported,
 	WebSocketsNotEnabled,
 } from "@/actor/errors";
+import type { Encoding } from "@/client/mod";
 import {
 	handleRouteError,
 	handleRouteNotFound,
 	loggerMiddleware,
 } from "@/common/router";
+import { deconstructError, noopNext } from "@/common/utils";
+import { HEADER_ACTOR_ID } from "@/driver-helpers/mod";
+import type {
+	TestInlineDriverCallRequest,
+	TestInlineDriverCallResponse,
+} from "@/driver-test-suite/test-inline-client-driver";
 import { createManagerInspectorRouter } from "@/inspector/manager";
 import { secureInspector } from "@/inspector/utils";
 import {
@@ -32,8 +41,10 @@ import {
 	ActorsGetOrCreateByIdResponseSchema,
 } from "@/manager-api/routes/actors-get-or-create-by-id";
 import { RivetIdSchema } from "@/manager-api/routes/common";
+import type { UniversalWebSocket, UpgradeWebSocketArgs } from "@/mod";
 import type { RegistryConfig } from "@/registry/config";
 import type { RunConfig } from "@/registry/run-config";
+import { stringifyError } from "@/utils";
 import type { ManagerDriver } from "./driver";
 import { logger } from "./log";
 
@@ -137,9 +148,10 @@ export function createManagerRouter(
 			const proxyUrl = new URL(`http://actor${url.pathname}${url.search}`);
 
 			const proxyRequest = new Request(proxyUrl, {
-				method: c.req.method,
+				method: c.req.raw.method,
 				headers: proxyHeaders,
 				body: c.req.raw.body,
+				signal: c.req.raw.signal,
 			});
 
 			return await managerDriver.proxyRequest(c, proxyRequest, actorId);
@@ -288,8 +300,8 @@ export function createManagerRouter(
 				actor_id: actorOutput.actorId,
 				name: actorOutput.name,
 				key: actorOutput.key,
-				namespace_id: "", // Not available from driver
-				runner_name_selector: "", // Not available from driver
+				namespace_id: "default", // Assert default namespace
+				runner_name_selector: "rivetkit", // Assert rivetkit runner
 				create_ts: Date.now(), // Not available from driver
 				connectable_ts: null,
 				destroy_ts: null,
@@ -350,8 +362,8 @@ export function createManagerRouter(
 				actor_id: actorOutput.actorId,
 				name: actorOutput.name,
 				key: actorOutput.key,
-				namespace_id: "", // Not available from driver
-				runner_name_selector: body.runner_name_selector,
+				namespace_id: "default", // Assert default namespace
+				runner_name_selector: "rivetkit", // Assert rivetkit runner
 				create_ts: Date.now(),
 				connectable_ts: null,
 				destroy_ts: null,
@@ -387,6 +399,139 @@ export function createManagerRouter(
 	// 	});
 	// }
 
+	if (registryConfig.test.enabled) {
+		// Add HTTP endpoint to test the inline client
+		//
+		// We have to do this in a router since this needs to run in the same server as the RivetKit registry. Some test contexts to not run in the same server.
+		router.post(".test/inline-driver/call", async (c) => {
+			// TODO: use openapi instead
+			const buffer = await c.req.arrayBuffer();
+			const { encoding, transport, method, args }: TestInlineDriverCallRequest =
+				cbor.decode(new Uint8Array(buffer));
+
+			logger().debug({
+				msg: "received inline request",
+				encoding,
+				transport,
+				method,
+				args,
+			});
+
+			// Forward inline driver request
+			let response: TestInlineDriverCallResponse<unknown>;
+			try {
+				const output = await ((managerDriver as any)[method] as any)(...args);
+				response = { ok: output };
+			} catch (rawErr) {
+				const err = deconstructError(rawErr, logger(), {}, true);
+				response = { err };
+			}
+
+			return c.body(cbor.encode(response));
+		});
+
+		router.get(".test/inline-driver/connect-websocket/*", async (c) => {
+			const upgradeWebSocket = runConfig.getUpgradeWebSocket?.();
+			invariant(upgradeWebSocket, "websockets not supported on this platform");
+
+			return upgradeWebSocket(async (c: any) => {
+				const {
+					path,
+					actorId,
+					params: paramsRaw,
+					encodingKind,
+				} = c.req.query() as {
+					path: string;
+					actorId: string;
+					params?: string;
+					encodingKind: Encoding;
+				};
+				const params =
+					paramsRaw !== undefined ? JSON.parse(paramsRaw) : undefined;
+
+				logger().debug({
+					msg: "received test inline driver websocket",
+					actorId,
+					params,
+					encodingKind,
+					path: path,
+				});
+
+				// Connect to the actor using the inline client driver - this returns a Promise<WebSocket>
+				const clientWsPromise = managerDriver.openWebSocket(
+					path,
+					actorId,
+					encodingKind,
+					params,
+				);
+
+				return await createTestWebSocketProxy(clientWsPromise, "standard");
+			})(c, noopNext());
+		});
+
+		router.all(".test/inline-driver/send-request/*", async (c) => {
+			// Extract parameters from headers
+			const actorId = c.req.header(HEADER_ACTOR_ID);
+
+			if (!actorId) {
+				return c.text("Missing required headers", 400);
+			}
+
+			// Extract the path after /send-request/
+			const pathOnly =
+				c.req.path.split("/.test/inline-driver/send-request/")[1] || "";
+
+			// Include query string
+			const url = new URL(c.req.url);
+			const pathWithQuery = pathOnly + url.search;
+
+			logger().debug({
+				msg: "received test inline driver raw http",
+				actorId,
+				path: pathWithQuery,
+				method: c.req.method,
+			});
+
+			try {
+				// Forward the request using the inline client driver
+				const response = await managerDriver.sendRequest(
+					actorId,
+					new Request(`http://actor/${pathWithQuery}`, {
+						method: c.req.method,
+						headers: c.req.raw.headers,
+						body: c.req.raw.body,
+					}),
+				);
+
+				// Return the response directly
+				return response;
+			} catch (error) {
+				logger().error({
+					msg: "error in test inline raw http",
+					error: stringifyError(error),
+				});
+
+				// Return error response
+				const err = deconstructError(error, logger(), {}, true);
+				return c.json(
+					{
+						error: {
+							code: err.code,
+							message: err.message,
+							metadata: err.metadata,
+						},
+					},
+					err.statusCode,
+				);
+			}
+		});
+	}
+
+	managerDriver.modifyManagerRouter?.(
+		registryConfig,
+		router as unknown as Hono,
+	);
+
 	if (runConfig.inspector?.enabled) {
 		if (!managerDriver.inspector) {
 			throw new Unsupported("inspector");
@@ -410,3 +555,226 @@ export function createManagerRouter(
 
 	return { router: router as Hono, openapi: router };
 }
+/**
+ * Creates a WebSocket proxy for test endpoints that forwards messages between server and client WebSockets
+ */
+async function createTestWebSocketProxy(
+	clientWsPromise: Promise<UniversalWebSocket>,
+	connectionType: string,
+): Promise<UpgradeWebSocketArgs> {
+	// Store a reference to the resolved WebSocket
+	let clientWs: UniversalWebSocket | null = null;
+	try {
+		// Resolve the client WebSocket promise
+		logger().debug({ msg: "awaiting client websocket promise" });
+		const ws = await clientWsPromise;
+		clientWs = ws;
+		logger().debug({
+			msg: "client websocket promise resolved",
+			constructor: ws?.constructor.name,
+		});
+
+		// Wait for ws to open
+		await new Promise<void>((resolve, reject) => {
+			const onOpen = () => {
+				logger().debug({ msg: "test websocket connection opened" });
+				resolve();
+			};
+			const onError = (error: any) => {
+				logger().error({ msg: "test websocket connection failed", error });
+				reject(
+					new Error(`Failed to open WebSocket: ${error.message || error}`),
+				);
+			};
+			ws.addEventListener("open", onOpen);
+			ws.addEventListener("error", onError);
+		});
+	} catch (error) {
+		logger().error({
+			msg: `failed to establish client ${connectionType} websocket connection`,
+			error,
+		});
+		return {
+			onOpen: (_evt, serverWs) => {
+				serverWs.close(1011, "Failed to establish connection");
+			},
+			onMessage: () => {},
+			onError: () => {},
+			onClose: () => {},
+		};
+	}
+
+	// Create WebSocket proxy handlers to relay messages between client and server
+	return {
+		onOpen: (_evt: any, serverWs: WSContext) => {
+			logger().debug({
+				msg: `test ${connectionType} websocket connection opened`,
+			});
+
+			// Check WebSocket type
+			logger().debug({
+				msg: "clientWs info",
+				constructor: clientWs.constructor.name,
+				hasAddEventListener: typeof clientWs.addEventListener === "function",
+				readyState: clientWs.readyState,
+			});
+
+			// Add message handler to forward messages from client to server
+			clientWs.addEventListener("message", (clientEvt: MessageEvent) => {
+				logger().debug({
+					msg: `test ${connectionType} websocket connection message from client`,
+					dataType: typeof clientEvt.data,
+					isBlob: clientEvt.data instanceof Blob,
+					isArrayBuffer: clientEvt.data instanceof ArrayBuffer,
+					dataConstructor: clientEvt.data?.constructor?.name,
+					dataStr:
+						typeof clientEvt.data === "string"
+							? clientEvt.data.substring(0, 100)
+							: undefined,
+				});
+
+				if (serverWs.readyState === 1) {
+					// OPEN
+					// Handle Blob data
+					if (clientEvt.data instanceof Blob) {
+						clientEvt.data
+							.arrayBuffer()
+							.then((buffer) => {
+								logger().debug({
+									msg: "converted client blob to arraybuffer, sending to server",
+									bufferSize: buffer.byteLength,
+								});
+								serverWs.send(buffer as any);
+							})
+							.catch((error) => {
+								logger().error({
+									msg: "failed to convert blob to arraybuffer",
+									error,
+								});
+							});
+					} else {
+						logger().debug({
+							msg: "sending client data directly to server",
+							dataType: typeof clientEvt.data,
+							dataLength:
+								typeof clientEvt.data === "string"
+									? clientEvt.data.length
+									: undefined,
+						});
+						serverWs.send(clientEvt.data as any);
+					}
+				}
+			});
+
+			// Add close handler to close server when client closes
+			clientWs.addEventListener("close", (clientEvt: any) => {
+				logger().debug({
+					msg: `test ${connectionType} websocket connection closed`,
+				});
+
+				if (serverWs.readyState !== 3) {
+					// Not CLOSED
+					serverWs.close(clientEvt.code, clientEvt.reason);
+				}
+			});
+
+			// Add error handler
+			clientWs.addEventListener("error", () => {
+				logger().debug({
+					msg: `test ${connectionType} websocket connection error`,
+				});
+
+				if (serverWs.readyState !== 3) {
+					// Not CLOSED
+					serverWs.close(1011, "Error in client websocket");
+				}
+			});
+		},
+		onMessage: (evt: { data: any }) => {
+			logger().debug({
+				msg: "received message from server",
+				dataType: typeof evt.data,
+				isBlob: evt.data instanceof Blob,
+				isArrayBuffer: evt.data instanceof ArrayBuffer,
+				dataConstructor: evt.data?.constructor?.name,
+				dataStr:
+					typeof evt.data === "string" ? evt.data.substring(0, 100) : undefined,
+			});
+
+			// Forward messages from server websocket to client websocket
+			if (clientWs.readyState === 1) {
+				// OPEN
+				// Handle Blob data
+				if (evt.data instanceof Blob) {
+					evt.data
+						.arrayBuffer()
+						.then((buffer) => {
+							logger().debug({
+								msg: "converted blob to arraybuffer, sending",
+								bufferSize: buffer.byteLength,
+							});
+							clientWs.send(buffer);
+						})
+						.catch((error) => {
+							logger().error({
+								msg: "failed to convert blob to arraybuffer",
+								error,
+							});
+						});
+				} else {
+					logger().debug({
+						msg: "sending data directly",
+						dataType: typeof evt.data,
+						dataLength:
+							typeof evt.data === "string" ? evt.data.length : undefined,
+					});
+					clientWs.send(evt.data);
+				}
+			}
+		},
+		onClose: (
+			event: {
+				wasClean: boolean;
+				code: number;
+				reason: string;
+			},
+			serverWs: WSContext,
+		) => {
+			logger().debug({
+				msg: `server ${connectionType} websocket closed`,
+				wasClean: event.wasClean,
+				code: event.code,
+				reason: event.reason,
+			});
+
+			// HACK: Close socket in order to fix bug with Cloudflare leaving WS in closing state
+			// https://github.com/cloudflare/workerd/issues/2569
+			serverWs.close(1000, "hack_force_close");
+
+			// Close the client websocket when the server websocket closes
+			if (
+				clientWs &&
+				clientWs.readyState !== clientWs.CLOSED &&
+				clientWs.readyState !== clientWs.CLOSING
+			) {
+				// Don't pass code/message since this may affect how close events are triggered
+				clientWs.close(1000, event.reason);
+			}
+		},
+		onError: (error: unknown) => {
+			logger().error({
+				msg: `error in server ${connectionType} websocket`,
+				error,
+			});
+
+			// Close the client websocket on error
+			if (
+				clientWs &&
+				clientWs.readyState !== clientWs.CLOSED &&
+				clientWs.readyState !== clientWs.CLOSING
+			) {
+				clientWs.close(1011, "Error in server websocket");
+			}
+		},
+	};
+}