fix(rust): sweep orphan route_channels when their model has no active account
The per-route loop in `rebuild_routes` only visits routes whose model_pattern appears in the current `by_model` set. When the last active account for a model goes away (account deleted, status flipped, model dropped from model_availability), the model falls out of `by_model` entirely and its existing auto-managed channels are never reconciled — they stay behind pointing at a route nothing can serve. Add a final orphan sweep over all auto-managed route_channels. Anything whose route wasn't processed in the per-route loop gets hard-deleted. Manual-override channels still get the same exemption as before. Covered by two new in-memory sqlite tests: - orphan_channel_is_cleaned_when_last_account_goes_inactive - orphan_sweep_preserves_manual_channels
This commit is contained in:
boos
@@ -0,0 +1,435 @@
|
||||
//! Route rebuild pipeline.
|
||||
//!
|
||||
//! Reads `model_availability` (filled by `POST /api/accounts/{id}/models/refresh`),
|
||||
//! aggregates by model name, and writes `token_routes` + `route_channels`. This is
|
||||
//! the pipeline that turns "I added an account and refreshed its models" into rows
|
||||
//! the frontend `/routes` page can render.
|
||||
//!
|
||||
//! Behavior:
|
||||
//! - One `token_routes` row per distinct model name found across active accounts.
|
||||
//! Existing rows for that model are reused (we update display fields where
|
||||
//! missing) so manually-edited routes survive a rebuild.
|
||||
//! - One `route_channels` row per `(route, account, model)` triple. Channels
|
||||
//! marked `manual_override = 1` are NEVER touched (created, updated, deleted).
|
||||
//! Other channels are reconciled to match the current model availability:
|
||||
//! missing combos get inserted, stale combos get hard-deleted.
|
||||
//! - Disabled site models (`site_disabled_models`) are filtered out before the
|
||||
//! write. So is anything from accounts whose status isn't `active`, or models
|
||||
//! where `model_availability.available != 1`.
|
||||
//!
|
||||
//! The handler returns `rebuild.createdRoutes` / `createdChannels` so the existing
|
||||
//! frontend toast (`TokenRoutes.tsx:411-412`) reads the right counts.
|
||||
|
||||
use crate::{error::AppResult, AppState};
|
||||
use serde::Serialize;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
/// Summary returned to the frontend after a rebuild.
|
||||
#[derive(Debug, Default, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RebuildSummary {
|
||||
pub created_routes: i64,
|
||||
pub reused_routes: i64,
|
||||
pub created_channels: i64,
|
||||
pub deleted_channels: i64,
|
||||
pub skipped_manual_channels: i64,
|
||||
pub scanned_models: i64,
|
||||
pub scanned_accounts: i64,
|
||||
}
|
||||
|
||||
/// One `(account_id, model_name)` pair eligible for routing.
|
||||
#[derive(Debug, Clone)]
|
||||
struct AvailRow {
|
||||
account_id: i64,
|
||||
model_name: String,
|
||||
}
|
||||
|
||||
/// Run the rebuild against the live database. The whole thing runs inside a
|
||||
/// single transaction so a partial failure doesn't leave `token_routes` and
|
||||
/// `route_channels` out of sync.
|
||||
pub async fn rebuild_routes(state: &AppState) -> AppResult<RebuildSummary> {
|
||||
rebuild_routes_inner(&state.db).await
|
||||
}
|
||||
|
||||
/// Pool-only entry point. Identical to `rebuild_routes`, just decoupled from
|
||||
/// `AppState` so unit tests can drive it with an in-memory SQLite pool
|
||||
/// instead of building the full app config + builtin-keys boilerplate.
|
||||
pub(crate) async fn rebuild_routes_inner(
|
||||
pool: &sqlx::Pool<sqlx::Sqlite>,
|
||||
) -> AppResult<RebuildSummary> {
|
||||
let mut tx = pool.begin().await?;
|
||||
let mut summary = RebuildSummary::default();
|
||||
|
||||
// 1. Pull every (account, model) pair we should route. Joining on accounts
|
||||
// lets us drop inactive accounts and apply site_disabled_models in one
|
||||
// query — keeping the in-memory set small for installs with hundreds of
|
||||
// accounts.
|
||||
let rows: Vec<(i64, String)> = sqlx::query_as(
|
||||
"SELECT ma.account_id, ma.model_name
|
||||
FROM model_availability ma
|
||||
JOIN accounts a ON a.id = ma.account_id
|
||||
WHERE COALESCE(ma.available, 0) = 1
|
||||
AND a.status = 'active'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM site_disabled_models sdm
|
||||
WHERE sdm.site_id = a.site_id
|
||||
AND sdm.model_name = ma.model_name
|
||||
)
|
||||
ORDER BY ma.model_name, ma.account_id",
|
||||
)
|
||||
.fetch_all(&mut *tx)
|
||||
.await?;
|
||||
|
||||
let avail: Vec<AvailRow> = rows
|
||||
.into_iter()
|
||||
.map(|(account_id, model_name)| AvailRow {
|
||||
account_id,
|
||||
model_name,
|
||||
})
|
||||
.collect();
|
||||
|
||||
summary.scanned_models = avail
|
||||
.iter()
|
||||
.map(|r| r.model_name.as_str())
|
||||
.collect::<HashSet<_>>()
|
||||
.len() as i64;
|
||||
summary.scanned_accounts = avail
|
||||
.iter()
|
||||
.map(|r| r.account_id)
|
||||
.collect::<HashSet<_>>()
|
||||
.len() as i64;
|
||||
|
||||
// 2. Group by model name. Each group becomes one token_routes row.
|
||||
let mut by_model: HashMap<String, Vec<i64>> = HashMap::new();
|
||||
for row in &avail {
|
||||
by_model
|
||||
.entry(row.model_name.clone())
|
||||
.or_default()
|
||||
.push(row.account_id);
|
||||
}
|
||||
|
||||
// 3. Pre-load existing routes by model_pattern so we can reuse rows. The
|
||||
// frontend lets users rename `display_name` / pick `display_icon`; we
|
||||
// must not overwrite those on rebuild.
|
||||
let existing_routes: Vec<(String, i64)> =
|
||||
sqlx::query_as("SELECT model_pattern, id FROM token_routes")
|
||||
.fetch_all(&mut *tx)
|
||||
.await?;
|
||||
let mut route_id_by_pattern: HashMap<String, i64> =
|
||||
existing_routes.into_iter().collect();
|
||||
|
||||
// 4. For each model, ensure a token_routes row, then reconcile channels.
|
||||
for (model_name, account_ids) in &by_model {
|
||||
let route_id = match route_id_by_pattern.get(model_name) {
|
||||
Some(id) => {
|
||||
summary.reused_routes += 1;
|
||||
*id
|
||||
}
|
||||
None => {
|
||||
// route_mode = 'pattern' is the default the frontend expects.
|
||||
// routing_strategy = 'weighted' matches the schema default and
|
||||
// the dropdown's first option.
|
||||
let result = sqlx::query(
|
||||
"INSERT INTO token_routes
|
||||
(model_pattern, routing_strategy, route_mode, enabled)
|
||||
VALUES (?, 'weighted', 'pattern', 1)",
|
||||
)
|
||||
.bind(model_name)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
let id = result.last_insert_rowid();
|
||||
route_id_by_pattern.insert(model_name.clone(), id);
|
||||
summary.created_routes += 1;
|
||||
id
|
||||
}
|
||||
};
|
||||
|
||||
// Existing channels for this route.
|
||||
let existing_channels: Vec<(i64, i64, Option<String>, Option<bool>)> = sqlx::query_as(
|
||||
"SELECT id, account_id, source_model, manual_override
|
||||
FROM route_channels
|
||||
WHERE route_id = ?",
|
||||
)
|
||||
.bind(route_id)
|
||||
.fetch_all(&mut *tx)
|
||||
.await?;
|
||||
|
||||
// Index the auto-managed (non-manual) channels by (account_id, source_model).
|
||||
// Manual-override channels are tracked separately so we can leave them alone.
|
||||
type AutoKey = (i64, String);
|
||||
let mut auto_by_key: HashMap<AutoKey, i64> = HashMap::new();
|
||||
let mut manual_keys: HashSet<AutoKey> = HashSet::new();
|
||||
for (ch_id, account_id, source_model, manual) in existing_channels {
|
||||
let key = (account_id, source_model.unwrap_or_else(|| model_name.clone()));
|
||||
if manual.unwrap_or(false) {
|
||||
manual_keys.insert(key);
|
||||
summary.skipped_manual_channels += 1;
|
||||
} else {
|
||||
auto_by_key.insert(key, ch_id);
|
||||
}
|
||||
}
|
||||
|
||||
// Desired auto-managed set: every account that has this model available,
|
||||
// minus any (account, model) combo already covered by a manual channel.
|
||||
let mut desired: HashSet<AutoKey> = HashSet::new();
|
||||
for &account_id in account_ids {
|
||||
let key = (account_id, model_name.clone());
|
||||
if !manual_keys.contains(&key) {
|
||||
desired.insert(key);
|
||||
}
|
||||
}
|
||||
|
||||
// Insert missing auto channels.
|
||||
for key in &desired {
|
||||
if auto_by_key.contains_key(key) {
|
||||
continue;
|
||||
}
|
||||
// Pick a default token for this account if one exists. token_id is
|
||||
// nullable in the schema; the proxy layer already handles the
|
||||
// "account-level token" fallback when token_id is NULL.
|
||||
let token_id: Option<i64> = sqlx::query_scalar(
|
||||
"SELECT id FROM account_tokens
|
||||
WHERE account_id = ?
|
||||
AND COALESCE(enabled, 1) = 1
|
||||
ORDER BY COALESCE(is_default, 0) DESC, id ASC
|
||||
LIMIT 1",
|
||||
)
|
||||
.bind(key.0)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
"INSERT INTO route_channels
|
||||
(route_id, account_id, token_id, source_model,
|
||||
priority, weight, enabled, manual_override)
|
||||
VALUES (?, ?, ?, ?, 0, 10, 1, 0)",
|
||||
)
|
||||
.bind(route_id)
|
||||
.bind(key.0)
|
||||
.bind(token_id)
|
||||
.bind(&key.1)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
summary.created_channels += 1;
|
||||
}
|
||||
|
||||
// Delete auto channels that are no longer desired (account dropped the
|
||||
// model, or the model is now in site_disabled_models).
|
||||
for (key, ch_id) in &auto_by_key {
|
||||
if !desired.contains(key) {
|
||||
sqlx::query("DELETE FROM route_channels WHERE id = ?")
|
||||
.bind(ch_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
summary.deleted_channels += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Orphan sweep. The per-route loop above only visits routes whose
|
||||
// `model_pattern` is in `by_model`. When the last active account for a
|
||||
// given model goes away (account deleted, status flipped to non-active,
|
||||
// or the model dropped from `model_availability`), that model is not in
|
||||
// `by_model` at all — so its existing auto-managed channels never get
|
||||
// reconciled and stay behind as orphans pointing at routes nothing can
|
||||
// serve. Sweep them here. Manual-override channels are still preserved.
|
||||
let processed_route_ids: HashSet<i64> = by_model
|
||||
.keys()
|
||||
.filter_map(|model_name| route_id_by_pattern.get(model_name).copied())
|
||||
.collect();
|
||||
|
||||
let orphan_candidates: Vec<(i64, i64)> = sqlx::query_as(
|
||||
"SELECT id, route_id FROM route_channels
|
||||
WHERE COALESCE(manual_override, 0) = 0",
|
||||
)
|
||||
.fetch_all(&mut *tx)
|
||||
.await?;
|
||||
|
||||
for (ch_id, route_id) in orphan_candidates {
|
||||
if !processed_route_ids.contains(&route_id) {
|
||||
sqlx::query("DELETE FROM route_channels WHERE id = ?")
|
||||
.bind(ch_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
summary.deleted_channels += 1;
|
||||
}
|
||||
}
|
||||
|
||||
tx.commit().await?;
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use sqlx::sqlite::SqlitePoolOptions;
|
||||
use sqlx::{Pool, Sqlite};
|
||||
|
||||
/// Minimal schema for the rebuild path. We only create the tables and
|
||||
/// columns `rebuild_routes_inner` actually reads/writes — the rest of the
|
||||
/// app's schema is irrelevant here and pulling in the full migration set
|
||||
/// would couple the test to unrelated migrations.
|
||||
async fn fixture_pool() -> Pool<Sqlite> {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.expect("open in-memory sqlite");
|
||||
|
||||
for ddl in [
|
||||
"CREATE TABLE accounts (
|
||||
id INTEGER PRIMARY KEY,
|
||||
site_id INTEGER NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'active'
|
||||
)",
|
||||
"CREATE TABLE model_availability (
|
||||
account_id INTEGER NOT NULL,
|
||||
model_name TEXT NOT NULL,
|
||||
available INTEGER NOT NULL DEFAULT 1,
|
||||
PRIMARY KEY (account_id, model_name)
|
||||
)",
|
||||
"CREATE TABLE site_disabled_models (
|
||||
site_id INTEGER NOT NULL,
|
||||
model_name TEXT NOT NULL,
|
||||
PRIMARY KEY (site_id, model_name)
|
||||
)",
|
||||
"CREATE TABLE token_routes (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
model_pattern TEXT NOT NULL UNIQUE,
|
||||
routing_strategy TEXT NOT NULL DEFAULT 'weighted',
|
||||
route_mode TEXT NOT NULL DEFAULT 'pattern',
|
||||
enabled INTEGER NOT NULL DEFAULT 1
|
||||
)",
|
||||
"CREATE TABLE route_channels (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
route_id INTEGER NOT NULL,
|
||||
account_id INTEGER NOT NULL,
|
||||
token_id INTEGER,
|
||||
source_model TEXT,
|
||||
priority INTEGER DEFAULT 0,
|
||||
weight INTEGER DEFAULT 10,
|
||||
enabled INTEGER DEFAULT 1,
|
||||
manual_override INTEGER DEFAULT 0
|
||||
)",
|
||||
"CREATE TABLE account_tokens (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
account_id INTEGER NOT NULL,
|
||||
enabled INTEGER DEFAULT 1,
|
||||
is_default INTEGER DEFAULT 0
|
||||
)",
|
||||
] {
|
||||
sqlx::query(ddl)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap_or_else(|err| panic!("failed to create schema: {err}\nddl: {ddl}"));
|
||||
}
|
||||
|
||||
pool
|
||||
}
|
||||
|
||||
/// Helper: count auto-managed (non-manual) channels. Manual ones are out
|
||||
/// of scope of the rebuild — we never touch them so they shouldn't be
|
||||
/// part of any assertion about cleanup.
|
||||
async fn auto_channel_count(pool: &Pool<Sqlite>) -> i64 {
|
||||
sqlx::query_scalar::<_, i64>(
|
||||
"SELECT COUNT(*) FROM route_channels WHERE COALESCE(manual_override, 0) = 0",
|
||||
)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Reproduces #34. Account 1 advertises model `gpt-x` and the first
|
||||
/// rebuild creates a route + channel for it. Then account 1 flips to
|
||||
/// `disabled` so the model is no longer routable from any active
|
||||
/// account. Before the orphan sweep, the channel survived because the
|
||||
/// per-route loop only visited routes that appeared in `by_model`, and
|
||||
/// `gpt-x` no longer did. After the sweep, the channel must be gone.
|
||||
#[tokio::test]
|
||||
async fn orphan_channel_is_cleaned_when_last_account_goes_inactive() {
|
||||
let pool = fixture_pool().await;
|
||||
|
||||
sqlx::query("INSERT INTO accounts (id, site_id, status) VALUES (1, 1, 'active')")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"INSERT INTO model_availability (account_id, model_name, available) VALUES (1, 'gpt-x', 1)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// First pass: route + channel get created.
|
||||
let summary = rebuild_routes_inner(&pool).await.expect("first rebuild");
|
||||
assert_eq!(summary.created_routes, 1, "route should be created");
|
||||
assert_eq!(summary.created_channels, 1, "channel should be created");
|
||||
assert_eq!(auto_channel_count(&pool).await, 1);
|
||||
|
||||
// Account flips to disabled. `rebuild_routes` filters on
|
||||
// `accounts.status = 'active'`, so no row from this account makes it
|
||||
// into the working set — `gpt-x` falls out of `by_model` entirely.
|
||||
sqlx::query("UPDATE accounts SET status = 'disabled' WHERE id = 1")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let summary = rebuild_routes_inner(&pool).await.expect("second rebuild");
|
||||
// Without the orphan sweep this would be 0 and the channel would
|
||||
// remain. With the sweep, the channel pointing at an account that's
|
||||
// no longer routable should be hard-deleted.
|
||||
assert_eq!(
|
||||
auto_channel_count(&pool).await,
|
||||
0,
|
||||
"orphan channel must be cleaned up after last active account drops out"
|
||||
);
|
||||
assert_eq!(
|
||||
summary.deleted_channels, 1,
|
||||
"summary should report exactly one deletion"
|
||||
);
|
||||
// The route row itself stays — operators may have edited display
|
||||
// metadata on it, and the next time an active account advertises
|
||||
// `gpt-x` we want to reuse it. The rebuild never deletes routes.
|
||||
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(route_count, 1, "token_routes row should be reused, not dropped");
|
||||
}
|
||||
|
||||
/// Sanity check: manual_override channels survive an orphan sweep. Even
|
||||
/// when the rebuild has no reason to keep them (no active account, no
|
||||
/// availability row), they're operator-curated and the rebuild contract
|
||||
/// says we never touch them.
|
||||
#[tokio::test]
|
||||
async fn orphan_sweep_preserves_manual_channels() {
|
||||
let pool = fixture_pool().await;
|
||||
|
||||
sqlx::query("INSERT INTO accounts (id, site_id, status) VALUES (1, 1, 'disabled')")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("INSERT INTO token_routes (id, model_pattern) VALUES (1, 'manual-only')")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"INSERT INTO route_channels (route_id, account_id, source_model, manual_override)
|
||||
VALUES (1, 1, 'manual-only', 1)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let summary = rebuild_routes_inner(&pool).await.expect("rebuild");
|
||||
assert_eq!(summary.deleted_channels, 0, "manual channel must survive");
|
||||
|
||||
let manual_count: i64 = sqlx::query_scalar(
|
||||
"SELECT COUNT(*) FROM route_channels WHERE manual_override = 1",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(manual_count, 1);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user