rivetkit 2.0.4 → 2.0.5

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/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");
+			}
+		},
+	};
+}

package/src/registry/mod.ts

@@ -21,12 +21,10 @@ import {
 import { crossPlatformServe } from "./serve";
 
 interface ServerOutput<A extends Registry<any>> {
-	config: RunConfig;
-	driver: DriverConfig;
+	/** Client to communicate with the actors. */
 	client: Client<A>;
-	hono: Hono;
-	handler: (req: Request) => Promise<Response>;
-	serve: (hono?: Hono) => void;
+	/** Fetch handler to manually route requests to the Rivet manager API. */
+	fetch: (request: Request, ...args: any) => Response | Promise<Response>;
 }
 
 export class Registry<A extends RegistryActors> {
@@ -43,7 +41,7 @@ export class Registry<A extends RegistryActors> {
 	/**
 	 * Runs the registry for a server.
 	 */
-	public createServer(inputConfig?: RunConfigInput): ServerOutput<this> {
+	public start(inputConfig?: RunConfigInput): ServerOutput<this> {
 		const config = RunConfigSchema.parse(inputConfig);
 
 		// Configure logger
@@ -62,6 +60,13 @@ export class Registry<A extends RegistryActors> {
 		// TODO: Find cleaner way of disabling by default
 		if (driver.name === "engine") {
 			config.inspector.enabled = false;
+			config.disableServer = true;
+		}
+		if (driver.name === "cloudflare-workers") {
+			config.inspector.enabled = false;
+			config.disableServer = true;
+			config.disableActorDriver = true;
+			config.noWelcome = true;
 		}
 
 		// Configure getUpgradeWebSocket lazily so we can assign it in crossPlatformServe
@@ -109,41 +114,32 @@ export class Registry<A extends RegistryActors> {
 			console.log();
 		}
 
+		// HACK: We need to find a better way to let the driver itself decide when to start the actor driver
 		// Create runner
 		//
 		// Even though we do not use the return value, this is required to start the code that will handle incoming actors
-		const _actorDriver = driver.actor(
-			this.#config,
-			config,
-			managerDriver,
-			client,
-		);
+		if (!config.disableActorDriver) {
+			const _actorDriver = driver.actor(
+				this.#config,
+				config,
+				managerDriver,
+				client,
+			);
+		}
+
+		// Start server
+		if (!config.disableServer) {
+			(async () => {
+				const out = await crossPlatformServe(hono, undefined);
+				upgradeWebSocket = out.upgradeWebSocket;
+			})();
+		}
 
 		return {
-			config,
-			driver,
 			client,
-			hono,
-			handler: async (req: Request) => await hono.fetch(req),
-			serve: async (app) => {
-				const out = await crossPlatformServe(hono, app);
-				upgradeWebSocket = out.upgradeWebSocket;
-			},
+			fetch: hono.fetch.bind(hono),
 		};
 	}
-
-	/**
-	 * Runs the registry as a standalone server.
-	 */
-	public async runServer(inputConfig?: RunConfigInput) {
-		const { driver, serve } = this.createServer(inputConfig);
-
-		// TODO: FInd better way of doing this
-		// Don't run server by default
-		if (driver.name !== "engine") {
-			serve();
-		}
-	}
 }
 
 export function setup<A extends RegistryActors>(

package/src/registry/run-config.ts

@@ -29,16 +29,31 @@ export const RunConfigSchema = ClientConfigSchema.extend({
 
 	inspector: InspectorConfigSchema,
 
+	/** @experimental */
+	disableServer: z.boolean().optional().default(false),
+
+	/** @experimental */
+	disableActorDriver: z.boolean().optional().default(false),
+
 	/**
+	 * @experimental
+	 *
 	 * Base path for the router. This is used to prefix all routes.
 	 * For example, if the base path is `/api`, then the route `/actors` will be
 	 * available at `/api/actors`.
 	 */
 	basePath: z.string().optional().default("/"),
 
-	/** Disable welcome message. */
+	/**
+	 * @experimental
+	 *
+	 * Disable welcome message.
+	 * */
 	noWelcome: z.boolean().optional().default(false),
 
+	/**
+	 * @experimental
+	 * */
 	logging: z
 		.object({
 			baseLogger: z.custom<Logger>().optional(),

package/src/remote-manager-driver/actor-http-client.ts

@@ -1,4 +1,5 @@
 import type { ClientConfig } from "@/client/config";
+import { combineUrlPath } from "@/utils";
 import { getEndpoint } from "./api-utils";
 
 export async function sendHttpRequestToActor(
@@ -9,7 +10,7 @@ export async function sendHttpRequestToActor(
 	// Route through guard port
 	const url = new URL(actorRequest.url);
 	const endpoint = getEndpoint(runConfig);
-	const guardUrl = `${endpoint}${url.pathname}${url.search}`;
+	const guardUrl = combineUrlPath(endpoint, url.pathname + url.search);
 
 	// Handle body properly based on method and presence
 	let bodyToSend: ArrayBuffer | null = null;
@@ -43,6 +44,7 @@ export async function sendHttpRequestToActor(
 		method: actorRequest.method,
 		headers: guardHeaders,
 		body: bodyToSend,
+		signal: actorRequest.signal,
 	});
 
 	return mutableResponse(await fetch(guardRequest));

package/src/remote-manager-driver/actor-websocket-client.ts

@@ -6,6 +6,7 @@ import {
 import type { ClientConfig } from "@/client/config";
 import { importWebSocket } from "@/common/websocket";
 import type { Encoding, UniversalWebSocket } from "@/mod";
+import { combineUrlPath } from "@/utils";
 import { getEndpoint } from "./api-utils";
 import { logger } from "./log";
 
@@ -20,7 +21,7 @@ export async function openWebSocketToActor(
 
 	// WebSocket connections go through guard
 	const endpoint = getEndpoint(runConfig);
-	const guardUrl = `${endpoint}${path}`;
+	const guardUrl = combineUrlPath(endpoint, path);
 
 	logger().debug({
 		msg: "opening websocket to actor via guard",