@rivetkit/engine 1.0.0

package/packages/config/src/lib.rs

@@ -0,0 +1,135 @@
+use std::{ops::Deref, path::Path, result::Result::Ok, sync::Arc};
+
+use ::config as config_loader;
+use anyhow::*;
+
+pub mod config;
+pub mod defaults;
+pub mod paths;
+pub mod secret;
+
+struct ConfigData {
+	config: config::Root,
+}
+
+#[derive(Clone)]
+pub struct Config(Arc<ConfigData>);
+
+impl Config {
+	pub async fn load<P: AsRef<Path>>(paths: &[P]) -> Result<Self> {
+		let mut settings = config_loader::Config::builder();
+
+		// Start with default values
+		settings = settings.add_source(config_loader::Config::try_from(&config::Root::default())?);
+
+		if paths.is_empty() {
+			let default_path = paths::system_config_dir();
+			if default_path.exists() {
+				// Add default config directory if it exists
+				settings = add_source(settings, default_path)?;
+			}
+		} else {
+			// Use provided paths
+			for path in paths {
+				settings = add_source(settings, path)?;
+			}
+		}
+
+		// Add env source for overrides
+		settings = settings.add_source(
+			config_loader::Environment::with_prefix("RIVET")
+				.separator("__")
+				.list_separator(","),
+		);
+
+		// Read config
+		let mut config_root = settings
+			.build()
+			.context("failed to build config")?
+			.try_deserialize::<config::Root>()
+			.context("failed to deserialize config")?;
+
+		// Validate configuration at load time
+		config_root.validate_and_set_defaults()?;
+
+		Ok(Self(Arc::new(ConfigData {
+			config: config_root,
+		})))
+	}
+
+	pub fn from_root(config: config::Root) -> Self {
+		Self(Arc::new(ConfigData { config }))
+	}
+}
+
+impl Deref for Config {
+	type Target = config::Root;
+
+	fn deref(&self) -> &Self::Target {
+		&self.0.config
+	}
+}
+
+/// Adds a source to the config builder. If the path is a directory, it reads all config files.
+/// If it's a file, it adds it directly. If the path doesn't exist, it's silently ignored.
+fn add_source<P: AsRef<Path>>(
+	mut settings: config_loader::ConfigBuilder<config_loader::builder::DefaultState>,
+	path: P,
+) -> Result<config_loader::ConfigBuilder<config_loader::builder::DefaultState>> {
+	let path = path.as_ref();
+
+	if !path.exists() {
+		// Silently ignore non-existent paths
+		return Ok(settings);
+	}
+
+	if path.is_dir() {
+		for entry in std::fs::read_dir(path)? {
+			let entry = entry?;
+			let path = entry.path();
+			if path.is_file() {
+				if let Some(extension) = path.extension().and_then(std::ffi::OsStr::to_str) {
+					if ["json", "json5", "jsonc", "yaml", "yml"].contains(&extension) {
+						settings = add_file_source(settings, &path)?;
+					}
+				}
+			}
+		}
+	} else if path.is_file() {
+		settings = add_file_source(settings, path)?;
+	} else {
+		bail!("Invalid path: {}", path.display());
+	}
+
+	Ok(settings)
+}
+
+/// Adds a single file source to the config builder.
+fn add_file_source<P: AsRef<Path>>(
+	settings: config_loader::ConfigBuilder<config_loader::builder::DefaultState>,
+	path: P,
+) -> Result<config_loader::ConfigBuilder<config_loader::builder::DefaultState>> {
+	let path = path.as_ref();
+	let content = std::fs::read_to_string(path)
+		.with_context(|| format!("failed to read file: {}", path.display()))?;
+
+	let format = match path.extension().and_then(std::ffi::OsStr::to_str) {
+		Some("json") => config_loader::FileFormat::Json,
+		Some("json5") | Some("jsonc") => {
+			// Parse JSON5/JSONC and convert to regular JSON
+			let value = match json5::from_str::<serde_json::Value>(&content) {
+				Ok(x) => x,
+				Err(err) => bail!("failed to parse config file at {}: {err}", path.display()),
+			};
+			let json = serde_json::to_string(&value)?;
+			return Ok(settings.add_source(config_loader::File::from_str(
+				&json,
+				config_loader::FileFormat::Json,
+			)));
+		}
+		Some("yaml") | Some("yml") => config_loader::FileFormat::Yaml,
+		_ => bail!("Unsupported file format: {}", path.display()),
+	};
+
+	Ok(settings.add_source(config_loader::File::from_str(&content, format)))
+}

