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

package/packages/engine/Cargo.toml

@@ -11,6 +11,7 @@ path = "src/main.rs"
 
 [dependencies]
 anyhow.workspace = true
+base64.workspace = true
 chrono.workspace = true
 clap.workspace = true
 colored_json.workspace = true
@@ -21,8 +22,9 @@ include_dir.workspace = true
 indoc.workspace = true
 lz4_flex.workspace = true
 once_cell.workspace = true
+pegboard-outbound.workspace = true
 pegboard-runner.workspace = true
-pegboard-serverless.workspace = true
+pegboard.workspace = true
 reqwest.workspace = true
 rivet-api-peer.workspace = true
 rivet-bootstrap.workspace = true
@@ -37,6 +39,7 @@ rivet-service-manager.workspace = true
 rivet-telemetry.workspace = true
 rivet-term.workspace = true
 rivet-tracing-reconfigure.workspace = true
+rivet-types.workspace = true
 rivet-util.workspace = true
 rivet-workflow-worker.workspace = true
 rustyline.workspace = true
@@ -55,20 +58,26 @@ url.workspace = true
 uuid.workspace = true
 
 [dev-dependencies]
+async-trait.workspace = true
 axum.workspace = true
 base64.workspace = true
 chrono.workspace = true
 epoxy.workspace = true
 futures-util.workspace = true
 namespace.workspace = true
-pegboard.workspace = true
 portpicker.workspace = true
 rand.workspace = true
 rivet-api-public.workspace = true
+rivet-api-types.workspace = true
+rivet-envoy-protocol.workspace = true
 rivet-runner-protocol.workspace = true
+rivet-test-envoy.workspace = true
 rivet-test-deps.workspace = true
 rivet-util.workspace = true
 rstest.workspace = true
+serde_bare.workspace = true
+serde_html_form.workspace = true
 tokio-tungstenite.workspace = true
 tracing-subscriber.workspace = true
+urlencoding.workspace = true
 vbare.workspace = true

package/packages/engine/src/commands/db/mod.rs

