rivetkit 2.0.38 → 2.0.39

package/src/actor/instance/kv.ts

@@ -1,15 +1,241 @@
-export const KEYS = {
-	PERSIST_DATA: Uint8Array.from([1]),
-	CONN_PREFIX: Uint8Array.from([2]), // Prefix for connection keys
-	INSPECTOR_TOKEN: Uint8Array.from([3]), // Inspector token key
-};
-
-// Helper to create a connection key
-export function makeConnKey(connId: string): Uint8Array {
-	const encoder = new TextEncoder();
-	const connIdBytes = encoder.encode(connId);
-	const key = new Uint8Array(KEYS.CONN_PREFIX.length + connIdBytes.length);
-	key.set(KEYS.CONN_PREFIX, 0);
-	key.set(connIdBytes, KEYS.CONN_PREFIX.length);
-	return key;
+import type { ActorDriver } from "../driver";
+import { makePrefixedKey, removePrefixFromKey } from "./keys";
+
+/**
+ * User-facing KV storage interface exposed on ActorContext.
+ */
+type KvValueType = "text" | "arrayBuffer" | "binary";
+type KvKeyType = "text" | "binary";
+type KvKey = Uint8Array | string;
+
+type KvValueTypeMap = {
+	text: string;
+	arrayBuffer: ArrayBuffer;
+	binary: Uint8Array;
+};
+
+type KvKeyTypeMap = {
+	text: string;
+	binary: Uint8Array;
+};
+
+type KvValueOptions<T extends KvValueType = "text"> = {
+	type?: T;
+};
+
+type KvListOptions<
+	T extends KvValueType = "text",
+	K extends KvKeyType = "text",
+> = KvValueOptions<T> & {
+	keyType?: K;
+};
+
+const textEncoder = new TextEncoder();
+const textDecoder = new TextDecoder();
+
+function encodeKey<K extends KvKeyType = KvKeyType>(
+	key: KvKeyTypeMap[K],
+	keyType?: K,
+): Uint8Array {
+	if (key instanceof Uint8Array) {
+		return key;
+	}
+	const resolvedKeyType = keyType ?? "text";
+	if (resolvedKeyType === "binary") {
+		throw new TypeError("Expected a Uint8Array when keyType is binary");
+	}
+	return textEncoder.encode(key);
+}
+
+function decodeKey<K extends KvKeyType = "text">(
+	key: Uint8Array,
+	keyType?: K,
+): KvKeyTypeMap[K] {
+	const resolvedKeyType = keyType ?? "text";
+	switch (resolvedKeyType) {
+		case "text":
+			return textDecoder.decode(key) as KvKeyTypeMap[K];
+		case "binary":
+			return key as KvKeyTypeMap[K];
+		default:
+			throw new TypeError("Invalid kv key type");
+	}
+}
+
+function resolveValueType(
+	value: string | Uint8Array | ArrayBuffer,
+): KvValueType {
+	if (typeof value === "string") {
+		return "text";
+	}
+	if (value instanceof Uint8Array) {
+		return "binary";
+	}
+	if (value instanceof ArrayBuffer) {
+		return "arrayBuffer";
+	}
+	throw new TypeError("Invalid kv value");
+}
+
+function encodeValue<T extends KvValueType = KvValueType>(
+	value: KvValueTypeMap[T],
+	options?: KvValueOptions<T>,
+): Uint8Array {
+	const type =
+		options?.type ??
+		resolveValueType(value as string | Uint8Array | ArrayBuffer);
+	switch (type) {
+		case "text":
+			if (typeof value !== "string") {
+				throw new TypeError("Expected a string when type is text");
+			}
+			return textEncoder.encode(value);
+		case "arrayBuffer":
+			if (!(value instanceof ArrayBuffer)) {
+				throw new TypeError("Expected an ArrayBuffer when type is arrayBuffer");
+			}
+			return new Uint8Array(value);
+		case "binary":
+			if (!(value instanceof Uint8Array)) {
+				throw new TypeError("Expected a Uint8Array when type is binary");
+			}
+			return value;
+		default:
+			throw new TypeError("Invalid kv value type");
+	}
+}
+
+function decodeValue<T extends KvValueType = "text">(
+	value: Uint8Array,
+	options?: KvValueOptions<T>,
+): KvValueTypeMap[T] {
+	const type = options?.type ?? "text";
+	switch (type) {
+		case "text":
+			return textDecoder.decode(value) as KvValueTypeMap[T];
+		case "arrayBuffer": {
+			const copy = new Uint8Array(value.byteLength);
+			copy.set(value);
+			return copy.buffer as KvValueTypeMap[T];
+		}
+		case "binary":
+			return value as KvValueTypeMap[T];
+		default:
+			throw new TypeError("Invalid kv value type");
+	}
+}
+
+export class ActorKv {
+	#driver: ActorDriver;
+	#actorId: string;
+
+	constructor(driver: ActorDriver, actorId: string) {
+		this.#driver = driver;
+		this.#actorId = actorId;
+	}
+
+	/**
+	 * Get a single value by key.
+	 */
+	async get<T extends KvValueType = "text">(
+		key: KvKey,
+		options?: KvValueOptions<T>,
+	): Promise<KvValueTypeMap[T] | null> {
+		const results = await this.#driver.kvBatchGet(this.#actorId, [
+			makePrefixedKey(encodeKey(key)),
+		]);
+		const result = results[0];
+		if (!result) {
+			return null;
+		}
+		return decodeValue(result, options);
+	}
+
+	/**
+	 * Get multiple values by keys.
+	 */
+	async getBatch<T extends KvValueType = "text">(
+		keys: KvKey[],
+		options?: KvValueOptions<T>,
+	): Promise<(KvValueTypeMap[T] | null)[]> {
+		const prefixedKeys = keys.map((key) =>
+			makePrefixedKey(encodeKey(key)),
+		);
+		const results = await this.#driver.kvBatchGet(
+			this.#actorId,
+			prefixedKeys,
+		);
+		return results.map((result) =>
+			result ? decodeValue(result, options) : null,
+		);
+	}
+
+	/**
+	 * Put a single key-value pair.
+	 */
+	async put<T extends KvValueType = KvValueType>(
+		key: KvKey,
+		value: KvValueTypeMap[T],
+		options?: KvValueOptions<T>,
+	): Promise<void> {
+		await this.#driver.kvBatchPut(this.#actorId, [
+			[makePrefixedKey(encodeKey(key)), encodeValue(value, options)],
+		]);
+	}
+
+	/**
+	 * Put multiple key-value pairs.
+	 */
+	async putBatch<T extends KvValueType = KvValueType>(
+		entries: [KvKey, KvValueTypeMap[T]][],
+		options?: KvValueOptions<T>,
+	): Promise<void> {
+		const prefixedEntries: [Uint8Array, Uint8Array][] = entries.map(
+			([key, value]) => [
+				makePrefixedKey(encodeKey(key)),
+				encodeValue(value, options),
+			],
+		);
+		await this.#driver.kvBatchPut(this.#actorId, prefixedEntries);
+	}
+
+	/**
+	 * Delete a single key.
+	 */
+	async delete(key: KvKey): Promise<void> {
+		await this.#driver.kvBatchDelete(this.#actorId, [
+			makePrefixedKey(encodeKey(key)),
+		]);
+	}
+
+	/**
+	 * Delete multiple keys.
+	 */
+	async deleteBatch(keys: KvKey[]): Promise<void> {
+		const prefixedKeys = keys.map((key) =>
+			makePrefixedKey(encodeKey(key)),
+		);
+		await this.#driver.kvBatchDelete(this.#actorId, prefixedKeys);
+	}
+
+	/**
+	 * List all keys with a given prefix.
+	 * Returns key-value pairs where keys have the user prefix removed.
+	 */
+	async list<T extends KvValueType = "text", K extends KvKeyType = "text">(
+		prefix: KvKeyTypeMap[K],
+		options?: KvListOptions<T, K>,
+	): Promise<[KvKeyTypeMap[K], KvValueTypeMap[T]][]> {
+		const prefixedPrefix = makePrefixedKey(
+			encodeKey(prefix, options?.keyType),
+		);
+		const results = await this.#driver.kvListPrefix(
+			this.#actorId,
+			prefixedPrefix,
+		);
+		return results.map(([key, value]) => [
+			decodeKey<K>(removePrefixFromKey(key), options?.keyType),
+			decodeValue<T>(value, options),
+		]);
+	}
 }

