rivetkit 2.0.5 → 2.0.7-rc.1

package/src/remote-manager-driver/mod.ts

@@ -1,8 +1,8 @@
 import * as cbor from "cbor-x";
-import type { Hono, Context as HonoContext } from "hono";
+import type { Context as HonoContext } from "hono";
 import invariant from "invariant";
-import { ActorAlreadyExists } from "@/actor/errors";
 import { deserializeActorKey, serializeActorKey } from "@/actor/keys";
+import { generateRandomString } from "@/actor/utils";
 import type { ClientConfig } from "@/client/client";
 import { noopNext } from "@/common/utils";
 import type {
@@ -14,21 +14,19 @@ import type {
 	ManagerDisplayInformation,
 	ManagerDriver,
 } from "@/driver-helpers/mod";
-import type { ManagerInspector } from "@/inspector/manager";
-import type { Encoding, RegistryConfig, UniversalWebSocket } from "@/mod";
-import type { RunConfig } from "@/registry/run-config";
+import type { Encoding, UniversalWebSocket } from "@/mod";
 import { combineUrlPath } from "@/utils";
 import { sendHttpRequestToActor } from "./actor-http-client";
 import {
-	buildGuardHeadersForWebSocket,
+	buildWebSocketProtocols,
 	openWebSocketToActor,
 } from "./actor-websocket-client";
 import {
 	createActor,
 	destroyActor,
 	getActor,
-	getActorById,
-	getOrCreateActorById,
+	getActorByKey,
+	getOrCreateActor,
 } from "./api-endpoints";
 import { EngineApiError, getEndpoint } from "./api-utils";
 import { logger } from "./log";
@@ -61,39 +59,30 @@ export class RemoteManagerDriver implements ManagerDriver {
 		actorId,
 	}: GetForIdInput): Promise<ActorOutput | undefined> {
 		// Fetch from API if not in cache
-		try {
-			const response = await getActor(this.#config, actorId);
-
-			// Validate name matches
-			if (response.actor.name !== name) {
-				logger().debug({
-					msg: "actor name mismatch from api",
-					actorId,
-					apiName: response.actor.name,
-					requestedName: name,
-				});
-				return undefined;
-			}
-
-			const keyRaw = response.actor.key;
-			invariant(keyRaw, `actor ${actorId} should have key`);
-			const key = deserializeActorKey(keyRaw);
+		const response = await getActor(this.#config, name, actorId);
+		const actor = response.actors[0];
+		if (!actor) return undefined;
 
-			return {
+		// Validate name matches
+		if (actor.name !== name) {
+			logger().debug({
+				msg: "actor name mismatch from api",
 				actorId,
-				name,
-				key,
-			};
-		} catch (error) {
-			if (
-				error instanceof EngineApiError &&
-				(error as EngineApiError).group === "actor" &&
-				(error as EngineApiError).code === "not_found"
-			) {
-				return undefined;
-			}
-			throw error;
+				apiName: actor.name,
+				requestedName: name,
+			});
+			return undefined;
 		}
+
+		const keyRaw = actor.key;
+		invariant(keyRaw, `actor ${actorId} should have key`);
+		const key = deserializeActorKey(keyRaw);
+
+		return {
+			actorId,
+			name,
+			key,
+		};
 	}
 
 	async getWithKey({
@@ -105,13 +94,11 @@ export class RemoteManagerDriver implements ManagerDriver {
 
 		// If not in local cache, fetch by key from API
 		try {
-			const response = await getActorById(this.#config, name, key);
-
-			if (!response.actor_id) {
-				return undefined;
-			}
+			const response = await getActorByKey(this.#config, name, key);
+			const actor = response.actors[0];
+			if (!actor) return undefined;
 
-			const actorId = response.actor_id;
+			const actorId = actor.actor_id;
 
 			logger().debug({
 				msg: "getWithKey: found actor via api",
@@ -148,22 +135,24 @@ export class RemoteManagerDriver implements ManagerDriver {
 			key,
 		});
 
-		const response = await getOrCreateActorById(this.#config, {
+		const { actor, created } = await getOrCreateActor(this.#config, {
 			name,
 			key: serializeActorKey(key),
 			runner_name_selector: this.#config.runnerName,
-			input: input ? cbor.encode(actorInput).toString("base64") : undefined,
+			input: actorInput
+				? cbor.encode(actorInput).toString("base64")
+				: undefined,
 			crash_policy: "sleep",
 		});
 
-		const actorId = response.actor_id;
+		const actorId = actor.actor_id;
 
 		logger().info({
 			msg: "getOrCreateWithKey: actor ready",
 			actorId,
 			name,
 			key,
-			created: response.created,
+			created,
 		});
 
 		return {
@@ -217,6 +206,8 @@ export class RemoteManagerDriver implements ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		params: unknown,
+		connId?: string,
+		connToken?: string,
 	): Promise<UniversalWebSocket> {
 		return await openWebSocketToActor(
 			this.#config,
@@ -224,6 +215,8 @@ export class RemoteManagerDriver implements ManagerDriver {
 			actorId,
 			encoding,
 			params,
+			connId,
+			connToken,
 		);
 	}
 
@@ -241,7 +234,8 @@ export class RemoteManagerDriver implements ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		params: unknown,
-		authData: unknown,
+		connId?: string,
+		connToken?: string,
 	): Promise<Response> {
 		const upgradeWebSocket = this.#config.getUpgradeWebSocket?.();
 		invariant(upgradeWebSocket, "missing getUpgradeWebSocket");
@@ -257,14 +251,16 @@ export class RemoteManagerDriver implements ManagerDriver {
 			guardUrl,
 		});
 
-		// Build headers
-		const headers = buildGuardHeadersForWebSocket(
+		// Build protocols
+		const protocols = buildWebSocketProtocols(
+			this.#config,
 			actorId,
 			encoding,
 			params,
-			authData,
+			connId,
+			connToken,
 		);
-		const args = await createWebSocketProxy(c, wsGuardUrl, headers);
+		const args = await createWebSocketProxy(c, wsGuardUrl, protocols);
 
 		return await upgradeWebSocket(() => args)(c, noopNext());
 	}
@@ -272,4 +268,8 @@ export class RemoteManagerDriver implements ManagerDriver {
 	displayInformation(): ManagerDisplayInformation {
 		return { name: "Remote", properties: {} };
 	}
+
+	getOrCreateInspectorAccessToken() {
+		return generateRandomString();
+	}
 }

package/src/remote-manager-driver/ws-proxy.ts

@@ -11,17 +11,10 @@ import { logger } from "./log";
 export async function createWebSocketProxy(
 	c: HonoContext,
 	targetUrl: string,
-	headers: Record<string, string>,
+	protocols: string[],
 ): Promise<UpgradeWebSocketArgs> {
 	const WebSocket = await importWebSocket();
 
-	// HACK: Sanitize WebSocket-specific headers. If we don't do this, some WebSocket implementations (i.e. native WebSocket in Node.js) will fail to connect.
-	for (const [k, v] of c.req.raw.headers.entries()) {
-		if (!k.startsWith("sec-") && k !== "connection" && k !== "upgrade") {
-			headers[k] = v;
-		}
-	}
-
 	// WebSocket state
 	interface WsState {
 		targetWs?: WebSocket;
@@ -43,7 +36,7 @@ export async function createWebSocketProxy(
 			}
 
 			// Create WebSocket
-			const targetWs = new WebSocket(targetUrl, { headers });
+			const targetWs = new WebSocket(targetUrl, protocols);
 			state.targetWs = targetWs;
 
 			// Setup connection promise

package/src/test/mod.ts

@@ -5,7 +5,10 @@ import { type TestContext, vi } from "vitest";
 import { type Client, createClient } from "@/client/mod";
 import { chooseDefaultDriver } from "@/drivers/default";
 import { createFileSystemOrMemoryDriver } from "@/drivers/file-system/mod";
-import { getInspectorUrl } from "@/inspector/utils";
+import {
+	configureInspectorAccessToken,
+	getInspectorUrl,
+} from "@/inspector/utils";
 import { createManagerRouter } from "@/manager/router";
 import type { Registry } from "@/registry/mod";
 import { RunConfigSchema } from "@/registry/run-config";
@@ -27,11 +30,12 @@ function serve(registry: Registry<any>, inputConfig?: InputConfig): ServerType {
 	const runConfig = RunConfigSchema.parse(inputConfig);
 	const driver = inputConfig.driver ?? createFileSystemOrMemoryDriver(false);
 	const managerDriver = driver.manager(registry.config, config);
+	configureInspectorAccessToken(config, managerDriver);
 	const { router } = createManagerRouter(
 		registry.config,
 		runConfig,
 		managerDriver,
-		false,
+		undefined,
 	);
 
 	// Inject WebSocket

package/src/utils.ts

@@ -77,6 +77,25 @@ const TIMEOUT_MAX = 2147483647; // 2^31-1
 
 export type LongTimeoutHandle = { abort: () => void };
 
+/**
+ * Polyfill for Promise.withResolvers().
+ *
+ * This is specifically for Cloudflare Workers. Their implementation of Promise.withResolvers does not work correctly.
+ */
+export function promiseWithResolvers<T>(): {
+	promise: Promise<T>;
+	resolve: (value: T | PromiseLike<T>) => void;
+	reject: (reason?: any) => void;
+} {
+	let resolve!: (value: T | PromiseLike<T>) => void;
+	let reject!: (reason?: any) => void;
+	const promise = new Promise<T>((res, rej) => {
+		resolve = res;
+		reject = rej;
+	});
+	return { promise, resolve, reject };
+}
+
 export function setLongTimeout(
 	listener: () => void,
 	after: number,
@@ -115,7 +134,7 @@ export class SinglePromiseQueue {
 	runningDrainLoop?: Promise<void>;
 
 	/** Pending resolver fro the currently queued entry. */
-	#pending?: PromiseWithResolvers<void>;
+	#pending?: ReturnType<typeof promiseWithResolvers<void>>;
 
 	/** Queue the next operation and return a promise that resolves when it flushes. */
 	enqueue(op: () => Promise<void>): Promise<void> {
@@ -124,7 +143,7 @@ export class SinglePromiseQueue {
 
 		// Ensure a shared resolver exists for all callers in this cycle
 		if (!this.#pending) {
-			this.#pending = Promise.withResolvers<void>();
+			this.#pending = promiseWithResolvers<void>();
 		}
 
 		const waitForThisCycle = this.#pending.promise;