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
2026-05-20 19:58:45 +08:00
parent 423912277b
commit c16479da7c
+435
View File
@@ -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);
}
}