@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

package/packages/api-public/src/runner_configs/utils.rs

@@ -1,148 +1,50 @@
-use std::{collections::HashMap, time::Duration};
+use anyhow::Result;
+use std::collections::HashMap;
 
-use anyhow::anyhow;
 use gas::prelude::*;
-use reqwest::header::{HeaderMap as ReqwestHeaderMap, HeaderName, HeaderValue};
-use serde::{Deserialize, Serialize};
-use utoipa::ToSchema;
 
 use crate::ctx::ApiCtx;
 
-const RESPONSE_BODY_MAX_LEN: usize = 1024;
-
-#[derive(Deserialize, Serialize, ToSchema, Clone, Debug, PartialEq, Eq)]
-#[serde(rename_all = "snake_case")]
-#[schema(as = RunnerConfigsServerlessMetadataError)]
-pub enum ServerlessMetadataError {
-	InvalidRequest {},
-	RequestFailed {},
-	RequestTimedOut {},
-	NonSuccessStatus { status_code: u16, body: String },
-	InvalidResponseJson { body: String },
-	InvalidResponseSchema { runtime: String, version: String },
-}
+// Re-export types from pegboard for API schema
+pub use pegboard::ops::serverless_metadata::fetch::ServerlessMetadataError;
 
+/// Serverless metadata returned from a runner.
 #[derive(Debug, Clone)]
 pub struct ServerlessMetadata {
 	pub runtime: String,
 	pub version: String,
 	pub actor_names: HashMap<String, serde_json::Value>,
+	pub envoy_version: Option<u32>,
 }
 
