rivetkit 2.0.24-rc.1 → 2.0.25-rc.1

package/src/client/client.ts

@@ -1,5 +1,4 @@
 import type { AnyActorDefinition } from "@/actor/definition";
-import type { Transport } from "@/actor/protocol/old";
 import type { Encoding } from "@/actor/protocol/serde";
 import type { ManagerDriver } from "@/driver-helpers/mod";
 import type { ActorQuery } from "@/manager/protocol/query";
@@ -149,7 +148,6 @@ export interface Region {
 
 export const ACTOR_CONNS_SYMBOL = Symbol("actorConns");
 export const CREATE_ACTOR_CONN_PROXY = Symbol("createActorConnProxy");
-export const TRANSPORT_SYMBOL = Symbol("transport");
 
 /**
  * Client for managing & connecting to actors.
@@ -164,7 +162,6 @@ export class ClientRaw {
 
 	#driver: ManagerDriver;
 	#encodingKind: Encoding;
-	[TRANSPORT_SYMBOL]: Transport;
 
 	/**
 	 * Creates an instance of Client.
@@ -173,7 +170,6 @@ export class ClientRaw {
 		this.#driver = driver;
 
 		this.#encodingKind = config.encoding ?? "bare";
-		this[TRANSPORT_SYMBOL] = config.transport ?? "websocket";
 	}
 
 	/**

package/src/client/config.ts

@@ -1,5 +1,4 @@
 import z from "zod";
-import { TransportSchema } from "@/actor/protocol/old";
 import { EncodingSchema } from "@/actor/protocol/serde";
 import { type GetUpgradeWebSocket, getEnvUniversal } from "@/utils";
 
@@ -33,15 +32,13 @@ export const ClientConfigSchema = z.object({
 
 	encoding: EncodingSchema.default("bare"),
 
-	transport: TransportSchema.default("websocket"),
-
 	headers: z.record(z.string()).optional().default({}),
 
 	// See RunConfig.getUpgradeWebSocket
 	getUpgradeWebSocket: z.custom<GetUpgradeWebSocket>().optional(),
 
 	/** Whether to automatically perform health checks when the client is created. */
-	disableHealthCheck: z.boolean().optional().default(false),
+	disableMetadataLookup: z.boolean().optional().default(false),
 });
 
 export type ClientConfig = z.infer<typeof ClientConfigSchema>;

package/src/client/mod.ts

@@ -11,7 +11,6 @@ export {
 	ActorDefinition,
 	AnyActorDefinition,
 } from "@/actor/definition";