package/packages/config/src/paths.rs

@@ -0,0 +1,25 @@
+use std::path::PathBuf;
+
+/// Name of the config directory holding the rivet config.
+const CONFIG_DIR_NAME: &str = "rivet";
+
+#[cfg(target_os = "linux")]
+pub fn system_config_dir() -> PathBuf {
+	PathBuf::from("/etc").join(CONFIG_DIR_NAME)
+}
+
+#[cfg(target_os = "macos")]
+pub fn system_config_dir() -> PathBuf {
+	PathBuf::from("/Library/Application Support").join(CONFIG_DIR_NAME)
+}
+
+#[cfg(target_os = "windows")]
+pub fn system_config_dir() -> PathBuf {
+	PathBuf::from("C:\\ProgramData").join(CONFIG_DIR_NAME)
+}
+
+// This will cause a compile-time error if an unsupported OS is targeted.
+#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
+pub fn system_config_dir() -> PathBuf {
+	compile_error!("Unsupported OS");
+}

package/packages/config/src/secret.rs

@@ -0,0 +1,72 @@
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+use std::fmt;
+
+/// Wraps a given value in a container that does not implement `Debug` or `Display` in order to
+/// prevent accidentally logging the inner secret.
+///
+/// In order to access the secret, explicitly use `read()`.
+#[derive(Clone)]
+pub struct Secret<T>(T)
+where
+	T: Clone;
+
+impl<T> Secret<T>
+where
+	T: Clone,
+{
+	pub fn new(v: T) -> Self {
+		Self(v)
+	}
+
+	/// Read the secret.
+	pub fn read(&self) -> &T {
+		&self.0
+	}
+}
+
+impl<T> fmt::Debug for Secret<T>
+where
+	T: Clone,
+{
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		f.debug_struct("Secret").finish()
+	}
+}
+
+impl<T> Serialize for Secret<T>
+where
+	T: Clone + Serialize,
+{
+	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+	where
+		S: serde::Serializer,
+	{
+		self.0.serialize(serializer)
+	}
+}
+
+impl<'de, T> Deserialize<'de> for Secret<T>
+where
+	T: Clone + Deserialize<'de>,
+{
+	fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+	where
+		D: serde::Deserializer<'de>,
+	{
+		T::deserialize(deserializer).map(Secret)
+	}
+}
+
+impl<T> JsonSchema for Secret<T>
+where
+	T: Clone + JsonSchema,
+{
+	fn schema_name() -> String {
+		format!("Secret<{}>", T::schema_name())
+	}
+
+	fn json_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::schema::Schema {
+		T::json_schema(generator)
+	}
+}

package/packages/dump-openapi/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "rivet-dump-openapi"
+version.workspace = true
+edition.workspace = true
+authors.workspace = true
+license.workspace = true
+
+[build-dependencies]
+rivet-api-public.workspace = true
+serde_json.workspace = true
+utoipa.workspace = true

package/packages/dump-openapi/build.rs

@@ -0,0 +1,26 @@
+use std::{fs, path::Path};
+use utoipa::OpenApi;
+
+fn main() {
+	let openapi = rivet_api_public::router::ApiDoc::openapi();
+
+	// Create out directory at workspace root
+	let workspace_root = std::env::var("CARGO_MANIFEST_DIR")
+		.map(|dir| {
+			Path::new(&dir)
+				.parent()
+				.unwrap()
+				.parent()
+				.unwrap()
+				.parent()
+				.unwrap()
+				.to_path_buf()
+		})
+		.unwrap();
+	let out_dir = workspace_root.join("engine").join("artifacts");
+	fs::create_dir_all(&out_dir).unwrap();
+
+	// Write pretty-formatted JSON to out/openapi.json
+	let json = serde_json::to_string_pretty(&openapi).expect("Failed to serialize OpenAPI spec");
+	fs::write(out_dir.join("openapi.json"), json).expect("Failed to write openapi.json");
+}

package/packages/dump-openapi/src/lib.rs

@@ -0,0 +1 @@
+// This library does nothing. See build.rs.

package/packages/engine/Cargo.toml

