fix(rust): trigger route rebuild on account writes; add segment filter

- spawn_route_rebuild after create / delete / model refresh / manual model
  upsert; gate update_account on status change so cosmetic edits don't
  burn a transaction
- list_accounts now accepts ?segment=apikey|session|oauth, filtered against
  extra_config.$.credentialMode (legacy null rows treated as apikey)
This commit is contained in:
boos
2026-05-20 19:35:53 +08:00
parent 7c9a1e2cc4
commit cf696c01a5
+311 -13
View File
@@ -16,6 +16,10 @@ use tracing::info;
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
pub struct ListAccountsQuery { pub struct ListAccountsQuery {
pub site_id: Option<i64>, pub site_id: Option<i64>,
/// Filter by credential mode stored in `accounts.extra_config.$.credentialMode`.
/// Frontend uses this as `?segment=apikey` / `?segment=session` / `?segment=oauth`
/// on the accounts list page so each segment only sees its own rows.
pub segment: Option<String>,
} }
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
@@ -89,20 +93,45 @@ pub async fn list_accounts(
State(state): State<AppState>, State(state): State<AppState>,
Query(query): Query<ListAccountsQuery>, Query(query): Query<ListAccountsQuery>,
) -> AppResult<Json<Value>> { ) -> AppResult<Json<Value>> {
info!("Listing accounts, site_id filter: {:?}", query.site_id); info!(
"Listing accounts, site_id filter: {:?}, segment filter: {:?}",
query.site_id, query.segment
);
let accounts: Vec<Account> = if let Some(site_id) = query.site_id { // Build a single dynamic query so site_id and segment compose. `segment`
sqlx::query_as::<_, Account>( // matches the JSON field `extra_config.$.credentialMode`; for legacy rows
"SELECT * FROM accounts WHERE site_id = ? ORDER BY sort_order, id", // that have no extra_config we treat them as 'apikey' (the historical
) // default before the credentialMode field existed) so the apikey segment
.bind(site_id) // doesn't accidentally hide them.
.fetch_all(&*state.db) let mut sql = String::from("SELECT * FROM accounts WHERE 1=1");
.await? if query.site_id.is_some() {
sql.push_str(" AND site_id = ?");
}
let segment = query
.segment
.as_deref()
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::to_ascii_lowercase);
if let Some(ref seg) = segment {
if seg == "apikey" {
sql.push_str(
" AND COALESCE(LOWER(json_extract(extra_config, '$.credentialMode')), 'apikey') = ?",
);
} else { } else {
sqlx::query_as::<_, Account>("SELECT * FROM accounts ORDER BY site_id, sort_order, id") sql.push_str(" AND LOWER(json_extract(extra_config, '$.credentialMode')) = ?");
.fetch_all(&*state.db) }
.await? }
}; sql.push_str(" ORDER BY site_id, sort_order, id");
let mut q = sqlx::query_as::<_, Account>(&sql);
if let Some(site_id) = query.site_id {
q = q.bind(site_id);
}
if let Some(ref seg) = segment {
q = q.bind(seg);
}
let accounts: Vec<Account> = q.fetch_all(&*state.db).await?;
let sites: Vec<Site> = let sites: Vec<Site> =
sqlx::query_as::<_, Site>("SELECT * FROM sites ORDER BY sort_order, name") sqlx::query_as::<_, Site>("SELECT * FROM sites ORDER BY sort_order, name")
@@ -226,6 +255,9 @@ pub async fn create_account(
} }
if created.len() + failed.len() > 1 { if created.len() + failed.len() > 1 {
if !created.is_empty() {
spawn_route_rebuild(state.clone());
}
return Ok(( return Ok((
StatusCode::CREATED, StatusCode::CREATED,
Json(json!({ Json(json!({
@@ -245,6 +277,8 @@ pub async fn create_account(
.next() .next()
.ok_or_else(|| AppError::Internal(anyhow::anyhow!("account was not created")))?; .ok_or_else(|| AppError::Internal(anyhow::anyhow!("account was not created")))?;
spawn_route_rebuild(state.clone());
Ok(( Ok((
StatusCode::CREATED, StatusCode::CREATED,
Json(json!({ Json(json!({
@@ -255,6 +289,32 @@ pub async fn create_account(
)) ))
} }
/// Fire a route rebuild in the background after account creation. We don't
/// block the HTTP response on this — `model_availability` is filled by the
/// separate `POST /api/accounts/{id}/models/refresh` call the frontend makes
/// next, so the rebuild that matters most happens after that refresh, not
/// here. This still covers the case where models were already populated for
/// re-added accounts, plus it keeps the proposal's "auto-rebuild" promise.
/// Errors are logged and swallowed.
fn spawn_route_rebuild(state: AppState) {
tokio::spawn(async move {
match crate::services::route_rebuild::rebuild_routes(&state).await {
Ok(summary) => {
info!(
created_routes = summary.created_routes,
reused_routes = summary.reused_routes,
created_channels = summary.created_channels,
deleted_channels = summary.deleted_channels,
"route rebuild completed after account creation"
);
}
Err(err) => {
tracing::warn!(error = %err, "route rebuild after account creation failed");
}
}
});
}
/// POST /api/accounts/verify-token /// POST /api/accounts/verify-token
pub async fn verify_token( pub async fn verify_token(
State(state): State<AppState>, State(state): State<AppState>,
@@ -394,13 +454,20 @@ pub async fn update_account(
.await? .await?
.ok_or_else(|| AppError::NotFound(format!("Account with id {} not found", id)))?; .ok_or_else(|| AppError::NotFound(format!("Account with id {} not found", id)))?;
// Resolve the final status before the UPDATE so we can compare it against
// the previous value and decide whether routing needs to be rebuilt.
// `rebuild_routes` filters on `accounts.status = 'active'`, so any flip
// into or out of `active` changes which channels should exist.
let new_status = payload.status.clone().unwrap_or_else(|| existing.status.clone());
let status_changed = new_status != existing.status;
sqlx::query( sqlx::query(
"UPDATE accounts SET username = ?, access_token = ?, api_token = ?, status = ?, checkin_enabled = ?, extra_config = ?, oauth_provider = ?, oauth_account_key = ?, oauth_project_id = ?, updated_at = datetime('now') WHERE id = ?" "UPDATE accounts SET username = ?, access_token = ?, api_token = ?, status = ?, checkin_enabled = ?, extra_config = ?, oauth_provider = ?, oauth_account_key = ?, oauth_project_id = ?, updated_at = datetime('now') WHERE id = ?"
) )
.bind(payload.username.or(existing.username)) .bind(payload.username.or(existing.username))
.bind(payload.access_token.as_ref().unwrap_or(&existing.access_token)) .bind(payload.access_token.as_ref().unwrap_or(&existing.access_token))
.bind(payload.api_token.or(existing.api_token)) .bind(payload.api_token.or(existing.api_token))
.bind(payload.status.unwrap_or(existing.status)) .bind(&new_status)
.bind(payload.checkin_enabled.or(existing.checkin_enabled)) .bind(payload.checkin_enabled.or(existing.checkin_enabled))
.bind(payload.extra_config.or(existing.extra_config)) .bind(payload.extra_config.or(existing.extra_config))
.bind(payload.oauth_provider.or(existing.oauth_provider)) .bind(payload.oauth_provider.or(existing.oauth_provider))
@@ -415,6 +482,15 @@ pub async fn update_account(
.fetch_one(&*state.db) .fetch_one(&*state.db)
.await?; .await?;
// Toggling status into or out of `active` changes the routing-eligible set;
// route_channels for the deactivated account need to be hard-deleted, and
// a reactivated account needs its channels (re)created. Other field edits
// (token, username, oauth metadata) don't affect routing — skip the rebuild
// there to avoid burning a transaction on every cosmetic edit.
if status_changed {
spawn_route_rebuild(state.clone());
}
Ok(Json(json!({ "data": account }))) Ok(Json(json!({ "data": account })))
} }
@@ -438,6 +514,13 @@ pub async fn delete_account(
))); )));
} }
// Mirror the Node implementation (`accounts.ts` DELETE handler): after the
// account row is gone its auto-managed channels become orphaned references
// — `rebuild_routes` will hard-delete them on the next pass. Manual-override
// channels are left alone (they can still target a deleted account if the
// operator chose to keep them).
spawn_route_rebuild(state.clone());
Ok(Json(json!({ "data": { "deleted": true, "id": id } }))) Ok(Json(json!({ "data": { "deleted": true, "id": id } })))
} }
@@ -650,5 +733,220 @@ pub async fn add_account_manual_models(
inserted += 1; inserted += 1;
} }
// Manual models just landed in `model_availability` — same trigger point as
// refresh_account_models. Mirrors the Node handler that calls
// `rebuildRoutesBestEffort()` after upserting manual model rows.
if inserted > 0 {
spawn_route_rebuild(state.clone());
}
Ok(Json(json!({ "inserted": inserted }))) Ok(Json(json!({ "inserted": inserted })))
} }
/// POST /api/accounts/{id}/models/refresh — re-discover models from the upstream
/// for one account. Calls `{site_base}/models` with the account's access_token,
/// upserts results into `model_availability`, and returns a shape compatible
/// with the frontend's `checkModels` flow:
///
/// ```json
/// {
/// "message": "...",
/// "refresh": {
/// "status": "success" | "error",
/// "modelCount": 12,
/// "modelsPreview": ["gpt-4o", "..."],
/// "errorCode": "timeout" | "unauthorized" | "empty_models" | "upstream_error",
/// "errorMessage": "..."
/// }
/// }
/// ```
pub async fn refresh_account_models(
State(state): State<AppState>,
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
info!("Refreshing upstream models for account: {}", id);
// 1) Account row (need site_id + access_token).
let account: Account = sqlx::query_as::<_, Account>("SELECT * FROM accounts WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Account {} not found", id)))?;
let access_token = account.access_token.trim();
if access_token.is_empty() {
return Ok(Json(json!({
"message": "账户未配置 access_token",
"refresh": {
"status": "error",
"errorCode": "unauthorized",
"errorMessage": "missing access_token",
"modelCount": 0,
"modelsPreview": [],
}
})));
}
// 2) Site + base url (reuses the same selector as verify_token so a site's
// active api endpoint is honored).
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
.bind(account.site_id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| {
AppError::NotFound(format!("Site {} not found for account {}", account.site_id, id))
})?;
let base_url = resolve_site_api_base_url(&state, &site).await?;
let models_url = format!("{}/models", base_url.trim_end_matches('/'));
// 3) Upstream GET /models with the account bearer.
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(20))
.build()
.map_err(|err| AppError::Internal(err.into()))?;
let response = client
.get(&models_url)
.bearer_auth(access_token)
.header(reqwest::header::ACCEPT, "application/json")
.send()
.await;
let response = match response {
Ok(resp) => resp,
Err(err) => {
let timed_out = err.is_timeout();
let code = if timed_out { "timeout" } else { "network" };
return Ok(Json(json!({
"message": format!("无法访问上游 /models:{}", err),
"refresh": {
"status": "error",
"errorCode": code,
"errorMessage": err.to_string(),
"modelCount": 0,
"modelsPreview": [],
"siteId": site.id,
"baseUrl": base_url,
}
})));
}
};
let status = response.status();
if !status.is_success() {
let status_code = status.as_u16();
let body_text = response.text().await.unwrap_or_default();
let code = match status_code {
401 | 403 => "unauthorized",
408 | 504 => "timeout",
_ => "upstream_error",
};
return Ok(Json(json!({
"message": format!("上游 /models 返回 HTTP {}", status_code),
"refresh": {
"status": "error",
"errorCode": code,
"errorMessage": body_text.chars().take(240).collect::<String>(),
"upstreamStatus": status_code,
"modelCount": 0,
"modelsPreview": [],
"siteId": site.id,
"baseUrl": base_url,
}
})));
}
// 4) Parse models from `data[].id` (with `name` fallback, matching verify_token).
let body: Value = response.json().await.unwrap_or_else(|_| json!({}));
let models: Vec<String> = body
.get("data")
.and_then(Value::as_array)
.map(|items| {
items
.iter()
.filter_map(|item| {
item.get("id")
.or_else(|| item.get("name"))
.and_then(Value::as_str)
.map(str::to_string)
})
.filter(|m| !m.trim().is_empty())
.collect::<Vec<_>>()
})
.unwrap_or_default();
if models.is_empty() {
return Ok(Json(json!({
"message": "上游未返回任何模型",
"refresh": {
"status": "error",
"errorCode": "empty_models",
"errorMessage": "upstream returned 0 models",
"modelCount": 0,
"modelsPreview": [],
"siteId": site.id,
"baseUrl": base_url,
}
})));
}
// 5) Upsert into model_availability (is_manual=0 — these came from the upstream).
// Manual entries are left untouched: the ON CONFLICT clause does not lower
// `is_manual` for rows that are already 1.
let mut upserted = 0usize;
for model_name in models.iter().map(|s| s.trim()).filter(|s| !s.is_empty()) {
sqlx::query(
"INSERT INTO model_availability (account_id, model_name, available, is_manual, checked_at)
VALUES (?, ?, 1, 0, datetime('now'))
ON CONFLICT (account_id, model_name) DO UPDATE SET
available = 1,
checked_at = datetime('now')",
)
.bind(id)
.bind(model_name)
.execute(&*state.db)
.await?;
upserted += 1;
}
// 6) Mark non-manual rows that disappeared from the upstream as unavailable,
// so stale entries don't keep showing up. Manual rows are preserved.
let placeholders = std::iter::repeat_n("?", models.len())
.collect::<Vec<_>>()
.join(",");
let stale_sql = format!(
"UPDATE model_availability
SET available = 0, checked_at = datetime('now')
WHERE account_id = ?
AND COALESCE(is_manual, 0) = 0
AND model_name NOT IN ({})",
placeholders
);
let mut stale = sqlx::query(&stale_sql).bind(id);
for m in &models {
stale = stale.bind(m);
}
stale.execute(&*state.db).await?;
let preview: Vec<String> = models.iter().take(8).cloned().collect();
// Models just landed in `model_availability` — this is the moment the
// rebuild can actually produce token_routes / route_channels. The fire-
// and-forget call in create_account ran too early (before refresh) so
// without this hook /routes stays empty until the user clicks rebuild
// manually. Errors are logged inside spawn_route_rebuild.
spawn_route_rebuild(state.clone());
Ok(Json(json!({
"message": "模型刷新成功",
"refresh": {
"status": "success",
"modelCount": models.len(),
"modelsPreview": preview,
"upserted": upserted,
"siteId": site.id,
"baseUrl": base_url,
}
})))
}