package/src/actor/instance/mod.ts

@@ -44,7 +44,7 @@ import {
 } from "../utils";
 import { ConnectionManager } from "./connection-manager";
 import { EventManager } from "./event-manager";
-import { KEYS } from "./kv";
+import { KEYS } from "./keys";
 import {
 	convertActorFromBarePersisted,
 	type PersistedActor,
@@ -203,7 +203,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	}
 
 	get actions(): string[] {
-		return Object.keys(this.#config.actions);
+		return Object.keys(this.#config.actions ?? {});
 	}
 
 	get config(): ActorConfig<S, CP, CS, V, I, DB> {
@@ -516,12 +516,13 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	): Promise<unknown> {
 		this.assertReady();
 
-		if (!(actionName in this.#config.actions)) {
+		const actions = this.#config.actions ?? {};
+		if (!(actionName in actions)) {
 			this.#rLog.warn({ msg: "action does not exist", actionName });
 			throw new errors.ActionNotFound(actionName);
 		}
 
-		const actionFunction = this.#config.actions[actionName];
+		const actionFunction = actions[actionName];
 		if (typeof actionFunction !== "function") {
 			this.#rLog.warn({
 				msg: "action is not a function",

package/src/actor/instance/state-manager.ts

@@ -12,7 +12,7 @@ import { convertConnToBarePersistedConn } from "../conn/persisted";
 import type { ActorDriver } from "../driver";
 import * as errors from "../errors";
 import { isConnStatePath, isStatePath } from "../utils";
-import { KEYS, makeConnKey } from "./kv";
+import { KEYS, makeConnKey } from "./keys";
 import type { ActorInstance } from "./mod";
 import { convertActorToBarePersisted, type PersistedActor } from "./persisted";
 

package/src/actor/mod.ts

@@ -76,7 +76,8 @@ export type { AnyConn, Conn } from "./conn/mod";
 export type { ActorDefinition, AnyActorDefinition } from "./definition";
 export { lookupInRegistry } from "./definition";
 export { UserError, type UserErrorOptions } from "./errors";
-export { KEYS as KV_KEYS } from "./instance/kv";
+export { KEYS as KV_KEYS } from "./instance/keys";
+export { ActorKv } from "./instance/kv";
 export type { AnyActorInstance } from "./instance/mod";
 export {
 	type ActorRouter,

package/src/actor/router-websocket-endpoints.ts

@@ -384,7 +384,8 @@ export function parseWebSocketProtocols(
 		}
 	}
 
-	const encoding = EncodingSchema.parse(encodingRaw);
+	// Default to "json" encoding for raw WebSocket connections without subprotocols
+	const encoding = EncodingSchema.parse(encodingRaw ?? "json");
 	const connParams = connParamsRaw ? JSON.parse(connParamsRaw) : undefined;
 
 	return { encoding, connParams };

package/src/client/actor-handle.ts

@@ -239,9 +239,7 @@ export class ActorHandleRaw {
 	}
 
 	/**
-	 * Resolves the actor to get its unique actor ID
-	 *
-	 * @returns {Promise<string>} - A promise that resolves to the actor's ID
+	 * Resolves the actor to get its unique actor ID.
 	 */
 	async resolve({ signal }: { signal?: AbortSignal } = {}): Promise<string> {
 		if (
@@ -277,6 +275,18 @@ export class ActorHandleRaw {
 			assertUnreachable(this.#actorQuery);
 		}
 	}
+
+	/**
+	 * Returns the raw URL for routing traffic to the actor.
+	 */
+	async getGatewayUrl(): Promise<string> {
+		const { actorId } = await queryActor(
+			undefined,
+			this.#actorQuery,
+			this.#driver,
+		);
+		return await this.#driver.buildGatewayUrl(actorId);
+	}
 }
 
 /**

package/src/client/mod.ts

@@ -22,7 +22,7 @@ export {
 	ManagerError,
 } from "@/client/errors";
 export type { CreateRequest } from "@/manager/protocol/query";
-export { KEYS as KV_KEYS } from "../actor/instance/kv";
+export { KEYS as KV_KEYS } from "../actor/instance/keys";
 export type { ActorActionFunction } from "./actor-common";
 export type {
 	ActorConn,

package/src/driver-helpers/utils.ts

@@ -1,5 +1,5 @@
 import * as cbor from "cbor-x";
-import { KEYS } from "@/actor/instance/kv";
+import { KEYS } from "@/actor/instance/keys";
 import type * as persistSchema from "@/schemas/actor-persist/mod";
 import {
 	ACTOR_VERSIONED,

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

@@ -22,6 +22,7 @@ import { runActorErrorHandlingTests } from "./tests/actor-error-handling";
 import { runActorHandleTests } from "./tests/actor-handle";
 import { runActorInlineClientTests } from "./tests/actor-inline-client";
 import { runActorInspectorTests } from "./tests/actor-inspector";
+import { runActorKvTests } from "./tests/actor-kv";
 import { runActorMetadataTests } from "./tests/actor-metadata";
 import { runActorOnStateChangeTests } from "./tests/actor-onstatechange";
 import { runActorVarsTests } from "./tests/actor-vars";
@@ -122,6 +123,8 @@ export function runDriverTests(
 
 					runActorInlineClientTests(driverTestConfig);
 
+					runActorKvTests(driverTestConfig);
+
 					runRawHttpTests(driverTestConfig);
 
 					runRawHttpRequestPropertiesTests(driverTestConfig);

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

@@ -237,6 +237,9 @@ export function createTestInlineClientDriver(
 			);
 			return upgradeWebSocket(() => wsHandler)(c, noopNext());
 		},
+		async buildGatewayUrl(actorId: string): Promise<string> {
+			return `${endpoint}/gateway/${actorId}`;
+		},
 		displayInformation(): ManagerDisplayInformation {
 			return { properties: {} };
 		},

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

@@ -0,0 +1,44 @@
+import type { DriverTestConfig } from "../mod";
+import { setupDriverTest } from "../utils";
+import { describe, expect, test, type TestContext } from "vitest";
+
+export function runActorKvTests(driverTestConfig: DriverTestConfig) {
+	describe("Actor KV Tests", () => {
+		test("supports text encoding and decoding", async (c: TestContext) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+			const kvHandle = client.kvActor.getOrCreate(["kv-text"]);
+
+			await kvHandle.putText("greeting", "hello");
+			const value = await kvHandle.getText("greeting");
+			expect(value).toBe("hello");
+
+			await kvHandle.putText("prefix-a", "alpha");
+			await kvHandle.putText("prefix-b", "beta");
+
+			const results = await kvHandle.listText("prefix-");
+			const sorted = results.sort((a, b) => a.key.localeCompare(b.key));
+			expect(sorted).toEqual([
+				{ key: "prefix-a", value: "alpha" },
+				{ key: "prefix-b", value: "beta" },
+			]);
+		});
+
+		test(
+			"supports arrayBuffer encoding and decoding",
+			async (c: TestContext) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+			const kvHandle = client.kvActor.getOrCreate(["kv-array-buffer"]);
+
+			const values = await kvHandle.roundtripArrayBuffer("bytes", [
+				4,
+				8,
+				15,
+				16,
+				23,
+				42,
+			]);
+			expect(values).toEqual([4, 8, 15, 16, 23, 42]);
+		},
+		);
+	});
+}

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

@@ -39,6 +39,10 @@ export async function setupDriverTest(
 			namespace,
 			runnerName,
 			encoding: driverTestConfig.encoding,
+			// Disable metadata lookup to prevent redirect to the wrong port.
+			// Each test starts a new server on a dynamic port, but the
+			// registry's publicEndpoint defaults to port 6420.
+			disableMetadataLookup: true,
 		});
 	} else if (driverTestConfig.clientType === "inline") {
 		// Use inline client from driver

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

@@ -11,7 +11,7 @@ import { WSContext, type WSContextInit } from "hono/ws";
 import invariant from "invariant";
 import { type AnyConn, CONN_STATE_MANAGER_SYMBOL } from "@/actor/conn/mod";
 import { lookupInRegistry } from "@/actor/definition";
-import { KEYS } from "@/actor/instance/kv";
+import { KEYS } from "@/actor/instance/keys";
 import { deserializeActorKey } from "@/actor/keys";
 import { getValueLength } from "@/actor/protocol/old";
 import { type ActorRouter, createActorRouter } from "@/actor/router";
@@ -153,7 +153,7 @@ export class EngineActorDriver implements ActorDriver {
 			onConnected: () => {
 				this.#runnerStarted.resolve(undefined);
 			},
-			onDisconnected: (_code, _reason) => { },
+			onDisconnected: (_code, _reason) => {},
 			onShutdown: () => {
 				this.#runnerStopped.resolve(undefined);
 				this.#isRunnerStopped = true;
@@ -358,7 +358,7 @@ export class EngineActorDriver implements ActorDriver {
 	async serverlessHandleStart(c: HonoContext): Promise<Response> {
 		return streamSSE(c, async (stream) => {
 			// NOTE: onAbort does not work reliably
-			stream.onAbort(() => { });
+			stream.onAbort(() => {});
 			c.req.raw.signal.addEventListener("abort", () => {
 				logger().debug("SSE aborted, shutting down runner");
 

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

@@ -147,6 +147,11 @@ export class FileSystemManagerDriver implements ManagerDriver {
 		return upgradeWebSocket(() => wsHandler)(c, noopNext());
 	}
 
+	async buildGatewayUrl(actorId: string): Promise<string> {
+		const port = this.#config.managerPort ?? 6420;
+		return `http://127.0.0.1:${port}/gateway/${encodeURIComponent(actorId)}`;
+	}
+
 	async getForId({
 		actorId,
 	}: GetForIdInput): Promise<ActorOutput | undefined> {

package/src/manager/driver.ts

@@ -32,6 +32,13 @@ export interface ManagerDriver {
 		params: unknown,
 	): Promise<Response>;
 
+	/**
+	 * Build a public gateway URL for a specific actor.
+	 *
+	 * This lives on the driver because the base endpoint varies by runtime.
+	 */
+	buildGatewayUrl(actorId: string): Promise<string>;
+
 	displayInformation(): ManagerDisplayInformation;
 
 	extraStartupLog?: () => Record<string, unknown>;

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

@@ -1,6 +1,6 @@
 import type { ClientConfig } from "@/client/config";
 import { HEADER_RIVET_TOKEN } from "@/common/actor-router-consts";
-import { combineUrlPath } from "@/utils";
+import { buildActorGatewayUrl } from "./actor-websocket-client";
 import { getEndpoint } from "./api-utils";
 
 export async function sendHttpRequestToActor(
@@ -11,9 +11,11 @@ export async function sendHttpRequestToActor(
 	// Route through guard port
 	const url = new URL(actorRequest.url);
 	const endpoint = getEndpoint(runConfig);
-	const guardUrl = combineUrlPath(
+	const guardUrl = buildActorGatewayUrl(
 		endpoint,
-		`/gateway/${actorId}${url.pathname}${url.search}`,
+		actorId,
+		runConfig.token,
+		`${url.pathname}${url.search}`,
 	);
 
 	// Handle body properly based on method and presence

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

@@ -13,6 +13,18 @@ import { combineUrlPath } from "@/utils";
 import { getEndpoint } from "./api-utils";
 import { logger } from "./log";
 
+export function buildActorGatewayUrl(
+	endpoint: string,
+	actorId: string,
+	token: string | undefined,
+	path = "",
+): string {
+	const tokenSegment =
+		token !== undefined ? `@${encodeURIComponent(token)}` : "";
+	const gatewayPath = `/gateway/${encodeURIComponent(actorId)}${tokenSegment}${path}`;
+	return combineUrlPath(endpoint, gatewayPath);
+}
+
 export async function openWebSocketToActor(
 	runConfig: ClientConfig,
 	path: string,
@@ -24,13 +36,12 @@ export async function openWebSocketToActor(
 
 	// WebSocket connections go through guard
 	const endpoint = getEndpoint(runConfig);
-	let gatewayPath;
-	if (runConfig.token !== undefined) {
-		gatewayPath = `/gateway/${encodeURIComponent(actorId)}@${encodeURIComponent(runConfig.token)}${path}`;
-	} else {
-		gatewayPath = `/gateway/${encodeURIComponent(actorId)}${path}`;
-	}
-	const guardUrl = combineUrlPath(endpoint, gatewayPath);
+	const guardUrl = buildActorGatewayUrl(
+		endpoint,
+		actorId,
+		runConfig.token,
+		path,
+	);
 
 	logger().debug({
 		msg: "opening websocket to actor via guard",

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

@@ -21,6 +21,7 @@ import { combineUrlPath, type GetUpgradeWebSocket } from "@/utils";
 import { getNextPhase } from "@/utils/env-vars";
 import { sendHttpRequestToActor } from "./actor-http-client";
 import {
+	buildActorGatewayUrl,
 	buildWebSocketProtocols,
 	openWebSocketToActor,
 } from "./actor-websocket-client";
@@ -309,6 +310,15 @@ export class RemoteManagerDriver implements ManagerDriver {
 		);
 	}
 
+	async buildGatewayUrl(actorId: string): Promise<string> {
+		if (this.#metadataPromise) {
+			await this.#metadataPromise;
+		}
+
+		const endpoint = getEndpoint(this.#config);
+		return buildActorGatewayUrl(endpoint, actorId, this.#config.token);
+	}
+
 	async proxyRequest(
 		_c: HonoContext,
 		actorRequest: Request,