rivet-envoy-protocol 2.3.0-rc.12

use std::{
	fs,
	path::{Path, PathBuf},
	process::Command,
};

fn main() -> Result<(), Box<dyn std::error::Error>> {
	let manifest_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR")?);
	let out_dir = PathBuf::from(std::env::var("OUT_DIR")?);
	let workspace_root = manifest_dir
		.parent()
		.and_then(|p| p.parent())
		.and_then(|p| p.parent())
		.ok_or("Failed to find workspace root")?;

	let schema_dir = manifest_dir.join("schemas");

	// Rust SDK generation
	let cfg = vbare_compiler::Config::default();
	vbare_compiler::process_schemas_with_config(&schema_dir, &cfg)?;

	// Append protocol version constant to generated file
	let (highest_version, _) = find_highest_version(&schema_dir);
	let combined_imports_path = out_dir.join("combined_imports.rs");
	let mut combined = fs::read_to_string(&combined_imports_path)?;
	combined.push_str(&format!(
		"\npub const PROTOCOL_VERSION: u16 = {};\n",
		highest_version
	));
	fs::write(combined_imports_path, combined)?;

	// TypeScript SDK generation
	let cli_js_path = workspace_root
		.parent()
		.unwrap()
		.join("node_modules/@bare-ts/tools/dist/bin/cli.js");
	if cli_js_path.exists() {
		typescript::generate_sdk(&schema_dir);
	} else {
		println!(
			"cargo:warning=TypeScript SDK generation skipped: cli.js not found at {}. Run `pnpm install` to install.",
			cli_js_path.display()
		);
	}

	Ok(())
}

mod typescript {
	use super::*;

	pub fn generate_sdk(schema_dir: &Path) {
		let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
		let workspace_root = Path::new(&manifest_dir)
			.parent()
			.and_then(|p| p.parent())
			.and_then(|p| p.parent())
			.expect("Failed to find workspace root");

		let sdk_dir = workspace_root
			.join("sdks")
			.join("typescript")
			.join("envoy-protocol");
		let src_dir = sdk_dir.join("src");

		let (highest_version, highest_version_path) = super::find_highest_version(schema_dir);

		let _ = fs::remove_dir_all(&src_dir);
		if let Err(e) = fs::create_dir_all(&src_dir) {
			panic!("Failed to create SDK directory: {}", e);
		}

		let output_path = src_dir.join("index.ts");

		let output = Command::new(
			workspace_root
				.parent()
				.unwrap()
				.join("node_modules/@bare-ts/tools/dist/bin/cli.js"),
		)
		.arg("compile")
		.arg("--generator")
		.arg("ts")
		.arg(highest_version_path)
		.arg("-o")
		.arg(&output_path)
		.output()
		.expect("Failed to execute bare compiler for TypeScript");

		if !output.status.success() {
			panic!(
				"BARE TypeScript generation failed: {}",
				String::from_utf8_lossy(&output.stderr),
			);
		}

		// Post-process the generated TypeScript file
		// IMPORTANT: Keep this in sync with rivetkit-typescript/packages/rivetkit/scripts/compile-bare.ts
		post_process_generated_ts(highest_version, &output_path);
	}

	const POST_PROCESS_MARKER: &str = "// @generated - post-processed by build.rs\n";

	/// Post-process the generated TypeScript file to:
	/// 1. Replace @bare-ts/lib import with @rivetkit/bare-ts
	/// 2. Replace Node.js assert import with a custom assert function
	///
	/// IMPORTANT: Keep this in sync with rivetkit-typescript/packages/rivetkit/scripts/compile-bare.ts
	fn post_process_generated_ts(highest_version: u32, path: &Path) {
		let content = fs::read_to_string(path).expect("Failed to read generated TypeScript file");

		// Skip if already post-processed
		if content.starts_with(POST_PROCESS_MARKER) {
			return;
		}

		// Replace @bare-ts/lib with @rivetkit/bare-ts
		let content = content.replace("@bare-ts/lib", "@rivetkit/bare-ts");

		// Replace Node.js assert import with custom assert function
		let content = content.replace("import assert from \"assert\"", "");
		let content = content.replace("import assert from \"node:assert\"", "");

		// Append custom assert function
		let assert_function = r#"
function assert(condition: boolean, message?: string): asserts condition {
    if (!condition) throw new Error(message ?? "Assertion failed")
}
"#;
		// Append current protocol version
		let version = format!(r#"export const VERSION = {};"#, highest_version);

		let content = format!(
			"{}{}\n{}\n{}",
			POST_PROCESS_MARKER, content, assert_function, version
		);

		// Validate post-processing succeeded
		assert!(
			!content.contains("@bare-ts/lib"),
			"Failed to replace @bare-ts/lib import"
		);
		assert!(
			!content.contains("import assert from"),
			"Failed to remove Node.js assert import"
		);

		fs::write(path, content).expect("Failed to write post-processed TypeScript file");
	}
}

fn find_highest_version(schema_dir: &Path) -> (u32, PathBuf) {
	let mut highest_version = 0;
	let mut highest_version_path = PathBuf::new();

	for entry in fs::read_dir(schema_dir).unwrap().flatten() {
		if !entry.path().is_dir() {
			let path = entry.path();
			let bare_name = path
				.file_name()
				.unwrap()
				.to_str()
				.unwrap()
				.split_once('.')
				.unwrap()
				.0;

			if let Ok(version) = bare_name[1..].parse::<u32>() {
				if version > highest_version {
					highest_version = version;
					highest_version_path = path;
				}
			}
		}
	}

	(highest_version, highest_version_path)
}