@@ -22,10 +22,6 @@ pub enum SubCommand {
 pub enum DatabaseType {
 	#[clap(alias = "ch")]
 	Clickhouse,
-	#[clap(alias = "wfd")]
-	WorkflowData,
-	#[clap(alias = "wfi")]
-	WorkflowInternal,
 }
 
 impl SubCommand {
@@ -48,12 +44,6 @@ impl SubCommand {
 					DatabaseType::Clickhouse => {
 						crate::util::db::clickhouse_shell(config, shell_ctx).await?
 					}
-					DatabaseType::WorkflowData => {
-						crate::util::db::wf_sqlite_shell(config, shell_ctx, false).await?
-					}
-					DatabaseType::WorkflowInternal => {
-						crate::util::db::wf_sqlite_shell(config, shell_ctx, true).await?
-					}
 				}
 
 				Ok(())

package/packages/engine/src/commands/epoxy.rs

@@ -0,0 +1,395 @@
+use anyhow::{Context, Result, bail};
+use base64::Engine;
+use clap::Parser;
+use serde::{Deserialize, Serialize};
+
+use crate::util::udb::{SimpleTuple, SimpleTupleValue};
+
+#[derive(Parser)]
+pub enum SubCommand {
+	/// Get debug information for the local replica
+	ReplicaDebug {
+		/// API peer endpoint (defaults to topology peer_url)
+		#[clap(long)]
+		endpoint: Option<String>,
+	},
+	/// Get coordinator cluster state
+	CoordinatorState {
+		/// API peer endpoint (defaults to topology peer_url)
+		#[clap(long)]
+		endpoint: Option<String>,
+	},
+	/// Trigger coordinator to reconfigure all replicas
+	CoordinatorReconfigure {
+		/// API peer endpoint (defaults to topology peer_url)
+		#[clap(long)]
+		endpoint: Option<String>,
+	},
+	/// Set coordinator cluster state
+	CoordinatorSetState {
+		/// JSON file containing the cluster config
+		#[clap(long)]
+		config_file: String,
+		/// API peer endpoint (defaults to topology peer_url)
+		#[clap(long)]
+		endpoint: Option<String>,
+	},
+	/// Get debug information for a specific key on the local replica
+	KeyDebug {
+		/// Key to inspect using tuple format (same as UDB CLI).
+		key: String,
+		/// Interpret the key as base64-encoded instead of tuple format
+		#[clap(long)]
+		base64: bool,
+		/// API peer endpoint (defaults to topology peer_url)
+		#[clap(long)]
+		endpoint: Option<String>,
+	},
+	/// Get debug information for a specific key across all replicas (fanout)
+	KeyDebugFanout {
+		/// Key to inspect using tuple format (same as UDB CLI).
+		key: String,
+		/// Interpret the key as base64-encoded instead of tuple format
+		#[clap(long)]
+		base64: bool,
+		/// API peer endpoint (defaults to topology peer_url)
+		#[clap(long)]
+		endpoint: Option<String>,
+	},
+	/// Get value at a key from local replica only
+	GetLocal {
+		/// Key to get using tuple format (same as UDB CLI).
+		key: String,
+		/// Interpret the key as base64-encoded instead of tuple format
+		#[clap(long)]
+		base64: bool,
+		/// Optional type hint for value parsing (u64, i64, f64, uuid, id, str, bytes, json, raw)
+		#[clap(short = 't', long = "type")]
+		type_hint: Option<String>,
+		/// API peer endpoint (defaults to topology peer_url)
+		#[clap(long)]
+		endpoint: Option<String>,
+	},
+	/// Get value at a key with optimistic read (local + cache + fanout)
+	GetOptimistic {
+		/// Key to get using tuple format (same as UDB CLI).
+		key: String,
+		/// Interpret the key as base64-encoded instead of tuple format
+		#[clap(long)]
+		base64: bool,
+		/// Optional type hint for value parsing (u64, i64, f64, uuid, id, str, bytes, json, raw)
+		#[clap(short = 't', long = "type")]
+		type_hint: Option<String>,
+		/// API peer endpoint (defaults to topology peer_url)
+		#[clap(long)]
+		endpoint: Option<String>,
+	},
+	/// Set value at a key in UDB (same as UDB CLI set)
+	Set {
+		/// Key to set using tuple format (same as UDB CLI).
+		key: String,
+		/// Value to set, with optional type prefix (e.g. "u64:123")
+		value: String,
+		/// Interpret the key as base64-encoded instead of tuple format
+		#[clap(long)]
+		base64: bool,
+		/// Optional type hint for the value
+		#[clap(short = 't', long = "type")]
+		type_hint: Option<String>,
+		/// API peer endpoint (defaults to topology peer_url)
+		#[clap(long)]
+		endpoint: Option<String>,
+	},
+}
+
+impl SubCommand {
+	pub async fn execute(self, config: rivet_config::Config) -> Result<()> {
+		match self {
+			Self::ReplicaDebug { endpoint } => {
+				let endpoint = get_endpoint(&config, endpoint)?;
+				let url = format!("{}/epoxy/replica/debug", endpoint);
+
+				let response = make_get_request(&url).await?;
+				print_json(&response)?;
+
+				Ok(())
+			}
+			Self::CoordinatorState { endpoint } => {
+				let endpoint = get_endpoint(&config, endpoint)?;
+				let url = format!("{}/epoxy/coordinator/state", endpoint);
+
+				let response = make_get_request(&url).await?;
+				print_json(&response)?;
+
+				Ok(())
+			}
+			Self::CoordinatorReconfigure { endpoint } => {
+				let endpoint = get_endpoint(&config, endpoint)?;
+				let url = format!("{}/epoxy/coordinator/replica-reconfigure", endpoint);
+
+				let client = rivet_pools::reqwest::client().await?;
+				let response = client
+					.post(&url)
+					.json(&serde_json::json!({}))
+					.send()
+					.await
+					.context("failed to send request")?;
+
+				if response.status().is_success() {
+					println!("Reconfigure signal sent successfully");
+				} else {
+					let status = response.status();
+					let body = response.text().await.unwrap_or_default();
+					anyhow::bail!("Failed to trigger reconfigure: {} - {}", status, body);
+				}
+
+				Ok(())
+			}
+			Self::CoordinatorSetState {
+				config_file,
+				endpoint,
+			} => {
+				let endpoint = get_endpoint(&config, endpoint)?;
+				let url = format!("{}/epoxy/coordinator/state", endpoint);
+
+				// Read config file
+				let config_content = tokio::fs::read_to_string(&config_file)
+					.await
+					.context("failed to read config file")?;
+				let cluster_config: serde_json::Value = serde_json::from_str(&config_content)
+					.context("failed to parse config file as JSON")?;
+
+				let client = rivet_pools::reqwest::client().await?;
+				let response = client
+					.post(&url)
+					.json(&SetEpoxyStateRequest {
+						config: cluster_config,
+					})
+					.send()
+					.await
+					.context("failed to send request")?;
+
+				if response.status().is_success() {
+					println!("Coordinator state updated successfully");
+				} else {
+					let status = response.status();
+					let body = response.text().await.unwrap_or_default();
+					anyhow::bail!("Failed to set coordinator state: {} - {}", status, body);
+				}
+
+				Ok(())
+			}
+			Self::KeyDebug {
+				key,
+				base64: is_base64,
+				endpoint,
+			} => {
+				let endpoint = get_endpoint(&config, endpoint)?;
+				let encoded_key = encode_key(&key, is_base64)?;
+				let url = format!("{}/epoxy/replica/key/{}", endpoint, encoded_key);
+
+				let response = make_get_request(&url).await?;
+				print_json(&response)?;
+
+				Ok(())
+			}
+			Self::KeyDebugFanout {
+				key,
+				base64: is_base64,
+				endpoint,
+			} => {
+				let endpoint = get_endpoint(&config, endpoint)?;
+				let encoded_key = encode_key(&key, is_base64)?;
+				let url = format!("{}/epoxy/replica/key/{}/fanout", endpoint, encoded_key);
+
+				let response = make_get_request(&url).await?;
+				print_json(&response)?;
+
+				Ok(())
+			}
+			Self::GetLocal {
+				key,
+				base64: is_base64,
+				type_hint,
+				endpoint,
+			} => {
+				let endpoint = get_endpoint(&config, endpoint)?;
+				let encoded_key = encode_key(&key, is_base64)?;
+				let url = format!("{}/epoxy/replica/kv/{}/local", endpoint, encoded_key);
+
+				print_kv_response(&url, type_hint.as_deref()).await
+			}
+			Self::GetOptimistic {
+				key,
+				base64: is_base64,
+				type_hint,
+				endpoint,
+			} => {
+				let endpoint = get_endpoint(&config, endpoint)?;
+				let encoded_key = encode_key(&key, is_base64)?;
+				let url = format!("{}/epoxy/replica/kv/{}/optimistic", endpoint, encoded_key);
+
+				print_kv_response(&url, type_hint.as_deref()).await
+			}
+			Self::Set {
+				key,
+				value,
+				base64: is_base64,
+				type_hint,
+				endpoint,
+			} => {
+				let endpoint = get_endpoint(&config, endpoint)?;
+				let encoded_key = encode_key(&key, is_base64)?;
+				let url = format!("{}/epoxy/replica/kv/{}", endpoint, encoded_key);
+
+				// Parse value using SimpleTupleValue (same as udb CLI)
+				let parsed_value =
+					SimpleTupleValue::parse_str_with_type_hint(type_hint.as_deref(), &value)?;
+				let value_bytes = parsed_value.serialize()?;
+				let value_b64 = base64::engine::general_purpose::STANDARD.encode(&value_bytes);
+
+				let client = rivet_pools::reqwest::client().await?;
+				let response = client
+					.put(&url)
+					.json(&SetKvRequest {
+						value: Some(value_b64),
+					})
+					.send()
+					.await
+					.context("failed to send request")?;
+
+				if response.status().is_success() {
+					let result: SetKvResponse =
+						response.json().await.context("failed to parse response")?;
+					println!("{}", result.result);
+				} else {
+					let status = response.status();
+					let body = response.text().await.unwrap_or_default();
+					anyhow::bail!("Failed to set value: {} - {}", status, body);
+				}
+
+				Ok(())
+			}
+		}
+	}
+}
+
+#[derive(Serialize, Deserialize)]
+struct SetEpoxyStateRequest {
+	config: serde_json::Value,
+}
+
+#[derive(Serialize, Deserialize)]
+struct GetKvResponse {
+	exists: bool,
+	value: Option<String>,
+}
+
+#[derive(Serialize, Deserialize)]
+struct SetKvRequest {
+	value: Option<String>,
+}
+
+#[derive(Serialize, Deserialize)]
+struct SetKvResponse {
+	result: String,
+}
+
+fn get_endpoint(config: &rivet_config::Config, endpoint: Option<String>) -> Result<String> {
+	match endpoint {
+		Some(e) => Ok(e),
+		None => {
+			let topology = config
+				.topology
+				.as_ref()
+				.context("topology not configured")?;
+			let dc = topology.current_dc()?;
+			Ok(dc.peer_url.to_string().trim_end_matches('/').to_string())
+		}
+	}
+}
+
+/// Encodes a key for use in the API URL.
+///
+/// If `is_base64` is true, the key is assumed to already be base64-encoded.
+/// Otherwise, the key is parsed as a tuple path (e.g., "/rivet/epoxy/replica/1/key_instance/b:deadbeef")
+/// and tuple-packed before being base64-encoded.
+fn encode_key(key: &str, is_base64: bool) -> Result<String> {
+	use base64::Engine;
+
+	if is_base64 {
+		// Validate it's valid base64
+		base64::engine::general_purpose::STANDARD
+			.decode(key)
+			.context("invalid base64 key")?;
+		Ok(key.to_string())
+	} else {
+		// Parse as tuple path and encode
+		let (tuple, _relative, back_count) =
+			SimpleTuple::parse(key).context("failed to parse key as tuple path")?;
+
+		if back_count > 0 {
+			bail!("relative paths with '..' are not supported for key lookup");
+		}
+
+		// Tuple-pack the key
+		let packed = universaldb::tuple::pack(&tuple);
+
+		// Base64 encode for URL
+		Ok(base64::engine::general_purpose::STANDARD.encode(&packed))
+	}
+}
+
+async fn make_get_request<T: serde::de::DeserializeOwned>(url: &str) -> Result<T> {
+	let client = rivet_pools::reqwest::client().await?;
+	let response = client
+		.get(url)
+		.send()
+		.await
+		.context("failed to send request")?;
+
+	if response.status().is_success() {
+		let body = response
+			.json::<T>()
+			.await
+			.context("failed to parse response")?;
+		Ok(body)
+	} else {
+		let status = response.status();
+		let body = response.text().await.unwrap_or_default();
+		anyhow::bail!("Request failed: {} - {}", status, body);
+	}
+}
+
+fn print_json(value: &serde_json::Value) -> Result<()> {
+	let output = colored_json::to_colored_json_auto(value)?;
+	println!("{}", output);
+	Ok(())
+}
+
+async fn print_kv_response(url: &str, type_hint: Option<&str>) -> Result<()> {
+	let response: GetKvResponse = make_get_request(url).await?;
+
+	if !response.exists {
+		println!("key does not exist");
+	} else if let Some(value_b64) = response.value {
+		let value_bytes = base64::engine::general_purpose::STANDARD
+			.decode(&value_b64)
+			.context("invalid base64 value from server")?;
+
+		if type_hint == Some("raw") {
+			println!("{}", SimpleTupleValue::Unknown(value_bytes));
+		} else {
+			match SimpleTupleValue::deserialize(type_hint, &value_bytes) {
+				Ok(parsed) => {
+					let mut s = String::new();
+					parsed.write(&mut s, false).unwrap();
+					println!("{s}");
+				}
+				Err(err) => println!("error: {err:#}"),
+			}
+		}
+	}
+
+	Ok(())
+}

package/packages/engine/src/commands/mod.rs

@@ -1,7 +1,7 @@
 pub mod config;
 pub mod db;
+pub mod epoxy;
 pub mod start;
 pub mod tracing;
 pub mod udb;
-pub mod udb_keys;
 pub mod wf;

package/packages/engine/src/commands/start.rs

@@ -1,9 +1,9 @@
 use std::time::Duration;
 
-use anyhow::{Context, Result, anyhow};
+use anyhow::{Context, Result, anyhow, bail};
 use clap::Parser;
 use indoc::formatdoc;
-use rivet_service_manager::{CronConfig, RunConfig};
+use rivet_service_manager::RunConfig;
 use universaldb::utils::IsolationLevel::*;
 
 use crate::keys;
@@ -13,36 +13,12 @@ const LOGS_RETENTION: Duration = Duration::from_secs(7 * 24 * 60 * 60);
 
 #[derive(Parser)]
 pub struct Opts {
-	#[arg(long, value_enum)]
-	services: Vec<ServiceKind>,
+	#[arg(short = 's', long, conflicts_with = "except_services")]
+	services: Vec<String>,
 
 	/// Exclude the specified services instead of including them
 	#[arg(long)]
-	except_services: Vec<ServiceKind>,
-}
-
-#[derive(clap::ValueEnum, Clone, PartialEq)]
-enum ServiceKind {
-	ApiPublic,
-	ApiPeer,
-	Standalone,
-	Singleton,
-	Oneshot,
-	Cron,
-}
-
-impl From<ServiceKind> for rivet_service_manager::ServiceKind {
-	fn from(val: ServiceKind) -> Self {
-		use ServiceKind::*;
-		match val {
-			ApiPublic => rivet_service_manager::ServiceKind::ApiPublic,
-			ApiPeer => rivet_service_manager::ServiceKind::ApiPeer,
-			Standalone => rivet_service_manager::ServiceKind::Standalone,
-			Singleton => rivet_service_manager::ServiceKind::Singleton,
-			Oneshot => rivet_service_manager::ServiceKind::Oneshot,
-			Cron => rivet_service_manager::ServiceKind::Cron(CronConfig::default()),
-		}
-	}
+	except_services: Vec<String>,
 }
 
 impl Opts {
@@ -63,33 +39,37 @@ impl Opts {
 			// Run all services
 			run_config.services.clone()
 		} else if !self.except_services.is_empty() {
-			// Exclude specified services
-			let except_service_kinds = self
-				.except_services
-				.iter()
-				.map(|x| x.clone().into())
-				.collect::<Vec<rivet_service_manager::ServiceKind>>();
-
-			run_config
-				.services
-				.iter()
-				.filter(|x| !except_service_kinds.iter().any(|y| y.eq(&x.kind)))
-				.cloned()
-				.collect::<Vec<_>>()
+			let mut services = run_config.services.clone();
+
+			for exclude_name in &self.except_services {
+				if !run_config
+					.services
+					.iter()
+					.any(|service| service.name == exclude_name)
+				{
+					bail!("service {exclude_name:?} not found");
+				}
+
+				services.retain(|service| service.name != exclude_name);
+			}
+
+			services
 		} else {
-			// Include only specified services
-			let service_kinds = self
-				.services
-				.iter()
-				.map(|x| x.clone().into())
-				.collect::<Vec<rivet_service_manager::ServiceKind>>();
-
-			run_config
-				.services
-				.iter()
-				.filter(|x| service_kinds.iter().any(|y| y.eq(&x.kind)))
-				.cloned()
-				.collect::<Vec<_>>()
+			let mut services = Vec::new();
+
+			for name in &self.services {
+				let Some(service) = run_config
+					.services
+					.iter()
+					.find(|service| service.name == name)
+				else {
+					bail!("service {name:?} not found");
+				};
+
+				services.push(service.clone());
+			}
+
+			services
 		};
 
 		let pools = rivet_pools::Pools::new(config.clone()).await?;
@@ -115,15 +95,14 @@ async fn verify_engine_version(
 	pools
 		.udb()?
 		.run(|tx| async move {
-			let current_version = semver::Version::parse(env!("CARGO_PKG_VERSION"))
-				.context("failed to parse cargo pkg version as semver")?;
+			let current_version = semver::Version::parse(env!("CARGO_PKG_VERSION")).context("failed to parse cargo pkg version as semver")?;
 
-			if let Some(existing_version) =
-				tx.read_opt(&keys::EngineVersionKey {}, Serializable).await?
-			{
+			if let Some(existing_version) = tx.read_opt(&keys::EngineVersionKey {}, Serializable).await? {
 				if current_version < existing_version {
-					return Ok(Err(anyhow!("{}", formatdoc!(
-						"
+					return Ok(Err(anyhow!(
+						"{}",
+						formatdoc!(
+							"
 						Rivet Engine has been rolled back to a previous version:
 						  - Last Used Version: {existing_version}
 						  - Current Version:   {current_version}
@@ -131,7 +110,8 @@ async fn verify_engine_version(
 						
 						(If you know what you're doing, this error can be disabled in the Rivet config via `allow_version_rollback: true`)
 						"
-					))));
+						)
+					)));
 				}
 			}