-#[derive(Deserialize)]
-struct ServerlessMetadataPayload {
-	runtime: String,
-	version: String,
-	#[serde(rename = "actorNames", default)]
-	actor_names: HashMap<String, serde_json::Value>,
-}
-
-fn truncate_response_body(body: &str) -> String {
-	let mut chars = body.chars();
-	let mut truncated: String = chars.by_ref().take(RESPONSE_BODY_MAX_LEN).collect();
-	if chars.next().is_some() {
-		truncated.push_str("...[truncated]");
+impl From<pegboard::ops::serverless_metadata::fetch::Output> for ServerlessMetadata {
+	fn from(output: pegboard::ops::serverless_metadata::fetch::Output) -> Self {
+		ServerlessMetadata {
+			runtime: output.runtime,
+			version: output.version,
+			actor_names: output
+				.actor_names
+				.into_iter()
+				.map(|a| (a.name, serde_json::Value::Object(a.metadata)))
+				.collect(),
+			envoy_version: output.envoy_version,
+		}
 	}
-
-	truncated
 }
 
 /// Fetches metadata from a serverless runner at the given URL.
 ///
 /// Returns metadata including runtime, version, and actor names if available.
 #[tracing::instrument(skip_all)]
-pub async fn fetch_serverless_runner_metadata(
+pub async fn fetch_serverless_metadata(
+	ctx: &ApiCtx,
 	url: String,
 	headers: HashMap<String, String>,
-) -> Result<ServerlessMetadata, ServerlessMetadataError> {
-	tracing::debug!(?url, "fetching serverless runner metadata");
-
-	let trimmed_url = url.trim();
-	if trimmed_url.is_empty() {
-		return Err(ServerlessMetadataError::InvalidRequest {});
-	}
-
-	let metadata_url = format!("{}/metadata", trimmed_url.trim_end_matches('/'));
-
-	if reqwest::Url::parse(&metadata_url).is_err() {
-		return Err(ServerlessMetadataError::InvalidRequest {});
-	}
-
-	let mut header_map = ReqwestHeaderMap::new();
-	for (name, value) in headers {
-		let header_name = HeaderName::from_bytes(name.trim().as_bytes())
-			.map_err(|_| ServerlessMetadataError::InvalidRequest {})?;
-
-		let header_value = HeaderValue::from_str(value.trim())
-			.map_err(|_| ServerlessMetadataError::InvalidRequest {})?;
-
-		header_map.insert(header_name, header_value);
-	}
-
-	let client = rivet_pools::reqwest::client()
+) -> std::result::Result<ServerlessMetadata, ServerlessMetadataError> {
+	ctx.op(pegboard::ops::serverless_metadata::fetch::Input { url, headers })
 		.await
-		.map_err(|_| ServerlessMetadataError::RequestFailed {})?;
-
-	tracing::debug!("sending metadata request");
-	let response = client
-		.get(&metadata_url)
-		.headers(header_map)
-		.timeout(Duration::from_secs(10))
-		.send()
-		.custom_instrument(tracing::info_span!("fetch_metadata_request"))
-		.await
-		.map_err(|err| {
-			if err.is_timeout() {
-				ServerlessMetadataError::RequestTimedOut {}
-			} else {
-				ServerlessMetadataError::RequestFailed {}
-			}
-		})?;
-
-	let status = response.status();
-	tracing::debug!(?status, "received metadata response");
-	let body_raw = response
-		.text()
-		.await
-		.unwrap_or_else(|_| String::from("<failed to read body>"));
-	let body_for_user = truncate_response_body(&body_raw);
-
-	if !status.is_success() {
-		return Err(ServerlessMetadataError::NonSuccessStatus {
-			status_code: status.as_u16(),
-			body: body_for_user,
-		});
-	}
-
-	let payload = serde_json::from_str::<ServerlessMetadataPayload>(&body_raw).map_err(|_| {
-		ServerlessMetadataError::InvalidResponseJson {
-			body: body_for_user,
-		}
-	})?;
-
-	let ServerlessMetadataPayload {
-		runtime,
-		version,
-		actor_names,
-	} = payload;
-
-	tracing::debug!(
-		?runtime,
-		?version,
-		actor_names_count = actor_names.len(),
-		"parsed metadata payload"
-	);
-
-	let trimmed_version = version.trim();
-	if runtime != "rivetkit" || trimmed_version.is_empty() {
-		return Err(ServerlessMetadataError::InvalidResponseSchema { runtime, version });
-	}
-
-	Ok(ServerlessMetadata {
-		runtime,
-		version: trimmed_version.to_owned(),
-		actor_names,
-	})
+		.map_err(|_| ServerlessMetadataError::RequestFailed {})?
+		.map(ServerlessMetadata::from)
 }
 
 /// Fetches metadata from the given URL and populates actor names in the database.
@@ -153,65 +55,23 @@ pub async fn refresh_runner_config_metadata(
 	runner_name: String,
 	url: String,
 	headers: HashMap<String, String>,
-) -> anyhow::Result<()> {
+) -> Result<()> {
 	tracing::debug!(
 		?namespace_id,
 		?runner_name,
 		"refreshing runner config metadata"
 	);
 
-	// Fetch metadata
-	let metadata = fetch_serverless_runner_metadata(url, headers)
-		.await
-		.map_err(|e| anyhow!("failed to fetch serverless runner metadata: {:?}", e))?;
-
-	if !metadata.actor_names.is_empty() {
-		tracing::debug!(
-			actor_names_count = metadata.actor_names.len(),
-			"storing actor names metadata"
-		);
-		// Convert actor names to the format needed for database operations
-		let actor_names: Vec<(String, serde_json::Map<String, serde_json::Value>)> = metadata
-			.actor_names
-			.into_iter()
-			.map(|(name, value)| {
-				if let serde_json::Value::Object(map) = value {
-					Ok((name, map))
-				} else {
-					Err(anyhow!(
-						"actor name '{}' metadata must be an object, got: {:?}",
-						name,
-						value
-					))
-				}
-			})
-			.collect::<anyhow::Result<Vec<_>>>()?;
-
-		// Store actor names metadata with the runner config
-		ctx.udb()?
-			.run(|tx| {
-				let actor_names = actor_names.clone();
-				async move {
-					let tx = tx.with_subspace(pegboard::keys::subspace());
-
-					// Write actor names
-					for (name, metadata) in actor_names {
-						tx.write(
-							&pegboard::keys::ns::ActorNameKey::new(namespace_id, name.clone()),
-							rivet_data::converted::ActorNameKeyData { metadata },
-						)?;
-					}
-
-					Ok(())
-				}
-			})
-			.custom_instrument(tracing::info_span!("runner_config_populate_actor_names_tx"))
-			.await?;
-
-		tracing::debug!("successfully stored actor names metadata");
-	} else {
-		tracing::debug!("no actor names to store");
-	}
+	ctx.op(pegboard::ops::runner_config::refresh_metadata::Input {
+		namespace_id,
+		runner_name,
+		url,
+		headers,
+	})
+	.await?
+	.map_err(|e| {
+		pegboard::errors::ServerlessRunnerPool::FailedToFetchMetadata { reason: e }.build()
+	})?;
 
 	Ok(())
 }

package/packages/api-public/src/runners.rs

@@ -1,24 +1,23 @@
 use anyhow::Result;
 use axum::response::{IntoResponse, Response};
+use indexmap::IndexSet;
 use rivet_api_builder::{
 	ApiError,
 	extract::{Extension, Json, Query},
 };
-use rivet_api_types::{pagination::Pagination, runners::list::*};
+use rivet_api_types::{pagination::Pagination, runners::list::*, runners::list_names::*};
 use rivet_api_util::fanout_to_datacenters;
-use serde::{Deserialize, Serialize};
-use utoipa::{IntoParams, ToSchema};
 
 use crate::ctx::ApiCtx;
 
 #[utoipa::path(
-    get,
+	get,
 	operation_id = "runners_list",
-    path = "/runners",
-    params(ListQuery),
-    responses(
-        (status = 200, body = ListResponse),
-    ),
+	path = "/runners",
+	params(ListQuery),
+	responses(
+		(status = 200, body = ListResponse),
+	),
 	security(("bearer_auth" = [])),
 )]
 #[tracing::instrument(skip_all)]
@@ -58,23 +57,6 @@ async fn list_inner(ctx: ApiCtx, query: ListQuery) -> Result<ListResponse> {
 	})
 }
 
