fix(rivetkit): fix race condition with websocket open events

engine/packages/api-peer/Cargo.toml

@@ -8,6 +8,7 @@ license.workspace = true
 [dependencies]
 anyhow.workspace = true
 axum.workspace = true
+base64.workspace = true
 gas.workspace = true
 epoxy.workspace = true
 futures-util.workspace = true
@@ -27,6 +28,7 @@ tokio.workspace = true
 tracing.workspace = true
 namespace.workspace = true
 pegboard.workspace = true
+pegboard-actor-kv.workspace = true
 universalpubsub.workspace = true
 uuid.workspace = true
 utoipa.workspace = true

engine/packages/api-peer/src/actors/kv_get.rs

@@ -0,0 +1,67 @@
+use anyhow::*;
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use pegboard_actor_kv as actor_kv;
+use rivet_api_builder::ApiCtx;
+use rivet_util::Id;
+use serde::{Deserialize, Serialize};
+use utoipa::ToSchema;
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct KvGetPath {
+	pub actor_id: Id,
+	pub key: String,
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(deny_unknown_fields)]
+pub struct KvGetQuery {}
+
+#[derive(Serialize, ToSchema)]
+#[schema(as = ActorsKvGetResponse)]
+pub struct KvGetResponse {
+	/// Value encoded in base 64.
+	pub value: String,
+	pub update_ts: i64,
+}
+
+#[utoipa::path(
+	get,
+	operation_id = "actors_kv_get",
+	path = "/actors/{actor_id}/kv/keys/{key}",
+	params(
+		("actor_id" = Id, Path),
+		("key" = String, Path),
+	),
+	responses(
+		(status = 200, body = KvGetResponse),
+	),
+)]
+#[tracing::instrument(skip_all)]
+pub async fn kv_get(ctx: ApiCtx, path: KvGetPath, _query: KvGetQuery) -> Result<KvGetResponse> {
+	// Decode base64 key
+	let key_bytes = BASE64_STANDARD
+		.decode(&path.key)
+		.context("failed to decode base64 key")?;
+
+	// Get the KV value
+	let udb = ctx.pools().udb()?;
+	let (keys, values, metadata) =
+		actor_kv::get(&*udb, path.actor_id, vec![key_bytes.clone()]).await?;
+
+	// Check if key was found
+	if keys.is_empty() {
+		return Err(pegboard::errors::Actor::KvKeyNotFound.build());
+	}
+
+	// Encode value as base64
+	let value_base64 = BASE64_STANDARD.encode(&values[0]);
+
+	Ok(KvGetResponse {
+		value: value_base64,
+		// NOTE: Intentionally uses different name in public API. `create_ts` is actually
+		// `update_ts`.
+		update_ts: metadata[0].create_ts,
+	})
+}

engine/packages/api-peer/src/actors/mod.rs

@@ -1,4 +1,5 @@
 pub mod create;
 pub mod delete;
+pub mod kv_get;
 pub mod list;
 pub mod list_names;

engine/packages/api-peer/src/router.rs

@@ -26,6 +26,10 @@ pub async fn router(
 			.route("/actors", post(actors::create::create))
 			.route("/actors/{actor_id}", delete(actors::delete::delete))
 			.route("/actors/names", get(actors::list_names::list_names))
+			.route(
+				"/actors/{actor_id}/kv/keys/{key}",
+				get(actors::kv_get::kv_get),
+			)
 			// MARK: Runners
 			.route("/runners", get(runners::list))
 			.route("/runners/names", get(runners::list_names))

engine/packages/api-public/Cargo.toml

@@ -30,6 +30,7 @@ serde.workspace = true
 tokio.workspace = true
 tower-http.workspace = true
 tracing.workspace = true
+urlencoding.workspace = true
 utoipa.workspace = true
 
 [build-dependencies]

engine/packages/api-public/src/actors/kv_get.rs

@@ -0,0 +1,75 @@
+use anyhow::Result;
+use axum::response::{IntoResponse, Response};
+use rivet_api_builder::{
+	ApiError,
+	extract::{Extension, Path},
+};
+use rivet_api_util::request_remote_datacenter_raw;
+use rivet_util::Id;
+use serde::{Deserialize, Serialize};
+use utoipa::ToSchema;
+
+use crate::ctx::ApiCtx;
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct KvGetPath {
+	pub actor_id: Id,
+	pub key: String,
+}
+
+#[derive(Serialize, ToSchema)]
+#[schema(as = ActorsKvGetResponse)]
+pub struct KvGetResponse {
+	pub value: String,
+	pub update_ts: i64,
+}
+
+#[utoipa::path(
+	get,
+	operation_id = "actors_kv_get",
+	path = "/actors/{actor_id}/kv/keys/{key}",
+	params(
+		("actor_id" = Id, Path),
+		("key" = String, Path),
+	),
+	responses(
+		(status = 200, body = KvGetResponse),
+	),
+	security(("bearer_auth" = [])),
+)]
+#[tracing::instrument(skip_all)]
+pub async fn kv_get(Extension(ctx): Extension<ApiCtx>, Path(path): Path<KvGetPath>) -> Response {
+	match kv_get_inner(ctx, path).await {
+		Ok(response) => response,
+		Err(err) => ApiError::from(err).into_response(),
+	}
+}
+
+#[tracing::instrument(skip_all)]
+async fn kv_get_inner(ctx: ApiCtx, path: KvGetPath) -> Result<Response> {
+	use axum::Json;
+
+	ctx.auth().await?;
+
+	if path.actor_id.label() == ctx.config().dc_label() {
+		let peer_path = rivet_api_peer::actors::kv_get::KvGetPath {
+			actor_id: path.actor_id,
+			key: path.key,
+		};
+		let peer_query = rivet_api_peer::actors::kv_get::KvGetQuery {};
+		let res = rivet_api_peer::actors::kv_get::kv_get(ctx.into(), peer_path, peer_query).await?;
+
+		Ok(Json(res).into_response())
+	} else {
+		request_remote_datacenter_raw(
+			&ctx,
+			path.actor_id.label(),
+			&format!("/actors/{}/kv/keys/{}", path.actor_id, urlencoding::encode(&path.key)),
+			axum::http::Method::GET,
+			Option::<&()>::None,
+			Option::<&()>::None,
+		)
+		.await
+	}
+}

