initial commit: OpenRouter free model radar Discord bot (WIP)
This commit is contained in:
commit
ee78de7923
7 changed files with 3091 additions and 0 deletions
1
.gitignore
vendored
Normal file
1
.gitignore
vendored
Normal file
|
|
@ -0,0 +1 @@
|
||||||
|
/target
|
||||||
2476
Cargo.lock
generated
Normal file
2476
Cargo.lock
generated
Normal file
File diff suppressed because it is too large
Load diff
19
Cargo.toml
Normal file
19
Cargo.toml
Normal file
|
|
@ -0,0 +1,19 @@
|
||||||
|
[package]
|
||||||
|
name = "model-radar"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2024"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
poise = "0.7"
|
||||||
|
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
|
||||||
|
serde = { version = "1", features = ["derive"] }
|
||||||
|
serde_json = "1"
|
||||||
|
tokio = { version = "1", features = ["macros", "rt-multi-thread", "fs", "signal", "time", "sync"] }
|
||||||
|
tracing = "0.1"
|
||||||
|
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||||
|
|
||||||
|
[profile.dev]
|
||||||
|
debug = 0
|
||||||
|
|
||||||
|
[profile.test]
|
||||||
|
debug = 0
|
||||||
16
README.md
Normal file
16
README.md
Normal file
|
|
@ -0,0 +1,16 @@
|
||||||
|
# model-radar
|
||||||
|
|
||||||
|
> **Status:** Work in Progress (WIP)
|
||||||
|
|
||||||
|
A Discord bot that monitors [OpenRouter](https://openrouter.ai) for free LLM models and alerts/pings users when new free models become available.
|
||||||
|
|
||||||
|
## Features (WIP)
|
||||||
|
|
||||||
|
- Polls OpenRouter model catalog API for free tiers (`:free`)
|
||||||
|
- Tracks pricing and availability diffs
|
||||||
|
- Discord slash commands via poise & serenity
|
||||||
|
- Automated role/user pings on model drops
|
||||||
|
|
||||||
|
## License
|
||||||
|
|
||||||
|
MIT
|
||||||
318
src/catalog.rs
Normal file
318
src/catalog.rs
Normal file
|
|
@ -0,0 +1,318 @@
|
||||||
|
use std::collections::BTreeMap;
|
||||||
|
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
use crate::Error;
|
||||||
|
|
||||||
|
const ENDPOINT: &str = "https://openrouter.ai/api/v1/models";
|
||||||
|
const MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024;
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||||
|
pub struct Model {
|
||||||
|
pub id: String,
|
||||||
|
pub name: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub description: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub context_length: u64,
|
||||||
|
#[serde(default)]
|
||||||
|
pub pricing: BTreeMap<String, String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||||
|
pub struct Event {
|
||||||
|
pub model: Model,
|
||||||
|
pub kind: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_decimal_zero(raw: &str) -> bool {
|
||||||
|
let body = raw.strip_prefix(['+', '-']).unwrap_or(raw);
|
||||||
|
let mantissa = if let Some((mantissa, exponent)) = body.split_once(['e', 'E']) {
|
||||||
|
let exponent = exponent.strip_prefix(['+', '-']).unwrap_or(exponent);
|
||||||
|
if exponent.is_empty() || !exponent.bytes().all(|b| b.is_ascii_digit()) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
mantissa
|
||||||
|
} else {
|
||||||
|
body
|
||||||
|
};
|
||||||
|
let mut seen_dot = false;
|
||||||
|
let mut digits = 0usize;
|
||||||
|
mantissa.bytes().all(|b| match b {
|
||||||
|
b'.' if !seen_dot => {
|
||||||
|
seen_dot = true;
|
||||||
|
true
|
||||||
|
}
|
||||||
|
b'0' => {
|
||||||
|
digits += 1;
|
||||||
|
true
|
||||||
|
}
|
||||||
|
_ => false,
|
||||||
|
}) && digits > 0
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Model {
|
||||||
|
pub fn is_free(&self) -> bool {
|
||||||
|
if !is_decimal_zero(
|
||||||
|
self.pricing
|
||||||
|
.get("prompt")
|
||||||
|
.map(String::as_str)
|
||||||
|
.unwrap_or_default(),
|
||||||
|
) || !is_decimal_zero(
|
||||||
|
self.pricing
|
||||||
|
.get("completion")
|
||||||
|
.map(String::as_str)
|
||||||
|
.unwrap_or_default(),
|
||||||
|
) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
self.pricing.values().all(|value| is_decimal_zero(value))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate(catalog: &BTreeMap<String, Model>) -> Result<(), Error> {
|
||||||
|
if catalog.is_empty() {
|
||||||
|
return Err("catalog is empty".into());
|
||||||
|
}
|
||||||
|
for (key, model) in catalog {
|
||||||
|
if model.id.trim().is_empty() || model.name.trim().is_empty() {
|
||||||
|
return Err("model identity is empty".into());
|
||||||
|
}
|
||||||
|
if key != &model.id {
|
||||||
|
return Err(
|
||||||
|
format!("catalog key {key:?} does not match model id {:?}", model.id).into(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_catalog(bytes: &[u8]) -> Result<BTreeMap<String, Model>, Error> {
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct Response {
|
||||||
|
data: Vec<Model>,
|
||||||
|
}
|
||||||
|
let response: Response = serde_json::from_slice(bytes)?;
|
||||||
|
let mut catalog = BTreeMap::new();
|
||||||
|
for model in response.data {
|
||||||
|
if model.id.trim().is_empty() || model.name.trim().is_empty() {
|
||||||
|
return Err("model identity is empty".into());
|
||||||
|
}
|
||||||
|
if catalog.insert(model.id.clone(), model).is_some() {
|
||||||
|
return Err("duplicate model id".into());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
validate(&catalog)?;
|
||||||
|
Ok(catalog)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn fetch(client: &reqwest::Client) -> Result<BTreeMap<String, Model>, Error> {
|
||||||
|
let mut response = client.get(ENDPOINT).send().await?.error_for_status()?;
|
||||||
|
let mut payload: Vec<u8> = Vec::new();
|
||||||
|
while let Some(chunk) = response.chunk().await? {
|
||||||
|
if chunk.len() > MAX_RESPONSE_BYTES - payload.len() {
|
||||||
|
return Err("response exceeds 8 MiB limit".into());
|
||||||
|
}
|
||||||
|
payload.extend_from_slice(&chunk);
|
||||||
|
}
|
||||||
|
parse_catalog(&payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn changes(
|
||||||
|
previous: &BTreeMap<String, Model>,
|
||||||
|
current: &BTreeMap<String, Model>,
|
||||||
|
) -> Vec<Event> {
|
||||||
|
let mut events = Vec::new();
|
||||||
|
for (id, model) in current {
|
||||||
|
match previous.get(id) {
|
||||||
|
None => events.push(Event {
|
||||||
|
model: model.clone(),
|
||||||
|
kind: "new".to_string(),
|
||||||
|
}),
|
||||||
|
Some(old) if !old.is_free() && model.is_free() => events.push(Event {
|
||||||
|
model: model.clone(),
|
||||||
|
kind: "paid_to_free".to_string(),
|
||||||
|
}),
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
events
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn model(id: &str, pricing: &[(&str, &str)]) -> Model {
|
||||||
|
Model {
|
||||||
|
id: id.to_string(),
|
||||||
|
name: id.to_string(),
|
||||||
|
description: String::new(),
|
||||||
|
context_length: 4096,
|
||||||
|
pricing: pricing
|
||||||
|
.iter()
|
||||||
|
.map(|(k, v)| (k.to_string(), v.to_string()))
|
||||||
|
.collect(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn map(models: Vec<Model>) -> BTreeMap<String, Model> {
|
||||||
|
models.into_iter().map(|m| (m.id.clone(), m)).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn free_requires_exact_decimal_zero() {
|
||||||
|
assert!(model("a", &[("prompt", "0"), ("completion", "0")]).is_free());
|
||||||
|
assert!(model("a", &[("prompt", "0.0"), ("completion", "-0")]).is_free());
|
||||||
|
assert!(model("a", &[("prompt", "+0.00000"), ("completion", "0")]).is_free());
|
||||||
|
assert!(
|
||||||
|
!model(
|
||||||
|
"a",
|
||||||
|
&[("prompt", "0.00000000000000000001"), ("completion", "0")]
|
||||||
|
)
|
||||||
|
.is_free()
|
||||||
|
);
|
||||||
|
assert!(!model("a", &[("prompt", "1e-9"), ("completion", "0")]).is_free());
|
||||||
|
assert!(!model("a", &[("prompt", ""), ("completion", "0")]).is_free());
|
||||||
|
assert!(!model("a", &[("prompt", "abc"), ("completion", "0")]).is_free());
|
||||||
|
assert!(
|
||||||
|
!model(
|
||||||
|
"a",
|
||||||
|
&[("prompt", "0"), ("completion", "0"), ("request", "0.1")]
|
||||||
|
)
|
||||||
|
.is_free()
|
||||||
|
);
|
||||||
|
assert!(!model("a", &[]).is_free());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn decimal_zero_edge_cases() {
|
||||||
|
for value in ["0", "00", "-0", "+0.00", "0e-9999", "0E+9999", ".0", "0."] {
|
||||||
|
assert!(
|
||||||
|
model("a", &[("prompt", value), ("completion", value)]).is_free(),
|
||||||
|
"{value}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
for value in [
|
||||||
|
"1e-9999", "-1e-9999", "NaN", "inf", "-1", ".", "+", "-", "", "0e", "0e+", "0e1.0",
|
||||||
|
"0e0e0", "0..0", "--0", "0x0", " 0", "0 ", "0",
|
||||||
|
] {
|
||||||
|
assert!(
|
||||||
|
!model("a", &[("prompt", value), ("completion", "0")]).is_free(),
|
||||||
|
"{value}"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!model("a", &[("prompt", "0"), ("completion", value)]).is_free(),
|
||||||
|
"{value}"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!model(
|
||||||
|
"a",
|
||||||
|
&[("prompt", "0"), ("completion", "0"), ("unknown", value)]
|
||||||
|
)
|
||||||
|
.is_free(),
|
||||||
|
"{value}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let tiny = format!("0.{}1", "0".repeat(400));
|
||||||
|
assert!(!model("a", &[("prompt", &tiny), ("completion", "0")]).is_free());
|
||||||
|
assert!(!model("a", &[("prompt", "0")]).is_free());
|
||||||
|
assert!(!model("a", &[("completion", "0")]).is_free());
|
||||||
|
assert!(
|
||||||
|
model(
|
||||||
|
"a",
|
||||||
|
&[("prompt", "0"), ("completion", "0"), ("unknown", "0.0")]
|
||||||
|
)
|
||||||
|
.is_free()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn changes_detects_new_and_paid_to_free() {
|
||||||
|
let prev = map(vec![
|
||||||
|
model("a", &[("prompt", "1"), ("completion", "1")]),
|
||||||
|
model("b", &[("prompt", "0"), ("completion", "0")]),
|
||||||
|
]);
|
||||||
|
let cur = map(vec![
|
||||||
|
model("a", &[("prompt", "0"), ("completion", "0")]),
|
||||||
|
model("b", &[("prompt", "0"), ("completion", "0")]),
|
||||||
|
model("c", &[("prompt", "0"), ("completion", "0")]),
|
||||||
|
model("d", &[("prompt", "2"), ("completion", "3")]),
|
||||||
|
]);
|
||||||
|
let events = changes(&prev, &cur);
|
||||||
|
let kinds: Vec<&str> = events.iter().map(|e| e.kind.as_str()).collect();
|
||||||
|
assert_eq!(kinds, vec!["paid_to_free", "new", "new"]);
|
||||||
|
assert_eq!(events[0].model.id, "a");
|
||||||
|
assert_eq!(events[1].model.id, "c");
|
||||||
|
assert_eq!(events[2].model.id, "d");
|
||||||
|
assert!(!events[2].model.is_free());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn changes_ignores_removals_and_unchanged() {
|
||||||
|
let prev = map(vec![model("a", &[("prompt", "0"), ("completion", "0")])]);
|
||||||
|
let cur = BTreeMap::new();
|
||||||
|
assert!(changes(&prev, &cur).is_empty());
|
||||||
|
assert_eq!(changes(&cur, &prev).len(), 1);
|
||||||
|
assert!(changes(&prev, &prev).is_empty());
|
||||||
|
let paid = map(vec![model("a", &[("prompt", "1"), ("completion", "1")])]);
|
||||||
|
assert!(changes(&prev, &paid).is_empty());
|
||||||
|
assert!(changes(&paid, &paid).is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_rejects_empty() {
|
||||||
|
assert!(parse_catalog(b"{\"data\":[]}").is_err());
|
||||||
|
assert!(parse_catalog(b"{}").is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_rejects_malformed() {
|
||||||
|
for body in [
|
||||||
|
"not json",
|
||||||
|
r#"{"data":null}"#,
|
||||||
|
r#"{"data":{}}"#,
|
||||||
|
r#"{"data":[{"id":"","name":"A"}]}"#,
|
||||||
|
r#"{"data":[{"id":" ","name":"A"}]}"#,
|
||||||
|
r#"{"data":[{"id":"a","name":" "}]}"#,
|
||||||
|
r#"{"data":[{"name":"A"}]}"#,
|
||||||
|
r#"{"data":[{"id":"a"}]}"#,
|
||||||
|
r#"{"data":[{"id":1,"name":"A"}]}"#,
|
||||||
|
r#"{"data":[{"id":"a","name":"A","context_length":-1}]}"#,
|
||||||
|
r#"{"data":[{"id":"a","name":"A","pricing":{"prompt":0}}]}"#,
|
||||||
|
r#"{"data":[{"id":"a","name":"A","description":null}]}"#,
|
||||||
|
r#"{"data":[{"id":"a","name":"A"}]} trailing"#,
|
||||||
|
] {
|
||||||
|
assert!(parse_catalog(body.as_bytes()).is_err(), "{body}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_rejects_duplicates() {
|
||||||
|
let body = r#"{"data":[{"id":"a","name":"A"},{"id":"a","name":"B"}]}"#;
|
||||||
|
assert!(parse_catalog(body.as_bytes()).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_accepts_valid() {
|
||||||
|
let body = r#"{"data":[{"id":"a","name":"A","context_length":8192,"pricing":{"prompt":"0","completion":"0"}}]}"#;
|
||||||
|
let catalog = parse_catalog(body.as_bytes()).unwrap();
|
||||||
|
assert_eq!(catalog.len(), 1);
|
||||||
|
assert!(catalog["a"].is_free());
|
||||||
|
assert_eq!(catalog["a"].context_length, 8192);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_defaults_optional_metadata_and_prices() {
|
||||||
|
let catalog = parse_catalog(br#"{"data":[{"id":"a","name":"A"}]}"#).unwrap();
|
||||||
|
assert_eq!(catalog["a"].description, "");
|
||||||
|
assert_eq!(catalog["a"].context_length, 0);
|
||||||
|
assert!(catalog["a"].pricing.is_empty());
|
||||||
|
assert!(!catalog["a"].is_free());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn validate_rejects_empty_catalog() {
|
||||||
|
assert!(validate(&BTreeMap::new()).is_err());
|
||||||
|
}
|
||||||
|
}
|
||||||
38
src/main.rs
Normal file
38
src/main.rs
Normal file
|
|
@ -0,0 +1,38 @@
|
||||||
|
mod catalog;
|
||||||
|
mod store;
|
||||||
|
|
||||||
|
use poise::serenity_prelude as serenity;
|
||||||
|
|
||||||
|
type Error = Box<dyn std::error::Error + Send + Sync>;
|
||||||
|
type Context<'a> = poise::Context<'a, (), Error>;
|
||||||
|
|
||||||
|
#[poise::command(slash_command)]
|
||||||
|
async fn ping(ctx: Context<'_>) -> Result<(), Error> {
|
||||||
|
ctx.say("Model Radar is online.").await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> Result<(), Error> {
|
||||||
|
let token = std::env::var("DISCORD_TOKEN")?;
|
||||||
|
let mut ping = ping();
|
||||||
|
ping.description = Some("Check whether Model Radar is online".into());
|
||||||
|
let framework = poise::Framework::builder()
|
||||||
|
.options(poise::FrameworkOptions {
|
||||||
|
commands: vec![ping],
|
||||||
|
..Default::default()
|
||||||
|
})
|
||||||
|
.setup(|ctx, _, framework| {
|
||||||
|
Box::pin(async move {
|
||||||
|
poise::builtins::register_globally(ctx, &framework.options().commands).await?;
|
||||||
|
Ok(())
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.build();
|
||||||
|
serenity::ClientBuilder::new(token, serenity::GatewayIntents::GUILDS)
|
||||||
|
.framework(framework)
|
||||||
|
.await?
|
||||||
|
.start()
|
||||||
|
.await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
223
src/store.rs
Normal file
223
src/store.rs
Normal file
|
|
@ -0,0 +1,223 @@
|
||||||
|
use std::collections::BTreeMap;
|
||||||
|
use std::io::{ErrorKind, Write};
|
||||||
|
use std::path::PathBuf;
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
use crate::Error;
|
||||||
|
use crate::catalog::{Event, Model};
|
||||||
|
|
||||||
|
static TEMP_SEQUENCE: AtomicU64 = AtomicU64::new(0);
|
||||||
|
|
||||||
|
#[derive(Clone, Default, Debug, Serialize, Deserialize)]
|
||||||
|
pub struct State {
|
||||||
|
pub initialized: bool,
|
||||||
|
pub catalog: BTreeMap<String, Model>,
|
||||||
|
pub subscriptions: BTreeMap<u64, Subscription>,
|
||||||
|
pub last_success: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||||
|
pub struct Subscription {
|
||||||
|
pub channel_id: u64,
|
||||||
|
pub role_id: Option<u64>,
|
||||||
|
pub free_only: bool,
|
||||||
|
pub pending: Vec<Event>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct Store {
|
||||||
|
pub state: State,
|
||||||
|
path: PathBuf,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Store {
|
||||||
|
pub async fn load(path: PathBuf) -> Result<Self, Error> {
|
||||||
|
let state = match tokio::fs::read(&path).await {
|
||||||
|
Ok(bytes) => serde_json::from_slice(&bytes)?,
|
||||||
|
Err(error) if error.kind() == ErrorKind::NotFound => State::default(),
|
||||||
|
Err(error) => return Err(error.into()),
|
||||||
|
};
|
||||||
|
Ok(Self { state, path })
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn commit(&mut self, next: State) -> Result<(), Error> {
|
||||||
|
let bytes = serde_json::to_vec(&next)?;
|
||||||
|
let path = self.path.clone();
|
||||||
|
tokio::task::spawn_blocking(move || -> Result<(), Error> {
|
||||||
|
let name = path.file_name().ok_or("state path has no file name")?;
|
||||||
|
let (temporary, mut file) = loop {
|
||||||
|
let sequence = TEMP_SEQUENCE.fetch_add(1, Ordering::Relaxed);
|
||||||
|
let mut temporary_name = name.to_os_string();
|
||||||
|
temporary_name.push(format!(".{}.{}.tmp", std::process::id(), sequence));
|
||||||
|
let temporary = path.with_file_name(temporary_name);
|
||||||
|
match std::fs::OpenOptions::new()
|
||||||
|
.write(true)
|
||||||
|
.create_new(true)
|
||||||
|
.open(&temporary)
|
||||||
|
{
|
||||||
|
Ok(file) => break (temporary, file),
|
||||||
|
Err(error) if error.kind() == ErrorKind::AlreadyExists => continue,
|
||||||
|
Err(error) => return Err(error.into()),
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let result = (|| -> Result<(), Error> {
|
||||||
|
file.write_all(&bytes)?;
|
||||||
|
file.sync_all()?;
|
||||||
|
drop(file);
|
||||||
|
std::fs::rename(&temporary, &path)?;
|
||||||
|
Ok(())
|
||||||
|
})();
|
||||||
|
if result.is_err() {
|
||||||
|
let _ = std::fs::remove_file(&temporary);
|
||||||
|
}
|
||||||
|
result
|
||||||
|
})
|
||||||
|
.await??;
|
||||||
|
self.state = next;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
struct TestDirectory(PathBuf);
|
||||||
|
|
||||||
|
impl TestDirectory {
|
||||||
|
fn new() -> Self {
|
||||||
|
loop {
|
||||||
|
let sequence = TEMP_SEQUENCE.fetch_add(1, Ordering::Relaxed);
|
||||||
|
let path = std::env::temp_dir().join(format!(
|
||||||
|
"model-radar-store-test-{}-{}",
|
||||||
|
std::process::id(),
|
||||||
|
sequence
|
||||||
|
));
|
||||||
|
match std::fs::create_dir(&path) {
|
||||||
|
Ok(()) => return Self(path),
|
||||||
|
Err(error) if error.kind() == ErrorKind::AlreadyExists => continue,
|
||||||
|
Err(error) => panic!("{error}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn state_path(&self) -> PathBuf {
|
||||||
|
self.0.join("state.json")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for TestDirectory {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
let _ = std::fs::remove_dir_all(&self.0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn populated_state() -> State {
|
||||||
|
let model = Model {
|
||||||
|
id: "example/model".into(),
|
||||||
|
name: "Example".into(),
|
||||||
|
description: "Description".into(),
|
||||||
|
context_length: 8192,
|
||||||
|
pricing: BTreeMap::from([
|
||||||
|
("prompt".into(), "0".into()),
|
||||||
|
("completion".into(), "0".into()),
|
||||||
|
]),
|
||||||
|
};
|
||||||
|
State {
|
||||||
|
initialized: true,
|
||||||
|
catalog: BTreeMap::from([(model.id.clone(), model.clone())]),
|
||||||
|
subscriptions: BTreeMap::from([
|
||||||
|
(
|
||||||
|
u64::MAX,
|
||||||
|
Subscription {
|
||||||
|
channel_id: 42,
|
||||||
|
role_id: Some(99),
|
||||||
|
free_only: true,
|
||||||
|
pending: vec![Event {
|
||||||
|
model,
|
||||||
|
kind: "new".into(),
|
||||||
|
}],
|
||||||
|
},
|
||||||
|
),
|
||||||
|
(
|
||||||
|
1,
|
||||||
|
Subscription {
|
||||||
|
channel_id: 43,
|
||||||
|
role_id: None,
|
||||||
|
free_only: false,
|
||||||
|
pending: Vec::new(),
|
||||||
|
},
|
||||||
|
),
|
||||||
|
]),
|
||||||
|
last_success: Some(123456),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn missing_file_defaults_without_writing() {
|
||||||
|
let directory = TestDirectory::new();
|
||||||
|
let path = directory.state_path();
|
||||||
|
let store = Store::load(path.clone()).await.unwrap();
|
||||||
|
assert!(!store.state.initialized);
|
||||||
|
assert!(store.state.catalog.is_empty());
|
||||||
|
assert!(store.state.subscriptions.is_empty());
|
||||||
|
assert_eq!(store.state.last_success, None);
|
||||||
|
assert!(!path.exists());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn persistence_round_trip_and_replacement() {
|
||||||
|
let directory = TestDirectory::new();
|
||||||
|
let path = directory.state_path();
|
||||||
|
let mut store = Store::load(path.clone()).await.unwrap();
|
||||||
|
let next = populated_state();
|
||||||
|
store.commit(next.clone()).await.unwrap();
|
||||||
|
let loaded = Store::load(path.clone()).await.unwrap();
|
||||||
|
let expected = serde_json::to_value(&next).unwrap();
|
||||||
|
assert_eq!(serde_json::to_value(&store.state).unwrap(), expected);
|
||||||
|
assert_eq!(serde_json::to_value(&loaded.state).unwrap(), expected);
|
||||||
|
store.commit(State::default()).await.unwrap();
|
||||||
|
let loaded = Store::load(path).await.unwrap();
|
||||||
|
assert!(!loaded.state.initialized);
|
||||||
|
assert!(loaded.state.catalog.is_empty());
|
||||||
|
assert!(loaded.state.subscriptions.is_empty());
|
||||||
|
assert_eq!(loaded.state.last_success, None);
|
||||||
|
assert_eq!(std::fs::read_dir(&directory.0).unwrap().count(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn corrupt_files_are_errors_and_not_overwritten() {
|
||||||
|
let directory = TestDirectory::new();
|
||||||
|
let path = directory.state_path();
|
||||||
|
for bytes in ["", "not json", "{}", "null", "{\"initialized\":false"] {
|
||||||
|
tokio::fs::write(&path, bytes).await.unwrap();
|
||||||
|
assert!(Store::load(path.clone()).await.is_err());
|
||||||
|
assert_eq!(tokio::fs::read(&path).await.unwrap(), bytes.as_bytes());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn failed_rename_preserves_state_and_cleans_temporary() {
|
||||||
|
let directory = TestDirectory::new();
|
||||||
|
let path = directory.state_path();
|
||||||
|
let mut store = Store::load(path.clone()).await.unwrap();
|
||||||
|
tokio::fs::create_dir(&path).await.unwrap();
|
||||||
|
assert!(store.commit(populated_state()).await.is_err());
|
||||||
|
assert!(!store.state.initialized);
|
||||||
|
assert!(store.state.catalog.is_empty());
|
||||||
|
assert!(path.is_dir());
|
||||||
|
assert_eq!(std::fs::read_dir(&directory.0).unwrap().count(), 1);
|
||||||
|
assert!(Store::load(path).await.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn missing_parent_commit_preserves_state() {
|
||||||
|
let directory = TestDirectory::new();
|
||||||
|
let path = directory.0.join("missing").join("state.json");
|
||||||
|
let mut store = Store::load(path).await.unwrap();
|
||||||
|
assert!(store.commit(populated_state()).await.is_err());
|
||||||
|
assert!(!store.state.initialized);
|
||||||
|
assert!(store.state.subscriptions.is_empty());
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Add table
Reference in a new issue