-#[derive(Debug, Deserialize, Serialize, Clone, IntoParams)]
-#[serde(deny_unknown_fields)]
-#[into_params(parameter_in = Query)]
-pub struct ListNamesQuery {
-	pub namespace: String,
-	pub limit: Option<usize>,
-	pub cursor: Option<String>,
-}
-
-#[derive(Deserialize, Serialize, ToSchema)]
-#[serde(deny_unknown_fields)]
-#[schema(as = RunnersListNamesResponse)]
-pub struct ListNamesResponse {
-	pub names: Vec<String>,
-	pub pagination: Pagination,
-}
-
 /// ## Datacenter Round Trips
 ///
 /// 2 round trips:
@@ -106,39 +88,29 @@ async fn list_names_inner(ctx: ApiCtx, query: ListNamesQuery) -> Result<ListName
 	ctx.auth().await?;
 
 	// Prepare peer query for local handler
-	let peer_query = rivet_api_peer::runners::ListNamesQuery {
-		namespace: query.namespace.clone(),
-		limit: query.limit,
-		cursor: query.cursor.clone(),
-	};
+	let limit = query.limit.unwrap_or(100);
 
 	// Fanout to all datacenters
-	let mut all_names = fanout_to_datacenters::<
-		rivet_api_peer::runners::ListNamesResponse,
-		_,
-		_,
-		_,
-		_,
-		Vec<String>,
-	>(
+	let mut all_names = fanout_to_datacenters::<ListNamesResponse, _, _, _, _, IndexSet<String>>(
 		ctx.into(),
 		"/runners/names",
-		peer_query,
+		query,
 		|ctx, query| async move { rivet_api_peer::runners::list_names(ctx, (), query).await },
 		|_, res, agg| agg.extend(res.names),
 	)
-	.await?;
+	.await?
+	.into_iter()
+	// Apply limit
+	.take(limit)
+	.collect::<IndexSet<_>>();
 
 	// Sort by name for consistency
 	all_names.sort();
 
-	// Truncate to the requested limit
-	all_names.truncate(query.limit.unwrap_or(100));
-
 	let cursor = all_names.last().map(|x: &String| x.to_string());
 
 	Ok(ListNamesResponse {
-		names: all_names,
+		names: all_names.into_iter().collect(),
 		pagination: Pagination { cursor },
 	})
 }