@@ -0,0 +1,74 @@
+[package]
+name = "rivet-engine"
+version.workspace = true
+authors.workspace = true
+license.workspace = true
+edition.workspace = true
+
+[[bin]]
+name = "rivet-engine"
+path = "src/main.rs"
+
+[dependencies]
+anyhow.workspace = true
+chrono.workspace = true
+clap.workspace = true
+colored_json.workspace = true
+futures-util.workspace = true
+gas.workspace = true
+hex.workspace = true
+include_dir.workspace = true
+indoc.workspace = true
+lz4_flex.workspace = true
+once_cell.workspace = true
+pegboard-runner.workspace = true
+pegboard-serverless.workspace = true
+reqwest.workspace = true
+rivet-api-peer.workspace = true
+rivet-bootstrap.workspace = true
+rivet-cache-purge.workspace = true
+rivet-cache.workspace = true
+rivet-config.workspace = true
+rivet-guard.workspace = true
+rivet-logs.workspace = true
+rivet-pools.workspace = true
+rivet-runtime.workspace = true
+rivet-service-manager.workspace = true
+rivet-telemetry.workspace = true
+rivet-term.workspace = true
+rivet-tracing-reconfigure.workspace = true
+rivet-util.workspace = true
+rivet-workflow-worker.workspace = true
+rustyline.workspace = true
+semver.workspace = true
+serde_json.workspace = true
+serde_yaml.workspace = true
+serde.workspace = true
+strum.workspace = true
+tabled.workspace = true
+tempfile.workspace = true
+thiserror.workspace = true
+tokio.workspace = true
+tracing.workspace = true
+universaldb.workspace = true
+url.workspace = true
+uuid.workspace = true
+
+[dev-dependencies]
+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-runner-protocol.workspace = true
+rivet-test-deps.workspace = true
+rivet-util.workspace = true
+rstest.workspace = true
+tokio-tungstenite.workspace = true
+tracing-subscriber.workspace = true
+vbare.workspace = true

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

@@ -0,0 +1,18 @@
+use anyhow::*;
+use clap::Parser;
+
+#[derive(Parser)]
+pub enum SubCommand {
+	Show,
+}
+
+impl SubCommand {
+	pub async fn execute(self, config: rivet_config::Config) -> Result<()> {
+		match self {
+			Self::Show => {
+				println!("{:#?}", *config);
+				Ok(())
+			}
+		}
+	}
+}

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

@@ -0,0 +1,80 @@
+use anyhow::*;
+use clap::{Parser, ValueEnum};
+
+#[derive(Parser)]
+pub enum SubCommand {
+	#[clap(alias = "sh")]
+	Shell {
+		#[clap(index = 1)]
+		database_type: DatabaseType,
+		#[clap(index = 2)]
+		service: String,
+		#[clap(short = 'q', long)]
+		query: Option<String>,
+	},
+	Nuke {
+		#[clap(long)]
+		yes: bool,
+	},
+}
+
+#[derive(ValueEnum, Clone, PartialEq)]
+pub enum DatabaseType {
+	#[clap(alias = "ch")]
+	Clickhouse,
+	#[clap(alias = "wfd")]
+	WorkflowData,
+	#[clap(alias = "wfi")]
+	WorkflowInternal,
+}
+
+impl SubCommand {
+	pub async fn execute(self, config: rivet_config::Config) -> Result<()> {
+		match self {
+			Self::Shell {
+				database_type: db_type,
+				service,
+				query,
+			} => {
+				let shell_query = crate::util::db::ShellQuery {
+					svc: service,
+					query,
+				};
+				let shell_ctx = crate::util::db::ShellContext {
+					queries: &[shell_query],
+				};
+
+				match db_type {
+					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(())
+			}
+			Self::Nuke { yes } => {
+				if !yes {
+					bail!("Please pass '--yes' flag to confirm you want to nuke the database");
+				}
+
+				match config.database() {
+					rivet_config::config::Database::Postgres(_) => {
+						bail!("nuke postgres not implemented");
+					}
+					rivet_config::config::Database::FileSystem(file_system) => {
+						println!("Removing {}", file_system.path.display());
+						tokio::fs::remove_dir_all(&file_system.path).await?;
+						println!("Complete");
+						Ok(())
+					}
+				}
+			}
+		}
+	}
+}

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

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

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

@@ -0,0 +1,143 @@
+use std::time::Duration;
+
+use anyhow::{Context, Result, anyhow};
+use clap::Parser;
+use indoc::formatdoc;
+use rivet_service_manager::{CronConfig, RunConfig};
+use universaldb::utils::IsolationLevel::*;
+
+use crate::keys;
+
+// 7 day logs retention
+const LOGS_RETENTION: Duration = Duration::from_secs(7 * 24 * 60 * 60);
+
+#[derive(Parser)]
+pub struct Opts {
+	#[arg(long, value_enum)]
+	services: Vec<ServiceKind>,
+
+	/// 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()),
+		}
+	}
+}
+
+impl Opts {
+	pub async fn execute(
+		&self,
+		config: rivet_config::Config,
+		run_config: &RunConfig,
+	) -> Result<()> {
+		// Redirect logs if enabled on the edge
+		if let Some(logs_dir) = config.logs().redirect_logs_dir.as_ref() {
+			rivet_logs::Logs::new(logs_dir.clone(), LOGS_RETENTION)
+				.start()
+				.await?;
+		}
+
+		// Select services to run
+		let services = if self.services.is_empty() && self.except_services.is_empty() {
+			// 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<_>>()
+		} 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 pools = rivet_pools::Pools::new(config.clone()).await?;
+
+		verify_engine_version(&config, &pools).await?;
+
+		// Start server
+		rivet_service_manager::start(config, pools, services).await?;
+
+		Ok(())
+	}
+}
+
+/// Verifies that no rollback has occurred (if allowing rollback is disabled).
+async fn verify_engine_version(
+	config: &rivet_config::Config,
+	pools: &rivet_pools::Pools,
+) -> Result<()> {
+	if config.runtime.allow_version_rollback() {
+		return Ok(());
+	}
+
+	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")?;
+
+			if let Some(existing_version) =
+				tx.read_opt(&keys::EngineVersionKey {}, Serializable).await?
+			{
+				if current_version < existing_version {
+					return Ok(Err(anyhow!("{}", formatdoc!(
+						"
+						Rivet Engine has been rolled back to a previous version:
+						  - Last Used Version: {existing_version}
+						  - Current Version:   {current_version}
+						Cannot proceed without potential data corruption.
+						
+						(If you know what you're doing, this error can be disabled in the Rivet config via `allow_version_rollback: true`)
+						"
+					))));
+				}
+			}
+
+			tx.write(&keys::EngineVersionKey {}, current_version)?;
+
+			Ok(Ok(()))
+		})
+		.await?
+}

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

