feat: kv api

Author: NathanFlurry

rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/actor.ts

@@ -6,6 +6,7 @@ import type { Conn, ConnId } from "../../conn/mod";
 import type { AnyDatabaseProvider, InferDatabaseClient } from "../../database";
 import type { ActorDefinition, AnyActorDefinition } from "../../definition";
 import type { ActorInstance, SaveStateOptions } from "../../instance/mod";
+import { ActorKv } from "../../instance/kv";
 import type { Schedule } from "../../schedule";
 
 /**
@@ -27,6 +28,7 @@ export class ActorContext<
 		TInput,
 		TDatabase
 	>;
+	#kv: ActorKv | undefined;
 
 	constructor(
 		actor: ActorInstance<
@@ -41,6 +43,16 @@ export class ActorContext<
 		this.#actor = actor;
 	}
 
+	/**
+	 * Gets the KV storage interface.
+	 */
+	get kv(): ActorKv {
+		if (!this.#kv) {
+			this.#kv = new ActorKv(this.#actor.driver, this.#actor.id);
+		}
+		return this.#kv;
+	}
+
 	/**
 	 * Get the actor state
 	 *

rivetkit-typescript/packages/rivetkit/src/actor/instance/kv.ts

@@ -1,9 +1,103 @@
+import type { ActorDriver } from "../driver";
+
 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
+	KV: Uint8Array.from([4]), // Prefix for user-facing KV storage
 };
 
+// Helper to create a prefixed key for user-facing KV storage
+function makePrefixedKey(key: Uint8Array): Uint8Array {
+	const prefixed = new Uint8Array(KEYS.KV.length + key.length);
+	prefixed.set(KEYS.KV, 0);
+	prefixed.set(key, KEYS.KV.length);
+	return prefixed;
+}
+
+// Helper to remove the prefix from a key
+function removePrefixFromKey(prefixedKey: Uint8Array): Uint8Array {
+	return prefixedKey.slice(KEYS.KV.length);
+}
+
+/**
+ * User-facing KV storage interface exposed on ActorContext.
+ */
+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(key: Uint8Array): Promise<Uint8Array | null> {
+		const results = await this.#driver.kvBatchGet(this.#actorId, [
+			makePrefixedKey(key),
+		]);
+		return results[0] ?? null;
+	}
+
+	/**
+	 * Get multiple values by keys.
+	 */
+	async getBatch(keys: Uint8Array[]): Promise<(Uint8Array | null)[]> {
+		const prefixedKeys = keys.map(makePrefixedKey);
+		return this.#driver.kvBatchGet(this.#actorId, prefixedKeys);
+	}
+
+	/**
+	 * Put a single key-value pair.
+	 */
+	async put(key: Uint8Array, value: Uint8Array): Promise<void> {
+		await this.#driver.kvBatchPut(this.#actorId, [
+			[makePrefixedKey(key), value],
+		]);
+	}
+
+	/**
+	 * Put multiple key-value pairs.
+	 */
+	async putBatch(entries: [Uint8Array, Uint8Array][]): Promise<void> {
+		const prefixedEntries: [Uint8Array, Uint8Array][] = entries.map(
+			([key, value]) => [makePrefixedKey(key), value],
+		);
+		await this.#driver.kvBatchPut(this.#actorId, prefixedEntries);
+	}
+
+	/**
+	 * Delete a single key.
+	 */
+	async delete(key: Uint8Array): Promise<void> {
+		await this.#driver.kvBatchDelete(this.#actorId, [makePrefixedKey(key)]);
+	}
+
+	/**
+	 * Delete multiple keys.
+	 */
+	async deleteBatch(keys: Uint8Array[]): Promise<void> {
+		const prefixedKeys = keys.map(makePrefixedKey);
+		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(prefix: Uint8Array): Promise<[Uint8Array, Uint8Array][]> {
+		const prefixedPrefix = makePrefixedKey(prefix);
+		const results = await this.#driver.kvListPrefix(
+			this.#actorId,
+			prefixedPrefix,
+		);
+		return results.map(([key, value]) => [removePrefixFromKey(key), value]);
+	}
+}
+
 // Helper to create a connection key
 export function makeConnKey(connId: string): Uint8Array {
 	const encoder = new TextEncoder();

rivetkit-typescript/packages/rivetkit/src/actor/mod.ts

@@ -76,7 +76,7 @@ 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 { ActorKv, KEYS as KV_KEYS } from "./instance/kv";
 export type { AnyActorInstance } from "./instance/mod";
 export {
 	type ActorRouter,