package/packages/{dump-openapi → api-public-openapi-gen}/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "rivet-dump-openapi"
+name = "rivet-api-public-openapi-gen"
 version.workspace = true
 edition.workspace = true
 authors.workspace = true

package/packages/api-types/src/actors/create.rs

@@ -16,6 +16,7 @@ pub struct CreateRequest {
 	pub datacenter: Option<String>,
 	pub name: String,
 	pub key: Option<String>,
+	/// Arbitrary base64 encoded binary data.
 	pub input: Option<String>,
 	pub runner_name_selector: String,
 	pub crash_policy: rivet_types::actors::CrashPolicy,

package/packages/api-types/src/actors/delete.rs

@@ -0,0 +1,20 @@
+use rivet_util::Id;
+use serde::{Deserialize, Serialize};
+use utoipa::{IntoParams, ToSchema};
+
+#[derive(Debug, Deserialize, Serialize, IntoParams)]
+#[serde(deny_unknown_fields)]
+#[into_params(parameter_in = Query)]
+pub struct DeleteQuery {
+	pub namespace: String,
+}
+
+#[derive(Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct DeletePath {
+	pub actor_id: Id,
+}
+
+#[derive(Serialize, Deserialize, ToSchema)]
+#[schema(as = ActorsDeleteResponse)]
+pub struct DeleteResponse {}

package/packages/api-types/src/actors/get_or_create.rs

@@ -0,0 +1,30 @@
+use serde::{Deserialize, Serialize};
+use utoipa::{IntoParams, ToSchema};
+
+#[derive(Debug, Serialize, Deserialize, IntoParams)]
+#[serde(deny_unknown_fields)]
+#[into_params(parameter_in = Query)]
+pub struct GetOrCreateQuery {
+	pub namespace: String,
+}
+
+#[derive(Serialize, Deserialize, ToSchema)]
+#[serde(deny_unknown_fields)]
+#[schema(as = ActorsGetOrCreateRequest)]
+pub struct GetOrCreateRequest {
+	// Ignored in api-peer
+	pub datacenter: Option<String>,
+	pub name: String,
+	pub key: String,
+	pub input: Option<String>,
+	pub runner_name_selector: String,
+	pub crash_policy: rivet_types::actors::CrashPolicy,
+}
+
+#[derive(Serialize, Deserialize, ToSchema)]
+#[serde(deny_unknown_fields)]
+#[schema(as = ActorsGetOrCreateResponse)]
+pub struct GetOrCreateResponse {
+	pub actor: rivet_types::actors::Actor,
+	pub created: bool,
+}

package/packages/api-types/src/actors/kv_get.rs

@@ -0,0 +1,25 @@
+use gas::prelude::*;
+use serde::{Deserialize, Serialize};
+use utoipa::{IntoParams, ToSchema};
+
+#[derive(Debug, Deserialize, Serialize, IntoParams)]
+#[serde(deny_unknown_fields)]
+#[into_params(parameter_in = Query)]
+pub struct KvGetQuery {
+	pub namespace: String,
+}
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct KvGetPath {
+	pub actor_id: Id,
+	pub key: String,
+}
+
+#[derive(Serialize, ToSchema)]
+#[schema(as = ActorsKvGetResponse)]
+#[serde(deny_unknown_fields)]
+pub struct KvGetResponse {
+	pub value: String,
+	pub update_ts: i64,
+}

package/packages/api-types/src/actors/list.rs

@@ -1,6 +1,9 @@
+use gas::prelude::*;
 use serde::{Deserialize, Serialize};
 use utoipa::{IntoParams, ToSchema};
 