@@ -0,0 +1,86 @@
+use anyhow::*;
+use clap::Parser;
+use serde::{Deserialize, Serialize};
+
+#[derive(Parser)]
+pub enum SubCommand {
+	/// Configure tracing settings (log filter and sampler ratio)
+	Config {
+		/// Log filter (e.g., "debug", "info", "rivet_api_peer=trace")
+		/// Set to null to reset to defaults
+		#[clap(short, long)]
+		filter: Option<String>,
+
+		/// OpenTelemetry sampler ratio (0.0-1.0)
+		/// Set to null to reset to default
+		#[clap(short, long)]
+		sampler_ratio: Option<f64>,
+
+		/// API peer endpoint
+		#[clap(long, default_value = "http://localhost:6421")]
+		endpoint: String,
+	},
+}
+
+#[derive(Serialize, Deserialize)]
+struct SetTracingConfigRequest {
+	#[serde(skip_serializing_if = "Option::is_none")]
+	pub filter: Option<Option<String>>,
+	#[serde(skip_serializing_if = "Option::is_none")]
+	pub sampler_ratio: Option<Option<f64>>,
+}
+
+impl SubCommand {
+	pub async fn execute(self, _config: rivet_config::Config) -> Result<()> {
+		match self {
+			Self::Config {
+				filter,
+				sampler_ratio,
+				endpoint,
+			} => {
+				// Build request body
+				let request = SetTracingConfigRequest {
+					filter: filter.map(|f| if f.is_empty() { None } else { Some(f) }),
+					sampler_ratio: sampler_ratio.map(Some),
+				};
+
+				// Send HTTP request
+				let client = rivet_pools::reqwest::client().await?;
+				let url = format!("{}/debug/tracing/config", endpoint);
+
+				let response = client
+					.put(&url)
+					.json(&request)
+					.send()
+					.await
+					.context("failed to send request")?;
+
+				if response.status().is_success() {
+					println!("Tracing configuration updated successfully");
+
+					if let Some(Some(f)) = &request.filter {
+						println!("  Filter: {}", f);
+					} else if let Some(None) = &request.filter {
+						println!("  Filter: reset to default");
+					}
+
+					if let Some(Some(r)) = request.sampler_ratio {
+						println!("  Sampler ratio: {}", r);
+					} else if let Some(None) = request.sampler_ratio {
+						println!("  Sampler ratio: reset to default (0.001)");
+					}
+				} else {
+					let status = response.status();
+					let body = response.text().await.unwrap_or_default();
+					bail!(
+						"Failed to update tracing configuration: {} - {}",
+						status,
+						body
+					);
+				}
+
+				Ok(())
+			}
+		}
+	}
+}