engine/packages/api-public/src/actors/mod.rs

@@ -1,6 +1,7 @@
 pub mod create;
 pub mod delete;
 pub mod get_or_create;
+pub mod kv_get;
 pub mod list;
 pub mod list_names;
 pub mod utils;

engine/packages/api-public/src/router.rs

@@ -18,6 +18,7 @@ use crate::{actors, ctx, datacenters, health, metadata, namespaces, runner_confi
 		actors::delete::delete,
 		actors::list_names::list_names,
 		actors::get_or_create::get_or_create,
+		actors::kv_get::kv_get,
 		runners::list,
 		runners::list_names,
 		namespaces::list,
@@ -88,6 +89,10 @@ pub async fn router(
 				"/actors/names",
 				axum::routing::get(actors::list_names::list_names),
 			)
+			.route(
+				"/actors/{actor_id}/kv/keys/{key}",
+				axum::routing::get(actors::kv_get::kv_get),
+			)
 			// MARK: Runners
 			.route("/runners", axum::routing::get(runners::list))
 			.route("/runners/names", axum::routing::get(runners::list_names))

engine/packages/pegboard/src/errors.rs

@@ -63,6 +63,9 @@ pub enum Actor {
 		namespace: String,
 		runner_name: String,
 	},
+
+	#[error("kv_key_not_found", "The KV key does not exist for this actor.")]
+	KvKeyNotFound,
 }
 
 #[derive(RivetError, Debug, Clone, Deserialize, Serialize)]

rivetkit-typescript/packages/cloudflare-workers/src/actor-driver.ts

@@ -6,10 +6,10 @@ import type {
 } from "rivetkit";
 import { lookupInRegistry } from "rivetkit";
 import type { Client } from "rivetkit/client";
-import {
-	type ActorDriver,
-	type AnyActorInstance,
-	type ManagerDriver,
+import type {
+	ActorDriver,
+	AnyActorInstance,
+	ManagerDriver,
 } from "rivetkit/driver-helpers";
 import { promiseWithResolvers } from "rivetkit/utils";
 import { KEYS } from "./actor-handler-do";
@@ -239,7 +239,6 @@ export class CloudflareActorsActorDriver implements ActorDriver {
 		// Persist data key
 		return Uint8Array.from([1]);
 	}
-
 }
 
 export function createCloudflareActorsActorDriverBuilder(

rivetkit-typescript/packages/cloudflare-workers/src/manager-driver.ts

@@ -7,6 +7,7 @@ import {
 	type GetOrCreateWithKeyInput,
 	type GetWithKeyInput,
 	generateRandomString,
+	type ListActorsInput,
 	type ManagerDisplayInformation,
 	type ManagerDriver,
 	WS_PROTOCOL_ACTOR,
@@ -348,6 +349,14 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 		};
 	}
 
+	async listActors({ c, name }: ListActorsInput): Promise<ActorOutput[]> {
+		logger().warn({
+			msg: "listActors not fully implemented for Cloudflare Workers",
+			name,
+		});
+		return [];
+	}
+
 	// Helper method to build actor output from an ID
 	async #buildActorOutput(
 		c: any,

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-params.ts

@@ -2,7 +2,7 @@ import { actor } from "rivetkit";
 
 export const counterWithParams = actor({
 	state: { count: 0, initializers: [] as string[] },
-	createConnState: (c, opts, params: { name?: string }) => {
+	createConnState: (c, params: { name?: string }) => {
 		return {
 			name: params.name || "anonymous",
 		};

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-state.ts

@@ -16,7 +16,6 @@ export const connStateActor = actor({
 	// Define connection state
 	createConnState: (
 		c,
-		opts,
 		params: { username?: string; role?: string; noCount?: boolean },
 	): ConnState => {
 		return {