+use crate::pagination::Pagination;
+
 #[derive(Debug, Serialize, Deserialize, Clone, IntoParams, Default)]
 #[serde(deny_unknown_fields)]
 #[into_params(parameter_in = Query)]
@@ -8,7 +11,11 @@ pub struct ListQuery {
 	pub namespace: String,
 	pub name: Option<String>,
 	pub key: Option<String>,
+	/// Deprecated.
+	#[serde(default)]
 	pub actor_ids: Option<String>,
+	#[serde(default)]
+	pub actor_id: Vec<Id>,
 	pub include_destroyed: Option<bool>,
 	pub limit: Option<usize>,
 	pub cursor: Option<String>,
@@ -19,5 +26,5 @@ pub struct ListQuery {
 #[schema(as = ActorsListResponse)]
 pub struct ListResponse {
 	pub actors: Vec<rivet_types::actors::Actor>,
-	pub pagination: crate::pagination::Pagination,
+	pub pagination: Pagination,
 }

package/packages/api-types/src/actors/mod.rs

@@ -1,3 +1,8 @@
 pub mod create;
+pub mod delete;
+pub mod get_or_create;
+pub mod kv_get;
 pub mod list;
 pub mod list_names;
+pub mod reschedule;
+pub mod sleep;

package/packages/api-types/src/actors/reschedule.rs

@@ -0,0 +1,26 @@
+use gas::prelude::*;
+use serde::{Deserialize, Serialize};
+use utoipa::{IntoParams, ToSchema};
+
+#[derive(Debug, Deserialize, Serialize, IntoParams)]
+#[serde(deny_unknown_fields)]
+#[into_params(parameter_in = Query)]
+pub struct RescheduleQuery {
+	pub namespace: String,
+}
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct ReschedulePath {
+	pub actor_id: Id,
+}
+
+#[derive(Serialize, Deserialize, ToSchema)]
+#[serde(deny_unknown_fields)]
+#[schema(as = ActorsRescheduleRequest)]
+pub struct RescheduleRequest {}
+
+#[derive(Serialize, ToSchema)]
+#[schema(as = ActorsRescheduleResponse)]
+#[serde(deny_unknown_fields)]
+pub struct RescheduleResponse {}

package/packages/api-types/src/actors/sleep.rs

@@ -0,0 +1,26 @@
+use gas::prelude::*;
+use serde::{Deserialize, Serialize};
+use utoipa::{IntoParams, ToSchema};
+
+#[derive(Debug, Deserialize, Serialize, IntoParams)]
+#[serde(deny_unknown_fields)]
+#[into_params(parameter_in = Query)]
+pub struct SleepQuery {
+	pub namespace: String,
+}
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct SleepPath {
+	pub actor_id: Id,
+}
+
+#[derive(Serialize, Deserialize, ToSchema)]
+#[serde(deny_unknown_fields)]
+#[schema(as = ActorsSleepRequest)]
+pub struct SleepRequest {}
+
+#[derive(Serialize, ToSchema)]
+#[schema(as = ActorsSleepResponse)]
+#[serde(deny_unknown_fields)]
+pub struct SleepResponse {}

package/packages/api-types/src/datacenters/list.rs

@@ -1,9 +1,9 @@
-use serde::Serialize;
+use serde::{Deserialize, Serialize};
 use utoipa::ToSchema;
 
 use crate::pagination::Pagination;
 
-#[derive(Serialize, ToSchema)]
+#[derive(Serialize, Deserialize, ToSchema)]
 #[serde(deny_unknown_fields)]
 #[schema(as = DatacentersListResponse)]
 pub struct ListResponse {

package/packages/api-types/src/envoys/list.rs

@@ -0,0 +1,24 @@
+use serde::{Deserialize, Serialize};
+use utoipa::{IntoParams, ToSchema};
+
+use crate::pagination::Pagination;
+
+#[derive(Debug, Serialize, Deserialize, Clone, IntoParams)]
+#[serde(deny_unknown_fields)]
+#[into_params(parameter_in = Query)]
+pub struct ListQuery {
+	pub namespace: String,
+	pub name: Option<String>,
+	#[serde(default)]
+	pub envoy_key: Vec<String>,
+	pub limit: Option<usize>,
+	pub cursor: Option<String>,
+}
+
+#[derive(Serialize, Deserialize, ToSchema)]
+#[serde(deny_unknown_fields)]
+#[schema(as = EnvoysListResponse)]
+pub struct ListResponse {
+	pub envoys: Vec<rivet_types::envoys::Envoy>,
+	pub pagination: Pagination,
+}

package/packages/api-types/src/envoys/mod.rs

@@ -0,0 +1 @@
+pub mod list;

package/packages/api-types/src/lib.rs

@@ -1,5 +1,6 @@
 pub mod actors;
 pub mod datacenters;
+pub mod envoys;
 pub mod namespaces;
 pub mod pagination;
 pub mod runner_configs;

package/packages/api-types/src/namespaces/list.rs

@@ -1,3 +1,4 @@
+use gas::prelude::*;
 use serde::{Deserialize, Serialize};
 use utoipa::{IntoParams, ToSchema};
 
@@ -8,8 +9,11 @@ pub struct ListQuery {
 	pub limit: Option<usize>,
 	pub cursor: Option<String>,
 	pub name: Option<String>,
+	/// Deprecated.
 	#[serde(default)]
 	pub namespace_ids: Option<String>,
+	#[serde(default)]
+	pub namespace_id: Vec<Id>,
 }
 
 #[derive(Serialize, Deserialize, ToSchema)]

package/packages/api-types/src/namespaces/runner_configs.rs

@@ -9,6 +9,8 @@ pub struct RunnerConfig {
 	pub kind: RunnerConfigKind,
 	#[serde(default, skip_serializing_if = "Option::is_none")]
 	pub metadata: Option<serde_json::Value>,
+	#[serde(default = "default_drain_on_version_upgrade")]
+	pub drain_on_version_upgrade: bool,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
@@ -20,36 +22,55 @@ pub enum RunnerConfigKind {
 		headers: Option<HashMap<String, String>>,
 		/// Seconds.
 		request_lifespan: u32,
+		max_concurrent_actors: Option<u64>,
 		slots_per_runner: u32,
 		min_runners: Option<u32>,
 		max_runners: u32,
 		runners_margin: Option<u32>,
+		/// Milliseconds between metadata polling. If not set, uses the global default.
+		metadata_poll_interval: Option<u64>,
 	},
 }
 
+fn default_drain_on_version_upgrade() -> bool {
+	false
+}
+
 impl Into<rivet_types::runner_configs::RunnerConfig> for RunnerConfig {
 	fn into(self) -> rivet_types::runner_configs::RunnerConfig {
-		let RunnerConfig { kind, metadata } = self;
+		let RunnerConfig {
+			kind,
+			metadata,
+			drain_on_version_upgrade,
+		} = self;
 		let kind = match kind {
 			RunnerConfigKind::Normal {} => rivet_types::runner_configs::RunnerConfigKind::Normal {},
 			RunnerConfigKind::Serverless {
 				url,
 				headers,
 				request_lifespan,
+				max_concurrent_actors,
 				slots_per_runner,
 				min_runners,
 				max_runners,
 				runners_margin,
+				metadata_poll_interval,
 			} => rivet_types::runner_configs::RunnerConfigKind::Serverless {
 				url,
 				headers: headers.unwrap_or_default(),
 				request_lifespan,
+				max_concurrent_actors: max_concurrent_actors.unwrap_or(max_runners as u64),
 				slots_per_runner,
 				min_runners: min_runners.unwrap_or_default(),
 				max_runners,
 				runners_margin: runners_margin.unwrap_or_default(),
+				metadata_poll_interval,
 			},
 		};
-		rivet_types::runner_configs::RunnerConfig { kind, metadata }
+		rivet_types::runner_configs::RunnerConfig {
+			kind,
+			metadata,
+			drain_on_version_upgrade,
+		}
 	}
 }