-export type { Transport } from "@/actor/protocol/old";
 export type { Encoding } from "@/actor/protocol/serde";
 export {
 	ActorClientError,

package/src/client/raw-utils.ts

@@ -1,5 +1,5 @@
 import invariant from "invariant";
-import { PATH_RAW_WEBSOCKET_PREFIX } from "@/common/actor-router-consts";
+import { PATH_WEBSOCKET_PREFIX } from "@/common/actor-router-consts";
 import { deconstructError } from "@/common/utils";
 import { HEADER_CONN_PARAMS, type ManagerDriver } from "@/driver-helpers/mod";
 import type { ActorQuery } from "@/manager/protocol/query";
@@ -69,7 +69,7 @@ export async function rawHttpFetch(
 
 		// Build the URL with normalized path
 		const normalizedPath = path.startsWith("/") ? path.slice(1) : path;
-		const url = new URL(`http://actor/raw/http/${normalizedPath}`);
+		const url = new URL(`http://actor/request/${normalizedPath}`);
 
 		// Forward conn params if provided
 		const proxyRequestHeaders = new Headers(mergedInit.headers);
@@ -132,7 +132,7 @@ export async function rawWebSocket(
 		}
 	}
 
-	const fullPath = `${PATH_RAW_WEBSOCKET_PREFIX}${pathPortion}${queryPortion}`;
+	const fullPath = `${PATH_WEBSOCKET_PREFIX}${pathPortion}${queryPortion}`;
 
 	logger().debug({
 		msg: "opening websocket",

package/src/client/utils.ts

@@ -1,10 +1,15 @@
 import * as cbor from "cbor-x";
 import invariant from "invariant";
+import type { z } from "zod";
 import type { Encoding } from "@/actor/protocol/serde";
 import { assertUnreachable } from "@/common/utils";
 import type { VersionedDataHandler } from "@/common/versioned-data";
 import type { HttpResponseError } from "@/schemas/client-protocol/mod";
 import { HTTP_RESPONSE_ERROR_VERSIONED } from "@/schemas/client-protocol/versioned";
+import {
+	type HttpResponseError as HttpResponseErrorJson,
+	HttpResponseErrorSchema,
+} from "@/schemas/client-protocol-zod/mod";
 import {
 	contentTypeForEncoding,
 	deserializeWithEncoding,
@@ -33,25 +38,51 @@ export function messageLength(message: WebSocketMessage): number {
 	assertUnreachable(message);
 }
 
-export interface HttpRequestOpts<RequestBody, ResponseBody> {
+export interface HttpRequestOpts<
+	RequestBare,
+	ResponseBare,
+	RequestJson = RequestBare,
+	ResponseJson = ResponseBare,
+	Request = RequestBare,
+	Response = ResponseBare,
+> {
 	method: string;
 	url: string;
 	headers: Record<string, string>;
-	body?: RequestBody;
+	body?: Request;
 	encoding: Encoding;
 	skipParseResponse?: boolean;
 	signal?: AbortSignal;
-	customFetch?: (req: Request) => Promise<Response>;
-	requestVersionedDataHandler: VersionedDataHandler<RequestBody> | undefined;
+	customFetch?: (req: globalThis.Request) => Promise<globalThis.Response>;
+	requestVersionedDataHandler: VersionedDataHandler<RequestBare> | undefined;
 	responseVersionedDataHandler:
-		| VersionedDataHandler<ResponseBody>
+		| VersionedDataHandler<ResponseBare>
 		| undefined;
+	requestZodSchema: z.ZodType<RequestJson>;
+	responseZodSchema: z.ZodType<ResponseJson>;
+	requestToJson: (value: Request) => RequestJson;
+	requestToBare: (value: Request) => RequestBare;
+	responseFromJson: (value: ResponseJson) => Response;
+	responseFromBare: (value: ResponseBare) => Response;
 }
 
 export async function sendHttpRequest<
-	RequestBody = unknown,
-	ResponseBody = unknown,
->(opts: HttpRequestOpts<RequestBody, ResponseBody>): Promise<ResponseBody> {
+	RequestBare = unknown,
+	ResponseBare = unknown,
+	RequestJson = RequestBare,
+	ResponseJson = ResponseBare,
+	Request = RequestBare,
+	Response = ResponseBare,
+>(
+	opts: HttpRequestOpts<
+		RequestBare,
+		ResponseBare,
+		RequestJson,
+		ResponseJson,
+		Request,
+		Response
+	>,
+): Promise<Response> {
 	logger().debug({
 		msg: "sending http request",
 		url: opts.url,
@@ -64,19 +95,22 @@ export async function sendHttpRequest<
 	if (opts.method === "POST" || opts.method === "PUT") {
 		invariant(opts.body !== undefined, "missing body");
 		contentType = contentTypeForEncoding(opts.encoding);
-		bodyData = serializeWithEncoding<RequestBody>(
+		bodyData = serializeWithEncoding<RequestBare, RequestJson, Request>(
 			opts.encoding,
 			opts.body,
 			opts.requestVersionedDataHandler,
+			opts.requestZodSchema,
+			opts.requestToJson,
+			opts.requestToBare,
 		);
 	}
 
 	// Send request
-	let response: Response;
+	let response: globalThis.Response;
 	try {
 		// Make the HTTP request
 		response = await (opts.customFetch ?? fetch)(
-			new Request(opts.url, {
+			new globalThis.Request(opts.url, {
 				method: opts.method,
 				headers: {
 					...opts.headers,
@@ -100,27 +134,52 @@ export async function sendHttpRequest<
 
 	// Parse response error
 	if (!response.ok) {
-		// Attempt to parse structured data
 		const bufferResponse = await response.arrayBuffer();
-		let responseData: HttpResponseError;
+		const contentType = response.headers.get("content-type");
+		const rayId = response.headers.get("x-rivet-ray-id");
+
+		// Determine encoding from Content-Type header, defaulting to provided encoding
+		const encoding: Encoding = contentType?.includes("application/json")
+			? "json"
+			: opts.encoding;
+
+		// Attempt to parse structured error data
 		try {
-			responseData = deserializeWithEncoding(
-				opts.encoding,
+			const responseData = deserializeWithEncoding(
+				encoding,
 				new Uint8Array(bufferResponse),
 				HTTP_RESPONSE_ERROR_VERSIONED,
+				HttpResponseErrorSchema,
+				// JSON: metadata is already unknown
+				(json): HttpResponseErrorJson => json as HttpResponseErrorJson,
+				// BARE: decode ArrayBuffer metadata to unknown
+				(bare): any => ({
+					group: bare.group,
+					code: bare.code,
+					message: bare.message,
+					metadata: bare.metadata
+						? cbor.decode(new Uint8Array(bare.metadata))
+						: undefined,
+				}),
+			);
+
+			throw new ActorError(
+				responseData.group,
+				responseData.code,
+				responseData.message,
+				responseData.metadata,
 			);
 		} catch (error) {
-			//logger().warn("failed to cleanly parse error, this is likely because a non-structured response is being served", {
-			//	error: stringifyError(error),
-			//});
+			// If it's already an ActorError, re-throw it
+			if (error instanceof ActorError) {
+				throw error;
+			}
 
-			// Error is not structured
+			// Otherwise, fall back to generic error with text response
 			const textResponse = new TextDecoder("utf-8", {
 				fatal: false,
 			}).decode(bufferResponse);
 
-			const rayId = response.headers.get("x-rivet-ray-id");
-
 			if (rayId) {
 				throw new HttpRequestError(
 					`${response.statusText} (${response.status}) (Ray ID: ${rayId}):\n${textResponse}`,
@@ -131,36 +190,23 @@ export async function sendHttpRequest<
 				);
 			}
 		}
-
-		// Decode metadata based on encoding - only binary encodings have CBOR-encoded metadata
-		let decodedMetadata: unknown;
-		if (responseData.metadata && encodingIsBinary(opts.encoding)) {
-			decodedMetadata = cbor.decode(
-				new Uint8Array(responseData.metadata),
-			);
-		}
-
-		// Throw structured error
-		throw new ActorError(
-			responseData.group,
-			responseData.code,
-			responseData.message,
-			decodedMetadata,
-		);
 	}
 
 	// Some requests don't need the success response to be parsed, so this can speed things up
 	if (opts.skipParseResponse) {
-		return undefined as ResponseBody;
+		return undefined as Response;
 	}
 
 	// Parse the response based on encoding
 	try {
 		const buffer = new Uint8Array(await response.arrayBuffer());
-		return deserializeWithEncoding(
+		return deserializeWithEncoding<ResponseBare, ResponseJson, Response>(
 			opts.encoding,
 			buffer,
 			opts.responseVersionedDataHandler,
+			opts.responseZodSchema,
+			opts.responseFromJson,
+			opts.responseFromBare,
 		);
 	} catch (error) {
 		throw new HttpRequestError(`Failed to parse response: ${error}`, {

package/src/common/actor-router-consts.ts

@@ -1,8 +1,10 @@
 // NOTE: This is in a separate file from the router since it needs to be shared between the client & the server. If this was in the router file, the client would end up importing the *entire* actor router and tree shaking would not work.
 
 // MARK: Paths
-export const PATH_CONNECT_WEBSOCKET = "/connect/websocket";
-export const PATH_RAW_WEBSOCKET_PREFIX = "/raw/websocket/";
+export const PATH_CONNECT = "/connect";
+export const PATH_WEBSOCKET_BASE = "/websocket";
+export const PATH_WEBSOCKET_PREFIX = "/websocket/";
+export const PATH_INSPECTOR_CONNECT = "/inspector/connect";
 
 // MARK: Headers
 export const HEADER_ACTOR_QUERY = "x-rivet-query";
@@ -14,10 +16,6 @@ export const HEADER_CONN_PARAMS = "x-rivet-conn-params";
 
 export const HEADER_ACTOR_ID = "x-rivet-actor";
 
-export const HEADER_CONN_ID = "x-rivet-conn";
-
-export const HEADER_CONN_TOKEN = "x-rivet-conn-token";
-
 export const HEADER_RIVET_TOKEN = "x-rivet-token";
 
 // MARK: Manager Gateway Headers
@@ -32,13 +30,10 @@ export const WS_PROTOCOL_TARGET = "rivet_target.";
 export const WS_PROTOCOL_ACTOR = "rivet_actor.";
 export const WS_PROTOCOL_ENCODING = "rivet_encoding.";
 export const WS_PROTOCOL_CONN_PARAMS = "rivet_conn_params.";
-export const WS_PROTOCOL_CONN_ID = "rivet_conn.";
-export const WS_PROTOCOL_CONN_TOKEN = "rivet_conn_token.";
 export const WS_PROTOCOL_TOKEN = "rivet_token.";
 
 // MARK: WebSocket Inline Test Protocol Prefixes
-export const WS_PROTOCOL_TRANSPORT = "test_transport.";
-export const WS_PROTOCOL_PATH = "test_path.";
+export const WS_TEST_PROTOCOL_PATH = "test_path.";
 
 /**
  * Headers that publics can send from public clients.
@@ -52,8 +47,6 @@ export const ALLOWED_PUBLIC_HEADERS = [
 	HEADER_ENCODING,
 	HEADER_CONN_PARAMS,
 	HEADER_ACTOR_ID,
-	HEADER_CONN_ID,
-	HEADER_CONN_TOKEN,
 	HEADER_RIVET_TARGET,
 	HEADER_RIVET_ACTOR,
 	HEADER_RIVET_NAMESPACE,

package/src/common/{inline-websocket-adapter2.ts → inline-websocket-adapter.ts}

@@ -1,4 +1,5 @@
 import { WSContext } from "hono/ws";
+import type { UpgradeWebSocketArgs } from "@/actor/router-websocket-endpoints";
 import type {
 	RivetCloseEvent,
 	RivetEvent,
@@ -11,20 +12,11 @@ export function logger() {
 	return getLogger("fake-event-source2");
 }
 
-// TODO: Merge with ConnectWebSocketOutput interface
-export interface UpgradeWebSocketArgs {
-	onOpen: (event: any, ws: WSContext) => void;
-	onMessage: (event: any, ws: WSContext) => void;
-	onClose: (event: any, ws: WSContext) => void;
-	onError: (error: any, ws: WSContext) => void;
-}
-
-// TODO: Remove `2` suffix
 /**
  * InlineWebSocketAdapter implements a WebSocket-like interface
  * that connects to a UpgradeWebSocketArgs handler
  */
-export class InlineWebSocketAdapter2 implements UniversalWebSocket {
+export class InlineWebSocketAdapter implements UniversalWebSocket {
 	// WebSocket readyState values
 	readonly CONNECTING = 0 as const;
 	readonly OPEN = 1 as const;
@@ -36,14 +28,8 @@ export class InlineWebSocketAdapter2 implements UniversalWebSocket {
 	#wsContext: WSContext;
 	#readyState: 0 | 1 | 2 | 3 = 0; // Start in CONNECTING state
 	#queuedMessages: Array<string | ArrayBuffer | Uint8Array> = [];
-	// Event buffering is needed since events can be fired
-	// before JavaScript has a chance to add event listeners (e.g. within the same tick)
-	#bufferedEvents: Array<{
-		type: string;
-		event: any;
-	}> = [];
-
-	// Event listeners with buffering
+
+	// Event listeners
 	#eventListeners: Map<string, ((ev: any) => void)[]> = new Map();
 
 	constructor(handler: UpgradeWebSocketArgs) {
@@ -65,7 +51,11 @@ export class InlineWebSocketAdapter2 implements UniversalWebSocket {
 		});
 
 		// Initialize the connection
-		this.#initialize();
+		//
+		// Defer initialization to allow event listeners to be attached first
+		setTimeout(() => {
+			this.#initialize();
+		}, 0);
 	}
 
 	get readyState(): 0 | 1 | 2 | 3 {
@@ -99,19 +89,28 @@ export class InlineWebSocketAdapter2 implements UniversalWebSocket {
 	send(data: string | ArrayBufferLike | Blob | ArrayBufferView): void {
 		logger().debug({ msg: "send called", readyState: this.readyState });
 
-		if (this.readyState !== this.OPEN) {
-			const error = new Error("WebSocket is not open");
-			logger().warn({
-				msg: "cannot send message, websocket not open",
+		// Handle different ready states
+		if (this.readyState === this.CONNECTING) {
+			// Throw InvalidStateError if still connecting
+			throw new DOMException(
+				"WebSocket is still in CONNECTING state",
+				"InvalidStateError",
+			);
+		}
+
+		if (
+			this.readyState === this.CLOSING ||
+			this.readyState === this.CLOSED
+		) {
+			// Silently ignore if closing or closed
+			logger().debug({
+				msg: "ignoring send, websocket is closing or closed",
 				readyState: this.readyState,
-				dataType: typeof data,
-				dataLength: typeof data === "string" ? data.length : "binary",
-				error,
 			});
-			this.#fireError(error);
 			return;
 		}
 
+		// Must be OPEN at this point
 		this.#handler.onMessage({ data }, this.#wsContext);
 	}
 
@@ -284,9 +283,6 @@ export class InlineWebSocketAdapter2 implements UniversalWebSocket {
 			this.#eventListeners.set(type, []);
 		}
 		this.#eventListeners.get(type)!.push(listener);
-
-		// Flush any buffered events for this type
-		this.#flushBufferedEvents(type);
 	}
 
 	removeEventListener(type: string, listener: (ev: any) => void): void {
@@ -315,11 +311,6 @@ export class InlineWebSocketAdapter2 implements UniversalWebSocket {
 					});
 				}
 			}
-		} else {
-			logger().debug({
-				msg: `no ${type} listeners registered, buffering event`,
-			});
-			this.#bufferedEvents.push({ type, event });
 		}
 
 		// Also check for on* properties
@@ -380,19 +371,6 @@ export class InlineWebSocketAdapter2 implements UniversalWebSocket {
 		return true;
 	}
 
-	#flushBufferedEvents(type: string): void {
-		const eventsToFlush = this.#bufferedEvents.filter(
-			(buffered) => buffered.type === type,
-		);
-		this.#bufferedEvents = this.#bufferedEvents.filter(
-			(buffered) => buffered.type !== type,
-		);
-
-		for (const { event } of eventsToFlush) {
-			this.#dispatchEvent(type, event);
-		}
-	}
-
 	#fireOpen(): void {
 		try {
 			// Create an Event-like object since Event constructor may not be available

package/src/common/log.ts

@@ -147,7 +147,7 @@ export async function configureDefaultLogger(
 		},
 		hooks: {
 			logMethod(inputArgs, method, level) {
-				// TODO: This is a hack to not implement our own transport target. We can get better perf if we have our own transport target.
+				// TODO: This is a hack to not implement our own Pino transport target. We can get better perf if we have our own transport target.
 
 				const levelMap: Record<number, string> = {
 					10: "trace",

package/src/common/router.ts

@@ -7,9 +7,12 @@ import {
 } from "@/actor/router-endpoints";
 import { buildActorNames, type RegistryConfig } from "@/registry/config";
 import type { RunnerConfig } from "@/registry/run-config";
-import { getEndpoint } from "@/remote-manager-driver/api-utils";
-import { HttpResponseError } from "@/schemas/client-protocol/mod";
+import type * as protocol from "@/schemas/client-protocol/mod";
 import { HTTP_RESPONSE_ERROR_VERSIONED } from "@/schemas/client-protocol/versioned";
+import {
+	type HttpResponseError as HttpResponseErrorJson,
+	HttpResponseErrorSchema,
+} from "@/schemas/client-protocol-zod/mod";
 import { encodingIsBinary, serializeWithEncoding } from "@/serde";
 import { bufferToArrayBuffer, getEnvUniversal, VERSION } from "@/utils";
 import { getLogger, type Logger } from "./log";
@@ -68,18 +71,28 @@ export function handleRouteError(error: unknown, c: HonoContext) {
 		encoding = "json";
 	}
 
+	const errorData = { group, code, message, metadata };
 	const output = serializeWithEncoding(
 		encoding,
-		{
-			group,
-			code,
-			message,
-			// TODO: Cannot serialize non-binary meta since it requires ArrayBuffer atm
-			metadata: encodingIsBinary(encoding)
-				? bufferToArrayBuffer(cbor.encode(metadata))
-				: null,
-		},
+		errorData,
 		HTTP_RESPONSE_ERROR_VERSIONED,
+		HttpResponseErrorSchema,
+		// JSON: metadata is the raw value (will be serialized by jsonStringifyCompat)
+		(value): HttpResponseErrorJson => ({
+			group: value.group,
+			code: value.code,
+			message: value.message,
+			metadata: value.metadata,
+		}),
+		// BARE/CBOR: metadata needs to be CBOR-encoded to ArrayBuffer
+		(value): protocol.HttpResponseError => ({
+			group: value.group,
+			code: value.code,
+			message: value.message,
+			metadata: value.metadata
+				? bufferToArrayBuffer(cbor.encode(value.metadata))
+				: null,
+		}),
 	);
 
 	// TODO: Remove any
@@ -125,12 +138,10 @@ export function handleMetadataRequest(
 					: { normal: {} },
 		},
 		actorNames: buildActorNames(registryConfig),
-		// Do not return client endpoint if default server disabled
-		clientEndpoint:
-			runConfig.overrideServerAddress ??
-			(runConfig.disableDefaultServer
-				? undefined
-				: getEndpoint(runConfig)),
+		// If server address is changed, return a different client endpoint.
+		// Otherwise, return null indicating the client should use the current
+		// endpoint it's already configured with.
+		clientEndpoint: runConfig.overrideServerAddress,
 	};
 
 	return c.json(response);

package/src/common/utils.ts

@@ -245,6 +245,7 @@ export function deconstructError(
 				group,
 				code,
 				message,
+				stack: (error as Error)?.stack,
 				...EXTRA_ERROR_LOG,
 				...extraLog,
 			});
@@ -260,6 +261,7 @@ export function deconstructError(
 				group,
 				code,
 				message,
+				stack: (error as Error)?.stack,
 				...EXTRA_ERROR_LOG,
 				...extraLog,
 			});

package/src/driver-helpers/mod.ts

@@ -1,27 +1,23 @@
 export type { ActorDriver } from "@/actor/driver";
-export type { ActorInstance, AnyActorInstance } from "@/actor/instance";
+export type { ActorInstance, AnyActorInstance } from "@/actor/instance/mod";
 export { generateRandomString } from "@/actor/utils";
 export {
 	ALLOWED_PUBLIC_HEADERS,
 	HEADER_ACTOR_ID,
 	HEADER_ACTOR_QUERY,
-	HEADER_CONN_ID,
 	HEADER_CONN_PARAMS,
-	HEADER_CONN_TOKEN,
 	HEADER_ENCODING,
 	HEADER_RIVET_ACTOR,
 	HEADER_RIVET_TARGET,
-	PATH_CONNECT_WEBSOCKET,
-	PATH_RAW_WEBSOCKET_PREFIX,
+	PATH_CONNECT,
+	PATH_WEBSOCKET_BASE,
+	PATH_WEBSOCKET_PREFIX,
 	WS_PROTOCOL_ACTOR,
-	WS_PROTOCOL_CONN_ID,
 	WS_PROTOCOL_CONN_PARAMS,
-	WS_PROTOCOL_CONN_TOKEN,
 	WS_PROTOCOL_ENCODING,
-	WS_PROTOCOL_PATH,
 	WS_PROTOCOL_STANDARD,
 	WS_PROTOCOL_TARGET,
-	WS_PROTOCOL_TRANSPORT,
+	WS_TEST_PROTOCOL_PATH as WS_PROTOCOL_PATH,
 } from "@/common/actor-router-consts";
 export type {
 	ActorOutput,
@@ -29,6 +25,7 @@ export type {
 	GetForIdInput,
 	GetOrCreateWithKeyInput,
 	GetWithKeyInput,
+	ListActorsInput,
 	ManagerDisplayInformation,
 	ManagerDriver,
 } from "@/manager/driver";
@@ -36,4 +33,4 @@ export {
 	DriverConfigSchema,
 	RunnerConfigSchema as RunConfigSchema,
 } from "@/registry/run-config";
-export { serializeEmptyPersistData } from "./utils";
+export { getInitialActorKvState } from "./utils";

package/src/driver-helpers/utils.ts

@@ -1,21 +1,30 @@
 import * as cbor from "cbor-x";
-import type * as schema from "@/schemas/actor-persist/mod";
-import { PERSISTED_ACTOR_VERSIONED } from "@/schemas/actor-persist/versioned";
+import { KEYS } from "@/actor/instance/kv";
+import type * as persistSchema from "@/schemas/actor-persist/mod";
+import { ACTOR_VERSIONED } from "@/schemas/actor-persist/versioned";
 import { bufferToArrayBuffer } from "@/utils";
+import type { ActorDriver } from "./mod";
 
-export function serializeEmptyPersistData(
-	input: unknown | undefined,
-): Uint8Array {
-	const persistData: schema.PersistedActor = {
+function serializeEmptyPersistData(input: unknown | undefined): Uint8Array {
+	const persistData: persistSchema.Actor = {
 		input:
 			input !== undefined
 				? bufferToArrayBuffer(cbor.encode(input))
 				: null,
 		hasInitialized: false,
 		state: bufferToArrayBuffer(cbor.encode(undefined)),
-		connections: [],
 		scheduledEvents: [],
-		hibernatableWebSocket: [],
 	};
-	return PERSISTED_ACTOR_VERSIONED.serializeWithEmbeddedVersion(persistData);
+	return ACTOR_VERSIONED.serializeWithEmbeddedVersion(persistData);
+}
+
+/**
+ * Returns the initial KV state for a new actor. This is ued by the drivers to
+ * write the initial state in to KV storage before starting the actor.
+ */
+export function getInitialActorKvState(
+	input: unknown | undefined,
+): [Uint8Array, Uint8Array][] {
+	const persistData = serializeEmptyPersistData(input);
+	return [[KEYS.PERSIST_DATA, persistData]];
 }