chore: sync current workspace changes
Desktop CI Build / Build macOS ARM64 DMG (push) Has been cancelled
Desktop CI Build / Build Windows x64 (push) Has been cancelled

This commit is contained in:
boos
2026-07-14 03:21:10 +08:00
parent 1512fa23ce
commit 557914fea8
745 changed files with 124917 additions and 11222 deletions
+13 -5
View File
@@ -1,6 +1,6 @@
use argon2::password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString};
use argon2::Argon2;
use sqlx::SqlitePool;
use sqlx::AnyPool;
use crate::db::models::downstream_api_key::DownstreamApiKey;
@@ -20,7 +20,7 @@ pub enum DownstreamAuthResult {
/// 2. Falls back to checking against the global proxy token
pub async fn authorize_downstream_token(
token: &str,
pool: &SqlitePool,
pool: &AnyPool,
global_proxy_token: Option<&str>,
) -> Result<DownstreamAuthResult, String> {
if global_proxy_token
@@ -34,8 +34,15 @@ pub async fn authorize_downstream_token(
// Try managed keys first
if token.starts_with("sk-") {
let key: Option<DownstreamApiKey> = sqlx::query_as::<_, DownstreamApiKey>(
"SELECT * FROM downstream_api_keys WHERE key = ? AND enabled = 1",
let key: Option<DownstreamApiKey> = crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT * FROM downstream_api_keys
WHERE key = ?
AND enabled = 1
AND (
expires_at IS NULL
OR TRIM(expires_at) = ''
OR datetime(expires_at) > CURRENT_TIMESTAMP
)",
)
.bind(token)
.fetch_optional(pool)
@@ -112,7 +119,8 @@ mod tests {
#[tokio::test]
async fn test_global_proxy_token_authorizes_without_db_lookup() {
let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
sqlx::any::install_default_drivers();
let pool = AnyPool::connect("sqlite::memory:").await.unwrap();
let result = authorize_downstream_token("sk-global", &pool, Some("sk-global"))
.await
.unwrap();
+623 -6
View File
@@ -3,7 +3,13 @@ use figment::{
Figment,
};
use serde::Deserialize;
use std::path::PathBuf;
use std::{
io::Write,
path::{Path, PathBuf},
sync::OnceLock,
};
static CONFIG_SOURCE_PATH: OnceLock<PathBuf> = OnceLock::new();
#[derive(Debug, Deserialize, Clone)]
pub struct AppConfig {
@@ -14,11 +20,123 @@ pub struct AppConfig {
#[serde(default)]
pub comfyui: ComfyUIConfig,
#[serde(default)]
pub agent: AgentConfig,
#[serde(default)]
pub pixelle_video: PixelleVideoConfig,
#[serde(default)]
pub payment: PaymentConfig,
#[serde(default)]
pub jimeng: JimengConfig,
#[serde(default)]
pub seedance: SeedanceConfig,
}
/// Optional Redis 7+ integration. The connection is deliberately opt-in so
/// a single-process local installation can run without an additional service.
#[derive(Debug, Deserialize, Clone)]
pub struct RedisConfig {
#[serde(default)]
pub enabled: bool,
#[serde(default = "default_redis_url")]
pub url: String,
#[serde(default = "default_redis_key_prefix")]
pub key_prefix: String,
#[serde(default = "default_redis_connection_timeout_ms")]
pub connection_timeout_ms: u64,
#[serde(default = "default_redis_response_timeout_ms")]
pub response_timeout_ms: u64,
/// Required only when Redis is enabled. A 32-byte secret encoded as 64 hex
/// characters, used solely for short-lived shared authentication state.
#[serde(default)]
pub auth_state_secret: Option<String>,
}
#[derive(Debug, Deserialize)]
struct RedisConfigFile {
#[serde(default)]
redis: RedisConfig,
}
fn default_redis_url() -> String {
"redis://127.0.0.1:6379/0".to_string()
}
fn default_redis_key_prefix() -> String {
"boosapi".to_string()
}
fn default_redis_connection_timeout_ms() -> u64 {
3_000
}
fn default_redis_response_timeout_ms() -> u64 {
2_000
}
impl Default for RedisConfig {
fn default() -> Self {
Self {
enabled: false,
url: default_redis_url(),
key_prefix: default_redis_key_prefix(),
connection_timeout_ms: default_redis_connection_timeout_ms(),
response_timeout_ms: default_redis_response_timeout_ms(),
auth_state_secret: None,
}
}
}
impl RedisConfig {
pub fn load_from_file(path: &str) -> anyhow::Result<Self> {
dotenvy::dotenv().ok();
let _ = CONFIG_SOURCE_PATH.set(PathBuf::from(path));
let config = Figment::new()
.merge(Toml::file(path))
.merge(Env::prefixed("BOOSAPI_").split("__"))
.extract::<RedisConfigFile>()?
.redis;
config.validate()?;
Ok(config)
}
pub fn validate(&self) -> anyhow::Result<()> {
anyhow::ensure!(
self.connection_timeout_ms > 0,
"redis.connection_timeout_ms must be greater than zero"
);
anyhow::ensure!(
self.response_timeout_ms > 0,
"redis.response_timeout_ms must be greater than zero"
);
Ok(())
}
}
#[derive(Debug, Deserialize, Clone)]
pub struct AgentConfig {
/// Optional dedicated default. Empty/unset falls back to `comfyui.agent_model`.
#[serde(default)]
pub default_model: Option<String>,
#[serde(default = "default_agent_max_turns")]
pub max_turns: u32,
#[serde(default = "default_agent_context_budget_tokens")]
pub context_budget_tokens: usize,
#[serde(default = "default_agent_interaction_timeout_secs")]
pub interaction_timeout_secs: u64,
#[serde(default = "default_agent_max_concurrent_runs")]
pub max_concurrent_runs_per_user: u32,
#[serde(default = "default_agent_tools")]
pub tools: Vec<String>,
}
impl AgentConfig {
pub fn resolved_default_model<'a>(&'a self, comfyui: &'a ComfyUIConfig) -> &'a str {
self.default_model
.as_deref()
.map(str::trim)
.filter(|model| !model.is_empty())
.unwrap_or(&comfyui.agent_model)
}
}
#[derive(Debug, Deserialize, Clone)]
@@ -33,12 +151,176 @@ pub struct ServerConfig {
pub struct DbConfig {
pub url: String,
pub max_connections: u32,
#[serde(default)]
pub ssl: bool,
#[serde(default)]
pub ssl_mode: Option<DbSslMode>,
#[serde(default)]
pub ca_cert_path: Option<String>,
pub runtime_db_url: Option<String>,
pub runtime_db_dialect: Option<String>,
pub runtime_db_ssl: Option<bool>,
}
#[derive(Debug, Deserialize, Clone)]
pub struct PgrustConfig {
#[serde(default = "default_pgrust_enabled")]
pub enabled: bool,
#[serde(default)]
pub runtime_dir: Option<PathBuf>,
#[serde(default = "default_pgrust_data_dir")]
pub data_dir: PathBuf,
#[serde(default = "default_pgrust_port")]
pub port: u16,
#[serde(default = "default_pgrust_startup_timeout_secs")]
pub startup_timeout_secs: u64,
#[serde(default = "default_pgrust_shutdown_timeout_secs")]
pub shutdown_timeout_secs: u64,
#[serde(default = "default_pgrust_run_uid")]
pub run_uid: u32,
#[serde(default = "default_pgrust_run_gid")]
pub run_gid: u32,
}
#[derive(Debug, Deserialize)]
struct PgrustConfigFile {
#[serde(default)]
pgrust: PgrustConfig,
}
fn default_pgrust_enabled() -> bool {
false
}
fn default_pgrust_data_dir() -> PathBuf {
PathBuf::from("data/pgrust")
}
fn default_pgrust_port() -> u16 {
6432
}
fn default_pgrust_startup_timeout_secs() -> u64 {
120
}
fn default_pgrust_shutdown_timeout_secs() -> u64 {
15
}
fn default_pgrust_run_uid() -> u32 {
65_534
}
fn default_pgrust_run_gid() -> u32 {
65_534
}
impl Default for PgrustConfig {
fn default() -> Self {
Self {
enabled: default_pgrust_enabled(),
runtime_dir: None,
data_dir: default_pgrust_data_dir(),
port: default_pgrust_port(),
startup_timeout_secs: default_pgrust_startup_timeout_secs(),
shutdown_timeout_secs: default_pgrust_shutdown_timeout_secs(),
run_uid: default_pgrust_run_uid(),
run_gid: default_pgrust_run_gid(),
}
}
}
impl PgrustConfig {
pub fn load_from_file(path: &str) -> anyhow::Result<Self> {
dotenvy::dotenv().ok();
let _ = CONFIG_SOURCE_PATH.set(PathBuf::from(path));
let config = Figment::new()
.merge(Toml::file(path))
.merge(Env::prefixed("BOOSAPI_").split("__"))
.extract::<PgrustConfigFile>()?
.pgrust;
config.validate()?;
Ok(config)
}
pub fn validate(&self) -> anyhow::Result<()> {
anyhow::ensure!(self.port > 0, "pgrust.port must be greater than zero");
anyhow::ensure!(
self.startup_timeout_secs > 0,
"pgrust.startup_timeout_secs must be greater than zero"
);
anyhow::ensure!(
self.shutdown_timeout_secs > 0,
"pgrust.shutdown_timeout_secs must be greater than zero"
);
Ok(())
}
}
#[derive(Debug, Deserialize, Clone, Copy, PartialEq, Eq)]
#[serde(rename_all = "kebab-case")]
pub enum DbSslMode {
Disable,
Prefer,
Require,
VerifyCa,
VerifyFull,
}
impl DbConfig {
pub fn resolved_url(&self) -> anyhow::Result<String> {
let normalized = self.url.trim().to_ascii_lowercase();
if normalized.starts_with("sqlite:") {
return Ok(self.url.clone());
}
let mut parsed = url::Url::parse(&self.url)
.map_err(|error| anyhow::anyhow!("invalid database URL: {error}"))?;
let mysql = normalized.starts_with("mysql:") || normalized.starts_with("mariadb:");
let mode = self.ssl_mode.unwrap_or(if self.ssl {
DbSslMode::Require
} else {
DbSslMode::Disable
});
let mode_key = if mysql { "ssl-mode" } else { "sslmode" };
let ca_key = if mysql { "ssl-ca" } else { "sslrootcert" };
let has_mode = parsed.query_pairs().any(|(key, _)| key == mode_key);
let has_ca = parsed.query_pairs().any(|(key, _)| key == ca_key);
if !has_mode {
let value = if mysql {
match mode {
DbSslMode::Disable => "DISABLED",
DbSslMode::Prefer => "PREFERRED",
DbSslMode::Require => "REQUIRED",
DbSslMode::VerifyCa => "VERIFY_CA",
DbSslMode::VerifyFull => "VERIFY_IDENTITY",
}
} else {
match mode {
DbSslMode::Disable => "disable",
DbSslMode::Prefer => "prefer",
DbSslMode::Require => "require",
DbSslMode::VerifyCa => "verify-ca",
DbSslMode::VerifyFull => "verify-full",
}
};
parsed.query_pairs_mut().append_pair(mode_key, value);
}
if !has_ca {
if let Some(path) = self
.ca_cert_path
.as_deref()
.map(str::trim)
.filter(|v| !v.is_empty())
{
parsed.query_pairs_mut().append_pair(ca_key, path);
}
}
Ok(parsed.into())
}
}
#[derive(Debug, Deserialize, Clone)]
pub struct AuthConfig {
pub jwt_secret: String,
@@ -86,6 +368,11 @@ pub struct ComfyUIConfig {
#[derive(Debug, Deserialize, Clone)]
pub struct PixelleVideoConfig {
pub enabled: bool,
/// Execution engine for Pixelle-compatible jobs. `external` keeps the
/// Python service adapter; `api-native` uses BoosAPI's configured cloud
/// providers and does not require a local GPU.
#[serde(default = "default_pixelle_execution_mode")]
pub execution_mode: String,
pub api_base_url: Option<String>,
pub api_key: Option<String>,
#[serde(default = "default_pixelle_backend_provider")]
@@ -152,21 +439,74 @@ pub struct JimengConfig {
pub video_cost: f64,
}
/// Volcengine Ark Seedance asynchronous video generation provider.
///
/// `api_key` must be supplied via `BOOSAPI_SEEDANCE__API_KEY` (or an active
/// `sites.platform='seedance'` account). It must never be committed to TOML.
#[derive(Debug, Deserialize, Clone)]
pub struct SeedanceConfig {
#[serde(default)]
pub enabled: bool,
#[serde(default)]
pub api_key: String,
#[serde(default = "default_seedance_base_url")]
pub base_url: String,
#[serde(default = "default_seedance_model")]
pub default_model: String,
#[serde(default = "default_seedance_poll_interval_secs")]
pub poll_interval_secs: u64,
#[serde(default = "default_seedance_timeout_secs")]
pub timeout_secs: u64,
#[serde(default = "default_seedance_execution_expires_after")]
pub execution_expires_after: u64,
#[serde(default = "default_seedance_max_output_bytes")]
pub max_output_bytes: usize,
#[serde(default = "default_seedance_output_host_allowlist")]
pub output_host_allowlist: Vec<String>,
#[serde(default)]
pub cost_per_million_tokens: f64,
}
#[derive(Debug, Deserialize, Clone)]
pub struct PaymentConfig {
pub enabled: bool,
#[serde(default)]
pub provider: PaymentProvider,
/// Independent 32-byte key used only for encrypting database-backed
/// payment-provider credentials. Configure through
/// BOOSAPI_PAYMENT__CREDENTIAL_MASTER_KEY (base64 or 64-char hex).
pub credential_master_key: Option<String>,
pub public_base_url: Option<String>,
pub product_name_prefix: String,
pub currency: String,
/// API quota granted for one unit of the configured payment currency.
///
/// USD payments default to 1:1 for backwards compatibility. Non-USD
/// payment providers must set this explicitly so a payment amount is never
/// silently treated as a USD quota amount.
pub credit_per_payment_unit: Option<f64>,
#[serde(default = "default_payment_credit_unit")]
pub payment_credit_unit: String,
pub min_amount_usd: f64,
pub max_amount_usd: Option<f64>,
pub daily_limit_usd: Option<f64>,
pub order_timeout_minutes: i64,
pub max_pending_orders: i64,
#[serde(default = "default_refund_request_window_days")]
pub refund_request_window_days: i64,
#[serde(default = "default_refund_provider_timeout_secs")]
pub refund_provider_timeout_secs: u64,
pub stripe_secret_key: Option<String>,
pub stripe_webhook_secret: Option<String>,
pub alipay_app_id: Option<String>,
pub alipay_private_key: Option<String>,
pub alipay_public_key: Option<String>,
/// Optional merchant/seller ID expected in Alipay notifications.
pub alipay_seller_id: Option<String>,
#[serde(default = "default_alipay_gateway_url")]
pub alipay_gateway_url: String,
#[serde(default = "default_alipay_notify_path")]
pub alipay_notify_path: String,
pub comfy_customer_token: Option<String>,
pub comfy_customer_token_endpoint: Option<String>,
pub comfy_firebase_api_key: Option<String>,
@@ -192,6 +532,7 @@ pub enum PaymentProvider {
#[default]
ComfyCloud,
Stripe,
Alipay,
}
#[derive(Debug, Deserialize, Clone, Default)]
@@ -213,10 +554,58 @@ fn default_comfyui_container_name() -> String {
"boosapi-comfyui-workflow".to_string()
}
fn default_agent_max_turns() -> u32 {
20
}
fn default_agent_context_budget_tokens() -> usize {
64_000
}
fn default_agent_interaction_timeout_secs() -> u64 {
300
}
fn default_agent_max_concurrent_runs() -> u32 {
2
}
fn default_agent_tools() -> Vec<String> {
[
"Read",
"Write",
"Edit",
"Glob",
"Grep",
"AskUserQuestion",
"Bash",
]
.into_iter()
.map(str::to_string)
.collect()
}
impl Default for AgentConfig {
fn default() -> Self {
Self {
default_model: None,
max_turns: default_agent_max_turns(),
context_budget_tokens: default_agent_context_budget_tokens(),
interaction_timeout_secs: default_agent_interaction_timeout_secs(),
max_concurrent_runs_per_user: default_agent_max_concurrent_runs(),
tools: default_agent_tools(),
}
}
}
fn default_pixelle_backend_provider() -> String {
"comfy".to_string()
}
fn default_pixelle_execution_mode() -> String {
"external".to_string()
}
fn default_pixelle_config_path() -> String {
"/opt/Pixelle-Video/config.yaml".to_string()
}
@@ -249,6 +638,37 @@ fn default_jimeng_video_timeout_secs() -> u64 {
600
}
fn default_seedance_base_url() -> String {
"https://ark.cn-beijing.volces.com/api/v3".to_string()
}
fn default_seedance_model() -> String {
"doubao-seedance-2-0-260128".to_string()
}
fn default_seedance_poll_interval_secs() -> u64 {
5
}
fn default_seedance_timeout_secs() -> u64 {
3_600
}
fn default_seedance_execution_expires_after() -> u64 {
3_600
}
fn default_seedance_max_output_bytes() -> usize {
512 * 1024 * 1024
}
fn default_seedance_output_host_allowlist() -> Vec<String> {
["volces.com", "volcengine.com", "volcengineapi.com"]
.into_iter()
.map(str::to_string)
.collect()
}
fn default_workflow_points_cost() -> f64 {
10.0
}
@@ -261,6 +681,26 @@ fn default_usd_to_points_rate() -> f64 {
100.0
}
fn default_alipay_gateway_url() -> String {
"https://openapi.alipay.com/gateway.do".to_string()
}
fn default_alipay_notify_path() -> String {
"/api/payment/alipay/notify".to_string()
}
fn default_payment_credit_unit() -> String {
"USD API quota".to_string()
}
fn default_refund_request_window_days() -> i64 {
7
}
fn default_refund_provider_timeout_secs() -> u64 {
60
}
impl Default for ComfyUIConfig {
fn default() -> Self {
Self {
@@ -288,6 +728,7 @@ impl Default for PixelleVideoConfig {
fn default() -> Self {
Self {
enabled: false,
execution_mode: default_pixelle_execution_mode(),
api_base_url: Some("http://127.0.0.1:8000".to_string()),
api_key: None,
backend_provider: default_pixelle_backend_provider(),
@@ -312,16 +753,27 @@ impl Default for PaymentConfig {
Self {
enabled: false,
provider: PaymentProvider::ComfyCloud,
credential_master_key: None,
public_base_url: None,
product_name_prefix: "BossAI 账户充值".to_string(),
currency: "usd".to_string(),
credit_per_payment_unit: None,
payment_credit_unit: default_payment_credit_unit(),
min_amount_usd: 1.0,
max_amount_usd: Some(100_000.0),
daily_limit_usd: None,
order_timeout_minutes: 30,
order_timeout_minutes: 35,
max_pending_orders: 3,
refund_request_window_days: default_refund_request_window_days(),
refund_provider_timeout_secs: default_refund_provider_timeout_secs(),
stripe_secret_key: None,
stripe_webhook_secret: None,
alipay_app_id: None,
alipay_private_key: None,
alipay_public_key: None,
alipay_seller_id: None,
alipay_gateway_url: default_alipay_gateway_url(),
alipay_notify_path: default_alipay_notify_path(),
comfy_customer_token: None,
comfy_customer_token_endpoint: None,
comfy_firebase_api_key: None,
@@ -357,6 +809,23 @@ impl Default for JimengConfig {
}
}
impl Default for SeedanceConfig {
fn default() -> Self {
Self {
enabled: false,
api_key: String::new(),
base_url: default_seedance_base_url(),
default_model: default_seedance_model(),
poll_interval_secs: default_seedance_poll_interval_secs(),
timeout_secs: default_seedance_timeout_secs(),
execution_expires_after: default_seedance_execution_expires_after(),
max_output_bytes: default_seedance_max_output_bytes(),
output_host_allowlist: default_seedance_output_host_allowlist(),
cost_per_million_tokens: 0.0,
}
}
}
impl Default for AppConfig {
fn default() -> Self {
Self {
@@ -367,9 +836,11 @@ impl Default for AppConfig {
cors_origins: vec!["*".to_string()],
},
db: DbConfig {
url: "sqlite::memory:".to_string(),
max_connections: 1,
url: "sqlite:boosapi.db?mode=rwc".to_string(),
max_connections: 10,
ssl: false,
ssl_mode: None,
ca_cert_path: None,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
@@ -388,9 +859,11 @@ impl Default for AppConfig {
template_proxy_url: None,
},
comfyui: ComfyUIConfig::default(),
agent: AgentConfig::default(),
pixelle_video: PixelleVideoConfig::default(),
payment: PaymentConfig::default(),
jimeng: JimengConfig::default(),
seedance: SeedanceConfig::default(),
}
}
}
@@ -434,9 +907,153 @@ impl AppConfig {
pub fn load_from_file(path: &str) -> anyhow::Result<Self> {
dotenvy::dotenv().ok();
Ok(Figment::new()
.merge(Toml::file(path))
let override_path = runtime_db_override_path(path);
let mut figment = Figment::new().merge(Toml::file(path));
if override_path.exists() {
figment = figment.merge(Toml::file(override_path));
}
Ok(figment
.merge(Env::prefixed("BOOSAPI_").split("__"))
.extract()?)
}
}
pub fn runtime_db_override_path(config_path: impl AsRef<Path>) -> PathBuf {
let config_path = config_path.as_ref();
let filename = config_path
.file_name()
.and_then(|value| value.to_str())
.unwrap_or("config.toml");
config_path.with_file_name(format!("{filename}.db-runtime.toml"))
}
pub fn save_runtime_database_url(target_url: &str) -> anyhow::Result<PathBuf> {
let config_path = CONFIG_SOURCE_PATH
.get()
.cloned()
.unwrap_or_else(|| PathBuf::from("config.toml"));
save_runtime_database_url_at(config_path, target_url)
}
pub fn save_runtime_database_url_at(
config_path: impl AsRef<Path>,
target_url: &str,
) -> anyhow::Result<PathBuf> {
anyhow::ensure!(
std::env::var_os("BOOSAPI_DB__URL").is_none(),
"cannot save runtime database configuration while BOOSAPI_DB__URL is set; update the deployment environment instead"
);
let path = runtime_db_override_path(config_path);
let parent = path.parent().unwrap_or_else(|| Path::new("."));
std::fs::create_dir_all(parent)?;
let escaped = target_url.replace('\\', "\\\\").replace('"', "\\\"");
let mut options = std::fs::OpenOptions::new();
options.create(true).truncate(true).write(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
options.mode(0o600);
}
let mut file = options.open(&path)?;
file.write_all(format!("[db]\nurl = \"{escaped}\"\n").as_bytes())?;
file.sync_all()?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?;
}
Ok(path)
}
#[cfg(test)]
mod redis_config_tests {
use super::*;
use std::{ffi::OsString, sync::Mutex};
static ENV_LOCK: Mutex<()> = Mutex::new(());
fn parse_redis_config(toml: &str) -> RedisConfig {
Figment::new()
.merge(Toml::string(toml))
.extract::<RedisConfigFile>()
.expect("valid Redis test config")
.redis
}
#[test]
fn redis_config_defaults_are_explicit_and_disabled() {
let config = parse_redis_config("");
assert!(!config.enabled);
assert_eq!(config.url, "redis://127.0.0.1:6379/0");
assert_eq!(config.key_prefix, "boosapi");
assert_eq!(config.connection_timeout_ms, 3_000);
assert_eq!(config.response_timeout_ms, 2_000);
assert!(config.auth_state_secret.is_none());
config.validate().expect("default config is valid");
}
#[test]
fn redis_config_accepts_toml_timeout_overrides() {
let config = parse_redis_config(
r#"
[redis]
connection_timeout_ms = 4500
response_timeout_ms = 2500
"#,
);
assert_eq!(config.connection_timeout_ms, 4_500);
assert_eq!(config.response_timeout_ms, 2_500);
config.validate().expect("positive timeouts are valid");
}
#[test]
fn redis_config_accepts_environment_timeout_overrides() {
let _lock = ENV_LOCK.lock().unwrap_or_else(|error| error.into_inner());
let connection_key = "BOOSAPI_REDIS__CONNECTION_TIMEOUT_MS";
let response_key = "BOOSAPI_REDIS__RESPONSE_TIMEOUT_MS";
let previous_connection = std::env::var_os(connection_key);
let previous_response = std::env::var_os(response_key);
std::env::set_var(connection_key, "5100");
std::env::set_var(response_key, "2600");
let result = std::panic::catch_unwind(|| {
Figment::new()
.merge(Env::prefixed("BOOSAPI_").split("__"))
.extract::<RedisConfigFile>()
.expect("valid Redis environment config")
.redis
});
restore_env(connection_key, previous_connection);
restore_env(response_key, previous_response);
let config = result.expect("environment config test did not panic");
assert_eq!(config.connection_timeout_ms, 5_100);
assert_eq!(config.response_timeout_ms, 2_600);
}
#[test]
fn redis_config_rejects_zero_timeouts() {
let mut config = RedisConfig::default();
config.connection_timeout_ms = 0;
assert_eq!(
config.validate().unwrap_err().to_string(),
"redis.connection_timeout_ms must be greater than zero"
);
config.connection_timeout_ms = 3_000;
config.response_timeout_ms = 0;
assert_eq!(
config.validate().unwrap_err().to_string(),
"redis.response_timeout_ms must be greater than zero"
);
}
fn restore_env(key: &str, previous: Option<OsString>) {
match previous {
Some(value) => std::env::set_var(key, value),
None => std::env::remove_var(key),
}
}
}
+949
View File
@@ -0,0 +1,949 @@
use std::{
collections::{BTreeMap, BTreeSet},
sync::{
atomic::{AtomicBool, Ordering},
Arc, RwLock,
},
};
use anyhow::{bail, Context};
use serde::Serialize;
use sha2::{Digest, Sha256};
use sqlx::{any::AnyRow, Any, AnyPool, Column, Row};
use sqlx_core::any::AnyValueKind;
use super::{create_any_pool, run_migrations, DbKind};
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ColumnManifest {
pub name: String,
pub data_type: String,
pub primary_key_position: u32,
pub auto_increment: bool,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct TableManifest {
pub name: String,
pub columns: Vec<ColumnManifest>,
pub dependencies: Vec<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct TableMigrationSummary {
pub table: String,
pub rows_copied: u64,
pub source_digest: String,
pub target_digest: String,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct MigrationSummary {
pub source: DbKind,
pub target: DbKind,
pub verified: bool,
pub tables: Vec<TableMigrationSummary>,
}
#[derive(Debug, Clone, Default, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct MigrationProgress {
pub phase: String,
pub current_table: Option<String>,
pub tables_completed: usize,
pub total_tables: usize,
pub rows_copied: u64,
}
#[derive(Debug, Default)]
pub struct MigrationControl {
cancelled: AtomicBool,
progress: RwLock<MigrationProgress>,
}
impl MigrationControl {
pub fn cancel(&self) {
self.cancelled.store(true, Ordering::Release);
}
pub fn progress(&self) -> MigrationProgress {
self.progress
.read()
.expect("migration progress poisoned")
.clone()
}
fn update(&self, update: impl FnOnce(&mut MigrationProgress)) {
update(&mut self.progress.write().expect("migration progress poisoned"));
}
fn check_cancelled(&self) -> anyhow::Result<()> {
anyhow::ensure!(
!self.cancelled.load(Ordering::Acquire),
"database migration cancelled"
);
Ok(())
}
}
pub async fn migrate(
source_url: &str,
target_url: &str,
force: bool,
batch_size: usize,
verify_only: bool,
) -> anyhow::Result<MigrationSummary> {
migrate_with_control(source_url, target_url, force, batch_size, verify_only, None).await
}
pub async fn migrate_with_control(
source_url: &str,
target_url: &str,
force: bool,
batch_size: usize,
verify_only: bool,
control: Option<Arc<MigrationControl>>,
) -> anyhow::Result<MigrationSummary> {
anyhow::ensure!(batch_size > 0, "batch size must be greater than zero");
anyhow::ensure!(
source_url.trim() != target_url.trim(),
"source and target database URLs must differ"
);
let source_kind = DbKind::from_url(source_url)?;
let target_kind = DbKind::from_url(target_url)?;
let source = create_any_pool(source_url, 2).await?;
let target = create_any_pool(target_url, 2).await?;
run_migrations(target_kind, &target)
.await
.context("apply target schema migrations")?;
let source_manifest = inspect_schema(&source, source_kind).await?;
let target_manifest = inspect_schema(&target, target_kind).await?;
validate_manifest(&source_manifest, &target_manifest)?;
// Copy order must satisfy the target's constraints. Older source schemas
// may predate foreign keys that are enforced by the current migrations.
let order = dependency_order(&target_manifest);
if let Some(control) = &control {
control.update(|progress| {
progress.phase = if verify_only { "verifying" } else { "copying" }.into();
progress.total_tables = order.len();
});
control.check_cancelled()?;
}
if !verify_only {
prepare_target(&target, target_kind, &target_manifest, &order, force).await?;
copy_all(
&source,
source_kind,
&target,
target_kind,
&source_manifest,
&target_manifest,
&order,
batch_size,
control.as_deref(),
)
.await?;
synchronize_sequences(&target, target_kind, &target_manifest).await?;
}
let tables = verify_all(
&source,
source_kind,
&target,
target_kind,
&source_manifest,
&target_manifest,
&order,
control.as_deref(),
)
.await?;
verify_database_integrity(&target, target_kind).await?;
source.close().await;
target.close().await;
Ok(MigrationSummary {
source: source_kind,
target: target_kind,
verified: true,
tables,
})
}
async fn verify_database_integrity(pool: &AnyPool, kind: DbKind) -> anyhow::Result<()> {
match kind {
DbKind::Sqlite => {
let violations = sqlx::query("PRAGMA foreign_key_check")
.fetch_all(pool)
.await?;
anyhow::ensure!(
violations.is_empty(),
"target SQLite foreign key check failed"
);
let result: String = sqlx::query_scalar("PRAGMA integrity_check")
.fetch_one(pool)
.await?;
anyhow::ensure!(
result == "ok",
"target SQLite integrity check failed: {result}"
);
}
DbKind::Postgres => {
let unvalidated: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM pg_constraint c
JOIN pg_namespace n ON n.oid = c.connamespace
WHERE n.nspname = 'public' AND NOT c.convalidated",
)
.fetch_one(pool)
.await?;
anyhow::ensure!(
unvalidated == 0,
"target PostgreSQL has {unvalidated} unvalidated constraints"
);
}
DbKind::MySql => {
// InnoDB foreign keys stay enabled during copying and reject orphan rows.
}
}
Ok(())
}
async fn inspect_schema(
pool: &AnyPool,
kind: DbKind,
) -> anyhow::Result<BTreeMap<String, TableManifest>> {
let table_sql = match kind {
DbKind::Sqlite => {
"SELECT name FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%' AND name <> '_sqlx_migrations' ORDER BY name"
}
DbKind::Postgres => {
"SELECT table_name::text AS name FROM information_schema.tables WHERE table_schema = 'public' AND table_type = 'BASE TABLE' AND table_name <> '_sqlx_migrations' ORDER BY table_name"
}
DbKind::MySql => {
"SELECT CAST(table_name AS CHAR) AS name FROM information_schema.tables WHERE table_schema = DATABASE() AND table_type = 'BASE TABLE' AND table_name <> '_sqlx_migrations' ORDER BY table_name"
}
};
let mut names: Vec<String> = sqlx::query_scalar(table_sql).fetch_all(pool).await?;
names.retain(|name| !is_sqlite_recovery_table(kind, name));
let mut manifest = BTreeMap::new();
for name in names {
let columns = inspect_columns(pool, kind, &name).await?;
let dependencies = inspect_dependencies(pool, kind, &name).await?;
manifest.insert(
name.clone(),
TableManifest {
name,
columns,
dependencies,
},
);
}
Ok(manifest)
}
fn is_sqlite_recovery_table(kind: DbKind, name: &str) -> bool {
kind == DbKind::Sqlite
&& (name == "lost_and_found"
|| name.starts_with("lost_and_found_")
|| name == "route_channels_rescue")
}
async fn inspect_columns(
pool: &AnyPool,
kind: DbKind,
table: &str,
) -> anyhow::Result<Vec<ColumnManifest>> {
match kind {
DbKind::Sqlite => {
let sql = format!("PRAGMA table_xinfo({})", quote_ident(kind, table));
let rows = sqlx::query(&sql).fetch_all(pool).await?;
let mut columns = Vec::new();
for row in rows {
let hidden: i64 = row.try_get("hidden")?;
if hidden != 0 {
continue;
}
let data_type: String = row.try_get("type")?;
let primary_key_position = row.try_get::<i64, _>("pk")? as u32;
columns.push(ColumnManifest {
name: row.try_get("name")?,
auto_increment: primary_key_position > 0
&& data_type.eq_ignore_ascii_case("INTEGER"),
data_type,
primary_key_position,
});
}
Ok(columns)
}
DbKind::Postgres => {
let sql = format!(
"SELECT c.column_name::text AS name, c.data_type::text AS data_type, \
COALESCE(pk.position, 0)::bigint AS pk_position, \
COALESCE(c.is_identity = 'YES' OR c.column_default LIKE 'nextval(%', false) AS auto_increment \
FROM information_schema.columns c \
LEFT JOIN (SELECT a.attname, array_position(i.indkey, a.attnum) AS position \
FROM pg_index i JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) \
WHERE i.indrelid = '{}'::regclass AND i.indisprimary) pk ON pk.attname = c.column_name \
WHERE c.table_schema = 'public' AND c.table_name = '{}' AND c.is_generated = 'NEVER' \
ORDER BY c.ordinal_position",
pg_literal(table),
pg_literal(table)
);
manifest_columns(sqlx::query(&sql).fetch_all(pool).await?)
}
DbKind::MySql => {
let sql = format!(
"SELECT CAST(c.column_name AS CHAR) AS name, CAST(c.data_type AS CHAR) AS data_type, \
CAST(COALESCE(k.ordinal_position, 0) AS SIGNED) AS pk_position, \
(c.extra LIKE '%auto_increment%') AS auto_increment \
FROM information_schema.columns c \
LEFT JOIN information_schema.key_column_usage k \
ON k.table_schema = c.table_schema AND k.table_name = c.table_name \
AND k.column_name = c.column_name AND k.constraint_name = 'PRIMARY' \
WHERE c.table_schema = DATABASE() AND c.table_name = '{}' \
AND c.extra NOT LIKE '%VIRTUAL GENERATED%' \
AND c.extra NOT LIKE '%STORED GENERATED%' ORDER BY c.ordinal_position",
mysql_literal(table)
);
manifest_columns(sqlx::query(&sql).fetch_all(pool).await?)
}
}
}
fn manifest_columns(rows: Vec<AnyRow>) -> anyhow::Result<Vec<ColumnManifest>> {
rows.into_iter()
.map(|row| {
Ok(ColumnManifest {
name: string_value(&row, "name")?,
data_type: string_value(&row, "data_type")?,
primary_key_position: integer_value(&row, "pk_position")? as u32,
auto_increment: bool_value(&row, "auto_increment")?,
})
})
.collect()
}
async fn inspect_dependencies(
pool: &AnyPool,
kind: DbKind,
table: &str,
) -> anyhow::Result<Vec<String>> {
let sql = match kind {
DbKind::Sqlite => format!("PRAGMA foreign_key_list({})", quote_ident(kind, table)),
DbKind::Postgres => format!(
"SELECT DISTINCT parent.relname::text AS parent_table \
FROM pg_constraint fk JOIN pg_class child ON child.oid = fk.conrelid \
JOIN pg_class parent ON parent.oid = fk.confrelid \
WHERE fk.contype = 'f' AND child.relname = '{}'",
pg_literal(table)
),
DbKind::MySql => format!(
"SELECT DISTINCT CAST(referenced_table_name AS CHAR) AS parent_table \
FROM information_schema.key_column_usage WHERE table_schema = DATABASE() \
AND table_name = '{}' AND referenced_table_name IS NOT NULL",
mysql_literal(table)
),
};
let rows = sqlx::query(&sql).fetch_all(pool).await?;
let key = if kind == DbKind::Sqlite {
"table"
} else {
"parent_table"
};
let mut dependencies = BTreeSet::new();
for row in rows {
let parent: String = row.try_get(key)?;
if parent != table {
dependencies.insert(parent);
}
}
Ok(dependencies.into_iter().collect())
}
fn validate_manifest(
source: &BTreeMap<String, TableManifest>,
target: &BTreeMap<String, TableManifest>,
) -> anyhow::Result<()> {
for (table, source_table) in source {
let target_table = target
.get(table)
.with_context(|| format!("target schema is missing table {table}"))?;
for source_column in &source_table.columns {
anyhow::ensure!(
target_table
.columns
.iter()
.any(|column| column.name == source_column.name),
"target schema is missing column {table}.{}",
source_column.name
);
}
}
Ok(())
}
fn dependency_order(manifest: &BTreeMap<String, TableManifest>) -> Vec<String> {
let mut pending: BTreeSet<String> = manifest.keys().cloned().collect();
let mut completed = BTreeSet::new();
let mut order = Vec::with_capacity(pending.len());
while !pending.is_empty() {
let ready: Vec<String> = pending
.iter()
.filter(|table| {
manifest[*table].dependencies.iter().all(|dependency| {
completed.contains(dependency) || !manifest.contains_key(dependency)
})
})
.cloned()
.collect();
if ready.is_empty() {
order.extend(pending.iter().cloned());
break;
}
for table in ready {
pending.remove(&table);
completed.insert(table.clone());
order.push(table);
}
}
order
}
async fn prepare_target(
target: &AnyPool,
kind: DbKind,
manifest: &BTreeMap<String, TableManifest>,
order: &[String],
force: bool,
) -> anyhow::Result<()> {
let mut nonempty = Vec::new();
for table in order {
let count = table_count(target, kind, table).await?;
if count > 0 {
nonempty.push(format!("{table} ({count})"));
}
}
if nonempty.is_empty() {
return Ok(());
}
if !force {
bail!(
"target database is not empty: {}; rerun with --force to overwrite",
nonempty.join(", ")
);
}
let mut connection = target.acquire().await?;
if kind == DbKind::MySql {
sqlx::query("SET FOREIGN_KEY_CHECKS = 0")
.execute(&mut *connection)
.await?;
}
for table in order.iter().rev() {
let sql = format!("DELETE FROM {}", quote_ident(kind, table));
sqlx::query(&sql).execute(&mut *connection).await?;
}
if kind == DbKind::MySql {
sqlx::query("SET FOREIGN_KEY_CHECKS = 1")
.execute(&mut *connection)
.await?;
}
let _ = manifest;
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn copy_all(
source: &AnyPool,
source_kind: DbKind,
target: &AnyPool,
target_kind: DbKind,
source_manifest: &BTreeMap<String, TableManifest>,
target_manifest: &BTreeMap<String, TableManifest>,
order: &[String],
batch_size: usize,
control: Option<&MigrationControl>,
) -> anyhow::Result<()> {
for (table_index, table) in order.iter().enumerate() {
if let Some(control) = control {
control.check_cancelled()?;
control.update(|progress| {
progress.current_table = Some(table.clone());
progress.tables_completed = table_index;
});
}
let source_table = &source_manifest[table];
let target_table = &target_manifest[table];
let columns: Vec<&ColumnManifest> = target_table
.columns
.iter()
.filter(|column| {
source_table
.columns
.iter()
.any(|item| item.name == column.name)
})
.collect();
let column_sql = columns
.iter()
.map(|column| quote_ident(source_kind, &column.name))
.collect::<Vec<_>>()
.join(", ");
let order_sql = primary_key_order(source_table, source_kind);
let mut offset = 0_usize;
loop {
if let Some(control) = control {
control.check_cancelled()?;
}
let select = format!(
"SELECT {column_sql} FROM {}{order_sql} LIMIT {batch_size} OFFSET {offset}",
quote_ident(source_kind, table)
);
let rows = sqlx::query(&select).fetch_all(source).await?;
if rows.is_empty() {
break;
}
let target_columns = columns
.iter()
.map(|column| quote_ident(target_kind, &column.name))
.collect::<Vec<_>>()
.join(", ");
let placeholders = (1..=columns.len())
.map(|index| placeholder(target_kind, index))
.collect::<Vec<_>>()
.join(", ");
let insert = format!(
"INSERT INTO {} ({target_columns}) VALUES ({placeholders})",
quote_ident(target_kind, table)
);
let mut transaction = target.begin().await?;
for row in &rows {
let mut query = sqlx::query(&insert);
for (column, value) in columns.iter().zip(&row.values) {
query = bind_value(query, &value.kind, column);
}
query
.execute(&mut *transaction)
.await
.with_context(|| format!("copy row into {table}"))?;
}
transaction.commit().await?;
offset += rows.len();
if let Some(control) = control {
control.update(|progress| progress.rows_copied += rows.len() as u64);
}
}
if let Some(control) = control {
control.update(|progress| progress.tables_completed = table_index + 1);
}
}
Ok(())
}
fn bind_value<'q>(
query: sqlx::query::Query<'q, Any, sqlx::any::AnyArguments<'q>>,
value: &AnyValueKind<'_>,
target_column: &ColumnManifest,
) -> sqlx::query::Query<'q, Any, sqlx::any::AnyArguments<'q>> {
let target_type = target_column.data_type.to_ascii_lowercase();
if target_type == "boolean" {
return match value {
AnyValueKind::Null(_) => query.bind(Option::<bool>::None),
AnyValueKind::Bool(value) => query.bind(*value),
AnyValueKind::SmallInt(value) => query.bind(*value != 0),
AnyValueKind::Integer(value) => query.bind(*value != 0),
AnyValueKind::BigInt(value) => query.bind(*value != 0),
_ => unreachable!("non-boolean source value for a boolean target column"),
};
}
if matches!(
target_type.as_str(),
"bigint" | "integer" | "int" | "smallint" | "tinyint"
) {
return match value {
AnyValueKind::Null(_) => query.bind(Option::<i64>::None),
AnyValueKind::Bool(value) => query.bind(if *value { 1_i64 } else { 0_i64 }),
AnyValueKind::SmallInt(value) => query.bind(i64::from(*value)),
AnyValueKind::Integer(value) => query.bind(i64::from(*value)),
AnyValueKind::BigInt(value) => query.bind(*value),
_ => unreachable!("non-integer source value for an integer target column"),
};
}
if matches!(
target_type.as_str(),
"double" | "double precision" | "float" | "real" | "numeric" | "decimal"
) {
return match value {
AnyValueKind::Null(_) => query.bind(Option::<f64>::None),
AnyValueKind::SmallInt(value) => query.bind(f64::from(*value)),
AnyValueKind::Integer(value) => query.bind(f64::from(*value)),
AnyValueKind::BigInt(value) => query.bind(*value as f64),
AnyValueKind::Real(value) => query.bind(f64::from(*value)),
AnyValueKind::Double(value) => query.bind(*value),
_ => unreachable!("non-numeric source value for a floating-point target column"),
};
}
match value {
AnyValueKind::Null(_) => bind_null(query, &target_type),
AnyValueKind::Bool(value) => query.bind(*value),
AnyValueKind::SmallInt(value) => query.bind(*value),
AnyValueKind::Integer(value) => query.bind(*value),
AnyValueKind::BigInt(value) => query.bind(*value),
AnyValueKind::Real(value) => query.bind(*value),
AnyValueKind::Double(value) => query.bind(*value),
AnyValueKind::Text(value) => query.bind(value.to_string()),
AnyValueKind::Blob(value) => query.bind(value.to_vec()),
_ => unreachable!("SQLx Any returned an unsupported value kind"),
}
}
fn bind_null<'q>(
query: sqlx::query::Query<'q, Any, sqlx::any::AnyArguments<'q>>,
target_type: &str,
) -> sqlx::query::Query<'q, Any, sqlx::any::AnyArguments<'q>> {
match target_type.to_ascii_lowercase().as_str() {
"bigint" | "integer" | "int" | "smallint" | "tinyint" => query.bind(Option::<i64>::None),
"double" | "double precision" | "float" | "real" | "numeric" | "decimal" => {
query.bind(Option::<f64>::None)
}
"blob" | "longblob" | "mediumblob" | "tinyblob" | "bytea" => {
query.bind(Option::<Vec<u8>>::None)
}
_ => query.bind(Option::<String>::None),
}
}
async fn synchronize_sequences(
target: &AnyPool,
kind: DbKind,
manifest: &BTreeMap<String, TableManifest>,
) -> anyhow::Result<()> {
for table in manifest.values() {
for column in table.columns.iter().filter(|column| column.auto_increment) {
let table_name = quote_ident(kind, &table.name);
let column_name = quote_ident(kind, &column.name);
match kind {
DbKind::Sqlite => {
sqlx::query("DELETE FROM sqlite_sequence WHERE name = ?")
.bind(&table.name)
.execute(target)
.await?;
let sql = format!(
"INSERT INTO sqlite_sequence(name, seq) VALUES (?, COALESCE((SELECT MAX({column_name}) FROM {table_name}), 0))"
);
sqlx::query(&sql).bind(&table.name).execute(target).await?;
}
DbKind::Postgres => {
let sql = format!(
"SELECT setval(pg_get_serial_sequence('{}', '{}'), \
COALESCE(MAX({column_name}), 1), COUNT(*) > 0) FROM {table_name}",
pg_literal(&table.name),
pg_literal(&column.name)
);
sqlx::query(&sql).execute(target).await?;
}
DbKind::MySql => {
let sql = format!("SELECT COALESCE(MAX({column_name}), 0) FROM {table_name}");
let max: i64 = sqlx::query_scalar(&sql).fetch_one(target).await?;
let sql = format!("ALTER TABLE {table_name} AUTO_INCREMENT = {}", max + 1);
sqlx::query(&sql).execute(target).await?;
}
}
}
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn verify_all(
source: &AnyPool,
source_kind: DbKind,
target: &AnyPool,
target_kind: DbKind,
source_manifest: &BTreeMap<String, TableManifest>,
target_manifest: &BTreeMap<String, TableManifest>,
order: &[String],
control: Option<&MigrationControl>,
) -> anyhow::Result<Vec<TableMigrationSummary>> {
let mut summaries = Vec::with_capacity(order.len());
if let Some(control) = control {
control.update(|progress| {
progress.phase = "verifying".into();
progress.tables_completed = 0;
});
}
for (table_index, table) in order.iter().enumerate() {
if let Some(control) = control {
control.check_cancelled()?;
control.update(|progress| progress.current_table = Some(table.clone()));
}
let source_count = table_count(source, source_kind, table).await?;
let target_count = table_count(target, target_kind, table).await?;
anyhow::ensure!(
source_count == target_count,
"row count mismatch for {table}: source={source_count}, target={target_count}"
);
let target_table = &target_manifest[table];
let columns: Vec<&ColumnManifest> = target_table
.columns
.iter()
.filter(|column| {
source_manifest[table]
.columns
.iter()
.any(|item| item.name == column.name)
})
.collect();
let source_digest = table_digest(
source,
source_kind,
table,
&columns,
&source_manifest[table],
)
.await?;
let target_digest =
table_digest(target, target_kind, table, &columns, target_table).await?;
if source_digest != target_digest {
let source_columns = column_digests(
source,
source_kind,
table,
&columns,
&source_manifest[table],
)
.await?;
let target_columns =
column_digests(target, target_kind, table, &columns, target_table).await?;
let mismatched = columns
.iter()
.filter(|column| {
source_columns.get(&column.name) != target_columns.get(&column.name)
})
.map(|column| column.name.as_str())
.collect::<Vec<_>>()
.join(", ");
bail!("content digest mismatch for table {table}; columns: {mismatched}");
}
summaries.push(TableMigrationSummary {
table: table.clone(),
rows_copied: source_count,
source_digest,
target_digest,
});
if let Some(control) = control {
control.update(|progress| progress.tables_completed = table_index + 1);
}
}
if let Some(control) = control {
control.update(|progress| {
progress.phase = "completed".into();
progress.current_table = None;
});
}
Ok(summaries)
}
async fn table_digest(
pool: &AnyPool,
kind: DbKind,
table: &str,
columns: &[&ColumnManifest],
table_manifest: &TableManifest,
) -> anyhow::Result<String> {
let projection = columns
.iter()
.map(|column| quote_ident(kind, &column.name))
.collect::<Vec<_>>()
.join(", ");
let sql = format!(
"SELECT {projection} FROM {}{}",
quote_ident(kind, table),
primary_key_order(table_manifest, kind)
);
let rows = sqlx::query(&sql).fetch_all(pool).await?;
let mut digest = Sha256::new();
for row in rows {
for (column, value) in columns.iter().zip(&row.values) {
digest.update((column.name.len() as u64).to_be_bytes());
digest.update(column.name.as_bytes());
hash_value(&mut digest, &value.kind);
}
}
Ok(hex::encode(digest.finalize()))
}
async fn column_digests(
pool: &AnyPool,
kind: DbKind,
table: &str,
columns: &[&ColumnManifest],
table_manifest: &TableManifest,
) -> anyhow::Result<BTreeMap<String, String>> {
let projection = columns
.iter()
.map(|column| quote_ident(kind, &column.name))
.collect::<Vec<_>>()
.join(", ");
let sql = format!(
"SELECT {projection} FROM {}{}",
quote_ident(kind, table),
primary_key_order(table_manifest, kind)
);
let rows = sqlx::query(&sql).fetch_all(pool).await?;
let mut digests: Vec<Sha256> = columns.iter().map(|_| Sha256::new()).collect();
for row in rows {
for (digest, value) in digests.iter_mut().zip(&row.values) {
hash_value(digest, &value.kind);
}
}
Ok(columns
.iter()
.zip(digests)
.map(|(column, digest)| (column.name.clone(), hex::encode(digest.finalize())))
.collect())
}
fn hash_value(digest: &mut Sha256, value: &AnyValueKind<'_>) {
match value {
AnyValueKind::Null(_) => digest.update([0]),
AnyValueKind::Bool(value) => {
hash_bytes(digest, 1, if *value { b"1" } else { b"0" });
}
AnyValueKind::SmallInt(value) => hash_bytes(digest, 1, value.to_string().as_bytes()),
AnyValueKind::Integer(value) => hash_bytes(digest, 1, value.to_string().as_bytes()),
AnyValueKind::BigInt(value) => hash_bytes(digest, 1, value.to_string().as_bytes()),
AnyValueKind::Real(value) => hash_bytes(digest, 2, (*value as f64).to_string().as_bytes()),
AnyValueKind::Double(value) => hash_bytes(digest, 2, value.to_string().as_bytes()),
AnyValueKind::Text(value) => hash_bytes(digest, 3, value.as_bytes()),
AnyValueKind::Blob(value) => hash_bytes(digest, 4, value),
_ => unreachable!("SQLx Any returned an unsupported value kind"),
}
}
fn hash_bytes(digest: &mut Sha256, tag: u8, value: &[u8]) {
digest.update([tag]);
digest.update((value.len() as u64).to_be_bytes());
digest.update(value);
}
async fn table_count(pool: &AnyPool, kind: DbKind, table: &str) -> anyhow::Result<u64> {
let sql = format!("SELECT COUNT(*) FROM {}", quote_ident(kind, table));
let count: i64 = sqlx::query_scalar(&sql).fetch_one(pool).await?;
Ok(count as u64)
}
fn primary_key_order(table: &TableManifest, kind: DbKind) -> String {
let mut primary_key: Vec<&ColumnManifest> = table
.columns
.iter()
.filter(|column| column.primary_key_position > 0)
.collect();
primary_key.sort_by_key(|column| column.primary_key_position);
if primary_key.is_empty() {
return table
.columns
.first()
.map(|column| format!(" ORDER BY {}", quote_ident(kind, &column.name)))
.unwrap_or_default();
}
format!(
" ORDER BY {}",
primary_key
.iter()
.map(|column| quote_ident(kind, &column.name))
.collect::<Vec<_>>()
.join(", ")
)
}
fn quote_ident(kind: DbKind, value: &str) -> String {
match kind {
DbKind::MySql => format!("`{}`", value.replace('`', "``")),
_ => format!("\"{}\"", value.replace('"', "\"\"")),
}
}
fn placeholder(kind: DbKind, position: usize) -> String {
if kind == DbKind::Postgres {
format!("${position}")
} else {
"?".to_string()
}
}
fn pg_literal(value: &str) -> String {
value.replace('\'', "''")
}
fn mysql_literal(value: &str) -> String {
value.replace('\\', "\\\\").replace('\'', "''")
}
fn integer_value(row: &AnyRow, column: &str) -> anyhow::Result<i64> {
let index = row
.columns()
.iter()
.position(|item| item.name() == column)
.with_context(|| format!("missing column {column}"))?;
match &row.values[index].kind {
AnyValueKind::SmallInt(value) => Ok(i64::from(*value)),
AnyValueKind::Integer(value) => Ok(i64::from(*value)),
AnyValueKind::BigInt(value) => Ok(*value),
other => bail!("column {column} is not an integer: {other:?}"),
}
}
fn bool_value(row: &AnyRow, column: &str) -> anyhow::Result<bool> {
let index = row
.columns()
.iter()
.position(|item| item.name() == column)
.with_context(|| format!("missing column {column}"))?;
match &row.values[index].kind {
AnyValueKind::Bool(value) => Ok(*value),
AnyValueKind::SmallInt(value) => Ok(*value != 0),
AnyValueKind::Integer(value) => Ok(*value != 0),
AnyValueKind::BigInt(value) => Ok(*value != 0),
other => bail!("column {column} is not boolean: {other:?}"),
}
}
fn string_value(row: &AnyRow, column: &str) -> anyhow::Result<String> {
let index = row
.columns()
.iter()
.position(|item| item.name() == column)
.with_context(|| format!("missing column {column}"))?;
match &row.values[index].kind {
AnyValueKind::Text(value) => Ok(value.to_string()),
AnyValueKind::Blob(value) => Ok(std::str::from_utf8(value)
.with_context(|| format!("column {column} is not valid UTF-8"))?
.to_string()),
other => bail!("column {column} is not text: {other:?}"),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn excludes_only_known_sqlite_recovery_tables() {
assert!(is_sqlite_recovery_table(DbKind::Sqlite, "lost_and_found"));
assert!(is_sqlite_recovery_table(DbKind::Sqlite, "lost_and_found_0"));
assert!(is_sqlite_recovery_table(
DbKind::Sqlite,
"route_channels_rescue"
));
assert!(!is_sqlite_recovery_table(DbKind::Sqlite, "route_channels"));
assert!(!is_sqlite_recovery_table(
DbKind::Postgres,
"lost_and_found"
));
}
}
+221 -54
View File
@@ -1,14 +1,19 @@
use std::borrow::Cow;
use anyhow::{bail, Context};
use serde::Serialize;
use sqlx::{
any::{install_default_drivers, AnyPoolOptions},
postgres::{PgPool, PgPoolOptions},
sqlite::{SqlitePool, SqlitePoolOptions},
AnyPool,
};
pub mod migrate;
pub mod models;
pub mod queries;
pub mod schema_compat;
pub mod sql;
pub mod sqlite_to_postgres;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
@@ -16,6 +21,27 @@ pub mod sqlite_to_postgres;
pub enum DbKind {
Sqlite,
Postgres,
MySql,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum DbFlavor {
Sqlite,
Postgres,
#[serde(rename = "pgrust")]
PgRust,
MySql,
#[serde(rename = "mariadb")]
MariaDb,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DbConnectionInfo {
pub kind: DbKind,
pub flavor: DbFlavor,
pub version: String,
}
impl DbKind {
@@ -25,8 +51,10 @@ impl DbKind {
Ok(Self::Sqlite)
} else if normalized.starts_with("postgres://") || normalized.starts_with("postgresql://") {
Ok(Self::Postgres)
} else if normalized.starts_with("mysql://") || normalized.starts_with("mariadb://") {
Ok(Self::MySql)
} else {
bail!("unsupported database URL; expected sqlite:, postgres:// or postgresql://")
bail!("unsupported database URL; expected sqlite:, postgres://, postgresql://, mysql:// or mariadb://")
}
}
@@ -34,30 +62,41 @@ impl DbKind {
match self {
Self::Sqlite => "sqlite",
Self::Postgres => "postgres",
Self::MySql => "mysql",
}
}
}
pub async fn create_pool(database_url: &str, max_connections: u32) -> anyhow::Result<SqlitePool> {
if DbKind::from_url(database_url)? != DbKind::Sqlite {
bail!(
"runtime PostgreSQL query support is not enabled yet; use migrate-sqlite-to-postgres to prepare PostgreSQL data"
);
fn sqlx_url(database_url: &str) -> Cow<'_, str> {
if database_url
.trim()
.to_ascii_lowercase()
.starts_with("mariadb://")
{
Cow::Owned(format!("mysql://{}", &database_url.trim()[10..]))
} else {
Cow::Borrowed(database_url)
}
}
pub async fn create_sqlite_pool(
database_url: &str,
max_connections: u32,
) -> anyhow::Result<SqlitePool> {
let pool = SqlitePoolOptions::new()
.max_connections(max_connections)
.after_connect(|conn, _meta| {
Box::pin(async move {
sqlx::query("PRAGMA journal_mode = WAL")
crate::db::sql::query("PRAGMA journal_mode = WAL")
.execute(&mut *conn)
.await?;
sqlx::query("PRAGMA synchronous = NORMAL")
crate::db::sql::query("PRAGMA synchronous = NORMAL")
.execute(&mut *conn)
.await?;
sqlx::query("PRAGMA busy_timeout = 5000")
crate::db::sql::query("PRAGMA busy_timeout = 5000")
.execute(&mut *conn)
.await?;
sqlx::query("PRAGMA foreign_keys = ON")
crate::db::sql::query("PRAGMA foreign_keys = ON")
.execute(&mut *conn)
.await?;
Ok(())
@@ -68,62 +107,190 @@ pub async fn create_pool(database_url: &str, max_connections: u32) -> anyhow::Re
Ok(pool)
}
pub async fn create_pg_pool(database_url: &str, max_connections: u32) -> anyhow::Result<PgPool> {
let pool = PgPoolOptions::new()
.max_connections(max_connections)
.connect(database_url)
.await
.context("connect PostgreSQL")?;
pub async fn create_pool(database_url: &str, max_connections: u32) -> anyhow::Result<AnyPool> {
let (pool, info) = connect_any_pool(database_url, max_connections).await?;
crate::db::sql::set_runtime_kind(info.kind)?;
Ok(pool)
}
pub async fn test_connection(database_url: &str, max_connections: u32) -> anyhow::Result<DbKind> {
match DbKind::from_url(database_url)? {
DbKind::Sqlite => {
SqlitePoolOptions::new()
.max_connections(max_connections)
.connect(database_url)
.await
.context("connect SQLite")?
.close()
.await;
Ok(DbKind::Sqlite)
}
DbKind::Postgres => {
create_pg_pool(database_url, max_connections)
.await?
.close()
.await;
Ok(DbKind::Postgres)
}
}
pub async fn create_pg_pool(database_url: &str, max_connections: u32) -> anyhow::Result<PgPool> {
PgPoolOptions::new()
.max_connections(max_connections)
.connect(database_url)
.await
.context("connect PostgreSQL")
}
pub async fn run_migrations(
kind: DbKind,
sqlite_pool: Option<&SqlitePool>,
pg_pool: Option<&PgPool>,
) -> anyhow::Result<()> {
pub async fn test_connection(
database_url: &str,
max_connections: u32,
) -> anyhow::Result<DbConnectionInfo> {
let (pool, info) = connect_any_pool(database_url, max_connections).await?;
pool.close().await;
Ok(info)
}
pub async fn run_migrations(kind: DbKind, pool: &AnyPool) -> anyhow::Result<()> {
match kind {
DbKind::Sqlite => {
let pool = sqlite_pool.context("SQLite pool required for SQLite migrations")?;
sqlx::migrate!("./migrations").run(pool).await?;
}
DbKind::Postgres => {
let pool = pg_pool.context("PostgreSQL pool required for PostgreSQL migrations")?;
sqlx::migrate!("./migrations_pg").run(pool).await?;
}
DbKind::Sqlite => sqlx::migrate!("./migrations").run(pool).await?,
DbKind::Postgres => sqlx::migrate!("./migrations_pg").run(pool).await?,
DbKind::MySql => sqlx::migrate!("./migrations_mysql").run(pool).await?,
}
Ok(())
}
pub async fn create_any_pool(
pub async fn create_any_pool(database_url: &str, max_connections: u32) -> anyhow::Result<AnyPool> {
connect_any_pool(database_url, max_connections)
.await
.map(|(pool, _)| pool)
}
pub async fn connect_any_pool(
database_url: &str,
max_connections: u32,
) -> anyhow::Result<sqlx::AnyPool> {
) -> anyhow::Result<(AnyPool, DbConnectionInfo)> {
install_default_drivers();
Ok(AnyPoolOptions::new()
let kind = DbKind::from_url(database_url)?;
let sqlite = kind == DbKind::Sqlite;
let mysql = kind == DbKind::MySql;
let pool = AnyPoolOptions::new()
.max_connections(max_connections)
.connect(database_url)
.await?)
.after_connect(move |connection, _| {
Box::pin(async move {
if sqlite {
sqlx::query("PRAGMA journal_mode = WAL")
.execute(&mut *connection)
.await?;
sqlx::query("PRAGMA synchronous = NORMAL")
.execute(&mut *connection)
.await?;
sqlx::query("PRAGMA busy_timeout = 5000")
.execute(&mut *connection)
.await?;
sqlx::query("PRAGMA foreign_keys = ON")
.execute(&mut *connection)
.await?;
}
if mysql {
sqlx::query(
"SET SESSION sql_mode = CONCAT_WS(',', @@SESSION.sql_mode, 'PIPES_AS_CONCAT')",
)
.execute(&mut *connection)
.await?;
}
Ok(())
})
})
.connect(sqlx_url(database_url).as_ref())
.await
.with_context(|| format!("connect {}", kind.as_str()))?;
let info = detect_connection_info(&pool, kind).await?;
validate_version(&info)?;
Ok((pool, info))
}
async fn detect_connection_info(pool: &AnyPool, kind: DbKind) -> anyhow::Result<DbConnectionInfo> {
let version: String = match kind {
DbKind::Sqlite => {
sqlx::query_scalar("SELECT sqlite_version()")
.fetch_one(pool)
.await?
}
DbKind::Postgres => {
sqlx::query_scalar("SELECT version()")
.fetch_one(pool)
.await?
}
DbKind::MySql => {
sqlx::query_scalar("SELECT VERSION()")
.fetch_one(pool)
.await?
}
};
let flavor = detect_flavor(kind, &version);
Ok(DbConnectionInfo {
kind,
flavor,
version,
})
}
fn detect_flavor(kind: DbKind, version: &str) -> DbFlavor {
match kind {
DbKind::Sqlite => DbFlavor::Sqlite,
DbKind::Postgres if version.to_ascii_lowercase().contains("pgrust") => DbFlavor::PgRust,
DbKind::Postgres => DbFlavor::Postgres,
DbKind::MySql if version.to_ascii_lowercase().contains("mariadb") => DbFlavor::MariaDb,
DbKind::MySql => DbFlavor::MySql,
}
}
fn validate_version(info: &DbConnectionInfo) -> anyhow::Result<()> {
let Some((major, minor)) = numeric_version(&info.version) else {
bail!("unable to parse database version: {}", info.version);
};
match info.flavor {
DbFlavor::MySql if (major, minor) < (8, 0) => {
bail!("MySQL 8.0 or newer is required; found {}", info.version)
}
DbFlavor::MariaDb if (major, minor) < (10, 6) => {
bail!("MariaDB 10.6 or newer is required; found {}", info.version)
}
_ => Ok(()),
}
}
fn numeric_version(raw: &str) -> Option<(u32, u32)> {
let start = raw.find(|ch: char| ch.is_ascii_digit())?;
let mut parts = raw[start..]
.split(|ch: char| !ch.is_ascii_digit() && ch != '.')
.next()?
.split('.');
Some((parts.next()?.parse().ok()?, parts.next()?.parse().ok()?))
}
#[cfg(test)]
mod tests {
use super::{detect_flavor, numeric_version, DbConnectionInfo, DbFlavor, DbKind};
#[test]
fn parses_mysql_and_mariadb_versions() {
assert_eq!(numeric_version("8.0.42"), Some((8, 0)));
assert_eq!(numeric_version("10.6.22-MariaDB"), Some((10, 6)));
}
#[test]
fn distinguishes_pgrust_from_postgresql() {
assert_eq!(
detect_flavor(
DbKind::Postgres,
"pgrust 18.3 on x86_64-linux-gnu, compiled by clang, 64-bit",
),
DbFlavor::PgRust,
);
assert_eq!(
serde_json::to_string(&DbFlavor::PgRust).unwrap(),
"\"pgrust\"",
);
assert_eq!(
detect_flavor(DbKind::Postgres, "PostgreSQL 17.5 on x86_64-pc-linux-gnu",),
DbFlavor::Postgres,
);
}
#[test]
fn rejects_unsupported_mysql_family_versions() {
let mysql = DbConnectionInfo {
kind: DbKind::MySql,
flavor: DbFlavor::MySql,
version: "5.7.44".into(),
};
let mariadb = DbConnectionInfo {
kind: DbKind::MySql,
flavor: DbFlavor::MariaDb,
version: "10.5.27-MariaDB".into(),
};
assert!(super::validate_version(&mysql).is_err());
assert!(super::validate_version(&mariadb).is_err());
}
}
+15 -13
View File
@@ -2,37 +2,39 @@ use crate::{
db::models::{downstream_api_key::DownstreamApiKey, user::User},
error::AppResult,
};
use sqlx::SqlitePool;
use sqlx::AnyPool;
pub async fn find_user_by_id(db: &SqlitePool, id: i64) -> AppResult<Option<User>> {
sqlx::query_as("SELECT * FROM users WHERE id = ?")
pub async fn find_user_by_id(db: &AnyPool, id: i64) -> AppResult<Option<User>> {
crate::db::sql::query_as("SELECT * FROM users WHERE id = ?")
.bind(id)
.fetch_optional(db)
.await
.map_err(Into::into)
}
pub async fn find_user_by_email(db: &SqlitePool, email: &str) -> AppResult<Option<User>> {
sqlx::query_as("SELECT * FROM users WHERE email = ?")
pub async fn find_user_by_email(db: &AnyPool, email: &str) -> AppResult<Option<User>> {
crate::db::sql::query_as("SELECT * FROM users WHERE email = ?")
.bind(email)
.fetch_optional(db)
.await
.map_err(Into::into)
}
pub async fn list_user_api_keys(db: &SqlitePool, user_id: i64) -> AppResult<Vec<DownstreamApiKey>> {
sqlx::query_as("SELECT * FROM downstream_api_keys WHERE user_id = ? ORDER BY created_at DESC")
.bind(user_id)
.fetch_all(db)
.await
.map_err(Into::into)
pub async fn list_user_api_keys(db: &AnyPool, user_id: i64) -> AppResult<Vec<DownstreamApiKey>> {
crate::db::sql::query_as(
"SELECT * FROM downstream_api_keys WHERE user_id = ? ORDER BY created_at DESC",
)
.bind(user_id)
.fetch_all(db)
.await
.map_err(Into::into)
}
pub async fn find_downstream_key_by_token(
db: &SqlitePool,
db: &AnyPool,
token: &str,
) -> AppResult<Option<DownstreamApiKey>> {
sqlx::query_as("SELECT * FROM downstream_api_keys WHERE key = ?")
crate::db::sql::query_as("SELECT * FROM downstream_api_keys WHERE key = ?")
.bind(token)
.fetch_optional(db)
.await
+61 -3
View File
@@ -49,6 +49,64 @@ pub async fn ensure_sqlite_upstream_compat(pool: &SqlitePool) -> anyhow::Result<
],
)
.await?;
ensure_columns(
pool,
"api_credit_grants",
&[("grant_kind", "TEXT NOT NULL DEFAULT 'balance'")],
)
.await?;
ensure_columns(
pool,
"api_entitlement_ledger",
&[
("package_amount_micros", "INTEGER NOT NULL DEFAULT 0"),
("balance_amount_micros", "INTEGER NOT NULL DEFAULT 0"),
],
)
.await?;
ensure_columns(
pool,
"commerce_products",
&[
("billing_track", "TEXT NOT NULL DEFAULT 'api'"),
("grant_kind", "TEXT NOT NULL DEFAULT 'package'"),
("benefit_type", "TEXT NOT NULL DEFAULT 'api_credit'"),
("validity_days", "INTEGER"),
],
)
.await?;
ensure_columns(
pool,
"commerce_orders",
&[
("billing_track_snapshot", "TEXT NOT NULL DEFAULT 'api'"),
("grant_kind_snapshot", "TEXT NOT NULL DEFAULT 'package'"),
(
"benefit_type_snapshot",
"TEXT NOT NULL DEFAULT 'api_credit'",
),
("validity_days_snapshot", "INTEGER"),
],
)
.await?;
ensure_columns(
pool,
"commerce_order_fulfillments",
&[("comfy_credit_grant_id", "INTEGER")],
)
.await?;
ensure_columns(
pool,
"payment_orders",
&[("billing_target", "TEXT NOT NULL DEFAULT 'api'")],
)
.await?;
ensure_columns(
pool,
"user_recharge_records",
&[("billing_target", "TEXT NOT NULL DEFAULT 'api'")],
)
.await?;
run_compat_ddl(pool).await?;
Ok(())
@@ -69,7 +127,7 @@ async fn ensure_columns(
continue;
}
let sql = format!("ALTER TABLE {table} ADD {name} {definition}");
sqlx::query(&sql)
crate::db::sql::query(&sql)
.execute(pool)
.await
.with_context(|| format!("add column {table}.{name}"))?;
@@ -79,7 +137,7 @@ async fn ensure_columns(
async fn table_columns(pool: &SqlitePool, table: &str) -> anyhow::Result<HashSet<String>> {
let sql = format!("PRAGMA table_info({table})");
let rows = sqlx::query(&sql).fetch_all(pool).await?;
let rows = crate::db::sql::query(&sql).fetch_all(pool).await?;
Ok(rows
.into_iter()
.filter_map(|row| row.try_get::<String, _>("name").ok())
@@ -88,7 +146,7 @@ async fn table_columns(pool: &SqlitePool, table: &str) -> anyhow::Result<HashSet
async fn run_compat_ddl(pool: &SqlitePool) -> anyhow::Result<()> {
for statement in SQLITE_COMPAT_DDL {
sqlx::query(statement)
crate::db::sql::query(statement)
.execute(pool)
.await
.with_context(|| format!("run SQLite compat DDL: {statement}"))?;
+825
View File
@@ -0,0 +1,825 @@
use std::{
collections::HashMap,
sync::{OnceLock, RwLock},
};
use sqlx::{
any::AnyArguments,
database::Database,
encode::Encode,
query::{Query, QueryAs, QueryScalar},
Any, Arguments, Executor, FromRow, Type,
};
use super::DbKind;
static RUNTIME_KIND: OnceLock<DbKind> = OnceLock::new();
static DIALECT_SQL: OnceLock<RwLock<HashMap<String, &'static str>>> = OnceLock::new();
static MYSQL_INSERT_SQL: OnceLock<RwLock<HashMap<String, &'static str>>> = OnceLock::new();
pub fn set_runtime_kind(kind: DbKind) -> anyhow::Result<()> {
if let Some(existing) = RUNTIME_KIND.get() {
anyhow::ensure!(
*existing == kind,
"runtime database dialect is already initialized as {}",
existing.as_str()
);
return Ok(());
}
RUNTIME_KIND
.set(kind)
.map_err(|_| anyhow::anyhow!("failed to initialize runtime database dialect"))
}
pub fn runtime_kind() -> DbKind {
RUNTIME_KIND.get().copied().unwrap_or(DbKind::Sqlite)
}
fn dialect_sql(kind: DbKind, sql: &str) -> &'static str {
let cache = DIALECT_SQL.get_or_init(|| RwLock::new(HashMap::new()));
let key = format!("{}\0{sql}", kind.as_str());
if let Some(cached) = cache.read().expect("SQL cache poisoned").get(&key) {
return cached;
}
let normalized = match kind {
DbKind::Postgres => rewrite_scalar_extrema(&rewrite_postgres_bigint_sums(
&rewrite_postgres_boolean_literals(&rewrite_misc_functions(
kind,
&rewrite_json_calls(
kind,
&rewrite_datetime_calls(
kind,
&sql.replace("CURRENT_TIMESTAMP", "(CURRENT_TIMESTAMP::TEXT)"),
),
),
)),
)),
DbKind::MySql if !sql.trim_end().ends_with("RETURNING id") => rewrite_scalar_extrema(
&mysql_compat_sql(&mysql_upsert_sql(&rewrite_misc_functions(
kind,
&rewrite_json_calls(kind, &rewrite_datetime_calls(kind, sql)),
))),
),
DbKind::MySql => rewrite_scalar_extrema(&rewrite_misc_functions(
kind,
&rewrite_json_calls(kind, &rewrite_datetime_calls(kind, sql)),
)),
DbKind::Sqlite => sql.to_string(),
};
let mut rewritten = String::with_capacity(normalized.len() + 8);
let mut chars = normalized.chars().peekable();
let mut parameter = 1_u32;
let mut single_quoted = false;
let mut double_quoted = false;
let mut line_comment = false;
let mut block_comment = false;
while let Some(ch) = chars.next() {
if line_comment {
rewritten.push(ch);
if ch == '\n' {
line_comment = false;
}
continue;
}
if block_comment {
rewritten.push(ch);
if ch == '*' && chars.peek() == Some(&'/') {
rewritten.push(chars.next().expect("peeked slash"));
block_comment = false;
}
continue;
}
if single_quoted {
rewritten.push(ch);
if ch == '\'' {
if chars.peek() == Some(&'\'') {
rewritten.push(chars.next().expect("peeked quote"));
} else {
single_quoted = false;
}
}
continue;
}
if double_quoted {
rewritten.push(ch);
if ch == '"' {
if chars.peek() == Some(&'"') {
rewritten.push(chars.next().expect("peeked quote"));
} else {
double_quoted = false;
}
}
continue;
}
match ch {
'\'' => {
single_quoted = true;
rewritten.push(ch);
}
'"' => {
double_quoted = true;
rewritten.push(ch);
}
'-' if chars.peek() == Some(&'-') => {
rewritten.push(ch);
rewritten.push(chars.next().expect("peeked dash"));
line_comment = true;
}
'/' if chars.peek() == Some(&'*') => {
rewritten.push(ch);
rewritten.push(chars.next().expect("peeked star"));
block_comment = true;
}
'?' if kind == DbKind::Postgres => {
rewritten.push('$');
rewritten.push_str(&parameter.to_string());
parameter += 1;
}
_ => rewritten.push(ch),
}
}
let value: &'static str = Box::leak(rewritten.into_boxed_str());
cache
.write()
.expect("SQL cache poisoned")
.insert(key, value);
value
}
pub fn query<'q, DB>(sql: &'q str) -> Query<'q, DB, <DB as Database>::Arguments<'q>>
where
DB: Database,
{
sqlx::query(dialect_sql(runtime_kind(), sql))
}
pub fn query_as<'q, DB, O>(sql: &'q str) -> QueryAs<'q, DB, O, <DB as Database>::Arguments<'q>>
where
DB: Database,
O: for<'r> FromRow<'r, DB::Row>,
{
sqlx::query_as(dialect_sql(runtime_kind(), sql))
}
pub fn query_scalar<'q, DB, O>(
sql: &'q str,
) -> QueryScalar<'q, DB, O, <DB as Database>::Arguments<'q>>
where
DB: Database,
(O,): for<'r> FromRow<'r, DB::Row>,
{
sqlx::query_scalar(dialect_sql(runtime_kind(), sql))
}
pub struct InsertId<'q> {
sql: &'static str,
arguments: Result<AnyArguments<'q>, sqlx::error::BoxDynError>,
}
impl<'q> InsertId<'q> {
pub fn bind<T>(mut self, value: T) -> Self
where
T: 'q + Encode<'q, Any> + Type<Any>,
{
if let Ok(arguments) = &mut self.arguments {
if let Err(error) = arguments.add(value) {
self.arguments = Err(error);
}
}
self
}
pub async fn fetch_one<'e, E>(self, executor: E) -> Result<i64, sqlx::Error>
where
'q: 'e,
E: Executor<'e, Database = Any>,
{
let arguments = self.arguments.map_err(sqlx::Error::Encode)?;
if runtime_kind() == DbKind::MySql {
let result = sqlx::query_with(mysql_insert_sql(self.sql, false)?, arguments)
.execute(executor)
.await?;
result.last_insert_id().ok_or_else(|| {
sqlx::Error::Protocol("MySQL did not return a last insert ID".into())
})
} else {
sqlx::query_scalar_with(self.sql, arguments)
.fetch_one(executor)
.await
}
}
pub async fn fetch_optional<'e, E>(self, executor: E) -> Result<Option<i64>, sqlx::Error>
where
'q: 'e,
E: Executor<'e, Database = Any>,
{
let arguments = self.arguments.map_err(sqlx::Error::Encode)?;
if runtime_kind() == DbKind::MySql {
let result = sqlx::query_with(mysql_insert_sql(self.sql, true)?, arguments)
.execute(executor)
.await?;
Ok(result.last_insert_id())
} else {
sqlx::query_scalar_with(self.sql, arguments)
.fetch_optional(executor)
.await
}
}
}
pub fn insert_id(sql: &str) -> InsertId<'_> {
InsertId {
sql: dialect_sql(runtime_kind(), sql),
arguments: Ok(AnyArguments::default()),
}
}
fn mysql_insert_sql(sql: &'static str, idempotent: bool) -> Result<&'static str, sqlx::Error> {
let suffix = if idempotent {
" ON CONFLICT (idempotency_key) DO NOTHING RETURNING id"
} else {
" RETURNING id"
};
let base = sql.strip_suffix(suffix).ok_or_else(|| {
sqlx::Error::Protocol(format!("insert-id SQL must end with {suffix}").into())
})?;
let cache = MYSQL_INSERT_SQL.get_or_init(|| RwLock::new(HashMap::new()));
if let Some(cached) = cache.read().expect("SQL cache poisoned").get(sql) {
return Ok(cached);
}
let base = mysql_compat_sql(base);
let rewritten: &'static str = Box::leak(
if idempotent {
format!("{base} ON DUPLICATE KEY UPDATE id = LAST_INSERT_ID(id)")
} else {
base
}
.into_boxed_str(),
);
cache
.write()
.expect("SQL cache poisoned")
.insert(sql.to_string(), rewritten);
Ok(rewritten)
}
fn mysql_upsert_sql(sql: &str) -> String {
let upper = sql.to_ascii_uppercase();
let Some(conflict_at) = upper.find(" ON CONFLICT") else {
return sql.to_string();
};
let Some(values_at) = upper.find("VALUES") else {
return sql.to_string();
};
let Some(open) = sql[..values_at].rfind('(') else {
return sql.to_string();
};
let Some(close) = matching_paren(sql, open) else {
return sql.to_string();
};
let first_column = sql[open + 1..close]
.split(',')
.next()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("id");
let conflict = &upper[conflict_at..];
if let Some(do_nothing) = conflict.find("DO NOTHING") {
let suffix_at = conflict_at + do_nothing + "DO NOTHING".len();
return format!(
"{} ON DUPLICATE KEY UPDATE {first_column} = {first_column}{}",
&sql[..conflict_at],
&sql[suffix_at..]
);
}
let Some(do_update) = conflict.find("DO UPDATE SET") else {
return sql.to_string();
};
let update_at = conflict_at + do_update + "DO UPDATE SET".len();
let assignments = replace_excluded_references(&sql[update_at..]);
format!(
"{} ON DUPLICATE KEY UPDATE{}",
&sql[..conflict_at],
assignments
)
}
fn mysql_compat_sql(sql: &str) -> String {
let lower = sql.to_ascii_lowercase();
if !lower.contains("settings") && !lower.contains("downstream_api_keys") {
return sql.to_string();
}
sql.replace("(key,", "(`key`,")
.replace("(key)", "(`key`)")
.replace(", key,", ", `key`,")
.replace(", key)", ", `key`)")
.replace("SELECT key,", "SELECT `key`,")
.replace(" key ", " `key` ")
.replace(".key", ".`key`")
.replace("ORDER BY key", "ORDER BY `key`")
}
fn rewrite_datetime_calls(kind: DbKind, sql: &str) -> String {
if kind == DbKind::Sqlite {
return sql.to_string();
}
let mut output = String::with_capacity(sql.len() + 32);
let mut remaining = sql;
loop {
let lower = remaining.to_ascii_lowercase();
let Some(start) = lower.find("datetime(") else {
output.push_str(remaining);
break;
};
output.push_str(&remaining[..start]);
let open = start + "datetime".len();
let Some(close) = matching_paren(remaining, open) else {
output.push_str(&remaining[start..]);
break;
};
let arguments = &remaining[open + 1..close];
let parts = split_top_level(arguments);
let replacement = match parts.as_slice() {
[value] => Some(match kind {
DbKind::Postgres => format!("(CAST({value} AS timestamptz)::text)"),
DbKind::MySql => {
format!("CAST(REPLACE(SUBSTRING({value}, 1, 19), 'T', ' ') AS DATETIME)")
}
DbKind::Sqlite => unreachable!(),
}),
[now, modifier] if now.trim().eq_ignore_ascii_case("'now'") => {
fixed_interval_expression(kind, modifier.trim())
}
_ => None,
};
if let Some(replacement) = replacement {
output.push_str(&replacement);
} else {
output.push_str(&remaining[start..=close]);
}
remaining = &remaining[close + 1..];
}
output
}
fn split_top_level(input: &str) -> Vec<&str> {
let mut parts = Vec::new();
let mut start = 0;
let mut depth = 0_u32;
let mut quoted = false;
let bytes = input.as_bytes();
let mut index = 0;
while index < bytes.len() {
match bytes[index] {
b'\'' => {
if quoted && bytes.get(index + 1) == Some(&b'\'') {
index += 1;
} else {
quoted = !quoted;
}
}
b'(' if !quoted => depth += 1,
b')' if !quoted => depth = depth.saturating_sub(1),
b',' if !quoted && depth == 0 => {
parts.push(input[start..index].trim());
start = index + 1;
}
_ => {}
}
index += 1;
}
parts.push(input[start..].trim());
parts
}
fn fixed_interval_expression(kind: DbKind, modifier: &str) -> Option<String> {
let modifier = modifier.strip_prefix('\'')?.strip_suffix('\'')?.trim();
let mut parts = modifier.split_whitespace();
let amount: i64 = parts.next()?.parse().ok()?;
let unit = parts.next()?.trim_end_matches('s').to_ascii_uppercase();
if parts.next().is_some() || !matches!(unit.as_str(), "SECOND" | "MINUTE" | "HOUR" | "DAY") {
return None;
}
Some(match kind {
DbKind::Postgres => format!("((now() + INTERVAL '{amount} {unit}')::text)"),
DbKind::MySql => format!("DATE_ADD(UTC_TIMESTAMP(), INTERVAL {amount} {unit})"),
DbKind::Sqlite => return None,
})
}
fn rewrite_scalar_extrema(sql: &str) -> String {
let mut output = String::with_capacity(sql.len() + 8);
let mut remaining = sql;
loop {
let lower = remaining.to_ascii_lowercase();
let max_at = lower.find("max(").map(|at| (at, "GREATEST", 3));
let min_at = lower.find("min(").map(|at| (at, "LEAST", 3));
let Some((start, replacement, name_len)) = (match (max_at, min_at) {
(Some(max), Some(min)) => Some(if max.0 <= min.0 { max } else { min }),
(Some(found), None) | (None, Some(found)) => Some(found),
(None, None) => None,
}) else {
output.push_str(remaining);
break;
};
output.push_str(&remaining[..start]);
let open = start + name_len;
let Some(close) = matching_paren(remaining, open) else {
output.push_str(&remaining[start..]);
break;
};
let arguments = &remaining[open + 1..close];
if split_top_level(arguments).len() > 1 {
output.push_str(replacement);
output.push('(');
output.push_str(arguments);
output.push(')');
} else {
output.push_str(&remaining[start..=close]);
}
remaining = &remaining[close + 1..];
}
output
}
fn rewrite_json_calls(kind: DbKind, sql: &str) -> String {
if kind == DbKind::Sqlite {
return sql.to_string();
}
let mut output = String::with_capacity(sql.len() + 32);
let mut remaining = sql;
loop {
let lower = remaining.to_ascii_lowercase();
let extract = lower.find("json_extract(").map(|at| (at, "json_extract"));
let json_type = lower.find("json_type(").map(|at| (at, "json_type"));
let Some((start, function)) = (match (extract, json_type) {
(Some(left), Some(right)) => Some(if left.0 <= right.0 { left } else { right }),
(Some(found), None) | (None, Some(found)) => Some(found),
(None, None) => None,
}) else {
output.push_str(remaining);
break;
};
output.push_str(&remaining[..start]);
let open = start + function.len();
let Some(close) = matching_paren(remaining, open) else {
output.push_str(&remaining[start..]);
break;
};
let arguments = split_top_level(&remaining[open + 1..close]);
let replacement = match arguments.as_slice() {
[document, path] => json_expression(kind, function, document, path),
_ => None,
};
if let Some(replacement) = replacement {
output.push_str(&replacement);
} else {
output.push_str(&remaining[start..=close]);
}
remaining = &remaining[close + 1..];
}
output
}
fn json_expression(kind: DbKind, function: &str, document: &str, path: &str) -> Option<String> {
let path = path.trim().strip_prefix("'$.")?.strip_suffix('\'')?;
let pg_path = path.split('.').collect::<Vec<_>>().join(",");
Some(match (kind, function) {
(DbKind::Postgres, "json_extract") => {
format!("(CAST({document} AS jsonb) #>> '{{{pg_path}}}')")
}
(DbKind::Postgres, "json_type") => {
format!("jsonb_typeof(CAST({document} AS jsonb) #> '{{{pg_path}}}')")
}
(DbKind::MySql, "json_extract") => {
format!("JSON_UNQUOTE(JSON_EXTRACT({document}, '$.{path}'))")
}
(DbKind::MySql, "json_type") => format!("JSON_TYPE(JSON_EXTRACT({document}, '$.{path}'))"),
_ => return None,
})
}
fn rewrite_misc_functions(kind: DbKind, sql: &str) -> String {
if kind == DbKind::Sqlite {
return sql.to_string();
}
static SQLITE_INDEX_HINT: OnceLock<regex::Regex> = OnceLock::new();
let sql = SQLITE_INDEX_HINT
.get_or_init(|| {
regex::Regex::new(r"(?i)\s+INDEXED\s+BY\s+[A-Za-z_][A-Za-z0-9_]*")
.expect("valid SQLite index hint regex")
})
.replace_all(sql, "")
.into_owned();
static SQLITE_NOCASE_COLLATION: OnceLock<regex::Regex> = OnceLock::new();
let sql = SQLITE_NOCASE_COLLATION
.get_or_init(|| {
regex::Regex::new(r"(?i)\s+COLLATE\s+NOCASE\b")
.expect("valid SQLite NOCASE collation regex")
})
.replace_all(&sql, "")
.into_owned();
let sql = if kind == DbKind::Postgres {
sql.replace(
"GROUP_CONCAT(DISTINCT s.name)",
"string_agg(DISTINCT s.name, ',')",
)
} else {
sql.to_string()
};
let mut output = String::with_capacity(sql.len() + 16);
let mut remaining = sql.as_str();
loop {
let lower = remaining.to_ascii_lowercase();
let Some(start) = lower.find("printf(") else {
output.push_str(remaining);
break;
};
output.push_str(&remaining[..start]);
let open = start + "printf".len();
let Some(close) = matching_paren(remaining, open) else {
output.push_str(&remaining[start..]);
break;
};
let arguments = split_top_level(&remaining[open + 1..close]);
let replacement = match arguments.as_slice() {
[format, value] if format.trim() == "'%.6f'" => Some(match kind {
DbKind::Postgres => format!("to_char({value}, 'FM999999990.000000')"),
DbKind::MySql => format!("FORMAT({value}, 6)"),
DbKind::Sqlite => unreachable!(),
}),
[format, value] if format.trim() == "'$%.6f'" => Some(match kind {
DbKind::Postgres => {
format!("('$' || to_char({value}, 'FM999999990.000000'))")
}
DbKind::MySql => format!("CONCAT('$', FORMAT({value}, 6))"),
DbKind::Sqlite => unreachable!(),
}),
_ => None,
};
if let Some(replacement) = replacement {
output.push_str(&replacement);
} else {
output.push_str(&remaining[start..=close]);
}
remaining = &remaining[close + 1..];
}
output
}
fn rewrite_postgres_boolean_literals(sql: &str) -> String {
const BOOLEAN_COLUMNS: &str = "enabled|is_default|checkin_enabled|is_pinned|notify_on_failure|notify_on_success|snapshot_sealed|read|is_manual|force_reversal|downgrade_decision|recover_applied|manual_override|post_refresh_probe_enabled|use_system_proxy";
static COALESCE: OnceLock<regex::Regex> = OnceLock::new();
static COMPARISON: OnceLock<regex::Regex> = OnceLock::new();
static COALESCE_COMPARISON: OnceLock<regex::Regex> = OnceLock::new();
let column = format!(r#"(?:[A-Za-z_][A-Za-z0-9_]*\.)?"?(?:{BOOLEAN_COLUMNS})"?"#);
let coalesce = COALESCE.get_or_init(|| {
regex::Regex::new(&format!(
r#"(?i)\bCOALESCE\(\s*({column})\s*,\s*([01])\s*\)"#
))
.expect("valid PostgreSQL boolean COALESCE regex")
});
let comparison = COMPARISON.get_or_init(|| {
regex::Regex::new(&format!(r#"(?i)\b({column}\s*(?:=|<>|!=)\s*)([01])\b"#))
.expect("valid PostgreSQL boolean comparison regex")
});
let coalesce_comparison = COALESCE_COMPARISON.get_or_init(|| {
regex::Regex::new(&format!(
r#"(?i)(COALESCE\(\s*{column}\s*,\s*(?:TRUE|FALSE)\s*\)\s*(?:=|<>|!=)\s*)([01])\b"#
))
.expect("valid PostgreSQL boolean COALESCE comparison regex")
});
let sql = coalesce.replace_all(sql, |captures: &regex::Captures<'_>| {
format!(
"COALESCE({}, {})",
&captures[1],
if &captures[2] == "1" { "TRUE" } else { "FALSE" }
)
});
let sql = comparison.replace_all(&sql, |captures: &regex::Captures<'_>| {
format!(
"{}{}",
&captures[1],
if &captures[2] == "1" { "TRUE" } else { "FALSE" }
)
});
coalesce_comparison
.replace_all(&sql, |captures: &regex::Captures<'_>| {
format!(
"{}{}",
&captures[1],
if &captures[2] == "1" { "TRUE" } else { "FALSE" }
)
})
.into_owned()
}
fn rewrite_postgres_bigint_sums(sql: &str) -> String {
const INTEGER_COLUMNS: &str = "total_tokens|prompt_tokens|completion_tokens|input_tokens|output_tokens|latency_ms|success_count|fail_count|total_calls|used_requests";
static SUM: OnceLock<regex::Regex> = OnceLock::new();
static AVG: OnceLock<regex::Regex> = OnceLock::new();
let sum = SUM.get_or_init(|| {
regex::Regex::new(&format!(
r#"(?i)COALESCE\(\s*SUM\(\s*((?:[A-Za-z_][A-Za-z0-9_]*\.)?(?:{INTEGER_COLUMNS}))\s*\)\s*,\s*0\s*\)"#
))
.expect("valid PostgreSQL BIGINT SUM regex")
});
let sql = sum.replace_all(sql, |captures: &regex::Captures<'_>| {
format!("CAST(COALESCE(SUM({}), 0) AS BIGINT)", &captures[1])
});
let avg = AVG.get_or_init(|| {
regex::Regex::new(&format!(
r#"(?i)AVG\(\s*((?:[A-Za-z_][A-Za-z0-9_]*\.)?(?:{INTEGER_COLUMNS}))\s*\)"#
))
.expect("valid PostgreSQL BIGINT AVG regex")
});
avg.replace_all(&sql, |captures: &regex::Captures<'_>| {
format!("CAST(AVG({}) AS DOUBLE PRECISION)", &captures[1])
})
.into_owned()
}
fn matching_paren(sql: &str, open: usize) -> Option<usize> {
let mut depth = 0_u32;
for (offset, ch) in sql[open..].char_indices() {
match ch {
'(' => depth += 1,
')' => {
depth = depth.checked_sub(1)?;
if depth == 0 {
return Some(open + offset);
}
}
_ => {}
}
}
None
}
fn replace_excluded_references(input: &str) -> String {
let mut output = String::with_capacity(input.len());
let mut remaining = input;
loop {
let lower = remaining.to_ascii_lowercase();
let Some(at) = lower.find("excluded.") else {
output.push_str(remaining);
break;
};
output.push_str(&remaining[..at]);
let identifier = &remaining[at + "excluded.".len()..];
let end = identifier
.find(|ch: char| !(ch.is_ascii_alphanumeric() || ch == '_'))
.unwrap_or(identifier.len());
output.push_str("VALUES(");
output.push_str(&identifier[..end]);
output.push(')');
remaining = &identifier[end..];
}
output
}
#[cfg(test)]
mod tests {
use super::{
dialect_sql, mysql_compat_sql, mysql_upsert_sql, rewrite_datetime_calls,
rewrite_postgres_bigint_sums, rewrite_postgres_boolean_literals, DbKind,
};
#[test]
fn rewrites_only_unquoted_placeholders() {
assert_eq!(
dialect_sql(
DbKind::Postgres,
"SELECT '?', \"?\" FROM users WHERE id = ? AND note = 'it''s ?' AND role = ? -- ?\n"
),
"SELECT '?', \"?\" FROM users WHERE id = $1 AND note = 'it''s ?' AND role = $2 -- ?\n"
);
}
#[test]
fn removes_sqlite_nocase_collation_for_server_databases() {
assert_eq!(
dialect_sql(
DbKind::Postgres,
"SELECT model_pattern FROM token_routes ORDER BY model_pattern COLLATE NOCASE"
),
"SELECT model_pattern FROM token_routes ORDER BY model_pattern"
);
assert_eq!(
dialect_sql(
DbKind::MySql,
"SELECT model_pattern FROM token_routes ORDER BY model_pattern collate nocase"
),
"SELECT model_pattern FROM token_routes ORDER BY model_pattern"
);
assert_eq!(
dialect_sql(
DbKind::Sqlite,
"SELECT model_pattern FROM token_routes ORDER BY model_pattern COLLATE NOCASE"
),
"SELECT model_pattern FROM token_routes ORDER BY model_pattern COLLATE NOCASE"
);
}
#[test]
fn rewrites_mysql_upserts() {
assert_eq!(
mysql_upsert_sql(
"INSERT INTO settings (key, value) VALUES (?, ?) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value"
),
"INSERT INTO settings (key, value) VALUES (?, ?) ON DUPLICATE KEY UPDATE value = VALUES(value)"
);
assert_eq!(
mysql_upsert_sql(
"INSERT INTO organization_members (org_id, user_id, role) VALUES (?, ?, ?) ON CONFLICT DO NOTHING"
),
"INSERT INTO organization_members (org_id, user_id, role) VALUES (?, ?, ?) ON DUPLICATE KEY UPDATE org_id = org_id"
);
}
#[test]
fn quotes_mysql_settings_key() {
assert_eq!(
mysql_compat_sql("SELECT value FROM settings WHERE key = ?"),
"SELECT value FROM settings WHERE `key` = ?"
);
}
#[test]
fn rewrites_portable_datetime_calls() {
assert_eq!(
rewrite_datetime_calls(
DbKind::MySql,
"expires_at IS NULL OR datetime(expires_at) > CURRENT_TIMESTAMP"
),
"expires_at IS NULL OR CAST(REPLACE(SUBSTRING(expires_at, 1, 19), 'T', ' ') AS DATETIME) > CURRENT_TIMESTAMP"
);
assert_eq!(
rewrite_datetime_calls(DbKind::Postgres, "datetime('now', '-1 day')"),
"((now() + INTERVAL '-1 DAY')::text)"
);
}
#[test]
fn rewrites_only_known_postgres_boolean_literals() {
assert_eq!(
rewrite_postgres_boolean_literals(
"UPDATE events SET read = 1 WHERE read = 0 AND id = 1"
),
"UPDATE events SET read = TRUE WHERE read = FALSE AND id = 1"
);
assert_eq!(
rewrite_postgres_boolean_literals(
"WHERE (? = 1 OR tr.enabled = 1) AND COALESCE(rc.manual_override, 0) = 0 ORDER BY COALESCE(at.is_default, 0) DESC"
),
"WHERE (? = 1 OR tr.enabled = TRUE) AND COALESCE(rc.manual_override, FALSE) = FALSE ORDER BY COALESCE(at.is_default, FALSE) DESC"
);
assert_eq!(
dialect_sql(
DbKind::Postgres,
"INSERT INTO model_availability (available, is_manual) VALUES (1, TRUE)"
),
"INSERT INTO model_availability (available, is_manual) VALUES (1, TRUE)"
);
assert_eq!(
dialect_sql(
DbKind::Postgres,
"INSERT INTO route_channels (weight, enabled, manual_override) VALUES (10, TRUE, FALSE)"
),
"INSERT INTO route_channels (weight, enabled, manual_override) VALUES (10, TRUE, FALSE)"
);
}
#[test]
fn casts_postgres_bigint_sums_without_touching_float_sums() {
assert_eq!(
rewrite_postgres_bigint_sums(
"SELECT COALESCE(SUM(l.total_tokens), 0), COALESCE(SUM(cost), 0.0) FROM usage_logs l"
),
"SELECT CAST(COALESCE(SUM(l.total_tokens), 0) AS BIGINT), COALESCE(SUM(cost), 0.0) FROM usage_logs l"
);
assert_eq!(
rewrite_postgres_bigint_sums("SELECT AVG(l.latency_ms), AVG(cost) FROM proxy_logs l"),
"SELECT CAST(AVG(l.latency_ms) AS DOUBLE PRECISION), AVG(cost) FROM proxy_logs l"
);
}
#[test]
fn removes_sqlite_index_hints_from_server_dialects() {
assert_eq!(
dialect_sql(
DbKind::Postgres,
"SELECT * FROM workspace_assets INDEXED BY workspace_assets_gallery_created_idx WHERE id = ?"
),
"SELECT * FROM workspace_assets WHERE id = $1"
);
}
}
+85 -12
View File
@@ -30,6 +30,20 @@ pub enum AppError {
retry_after_secs: u64,
},
#[error("Request rate limited: {message}")]
RequestRateLimited {
message: String,
retry_after_secs: u64,
},
#[error("Billing required: {message}")]
BillingRequired {
message: String,
billing_track: &'static str,
required_micros: i64,
available_micros: i64,
},
#[error("Service unavailable: {0}")]
ServiceUnavailable(String),
@@ -55,6 +69,12 @@ impl IntoResponse for AppError {
AppError::RateLimitedWithRetry { message, .. } => {
(StatusCode::TOO_MANY_REQUESTS, message.clone())
}
AppError::RequestRateLimited { message, .. } => {
(StatusCode::TOO_MANY_REQUESTS, message.clone())
}
AppError::BillingRequired { message, .. } => {
(StatusCode::PAYMENT_REQUIRED, message.clone())
}
AppError::ServiceUnavailable(msg) => (StatusCode::SERVICE_UNAVAILABLE, msg.clone()),
AppError::GatewayTimeout(msg) => (StatusCode::GATEWAY_TIMEOUT, msg.clone()),
AppError::Internal(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
@@ -63,30 +83,83 @@ impl IntoResponse for AppError {
let retry_after_secs = match &self {
AppError::RateLimitedWithRetry {
retry_after_secs, ..
}
| AppError::RequestRateLimited {
retry_after_secs, ..
} => Some(*retry_after_secs),
_ => None,
};
let billing = match &self {
AppError::RateLimited | AppError::RateLimitedWithRetry { .. } => Some(json!({
"reason": "insufficient_quota",
"billingTrack": "api",
"rechargePath": "/user/recharge",
"packagePath": "/user/recharge?mode=api-package",
"usagePath": "/user/usage",
})),
AppError::BillingRequired {
billing_track,
required_micros,
available_micros,
..
} => Some(json!({
"reason": "insufficient_balance",
"billingTrack": billing_track,
"requiredMicros": required_micros,
"availableMicros": available_micros,
"shortfallMicros": required_micros.saturating_sub(*available_micros),
"rechargePath": format!("/user/recharge?mode={billing_track}-balance"),
"packagePath": format!("/user/recharge?mode={billing_track}-package"),
"usagePath": "/user/usage",
})),
_ => None,
};
(
status,
Json(json!({
"error": {
"message": message,
"type": status.canonical_reason(),
},
"retry_after_secs": retry_after_secs,
"billing": billing,
})),
)
.into_response()
let mut body = json!({
"error": {
"message": message,
"type": status.canonical_reason(),
},
"retry_after_secs": retry_after_secs,
});
if let Some(billing) = billing {
body["billing"] = billing;
}
(status, Json(body)).into_response()
}
}
pub type AppResult<T> = Result<T, AppError>;
#[cfg(test)]
mod tests {
use super::*;
use axum::body::to_bytes;
#[tokio::test]
async fn request_rate_limit_has_retry_without_billing_metadata() {
let response = AppError::RequestRateLimited {
message: "too many attempts".into(),
retry_after_secs: 45,
}
.into_response();
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(payload["retry_after_secs"], 45);
assert!(payload.get("billing").is_none());
}
#[tokio::test]
async fn quota_rate_limit_keeps_billing_metadata() {
let response = AppError::RateLimitedWithRetry {
message: "quota exhausted".into(),
retry_after_secs: 60,
}
.into_response();
let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(payload["billing"]["reason"], "insufficient_quota");
}
}
+46 -6
View File
@@ -10,27 +10,53 @@ pub mod template_http;
pub mod transformers;
pub mod workflow;
use sqlx::SqlitePool;
use std::{sync::Arc, time::Duration};
use sqlx::AnyPool;
use std::{
sync::{
atomic::{AtomicBool, Ordering},
Arc,
},
time::Duration,
};
#[derive(Clone)]
pub struct AppState {
pub config: config::AppConfig,
pub db: Arc<SqlitePool>,
pub db: Arc<AnyPool>,
pub http: reqwest::Client,
pub comfyui: Arc<services::comfyui_process::ComfyUIProcess>,
pub builtin_keys: Arc<services::builtin_keys::BuiltinKeys>,
pub runtime_auth: Arc<services::runtime_auth::RuntimeAuthSettings>,
pub register_captcha: Arc<services::register_captcha::RegisterCaptchaStore>,
pub registration_limiter: Arc<services::registration_limiter::RegistrationLimiter>,
pub registration_email_codes: Arc<services::registration_email::RegistrationEmailCodeStore>,
pub user_oauth_flows: Arc<services::user_oauth_flow::UserOAuthFlowStore>,
pub redis: Arc<services::redis::RedisStore>,
maintenance_mode: Arc<AtomicBool>,
}
impl AppState {
pub fn new(
config: config::AppConfig,
db: SqlitePool,
db: AnyPool,
builtin_keys: services::builtin_keys::BuiltinKeys,
runtime_auth: services::runtime_auth::RuntimeAuthSettings,
) -> Self {
Self::new_with_redis(
config,
db,
builtin_keys,
runtime_auth,
services::redis::RedisStore::disabled(),
)
}
pub fn new_with_redis(
config: config::AppConfig,
db: AnyPool,
builtin_keys: services::builtin_keys::BuiltinKeys,
runtime_auth: services::runtime_auth::RuntimeAuthSettings,
redis: services::redis::RedisStore,
) -> Self {
let boosapi_host = if config.server.host == "0.0.0.0" || config.server.host == "::" {
"127.0.0.1".to_string()
@@ -60,7 +86,7 @@ impl AppState {
let db = Arc::new(db);
// Hand the payment module a process-wide DB handle so the Comfy Firebase
// refresh path can persist a rotated refresh token without threading
// &SqlitePool through every balance/checkout call site.
// &AnyPool through every balance/checkout call site.
services::payment::init_payment_db(db.clone());
Self {
@@ -72,12 +98,26 @@ impl AppState {
runtime_auth: Arc::new(runtime_auth),
register_captcha: Arc::new(services::register_captcha::RegisterCaptchaStore::new()),
registration_limiter: Arc::new(
services::registration_limiter::RegistrationLimiter::new(),
services::registration_limiter::RegistrationLimiter::new(redis.clone()),
),
registration_email_codes: Arc::new(
services::registration_email::RegistrationEmailCodeStore::new(),
),
user_oauth_flows: Arc::new(services::user_oauth_flow::UserOAuthFlowStore::new()),
redis: Arc::new(redis),
maintenance_mode: Arc::new(AtomicBool::new(false)),
}
}
pub fn runtime_proxy_token(&self) -> String {
self.runtime_auth.proxy_token()
}
pub fn maintenance_mode(&self) -> bool {
self.maintenance_mode.load(Ordering::Acquire)
}
pub fn set_maintenance_mode(&self, enabled: bool) {
self.maintenance_mode.store(enabled, Ordering::Release);
}
}
+183 -11
View File
@@ -1,7 +1,7 @@
use anyhow::Context;
use anyhow::{bail, Context};
use axum::serve;
use boosapi::{
config::AppConfig,
config::{AppConfig, PgrustConfig, RedisConfig},
db,
services::{builtin_keys, runtime_auth},
AppState,
@@ -9,6 +9,7 @@ use boosapi::{
use clap::{Parser, Subcommand};
use std::ffi::OsString;
use std::net::SocketAddr;
use std::path::Path;
use tokio::signal;
use tracing_subscriber::{prelude::*, EnvFilter};
@@ -50,6 +51,17 @@ struct Cli {
#[derive(Subcommand, Debug)]
enum Commands {
/// Database schema and full-data migration tools.
Database {
#[command(subcommand)]
command: DatabaseCommands,
},
/// Apply PostgreSQL schema migrations without copying or truncating data.
MigratePostgres {
/// PostgreSQL URL, e.g. postgres://user:password@host:5432/boosapi
#[arg(long)]
postgres_url: String,
},
/// Copy core data from SQLite into PostgreSQL and exit
MigrateSqliteToPostgres {
/// Source SQLite URL, e.g. sqlite:boosapi.db?mode=rwc
@@ -66,19 +78,70 @@ enum Commands {
},
}
#[derive(Subcommand, Debug)]
enum DatabaseCommands {
/// Copy and verify all business data between supported databases.
Migrate {
#[arg(long)]
source_url: String,
#[arg(long)]
target_url: String,
/// Delete existing target business data before copying.
#[arg(long)]
force: bool,
#[arg(long, default_value_t = 500)]
batch_size: usize,
/// Compare row counts and content digests without copying.
#[arg(long)]
verify_only: bool,
/// Save the verified target URL to the 0600 runtime override file.
#[arg(long)]
save_runtime_config: bool,
},
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let args = Cli::parse_from(normalize_legacy_args());
if let Some(Commands::Database {
command:
DatabaseCommands::Migrate {
source_url,
target_url,
force,
batch_size,
verify_only,
save_runtime_config,
},
}) = &args.command
{
let summary =
db::migrate::migrate(source_url, target_url, *force, *batch_size, *verify_only).await?;
if *save_runtime_config {
save_runtime_database_config(&args.config, target_url)?;
}
println!("{}", serde_json::to_string_pretty(&summary)?);
return Ok(());
}
if let Some(Commands::MigratePostgres { postgres_url }) = &args.command {
let (pool, _) = db::connect_any_pool(postgres_url, 2).await?;
db::run_migrations(db::DbKind::Postgres, &pool).await?;
println!("PostgreSQL migrations applied.");
return Ok(());
}
if let Some(Commands::MigrateSqliteToPostgres {
sqlite_url,
postgres_url,
force_truncate,
}) = &args.command
{
let sqlite = db::create_pool(sqlite_url, 1).await?;
let sqlite = db::create_sqlite_pool(sqlite_url, 1).await?;
let pg = db::create_pg_pool(postgres_url, 2).await?;
db::run_migrations(db::DbKind::Postgres, None, Some(&pg)).await?;
let (migration_pool, _) = db::connect_any_pool(postgres_url, 2).await?;
db::run_migrations(db::DbKind::Postgres, &migration_pool).await?;
let summary = db::sqlite_to_postgres::migrate(&sqlite, &pg, *force_truncate).await?;
println!("{}", serde_json::to_string_pretty(&summary)?);
return Ok(());
@@ -86,6 +149,8 @@ async fn main() -> anyhow::Result<()> {
// Load config
let mut config = AppConfig::load_from_file(&args.config)?;
let redis_config = RedisConfig::load_from_file(&args.config)?;
let pgrust_config = PgrustConfig::load_from_file(&args.config)?;
// Apply CLI overrides
if let Some(host) = args.host {
@@ -97,17 +162,33 @@ async fn main() -> anyhow::Result<()> {
if let Some(level) = args.log_level {
config.server.log_level = level;
}
let config_path = Path::new(&args.config);
let use_embedded_pgrust = boosapi::services::pgrust::should_start(&pgrust_config, config_path);
if !use_embedded_pgrust {
config.db.url = config.db.resolved_url()?;
}
if args.check_config {
println!("Configuration is valid.");
println!("Server: {}:{}", config.server.host, config.server.port);
println!("Database URL: {}", config.db.url);
if use_embedded_pgrust {
println!("Database: embedded pgrust");
} else {
println!("Database URL: {}", config.db.url);
}
println!(
"Redis: {}",
if redis_config.enabled {
"enabled"
} else {
"disabled"
}
);
return Ok(());
}
// Initialize tracing
let filter = EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new(&config.server.log_level));
let filter = build_log_filter(&config.server.log_level)?;
tracing_subscriber::registry()
.with(filter)
@@ -117,19 +198,48 @@ async fn main() -> anyhow::Result<()> {
tracing::info!("Loaded configuration from {}", args.config);
let db_kind = db::DbKind::from_url(&config.db.url)?;
let mut pgrust_process = if use_embedded_pgrust {
let process = boosapi::services::pgrust::PgrustProcess::start(&pgrust_config, config_path)
.await
.context("failed to start embedded pgrust")?;
config.db.url = process.database_url().to_string();
tracing::info!(pid = ?process.pid(), "Embedded pgrust ready");
Some(process)
} else {
None
};
// Create database pool
let db_kind = db::DbKind::from_url(&config.db.url)?;
let pool = db::create_pool(&config.db.url, config.db.max_connections).await?;
tracing::info!("Database connected");
// Run migrations
db::run_migrations(db_kind, Some(&pool), None).await?;
db::schema_compat::ensure_sqlite_upstream_compat(&pool).await?;
db::run_migrations(db_kind, &pool).await?;
tracing::info!("Database migrations applied");
let interrupted =
boosapi::services::agent_workbench::session_store::interrupt_orphaned_runs(&pool)
.await
.context("failed to interrupt orphaned Agent runs")?;
if interrupted > 0 {
tracing::warn!(
interrupted,
"interrupted Agent runs left active by a prior process"
);
}
if boosapi::services::agent_workbench::bash_sandbox::sandbox_available() {
tracing::info!("Agent Bash sandbox self-check passed");
} else {
tracing::warn!("Agent Bash disabled: bwrap/prlimit sandbox self-check failed");
}
if args.migrate_only {
tracing::info!("Migration only mode requested. Exiting.");
pool.close().await;
if let Some(process) = pgrust_process.as_mut() {
process.stop().await?;
}
return Ok(());
}
@@ -148,9 +258,24 @@ async fn main() -> anyhow::Result<()> {
runtime_auth::RuntimeAuthSettings::load(&pool, config.auth.proxy_token.clone())
.await
.context("failed to load runtime auth settings")?;
let redis = boosapi::services::redis::RedisStore::connect(&redis_config)
.await
.context("failed to initialize Redis")?;
if redis.is_enabled() {
tracing::info!(
key_prefix = %redis.key_prefix(),
connection_timeout_ms = redis_config.connection_timeout_ms,
response_timeout_ms = redis_config.response_timeout_ms,
"Redis 7+ connected"
);
} else {
tracing::info!("Redis integration disabled; using process-local registration limiter");
}
// Build app state
let state = AppState::new(config.clone(), pool, builtin, runtime_auth);
let state = AppState::new_with_redis(config.clone(), pool, builtin, runtime_auth, redis);
let payment_job_worker = boosapi::services::payment_jobs::spawn(state.clone());
let account_health_worker = boosapi::services::account_health::spawn(state.clone());
if let Err(err) = state.comfyui.start().await {
tracing::warn!(%err, "ComfyUI embedded process did not start");
}
@@ -170,13 +295,60 @@ async fn main() -> anyhow::Result<()> {
.with_graceful_shutdown(shutdown_signal())
.await?;
if let Some(worker) = payment_job_worker {
worker.abort();
let _ = worker.await;
}
account_health_worker.abort();
let _ = account_health_worker.await;
if let Err(err) = comfyui_process.stop().await {
tracing::warn!(%err, "ComfyUI embedded process did not stop cleanly");
}
if let Some(process) = pgrust_process.as_mut() {
if let Err(err) = process.stop().await {
tracing::warn!(%err, "Embedded pgrust did not stop cleanly");
}
}
Ok(())
}
fn build_log_filter(log_level: &str) -> anyhow::Result<EnvFilter> {
if let Some(raw) = std::env::var_os("RUST_LOG") {
let raw = raw
.into_string()
.map_err(|_| anyhow::anyhow!("RUST_LOG must be valid UTF-8"))?;
return EnvFilter::try_new(raw).context("invalid RUST_LOG filter");
}
let app_level = log_level.trim().to_ascii_lowercase();
if !matches!(
app_level.as_str(),
"trace" | "debug" | "info" | "warn" | "error"
) {
bail!("invalid server.log_level {log_level:?}; expected trace, debug, info, warn or error");
}
let request_level = if matches!(app_level.as_str(), "trace" | "debug") {
app_level.as_str()
} else {
"warn"
};
let directives = format!(
"warn,boosapi={app_level},boosapi::routes={request_level},boosapi::proxy={request_level},sqlx=warn,sqlx::query=error"
);
EnvFilter::try_new(directives).context("build default log filter")
}
fn save_runtime_database_config(config_path: &str, target_url: &str) -> anyhow::Result<()> {
let path = boosapi::config::save_runtime_database_url_at(config_path, target_url)?;
eprintln!(
"Saved verified database URL to {}; restart BoosAPI to activate it.",
path.display()
);
Ok(())
}
fn normalize_legacy_args() -> Vec<OsString> {
std::env::args_os()
.map(|arg| {
+11 -3
View File
@@ -104,7 +104,10 @@ mod tests {
fn lists_all_four_templates() {
let templates = list_templates("https://api.boos.lat");
let ids: Vec<&str> = templates.iter().map(|t| t.id.as_str()).collect();
assert_eq!(ids, ["cursor", "vscode", "claude_code", "openai_compatible"]);
assert_eq!(
ids,
["cursor", "vscode", "claude_code", "openai_compatible"]
);
}
#[test]
@@ -124,7 +127,9 @@ mod tests {
fn claude_code_base_has_no_v1_suffix() {
// Claude Code appends /v1/messages itself; a /v1 base would double it.
let t = find_template("claude_code", "https://api.boos.lat").unwrap();
assert!(t.content.contains("ANTHROPIC_BASE_URL=\"https://api.boos.lat\""));
assert!(t
.content
.contains("ANTHROPIC_BASE_URL=\"https://api.boos.lat\""));
assert!(!t.content.contains("api.boos.lat/v1"));
}
@@ -146,7 +151,10 @@ mod tests {
let value = serde_json::to_value(&t).unwrap();
assert_eq!(value["id"], "cursor");
assert_eq!(value["format"], "json");
assert!(value["content"].as_str().unwrap().contains("localhost:4000/v1"));
assert!(value["content"]
.as_str()
.unwrap()
.contains("localhost:4000/v1"));
assert!(value["copy_hint"].is_string());
}
}
+1
View File
@@ -269,6 +269,7 @@ mod tests {
account_id: Some(30),
account_unit_cost: None,
base_url: "https://example.test/v1".into(),
proxy_url: None,
api_key: "key".into(),
model: "gpt-5.5".into(),
actual_model: "gpt-5.5".into(),
+1
View File
@@ -13,6 +13,7 @@ pub struct SelectedChannel {
pub account_id: Option<i64>,
pub account_unit_cost: Option<f64>,
pub base_url: String,
pub proxy_url: Option<String>,
pub api_key: String,
pub model: String,
pub actual_model: String,
+103 -16
View File
@@ -14,16 +14,52 @@ use crate::transformers::{
canonical::CanonicalRequest,
};
use crate::AppState;
use once_cell::sync::Lazy;
use parking_lot::RwLock;
use serde_json::{json, Value};
use std::{collections::HashMap, time::Duration};
static PROXY_HTTP_CLIENTS: Lazy<RwLock<HashMap<String, reqwest::Client>>> =
Lazy::new(|| RwLock::new(HashMap::new()));
/// Resolved upstream info for a given model request.
#[derive(Debug, Clone)]
pub struct UpstreamRoute {
pub base_url: String,
pub proxy_url: Option<String>,
pub api_key: String,
pub model: String,
}
fn http_client_for_route(state: &AppState, route: &UpstreamRoute) -> AppResult<reqwest::Client> {
let Some(proxy_url) = route
.proxy_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(state.http.clone());
};
let timeout_secs = state.config.proxy.upstream_timeout_secs.max(30);
let cache_key = format!("{timeout_secs}\0{proxy_url}");
if let Some(client) = PROXY_HTTP_CLIENTS.read().get(&cache_key) {
return Ok(client.clone());
}
let proxy = reqwest::Proxy::all(proxy_url)
.map_err(|err| AppError::BadRequest(format!("Invalid outbound proxy URL: {err}")))?;
let client = reqwest::Client::builder()
.proxy(proxy)
.timeout(Duration::from_secs(timeout_secs))
.connect_timeout(Duration::from_secs(15))
.pool_idle_timeout(Duration::from_secs(90))
.pool_max_idle_per_host(32)
.tcp_nodelay(true)
.build()
.map_err(|err| AppError::Internal(err.into()))?;
PROXY_HTTP_CLIENTS.write().insert(cache_key, client.clone());
Ok(client)
}
fn parse_upstream_json(
text: &str,
upstream_url: &str,
@@ -65,16 +101,30 @@ fn looks_like_sse(text: &str) -> bool {
text.trim_start().starts_with("data:")
}
fn openai_compatible_base_url(base_url: &str) -> String {
pub(crate) fn openai_compatible_base_url(base_url: &str) -> String {
let trimmed = base_url.trim().trim_end_matches('/');
if trimmed.ends_with("/v1") || trimmed.contains("/v1/") {
let has_version_segment = reqwest::Url::parse(trimmed)
.ok()
.and_then(|url| {
url.path_segments().map(|segments| {
segments
.filter(|segment| !segment.is_empty())
.any(|segment| {
segment.strip_prefix('v').is_some_and(|version| {
!version.is_empty() && version.chars().all(|ch| ch.is_ascii_digit())
})
})
})
})
.unwrap_or(false);
if has_version_segment {
trimmed.to_string()
} else {
format!("{trimmed}/v1")
}
}
fn upstream_url(base_url: &str, path_suffix: &str) -> String {
pub(crate) fn openai_compatible_url(base_url: &str, path_suffix: &str) -> String {
format!(
"{}/{}",
openai_compatible_base_url(base_url),
@@ -82,6 +132,10 @@ fn upstream_url(base_url: &str, path_suffix: &str) -> String {
)
}
fn upstream_url(base_url: &str, path_suffix: &str) -> String {
openai_compatible_url(base_url, path_suffix)
}
fn aggregate_chat_sse(text: &str, fallback_model: &str) -> Option<Value> {
let mut id: Option<String> = None;
let mut created: Option<i64> = None;
@@ -167,7 +221,7 @@ pub async fn execute_chat_completions(
route: UpstreamRoute,
body: Value,
) -> AppResult<Value> {
let client = state.http.clone();
let client = http_client_for_route(&state, &route)?;
let upstream_url = upstream_url(&route.base_url, "/chat/completions");
let resp = client
@@ -221,7 +275,7 @@ pub async fn execute_anthropic_messages(
route: UpstreamRoute,
request: CanonicalRequest,
) -> AppResult<Value> {
let client = state.http.clone();
let client = http_client_for_route(&state, &route)?;
let upstream_url = format!("{}/messages", route.base_url.trim_end_matches('/'));
let compatibility = AnthropicMessagesCompatibility::default();
let body = format_messages_request_with_compat(&request, &compatibility);
@@ -273,7 +327,7 @@ pub async fn execute_anthropic_messages_stream_response(
route: UpstreamRoute,
request: CanonicalRequest,
) -> AppResult<reqwest::Response> {
let client = state.http.clone();
let client = http_client_for_route(&state, &route)?;
let upstream_url = format!("{}/messages", route.base_url.trim_end_matches('/'));
let compatibility = AnthropicMessagesCompatibility::default();
let body = format_messages_request_with_compat(&request, &compatibility);
@@ -459,7 +513,7 @@ pub async fn execute_streaming_response(
path_suffix: &str,
body: Value,
) -> AppResult<reqwest::Response> {
let client = state.http.clone();
let client = http_client_for_route(&state, &route)?;
let upstream_url = upstream_url(&route.base_url, path_suffix);
let resp = client
@@ -492,7 +546,7 @@ pub async fn execute_proxy_request(
path_suffix: &str,
body: Value,
) -> AppResult<Value> {
let client = state.http.clone();
let client = http_client_for_route(&state, &route)?;
let upstream_url = upstream_url(&route.base_url, path_suffix);
let resp = client
@@ -527,7 +581,7 @@ pub async fn execute_proxy_get(
route: UpstreamRoute,
path_suffix: &str,
) -> AppResult<Value> {
let client = state.http.clone();
let client = http_client_for_route(&state, &route)?;
let upstream_url = upstream_url(&route.base_url, path_suffix);
let resp = client
@@ -1141,7 +1195,7 @@ pub async fn execute_image_generation(
route: UpstreamRoute,
body: Value,
) -> AppResult<Value> {
let client = state.http.clone();
let client = http_client_for_route(&state, &route)?;
let upstream_url = upstream_url(&route.base_url, "/images/generations");
let resp = client
@@ -1200,7 +1254,7 @@ pub async fn execute_image_edit(
) -> AppResult<Value> {
use reqwest::multipart::{Form, Part};
let client = state.http.clone();
let client = http_client_for_route(&state, &route)?;
let upstream_url = upstream_url(&route.base_url, "/images/edits");
let mut form = Form::new();
@@ -1304,9 +1358,14 @@ pub async fn execute_gemini_operation(
body: Value,
extra_query: Option<&str>,
) -> AppResult<Value> {
let client = state.http.clone();
let upstream_url =
gemini_operation_url(&route.base_url, &route.model, operation, &route.api_key, extra_query);
let client = http_client_for_route(&state, &route)?;
let upstream_url = gemini_operation_url(
&route.base_url,
&route.model,
operation,
&route.api_key,
extra_query,
);
let resp = client
.post(&upstream_url)
@@ -1341,7 +1400,7 @@ pub async fn execute_gemini_stream_response(
body: Value,
extra_query: Option<&str>,
) -> AppResult<reqwest::Response> {
let client = state.http.clone();
let client = http_client_for_route(&state, &route)?;
let query = merge_alt_sse(extra_query);
let upstream_url = gemini_operation_url(
&route.base_url,
@@ -1377,7 +1436,12 @@ pub async fn execute_gemini_stream_response(
/// arrays without it). Preserves any other client-supplied params.
fn merge_alt_sse(extra_query: Option<&str>) -> String {
match extra_query.filter(|q| !q.is_empty()) {
Some(q) if q.split('&').any(|p| p == "alt=sse" || p.starts_with("alt=")) => q.to_string(),
Some(q)
if q.split('&')
.any(|p| p == "alt=sse" || p.starts_with("alt=")) =>
{
q.to_string()
}
Some(q) => format!("{q}&alt=sse"),
None => "alt=sse".to_string(),
}
@@ -1387,6 +1451,29 @@ fn merge_alt_sse(extra_query: Option<&str>) -> String {
mod gemini_tests {
use super::*;
#[test]
fn openai_base_appends_v1_to_an_origin() {
assert_eq!(
openai_compatible_base_url("https://anyrouter.top/"),
"https://anyrouter.top/v1"
);
assert_eq!(
openai_compatible_url("https://anyrouter.top", "/models"),
"https://anyrouter.top/v1/models"
);
}
#[test]
fn openai_base_preserves_existing_versioned_api_paths() {
for base in [
"https://api.example.com/v1",
"https://ark.example.com/api/v3",
"https://open.example.com/api/coding/paas/v4",
] {
assert_eq!(openai_compatible_base_url(base), base);
}
}
#[test]
fn base_url_appends_v1beta_when_missing() {
assert_eq!(
+2 -2
View File
@@ -1,7 +1,7 @@
use crate::proxy::conductor::{FailureInfo, ProxyConductor, SelectedChannel, SuccessMetrics};
use crate::proxy::token_router::TokenRouter;
use futures::FutureExt;
use sqlx::SqlitePool;
use sqlx::AnyPool;
use std::sync::Arc;
#[derive(Clone)]
@@ -10,7 +10,7 @@ pub struct ProxyRuntime {
}
impl ProxyRuntime {
pub fn new(db: Arc<SqlitePool>) -> Self {
pub fn new(db: Arc<AnyPool>) -> Self {
Self {
router: TokenRouter::new(db),
}
+181 -89
View File
@@ -6,7 +6,8 @@ use crate::services::oauth::OAuthService;
use once_cell::sync::Lazy;
use parking_lot::RwLock;
use regex::Regex;
use sqlx::{FromRow, SqlitePool};
use serde_json::Value;
use sqlx::{AnyPool, FromRow};
use std::{
collections::{HashMap, HashSet},
sync::Arc,
@@ -43,7 +44,7 @@ pub fn invalidate_route_cache() {
#[derive(Clone)]
pub struct TokenRouter {
db: Arc<SqlitePool>,
db: Arc<AnyPool>,
}
#[derive(Debug, Clone, FromRow)]
@@ -63,6 +64,10 @@ struct ChannelRow {
oauth_route_unit_strategy: Option<String>,
site_id: i64,
base_url: String,
account_extra_config: Option<String>,
site_proxy_url: Option<String>,
site_use_system_proxy: Option<bool>,
system_proxy_url: Option<String>,
site_api_key: Option<String>,
account_access_token: String,
account_api_token: Option<String>,
@@ -89,6 +94,49 @@ enum RouteUnitStrategy {
StickUntilUnavailable,
}
pub(crate) fn resolve_proxy_url_from_parts(
account_extra_config: Option<&str>,
site_proxy_url: Option<&str>,
site_use_system_proxy: bool,
system_proxy_url: Option<&str>,
) -> Option<String> {
let account_proxy = account_extra_config
.and_then(|raw| serde_json::from_str::<Value>(raw).ok())
.and_then(|value| {
value
.get("proxyUrl")
.or_else(|| value.get("proxy_url"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
});
account_proxy
.or_else(|| {
site_proxy_url
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
})
.or_else(|| {
site_use_system_proxy.then(|| {
system_proxy_url
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
})?
})
}
fn resolved_channel_proxy_url(row: &ChannelRow) -> Option<String> {
resolve_proxy_url_from_parts(
row.account_extra_config.as_deref(),
row.site_proxy_url.as_deref(),
row.site_use_system_proxy.unwrap_or(false),
row.system_proxy_url.as_deref(),
)
}
impl RouteUnitStrategy {
fn from_db(value: Option<&str>) -> Self {
match value.unwrap_or("round_robin") {
@@ -99,7 +147,7 @@ impl RouteUnitStrategy {
}
impl TokenRouter {
pub fn new(db: Arc<SqlitePool>) -> Self {
pub fn new(db: Arc<AnyPool>) -> Self {
Self { db }
}
@@ -121,6 +169,7 @@ impl TokenRouter {
let decision = self.resolve_decision(model, &[]).await?;
Ok(UpstreamRoute {
base_url: decision.selected.base_url,
proxy_url: decision.selected.proxy_url,
api_key: decision.selected.api_key,
model: decision.selected.actual_model,
})
@@ -240,6 +289,7 @@ impl TokenRouter {
.unwrap_or(row.account_access_token.clone())
})
});
let proxy_url = resolved_channel_proxy_url(&row);
SelectedChannel {
channel_id: row.channel_id,
@@ -248,6 +298,7 @@ impl TokenRouter {
account_id: Some(row.account_id),
account_unit_cost: row.account_unit_cost,
base_url: row.base_url,
proxy_url,
api_key: api_key
.trim()
.to_string()
@@ -260,6 +311,7 @@ impl TokenRouter {
}))
}
async fn route_rows(&self) -> AppResult<Vec<RouteRow>> {
let db_cache_key = self.db_cache_key();
if let Some(cache) = ROUTE_ROWS_CACHE.read().get(&db_cache_key) {
@@ -268,7 +320,7 @@ impl TokenRouter {
}
}
let rows: Vec<RouteRow> = sqlx::query_as(
let rows: Vec<RouteRow> = crate::db::sql::query_as(
"SELECT id, model_pattern, model_mapping, route_mode
FROM token_routes
WHERE enabled = 1
@@ -295,7 +347,7 @@ impl TokenRouter {
}
}
let rows: Vec<ChannelRow> = sqlx::query_as(
let rows: Vec<ChannelRow> = crate::db::sql::query_as(
"SELECT
rc.id AS channel_id,
rc.route_id AS route_id,
@@ -308,10 +360,14 @@ impl TokenRouter {
FROM site_api_endpoints sae
WHERE sae.site_id = s.id
AND COALESCE(sae.enabled, 1) = 1
AND (sae.cooldown_until IS NULL OR sae.cooldown_until <= datetime('now'))
AND (sae.cooldown_until IS NULL OR sae.cooldown_until <= CURRENT_TIMESTAMP)
ORDER BY COALESCE(sae.sort_order, 0), sae.id
LIMIT 1
), s.url) AS base_url,
a.extra_config AS account_extra_config,
s.proxy_url AS site_proxy_url,
s.use_system_proxy AS site_use_system_proxy,
(SELECT value FROM settings WHERE key = 'runtime.system_proxy_url' LIMIT 1) AS system_proxy_url,
s.api_key AS site_api_key,
a.access_token AS account_access_token,
a.api_token AS account_api_token,
@@ -334,7 +390,7 @@ impl TokenRouter {
AND a.status = 'active'
AND s.status = 'active'
AND (rc.oauth_route_unit_id IS NULL OR oru.id IS NOT NULL)
AND (rc.cooldown_until IS NULL OR rc.cooldown_until <= datetime('now'))
AND (rc.cooldown_until IS NULL OR rc.cooldown_until <= CURRENT_TIMESTAMP)
AND (rc.token_id IS NULL OR at.id IS NOT NULL)
ORDER BY
COALESCE(rc.priority, 0) DESC,
@@ -358,7 +414,7 @@ impl TokenRouter {
}
pub async fn record_success(&self, channel_id: i64, metrics: SuccessMetrics) -> AppResult<()> {
sqlx::query(
crate::db::sql::query(
"UPDATE route_channels
SET success_count = COALESCE(success_count, 0) + 1,
total_latency_ms = COALESCE(total_latency_ms, 0) + ?,
@@ -366,8 +422,8 @@ impl TokenRouter {
consecutive_fail_count = 0,
cooldown_level = 0,
cooldown_until = NULL,
last_used_at = datetime('now'),
last_selected_at = datetime('now')
last_used_at = CURRENT_TIMESTAMP,
last_selected_at = CURRENT_TIMESTAMP
WHERE id = ?",
)
.bind(metrics.latency_ms.unwrap_or(0) as i64)
@@ -383,20 +439,19 @@ impl TokenRouter {
pub async fn record_failure(&self, channel_id: i64, info: FailureInfo) -> AppResult<()> {
let cooldown_seconds = self.cooldown_seconds_for_failure(&info).await?;
sqlx::query(
let cooldown_until = (cooldown_seconds > 0).then(|| {
(chrono::Utc::now() + chrono::Duration::seconds(cooldown_seconds)).to_rfc3339()
});
crate::db::sql::query(
"UPDATE route_channels
SET fail_count = COALESCE(fail_count, 0) + 1,
consecutive_fail_count = consecutive_fail_count + 1,
cooldown_level = cooldown_level + 1,
cooldown_until = CASE
WHEN ? > 0 THEN datetime('now', '+' || ? || ' seconds')
ELSE cooldown_until
END,
last_fail_at = datetime('now')
cooldown_until = COALESCE(?, cooldown_until),
last_fail_at = CURRENT_TIMESTAMP
WHERE id = ?",
)
.bind(cooldown_seconds)
.bind(cooldown_seconds)
.bind(cooldown_until)
.bind(channel_id)
.execute(&*self.db)
.await?;
@@ -408,7 +463,7 @@ impl TokenRouter {
async fn cooldown_seconds_for_failure(&self, info: &FailureInfo) -> AppResult<i64> {
let raw_max: Option<String> =
sqlx::query_scalar("SELECT value FROM settings WHERE key = ?")
crate::db::sql::query_scalar("SELECT value FROM settings WHERE key = ?")
.bind(FAILURE_COOLDOWN_MAX_SETTING_KEY)
.fetch_optional(&*self.db)
.await?;
@@ -424,7 +479,7 @@ impl TokenRouter {
channel_id: i64,
metrics: &SuccessMetrics,
) -> AppResult<()> {
sqlx::query(
crate::db::sql::query(
"UPDATE oauth_route_unit_members
SET success_count = COALESCE(success_count, 0) + 1,
total_latency_ms = COALESCE(total_latency_ms, 0) + ?,
@@ -432,9 +487,9 @@ impl TokenRouter {
consecutive_fail_count = 0,
cooldown_level = 0,
cooldown_until = NULL,
last_used_at = datetime('now'),
last_selected_at = datetime('now'),
updated_at = datetime('now')
last_used_at = CURRENT_TIMESTAMP,
last_selected_at = CURRENT_TIMESTAMP,
updated_at = CURRENT_TIMESTAMP
WHERE account_id = (
SELECT account_id FROM route_channels WHERE id = ?
)
@@ -457,17 +512,17 @@ impl TokenRouter {
_info: FailureInfo,
cooldown_seconds: i64,
) -> AppResult<()> {
sqlx::query(
let cooldown_until = (cooldown_seconds > 0).then(|| {
(chrono::Utc::now() + chrono::Duration::seconds(cooldown_seconds)).to_rfc3339()
});
crate::db::sql::query(
"UPDATE oauth_route_unit_members
SET fail_count = COALESCE(fail_count, 0) + 1,
consecutive_fail_count = consecutive_fail_count + 1,
cooldown_level = cooldown_level + 1,
cooldown_until = CASE
WHEN ? > 0 THEN datetime('now', '+' || ? || ' seconds')
ELSE cooldown_until
END,
last_fail_at = datetime('now'),
updated_at = datetime('now')
cooldown_until = COALESCE(?, cooldown_until),
last_fail_at = CURRENT_TIMESTAMP,
updated_at = CURRENT_TIMESTAMP
WHERE account_id = (
SELECT account_id FROM route_channels WHERE id = ?
)
@@ -475,8 +530,7 @@ impl TokenRouter {
SELECT oauth_route_unit_id FROM route_channels WHERE id = ?
)",
)
.bind(cooldown_seconds)
.bind(cooldown_seconds)
.bind(cooldown_until)
.bind(channel_id)
.bind(channel_id)
.execute(&*self.db)
@@ -613,20 +667,56 @@ fn wildcard_matches(model: &str, pattern: &str) -> bool {
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
use tokio::sync::Mutex;
static TEST_CACHE_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
async fn test_pool() -> SqlitePool {
invalidate_route_cache();
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
#[test]
fn account_proxy_overrides_site_and_system_proxy() {
let resolved = resolve_proxy_url_from_parts(
Some(r#"{"credentialMode":"apikey","proxyUrl":"socks5://10.0.0.3:1080"}"#),
Some("http://10.0.0.2:8080"),
true,
Some("http://10.0.0.1:7890"),
);
assert_eq!(resolved.as_deref(), Some("socks5://10.0.0.3:1080"));
}
sqlx::query(
#[test]
fn site_custom_proxy_overrides_system_proxy() {
let resolved = resolve_proxy_url_from_parts(
Some(r#"{"credentialMode":"apikey"}"#),
Some("http://10.0.0.2:8080"),
true,
Some("http://10.0.0.1:7890"),
);
assert_eq!(resolved.as_deref(), Some("http://10.0.0.2:8080"));
}
#[test]
fn system_proxy_requires_site_opt_in() {
assert_eq!(
resolve_proxy_url_from_parts(None, None, true, Some("http://10.0.0.1:7890")).as_deref(),
Some("http://10.0.0.1:7890")
);
assert!(
resolve_proxy_url_from_parts(None, None, false, Some("http://10.0.0.1:7890")).is_none()
);
}
async fn test_pool() -> AnyPool {
invalidate_route_cache();
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
crate::db::sql::query(
"CREATE TABLE settings (
key TEXT PRIMARY KEY NOT NULL,
value TEXT
@@ -635,7 +725,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE sites (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
@@ -643,13 +733,15 @@ mod tests {
platform TEXT NOT NULL DEFAULT 'openai',
status TEXT NOT NULL DEFAULT 'active',
api_key TEXT,
proxy_url TEXT,
use_system_proxy INTEGER DEFAULT 0,
custom_headers TEXT
)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE token_routes (
id INTEGER PRIMARY KEY,
model_pattern TEXT NOT NULL,
@@ -661,7 +753,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE accounts (
id INTEGER PRIMARY KEY,
site_id INTEGER NOT NULL,
@@ -677,7 +769,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE account_tokens (
id INTEGER PRIMARY KEY,
account_id INTEGER NOT NULL,
@@ -689,7 +781,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE route_channels (
id INTEGER PRIMARY KEY,
route_id INTEGER NOT NULL,
@@ -715,7 +807,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE oauth_route_units (
id INTEGER PRIMARY KEY,
site_id INTEGER NOT NULL,
@@ -728,7 +820,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE oauth_route_unit_members (
id INTEGER PRIMARY KEY,
unit_id INTEGER NOT NULL,
@@ -753,7 +845,7 @@ mod tests {
// select_channel_for_route's COALESCE subquery joins through
// site_api_endpoints. The table must exist or the SELECT errors
// out — we leave it empty and the COALESCE falls back to sites.url.
sqlx::query(
crate::db::sql::query(
"CREATE TABLE site_api_endpoints (
id INTEGER PRIMARY KEY,
site_id INTEGER NOT NULL,
@@ -767,23 +859,23 @@ mod tests {
.await
.unwrap();
sqlx::query("INSERT INTO sites (id, name, url, api_key) VALUES (1, 'OpenAI', 'https://api.example/v1', 'site-key')")
crate::db::sql::query("INSERT INTO sites (id, name, url, api_key) VALUES (1, 'OpenAI', 'https://api.example/v1', 'site-key')")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO accounts (id, site_id, access_token, api_token) VALUES (1, 1, 'access-key', NULL)")
crate::db::sql::query("INSERT INTO accounts (id, site_id, access_token, api_token) VALUES (1, 1, 'access-key', NULL)")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO accounts (id, site_id, access_token, api_token) VALUES (2, 1, 'access-key-2', NULL)")
crate::db::sql::query("INSERT INTO accounts (id, site_id, access_token, api_token) VALUES (2, 1, 'access-key-2', NULL)")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO account_tokens (id, account_id, name, token) VALUES (1, 1, 'primary', 'token-key')")
crate::db::sql::query("INSERT INTO account_tokens (id, account_id, name, token) VALUES (1, 1, 'primary', 'token-key')")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO account_tokens (id, account_id, name, token) VALUES (2, 2, 'secondary', 'token-key-2')")
crate::db::sql::query("INSERT INTO account_tokens (id, account_id, name, token) VALUES (2, 2, 'secondary', 'token-key-2')")
.execute(&pool)
.await
.unwrap();
@@ -795,11 +887,11 @@ mod tests {
async fn resolves_exact_model_route() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, model_mapping, route_mode) VALUES (1, 'gpt-5.5', 'upstream-gpt-5.5', 'exact')")
crate::db::sql::query("INSERT INTO token_routes (id, model_pattern, model_mapping, route_mode) VALUES (1, 'gpt-5.5', 'upstream-gpt-5.5', 'exact')")
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, 1)",
)
.execute(&pool)
@@ -819,22 +911,22 @@ mod tests {
let first_pool = test_pool().await;
let second_pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-alpha', 'exact')")
crate::db::sql::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-alpha', 'exact')")
.execute(&first_pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, 1)",
)
.execute(&first_pool)
.await
.unwrap();
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-beta', 'exact')")
crate::db::sql::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-beta', 'exact')")
.execute(&second_pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, 1)",
)
.execute(&second_pool)
@@ -865,15 +957,15 @@ mod tests {
async fn resolves_wildcard_route_and_skips_excluded_channel() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
crate::db::sql::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO route_channels (id, route_id, account_id, token_id, priority) VALUES (1, 1, 1, 1, 10)")
crate::db::sql::query("INSERT INTO route_channels (id, route_id, account_id, token_id, priority) VALUES (1, 1, 1, 1, 10)")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO route_channels (id, route_id, account_id, token_id, priority) VALUES (2, 1, 1, 1, 1)")
crate::db::sql::query("INSERT INTO route_channels (id, route_id, account_id, token_id, priority) VALUES (2, 1, 1, 1, 1)")
.execute(&pool)
.await
.unwrap();
@@ -888,13 +980,13 @@ mod tests {
async fn oauth_channel_uses_access_token_instead_of_refresh_token() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query(
crate::db::sql::query(
"INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"UPDATE accounts
SET access_token = 'oauth-access',
api_token = 'oauth-refresh',
@@ -912,7 +1004,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, NULL)")
crate::db::sql::query("INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, NULL)")
.execute(&pool)
.await
.unwrap();
@@ -940,7 +1032,7 @@ mod tests {
})))
.mount(&mock)
.await;
sqlx::query(
crate::db::sql::query(
"UPDATE sites
SET api_key = NULL,
platform = 'codex',
@@ -956,7 +1048,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"UPDATE accounts
SET access_token = 'oauth-access-old',
api_token = 'oauth-refresh-old',
@@ -974,13 +1066,13 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, NULL)")
crate::db::sql::query("INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, NULL)")
.execute(&pool)
.await
.unwrap();
@@ -989,7 +1081,7 @@ mod tests {
let decision = router.resolve_decision("gpt-4o", &[]).await.unwrap();
assert_eq!(decision.selected.api_key, "oauth-access-new");
let row: (String, Option<String>, Option<String>) = sqlx::query_as(
let row: (String, Option<String>, Option<String>) = crate::db::sql::query_as(
"SELECT access_token, api_token, extra_config FROM accounts WHERE id = 1",
)
.fetch_one(&pool)
@@ -1010,18 +1102,18 @@ mod tests {
async fn route_unit_selects_one_member_and_failover_stays_inside_pool() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
crate::db::sql::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO oauth_route_units (id, site_id, provider, name, strategy)
VALUES (9, 1, 'codex', 'Codex Pool', 'round_robin')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO oauth_route_unit_members (id, unit_id, account_id, sort_order) VALUES
(91, 9, 1, 0),
(92, 9, 2, 1)",
@@ -1029,7 +1121,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO route_channels (id, route_id, account_id, token_id, oauth_route_unit_id, priority) VALUES
(1, 1, 1, 1, 9, 10),
(2, 1, 2, 2, 9, 10),
@@ -1061,7 +1153,7 @@ mod tests {
)
.await
.unwrap();
let member_success: (i64, i64, f64) = sqlx::query_as(
let member_success: (i64, i64, f64) = crate::db::sql::query_as(
"SELECT success_count, total_latency_ms, total_cost
FROM oauth_route_unit_members
WHERE unit_id = 9 AND account_id = 2",
@@ -1082,7 +1174,7 @@ mod tests {
)
.await
.unwrap();
let member_failure: (i64, i64) = sqlx::query_as(
let member_failure: (i64, i64) = crate::db::sql::query_as(
"SELECT fail_count, consecutive_fail_count
FROM oauth_route_unit_members
WHERE unit_id = 9 AND account_id = 1",
@@ -1097,18 +1189,18 @@ mod tests {
async fn route_unit_round_robin_prefers_least_recent_member() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
crate::db::sql::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO oauth_route_units (id, site_id, provider, name, strategy)
VALUES (9, 1, 'codex', 'Codex Pool', 'round_robin')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO oauth_route_unit_members (id, unit_id, account_id, sort_order) VALUES
(91, 9, 1, 0),
(92, 9, 2, 1)",
@@ -1116,7 +1208,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO route_channels (
id, route_id, account_id, token_id, oauth_route_unit_id, priority, last_selected_at
) VALUES
@@ -1150,18 +1242,18 @@ mod tests {
async fn route_unit_stick_strategy_keeps_primary_until_it_cools_down() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
crate::db::sql::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO oauth_route_units (id, site_id, provider, name, strategy)
VALUES (9, 1, 'codex', 'Codex Pool', 'stick_until_unavailable')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO oauth_route_unit_members (id, unit_id, account_id, sort_order) VALUES
(91, 9, 1, 0),
(92, 9, 2, 1)",
@@ -1169,7 +1261,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO route_channels (
id, route_id, account_id, token_id, oauth_route_unit_id, priority, weight, last_selected_at
) VALUES
@@ -1207,19 +1299,19 @@ mod tests {
async fn record_failure_clamps_cooldown_to_runtime_setting() {
let _guard = TEST_CACHE_LOCK.lock().await;
let pool = test_pool().await;
sqlx::query("INSERT INTO settings (key, value) VALUES (?, ?)")
crate::db::sql::query("INSERT INTO settings (key, value) VALUES (?, ?)")
.bind(FAILURE_COOLDOWN_MAX_SETTING_KEY)
.bind("30")
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-4o', 'exact')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, 1)",
)
.execute(&pool)
@@ -1239,7 +1331,7 @@ mod tests {
.await
.unwrap();
let seconds_until: i64 = sqlx::query_scalar(
let seconds_until: i64 = crate::db::sql::query_scalar(
"SELECT CAST(strftime('%s', cooldown_until) AS INTEGER) - CAST(strftime('%s', 'now') AS INTEGER)
FROM route_channels
WHERE id = 1",
+38 -34
View File
@@ -111,14 +111,14 @@ pub async fn list_account_tokens(
);
let tokens: Vec<AccountToken> = if let Some(account_id) = query.account_id {
sqlx::query_as::<_, AccountToken>(
crate::db::sql::query_as::<_, AccountToken>(
"SELECT * FROM account_tokens WHERE account_id = ? ORDER BY token_group, id",
)
.bind(account_id)
.fetch_all(&*state.db)
.await?
} else {
sqlx::query_as::<_, AccountToken>(
crate::db::sql::query_as::<_, AccountToken>(
"SELECT * FROM account_tokens ORDER BY account_id, token_group, id",
)
.fetch_all(&*state.db)
@@ -168,10 +168,10 @@ pub async fn create_account_token(
let now = Utc::now().to_rfc3339();
let result = sqlx::query(
let token_id: i64 = crate::db::sql::insert_id(
"INSERT INTO account_tokens \
(account_id, name, token, source, token_group, value_status, enabled, is_default, created_at, updated_at) \
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) RETURNING id",
)
.bind(payload.account_id)
.bind(&name)
@@ -183,12 +183,11 @@ pub async fn create_account_token(
.bind(payload.is_default.unwrap_or(false))
.bind(&now)
.bind(&now)
.execute(&*state.db)
.fetch_one(&*state.db)
.await?;
let token_id = result.last_insert_rowid();
let account_token: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
crate::db::sql::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(token_id)
.fetch_one(&*state.db)
.await?;
@@ -204,7 +203,7 @@ pub async fn get_account_token(
info!("Getting account token: {}", id);
let account_token: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
crate::db::sql::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
@@ -222,7 +221,7 @@ pub async fn update_account_token(
info!("Updating account token: {}", id);
let existing: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
crate::db::sql::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
@@ -265,7 +264,7 @@ pub async fn update_account_token(
let now = Utc::now().to_rfc3339();
sqlx::query(
crate::db::sql::query(
"UPDATE account_tokens SET \
account_id = ?, name = ?, token = ?, source = ?, token_group = ?, value_status = ?, \
enabled = ?, is_default = ?, updated_at = ? \
@@ -285,7 +284,7 @@ pub async fn update_account_token(
.await?;
let account_token: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
crate::db::sql::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_one(&*state.db)
.await?;
@@ -300,7 +299,7 @@ pub async fn delete_account_token(
) -> AppResult<Json<Value>> {
info!("Deleting account token: {}", id);
let affected = sqlx::query("DELETE FROM account_tokens WHERE id = ?")
let affected = crate::db::sql::query("DELETE FROM account_tokens WHERE id = ?")
.bind(id)
.execute(&*state.db)
.await?
@@ -336,21 +335,25 @@ pub async fn batch_update_account_tokens(
for id in payload.ids.into_iter().filter(|id| *id > 0) {
let result = match action.as_str() {
"enable" => {
sqlx::query("UPDATE account_tokens SET enabled = 1, updated_at = ? WHERE id = ?")
.bind(&now)
.bind(id)
.execute(&*state.db)
.await
crate::db::sql::query(
"UPDATE account_tokens SET enabled = 1, updated_at = ? WHERE id = ?",
)
.bind(&now)
.bind(id)
.execute(&*state.db)
.await
}
"disable" => {
sqlx::query("UPDATE account_tokens SET enabled = 0, updated_at = ? WHERE id = ?")
.bind(&now)
.bind(id)
.execute(&*state.db)
.await
crate::db::sql::query(
"UPDATE account_tokens SET enabled = 0, updated_at = ? WHERE id = ?",
)
.bind(&now)
.bind(id)
.execute(&*state.db)
.await
}
"delete" => {
sqlx::query("DELETE FROM account_tokens WHERE id = ?")
crate::db::sql::query("DELETE FROM account_tokens WHERE id = ?")
.bind(id)
.execute(&*state.db)
.await
@@ -385,7 +388,7 @@ pub async fn list_account_token_groups(
) -> AppResult<Json<Value>> {
ensure_account_exists(&state, account_id).await?;
let rows: Vec<(Option<String>,)> = sqlx::query_as(
let rows: Vec<(Option<String>,)> = crate::db::sql::query_as(
"SELECT DISTINCT token_group FROM account_tokens WHERE account_id = ? ORDER BY token_group",
)
.bind(account_id)
@@ -409,7 +412,7 @@ pub async fn set_default_account_token(
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
let token: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
crate::db::sql::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
@@ -418,7 +421,7 @@ pub async fn set_default_account_token(
let group = normalize_token_group(token.token_group.clone());
let now = Utc::now().to_rfc3339();
sqlx::query(
crate::db::sql::query(
"UPDATE account_tokens SET is_default = 0, updated_at = ? \
WHERE account_id = ? AND COALESCE(NULLIF(TRIM(token_group), ''), 'default') = ?",
)
@@ -428,14 +431,14 @@ pub async fn set_default_account_token(
.execute(&*state.db)
.await?;
sqlx::query("UPDATE account_tokens SET is_default = 1, updated_at = ? WHERE id = ?")
crate::db::sql::query("UPDATE account_tokens SET is_default = 1, updated_at = ? WHERE id = ?")
.bind(&now)
.bind(id)
.execute(&*state.db)
.await?;
let updated: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
crate::db::sql::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_one(&*state.db)
.await?;
@@ -449,7 +452,7 @@ pub async fn get_account_token_value(
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
let token: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
crate::db::sql::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
@@ -492,7 +495,7 @@ pub async fn sync_all_account_tokens(State(state): State<AppState>) -> (StatusCo
info!("Sync all account tokens requested");
let rows: Vec<(i64, String)> =
sqlx::query_as("SELECT id, COALESCE(name, '') FROM accounts ORDER BY id")
crate::db::sql::query_as("SELECT id, COALESCE(name, '') FROM accounts ORDER BY id")
.fetch_all(&*state.db)
.await
.unwrap_or_default();
@@ -524,10 +527,11 @@ pub async fn sync_all_account_tokens(State(state): State<AppState>) -> (StatusCo
}
async fn ensure_account_exists(state: &AppState, account_id: i64) -> AppResult<()> {
let account_exists: Option<(i64,)> = sqlx::query_as("SELECT id FROM accounts WHERE id = ?")
.bind(account_id)
.fetch_optional(&*state.db)
.await?;
let account_exists: Option<(i64,)> =
crate::db::sql::query_as("SELECT id FROM accounts WHERE id = ?")
.bind(account_id)
.fetch_optional(&*state.db)
.await?;
if account_exists.is_none() {
return Err(AppError::BadRequest(format!(
File diff suppressed because it is too large Load Diff
+10 -2
View File
@@ -33,7 +33,12 @@ pub struct TemplatesQuery {
/// embed in the templates.
fn resolve_origin(state: &AppState, headers: &HeaderMap, query: &TemplatesQuery) -> String {
// 1. Explicit override wins. Normalize away a trailing `/v1` or `/`.
if let Some(raw) = query.base_url.as_ref().map(|s| s.trim()).filter(|s| !s.is_empty()) {
if let Some(raw) = query
.base_url
.as_ref()
.map(|s| s.trim())
.filter(|s| !s.is_empty())
{
let trimmed = raw.trim_end_matches('/');
let trimmed = trimmed.strip_suffix("/v1").unwrap_or(trimmed);
return trimmed.trim_end_matches('/').to_string();
@@ -58,7 +63,10 @@ fn resolve_origin(state: &AppState, headers: &HeaderMap, query: &TemplatesQuery)
}
fn header_str<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> {
headers.get(name).and_then(|v| v.to_str().ok()).map(str::trim)
headers
.get(name)
.and_then(|v| v.to_str().ok())
.map(str::trim)
}
/// `GET /api/cli-profiles/templates` — list copy-ready integration snippets.
+45
View File
@@ -0,0 +1,45 @@
use crate::{error::AppResult, services::comfy_billing, AppState};
use axum::{
extract::{Path, State},
routing::{get, put},
Json, Router,
};
use serde::Deserialize;
use serde_json::{json, Value};
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct PriceUpdate {
points_micros: i64,
enabled: bool,
}
pub fn routes() -> Router<AppState> {
Router::new()
.route("/api/admin/billing/comfy-prices", get(list_prices))
.route(
"/api/admin/billing/comfy-prices/{task_type}",
put(update_price),
)
}
async fn list_prices(State(state): State<AppState>) -> AppResult<Json<Value>> {
Ok(Json(
json!({ "data": comfy_billing::list_prices(&state.db).await? }),
))
}
async fn update_price(
State(state): State<AppState>,
Path(task_type): Path<String>,
Json(payload): Json<PriceUpdate>,
) -> AppResult<Json<Value>> {
let value = comfy_billing::set_price(
&state.db,
task_type.trim(),
payload.points_micros,
payload.enabled,
)
.await?;
Ok(Json(json!({ "data": value })))
}
+830 -167
View File
File diff suppressed because it is too large Load Diff
+162 -81
View File
@@ -18,11 +18,20 @@ use uuid::Uuid;
/// The system row is uniquely identified by `name = SYSTEM_KEY_NAME AND
/// user_id IS NULL` — see `services::builtin_keys::ensure_comfyui_agent_key`.
const SYSTEM_KEY_NAME: &str = "comfyui-agent";
const USER_KEY_MANAGEMENT_ERROR: &str =
"user-bound keys must be managed through /api/admin/users/{user_id}/keys";
fn is_system_key(name: &str, user_id: Option<i64>) -> bool {
name == SYSTEM_KEY_NAME && user_id.is_none()
}
fn ensure_admin_key(key: &DownstreamApiKey) -> AppResult<()> {
if key.user_id.is_some() {
return Err(AppError::Forbidden(USER_KEY_MANAGEMENT_ERROR.into()));
}
Ok(())
}
/// Mask a key down to `••••••••<last4>`. Operates on chars so non-ASCII
/// keys (unlikely but possible) do not panic on byte slicing.
fn mask_key(key: &str) -> String {
@@ -63,6 +72,7 @@ pub struct CreateDownstreamKeyRequest {
pub max_requests: Option<i64>,
pub supported_models: Option<Vec<String>>,
pub allowed_route_ids: Option<String>,
#[serde(alias = "userId")]
pub user_id: Option<i64>,
}
@@ -78,6 +88,7 @@ pub struct UpdateDownstreamKeyRequest {
pub max_requests: Option<i64>,
pub supported_models: Option<Vec<String>>,
pub allowed_route_ids: Option<String>,
#[serde(alias = "userId")]
pub user_id: Option<i64>,
}
@@ -100,15 +111,14 @@ pub async fn list_downstream_keys(
) -> AppResult<Json<Value>> {
info!("Listing downstream API keys with filter: {:?}", query);
let _legacy_user_id = query.user_id;
let q = query.q.unwrap_or_default().trim().to_string();
let q_like = format!("%{}%", q);
let user_id = query.user_id;
let keys: Vec<DownstreamApiKey> = sqlx::query_as::<_, DownstreamApiKey>(
let keys: Vec<DownstreamApiKey> = crate::db::sql::query_as::<_, DownstreamApiKey>(
r#"
SELECT * FROM downstream_api_keys
WHERE (? = '' OR name LIKE ? OR key LIKE ? OR tags LIKE ?)
AND (? IS NULL OR user_id = ?)
WHERE user_id IS NULL
AND (? = '' OR name LIKE ? OR key LIKE ? OR tags LIKE ?)
ORDER BY id DESC
"#,
)
@@ -116,8 +126,6 @@ pub async fn list_downstream_keys(
.bind(&q_like)
.bind(&q_like)
.bind(&q_like)
.bind(user_id)
.bind(user_id)
.fetch_all(&*state.db)
.await?;
@@ -137,13 +145,12 @@ pub async fn create_downstream_key(
if payload.name.trim().is_empty() {
return Err(AppError::BadRequest("name is required".into()));
}
if payload.user_id.is_some() {
return Err(AppError::BadRequest(USER_KEY_MANAGEMENT_ERROR.into()));
}
// Auto-generate key
let key = if payload.user_id.is_some() {
format!("sk-user-{}", Uuid::new_v4().to_string().replace('-', ""))
} else {
format!("sk-{}", Uuid::new_v4().to_string().replace('-', ""))
};
let key = format!("sk-{}", Uuid::new_v4().to_string().replace('-', ""));
// Serialize supported_models to JSON string if provided
let supported_models: Option<String> = if let Some(ref models) = payload.supported_models {
@@ -158,8 +165,8 @@ pub async fn create_downstream_key(
let enabled = payload.enabled.unwrap_or(true);
let result = sqlx::query(
"INSERT INTO downstream_api_keys (name, key, description, group_name, tags, enabled, expires_at, max_cost, max_requests, supported_models, allowed_route_ids, user_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
let key_id: i64 = crate::db::sql::insert_id(
"INSERT INTO downstream_api_keys (name, key, description, group_name, tags, enabled, expires_at, max_cost, max_requests, supported_models, allowed_route_ids, user_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) RETURNING id"
)
.bind(&payload.name)
.bind(&key)
@@ -172,16 +179,16 @@ pub async fn create_downstream_key(
.bind(payload.max_requests)
.bind(&supported_models)
.bind(&payload.allowed_route_ids)
.bind(payload.user_id)
.execute(&*state.db)
.bind(None::<i64>)
.fetch_one(&*state.db)
.await?;
let key_id = result.last_insert_rowid();
let api_key: DownstreamApiKey =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
.bind(key_id)
.fetch_one(&*state.db)
.await?;
let api_key: DownstreamApiKey = crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT * FROM downstream_api_keys WHERE id = ?",
)
.bind(key_id)
.fetch_one(&*state.db)
.await?;
Ok((StatusCode::CREATED, Json(json!({ "data": api_key }))))
}
@@ -194,14 +201,18 @@ pub async fn update_downstream_key(
) -> AppResult<Json<Value>> {
info!("Updating downstream API key: {}", id);
let existing: DownstreamApiKey =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| {
AppError::NotFound(format!("Downstream key with id {} not found", id))
})?;
let existing: DownstreamApiKey = crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT * FROM downstream_api_keys WHERE id = ?",
)
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Downstream key with id {} not found", id)))?;
ensure_admin_key(&existing)?;
if payload.user_id.is_some() {
return Err(AppError::BadRequest(USER_KEY_MANAGEMENT_ERROR.into()));
}
// Guard: refuse any request that tries to disable the built-in
// `comfyui-agent` system key (`name='comfyui-agent' AND user_id IS NULL`).
@@ -221,8 +232,8 @@ pub async fn update_downstream_key(
existing.supported_models.clone()
};
sqlx::query(
"UPDATE downstream_api_keys SET name = ?, description = ?, group_name = ?, tags = ?, enabled = ?, expires_at = ?, max_cost = ?, max_requests = ?, supported_models = ?, allowed_route_ids = ?, user_id = ?, updated_at = datetime('now') WHERE id = ?"
crate::db::sql::query(
"UPDATE downstream_api_keys SET name = ?, description = ?, group_name = ?, tags = ?, enabled = ?, expires_at = ?, max_cost = ?, max_requests = ?, supported_models = ?, allowed_route_ids = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ? AND user_id IS NULL"
)
.bind(payload.name.unwrap_or(existing.name))
.bind(payload.description.or(existing.description))
@@ -234,16 +245,16 @@ pub async fn update_downstream_key(
.bind(payload.max_requests.or(existing.max_requests))
.bind(&supported_models)
.bind(payload.allowed_route_ids.or(existing.allowed_route_ids))
.bind(payload.user_id.or(existing.user_id))
.bind(id)
.execute(&*state.db)
.await?;
let api_key: DownstreamApiKey =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
.bind(id)
.fetch_one(&*state.db)
.await?;
let api_key: DownstreamApiKey = crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT * FROM downstream_api_keys WHERE id = ?",
)
.bind(id)
.fetch_one(&*state.db)
.await?;
Ok(Json(json!({ "data": api_key })))
}
@@ -258,24 +269,26 @@ pub async fn delete_downstream_key(
// Pre-fetch so we can refuse the built-in `comfyui-agent` row before
// running the DELETE. Returning 403 instead of silently filtering keeps
// the failure observable to operators who attempt the action.
let existing: DownstreamApiKey =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| {
AppError::NotFound(format!("Downstream key with id {} not found", id))
})?;
let existing: DownstreamApiKey = crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT * FROM downstream_api_keys WHERE id = ?",
)
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Downstream key with id {} not found", id)))?;
ensure_admin_key(&existing)?;
if is_system_key(&existing.name, existing.user_id) {
return Err(AppError::Forbidden("system key cannot be deleted".into()));
}
let affected = sqlx::query("DELETE FROM downstream_api_keys WHERE id = ?")
.bind(id)
.execute(&*state.db)
.await?
.rows_affected();
let affected =
crate::db::sql::query("DELETE FROM downstream_api_keys WHERE id = ? AND user_id IS NULL")
.bind(id)
.execute(&*state.db)
.await?
.rows_affected();
if affected == 0 {
return Err(AppError::NotFound(format!(
@@ -294,32 +307,32 @@ pub async fn reset_downstream_key(
) -> AppResult<Json<Value>> {
info!("Resetting downstream API key: {}", id);
let existing: DownstreamApiKey =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Downstream key {} not found", id)))?;
let existing: DownstreamApiKey = crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT * FROM downstream_api_keys WHERE id = ?",
)
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Downstream key {} not found", id)))?;
let new_key = if existing.user_id.is_some() {
format!("sk-user-{}", Uuid::new_v4().to_string().replace('-', ""))
} else {
format!("sk-{}", Uuid::new_v4().to_string().replace('-', ""))
};
ensure_admin_key(&existing)?;
sqlx::query(
"UPDATE downstream_api_keys SET key = ?, updated_at = datetime('now') WHERE id = ?",
let new_key = format!("sk-{}", Uuid::new_v4().to_string().replace('-', ""));
crate::db::sql::query(
"UPDATE downstream_api_keys SET key = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ? AND user_id IS NULL",
)
.bind(&new_key)
.bind(id)
.execute(&*state.db)
.await?;
let api_key: DownstreamApiKey =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
.bind(id)
.fetch_one(&*state.db)
.await?;
let api_key: DownstreamApiKey = crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT * FROM downstream_api_keys WHERE id = ?",
)
.bind(id)
.fetch_one(&*state.db)
.await?;
Ok(Json(json!({ "data": api_key })))
}
@@ -329,16 +342,19 @@ mod tests {
use super::*;
use crate::config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig};
use crate::services::builtin_keys::BuiltinKeys;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
async fn test_state() -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE downstream_api_keys (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
@@ -359,8 +375,8 @@ mod tests {
excluded_credential_refs TEXT,
last_used_at TEXT,
user_id INTEGER,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now'))
created_at TEXT DEFAULT (CURRENT_TIMESTAMP),
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP)
)",
)
.execute(&pool)
@@ -368,19 +384,26 @@ mod tests {
.unwrap();
// Row 1: built-in system key (name='comfyui-agent', user_id IS NULL)
sqlx::query(
crate::db::sql::query(
"INSERT INTO downstream_api_keys (id, name, key, enabled, user_id) VALUES (1, 'comfyui-agent', 'sk-comfyui-abcdef0123456789', 1, NULL)"
)
.execute(&pool)
.await
.unwrap();
// Row 2: regular admin-scoped key
sqlx::query(
crate::db::sql::query(
"INSERT INTO downstream_api_keys (id, name, key, enabled, user_id) VALUES (2, 'team-prod', 'sk-prod-xyz123', 1, NULL)"
)
.execute(&pool)
.await
.unwrap();
// Row 3: user-bound key, managed only through the user admin routes.
crate::db::sql::query(
"INSERT INTO downstream_api_keys (id, name, key, enabled, user_id) VALUES (3, 'user-prod', 'sk-user-xyz123', 1, 42)"
)
.execute(&pool)
.await
.unwrap();
AppState::new(
AppConfig {
@@ -394,6 +417,8 @@ mod tests {
url: "sqlite::memory:".into(),
max_connections: 1,
ssl: false,
ssl_mode: None,
ca_cert_path: None,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
@@ -412,9 +437,11 @@ mod tests {
template_proxy_url: None,
},
comfyui: Default::default(),
agent: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
jimeng: Default::default(),
seedance: Default::default(),
},
pool,
BuiltinKeys {
@@ -440,6 +467,7 @@ mod tests {
.expect("list should succeed");
let items = resp.0.get("items").and_then(Value::as_array).unwrap();
assert_eq!(items.len(), 2);
assert!(items.iter().all(|item| item["id"] != 3));
let system = items
.iter()
@@ -507,7 +535,7 @@ mod tests {
// Confirm the row was not touched.
let still_enabled: Option<bool> =
sqlx::query_scalar("SELECT enabled FROM downstream_api_keys WHERE id = 1")
crate::db::sql::query_scalar("SELECT enabled FROM downstream_api_keys WHERE id = 1")
.fetch_one(&*state.db)
.await
.unwrap();
@@ -537,4 +565,57 @@ mod tests {
.expect("non-disable update should pass");
assert_eq!(resp.0["data"]["description"], "internal: do not delete");
}
#[tokio::test]
async fn generic_create_rejects_user_bound_keys() {
let state = test_state().await;
let payload = CreateDownstreamKeyRequest {
name: "user key".into(),
description: None,
group_name: None,
tags: None,
enabled: Some(true),
expires_at: None,
max_cost: None,
max_requests: None,
supported_models: None,
allowed_route_ids: None,
user_id: Some(42),
};
let result = create_downstream_key(State(state), Json(payload)).await;
assert!(
matches!(result, Err(AppError::BadRequest(message)) if message == USER_KEY_MANAGEMENT_ERROR)
);
}
#[tokio::test]
async fn generic_mutations_reject_user_bound_keys() {
let state = test_state().await;
let payload = UpdateDownstreamKeyRequest {
name: Some("renamed".into()),
description: None,
group_name: None,
tags: None,
enabled: None,
expires_at: None,
max_cost: None,
max_requests: None,
supported_models: None,
allowed_route_ids: None,
user_id: None,
};
let update = update_downstream_key(State(state.clone()), Path(3), Json(payload)).await;
assert!(
matches!(update, Err(AppError::Forbidden(message)) if message == USER_KEY_MANAGEMENT_ERROR)
);
let delete = delete_downstream_key(State(state.clone()), Path(3)).await;
assert!(
matches!(delete, Err(AppError::Forbidden(message)) if message == USER_KEY_MANAGEMENT_ERROR)
);
let reset = reset_downstream_key(State(state), Path(3)).await;
assert!(
matches!(reset, Err(AppError::Forbidden(message)) if message == USER_KEY_MANAGEMENT_ERROR)
);
}
}
+106 -24
View File
@@ -18,6 +18,7 @@ use tracing::info;
mod account_tokens;
pub(crate) mod accounts;
mod cli_templates;
mod comfy_billing;
mod compat;
mod downstream_keys;
mod model_candidates;
@@ -25,6 +26,7 @@ pub(crate) mod models_marketplace;
mod pixelle_connection;
mod sites;
mod tokens;
mod user_oauth;
mod users;
/// Middleware to require admin authentication via the configured admin API token
@@ -195,9 +197,10 @@ pub fn routes(state: AppState) -> Router<AppState> {
)
.route("/api/admin/users/{id}/usage", get(users::get_user_usage))
.route(
"/api/admin/user-balances",
get(users::list_user_balances),
"/api/admin/users/{id}/balance",
get(users::get_user_balance),
)
.route("/api/admin/user-balances", get(users::list_user_balances))
.route(
"/api/admin/pixelle/connection",
get(pixelle_connection::get_connection).put(pixelle_connection::put_connection),
@@ -206,6 +209,10 @@ pub fn routes(state: AppState) -> Router<AppState> {
"/api/admin/users/{id}/balance/adjust",
post(users::adjust_user_balance),
)
.route(
"/api/admin/users/{id}/billing-hold/resolve",
post(users::resolve_user_billing_hold),
)
.route(
"/api/admin/users/{id}/credits/adjust",
post(users::adjust_user_credits),
@@ -234,32 +241,36 @@ pub fn routes(state: AppState) -> Router<AppState> {
.merge(models_marketplace_routes)
.merge(cli_templates_routes)
.merge(account_tokens::routes())
.merge(user_oauth::routes())
.merge(compat::routes())
.merge(comfy_billing::routes())
.merge(super::commerce::admin_routes())
.merge(super::payment_providers::admin_routes())
.merge(super::payment_refunds::admin_routes())
.route_layer(middleware::from_fn_with_state(state, require_admin_auth))
}
// ── Settings handlers ────────────────────────────────────────────────
/// Settings rows whose keys begin with this prefix are managed by
/// `services::builtin_keys` and contain server-issued secrets (e.g. the
/// `comfyui-agent` downstream key). They must never appear in admin
/// list/get/edit responses.
/// Settings managed by dedicated secret-aware paths must never appear in the
/// generic admin list/get/edit responses.
const HIDDEN_SETTINGS_PREFIX: &str = "builtin_key.";
fn is_hidden_setting_key(key: &str) -> bool {
key.starts_with(HIDDEN_SETTINGS_PREFIX)
|| key.starts_with("user_oauth.")
|| key == crate::services::runtime_auth::PROXY_TOKEN_SETTING_KEY
|| key == "notify_smtp_pass"
}
/// GET /api/settings — return all settings as a key-value map.
///
/// Filters out any keys with the `builtin_key.` prefix so the admin UI
/// never sees secrets bootstrapped by `services::builtin_keys`.
/// Filters out bootstrap, runtime-auth, SMTP and dedicated user OAuth settings.
async fn get_settings(State(state): State<AppState>) -> AppResult<Json<Value>> {
info!("Getting all settings");
let rows: Vec<Setting> =
sqlx::query_as::<_, Setting>("SELECT key, value FROM settings ORDER BY key")
crate::db::sql::query_as::<_, Setting>("SELECT key, value FROM settings ORDER BY key")
.fetch_all(&*state.db)
.await?;
@@ -275,9 +286,7 @@ async fn get_settings(State(state): State<AppState>) -> AppResult<Json<Value>> {
/// PUT /api/settings — upsert a setting by key.
/// Expects JSON body: { "key": "...", "value": "..." }
///
/// Refuses any key with the `builtin_key.` prefix: those rows are owned by
/// the bootstrap path and a stray PUT would clobber the secret used by the
/// ComfyUI Agent's outbound calls.
/// Refuses keys owned by dedicated secret-aware update paths.
async fn update_setting(
State(state): State<AppState>,
Json(payload): Json<serde_json::Value>,
@@ -303,7 +312,7 @@ async fn update_setting(
info!("Upserting setting: {}", key);
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)")
crate::db::sql::query("INSERT INTO settings (key, value) VALUES (?, ?) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value")
.bind(key)
.bind(&value)
.execute(&*state.db)
@@ -327,36 +336,57 @@ mod tests {
body::Body,
http::{Request, StatusCode},
};
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
use tower::ServiceExt as _;
async fn settings_state() -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
sqlx::query("CREATE TABLE settings (key TEXT PRIMARY KEY, value TEXT)")
crate::db::sql::query("CREATE TABLE settings (key TEXT PRIMARY KEY, value TEXT)")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO settings (key, value) VALUES ('site.title', 'BoosAPI')")
crate::db::sql::query("INSERT INTO settings (key, value) VALUES ('site.title', 'BoosAPI')")
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO settings (key, value) VALUES ('builtin_key.comfyui_agent.secret', 'sk-comfyui-secret')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO settings (key, value) VALUES ('runtime.auth.proxy_token', 'sk-runtime-secret')",
)
.execute(&pool)
.await
.unwrap();
crate::db::sql::query(
"INSERT INTO settings (key, value) VALUES ('notify_smtp_pass', 'smtp-secret')",
)
.execute(&pool)
.await
.unwrap();
crate::db::sql::query(
"INSERT INTO settings (key, value) VALUES ('user_oauth.google.client_secret', 'oauth-secret')",
)
.execute(&pool)
.await
.unwrap();
crate::db::sql::query(
"INSERT INTO settings (key, value) VALUES ('user_oauth.google.client_id', 'oauth-client-id')",
)
.execute(&pool)
.await
.unwrap();
AppState::new(
AppConfig {
@@ -370,6 +400,8 @@ mod tests {
url: "sqlite::memory:".into(),
max_connections: 1,
ssl: false,
ssl_mode: None,
ca_cert_path: None,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
@@ -388,9 +420,11 @@ mod tests {
template_proxy_url: None,
},
comfyui: Default::default(),
agent: Default::default(),
pixelle_video: Default::default(),
payment: Default::default(),
jimeng: Default::default(),
seedance: Default::default(),
},
pool,
BuiltinKeys {
@@ -416,6 +450,14 @@ mod tests {
!data.contains_key("runtime.auth.proxy_token"),
"runtime proxy token must not appear in /api/settings response"
);
assert!(
!data.contains_key("notify_smtp_pass"),
"SMTP password must not appear in /api/settings response"
);
assert!(
!data.keys().any(|key| key.starts_with("user_oauth.")),
"user_oauth.* rows must only be exposed through the dedicated masked endpoint"
);
}
/// 5.3 — `PUT /api/settings` rejects writes targeting `builtin_key.*`
@@ -435,7 +477,7 @@ mod tests {
other => panic!("expected Forbidden, got {:?}", other),
}
let stored: Option<String> = sqlx::query_scalar(
let stored: Option<String> = crate::db::sql::query_scalar(
"SELECT value FROM settings WHERE key = 'builtin_key.comfyui_agent.secret'",
)
.fetch_one(&*state.db)
@@ -444,6 +486,46 @@ mod tests {
assert_eq!(stored.as_deref(), Some("sk-comfyui-secret"));
}
#[tokio::test]
async fn update_settings_rejects_smtp_password_writes() {
let state = settings_state().await;
let payload = json!({
"key": "notify_smtp_pass",
"value": "attacker-controlled"
});
let result = update_setting(State(state.clone()), Json(payload)).await;
assert!(matches!(result, Err(AppError::Forbidden(_))));
let stored: Option<String> = crate::db::sql::query_scalar(
"SELECT value FROM settings WHERE key = 'notify_smtp_pass'",
)
.fetch_one(&*state.db)
.await
.unwrap();
assert_eq!(stored.as_deref(), Some("smtp-secret"));
}
#[tokio::test]
async fn update_settings_rejects_user_oauth_writes() {
let state = settings_state().await;
let payload = json!({
"key": "user_oauth.google.client_secret",
"value": "attacker-controlled"
});
let result = update_setting(State(state.clone()), Json(payload)).await;
assert!(matches!(result, Err(AppError::Forbidden(_))));
let stored: Option<String> = crate::db::sql::query_scalar(
"SELECT value FROM settings WHERE key = 'user_oauth.google.client_secret'",
)
.fetch_one(&*state.db)
.await
.unwrap();
assert_eq!(stored.as_deref(), Some("oauth-secret"));
}
/// 5.3 — Regular settings still upsert as before.
#[tokio::test]
async fn update_settings_allows_regular_keys() {
+2 -2
View File
@@ -60,7 +60,7 @@ pub async fn token_candidates(State(state): State<AppState>) -> AppResult<Json<V
// 1. Models that already have at least one usable token. The joins drop:
// inactive accounts/sites, disabled tokens, and tokens whose value
// isn't ready (e.g. mid-rotation).
let covered: Vec<CoveredRow> = sqlx::query_as(
let covered: Vec<CoveredRow> = crate::db::sql::query_as(
"SELECT ma.model_name AS model_name,
a.id AS account_id,
at.id AS token_id,
@@ -84,7 +84,7 @@ pub async fn token_candidates(State(state): State<AppState>) -> AppResult<Json<V
// 2. All available `(account, model)` pairs — used to compute which
// accounts have the model but no usable token (`modelsWithoutToken`).
let account_models: Vec<AccountModelRow> = sqlx::query_as(
let account_models: Vec<AccountModelRow> = crate::db::sql::query_as(
"SELECT ma.model_name AS model_name,
a.id AS account_id,
a.username AS username,
+189 -42
View File
@@ -1,12 +1,11 @@
//! `GET /api/models/marketplace` — model marketplace listing for the admin
//! `/models` page (frontend: `Models.tsx`).
//!
//! Stage 1: aggregates `model_availability` × `accounts` × `account_tokens`
//! × `sites` into the shape `ModelsMarketplaceResponse` expects. Per-model
//! latency / success rate / description / tags / supportedEndpointTypes /
//! pricingSources are returned as null/empty placeholders — the frontend
//! falls back gracefully and stage 2 will hydrate them from `proxy_logs`
//! and pricing tables.
//! Aggregates `model_availability` × `accounts` × `account_tokens` × `sites`
//! into the shape `ModelsMarketplaceResponse` expects. Probe latency comes
//! from `model_availability`; the seven-day success rate comes from
//! `proxy_logs`. Description / tags / supportedEndpointTypes / pricingSources
//! remain placeholders until the pricing metadata stage is implemented.
//!
//! `?refresh=1` and `?includePricing=1` are accepted but currently no-ops:
//! we always return the in-memory aggregation and an empty `meta`. Stage 2
@@ -17,6 +16,7 @@ use axum::{
extract::{Query, State},
Json,
};
use chrono::{Duration, Utc};
use serde::Deserialize;
use serde_json::{json, Value};
use sqlx::FromRow;
@@ -34,6 +34,7 @@ pub struct MarketplaceQuery {
struct AvailRow {
account_id: i64,
model_name: String,
latency_ms: Option<i64>,
}
#[derive(FromRow)]
@@ -52,6 +53,13 @@ struct TokenRow {
is_default: Option<bool>,
}
#[derive(FromRow)]
struct ModelLogStatRow {
model_name: String,
total_calls: i64,
success_calls: i64,
}
pub async fn get_marketplace(
State(state): State<AppState>,
Query(_query): Query<MarketplaceQuery>,
@@ -64,11 +72,11 @@ pub async fn get_marketplace(
/// `AppState` and axum extractors so unit tests can drive it with an
/// in-memory SQLite pool (mirrors the `rebuild_routes_inner` pattern in
/// `services::route_rebuild`).
pub(crate) async fn marketplace_inner(pool: &sqlx::Pool<sqlx::Sqlite>) -> AppResult<Value> {
pub(crate) async fn marketplace_inner(pool: &sqlx::Pool<sqlx::Any>) -> AppResult<Value> {
// 1. Active (account, model) pairs. Filtering on accounts.status =
// 'active' here keeps the rest of the aggregation small.
let avail_rows: Vec<AvailRow> = sqlx::query_as(
"SELECT ma.account_id, ma.model_name
let avail_rows: Vec<AvailRow> = crate::db::sql::query_as(
"SELECT ma.account_id, ma.model_name, ma.latency_ms
FROM model_availability ma
JOIN accounts a ON a.id = ma.account_id
WHERE COALESCE(ma.available, 0) = 1
@@ -84,7 +92,7 @@ pub(crate) async fn marketplace_inner(pool: &sqlx::Pool<sqlx::Sqlite>) -> AppRes
// 2. One-shot lookups — accounts, sites, tokens. Bounded by the active
// account count (typically dozens), so a full scan beats per-model
// in-loop selects.
let accounts: Vec<AccountRow> = sqlx::query_as(
let accounts: Vec<AccountRow> = crate::db::sql::query_as(
"SELECT id, site_id, username, balance FROM accounts WHERE status = 'active'",
)
.fetch_all(pool)
@@ -92,12 +100,12 @@ pub(crate) async fn marketplace_inner(pool: &sqlx::Pool<sqlx::Sqlite>) -> AppRes
let accounts_by_id: HashMap<i64, AccountRow> =
accounts.into_iter().map(|a| (a.id, a)).collect();
let sites: Vec<(i64, String)> = sqlx::query_as("SELECT id, name FROM sites")
let sites: Vec<(i64, String)> = crate::db::sql::query_as("SELECT id, name FROM sites")
.fetch_all(pool)
.await?;
let site_name_by_id: HashMap<i64, String> = sites.into_iter().collect();
let tokens: Vec<TokenRow> = sqlx::query_as(
let tokens: Vec<TokenRow> = crate::db::sql::query_as(
"SELECT id, account_id, name, is_default
FROM account_tokens
WHERE COALESCE(enabled, 1) = 1",
@@ -109,28 +117,56 @@ pub(crate) async fn marketplace_inner(pool: &sqlx::Pool<sqlx::Sqlite>) -> AppRes
tokens_by_account.entry(t.account_id).or_default().push(t);
}
let log_since = (Utc::now() - Duration::days(7))
.naive_utc()
.format("%Y-%m-%d %H:%M:%S")
.to_string();
let log_rows: Vec<ModelLogStatRow> = crate::db::sql::query_as(
"SELECT
COALESCE(NULLIF(model_actual, ''), NULLIF(model_requested, '')) AS model_name,
COUNT(*) AS total_calls,
COALESCE(SUM(CASE WHEN status = 'success' THEN 1 ELSE 0 END), 0) AS success_calls
FROM proxy_logs
WHERE created_at >= ?
AND COALESCE(NULLIF(model_actual, ''), NULLIF(model_requested, '')) IS NOT NULL
GROUP BY model_name",
)
.bind(log_since)
.fetch_all(pool)
.await?;
let log_stats_by_model: HashMap<String, (i64, i64)> = log_rows
.into_iter()
.map(|row| (row.model_name, (row.total_calls, row.success_calls)))
.collect();
// 3. Group by model name. BTreeMap → output sorted by name (matches the
// deterministic ordering the frontend's table view expects).
let mut by_model: BTreeMap<String, Vec<i64>> = BTreeMap::new();
let mut by_model: BTreeMap<String, Vec<(i64, Option<i64>)>> = BTreeMap::new();
for row in avail_rows {
by_model
.entry(row.model_name)
.or_default()
.push(row.account_id);
.push((row.account_id, row.latency_ms));
}
// 4. Render. Stage 1 returns null/empty for the metadata fields; the
// frontend has fallback handling for each one.
// 4. Render. Metadata fields remain empty until pricing hydration lands;
// runtime metrics are populated whenever probes/logs exist.
let mut models: Vec<Value> = Vec::with_capacity(by_model.len());
for (model_name, account_ids) in by_model {
let account_count = account_ids.len();
for (model_name, account_metrics) in by_model {
let account_count = account_metrics.len();
let mut token_count: usize = 0;
let mut accounts_json: Vec<Value> = Vec::with_capacity(account_count);
for aid in &account_ids {
let mut latency_total = 0_f64;
let mut latency_count = 0_u64;
for (aid, latency_ms) in &account_metrics {
let acc = match accounts_by_id.get(aid) {
Some(a) => a,
None => continue,
};
if let Some(latency) = latency_ms.filter(|value| *value >= 0) {
latency_total += latency as f64;
latency_count += 1;
}
let site_name = site_name_by_id
.get(&acc.site_id)
.cloned()
@@ -152,18 +188,28 @@ pub(crate) async fn marketplace_inner(pool: &sqlx::Pool<sqlx::Sqlite>) -> AppRes
"id": acc.id,
"site": site_name,
"username": acc.username,
"latency": Value::Null,
"latency": latency_ms.filter(|value| *value >= 0),
"balance": acc.balance.unwrap_or(0.0),
"tokens": token_json,
}));
}
let avg_latency = (latency_count > 0)
.then(|| (latency_total / latency_count as f64).round() as i64);
let success_rate = log_stats_by_model
.get(&model_name)
.and_then(|(total_calls, success_calls)| {
(*total_calls > 0).then(|| {
((*success_calls as f64 / *total_calls as f64) * 1000.0).round() / 10.0
})
});
models.push(json!({
"name": model_name,
"accountCount": account_count,
"tokenCount": token_count,
"avgLatency": Value::Null,
"successRate": Value::Null,
"avgLatency": avg_latency,
"successRate": success_rate,
"description": Value::Null,
"tags": [],
"supportedEndpointTypes": [],
@@ -181,19 +227,22 @@ pub(crate) async fn marketplace_inner(pool: &sqlx::Pool<sqlx::Sqlite>) -> AppRes
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::{Pool, Sqlite};
use sqlx::any::{install_default_drivers, AnyPoolOptions};
use sqlx::{Any, Pool};
/// Minimal schema for the marketplace aggregation. Mirrors
/// `services::route_rebuild::tests::fixture_pool` — we create only the
/// columns `marketplace_inner` actually reads, so adding unrelated
/// migrations later doesn't break these tests.
async fn fixture_pool() -> Pool<Sqlite> {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("open in-memory sqlite");
async fn fixture_pool() -> Pool<Any> {
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("open in-memory sqlite");
for ddl in [
"CREATE TABLE sites (
@@ -211,6 +260,7 @@ mod tests {
account_id INTEGER NOT NULL,
model_name TEXT NOT NULL,
available INTEGER NOT NULL DEFAULT 1,
latency_ms INTEGER,
PRIMARY KEY (account_id, model_name)
)",
"CREATE TABLE account_tokens (
@@ -220,8 +270,15 @@ mod tests {
enabled INTEGER DEFAULT 1,
is_default INTEGER DEFAULT 0
)",
"CREATE TABLE proxy_logs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
model_requested TEXT,
model_actual TEXT,
status TEXT,
created_at TEXT NOT NULL
)",
] {
sqlx::query(ddl)
crate::db::sql::query(ddl)
.execute(&pool)
.await
.unwrap_or_else(|err| panic!("failed to create schema: {err}\nddl: {ddl}"));
@@ -248,11 +305,11 @@ mod tests {
async fn filters_inactive_accounts_and_unavailable_models() {
let pool = fixture_pool().await;
sqlx::query("INSERT INTO sites (id, name) VALUES (1, 'site-a')")
crate::db::sql::query("INSERT INTO sites (id, name) VALUES (1, 'site-a')")
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO accounts (id, site_id, username, balance, status) VALUES
(1, 1, 'alice', 10.0, 'active'),
(2, 1, 'bob', 5.0, 'disabled'),
@@ -261,7 +318,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO model_availability (account_id, model_name, available) VALUES
(1, 'gpt-x', 1),
(2, 'gpt-x', 1),
@@ -299,13 +356,13 @@ mod tests {
async fn aggregates_accounts_and_tokens_per_model() {
let pool = fixture_pool().await;
sqlx::query(
crate::db::sql::query(
"INSERT INTO sites (id, name) VALUES (1, 'openai-mirror'), (2, 'anthropic-mirror')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO accounts (id, site_id, username, balance, status) VALUES
(1, 1, 'alice', 100.5, 'active'),
(2, 2, 'bob', 50.0, 'active')",
@@ -313,7 +370,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO model_availability (account_id, model_name, available) VALUES
(1, 'gpt-x', 1),
(2, 'gpt-x', 1)",
@@ -323,7 +380,7 @@ mod tests {
.unwrap();
// Account 1: two enabled tokens (one default), one disabled token (must
// be excluded). Account 2: one enabled non-default token.
sqlx::query(
crate::db::sql::query(
"INSERT INTO account_tokens (id, account_id, name, enabled, is_default) VALUES
(10, 1, 'primary', 1, 1),
(11, 1, 'backup', 1, 0),
@@ -344,7 +401,7 @@ mod tests {
m["tokenCount"], 3,
"2 from account 1 + 1 from account 2; disabled token excluded"
);
// Stage 1 placeholders.
// No probe/log data was inserted, so runtime metrics remain unknown.
assert!(m["avgLatency"].is_null());
assert!(m["successRate"].is_null());
assert!(m["description"].is_null());
@@ -382,6 +439,96 @@ mod tests {
assert_eq!(bob["tokens"].as_array().unwrap().len(), 1);
}
#[tokio::test]
async fn aggregates_probe_latency_and_recent_success_rate() {
let pool = fixture_pool().await;
let recent = chrono::Utc::now()
.naive_utc()
.format("%Y-%m-%d %H:%M:%S")
.to_string();
let stale = (chrono::Utc::now() - chrono::Duration::days(8))
.naive_utc()
.format("%Y-%m-%d %H:%M:%S")
.to_string();
crate::db::sql::query("INSERT INTO sites (id, name) VALUES (1, 'site-a')")
.execute(&pool)
.await
.unwrap();
crate::db::sql::query(
"INSERT INTO accounts (id, site_id, username, status) VALUES
(1, 1, 'alice', 'active'),
(2, 1, 'bob', 'active')",
)
.execute(&pool)
.await
.unwrap();
crate::db::sql::query(
"INSERT INTO model_availability
(account_id, model_name, available, latency_ms) VALUES
(1, 'gpt-x', 1, 100),
(2, 'gpt-x', 1, 300),
(1, 'no-metrics', 1, NULL)",
)
.execute(&pool)
.await
.unwrap();
for (requested, actual, status, created_at) in [
("gpt-x", None, "success", recent.as_str()),
("alias-x", Some("gpt-x"), "success", recent.as_str()),
("gpt-x", Some(""), "failed", recent.as_str()),
("gpt-x", None, "success", stale.as_str()),
("other-model", None, "success", recent.as_str()),
] {
crate::db::sql::query(
"INSERT INTO proxy_logs
(model_requested, model_actual, status, created_at)
VALUES (?, ?, ?, ?)",
)
.bind(requested)
.bind(actual)
.bind(status)
.bind(created_at)
.execute(&pool)
.await
.unwrap();
}
let value = marketplace_inner(&pool).await.expect("inner ok");
let models = value["models"].as_array().unwrap();
let gpt = models
.iter()
.find(|model| model["name"] == "gpt-x")
.expect("gpt-x present");
assert_eq!(gpt["avgLatency"], 200);
assert_eq!(gpt["successRate"], 66.7);
let accounts = gpt["accounts"].as_array().unwrap();
assert_eq!(
accounts
.iter()
.find(|account| account["id"] == 1)
.unwrap()["latency"],
100
);
assert_eq!(
accounts
.iter()
.find(|account| account["id"] == 2)
.unwrap()["latency"],
300
);
let no_metrics = models
.iter()
.find(|model| model["name"] == "no-metrics")
.expect("no-metrics present");
assert!(no_metrics["avgLatency"].is_null());
assert!(no_metrics["successRate"].is_null());
assert!(no_metrics["accounts"][0]["latency"].is_null());
}
/// Models must come back in deterministic name order (BTreeMap). The
/// frontend's table view assumes a stable sort so toggling refresh
/// doesn't reshuffle rows; lock that contract here.
@@ -389,16 +536,16 @@ mod tests {
async fn models_are_sorted_by_name() {
let pool = fixture_pool().await;
sqlx::query("INSERT INTO sites (id, name) VALUES (1, 's')")
crate::db::sql::query("INSERT INTO sites (id, name) VALUES (1, 's')")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO accounts (id, site_id, status) VALUES (1, 1, 'active')")
crate::db::sql::query("INSERT INTO accounts (id, site_id, status) VALUES (1, 1, 'active')")
.execute(&pool)
.await
.unwrap();
// Insert in scrambled order; expect output sorted lexicographically.
sqlx::query(
crate::db::sql::query(
"INSERT INTO model_availability (account_id, model_name, available) VALUES
(1, 'gpt-z', 1),
(1, 'gpt-a', 1),
+15 -4
View File
@@ -10,7 +10,14 @@ const MASK_SENTINEL: &str = "••••";
fn mask(v: &Option<String>) -> Value {
match v {
Some(s) if s.chars().count() > 4 => {
let last4: String = s.chars().rev().take(4).collect::<Vec<_>>().into_iter().rev().collect();
let last4: String = s
.chars()
.rev()
.take(4)
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect();
json!(format!("{}{}", MASK_SENTINEL, last4))
}
Some(s) if !s.is_empty() => json!(MASK_SENTINEL),
@@ -55,10 +62,14 @@ pub async fn put_connection(
comfyui_api_key: keep_if_not_mask(p.comfyui_api_key),
};
pixelle_connection::save_connection(&state.db, &conn).await?;
let backup = pixelle_connection::sync_to_pixelle(&state.config.pixelle_video.config_path, &conn)?;
let restart = pixelle_connection::restart_pixelle(&state.config.pixelle_video.compose_dir).await;
let backup =
pixelle_connection::sync_to_pixelle(&state.config.pixelle_video.config_path, &conn)?;
let restart =
pixelle_connection::restart_pixelle(&state.config.pixelle_video.compose_dir).await;
match restart {
Ok(()) => Ok(Json(json!({ "success": true, "backup": backup, "restarted": true }))),
Ok(()) => Ok(Json(
json!({ "success": true, "backup": backup, "restarted": true }),
)),
Err(e) => Ok(Json(json!({
"success": true, "backup": backup, "restarted": false,
"warning": format!("配置已写入,但 Pixelle 重启失败,请手动重启: {e}")
+172 -50
View File
@@ -34,12 +34,17 @@ pub struct CreateSiteRequest {
pub url: String,
pub platform: String,
pub status: Option<String>,
#[serde(alias = "apiKey")]
pub api_key: Option<String>,
#[serde(alias = "proxyUrl")]
pub proxy_url: Option<String>,
#[serde(alias = "externalCheckinUrl")]
pub external_checkin_url: Option<String>,
#[serde(alias = "useSystemProxy")]
pub use_system_proxy: Option<bool>,
#[serde(alias = "customHeaders")]
pub custom_headers: Option<String>,
#[serde(alias = "globalWeight")]
pub global_weight: Option<f64>,
#[serde(alias = "postRefreshProbeEnabled")]
pub post_refresh_probe_enabled: Option<bool>,
@@ -63,12 +68,17 @@ pub struct UpdateSiteRequest {
pub url: Option<String>,
pub platform: Option<String>,
pub status: Option<String>,
#[serde(alias = "apiKey")]
pub api_key: Option<String>,
#[serde(alias = "proxyUrl")]
pub proxy_url: Option<String>,
#[serde(alias = "externalCheckinUrl")]
pub external_checkin_url: Option<String>,
#[serde(alias = "useSystemProxy")]
pub use_system_proxy: Option<bool>,
#[serde(alias = "customHeaders")]
pub custom_headers: Option<String>,
#[serde(alias = "globalWeight")]
pub global_weight: Option<f64>,
#[serde(alias = "postRefreshProbeEnabled")]
pub post_refresh_probe_enabled: Option<bool>,
@@ -123,7 +133,7 @@ pub async fn list_sites(State(state): State<AppState>) -> AppResult<Json<Value>>
info!("Listing all sites");
let sites: Vec<Site> =
sqlx::query_as::<_, Site>("SELECT * FROM sites ORDER BY sort_order, name")
crate::db::sql::query_as::<_, Site>("SELECT * FROM sites ORDER BY sort_order, name")
.fetch_all(&*state.db)
.await?;
let mut data = Vec::with_capacity(sites.len());
@@ -159,14 +169,14 @@ pub async fn create_site(
let mut tx = state.db.begin().await?;
let result = sqlx::query(
let site_id: i64 = crate::db::sql::insert_id(
"INSERT INTO sites (
name, url, platform, status, api_key, proxy_url, external_checkin_url,
use_system_proxy, custom_headers, global_weight,
post_refresh_probe_enabled, post_refresh_probe_model,
post_refresh_probe_scope, post_refresh_probe_latency_threshold_ms,
is_pinned, sort_order
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) RETURNING id",
)
.bind(&payload.name)
.bind(&site_url)
@@ -184,14 +194,14 @@ pub async fn create_site(
.bind(payload.post_refresh_probe_latency_threshold_ms)
.bind(payload.is_pinned)
.bind(payload.sort_order)
.execute(&mut *tx)
.fetch_one(&mut *tx)
.await?;
let site_id = result.last_insert_rowid();
replace_site_api_endpoints(&mut tx, site_id, payload.api_endpoints.as_deref()).await?;
tx.commit().await?;
crate::proxy::token_router::router::invalidate_route_cache();
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
let site: Site = crate::db::sql::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
.bind(site_id)
.fetch_one(&*state.db)
.await?;
@@ -209,7 +219,7 @@ pub async fn get_site(
) -> AppResult<Json<Value>> {
info!("Getting site: {}", id);
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
let site: Site = crate::db::sql::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
@@ -226,7 +236,7 @@ pub async fn update_site(
) -> AppResult<Json<Value>> {
info!("Updating site: {}", id);
let existing: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
let existing: Site = crate::db::sql::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
@@ -255,14 +265,14 @@ pub async fn update_site(
let mut tx = state.db.begin().await?;
sqlx::query(
crate::db::sql::query(
"UPDATE sites SET
name = ?, url = ?, platform = ?, status = ?, api_key = ?,
proxy_url = ?, external_checkin_url = ?, use_system_proxy = ?,
custom_headers = ?, global_weight = ?,
post_refresh_probe_enabled = ?, post_refresh_probe_model = ?,
post_refresh_probe_scope = ?, post_refresh_probe_latency_threshold_ms = ?,
is_pinned = ?, sort_order = ?, updated_at = datetime('now')
is_pinned = ?, sort_order = ?, updated_at = CURRENT_TIMESTAMP
WHERE id = ?",
)
.bind(payload.name.unwrap_or(existing.name))
@@ -309,8 +319,11 @@ pub async fn update_site(
replace_site_api_endpoints(&mut tx, id, Some(api_endpoints)).await?;
}
tx.commit().await?;
// The token router caches the site's outbound proxy alongside route data.
// A site-level proxy edit must take effect on the next proxied request.
crate::proxy::token_router::router::invalidate_route_cache();
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
let site: Site = crate::db::sql::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
.bind(id)
.fetch_one(&*state.db)
.await?;
@@ -325,7 +338,7 @@ pub async fn delete_site(
) -> AppResult<Json<Value>> {
info!("Deleting site: {}", id);
let affected = sqlx::query("DELETE FROM sites WHERE id = ?")
let affected = crate::db::sql::query("DELETE FROM sites WHERE id = ?")
.bind(id)
.execute(&*state.db)
.await?
@@ -377,7 +390,7 @@ async fn site_response(state: &AppState, site: Site) -> AppResult<Value> {
}
async fn site_api_endpoints(state: &AppState, site_id: i64) -> AppResult<Vec<Value>> {
let rows: Vec<SiteApiEndpointRow> = sqlx::query_as(
let rows: Vec<SiteApiEndpointRow> = crate::db::sql::query_as(
"SELECT id, url, enabled, sort_order, cooldown_until, last_selected_at, last_failed_at, last_failure_reason
FROM site_api_endpoints
WHERE site_id = ?
@@ -410,12 +423,12 @@ async fn site_api_endpoints(state: &AppState, site_id: i64) -> AppResult<Vec<Val
}
async fn resolve_site_api_base_url(state: &AppState, site: &Site) -> AppResult<String> {
let endpoint: Option<String> = sqlx::query_scalar(
let endpoint: Option<String> = crate::db::sql::query_scalar(
"SELECT url
FROM site_api_endpoints
WHERE site_id = ?
AND COALESCE(enabled, 1) = 1
AND (cooldown_until IS NULL OR cooldown_until <= datetime('now'))
AND (cooldown_until IS NULL OR cooldown_until <= CURRENT_TIMESTAMP)
ORDER BY COALESCE(sort_order, 0), id
LIMIT 1",
)
@@ -499,7 +512,7 @@ pub async fn list_site_disabled_models(
) -> AppResult<Json<Value>> {
// Validate the site exists so we return a clean 404 instead of an empty list
// for unknown ids.
let exists: Option<i64> = sqlx::query_scalar("SELECT id FROM sites WHERE id = ?")
let exists: Option<i64> = crate::db::sql::query_scalar("SELECT id FROM sites WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?;
@@ -507,7 +520,7 @@ pub async fn list_site_disabled_models(
return Err(AppError::NotFound("site not found".into()));
}
let models: Vec<String> = sqlx::query_scalar(
let models: Vec<String> = crate::db::sql::query_scalar(
"SELECT model_name FROM site_disabled_models WHERE site_id = ? ORDER BY model_name",
)
.bind(id)
@@ -523,7 +536,7 @@ pub async fn update_site_disabled_models(
Path(id): Path<i64>,
Json(payload): Json<UpdateDisabledModelsRequest>,
) -> AppResult<Json<Value>> {
let exists: Option<i64> = sqlx::query_scalar("SELECT id FROM sites WHERE id = ?")
let exists: Option<i64> = crate::db::sql::query_scalar("SELECT id FROM sites WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?;
@@ -542,12 +555,12 @@ pub async fn update_site_disabled_models(
normalised.dedup();
let mut tx = state.db.begin().await?;
sqlx::query("DELETE FROM site_disabled_models WHERE site_id = ?")
crate::db::sql::query("DELETE FROM site_disabled_models WHERE site_id = ?")
.bind(id)
.execute(&mut *tx)
.await?;
for model in &normalised {
sqlx::query(
crate::db::sql::query(
"INSERT INTO site_disabled_models (site_id, model_name) VALUES (?, ?)\n ON CONFLICT(site_id, model_name) DO NOTHING",
)
.bind(id)
@@ -605,7 +618,7 @@ pub async fn probe_site_stream(
}
async fn ensure_site_exists(state: &AppState, id: i64) -> AppResult<()> {
let exists: Option<i64> = sqlx::query_scalar("SELECT id FROM sites WHERE id = ?")
let exists: Option<i64> = crate::db::sql::query_scalar("SELECT id FROM sites WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?;
@@ -616,7 +629,7 @@ async fn ensure_site_exists(state: &AppState, id: i64) -> AppResult<()> {
}
async fn site_available_model_names(state: &AppState, site_id: i64) -> AppResult<Vec<String>> {
let rows: Vec<String> = sqlx::query_scalar(
let rows: Vec<String> = crate::db::sql::query_scalar(
"SELECT DISTINCT ma.model_name
FROM model_availability ma
JOIN accounts a ON a.id = ma.account_id
@@ -637,7 +650,7 @@ async fn build_probe_plan(
site_id: i64,
options: ProbeSiteOptions,
) -> AppResult<SiteProbePlan> {
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
let site: Site = crate::db::sql::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
.bind(site_id)
.fetch_optional(&*state.db)
.await?
@@ -645,10 +658,15 @@ async fn build_probe_plan(
// Jimeng has no `/models` endpoint and no per-model chat semantics, so the
// per-model probe is meaningless for it. Short-circuit with a clear message
// (covers both probe-now and probe-stream, which funnel through here).
if crate::routes::admin::accounts::is_jimeng_platform(&site.platform) {
return Err(AppError::BadRequest(
"即梦平台无需逐模型探测,请在连接管理用「验证」做签名连通性检查".into(),
));
if crate::routes::admin::accounts::is_dedicated_generation_platform(&site.platform) {
let label = if crate::routes::admin::accounts::is_seedance_platform(&site.platform) {
"Seedance"
} else {
"即梦"
};
return Err(AppError::BadRequest(format!(
"{label} 平台无需逐模型探测,请在连接管理使用「验证」检查凭据"
)));
}
if site.status != "active" {
return Err(AppError::BadRequest("站点未启用,无法执行探测".into()));
@@ -717,7 +735,7 @@ async fn site_probe_models(
model_name: Option<&str>,
) -> AppResult<Vec<SiteProbeModel>> {
let rows: Vec<(String, i64)> = if let Some(model_name) = model_name {
sqlx::query_as(
crate::db::sql::query_as(
"SELECT ma.model_name, MIN(a.id) AS source_account_id
FROM model_availability ma
JOIN accounts a ON a.id = ma.account_id
@@ -733,7 +751,7 @@ async fn site_probe_models(
.fetch_all(&*state.db)
.await?
} else {
sqlx::query_as(
crate::db::sql::query_as(
"SELECT ma.model_name, MIN(a.id) AS source_account_id
FROM model_availability ma
JOIN accounts a ON a.id = ma.account_id
@@ -932,6 +950,19 @@ async fn probe_site_model(
{
warn!(error = %err, model = %model.name, "failed to save probe latency");
}
if let Err(err) = record_site_probe_health(
state,
&selected,
&model.name,
true,
Some(200),
latency_ms,
None,
)
.await
{
warn!(error = %err, model = %model.name, "failed to save probe health");
}
SiteProbeModelResult {
model_name: model.name.clone(),
status: "unsupported".to_string(),
@@ -974,6 +1005,19 @@ async fn probe_site_model(
{
warn!(error = %err, model = %model.name, "failed to save probe latency");
}
if let Err(err) = record_site_probe_health(
state,
&selected,
&model.name,
true,
Some(200),
latency_ms,
None,
)
.await
{
warn!(error = %err, model = %model.name, "failed to save probe health");
}
SiteProbeModelResult {
model_name: model.name.clone(),
status: "supported".to_string(),
@@ -1014,12 +1058,26 @@ async fn probe_site_model(
{
warn!(error = %update_err, model = %model.name, "failed to save failed probe state");
}
let error_message = probe_error_message(&err);
if let Err(record_err) = record_site_probe_health(
state,
&selected,
&model.name,
false,
diagnostics.http_status,
latency_ms,
Some(&error_message),
)
.await
{
warn!(error = %record_err, model = %model.name, "failed to save failed probe health");
}
SiteProbeModelResult {
model_name: model.name.clone(),
status: "unsupported".to_string(),
latency_ms: Some(latency_ms),
latency_exceeded: false,
reason: Some(probe_error_message(&err)),
reason: Some(error_message),
http_status: diagnostics.http_status,
stage: Some(diagnostics.stage),
source: Some(diagnostics.source),
@@ -1040,7 +1098,7 @@ async fn resolve_probe_channel(
}
}
let row: Option<(i64, String, Option<String>, Option<String>)> = sqlx::query_as(
let row: Option<(i64, String, Option<String>, Option<String>)> = crate::db::sql::query_as(
"SELECT id, access_token, api_token, oauth_provider
FROM accounts
WHERE id = ?
@@ -1056,17 +1114,18 @@ async fn resolve_probe_channel(
let (account_id, access_token, api_token, oauth_provider) = match row {
Some(row) => row,
None => {
let fallback: Option<(i64, String, Option<String>, Option<String>)> = sqlx::query_as(
"SELECT id, access_token, api_token, oauth_provider
let fallback: Option<(i64, String, Option<String>, Option<String>)> =
crate::db::sql::query_as(
"SELECT id, access_token, api_token, oauth_provider
FROM accounts
WHERE site_id = ?
AND status = 'active'
ORDER BY sort_order, id
LIMIT 1",
)
.bind(plan.site.id)
.fetch_optional(&*state.db)
.await?;
)
.bind(plan.site.id)
.fetch_optional(&*state.db)
.await?;
fallback
.ok_or_else(|| AppError::BadRequest("当前站点没有启用中的账号,无法探测".into()))?
}
@@ -1106,6 +1165,7 @@ async fn resolve_probe_channel(
account_id: Some(account_id),
account_unit_cost: None,
base_url: plan.base_url.clone(),
proxy_url: None,
api_key,
model: model.name.clone(),
actual_model: model.name.clone(),
@@ -1137,6 +1197,7 @@ fn selected_channel_to_upstream_route(
) -> crate::proxy::executors::UpstreamRoute {
crate::proxy::executors::UpstreamRoute {
base_url: selected.base_url.clone(),
proxy_url: selected.proxy_url.clone(),
api_key: selected.api_key.clone(),
model: selected.actual_model.clone(),
}
@@ -1156,16 +1217,16 @@ async fn update_model_probe_latency(
let Some(account_id) = account_id else {
return Ok(());
};
sqlx::query(
crate::db::sql::query(
"UPDATE model_availability
SET latency_ms = ?,
available = ?,
checked_at = datetime('now')
checked_at = CURRENT_TIMESTAMP
WHERE account_id = ?
AND model_name = ?",
)
.bind(latency_ms)
.bind(available)
.bind(i64::from(available))
.bind(account_id)
.bind(model_name)
.execute(&*state.db)
@@ -1173,8 +1234,44 @@ async fn update_model_probe_latency(
Ok(())
}
async fn record_site_probe_health(
state: &AppState,
selected: &crate::proxy::conductor::SelectedChannel,
model_name: &str,
success: bool,
http_status: Option<u16>,
latency_ms: i64,
error_message: Option<&str>,
) -> AppResult<()> {
let Some(account_id) = selected.account_id else {
return Ok(());
};
crate::db::sql::query(
"INSERT INTO proxy_logs (
route_id, channel_id, account_id, model_requested, model_actual,
status, http_status, latency_ms, error_message, retry_count,
is_stream, client_family, client_app_id, client_app_name,
client_confidence, billing_details
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 0, 0, 'site-probe',
'boosapi-site-probe', 'BoosAPI 模型探测', 'high', ?)",
)
.bind(selected.route_id)
.bind(selected_channel_id(selected))
.bind(account_id)
.bind(model_name)
.bind(&selected.actual_model)
.bind(if success { "success" } else { "failed" })
.bind(http_status.map(i64::from))
.bind(latency_ms)
.bind(error_message)
.bind(r#"{"source":"site-probe","billable":false}"#)
.execute(&*state.db)
.await?;
Ok(())
}
async fn disable_site_model(state: &AppState, site_id: i64, model_name: &str) -> AppResult<()> {
sqlx::query(
crate::db::sql::query(
"INSERT INTO site_disabled_models (site_id, model_name)
VALUES (?, ?)
ON CONFLICT(site_id, model_name) DO NOTHING",
@@ -1211,6 +1308,8 @@ fn probe_error_message(err: &AppError) -> String {
AppError::UpstreamError(message) => simplify_upstream_error_message(message),
AppError::RateLimited => "触发频率限制".to_string(),
AppError::RateLimitedWithRetry { message, .. } => message.clone(),
AppError::RequestRateLimited { message, .. } => message.clone(),
AppError::BillingRequired { message, .. } => message.clone(),
AppError::Internal(err) => err.to_string(),
AppError::Database(err) => err.to_string(),
}
@@ -1338,7 +1437,7 @@ async fn ensure_site_unique(
except_id: Option<i64>,
) -> AppResult<()> {
let existing: Option<(i64, String)> = if let Some(id) = except_id {
sqlx::query_as(
crate::db::sql::query_as(
"SELECT id, name FROM sites WHERE platform = ? AND url = ? AND id != ? LIMIT 1",
)
.bind(platform)
@@ -1347,11 +1446,13 @@ async fn ensure_site_unique(
.fetch_optional(&*state.db)
.await?
} else {
sqlx::query_as("SELECT id, name FROM sites WHERE platform = ? AND url = ? LIMIT 1")
.bind(platform)
.bind(url)
.fetch_optional(&*state.db)
.await?
crate::db::sql::query_as(
"SELECT id, name FROM sites WHERE platform = ? AND url = ? LIMIT 1",
)
.bind(platform)
.bind(url)
.fetch_optional(&*state.db)
.await?
};
if let Some((id, name)) = existing {
@@ -1456,11 +1557,11 @@ fn detect_site_platform(url: &str) -> (&'static str, Option<&'static str>) {
}
async fn replace_site_api_endpoints(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
tx: &mut sqlx::Transaction<'_, sqlx::Any>,
site_id: i64,
endpoints: Option<&[SiteApiEndpointPayload]>,
) -> AppResult<()> {
sqlx::query("DELETE FROM site_api_endpoints WHERE site_id = ?")
crate::db::sql::query("DELETE FROM site_api_endpoints WHERE site_id = ?")
.bind(site_id)
.execute(&mut **tx)
.await?;
@@ -1474,7 +1575,7 @@ async fn replace_site_api_endpoints(
if url.is_empty() {
continue;
}
sqlx::query(
crate::db::sql::query(
"INSERT INTO site_api_endpoints (site_id, url, enabled, sort_order)
VALUES (?, ?, ?, ?)",
)
@@ -1493,6 +1594,27 @@ async fn replace_site_api_endpoints(
mod tests {
use super::*;
#[test]
fn site_proxy_fields_accept_the_frontend_camel_case_payload() {
let payload: UpdateSiteRequest = serde_json::from_value(json!({
"proxyUrl": "socks5://proxy-user:proxy-pass@127.0.0.1:1080",
"useSystemProxy": true,
"apiKey": "site-key",
"customHeaders": "{\"x-test\":\"1\"}",
"globalWeight": 1.5,
}))
.expect("site proxy payload should deserialize");
assert_eq!(
payload.proxy_url.as_deref(),
Some("socks5://proxy-user:proxy-pass@127.0.0.1:1080"),
);
assert_eq!(payload.use_system_proxy, Some(true));
assert_eq!(payload.api_key.as_deref(), Some("site-key"));
assert_eq!(payload.custom_headers.as_deref(), Some("{\"x-test\":\"1\"}"));
assert_eq!(payload.global_weight, Some(1.5));
}
#[test]
fn probe_options_accept_camel_case_payload() {
let options: ProbeSiteOptions = serde_json::from_value(json!({
+40 -34
View File
@@ -52,16 +52,18 @@ pub async fn list_tokens(
info!("Listing tokens, account_id filter: {:?}", query.account_id);
let tokens: Vec<AccountToken> = if let Some(account_id) = query.account_id {
sqlx::query_as::<_, AccountToken>(
crate::db::sql::query_as::<_, AccountToken>(
"SELECT * FROM account_tokens WHERE account_id = ? ORDER BY id",
)
.bind(account_id)
.fetch_all(&*state.db)
.await?
} else {
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens ORDER BY account_id, id")
.fetch_all(&*state.db)
.await?
crate::db::sql::query_as::<_, AccountToken>(
"SELECT * FROM account_tokens ORDER BY account_id, id",
)
.fetch_all(&*state.db)
.await?
};
Ok(Json(json!({ "data": tokens })))
@@ -83,10 +85,11 @@ pub async fn create_token(
}
// Verify account exists
let account_exists: Option<(i64,)> = sqlx::query_as("SELECT id FROM accounts WHERE id = ?")
.bind(payload.account_id)
.fetch_optional(&*state.db)
.await?;
let account_exists: Option<(i64,)> =
crate::db::sql::query_as("SELECT id FROM accounts WHERE id = ?")
.bind(payload.account_id)
.fetch_optional(&*state.db)
.await?;
if account_exists.is_none() {
return Err(AppError::BadRequest(format!(
@@ -99,8 +102,8 @@ pub async fn create_token(
let enabled = payload.enabled.unwrap_or(true);
let is_default = payload.is_default.unwrap_or(false);
let result = sqlx::query(
"INSERT INTO account_tokens (account_id, name, token, source, token_group, enabled, is_default) VALUES (?, ?, ?, ?, ?, ?, ?)"
let token_id: i64 = crate::db::sql::insert_id(
"INSERT INTO account_tokens (account_id, name, token, source, token_group, enabled, is_default) VALUES (?, ?, ?, ?, ?, ?, ?) RETURNING id"
)
.bind(payload.account_id)
.bind(&payload.name)
@@ -109,12 +112,11 @@ pub async fn create_token(
.bind(&payload.token_group)
.bind(enabled)
.bind(is_default)
.execute(&*state.db)
.fetch_one(&*state.db)
.await?;
let token_id = result.last_insert_rowid();
let token: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
crate::db::sql::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(token_id)
.fetch_one(&*state.db)
.await?;
@@ -130,7 +132,7 @@ pub async fn get_token(
info!("Getting token: {}", id);
let token: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
crate::db::sql::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
@@ -148,14 +150,14 @@ pub async fn update_token(
info!("Updating token: {}", id);
let existing: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
crate::db::sql::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound(format!("Token with id {} not found", id)))?;
sqlx::query(
"UPDATE account_tokens SET name = ?, token = ?, source = ?, token_group = ?, enabled = ?, is_default = ?, updated_at = datetime('now') WHERE id = ?"
crate::db::sql::query(
"UPDATE account_tokens SET name = ?, token = ?, source = ?, token_group = ?, enabled = ?, is_default = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?"
)
.bind(payload.name.unwrap_or(existing.name))
.bind(payload.token.as_ref().unwrap_or(&existing.token))
@@ -168,7 +170,7 @@ pub async fn update_token(
.await?;
let token: AccountToken =
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
crate::db::sql::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
.bind(id)
.fetch_one(&*state.db)
.await?;
@@ -183,7 +185,7 @@ pub async fn delete_token(
) -> AppResult<Json<Value>> {
info!("Deleting token: {}", id);
let affected = sqlx::query("DELETE FROM account_tokens WHERE id = ?")
let affected = crate::db::sql::query("DELETE FROM account_tokens WHERE id = ?")
.bind(id)
.execute(&*state.db)
.await?
@@ -207,15 +209,15 @@ pub async fn refresh_route(
info!("Refreshing route decision");
if let Some(route_id) = payload.route_id {
sqlx::query(
"UPDATE token_routes SET decision_refreshed_at = datetime('now'), updated_at = datetime('now') WHERE id = ?"
crate::db::sql::query(
"UPDATE token_routes SET decision_refreshed_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP WHERE id = ?"
)
.bind(route_id)
.execute(&*state.db)
.await?;
let route: Option<crate::db::models::token_route::TokenRoute> =
sqlx::query_as("SELECT * FROM token_routes WHERE id = ?")
crate::db::sql::query_as("SELECT * FROM token_routes WHERE id = ?")
.bind(route_id)
.fetch_optional(&*state.db)
.await?;
@@ -226,8 +228,8 @@ pub async fn refresh_route(
}
if let Some(ref model_pattern) = payload.model_pattern {
let result = sqlx::query(
"UPDATE token_routes SET decision_refreshed_at = datetime('now'), updated_at = datetime('now') WHERE model_pattern LIKE ?"
let result = crate::db::sql::query(
"UPDATE token_routes SET decision_refreshed_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP WHERE model_pattern LIKE ?"
)
.bind(format!("%{}%", model_pattern))
.execute(&*state.db)
@@ -243,8 +245,8 @@ pub async fn refresh_route(
}
// Refresh all routes
let result = sqlx::query(
"UPDATE token_routes SET decision_refreshed_at = datetime('now'), updated_at = datetime('now')"
let result = crate::db::sql::query(
"UPDATE token_routes SET decision_refreshed_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP"
)
.execute(&*state.db)
.await?;
@@ -267,21 +269,25 @@ pub async fn route_decision(
let routes: Vec<crate::db::models::token_route::TokenRoute> =
if let Some(route_id) = query.route_id {
sqlx::query_as("SELECT * FROM token_routes WHERE id = ? AND enabled = 1")
crate::db::sql::query_as("SELECT * FROM token_routes WHERE id = ? AND enabled = 1")
.bind(route_id)
.fetch_optional(&*state.db)
.await?
.map(|r| vec![r])
.unwrap_or_default()
} else if let Some(ref model_pattern) = query.model_pattern {
sqlx::query_as("SELECT * FROM token_routes WHERE model_pattern LIKE ? AND enabled = 1")
.bind(format!("%{}%", model_pattern))
.fetch_all(&*state.db)
.await?
crate::db::sql::query_as(
"SELECT * FROM token_routes WHERE model_pattern LIKE ? AND enabled = 1",
)
.bind(format!("%{}%", model_pattern))
.fetch_all(&*state.db)
.await?
} else {
sqlx::query_as("SELECT * FROM token_routes WHERE enabled = 1 ORDER BY model_pattern")
.fetch_all(&*state.db)
.await?
crate::db::sql::query_as(
"SELECT * FROM token_routes WHERE enabled = 1 ORDER BY model_pattern",
)
.fetch_all(&*state.db)
.await?
};
Ok(Json(json!({ "data": routes })))
+382
View File
@@ -0,0 +1,382 @@
use std::collections::HashMap;
use axum::{extract::State, routing::get, Json, Router};
use serde::{Deserialize, Serialize};
use sqlx::{Any, Transaction};
use url::{Host, Url};
use crate::{
error::{AppError, AppResult},
services::email::load_smtp_mail_config,
AppState,
};
const PUBLIC_BASE_URL_KEY: &str = "user_oauth.public_base_url";
const REGISTRATION_EMAIL_ENABLED_KEY: &str = "runtime.registration_email_verification_enabled";
pub fn routes() -> Router<AppState> {
Router::new().route(
"/api/admin/user-oauth/settings",
get(get_settings).put(update_settings),
)
}
#[derive(Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ProviderSettingsUpdate {
enabled: bool,
client_id: String,
client_secret: Option<String>,
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct UserOAuthSettingsUpdate {
public_base_url: String,
google: ProviderSettingsUpdate,
github: ProviderSettingsUpdate,
wechat: ProviderSettingsUpdate,
}
#[derive(Clone, Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct ProviderSettingsResponse {
enabled: bool,
client_id: String,
client_secret_configured: bool,
client_secret_masked: Option<String>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct UserOAuthSettingsResponse {
public_base_url: String,
google: ProviderSettingsResponse,
github: ProviderSettingsResponse,
wechat: ProviderSettingsResponse,
}
#[derive(Clone)]
struct EffectiveProviderSettings {
enabled: bool,
client_id: String,
client_secret: Option<String>,
}
async fn get_settings(State(state): State<AppState>) -> AppResult<Json<UserOAuthSettingsResponse>> {
let settings = load_user_oauth_settings(&state).await?;
Ok(Json(settings_response(&settings)))
}
async fn update_settings(
State(state): State<AppState>,
Json(request): Json<UserOAuthSettingsUpdate>,
) -> AppResult<Json<UserOAuthSettingsResponse>> {
let current = load_user_oauth_settings(&state).await?;
let public_base_url = normalize_public_base_url(&request.public_base_url)?;
let google = effective_provider("google", request.google, &current);
let github = effective_provider("github", request.github, &current);
let wechat = effective_provider("wechat", request.wechat, &current);
let enabled = [(&google, "Google"), (&github, "GitHub"), (&wechat, "微信")]
.into_iter()
.filter(|(provider, _)| provider.enabled)
.collect::<Vec<_>>();
if !enabled.is_empty() {
if public_base_url.is_empty() {
return Err(AppError::BadRequest(
"启用第三方登录前必须配置公开访问地址".into(),
));
}
for (provider, label) in &enabled {
if provider.client_id.is_empty() || provider.client_secret.is_none() {
return Err(AppError::BadRequest(format!(
"启用 {label} 登录前必须完整填写 Client ID 和 Client Secret"
)));
}
}
let registration_email_enabled =
setting_bool(&current, REGISTRATION_EMAIL_ENABLED_KEY, false);
let smtp_ready = load_smtp_mail_config(&state.db).await?.is_ready();
if !registration_email_enabled || !smtp_ready {
return Err(AppError::BadRequest(
"启用第三方登录前必须先启用注册邮箱验证码,并确保 SMTP 配置就绪".into(),
));
}
}
let mut transaction = state.db.begin().await?;
upsert(&mut transaction, PUBLIC_BASE_URL_KEY, &public_base_url).await?;
save_provider(&mut transaction, "google", &google).await?;
save_provider(&mut transaction, "github", &github).await?;
save_provider(&mut transaction, "wechat", &wechat).await?;
transaction.commit().await?;
let saved = load_user_oauth_settings(&state).await?;
Ok(Json(settings_response(&saved)))
}
fn effective_provider(
provider: &str,
update: ProviderSettingsUpdate,
current: &HashMap<String, String>,
) -> EffectiveProviderSettings {
let secret_key = provider_key(provider, "client_secret");
let supplied_secret = update
.client_secret
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty() && !is_masked_secret(value));
EffectiveProviderSettings {
enabled: update.enabled,
client_id: update.client_id.trim().to_string(),
client_secret: supplied_secret.or_else(|| setting(current, &secret_key)),
}
}
async fn save_provider(
transaction: &mut Transaction<'_, Any>,
provider: &str,
settings: &EffectiveProviderSettings,
) -> AppResult<()> {
upsert(
transaction,
&provider_key(provider, "enabled"),
if settings.enabled { "true" } else { "false" },
)
.await?;
upsert(
transaction,
&provider_key(provider, "client_id"),
&settings.client_id,
)
.await?;
if let Some(secret) = settings.client_secret.as_deref() {
upsert(
transaction,
&provider_key(provider, "client_secret"),
secret,
)
.await?;
}
Ok(())
}
async fn upsert(transaction: &mut Transaction<'_, Any>, key: &str, value: &str) -> AppResult<()> {
crate::db::sql::query("INSERT INTO settings (key, value) VALUES (?, ?) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value")
.bind(key)
.bind(value)
.execute(&mut **transaction)
.await?;
Ok(())
}
async fn load_user_oauth_settings(state: &AppState) -> AppResult<HashMap<String, String>> {
let rows: Vec<(String, Option<String>)> = crate::db::sql::query_as(
"SELECT key, value FROM settings
WHERE key LIKE 'user_oauth.%'
OR key = 'runtime.registration_email_verification_enabled'",
)
.fetch_all(&*state.db)
.await?;
Ok(rows
.into_iter()
.filter_map(|(key, value)| value.map(|value| (key, value)))
.collect())
}
fn settings_response(settings: &HashMap<String, String>) -> UserOAuthSettingsResponse {
UserOAuthSettingsResponse {
public_base_url: setting(settings, PUBLIC_BASE_URL_KEY).unwrap_or_default(),
google: provider_response(settings, "google"),
github: provider_response(settings, "github"),
wechat: provider_response(settings, "wechat"),
}
}
fn provider_response(
settings: &HashMap<String, String>,
provider: &str,
) -> ProviderSettingsResponse {
let secret = setting(settings, &provider_key(provider, "client_secret"));
ProviderSettingsResponse {
enabled: setting_bool(settings, &provider_key(provider, "enabled"), false),
client_id: setting(settings, &provider_key(provider, "client_id")).unwrap_or_default(),
client_secret_configured: secret.is_some(),
client_secret_masked: secret.as_deref().map(mask_secret),
}
}
fn provider_key(provider: &str, field: &str) -> String {
format!("user_oauth.{provider}.{field}")
}
fn setting(settings: &HashMap<String, String>, key: &str) -> Option<String> {
settings
.get(key)
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
fn setting_bool(settings: &HashMap<String, String>, key: &str, default: bool) -> bool {
setting(settings, key)
.map(|value| {
matches!(
value.to_ascii_lowercase().as_str(),
"1" | "true" | "yes" | "on" | "enabled"
)
})
.unwrap_or(default)
}
fn is_masked_secret(value: &str) -> bool {
let value = value.trim();
value.starts_with("••••")
|| value.starts_with("****")
|| (!value.is_empty() && value.chars().all(|ch| ch == '•' || ch == '*'))
}
fn mask_secret(value: &str) -> String {
let chars: Vec<char> = value.chars().collect();
if chars.len() <= 8 {
return "••••".to_string();
}
let suffix: String = chars[chars.len().saturating_sub(4)..].iter().collect();
format!("••••{suffix}")
}
fn normalize_public_base_url(value: &str) -> AppResult<String> {
let value = value.trim();
if value.is_empty() {
return Ok(String::new());
}
let parsed =
Url::parse(value).map_err(|_| AppError::BadRequest("公开访问地址不是有效 URL".into()))?;
if !parsed.username().is_empty() || parsed.password().is_some() {
return Err(AppError::BadRequest(
"公开访问地址不能包含账号或密码".into(),
));
}
if parsed.query().is_some() || parsed.fragment().is_some() {
return Err(AppError::BadRequest(
"公开访问地址不能包含查询参数或片段".into(),
));
}
if parsed.path() != "/" && !parsed.path().is_empty() {
return Err(AppError::BadRequest(
"公开访问地址只能填写站点 origin,不能包含路径".into(),
));
}
match parsed.scheme() {
"https" => {}
"http" if is_loopback_host(parsed.host()) => {}
"http" => {
return Err(AppError::BadRequest(
"公开访问地址必须使用 HTTPS;仅本机开发地址允许 HTTP".into(),
));
}
_ => {
return Err(AppError::BadRequest(
"公开访问地址必须使用 HTTP 或 HTTPS".into(),
));
}
}
if parsed.host().is_none() {
return Err(AppError::BadRequest("公开访问地址缺少主机名".into()));
}
Ok(parsed.origin().ascii_serialization())
}
fn is_loopback_host(host: Option<Host<&str>>) -> bool {
match host {
Some(Host::Domain(host)) => host.eq_ignore_ascii_case("localhost"),
Some(Host::Ipv4(host)) => host.is_loopback(),
Some(Host::Ipv6(host)) => host.is_loopback(),
None => false,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn canonicalizes_public_origins_and_rejects_unsafe_values() {
assert_eq!(
normalize_public_base_url(" https://Example.com:8443/ ").unwrap(),
"https://example.com:8443"
);
assert_eq!(
normalize_public_base_url("http://127.0.0.1:4000/").unwrap(),
"http://127.0.0.1:4000"
);
assert_eq!(
normalize_public_base_url("http://[::1]:4000").unwrap(),
"http://[::1]:4000"
);
for invalid in [
"http://example.com",
"https://user:pass@example.com",
"https://example.com/app",
"https://example.com/?next=/admin",
"https://example.com/#settings",
"ftp://example.com",
] {
assert!(
normalize_public_base_url(invalid).is_err(),
"{invalid} must be rejected"
);
}
}
#[test]
fn response_masks_secrets_and_masked_updates_preserve_them() {
let current = HashMap::from([
(
"user_oauth.google.client_secret".to_string(),
"oauth-super-secret".to_string(),
),
(
"user_oauth.google.client_id".to_string(),
"client-id".to_string(),
),
]);
let response = provider_response(&current, "google");
assert!(response.client_secret_configured);
assert_eq!(response.client_secret_masked.as_deref(), Some("••••cret"));
let serialized = serde_json::to_string(&response).unwrap();
assert!(!serialized.contains("oauth-super-secret"));
assert!(!serialized.contains("clientSecret\""));
assert_eq!(mask_secret("short"), "••••");
let effective = effective_provider(
"google",
ProviderSettingsUpdate {
enabled: true,
client_id: "client-id".into(),
client_secret: Some("••••cret".into()),
},
&current,
);
assert_eq!(
effective.client_secret.as_deref(),
Some("oauth-super-secret")
);
}
#[test]
fn all_supported_provider_keys_use_the_protected_prefix() {
for provider in ["google", "github", "wechat"] {
for field in ["enabled", "client_id", "client_secret"] {
assert!(provider_key(provider, field).starts_with("user_oauth."));
}
}
}
}
+841 -205
View File
File diff suppressed because it is too large Load Diff
+1543 -302
View File
File diff suppressed because it is too large Load Diff
+1038 -144
View File
File diff suppressed because it is too large Load Diff
+23 -21
View File
@@ -85,7 +85,7 @@ async fn list_checkin_logs(
sql.push_str(" ORDER BY created_at DESC");
sql.push_str(&format!(" LIMIT {} OFFSET {}", limit, offset));
let mut query = sqlx::query_as::<_, CheckinLog>(&sql);
let mut query = crate::db::sql::query_as::<_, CheckinLog>(&sql);
for p in &params {
query = query.bind(p);
}
@@ -98,8 +98,8 @@ async fn list_checkin_logs(
async fn checkin_stats(State(state): State<AppState>) -> AppResult<Json<Value>> {
info!("Checkin stats");
let rows: Vec<CheckinStats> = sqlx::query_as::<_, CheckinStats>(
"SELECT a.id as account_id, s.name as site_name, a.username, a.status,
let rows: Vec<CheckinStats> = crate::db::sql::query_as::<_, CheckinStats>(
"SELECT a.id as account_id, s.name as site_name, COALESCE(a.username, '') AS username, a.status,
a.last_checkin_at as last_checkin,
COALESCE(a.checkin_enabled, 0) as checkin_enabled,
a.balance
@@ -132,7 +132,7 @@ async fn checkin_stats(State(state): State<AppState>) -> AppResult<Json<Value>>
/// GET /api/checkin/accounts — list accounts with checkin info
async fn list_checkin_accounts(State(state): State<AppState>) -> AppResult<Json<Value>> {
let accounts: Vec<Value> =
sqlx::query_as::<_, Account>("SELECT * FROM accounts ORDER BY site_id, username")
crate::db::sql::query_as::<_, Account>("SELECT * FROM accounts ORDER BY site_id, username")
.fetch_all(&*state.db)
.await?
.into_iter()
@@ -159,15 +159,16 @@ async fn toggle_checkin_account(
) -> AppResult<Json<Value>> {
info!("Toggling checkin for account: {}", id);
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 account: Account =
crate::db::sql::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 new_value = !account.checkin_enabled.unwrap_or(true);
sqlx::query(
"UPDATE accounts SET checkin_enabled = ?, updated_at = datetime('now') WHERE id = ?",
crate::db::sql::query(
"UPDATE accounts SET checkin_enabled = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?",
)
.bind(new_value)
.bind(id)
@@ -183,7 +184,7 @@ async fn toggle_checkin_account(
async fn run_checkin(State(state): State<AppState>) -> AppResult<Json<Value>> {
info!("Running checkin for all enabled accounts");
let accounts: Vec<Account> = sqlx::query_as::<_, Account>(
let accounts: Vec<Account> = crate::db::sql::query_as::<_, Account>(
"SELECT * FROM accounts WHERE status = 'active' AND COALESCE(checkin_enabled, 1) = 1",
)
.fetch_all(&*state.db)
@@ -193,13 +194,13 @@ async fn run_checkin(State(state): State<AppState>) -> AppResult<Json<Value>> {
for account in &accounts {
// Simple balance bonus checkin (matches Node.js basic behavior)
let bonus = 0.01; // $0.01 per day
sqlx::query("UPDATE accounts SET balance = COALESCE(balance, 0) + ?, last_checkin_at = datetime('now'), updated_at = datetime('now') WHERE id = ?")
crate::db::sql::query("UPDATE accounts SET balance = COALESCE(balance, 0) + ?, last_checkin_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP WHERE id = ?")
.bind(bonus)
.bind(account.id)
.execute(&*state.db)
.await?;
sqlx::query("INSERT INTO checkin_logs (account_id, status, message, reward) VALUES (?, 'success', 'Daily checkin completed', ?)")
crate::db::sql::query("INSERT INTO checkin_logs (account_id, status, message, reward) VALUES (?, 'success', 'Daily checkin completed', ?)")
.bind(account.id)
.bind(format!("+${:.4}", bonus))
.execute(&*state.db)
@@ -227,20 +228,21 @@ async fn run_checkin_account(
) -> AppResult<Json<Value>> {
info!("Running checkin for account: {}", id);
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 account: Account =
crate::db::sql::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 bonus = 0.01;
sqlx::query("UPDATE accounts SET balance = COALESCE(balance, 0) + ?, last_checkin_at = datetime('now'), updated_at = datetime('now') WHERE id = ?")
crate::db::sql::query("UPDATE accounts SET balance = COALESCE(balance, 0) + ?, last_checkin_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP WHERE id = ?")
.bind(bonus)
.bind(account.id)
.execute(&*state.db)
.await?;
sqlx::query("INSERT INTO checkin_logs (account_id, status, message, reward) VALUES (?, 'success', 'Daily checkin completed', ?)")
crate::db::sql::query("INSERT INTO checkin_logs (account_id, status, message, reward) VALUES (?, 'success', 'Daily checkin completed', ?)")
.bind(account.id)
.bind(format!("+${:.4}", bonus))
.execute(&*state.db)
+104 -42
View File
@@ -39,29 +39,44 @@ async fn submit_prompt(
.await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
let status = resp.status();
let text = resp.text().await.map_err(|e| AppError::UpstreamError(e.to_string()))?;
let text = resp
.text()
.await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
if !status.is_success() {
return Err(AppError::UpstreamError(format!("Comfy Cloud prompt [{status}]: {text}")));
return Err(AppError::UpstreamError(format!(
"Comfy Cloud prompt [{status}]: {text}"
)));
}
let value: Value = serde_json::from_str(&text).unwrap_or_else(|_| serde_json::json!({ "body": text }));
let value: Value =
serde_json::from_str(&text).unwrap_or_else(|_| serde_json::json!({ "body": text }));
Ok(Json(value))
}
async fn fetch_history(db: &sqlx::SqlitePool, prompt_id: &str) -> AppResult<Value> {
async fn fetch_history(db: &sqlx::AnyPool, prompt_id: &str) -> AppResult<Value> {
let (base_url, api_key) = resolve_comfy_cloud_site(db).await?;
let http = client();
let status_resp = http
.get(format!("{base_url}/api/job/{prompt_id}/status"))
.header("X-API-Key", &api_key)
.send().await
.send()
.await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
if !status_resp.status().is_success() {
let st = status_resp.status();
let body = status_resp.text().await.unwrap_or_default();
return Err(AppError::UpstreamError(format!("Comfy Cloud status [{st}]: {body}")));
return Err(AppError::UpstreamError(format!(
"Comfy Cloud status [{st}]: {body}"
)));
}
let status_json: Value = status_resp.json().await.map_err(|e| AppError::UpstreamError(e.to_string()))?;
let status = status_json.get("status").and_then(Value::as_str).unwrap_or("in_progress");
let status_json: Value = status_resp
.json()
.await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
let status = status_json
.get("status")
.and_then(Value::as_str)
.unwrap_or("in_progress");
if status != "completed" {
return Ok(comfy_cloud_job_to_history(prompt_id, status, None));
@@ -69,14 +84,20 @@ async fn fetch_history(db: &sqlx::SqlitePool, prompt_id: &str) -> AppResult<Valu
let detail_resp = http
.get(format!("{base_url}/api/jobs/{prompt_id}"))
.header("X-API-Key", &api_key)
.send().await
.send()
.await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
if !detail_resp.status().is_success() {
let st = detail_resp.status();
let body = detail_resp.text().await.unwrap_or_default();
return Err(AppError::UpstreamError(format!("Comfy Cloud jobs [{st}]: {body}")));
return Err(AppError::UpstreamError(format!(
"Comfy Cloud jobs [{st}]: {body}"
)));
}
let detail_json: Value = detail_resp.json().await.map_err(|e| AppError::UpstreamError(e.to_string()))?;
let detail_json: Value = detail_resp
.json()
.await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
let outputs = detail_json.get("outputs");
Ok(comfy_cloud_job_to_history(prompt_id, status, outputs))
}
@@ -88,14 +109,19 @@ async fn get_history(
Ok(Json(fetch_history(&state.db, &prompt_id).await?))
}
async fn fetch_view(db: &sqlx::SqlitePool, query: Option<&str>) -> AppResult<(String, Vec<u8>)> {
async fn fetch_view(db: &sqlx::AnyPool, query: Option<&str>) -> AppResult<(String, Vec<u8>)> {
let (base_url, api_key) = resolve_comfy_cloud_site(db).await?;
let q = query.unwrap_or("");
let url = if q.is_empty() { format!("{base_url}/api/view") } else { format!("{base_url}/api/view?{q}") };
let url = if q.is_empty() {
format!("{base_url}/api/view")
} else {
format!("{base_url}/api/view?{q}")
};
let resp = client()
.get(url)
.header("X-API-Key", api_key)
.send().await
.send()
.await
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
if !resp.status().is_success() {
let st = resp.status();
@@ -107,14 +133,15 @@ async fn fetch_view(db: &sqlx::SqlitePool, query: Option<&str>) -> AppResult<(St
.and_then(|v| v.to_str().ok())
.unwrap_or("application/octet-stream")
.to_string();
let bytes = resp.bytes().await.map_err(|e| AppError::UpstreamError(e.to_string()))?.to_vec();
let bytes = resp
.bytes()
.await
.map_err(|e| AppError::UpstreamError(e.to_string()))?
.to_vec();
Ok((mime, bytes))
}
async fn get_view(
State(state): State<AppState>,
RawQuery(query): RawQuery,
) -> AppResult<Response> {
async fn get_view(State(state): State<AppState>, RawQuery(query): RawQuery) -> AppResult<Response> {
let (mime, bytes) = fetch_view(&state.db, query.as_deref()).await?;
Response::builder()
.status(StatusCode::OK)
@@ -126,15 +153,22 @@ async fn get_view(
#[cfg(test)]
mod tests {
use super::*;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
use wiremock::matchers::{header as match_header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use sqlx::sqlite::SqlitePoolOptions;
async fn db_with_site(base_url: &str) -> sqlx::SqlitePool {
let pool = SqlitePoolOptions::new().max_connections(1).connect("sqlite::memory:").await.unwrap();
sqlx::query("CREATE TABLE sites (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT, url TEXT NOT NULL, platform TEXT NOT NULL, status TEXT NOT NULL, api_key TEXT, sort_order INTEGER)")
async fn db_with_site(base_url: &str) -> sqlx::AnyPool {
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
crate::db::sql::query("CREATE TABLE sites (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT, url TEXT NOT NULL, platform TEXT NOT NULL, status TEXT NOT NULL, api_key TEXT, sort_order INTEGER)")
.execute(&pool).await.unwrap();
sqlx::query("INSERT INTO sites (name,url,platform,status,api_key) VALUES ('cc',?,'comfy_cloud','active','key-123')")
crate::db::sql::query("INSERT INTO sites (name,url,platform,status,api_key) VALUES ('cc',?,'comfy_cloud','active','key-123')")
.bind(base_url).execute(&pool).await.unwrap();
pool
}
@@ -145,15 +179,21 @@ mod tests {
Mock::given(method("POST"))
.and(path("/api/prompt"))
.and(match_header("X-API-Key", "key-123"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "prompt_id": "p1" })))
.mount(&server).await;
.respond_with(
ResponseTemplate::new(200).set_body_json(serde_json::json!({ "prompt_id": "p1" })),
)
.mount(&server)
.await;
let pool = db_with_site(&server.uri()).await;
let (base_url, api_key) = resolve_comfy_cloud_site(&pool).await.unwrap();
let resp = client().post(format!("{base_url}/api/prompt"))
let resp = client()
.post(format!("{base_url}/api/prompt"))
.header("X-API-Key", api_key)
.json(&serde_json::json!({ "prompt": {} }))
.send().await.unwrap();
.send()
.await
.unwrap();
let v: Value = resp.json().await.unwrap();
assert_eq!(v["prompt_id"], "p1");
}
@@ -164,8 +204,12 @@ mod tests {
Mock::given(method("GET"))
.and(path("/api/job/p1/status"))
.and(match_header("X-API-Key", "key-123"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "status": "completed" })))
.mount(&server).await;
.respond_with(
ResponseTemplate::new(200)
.set_body_json(serde_json::json!({ "status": "completed" })),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/api/jobs/p1"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
@@ -176,7 +220,10 @@ mod tests {
let pool = db_with_site(&server.uri()).await;
let history = fetch_history(&pool, "p1").await.unwrap();
assert_eq!(history["p1"]["status"]["completed"], true);
assert_eq!(history["p1"]["outputs"]["9"]["images"][0]["filename"], "a.png");
assert_eq!(
history["p1"]["outputs"]["9"]["images"][0]["filename"],
"a.png"
);
}
#[tokio::test]
@@ -185,12 +232,19 @@ mod tests {
Mock::given(method("GET"))
.and(path("/api/job/p1/status"))
.and(match_header("X-API-Key", "key-123"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "status": "completed" })))
.mount(&server).await;
.respond_with(
ResponseTemplate::new(200)
.set_body_json(serde_json::json!({ "status": "completed" })),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/api/jobs/p1"))
.respond_with(ResponseTemplate::new(500).set_body_json(serde_json::json!({ "error": "boom" })))
.mount(&server).await;
.respond_with(
ResponseTemplate::new(500).set_body_json(serde_json::json!({ "error": "boom" })),
)
.mount(&server)
.await;
let pool = db_with_site(&server.uri()).await;
assert!(fetch_history(&pool, "p1").await.is_err());
@@ -202,8 +256,11 @@ mod tests {
Mock::given(method("GET"))
.and(path("/api/job/p1/status"))
.and(match_header("X-API-Key", "key-123"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "status": "failed" })))
.mount(&server).await;
.respond_with(
ResponseTemplate::new(200).set_body_json(serde_json::json!({ "status": "failed" })),
)
.mount(&server)
.await;
let pool = db_with_site(&server.uri()).await;
let history = fetch_history(&pool, "p1").await.unwrap();
@@ -216,13 +273,18 @@ mod tests {
Mock::given(method("GET"))
.and(path("/api/view"))
.and(match_header("X-API-Key", "key-123"))
.respond_with(ResponseTemplate::new(200)
.insert_header("content-type", "image/png")
.set_body_bytes(b"PNGDATA".to_vec()))
.mount(&server).await;
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "image/png")
.set_body_bytes(b"PNGDATA".to_vec()),
)
.mount(&server)
.await;
let pool = db_with_site(&server.uri()).await;
let (mime, bytes) = fetch_view(&pool, Some("filename=a.png&type=output")).await.unwrap();
let (mime, bytes) = fetch_view(&pool, Some("filename=a.png&type=output"))
.await
.unwrap();
assert_eq!(mime, "image/png");
assert_eq!(bytes, b"PNGDATA");
}
+22 -12
View File
@@ -87,14 +87,17 @@ mod tests {
services::{builtin_keys::BuiltinKeys, runtime_auth::RuntimeAuthSettings},
};
use axum::http::HeaderValue;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
async fn test_state() -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool");
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool");
AppState::new(
AppConfig {
@@ -108,6 +111,8 @@ mod tests {
url: "sqlite::memory:".to_string(),
max_connections: 1,
ssl: false,
ssl_mode: None,
ca_cert_path: None,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
@@ -126,9 +131,11 @@ mod tests {
template_proxy_url: None,
},
comfyui: ComfyUIConfig::default(),
agent: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
jimeng: Default::default(),
seedance: Default::default(),
},
pool,
BuiltinKeys {
@@ -629,12 +636,15 @@ async fn proxy_to_comfyui_response(
}
async fn list_boosapi_models(state: &AppState) -> AppResult<Vec<String>> {
let rows: Vec<String> = sqlx::query_scalar(
"SELECT DISTINCT COALESCE(NULLIF(TRIM(model_mapping), ''), model_pattern)
FROM token_routes
WHERE COALESCE(enabled, 1) = 1
AND TRIM(COALESCE(model_pattern, '')) <> ''
ORDER BY 1 COLLATE NOCASE",
let rows: Vec<String> = crate::db::sql::query_scalar(
"SELECT model
FROM (
SELECT DISTINCT COALESCE(NULLIF(TRIM(model_mapping), ''), model_pattern) AS model
FROM token_routes
WHERE COALESCE(enabled, 1) = 1
AND TRIM(COALESCE(model_pattern, '')) <> ''
) available_models
ORDER BY LOWER(model), model",
)
.fetch_all(&*state.db)
.await?;
File diff suppressed because it is too large Load Diff
+417
View File
@@ -0,0 +1,417 @@
use crate::{
auth::jwt::validate_jwt,
error::{AppError, AppResult},
services::commerce::{self, ProductInput},
AppState,
};
use axum::{
extract::{Path, State},
http::{HeaderMap, StatusCode},
routing::get,
Json, Router,
};
use serde::Deserialize;
use serde_json::{json, Value};
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct CreateOrderPayload {
product_id: String,
currency: String,
}
pub(super) fn routes() -> Router<AppState> {
Router::new()
.route("/api/commerce/catalog", get(list_catalog))
.route("/api/user/payment/catalog", get(list_catalog))
.route("/api/commerce/orders", get(list_orders).post(create_order))
.route(
"/api/user/payment/package-orders",
get(list_orders).post(create_order),
)
.route("/api/commerce/orders/{public_id}", get(get_order))
.route(
"/api/user/payment/package-orders/{public_id}",
get(get_order),
)
.route("/api/commerce/subscriptions", get(list_subscriptions))
.route("/api/user/payment/subscriptions", get(list_subscriptions))
}
pub(super) fn admin_routes() -> Router<AppState> {
Router::new()
.route(
"/api/admin/commerce/products",
get(admin_list_products).post(admin_create_product),
)
.route(
"/api/admin/commerce/products/{public_id}",
get(admin_get_product)
.put(admin_update_product)
.delete(admin_delete_product),
)
}
fn current_user_id(headers: &HeaderMap, state: &AppState) -> AppResult<i64> {
let token = headers
.get("Authorization")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.strip_prefix("Bearer "))
.ok_or_else(|| AppError::Unauthorized("Missing or invalid Authorization header".into()))?;
let claims = validate_jwt(&state.config.auth.jwt_secret, token)
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
claims
.sub
.parse()
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))
}
fn idempotency_key(headers: &HeaderMap) -> AppResult<&str> {
headers
.get("Idempotency-Key")
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| AppError::BadRequest("Idempotency-Key header is required".into()))
}
async fn list_catalog(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
current_user_id(&headers, &state)?;
let products = commerce::list_products(&state.db, false).await?;
Ok(Json(json!({ "success": true, "products": products })))
}
async fn list_orders(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
let user_id = current_user_id(&headers, &state)?;
let orders = commerce::list_orders(&state.db, user_id).await?;
Ok(Json(json!({ "success": true, "orders": orders })))
}
async fn create_order(
State(state): State<AppState>,
headers: HeaderMap,
Json(payload): Json<CreateOrderPayload>,
) -> AppResult<(StatusCode, Json<Value>)> {
let user_id = current_user_id(&headers, &state)?;
let key = idempotency_key(&headers)?;
let order = commerce::create_order_with_timeout(
&state.db,
user_id,
&payload.product_id,
&payload.currency,
key,
state.config.payment.order_timeout_minutes,
)
.await?;
Ok((
StatusCode::CREATED,
Json(json!({ "success": true, "order": order })),
))
}
async fn get_order(
State(state): State<AppState>,
headers: HeaderMap,
Path(public_id): Path<String>,
) -> AppResult<Json<Value>> {
let user_id = current_user_id(&headers, &state)?;
let order = commerce::get_order(&state.db, user_id, &public_id).await?;
Ok(Json(json!({ "success": true, "order": order })))
}
async fn list_subscriptions(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Json<Value>> {
let user_id = current_user_id(&headers, &state)?;
let subscriptions = commerce::list_subscription_periods(&state.db, user_id).await?;
Ok(Json(json!({
"success": true,
"subscriptions": subscriptions,
})))
}
async fn admin_list_products(State(state): State<AppState>) -> AppResult<Json<Value>> {
let products = commerce::list_products(&state.db, true).await?;
Ok(Json(json!({ "success": true, "products": products })))
}
async fn admin_create_product(
State(state): State<AppState>,
Json(payload): Json<ProductInput>,
) -> AppResult<(StatusCode, Json<Value>)> {
let product = commerce::create_product(&state.db, payload).await?;
Ok((
StatusCode::CREATED,
Json(json!({ "success": true, "product": product })),
))
}
async fn admin_get_product(
State(state): State<AppState>,
Path(public_id): Path<String>,
) -> AppResult<Json<Value>> {
let product = commerce::get_product(&state.db, &public_id, true).await?;
Ok(Json(json!({ "success": true, "product": product })))
}
async fn admin_update_product(
State(state): State<AppState>,
Path(public_id): Path<String>,
Json(payload): Json<ProductInput>,
) -> AppResult<Json<Value>> {
let product = commerce::update_product(&state.db, &public_id, payload).await?;
Ok(Json(json!({ "success": true, "product": product })))
}
async fn admin_delete_product(
State(state): State<AppState>,
Path(public_id): Path<String>,
) -> AppResult<Json<Value>> {
let product = commerce::archive_product(&state.db, &public_id).await?;
Ok(Json(json!({ "success": true, "product": product })))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
auth::jwt::generate_jwt,
config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig},
services::{
builtin_keys::BuiltinKeys,
commerce::{ProductBenefitInput, ProductPriceInput},
runtime_auth::RuntimeAuthSettings,
},
};
use axum::{
body::{to_bytes, Body},
http::Request,
};
use sqlx::any::{install_default_drivers, AnyPoolOptions};
use tower::ServiceExt as _;
async fn test_state() -> AppState {
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool");
sqlx::migrate!("./migrations")
.run(&pool)
.await
.expect("migrations");
AppState::new(
AppConfig {
server: ServerConfig {
host: "127.0.0.1".into(),
port: 0,
log_level: "info".into(),
cors_origins: vec![],
},
db: DbConfig {
url: "sqlite::memory:".into(),
max_connections: 1,
ssl: false,
ssl_mode: None,
ca_cert_path: None,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
},
auth: AuthConfig {
jwt_secret: "commerce-test-secret".into(),
jwt_expiry_hours: 1,
auth_token: "commerce-admin-token".into(),
proxy_token: "commerce-proxy-token".into(),
},
proxy: ProxyConfig {
upstream_timeout_secs: 30,
first_byte_timeout_secs: 30,
max_retries: 1,
proxy_log_retention_days: 7,
template_proxy_url: None,
},
comfyui: Default::default(),
agent: Default::default(),
pixelle_video: Default::default(),
payment: Default::default(),
jimeng: Default::default(),
seedance: Default::default(),
},
pool,
BuiltinKeys {
comfyui_agent: (1, "sk-commerce-test".into()),
},
RuntimeAuthSettings::from_token_for_test("commerce-proxy-token"),
)
}
fn user_token(state: &AppState, user_id: i64) -> String {
generate_jwt(
&state.config.auth.jwt_secret,
&user_id.to_string(),
"user",
1,
)
.expect("user jwt")
}
#[tokio::test]
async fn user_catalog_requires_jwt() {
let state = test_state().await;
let app = routes().with_state(state.clone());
let missing = app
.clone()
.oneshot(
Request::builder()
.uri("/api/commerce/catalog")
.body(Body::empty())
.unwrap(),
)
.await
.expect("catalog response");
assert_eq!(missing.status(), StatusCode::UNAUTHORIZED);
let authorized = app
.oneshot(
Request::builder()
.uri("/api/commerce/catalog")
.header("Authorization", format!("Bearer {}", user_token(&state, 7)))
.body(Body::empty())
.unwrap(),
)
.await
.expect("authorized catalog response");
assert_eq!(authorized.status(), StatusCode::OK);
}
#[tokio::test]
async fn admin_product_routes_use_admin_middleware() {
let state = test_state().await;
let app = super::super::admin::routes(state.clone()).with_state(state.clone());
let missing = app
.clone()
.oneshot(
Request::builder()
.uri("/api/admin/commerce/products")
.body(Body::empty())
.unwrap(),
)
.await
.expect("missing admin auth response");
assert_eq!(missing.status(), StatusCode::UNAUTHORIZED);
let user = app
.clone()
.oneshot(
Request::builder()
.uri("/api/admin/commerce/products")
.header("Authorization", format!("Bearer {}", user_token(&state, 7)))
.body(Body::empty())
.unwrap(),
)
.await
.expect("user auth response");
assert_eq!(user.status(), StatusCode::FORBIDDEN);
let admin = app
.oneshot(
Request::builder()
.uri("/api/admin/commerce/products")
.header("Authorization", "Bearer commerce-admin-token")
.body(Body::empty())
.unwrap(),
)
.await
.expect("admin response");
assert_eq!(admin.status(), StatusCode::OK);
}
#[tokio::test]
async fn create_order_uses_catalog_price_and_rejects_client_price() {
let state = test_state().await;
crate::db::sql::query(
"INSERT INTO users (id, username, email, password_hash) VALUES (7, 'buyer', 'buyer@example.test', 'hash')",
)
.execute(&*state.db)
.await
.expect("insert user");
let product = commerce::create_product(
&state.db,
ProductInput {
code: "api-starter".into(),
name: "API Starter".into(),
description: None,
product_type: "credit_pack".into(),
billing_track: None,
grant_kind: None,
validity_days: None,
subscription_days: None,
enabled: true,
sort_order: 0,
prices: vec![ProductPriceInput {
currency: "CNY".into(),
amount_minor: 2990,
}],
benefits: vec![ProductBenefitInput {
benefit_type: "api_credit".into(),
amount_micros: 5_000_000,
}],
credit_scopes: vec![],
unlock_scopes: vec![],
},
)
.await
.expect("create product");
let token = user_token(&state, 7);
let app = routes().with_state(state);
let response = app
.clone()
.oneshot(
Request::builder()
.method("POST")
.uri("/api/commerce/orders")
.header("Authorization", format!("Bearer {token}"))
.header("Idempotency-Key", "route-price-test")
.header("Content-Type", "application/json")
.body(Body::from(
json!({ "productId": product.public_id, "currency": "cny" }).to_string(),
))
.unwrap(),
)
.await
.expect("create order response");
assert_eq!(response.status(), StatusCode::CREATED);
let body = to_bytes(response.into_body(), 1024 * 64).await.unwrap();
let value: Value = serde_json::from_slice(&body).unwrap();
assert_eq!(value["order"]["amountMinor"], 2990);
assert_eq!(value["order"]["currency"], "CNY");
let rejected = app
.oneshot(
Request::builder()
.method("POST")
.uri("/api/commerce/orders")
.header("Authorization", format!("Bearer {token}"))
.header("Idempotency-Key", "route-price-override")
.header("Content-Type", "application/json")
.body(Body::from(
json!({
"productId": product.public_id,
"currency": "CNY",
"amountMinor": 1
})
.to_string(),
))
.unwrap(),
)
.await
.expect("price override response");
assert_eq!(rejected.status(), StatusCode::UNPROCESSABLE_ENTITY);
}
}
File diff suppressed because it is too large Load Diff
+92 -57
View File
@@ -17,6 +17,7 @@ use axum::{
use serde::Deserialize;
use serde::Serialize;
use serde_json::{json, Value};
use std::future::Future;
use tracing::{debug, info};
#[derive(Debug, Deserialize)]
@@ -88,15 +89,15 @@ async fn list_events(
}
if let Some(r) = read_filter {
if r == "true" {
sql.push_str(" AND read = 1");
sql.push_str(" AND read = TRUE");
} else if r == "false" {
sql.push_str(" AND read = 0");
sql.push_str(" AND read = FALSE");
}
}
sql.push_str(" ORDER BY created_at DESC LIMIT 1000");
let mut query = sqlx::query_as::<_, Event>(&sql);
let mut query = crate::db::sql::query_as::<_, Event>(&sql);
for p in &params {
query = query.bind(p);
}
@@ -112,24 +113,40 @@ async fn list_events(
);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "checkin") {
rows.extend(load_checkin_log_events(&state, source_limit).await?);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "proxy") {
rows.extend(load_proxy_log_events(&state, source_limit).await?);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "balance") {
rows.extend(load_balance_log_events(&state, source_limit).await?);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "token") {
rows.extend(load_token_log_events(&state, source_limit).await?);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "site_notice") {
rows.extend(load_site_notice_events(&state, source_limit).await?);
}
if should_include_read_only_source(filter_type.as_deref(), read_filter, "status") {
rows.extend(load_status_log_events(&state, source_limit).await?);
}
let include =
|source| should_include_read_only_source(filter_type.as_deref(), read_filter, source);
let (checkins, proxy, balances, tokens, notices, statuses) = tokio::try_join!(
load_event_source(
include("checkin"),
load_checkin_log_events(&state, source_limit)
),
load_event_source(
include("proxy"),
load_proxy_log_events(&state, source_limit)
),
load_event_source(
include("balance"),
load_balance_log_events(&state, source_limit)
),
load_event_source(
include("token"),
load_token_log_events(&state, source_limit)
),
load_event_source(
include("site_notice"),
load_site_notice_events(&state, source_limit)
),
load_event_source(
include("status"),
load_status_log_events(&state, source_limit)
),
)?;
rows.extend(checkins);
rows.extend(proxy);
rows.extend(balances);
rows.extend(tokens);
rows.extend(notices);
rows.extend(statuses);
rows.sort_by(|a, b| {
b.created_at
@@ -159,8 +176,19 @@ fn should_include_read_only_source(
should_include_source(filter_type, source) && read_filter != Some("false")
}
async fn load_event_source<F>(enabled: bool, future: F) -> AppResult<Vec<EventListItem>>
where
F: Future<Output = AppResult<Vec<EventListItem>>>,
{
if enabled {
future.await
} else {
Ok(Vec::new())
}
}
async fn load_checkin_log_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let rows = sqlx::query_as::<_, EventListItem>(
let rows = crate::db::sql::query_as::<_, EventListItem>(
r#"
SELECT
(-10000000000 - c.id) AS id,
@@ -179,7 +207,7 @@ async fn load_checkin_log_events(state: &AppState, limit: i64) -> AppResult<Vec<
WHEN lower(COALESCE(c.status, '')) IN ('skipped', 'skip') THEN 'warning'
ELSE 'error'
END AS level,
1 AS is_read,
TRUE AS is_read,
c.id AS related_id,
'checkin_log' AS related_type,
c.created_at AS created_at
@@ -197,7 +225,7 @@ async fn load_checkin_log_events(state: &AppState, limit: i64) -> AppResult<Vec<
}
async fn load_proxy_log_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let rows = sqlx::query_as::<_, EventListItem>(
let rows = crate::db::sql::query_as::<_, EventListItem>(
r#"
SELECT
(-20000000000 - l.id) AS id,
@@ -215,7 +243,7 @@ async fn load_proxy_log_events(state: &AppState, limit: i64) -> AppResult<Vec<Ev
|| CASE WHEN l.estimated_cost IS NOT NULL THEN ' · 成本=' || printf('$%.6f', l.estimated_cost) ELSE '' END
|| CASE WHEN l.error_message IS NOT NULL AND trim(l.error_message) != '' THEN ' · ' || l.error_message ELSE '' END AS message,
CASE WHEN lower(COALESCE(l.status, '')) = 'success' THEN 'info' ELSE 'error' END AS level,
1 AS is_read,
TRUE AS is_read,
l.id AS related_id,
'proxy_log' AS related_type,
l.created_at AS created_at
@@ -237,7 +265,7 @@ async fn load_proxy_log_events(state: &AppState, limit: i64) -> AppResult<Vec<Ev
}
async fn load_balance_log_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let rows = sqlx::query_as::<_, EventListItem>(
let rows = crate::db::sql::query_as::<_, EventListItem>(
r#"
SELECT
(-30000000000 - a.id) AS id,
@@ -248,7 +276,7 @@ async fn load_balance_log_events(state: &AppState, limit: i64) -> AppResult<Vec<
|| ' · 已用=' || printf('%.6f', COALESCE(a.balance_used, 0))
|| ' · 额度=' || printf('%.6f', COALESCE(a.quota, 0)) AS message,
'info' AS level,
1 AS is_read,
TRUE AS is_read,
a.id AS related_id,
'account' AS related_type,
a.last_balance_refresh AS created_at
@@ -267,7 +295,7 @@ async fn load_balance_log_events(state: &AppState, limit: i64) -> AppResult<Vec<
}
async fn load_token_log_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let rows = sqlx::query_as::<_, EventListItem>(
let rows = crate::db::sql::query_as::<_, EventListItem>(
r#"
SELECT
(-40000000000 - t.id) AS id,
@@ -287,7 +315,7 @@ async fn load_token_log_events(state: &AppState, limit: i64) -> AppResult<Vec<Ev
AND lower(COALESCE(t.value_status, 'ready')) IN ('ready', 'valid') THEN 'info'
ELSE 'warning'
END AS level,
1 AS is_read,
TRUE AS is_read,
t.id AS related_id,
'account_token' AS related_type,
COALESCE(t.updated_at, t.created_at) AS created_at
@@ -307,7 +335,7 @@ async fn load_token_log_events(state: &AppState, limit: i64) -> AppResult<Vec<Ev
}
async fn load_site_notice_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let rows = sqlx::query_as::<_, EventListItem>(
let rows = crate::db::sql::query_as::<_, EventListItem>(
r#"
SELECT
(-50000000000 - n.id) AS id,
@@ -316,7 +344,7 @@ async fn load_site_notice_events(state: &AppState, limit: i64) -> AppResult<Vec<
COALESCE(NULLIF(s.name, ''), '站点 #' || n.site_id)
|| CASE WHEN n.content IS NOT NULL AND trim(n.content) != '' THEN ' · ' || substr(n.content, 1, 420) ELSE '' END AS message,
COALESCE(NULLIF(n.level, ''), 'info') AS level,
1 AS is_read,
TRUE AS is_read,
n.id AS related_id,
'site_announcement' AS related_type,
COALESCE(n.upstream_created_at, n.first_seen_at) AS created_at
@@ -334,7 +362,7 @@ async fn load_site_notice_events(state: &AppState, limit: i64) -> AppResult<Vec<
async fn load_status_log_events(state: &AppState, limit: i64) -> AppResult<Vec<EventListItem>> {
let per_source_limit = limit.clamp(50, 1000);
let mut rows = sqlx::query_as::<_, EventListItem>(
let site_rows = crate::db::sql::query_as::<_, EventListItem>(
r#"
SELECT
(-60000000000 - s.id) AS id,
@@ -348,7 +376,7 @@ async fn load_status_log_events(state: &AppState, limit: i64) -> AppResult<Vec<E
WHEN lower(COALESCE(s.status, '')) IN ('disabled', 'inactive') THEN 'warning'
ELSE 'error'
END AS level,
1 AS is_read,
TRUE AS is_read,
s.id AS related_id,
'site' AS related_type,
COALESCE(s.updated_at, s.created_at) AS created_at
@@ -360,11 +388,9 @@ async fn load_status_log_events(state: &AppState, limit: i64) -> AppResult<Vec<E
"#,
)
.bind(per_source_limit)
.fetch_all(&*state.db)
.await?;
.fetch_all(&*state.db);
rows.extend(
sqlx::query_as::<_, EventListItem>(
let account_rows = crate::db::sql::query_as::<_, EventListItem>(
r#"
SELECT
(-70000000000 - a.id) AS id,
@@ -378,7 +404,7 @@ async fn load_status_log_events(state: &AppState, limit: i64) -> AppResult<Vec<E
WHEN lower(COALESCE(a.status, '')) IN ('disabled', 'inactive') THEN 'warning'
ELSE 'error'
END AS level,
1 AS is_read,
TRUE AS is_read,
a.id AS related_id,
'account' AS related_type,
COALESCE(a.updated_at, a.created_at) AS created_at
@@ -391,16 +417,18 @@ async fn load_status_log_events(state: &AppState, limit: i64) -> AppResult<Vec<E
"#,
)
.bind(per_source_limit)
.fetch_all(&*state.db)
.await?,
);
.fetch_all(&*state.db);
let (mut rows, account_rows) = tokio::try_join!(site_rows, account_rows)?;
rows.extend(account_rows);
Ok(rows)
}
/// GET /api/events/count — unread count
async fn unread_count(State(state): State<AppState>) -> AppResult<Json<Value>> {
let count: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM events WHERE read = 0")
let count: (i64,) =
crate::db::sql::query_as("SELECT COUNT(*) FROM events WHERE read = FALSE")
.fetch_one(&*state.db)
.await?;
Ok(Json(json!({ "count": count.0 })))
@@ -409,7 +437,7 @@ async fn unread_count(State(state): State<AppState>) -> AppResult<Json<Value>> {
/// POST /api/events/{id}/read
async fn mark_read(State(state): State<AppState>, Path(id): Path<i64>) -> AppResult<Json<Value>> {
info!("Marking event {} as read", id);
sqlx::query("UPDATE events SET read = 1 WHERE id = ?")
crate::db::sql::query("UPDATE events SET read = TRUE WHERE id = ?")
.bind(id)
.execute(&*state.db)
.await?;
@@ -419,7 +447,7 @@ async fn mark_read(State(state): State<AppState>, Path(id): Path<i64>) -> AppRes
/// POST /api/events/read-all
async fn mark_all_read(State(state): State<AppState>) -> AppResult<Json<Value>> {
info!("Marking all events as read");
sqlx::query("UPDATE events SET read = 1 WHERE read = 0")
crate::db::sql::query("UPDATE events SET read = TRUE WHERE read = FALSE")
.execute(&*state.db)
.await?;
Ok(Json(json!({ "success": true })))
@@ -428,7 +456,7 @@ async fn mark_all_read(State(state): State<AppState>) -> AppResult<Json<Value>>
/// DELETE /api/events
async fn clear_events(State(state): State<AppState>) -> AppResult<Json<Value>> {
info!("Clearing all events");
sqlx::query("DELETE FROM events")
crate::db::sql::query("DELETE FROM events")
.execute(&*state.db)
.await?;
runtime_logs::clear();
@@ -540,15 +568,18 @@ mod tests {
body::Body,
http::{Request, StatusCode},
};
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
use tower::ServiceExt as _;
async fn test_state() -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("memory sqlite");
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("memory sqlite");
sqlx::migrate!("./migrations")
.run(&pool)
.await
@@ -566,6 +597,8 @@ mod tests {
url: "sqlite::memory:".into(),
max_connections: 1,
ssl: false,
ssl_mode: None,
ca_cert_path: None,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
@@ -584,9 +617,11 @@ mod tests {
template_proxy_url: None,
},
comfyui: Default::default(),
agent: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
jimeng: Default::default(),
seedance: Default::default(),
},
pool,
BuiltinKeys {
@@ -597,14 +632,14 @@ mod tests {
}
async fn seed_event_sources(state: &AppState) {
sqlx::query(
crate::db::sql::query(
"INSERT INTO sites (id, name, url, platform, status)
VALUES (1, '主站', 'https://example.test', 'new-api', 'active')",
)
.execute(&*state.db)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO accounts
(id, site_id, username, access_token, balance, balance_used, quota, last_balance_refresh, status)
VALUES
@@ -613,7 +648,7 @@ mod tests {
.execute(&*state.db)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO account_tokens
(id, account_id, name, token, source, token_group, value_status, enabled, created_at, updated_at)
VALUES
@@ -622,14 +657,14 @@ mod tests {
.execute(&*state.db)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO checkin_logs (id, account_id, status, message, reward, created_at)
VALUES (30, 10, 'success', 'Daily checkin completed', '+$0.0100', '2026-06-08T09:50:00Z')",
)
.execute(&*state.db)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO proxy_logs
(id, account_id, model_requested, model_actual, status, http_status, latency_ms, total_tokens, estimated_cost, error_message, created_at)
VALUES
@@ -638,7 +673,7 @@ mod tests {
.execute(&*state.db)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO site_announcements
(id, site_id, platform, source_key, title, content, level, first_seen_at, read_at)
VALUES
+3 -1
View File
@@ -4,7 +4,9 @@ use serde_json::{json, Value};
pub async fn health_check(State(state): State<AppState>) -> AppResult<Json<Value>> {
// Simple DB ping
sqlx::query("SELECT 1").execute(&*state.db).await?;
crate::db::sql::query("SELECT 1")
.execute(&*state.db)
.await?;
Ok(Json(json!({
"status": "ok",
+556 -43
View File
@@ -9,6 +9,7 @@
//! the persisted local asset is exposed.
use crate::{
auth::api_key::{authorize_downstream_token, DownstreamAuthResult},
auth::jwt::validate_jwt,
db::models::background_task::BackgroundTask,
db::models::downstream_api_key::DownstreamApiKey,
@@ -21,18 +22,26 @@ use crate::{
AppState,
};
use axum::{
body::Body,
extract::{Path, Query, State},
http::{header, HeaderMap},
http::{header, HeaderMap, Response, StatusCode},
routing::{get, post},
Json, Router,
};
use hmac::{Hmac, Mac};
use serde::Deserialize;
use serde_json::{json, Value};
use std::time::Duration;
use sha2::Sha256;
use sqlx::Row;
use std::{path::PathBuf, time::Duration};
use tokio_util::io::ReaderStream;
const TASK_TYPE_IMAGE: &str = "jimeng-image-generation";
const TASK_TYPE_VIDEO: &str = "jimeng-video-generation";
const PROGRESS_TOTAL: i64 = 4;
const INPUT_ASSET_URL_TTL_SECS: i64 = 60 * 60 * 2;
type HmacSha256 = Hmac<Sha256>;
#[derive(Debug, Clone)]
struct JimengAuth {
@@ -42,6 +51,18 @@ struct JimengAuth {
key_name: String,
}
#[derive(Debug, sqlx::FromRow)]
struct JimengAccountCredentials {
access_key: String,
extra_config: Option<String>,
}
#[derive(Clone)]
struct ResolvedJimengProvider {
provider: JimengProvider,
source: &'static str,
}
pub fn routes() -> Router<AppState> {
Router::new()
.route("/api/jimeng/status", get(status_handler))
@@ -50,6 +71,10 @@ pub fn routes() -> Router<AppState> {
.route("/api/jimeng/jobs", get(list_jobs))
.route("/api/jimeng/jobs/{id}", get(get_job))
.route("/api/jimeng/jobs/{id}/cancel", post(cancel_job))
.route(
"/api/jimeng/input-assets/{public_id}/file",
get(get_signed_input_asset_file),
)
}
fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
@@ -64,8 +89,195 @@ fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
.filter(|token| !token.is_empty())
}
fn header_str<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> {
headers
.get(name)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn public_base_url_from_headers(state: &AppState, headers: &HeaderMap) -> String {
if let (Some(proto), Some(host)) = (
header_str(headers, "x-forwarded-proto"),
header_str(headers, "x-forwarded-host").or_else(|| header_str(headers, "host")),
) {
let proto = proto.split(',').next().unwrap_or(proto).trim();
let host = host.split(',').next().unwrap_or(host).trim();
if matches!(proto, "http" | "https") && !host.is_empty() {
return format!("{proto}://{host}")
.trim_end_matches('/')
.to_string();
}
}
if let Some(origin) = header_str(headers, header::ORIGIN.as_str()) {
if origin.starts_with("http://") || origin.starts_with("https://") {
return origin.trim_end_matches('/').to_string();
}
}
if let Some(host) = header_str(headers, "host") {
let host_lower = host.to_ascii_lowercase();
let scheme = if host_lower.starts_with("127.")
|| host_lower.starts_with("localhost")
|| host_lower.starts_with("[::1]")
{
"http"
} else {
"https"
};
return format!("{scheme}://{host}")
.trim_end_matches('/')
.to_string();
}
let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
"127.0.0.1".to_string()
} else {
state.config.server.host.clone()
};
format!("http://{}:{}", host, state.config.server.port)
}
fn input_asset_signature(secret: &str, public_id: &str, expires_at: i64) -> AppResult<String> {
let mut mac = HmacSha256::new_from_slice(secret.as_bytes())
.map_err(|err| AppError::Internal(anyhow::anyhow!("create input asset signer: {err}")))?;
mac.update(format!("jimeng-input-asset:{public_id}:{expires_at}").as_bytes());
Ok(hex::encode(mac.finalize().into_bytes()))
}
fn verify_input_asset_signature(
secret: &str,
public_id: &str,
expires_at: i64,
signature: &str,
) -> AppResult<()> {
if expires_at < chrono::Utc::now().timestamp() {
return Err(AppError::Forbidden(
"Jimeng input asset link expired".into(),
));
}
let signature = hex::decode(signature.trim())
.map_err(|_| AppError::Forbidden("Invalid Jimeng input asset signature".into()))?;
let mut mac = HmacSha256::new_from_slice(secret.as_bytes())
.map_err(|err| AppError::Internal(anyhow::anyhow!("create input asset verifier: {err}")))?;
mac.update(format!("jimeng-input-asset:{public_id}:{expires_at}").as_bytes());
mac.verify_slice(&signature)
.map_err(|_| AppError::Forbidden("Invalid Jimeng input asset signature".into()))
}
fn workspace_public_id_from_ref(value: &str) -> Option<String> {
const MARKERS: [&str; 3] = [
"/api/user/workspace/assets/",
"/api/admin/workspace/assets/",
"/api/workspace/assets/",
];
let trimmed = value.trim();
let (_, rest) = MARKERS
.iter()
.find_map(|marker| trimmed.split_once(marker).map(|parts| (*marker, parts.1)))?;
let public_id = rest
.split(['/', '?', '#'])
.next()
.unwrap_or_default()
.trim();
if public_id.is_empty() {
None
} else {
Some(public_id.to_string())
}
}
async fn ensure_auth_can_use_workspace_asset(
state: &AppState,
auth: &JimengAuth,
public_id: &str,
) -> AppResult<()> {
let row = crate::db::sql::query(
"SELECT user_id, mime_type
FROM workspace_assets
WHERE public_id = ?
AND deleted_at IS NULL
AND status = 'completed'
LIMIT 1",
)
.bind(public_id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("Workspace input asset not found".into()))?;
let mime_type = row
.try_get::<Option<String>, _>("mime_type")
.ok()
.flatten()
.unwrap_or_default();
if !mime_type.to_ascii_lowercase().starts_with("image/") {
return Err(AppError::BadRequest("即梦参考素材必须是图片文件".into()));
}
if auth.role == "admin" && auth.user_id.is_none() {
return Ok(());
}
let asset_user_id = row.try_get::<Option<i64>, _>("user_id").ok().flatten();
if asset_user_id == auth.user_id {
Ok(())
} else {
Err(AppError::Forbidden(
"Workspace input asset access denied".into(),
))
}
}
fn signed_input_asset_url(
state: &AppState,
headers: &HeaderMap,
public_id: &str,
) -> AppResult<String> {
let expires_at = chrono::Utc::now().timestamp() + INPUT_ASSET_URL_TTL_SECS;
let sig = input_asset_signature(&state.config.auth.jwt_secret, public_id, expires_at)?;
Ok(format!(
"{}/api/jimeng/input-assets/{}/file?exp={}&sig={}",
public_base_url_from_headers(state, headers),
public_id,
expires_at,
sig
))
}
async fn resolve_jimeng_image_refs(
state: &AppState,
headers: &HeaderMap,
auth: &JimengAuth,
refs: Vec<String>,
) -> AppResult<Vec<String>> {
let mut out = Vec::new();
let mut seen = std::collections::HashSet::new();
for raw in refs {
let value = raw.trim();
if value.is_empty() {
continue;
}
if let Some(public_id) = workspace_public_id_from_ref(value) {
ensure_auth_can_use_workspace_asset(state, auth, &public_id).await?;
let url = signed_input_asset_url(state, headers, &public_id)?;
if seen.insert(url.clone()) {
out.push(url);
}
continue;
}
if value.starts_with("http://") || value.starts_with("https://") {
let url = value.to_string();
if seen.insert(url.clone()) {
out.push(url);
}
}
}
Ok(out)
}
async fn user_downstream_key(state: &AppState, user_id: i64) -> Option<DownstreamApiKey> {
sqlx::query_as::<_, DownstreamApiKey>(
crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT *
FROM downstream_api_keys
WHERE user_id = ?
@@ -73,7 +285,7 @@ async fn user_downstream_key(state: &AppState, user_id: i64) -> Option<Downstrea
AND (
expires_at IS NULL
OR expires_at = ''
OR datetime(expires_at) > datetime('now')
OR datetime(expires_at) > CURRENT_TIMESTAMP
)
ORDER BY updated_at DESC, created_at DESC, id DESC
LIMIT 1",
@@ -102,6 +314,24 @@ async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<JimengAut
});
}
if token.starts_with("sk-") {
let result = authorize_downstream_token(token, &state.db, Some(&proxy_token))
.await
.map_err(AppError::Unauthorized)?;
if let DownstreamAuthResult::Managed { key } = result {
return Ok(JimengAuth {
user_id: key.user_id,
role: if key.user_id.is_some() {
"user".to_string()
} else {
"admin".to_string()
},
key_id: Some(key.id),
key_name: key.name,
});
}
}
let claims = validate_jwt(&state.config.auth.jwt_secret, token)
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
if claims.role != "admin" && claims.role != "user" {
@@ -131,23 +361,109 @@ async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<JimengAut
}
}
fn provider_or_unavailable(state: &AppState) -> AppResult<JimengProvider> {
JimengProvider::from_config(&state.config.jimeng).map_err(|err| match err {
JimengError::Disabled => {
AppError::ServiceUnavailable("即梦能力未启用,请在配置中开启 [jimeng]".into())
}
other => AppError::Internal(anyhow::anyhow!(other.to_string())),
})
fn string_from_value(value: &Value, keys: &[&str]) -> Option<String> {
keys.iter()
.find_map(|key| value.get(*key).and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn jimeng_provider_from_account(row: &JimengAccountCredentials) -> AppResult<JimengProvider> {
let value = row
.extra_config
.as_deref()
.map(str::trim)
.filter(|raw| !raw.is_empty())
.map(|raw| serde_json::from_str::<Value>(raw))
.transpose()
.map_err(|err| AppError::Internal(anyhow::anyhow!("invalid jimeng extra_config: {err}")))?
.unwrap_or_else(|| json!({}));
let jimeng = value.get("jimeng").unwrap_or(&Value::Null);
let secret_key = string_from_value(jimeng, &["secretKey", "secret_key"]).ok_or_else(|| {
AppError::ServiceUnavailable(
"即梦连接缺少 Secret Key(SK),请在站点管理的连接里补全".into(),
)
})?;
let region = string_from_value(jimeng, &["region"]);
let host = string_from_value(jimeng, &["host"]);
let session_token = string_from_value(jimeng, &["sessionToken", "session_token"]);
JimengProvider::from_credentials(
&row.access_key,
&secret_key,
region.as_deref(),
host.as_deref(),
session_token.as_deref(),
)
.map_err(|err| AppError::ServiceUnavailable(format!("即梦连接凭证不可用:{err}")))
}
async fn provider_or_unavailable(state: &AppState) -> AppResult<ResolvedJimengProvider> {
let account = crate::db::sql::query_as::<_, JimengAccountCredentials>(
"SELECT a.access_token AS access_key, a.extra_config AS extra_config
FROM accounts a
JOIN sites s ON s.id = a.site_id
WHERE lower(trim(COALESCE(s.platform, ''))) = 'jimeng'
AND lower(COALESCE(s.status, '')) = 'active'
AND lower(COALESCE(a.status, '')) = 'active'
AND trim(COALESCE(a.access_token, '')) <> ''
ORDER BY COALESCE(a.is_pinned, 0) DESC,
COALESCE(a.sort_order, 9223372036854775807) ASC,
a.updated_at DESC,
a.created_at DESC,
a.id DESC
LIMIT 1",
)
.fetch_optional(&*state.db)
.await?;
if let Some(account) = account {
return Ok(ResolvedJimengProvider {
provider: jimeng_provider_from_account(&account)?,
source: "site_account",
});
}
JimengProvider::from_config(&state.config.jimeng)
.map(|provider| ResolvedJimengProvider {
provider,
source: "config",
})
.map_err(|err| match err {
JimengError::Disabled => AppError::ServiceUnavailable(
"即梦能力未启用:请先在站点管理添加平台为「即梦 (火山视觉)」的站点,并在连接管理里保存 AK/SK;或在配置中开启 [jimeng]".into(),
),
other => AppError::Internal(anyhow::anyhow!(other.to_string())),
})
}
// PLACEHOLDER_REST
async fn status_handler(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
async fn status_handler(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let cfg = &state.config.jimeng;
let site_account_enabled = crate::db::sql::query_scalar::<_, i64>(
"SELECT COUNT(*)
FROM accounts a
JOIN sites s ON s.id = a.site_id
WHERE lower(trim(COALESCE(s.platform, ''))) = 'jimeng'
AND lower(COALESCE(s.status, '')) = 'active'
AND lower(COALESCE(a.status, '')) = 'active'
AND trim(COALESCE(a.access_token, '')) <> ''",
)
.fetch_one(&*state.db)
.await
.unwrap_or(0)
> 0;
Ok(Json(json!({
"success": true,
"enabled": cfg.enabled,
"enabled": cfg.enabled || site_account_enabled,
"providerSource": if site_account_enabled { "site_account" } else if cfg.enabled { "config" } else { "none" },
"configEnabled": cfg.enabled,
"siteAccountEnabled": site_account_enabled,
"role": auth.role,
"userId": auth.user_id,
"keyName": auth.key_name,
@@ -163,6 +479,82 @@ struct JobListQuery {
limit: Option<i64>,
}
#[derive(Debug, Deserialize)]
struct SignedInputAssetQuery {
exp: i64,
sig: String,
}
async fn get_signed_input_asset_file(
State(state): State<AppState>,
Path(public_id): Path<String>,
Query(query): Query<SignedInputAssetQuery>,
) -> AppResult<Response<Body>> {
verify_input_asset_signature(
&state.config.auth.jwt_secret,
&public_id,
query.exp,
&query.sig,
)?;
let row = crate::db::sql::query(
"SELECT file_path, mime_type, file_name
FROM workspace_assets
WHERE public_id = ?
AND deleted_at IS NULL
AND status = 'completed'
LIMIT 1",
)
.bind(&public_id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("Jimeng input asset not found".into()))?;
let mime_type = row
.try_get::<Option<String>, _>("mime_type")
.ok()
.flatten()
.unwrap_or_else(|| "application/octet-stream".to_string());
if !mime_type.to_ascii_lowercase().starts_with("image/") {
return Err(AppError::BadRequest(
"Jimeng input asset must be an image".into(),
));
}
let file_path = row
.try_get::<Option<String>, _>("file_path")
.ok()
.flatten()
.filter(|value| !value.trim().is_empty())
.ok_or_else(|| {
AppError::NotFound("Jimeng input asset file is not archived locally".into())
})?;
let path = PathBuf::from(file_path);
let file = tokio::fs::File::open(&path)
.await
.map_err(|_| AppError::NotFound("Jimeng input asset file not found".into()))?;
let metadata = file
.metadata()
.await
.map_err(|_| AppError::NotFound("Jimeng input asset file not found".into()))?;
let file_name = row
.try_get::<Option<String>, _>("file_name")
.ok()
.flatten()
.unwrap_or_else(|| format!("{public_id}.png"));
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, mime_type)
.header(header::CONTENT_LENGTH, metadata.len().to_string())
.header(header::ACCEPT_RANGES, "bytes")
.header(
header::CONTENT_DISPOSITION,
format!("inline; filename=\"{}\"", file_name.replace('"', "")),
)
.body(Body::from_stream(ReaderStream::new(file)))
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))
}
/// A task is visible to the caller when they are the global admin, or when the
/// task's stored `userId` matches the caller's id.
fn job_visible_to_auth(task: &BackgroundTask, auth: &JimengAuth) -> bool {
@@ -202,7 +594,7 @@ fn task_to_json(task: BackgroundTask) -> Value {
}
async fn load_job(state: &AppState, id: &str, auth: &JimengAuth) -> AppResult<BackgroundTask> {
let task: BackgroundTask = sqlx::query_as(
let task: BackgroundTask = crate::db::sql::query_as(
"SELECT * FROM background_tasks
WHERE id = ? AND task_type IN (?, ?) LIMIT 1",
)
@@ -219,9 +611,9 @@ async fn load_job(state: &AppState, id: &str, auth: &JimengAuth) -> AppResult<Ba
}
async fn update_job_progress(state: &AppState, id: &str, progress: i64, message: &str) {
let _ = sqlx::query(
let _ = crate::db::sql::query(
"UPDATE background_tasks
SET progress = ?, message = ?, updated_at = datetime('now')
SET progress = ?, message = ?, updated_at = CURRENT_TIMESTAMP
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(progress)
@@ -233,7 +625,7 @@ async fn update_job_progress(state: &AppState, id: &str, progress: i64, message:
async fn is_job_cancelled(state: &AppState, id: &str) -> bool {
let status: Option<String> =
sqlx::query_scalar("SELECT status FROM background_tasks WHERE id = ? LIMIT 1")
crate::db::sql::query_scalar("SELECT status FROM background_tasks WHERE id = ? LIMIT 1")
.bind(id)
.fetch_optional(&*state.db)
.await
@@ -251,6 +643,8 @@ struct ImageJobRequest {
model: Option<String>,
#[serde(default)]
image_urls: Vec<String>,
#[serde(default, alias = "aspect_ratio", alias = "aspectRatio")]
ratio: Option<String>,
#[serde(default)]
width: Option<i64>,
#[serde(default)]
@@ -293,12 +687,31 @@ async fn create_image_job(
Json(req): Json<ImageJobRequest>,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let provider = provider_or_unavailable(&state)?;
if let Some(user_id) = auth.user_id.filter(|_| auth.role == "user") {
let billing_model = req
.model
.as_deref()
.unwrap_or(JimengCapability::Image.req_key());
crate::services::billing::account::ensure_model_can_spend(
&state.db,
user_id,
billing_model,
)
.await?;
}
if let Some(key_id) = auth.key_id {
UsageTracker::new(state.db.clone())
.check_rate_limit(key_id)
.await?;
}
let resolved_provider = provider_or_unavailable(&state).await?;
let image_urls = resolve_jimeng_image_refs(&state, &headers, &auth, req.image_urls).await?;
let image_req = ImageRequest {
prompt: req.prompt,
model: req.model,
image_urls: req.image_urls,
image_urls,
ratio: req.ratio,
width: req.width,
height: req.height,
size: req.size,
@@ -314,14 +727,16 @@ async fn create_image_job(
let task_id = spawn_job(
&state,
&auth,
provider,
resolved_provider.provider,
JimengCapability::Image,
body,
title,
json!({
"providerSource": resolved_provider.source,
"prompt": image_req.prompt,
"model": image_req.model,
"imageUrls": image_req.image_urls,
"ratio": image_req.ratio,
"width": image_req.width,
"height": image_req.height,
"size": image_req.size,
@@ -342,12 +757,30 @@ async fn create_video_job(
Json(req): Json<VideoJobRequest>,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let provider = provider_or_unavailable(&state)?;
if let Some(user_id) = auth.user_id.filter(|_| auth.role == "user") {
let billing_model = req
.model
.as_deref()
.unwrap_or(JimengCapability::Video.req_key());
crate::services::billing::account::ensure_model_can_spend(
&state.db,
user_id,
billing_model,
)
.await?;
}
if let Some(key_id) = auth.key_id {
UsageTracker::new(state.db.clone())
.check_rate_limit(key_id)
.await?;
}
let resolved_provider = provider_or_unavailable(&state).await?;
let image_urls = resolve_jimeng_image_refs(&state, &headers, &auth, req.image_urls).await?;
let video_req = VideoRequest {
prompt: req.prompt,
model: req.model,
image_urls: req.image_urls,
image_urls,
frames: req.frames,
aspect_ratio: req.aspect_ratio,
seed: req.seed,
@@ -364,11 +797,12 @@ async fn create_video_job(
let task_id = spawn_job(
&state,
&auth,
provider,
resolved_provider.provider,
JimengCapability::Video,
body,
title,
json!({
"providerSource": resolved_provider.source,
"prompt": video_req.prompt,
"model": video_req.model,
"imageUrls": video_req.image_urls,
@@ -396,7 +830,7 @@ async fn list_jobs(
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let limit = query.limit.unwrap_or(20).clamp(1, 100);
let tasks: Vec<BackgroundTask> = sqlx::query_as(
let tasks: Vec<BackgroundTask> = crate::db::sql::query_as(
"SELECT * FROM background_tasks
WHERE task_type IN (?, ?)
ORDER BY created_at DESC
@@ -434,15 +868,22 @@ async fn cancel_job(
let auth = resolve_auth(&headers, &state).await?;
let task = load_job(&state, &id, &auth).await?;
if task.status == "pending" || task.status == "running" {
sqlx::query(
crate::db::sql::query(
"UPDATE background_tasks
SET status = 'cancelled', message = '已取消', error = NULL,
updated_at = datetime('now'), finished_at = datetime('now')
updated_at = CURRENT_TIMESTAMP, finished_at = CURRENT_TIMESTAMP
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(&id)
.execute(&*state.db)
.await?;
let _ = crate::services::comfy_billing::release_task(
&state.db,
&task.task_type,
&id,
Some("cancelled"),
)
.await;
}
let task = load_job(&state, &id, &auth).await?;
Ok(Json(json!({ "success": true, "job": task_to_json(task) })))
@@ -460,8 +901,24 @@ async fn spawn_job(
title: String,
request_snapshot: Value,
) -> AppResult<String> {
let task_id = format!("{}_{}", capability.id_prefix(), uuid::Uuid::new_v4().simple());
let task_id = format!(
"{}_{}",
capability.id_prefix(),
uuid::Uuid::new_v4().simple()
);
let task_type = capability.task_type();
crate::services::comfy_billing::reserve_task(
&state.db,
if auth.role == "admin" {
None
} else {
auth.user_id
},
task_type,
&task_id,
None,
)
.await?;
let model = body
.get("req_key")
.and_then(Value::as_str)
@@ -475,10 +932,10 @@ async fn spawn_job(
"capability": model,
"request": request_snapshot,
});
sqlx::query(
let inserted = crate::db::sql::query(
"INSERT INTO background_tasks
(id, task_type, title, status, message, dedupe_key, result_json, progress, progress_total, started_at)
VALUES (?, ?, ?, 'running', '任务已提交,等待火山', ?, ?, 1, ?, datetime('now'))",
VALUES (?, ?, ?, 'running', '任务已提交,等待火山', ?, ?, 1, ?, CURRENT_TIMESTAMP)",
)
.bind(&task_id)
.bind(task_type)
@@ -491,7 +948,17 @@ async fn spawn_job(
.bind(initial_result.to_string())
.bind(PROGRESS_TOTAL)
.execute(&*state.db)
.await?;
.await;
if let Err(err) = inserted {
let _ = crate::services::comfy_billing::release_task(
&state.db,
task_type,
&task_id,
Some("task_insert_failed"),
)
.await;
return Err(err.into());
}
let state_for_task = state.clone();
let auth_for_task = auth.clone();
@@ -540,10 +1007,10 @@ async fn run_worker(
}
if !assets_have_completed(&assets) {
let error = "产物转存失败,未生成可用工作空间文件";
let _ = sqlx::query(
let _ = crate::db::sql::query(
"UPDATE background_tasks
SET status = 'failed', message = '生成完成但转存失败', error = ?,
result_json = ?, updated_at = datetime('now'), finished_at = datetime('now')
result_json = ?, updated_at = CURRENT_TIMESTAMP, finished_at = CURRENT_TIMESTAMP
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(error)
@@ -551,14 +1018,28 @@ async fn run_worker(
.bind(&task_id)
.execute(&*state.db)
.await;
record_billing(&state, &auth, capability, &model, false, Some(error.to_string()))
.await;
record_billing(
&state,
&auth,
capability,
&model,
false,
Some(error.to_string()),
)
.await;
let _ = crate::services::comfy_billing::release_task(
&state.db,
capability.task_type(),
&task_id,
Some("archive_failed"),
)
.await;
return;
}
let _ = sqlx::query(
let _ = crate::db::sql::query(
"UPDATE background_tasks
SET status = 'succeeded', message = '生成完成', result_json = ?, progress = 4,
updated_at = datetime('now'), finished_at = datetime('now')
updated_at = CURRENT_TIMESTAMP, finished_at = CURRENT_TIMESTAMP
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(final_result.to_string())
@@ -566,22 +1047,52 @@ async fn run_worker(
.execute(&*state.db)
.await;
record_billing(&state, &auth, capability, &model, true, None).await;
if let Err(error) = crate::services::comfy_billing::settle_task(
&state.db,
capability.task_type(),
&task_id,
)
.await
{
tracing::warn!(%error, task_id, "failed to settle Jimeng Comfy points");
}
}
Err(JimengError::Cancelled) => {
// Cancellation already set the terminal status; nothing to bill.
let _ = crate::services::comfy_billing::release_task(
&state.db,
capability.task_type(),
&task_id,
Some("cancelled"),
)
.await;
}
Err(err) => {
let _ = sqlx::query(
let _ = crate::db::sql::query(
"UPDATE background_tasks
SET status = 'failed', message = '生成失败', error = ?,
updated_at = datetime('now'), finished_at = datetime('now')
updated_at = CURRENT_TIMESTAMP, finished_at = CURRENT_TIMESTAMP
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(err.to_string())
.bind(&task_id)
.execute(&*state.db)
.await;
record_billing(&state, &auth, capability, &model, false, Some(err.to_string())).await;
record_billing(
&state,
&auth,
capability,
&model,
false,
Some(err.to_string()),
)
.await;
let _ = crate::services::comfy_billing::release_task(
&state.db,
capability.task_type(),
&task_id,
Some("failed"),
)
.await;
}
}
}
@@ -598,10 +1109,10 @@ async fn drive(
update_job_progress(state, task_id, 2, "已提交火山,正在生成").await;
// Record the upstream task id so it survives a poll restart / inspection.
let _ = sqlx::query(
let _ = crate::db::sql::query(
"UPDATE background_tasks
SET result_json = json_set(COALESCE(result_json, '{}'), '$.volcTaskId', ?),
updated_at = datetime('now')
updated_at = CURRENT_TIMESTAMP
WHERE id = ?",
)
.bind(&volc_task_id)
@@ -654,7 +1165,9 @@ async fn store_outputs(
let input = asset_input(auth, capability, model, None);
match record_base64_asset(state, input, data).await {
Ok(stored) => assets.push(asset_to_json(Ok(stored), auth.user_id.is_some())),
Err(err) => assets.push(json!({ "status": "failed", "error": err.to_string() })),
Err(err) => {
assets.push(json!({ "status": "failed", "error": err.to_string() }))
}
}
}
}
+40
View File
@@ -1,7 +1,9 @@
use crate::AppState;
use axum::{
body::Body,
extract::State,
http::{header, HeaderValue, Request, StatusCode},
middleware::{self, Next},
response::{IntoResponse, Response},
routing::{any, get},
Json, Router,
@@ -28,6 +30,7 @@ mod checkin;
mod comfy_cloud_adapter;
mod comfyui;
mod comfyui_agent;
mod commerce;
mod digital_human;
mod events;
mod health;
@@ -36,14 +39,19 @@ mod money_printer;
mod oauth;
mod organizations;
mod payment;
mod payment_providers;
mod payment_refunds;
mod picture_agent;
mod profile_avatar;
mod proxy;
mod search;
mod seedance;
mod site_announcements;
pub(crate) mod stats;
mod tasks;
mod user;
mod user_api_keys;
mod user_oauth;
#[path = "video_agent_legacy.rs"]
mod video_agent;
mod workflow_ui;
@@ -63,6 +71,7 @@ pub fn build_router(state: AppState) -> Router {
.merge(digital_human::routes())
.merge(comfyui::routes())
.merge(jimeng::routes())
.merge(seedance::routes())
.merge(crate::workflow::routes::routes())
.route("/comfy", any(removed_comfy_workbench))
.route("/comfy/", any(removed_comfy_workbench))
@@ -70,11 +79,15 @@ pub fn build_router(state: AppState) -> Router {
let api_routes = api_routes
.merge(auth::routes())
.merge(profile_avatar::routes())
.merge(events::routes(state.clone()))
.merge(checkin::routes())
.merge(oauth::routes())
.merge(organizations::routes())
.merge(payment::routes())
.merge(payment_providers::routes())
.merge(payment_refunds::routes())
.merge(commerce::routes())
.merge(search::routes())
.merge(site_announcements::routes())
.merge(workflow_ui::routes())
@@ -82,8 +95,13 @@ pub fn build_router(state: AppState) -> Router {
.merge(tasks::routes())
.merge(user::routes())
.merge(user_api_keys::routes())
.merge(user_oauth::routes())
.merge(workspace::routes())
.merge(video_agent::routes())
.route_layer(middleware::from_fn_with_state(
state.clone(),
maintenance_guard,
))
.with_state(state);
let router = Router::new().merge(api_routes);
@@ -122,6 +140,28 @@ pub fn build_router(state: AppState) -> Router {
router.layer(CompressionLayer::new())
}
async fn maintenance_guard(
State(state): State<AppState>,
request: Request<Body>,
next: Next,
) -> Response {
let path = request.uri().path();
if !state.maintenance_mode()
|| path == "/healthz"
|| path.starts_with("/api/settings/database/")
{
return next.run(request).await;
}
(
StatusCode::SERVICE_UNAVAILABLE,
Json(json!({
"error": "database_maintenance",
"message": "数据库迁移维护中,请稍后重试"
})),
)
.into_response()
}
fn api_not_found_response(path: &str) -> Response {
Response::builder()
.status(StatusCode::NOT_FOUND)
+96 -44
View File
@@ -118,7 +118,7 @@ fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
}
async fn user_downstream_key(state: &AppState, user_id: i64) -> Result<DownstreamApiKey, String> {
sqlx::query_as::<_, DownstreamApiKey>(
crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT *
FROM downstream_api_keys
WHERE user_id = ?
@@ -126,7 +126,7 @@ async fn user_downstream_key(state: &AppState, user_id: i64) -> Result<Downstrea
AND (
expires_at IS NULL
OR expires_at = ''
OR datetime(expires_at) > datetime('now')
OR datetime(expires_at) > CURRENT_TIMESTAMP
)
ORDER BY updated_at DESC, created_at DESC, id DESC
LIMIT 1",
@@ -270,13 +270,14 @@ fn task_to_json(task: BackgroundTask) -> Value {
}
async fn load_job(state: &AppState, id: &str, auth: &MoneyPrinterAuth) -> AppResult<Value> {
let task: BackgroundTask =
sqlx::query_as("SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1")
.bind(id)
.bind(TASK_TYPE_MONEY_PRINTER)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("MoneyPrinter job not found".into()))?;
let task: BackgroundTask = crate::db::sql::query_as(
"SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1",
)
.bind(id)
.bind(TASK_TYPE_MONEY_PRINTER)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("MoneyPrinter job not found".into()))?;
if !task_visible_to_auth(&task, auth) {
return Err(AppError::NotFound("MoneyPrinter job not found".into()));
}
@@ -289,9 +290,9 @@ async fn update_progress(
progress: i64,
message: &str,
) -> AppResult<()> {
sqlx::query(
crate::db::sql::query(
"UPDATE background_tasks
SET progress = ?, message = ?, updated_at = datetime('now')
SET progress = ?, message = ?, updated_at = CURRENT_TIMESTAMP
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(progress)
@@ -304,7 +305,7 @@ async fn update_progress(
async fn is_cancelled(state: &AppState, id: &str) -> AppResult<bool> {
let status: Option<String> =
sqlx::query_scalar("SELECT status FROM background_tasks WHERE id = ? LIMIT 1")
crate::db::sql::query_scalar("SELECT status FROM background_tasks WHERE id = ? LIMIT 1")
.bind(id)
.fetch_optional(&*state.db)
.await?;
@@ -1818,6 +1819,18 @@ async fn create_job(
let topic = string_from(&body, &["topic", "prompt"])
.ok_or_else(|| AppError::BadRequest("topic is required".into()))?;
let task_id = format!("mp_{}", uuid::Uuid::new_v4().simple());
crate::services::comfy_billing::reserve_task(
&state.db,
if auth.role == "admin" {
None
} else {
auth.user_id
},
crate::services::comfy_billing::TASK_MONEY_PRINTER,
&task_id,
None,
)
.await?;
let title = topic.chars().take(40).collect::<String>();
let initial_result = json!({
"userId": auth.user_id,
@@ -1827,10 +1840,10 @@ async fn create_job(
"topic": topic,
"request": body,
});
sqlx::query(
let inserted = crate::db::sql::query(
"INSERT INTO background_tasks
(id, task_type, title, status, message, dedupe_key, result_json, progress, progress_total, started_at)
VALUES (?, ?, ?, 'running', '短视频任务已创建', ?, ?, 1, 5, datetime('now'))",
VALUES (?, ?, ?, 'running', '短视频任务已创建', ?, ?, 1, 5, CURRENT_TIMESTAMP)",
)
.bind(&task_id)
.bind(TASK_TYPE_MONEY_PRINTER)
@@ -1838,7 +1851,17 @@ async fn create_job(
.bind(format!("money-printer:{}", auth.user_id.map(|id| id.to_string()).unwrap_or_else(|| "admin".into())))
.bind(initial_result.to_string())
.execute(&*state.db)
.await?;
.await;
if let Err(err) = inserted {
let _ = crate::services::comfy_billing::release_task(
&state.db,
TASK_TYPE_MONEY_PRINTER,
&task_id,
Some("task_insert_failed"),
)
.await;
return Err(err.into());
}
let state_for_task = state.clone();
let task_id_for_spawn = task_id.clone();
@@ -1885,7 +1908,7 @@ async fn create_job(
"request": body,
"script": script,
});
sqlx::query("UPDATE background_tasks SET result_json = ?, updated_at = datetime('now') WHERE id = ?")
crate::db::sql::query("UPDATE background_tasks SET result_json = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?")
.bind(mid_result.to_string())
.bind(&task_id_for_spawn)
.execute(&*state_for_task.db)
@@ -1934,7 +1957,7 @@ async fn create_job(
"prompt": prompt,
"video": submitted,
});
sqlx::query("UPDATE background_tasks SET result_json = ?, updated_at = datetime('now') WHERE id = ? AND status IN ('pending', 'running')")
crate::db::sql::query("UPDATE background_tasks SET result_json = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ? AND status IN ('pending', 'running')")
.bind(submitted_result.to_string())
.bind(&task_id_for_spawn)
.execute(&*state_for_task.db)
@@ -2077,16 +2100,22 @@ async fn create_job(
"assets": assets,
"billing": billing,
});
sqlx::query(
crate::db::sql::query(
"UPDATE background_tasks
SET status = 'succeeded', message = '短视频生成完成', result_json = ?, progress = 5,
updated_at = datetime('now'), finished_at = datetime('now')
updated_at = CURRENT_TIMESTAMP, finished_at = CURRENT_TIMESTAMP
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(final_result.to_string())
.bind(&task_id_for_spawn)
.execute(&*state_for_task.db)
.await?;
crate::services::comfy_billing::settle_task(
&state_for_task.db,
TASK_TYPE_MONEY_PRINTER,
&task_id_for_spawn,
)
.await?;
Ok::<(), AppError>(())
}
.await;
@@ -2112,15 +2141,22 @@ async fn create_job(
);
}
}
let _ = sqlx::query(
let _ = crate::db::sql::query(
"UPDATE background_tasks
SET status = 'failed', message = '短视频生成失败', error = ?, updated_at = datetime('now'), finished_at = datetime('now')
SET status = 'failed', message = '短视频生成失败', error = ?, updated_at = CURRENT_TIMESTAMP, finished_at = CURRENT_TIMESTAMP
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(err.to_string())
.bind(&task_id_for_spawn)
.execute(&*state_for_task.db)
.await;
let _ = crate::services::comfy_billing::release_task(
&state_for_task.db,
TASK_TYPE_MONEY_PRINTER,
&task_id_for_spawn,
Some("failed"),
)
.await;
}
});
@@ -2145,7 +2181,7 @@ async fn list_jobs(
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let limit = query.limit.unwrap_or(20).clamp(1, 100);
let tasks: Vec<BackgroundTask> = sqlx::query_as(
let tasks: Vec<BackgroundTask> = crate::db::sql::query_as(
"SELECT * FROM background_tasks
WHERE task_type = ?
ORDER BY created_at DESC
@@ -2170,35 +2206,43 @@ async fn cancel_job(
Path(id): Path<String>,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let task: BackgroundTask =
sqlx::query_as("SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1")
.bind(&id)
.bind(TASK_TYPE_MONEY_PRINTER)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("MoneyPrinter job not found".into()))?;
let task: BackgroundTask = crate::db::sql::query_as(
"SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1",
)
.bind(&id)
.bind(TASK_TYPE_MONEY_PRINTER)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("MoneyPrinter job not found".into()))?;
if !task_visible_to_auth(&task, &auth) {
return Err(AppError::NotFound("MoneyPrinter job not found".into()));
}
if terminal_status(&task.status) {
return Ok(Json(json!({ "success": true, "job": task_to_json(task) })));
}
sqlx::query(
crate::db::sql::query(
"UPDATE background_tasks
SET status = 'cancelled', message = '已取消', updated_at = datetime('now'), finished_at = datetime('now')
SET status = 'cancelled', message = '已取消', updated_at = CURRENT_TIMESTAMP, finished_at = CURRENT_TIMESTAMP
WHERE id = ? AND task_type = ? AND status IN ('pending', 'running')",
)
.bind(&id)
.bind(TASK_TYPE_MONEY_PRINTER)
.execute(&*state.db)
.await?;
let _ = crate::services::comfy_billing::release_task(
&state.db,
TASK_TYPE_MONEY_PRINTER,
&id,
Some("cancelled"),
)
.await;
let job = load_job(&state, &id, &auth).await?;
Ok(Json(json!({ "success": true, "job": job })))
}
async fn clear_jobs(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let tasks: Vec<BackgroundTask> = sqlx::query_as(
let tasks: Vec<BackgroundTask> = crate::db::sql::query_as(
"SELECT * FROM background_tasks WHERE task_type = ? ORDER BY created_at DESC",
)
.bind(TASK_TYPE_MONEY_PRINTER)
@@ -2209,12 +2253,13 @@ async fn clear_jobs(State(state): State<AppState>, headers: HeaderMap) -> AppRes
if !task_visible_to_auth(&task, &auth) || !terminal_status(&task.status) {
continue;
}
let affected = sqlx::query("DELETE FROM background_tasks WHERE id = ? AND task_type = ?")
.bind(&task.id)
.bind(TASK_TYPE_MONEY_PRINTER)
.execute(&*state.db)
.await?
.rows_affected();
let affected =
crate::db::sql::query("DELETE FROM background_tasks WHERE id = ? AND task_type = ?")
.bind(&task.id)
.bind(TASK_TYPE_MONEY_PRINTER)
.execute(&*state.db)
.await?
.rows_affected();
cleared += i64::try_from(affected).unwrap_or(0);
}
Ok(Json(json!({ "success": true, "cleared": cleared })))
@@ -2231,14 +2276,17 @@ mod tests {
services::{builtin_keys::BuiltinKeys, runtime_auth::RuntimeAuthSettings},
};
use axum::http::HeaderValue;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
async fn test_state_with_comfy_backend(backend: ComfyUIBackend) -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("in-memory database");
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("in-memory database");
let comfyui = ComfyUIConfig {
enabled: true,
backend,
@@ -2259,6 +2307,8 @@ mod tests {
url: "sqlite::memory:".into(),
max_connections: 1,
ssl: false,
ssl_mode: None,
ca_cert_path: None,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
@@ -2277,9 +2327,11 @@ mod tests {
template_proxy_url: None,
},
comfyui,
agent: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
jimeng: Default::default(),
seedance: Default::default(),
},
pool,
BuiltinKeys {
+258 -214
View File
File diff suppressed because it is too large Load Diff
+109 -92
View File
@@ -36,13 +36,14 @@ fn extract_jwt_user(headers: &HeaderMap, state: &AppState) -> AppResult<(i64, St
}
/// Check if user is org admin (owner or admin role).
async fn require_org_admin(db: &sqlx::SqlitePool, org_id: i64, user_id: i64) -> AppResult<()> {
let member: Option<OrganizationMember> =
sqlx::query_as("SELECT * FROM organization_members WHERE org_id = ? AND user_id = ?")
.bind(org_id)
.bind(user_id)
.fetch_optional(db)
.await?;
async fn require_org_admin(db: &sqlx::AnyPool, org_id: i64, user_id: i64) -> AppResult<()> {
let member: Option<OrganizationMember> = crate::db::sql::query_as(
"SELECT * FROM organization_members WHERE org_id = ? AND user_id = ?",
)
.bind(org_id)
.bind(user_id)
.fetch_optional(db)
.await?;
match member {
Some(m) if m.role.as_deref() == Some("owner") || m.role.as_deref() == Some("admin") => {
@@ -53,13 +54,14 @@ async fn require_org_admin(db: &sqlx::SqlitePool, org_id: i64, user_id: i64) ->
}
/// Check if user is a member of an org.
async fn is_org_member(db: &sqlx::SqlitePool, org_id: i64, user_id: i64) -> Result<bool, AppError> {
let member: Option<OrganizationMember> =
sqlx::query_as("SELECT * FROM organization_members WHERE org_id = ? AND user_id = ?")
.bind(org_id)
.bind(user_id)
.fetch_optional(db)
.await?;
async fn is_org_member(db: &sqlx::AnyPool, org_id: i64, user_id: i64) -> Result<bool, AppError> {
let member: Option<OrganizationMember> = crate::db::sql::query_as(
"SELECT * FROM organization_members WHERE org_id = ? AND user_id = ?",
)
.bind(org_id)
.bind(user_id)
.fetch_optional(db)
.await?;
Ok(member.is_some())
}
@@ -113,7 +115,7 @@ async fn list_organizations(
let (user_id, _) = extract_jwt_user(&headers, &state)?;
let limit = q.limit.unwrap_or(50).clamp(1, 500);
let orgs: Vec<Organization> = sqlx::query_as(
let orgs: Vec<Organization> = crate::db::sql::query_as(
r#"SELECT o.* FROM organizations o
INNER JOIN organization_members om ON om.org_id = o.id
WHERE om.user_id = ?
@@ -145,30 +147,32 @@ async fn create_organization(
.and_then(|v| v.as_str())
.ok_or_else(|| AppError::BadRequest("slug is required".into()))?;
let result = sqlx::query("INSERT INTO organizations (name, slug) VALUES (?, ?)")
.bind(name)
.bind(slug)
.execute(&*state.db)
.await
.map_err(|e| {
let msg = e.to_string();
if msg.contains("UNIQUE") {
AppError::BadRequest("Organization slug already exists".into())
} else {
AppError::Internal(anyhow::anyhow!("Failed to create organization: {e}"))
}
})?;
let org_id = result.last_insert_rowid();
let org_id: i64 = crate::db::sql::insert_id(
"INSERT INTO organizations (name, slug) VALUES (?, ?) RETURNING id",
)
.bind(name)
.bind(slug)
.fetch_one(&*state.db)
.await
.map_err(|e| {
let msg = e.to_string();
if msg.contains("UNIQUE") {
AppError::BadRequest("Organization slug already exists".into())
} else {
AppError::Internal(anyhow::anyhow!("Failed to create organization: {e}"))
}
})?;
// Add creator as owner
sqlx::query("INSERT INTO organization_members (org_id, user_id, role) VALUES (?, ?, 'owner')")
.bind(org_id)
.bind(user_id)
.execute(&*state.db)
.await?;
crate::db::sql::query(
"INSERT INTO organization_members (org_id, user_id, role) VALUES (?, ?, 'owner')",
)
.bind(org_id)
.bind(user_id)
.execute(&*state.db)
.await?;
let org: Organization = sqlx::query_as("SELECT * FROM organizations WHERE id = ?")
let org: Organization = crate::db::sql::query_as("SELECT * FROM organizations WHERE id = ?")
.bind(org_id)
.fetch_one(&*state.db)
.await?;
@@ -219,10 +223,11 @@ async fn get_organization(
));
}
let org: Option<Organization> = sqlx::query_as("SELECT * FROM organizations WHERE id = ?")
.bind(org_id)
.fetch_optional(&*state.db)
.await?;
let org: Option<Organization> =
crate::db::sql::query_as("SELECT * FROM organizations WHERE id = ?")
.bind(org_id)
.fetch_optional(&*state.db)
.await?;
match org {
Some(o) => Ok(Json(json!({ "organization": o }))),
@@ -248,21 +253,25 @@ async fn update_organization(
}
if let Some(s) = slug {
sqlx::query("UPDATE organizations SET slug = ?, updated_at = datetime('now') WHERE id = ?")
.bind(s)
.bind(org_id)
.execute(&*state.db)
.await?;
crate::db::sql::query(
"UPDATE organizations SET slug = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?",
)
.bind(s)
.bind(org_id)
.execute(&*state.db)
.await?;
}
if let Some(n) = name {
sqlx::query("UPDATE organizations SET name = ?, updated_at = datetime('now') WHERE id = ?")
.bind(n)
.bind(org_id)
.execute(&*state.db)
.await?;
crate::db::sql::query(
"UPDATE organizations SET name = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?",
)
.bind(n)
.bind(org_id)
.execute(&*state.db)
.await?;
}
let org: Organization = sqlx::query_as("SELECT * FROM organizations WHERE id = ?")
let org: Organization = crate::db::sql::query_as("SELECT * FROM organizations WHERE id = ?")
.bind(org_id)
.fetch_one(&*state.db)
.await?;
@@ -286,7 +295,7 @@ async fn list_members(
));
}
let members: Vec<Value> = sqlx::query_as::<_, OrganizationMember>(
let members: Vec<Value> = crate::db::sql::query_as::<_, OrganizationMember>(
"SELECT * FROM organization_members WHERE org_id = ? ORDER BY joined_at ASC",
)
.bind(org_id)
@@ -329,12 +338,14 @@ async fn update_member_role(
));
}
sqlx::query("UPDATE organization_members SET role = ? WHERE org_id = ? AND user_id = ?")
.bind(role)
.bind(org_id)
.bind(target_user_id)
.execute(&*state.db)
.await?;
crate::db::sql::query(
"UPDATE organization_members SET role = ? WHERE org_id = ? AND user_id = ?",
)
.bind(role)
.bind(org_id)
.bind(target_user_id)
.execute(&*state.db)
.await?;
Ok(Json(json!({ "success": true })))
}
@@ -348,7 +359,7 @@ async fn remove_member(
let (user_id, _) = extract_jwt_user(&headers, &state)?;
require_org_admin(&state.db, org_id, user_id).await?;
sqlx::query("DELETE FROM organization_members WHERE org_id = ? AND user_id = ?")
crate::db::sql::query("DELETE FROM organization_members WHERE org_id = ? AND user_id = ?")
.bind(org_id)
.bind(target_user_id)
.execute(&*state.db)
@@ -368,11 +379,12 @@ async fn list_invitations(
let (user_id, _) = extract_jwt_user(&headers, &state)?;
require_org_admin(&state.db, org_id, user_id).await?;
let invitations: Vec<Invitation> =
sqlx::query_as("SELECT * FROM invitations WHERE org_id = ? ORDER BY created_at DESC")
.bind(org_id)
.fetch_all(&*state.db)
.await?;
let invitations: Vec<Invitation> = crate::db::sql::query_as(
"SELECT * FROM invitations WHERE org_id = ? ORDER BY created_at DESC",
)
.bind(org_id)
.fetch_all(&*state.db)
.await?;
Ok(Json(json!({ "invitations": invitations })))
}
@@ -410,8 +422,8 @@ async fn create_invitation(
.format("%Y-%m-%d %H:%M:%S")
.to_string();
let result = sqlx::query(
"INSERT INTO invitations (org_id, email, token, role, invited_by, expires_at) VALUES (?, ?, ?, ?, ?, ?)"
let invitation_id: i64 = crate::db::sql::insert_id(
"INSERT INTO invitations (org_id, email, token, role, invited_by, expires_at) VALUES (?, ?, ?, ?, ?, ?) RETURNING id"
)
.bind(org_id)
.bind(email)
@@ -419,11 +431,11 @@ async fn create_invitation(
.bind(role)
.bind(user_id)
.bind(&expires_at)
.execute(&*state.db)
.fetch_one(&*state.db)
.await?;
let inv: Invitation = sqlx::query_as("SELECT * FROM invitations WHERE id = ?")
.bind(result.last_insert_rowid())
let inv: Invitation = crate::db::sql::query_as("SELECT * FROM invitations WHERE id = ?")
.bind(invitation_id)
.fetch_one(&*state.db)
.await?;
@@ -439,11 +451,13 @@ async fn cancel_invitation(
let (user_id, _) = extract_jwt_user(&headers, &state)?;
require_org_admin(&state.db, org_id, user_id).await?;
sqlx::query("UPDATE invitations SET status = 'cancelled' WHERE id = ? AND org_id = ?")
.bind(inv_id)
.bind(org_id)
.execute(&*state.db)
.await?;
crate::db::sql::query(
"UPDATE invitations SET status = 'cancelled' WHERE id = ? AND org_id = ?",
)
.bind(inv_id)
.bind(org_id)
.execute(&*state.db)
.await?;
Ok(StatusCode::NO_CONTENT)
}
@@ -455,11 +469,12 @@ async fn get_invitation(
State(state): State<AppState>,
Path(token): Path<String>,
) -> AppResult<Json<Value>> {
let invitation: Option<Invitation> =
sqlx::query_as("SELECT * FROM invitations WHERE token = ? AND status = 'pending'")
.bind(&token)
.fetch_optional(&*state.db)
.await?;
let invitation: Option<Invitation> = crate::db::sql::query_as(
"SELECT * FROM invitations WHERE token = ? AND status = 'pending'",
)
.bind(&token)
.fetch_optional(&*state.db)
.await?;
match invitation {
Some(i) => {
@@ -476,7 +491,7 @@ async fn get_invitation(
// Get org name
let org_name: Option<(String,)> =
sqlx::query_as("SELECT name FROM organizations WHERE id = ?")
crate::db::sql::query_as("SELECT name FROM organizations WHERE id = ?")
.bind(i.org_id)
.fetch_optional(&*state.db)
.await?;
@@ -508,11 +523,12 @@ async fn accept_invitation(
) -> AppResult<Json<Value>> {
let (user_id, _) = extract_jwt_user(&headers, &state)?;
let invitation: Option<Invitation> =
sqlx::query_as("SELECT * FROM invitations WHERE token = ? AND status = 'pending'")
.bind(&token)
.fetch_optional(&*state.db)
.await?;
let invitation: Option<Invitation> = crate::db::sql::query_as(
"SELECT * FROM invitations WHERE token = ? AND status = 'pending'",
)
.bind(&token)
.fetch_optional(&*state.db)
.await?;
let inv = match invitation {
Some(i) => i,
@@ -537,14 +553,14 @@ async fn accept_invitation(
}
// Mark invitation as accepted
sqlx::query("UPDATE invitations SET status = 'accepted' WHERE id = ?")
crate::db::sql::query("UPDATE invitations SET status = 'accepted' WHERE id = ?")
.bind(inv.id)
.execute(&*state.db)
.await?;
// Add user as member
sqlx::query(
"INSERT OR IGNORE INTO organization_members (org_id, user_id, role) VALUES (?, ?, ?)",
crate::db::sql::query(
"INSERT INTO organization_members (org_id, user_id, role) VALUES (?, ?, ?) ON CONFLICT DO NOTHING",
)
.bind(inv.org_id)
.bind(user_id)
@@ -552,10 +568,11 @@ async fn accept_invitation(
.execute(&*state.db)
.await?;
let org: Option<Organization> = sqlx::query_as("SELECT * FROM organizations WHERE id = ?")
.bind(inv.org_id)
.fetch_optional(&*state.db)
.await?;
let org: Option<Organization> =
crate::db::sql::query_as("SELECT * FROM organizations WHERE id = ?")
.bind(inv.org_id)
.fetch_optional(&*state.db)
.await?;
Ok(Json(json!({ "organization": org })))
}
+103 -14
View File
@@ -6,7 +6,7 @@ use crate::{
};
use axum::{
body::Bytes,
extract::{Path, Query, State},
extract::{Form, Path, Query, State},
http::{HeaderMap, StatusCode, Uri},
response::{Html, IntoResponse, Response},
routing::{get, post, put},
@@ -70,6 +70,11 @@ pub fn routes() -> Router<AppState> {
"/api/user/payment/orders/{public_id}",
get(get_payment_order),
)
.route(
"/api/payment/alipay/orders/{public_id}/pay",
get(alipay_pay_page),
)
.route("/api/payment/alipay/notify", post(alipay_notify))
.route("/api/v1/payment/webhook/stripe", post(stripe_webhook))
}
@@ -138,6 +143,7 @@ async fn get_payment_config(State(state): State<AppState>) -> AppResult<Json<Val
payment::comfy_checkout_endpoint_source(&runtime_payment, &state.config.payment);
Ok(Json(json!({
"success": true,
"publicBaseUrl": runtime_payment.public_base_url.as_deref(),
"payment": payment::public_config(&runtime_payment),
"comfyCloud": {
"balanceDiagnosticsPath": "/api/admin/payment/comfy-cloud/balance",
@@ -165,10 +171,42 @@ async fn get_payment_config(State(state): State<AppState>) -> AppResult<Json<Val
},
"webhooks": {
"stripe": "/api/v1/payment/webhook/stripe",
"alipay": "/api/payment/alipay/notify",
},
})))
}
async fn alipay_pay_page(
State(state): State<AppState>,
Path(public_id): Path<String>,
) -> AppResult<Html<String>> {
let runtime_payment = payment::runtime_payment_config(&state.db, &state.config.payment).await?;
let html =
payment::alipay_checkout_form_html(&state.db, &runtime_payment, public_id.trim()).await?;
Ok(Html(html))
}
async fn alipay_notify(
State(state): State<AppState>,
Form(params): Form<std::collections::HashMap<String, String>>,
) -> Response {
let runtime_payment =
match payment::runtime_payment_config(&state.db, &state.config.payment).await {
Ok(config) => config,
Err(err) => {
tracing::warn!(%err, "failed to load runtime payment config for alipay notify");
return "fail".into_response();
}
};
match payment::handle_alipay_notify(&state.db, &runtime_payment, &params).await {
Ok(reply) => reply.into_response(),
Err(err) => {
tracing::warn!(%err, "alipay notify handling failed");
"fail".into_response()
}
}
}
async fn get_comfy_cloud_balance(
State(state): State<AppState>,
headers: HeaderMap,
@@ -425,6 +463,12 @@ fn oauth_callback_error_status_message(error: &AppError) -> (StatusCode, String)
AppError::RateLimitedWithRetry { message, .. } => {
(StatusCode::TOO_MANY_REQUESTS, message.clone())
}
AppError::RequestRateLimited { message, .. } => {
(StatusCode::TOO_MANY_REQUESTS, message.clone())
}
AppError::BillingRequired { message, .. } => {
(StatusCode::PAYMENT_REQUIRED, message.clone())
}
AppError::Internal(error) => (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()),
AppError::Database(error) => (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()),
}
@@ -527,8 +571,11 @@ async fn admin_list_recharge_records(
Query(query): Query<AdminRechargeListQuery>,
) -> AppResult<Json<Value>> {
require_admin(&headers, &state)?;
let result = payment::admin_list_user_recharge_records(&state.db, query.limit.unwrap_or(100)).await?;
Ok(Json(json!({ "success": true, "records": result.get("records").cloned().unwrap_or(json!([])) })))
let result =
payment::admin_list_user_recharge_records(&state.db, query.limit.unwrap_or(100)).await?;
Ok(Json(
json!({ "success": true, "records": result.get("records").cloned().unwrap_or(json!([])) }),
))
}
async fn admin_confirm_recharge(
@@ -537,7 +584,8 @@ async fn admin_confirm_recharge(
Path(public_id): Path<String>,
) -> AppResult<Json<Value>> {
require_admin(&headers, &state)?;
let record = payment::admin_confirm_user_recharge(&state.db, &state.config.payment, &public_id).await?;
let record =
payment::admin_confirm_user_recharge(&state.db, &state.config.payment, &public_id).await?;
Ok(Json(json!({ "success": true, "record": record })))
}
@@ -582,6 +630,7 @@ async fn admin_comfy_cloud_recharge(
.unwrap_or("admin")
.to_string(),
user_agent: header_value(&headers, "user-agent").map(str::to_string),
billing_target: None,
};
let session = payment::create_comfy_cloud_checkout(
&state.http,
@@ -867,13 +916,46 @@ async fn stripe_webhook(
};
match event_type.as_str() {
"checkout.session.completed" => {
let completion = payment::complete_stripe_order_by_session(
"checkout.session.completed" | "checkout.session.async_payment_succeeded" => {
let Some((payment_status, amount_total, currency)) =
payment::stripe_checkout_session_payment_details(&event)
else {
return Ok((
StatusCode::OK,
Json(json!({ "received": true, "ignored": true })),
));
};
if payment_status != "paid" {
return Ok((
StatusCode::OK,
Json(
json!({ "received": true, "ignored": true, "paymentStatus": payment_status }),
),
));
}
let completion = match payment::complete_stripe_order_by_session(
&state.db,
&state.config.payment,
&session_id,
payment_intent.as_deref(),
amount_total,
&currency,
)
.await?;
.await
{
Ok(completion) => completion,
Err(AppError::NotFound(_)) => {
return Ok((
StatusCode::OK,
Json(json!({
"received": true,
"ignored": true,
"reason": "unknown_order",
})),
));
}
Err(err) => return Err(err),
};
Ok((
StatusCode::OK,
Json(json!({
@@ -885,7 +967,7 @@ async fn stripe_webhook(
})),
))
}
"checkout.session.expired" => {
"checkout.session.expired" | "checkout.session.async_payment_failed" => {
let order =
payment::mark_stripe_order_failed_by_session(&state.db, &session_id).await?;
Ok((
@@ -920,16 +1002,19 @@ mod tests {
};
use hmac::{Hmac, Mac};
use sha2::Sha256;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
use std::time::{SystemTime, UNIX_EPOCH};
use tower::ServiceExt;
async fn test_state_with_payment(payment: PaymentConfig) -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool");
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool");
AppState::new(
AppConfig {
@@ -943,6 +1028,8 @@ mod tests {
url: "sqlite::memory:".to_string(),
max_connections: 1,
ssl: false,
ssl_mode: None,
ca_cert_path: None,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
@@ -961,8 +1048,10 @@ mod tests {
template_proxy_url: None,
},
comfyui: ComfyUIConfig::default(),
agent: Default::default(),
pixelle_video: Default::default(),
jimeng: Default::default(),
seedance: Default::default(),
payment,
},
pool,
+820
View File
@@ -0,0 +1,820 @@
use crate::{
auth::jwt::validate_jwt,
error::{AppError, AppResult},
services::payment_providers::{
self, live_adapters::adapter_for_instance, AddRouteMember, BindingSettlement,
CallbackRequest, CheckoutRequest, CreateMethodRoute, CreateProviderInstance, Currency,
PaymentQuery, ProviderKind, ProviderPaymentStatus, ProviderRepository,
UpdateProviderInstance,
},
AppState,
};
use axum::{
body::Bytes,
extract::{Path, Query, State},
http::{HeaderMap, StatusCode, Uri},
response::{IntoResponse, Response},
routing::{get, post, put},
Json, Router,
};
use serde::Deserialize;
use serde_json::{json, Value};
use sqlx::Row;
pub(super) fn routes() -> Router<AppState> {
Router::new()
.route("/api/user/payment/methods", get(list_payment_methods))
.route(
"/api/user/payment/package-orders/{public_id}/checkout",
post(create_package_checkout),
)
.route(
"/api/user/payment/package-orders/{public_id}/refresh",
post(refresh_package_payment),
)
.route(
"/api/payment/callback/{provider}/{callback_token}",
post(provider_callback),
)
}
pub(super) fn admin_routes() -> Router<AppState> {
Router::new()
.route(
"/api/admin/payment/providers",
get(list_provider_instances).post(create_provider_instance),
)
.route(
"/api/admin/payment/providers/{public_id}/credentials",
put(update_provider_credentials),
)
.route(
"/api/admin/payment/providers/{public_id}",
put(update_provider_instance),
)
.route(
"/api/admin/payment/providers/{public_id}/callback-token",
post(rotate_provider_callback_token),
)
.route(
"/api/admin/payment/provider-routes",
get(list_provider_routes).post(create_provider_route),
)
.route(
"/api/admin/payment/provider-routes/{public_id}/members",
post(add_provider_route_member),
)
}
fn repository(state: &AppState) -> AppResult<ProviderRepository> {
let key =
payment_providers::parse_master_key(state.config.payment.credential_master_key.as_deref())?;
ProviderRepository::new((*state.db).clone(), &key)
}
fn current_user_id(headers: &HeaderMap, state: &AppState) -> AppResult<i64> {
let token = headers
.get("Authorization")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.strip_prefix("Bearer "))
.ok_or_else(|| AppError::Unauthorized("Missing or invalid Authorization header".into()))?;
let claims = validate_jwt(&state.config.auth.jwt_secret, token)
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
if claims.role != "user" && claims.role != "admin" {
return Err(AppError::Forbidden("User account is required".into()));
}
claims
.sub
.parse()
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))
}
fn generated_callback_token() -> String {
format!(
"pcb_{}{}",
uuid::Uuid::new_v4().simple(),
uuid::Uuid::new_v4().simple()
)
}
async fn list_provider_instances(State(state): State<AppState>) -> AppResult<Json<Value>> {
let instances = repository(&state)?.list_instances().await?;
Ok(Json(json!({ "success": true, "instances": instances })))
}
async fn create_provider_instance(
State(state): State<AppState>,
Json(mut payload): Json<CreateProviderInstance>,
) -> AppResult<(StatusCode, Json<Value>)> {
if payload.provider == ProviderKind::ComfyCloud {
return Err(AppError::BadRequest(
"ComfyCloud uses the existing managed checkout integration and cannot be added as an aggregate provider instance".into(),
));
}
let generated = payload.callback_token.is_none();
let callback_token = payload
.callback_token
.clone()
.unwrap_or_else(generated_callback_token);
payload.callback_token = Some(callback_token.clone());
let instance = repository(&state)?.create_instance(payload).await?;
let callback_path = format!(
"/api/payment/callback/{}/{}",
instance.provider.as_str(),
callback_token
);
Ok((
StatusCode::CREATED,
Json(json!({
"success": true,
"instance": instance,
"callbackToken": generated.then_some(callback_token),
"callbackPath": callback_path,
})),
))
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct CredentialsPayload {
credentials: Value,
}
async fn update_provider_credentials(
State(state): State<AppState>,
Path(public_id): Path<String>,
Json(payload): Json<CredentialsPayload>,
) -> AppResult<Json<Value>> {
let repo = repository(&state)?;
repo.update_credentials(&public_id, payload.credentials)
.await?;
Ok(Json(json!({
"success": true,
"instance": repo.get_instance(&public_id).await?,
})))
}
async fn update_provider_instance(
State(state): State<AppState>,
Path(public_id): Path<String>,
Json(payload): Json<UpdateProviderInstance>,
) -> AppResult<Json<Value>> {
let instance = repository(&state)?
.update_instance(&public_id, payload)
.await?;
Ok(Json(json!({
"success": true,
"instance": instance,
})))
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct RotateCallbackTokenPayload {
callback_token: Option<String>,
}
async fn rotate_provider_callback_token(
State(state): State<AppState>,
Path(public_id): Path<String>,
Json(payload): Json<RotateCallbackTokenPayload>,
) -> AppResult<Json<Value>> {
let callback_token = payload
.callback_token
.unwrap_or_else(generated_callback_token);
let repo = repository(&state)?;
repo.rotate_callback_token(&public_id, Some(&callback_token))
.await?;
let instance = repo.get_instance(&public_id).await?;
Ok(Json(json!({
"success": true,
"callbackToken": callback_token,
"callbackPath": format!(
"/api/payment/callback/{}/{}",
instance.provider.as_str(),
callback_token
),
})))
}
async fn create_provider_route(
State(state): State<AppState>,
Json(payload): Json<CreateMethodRoute>,
) -> AppResult<(StatusCode, Json<Value>)> {
if payload.payment_method == payment_providers::PaymentMethod::ComfyCloud {
return Err(AppError::BadRequest(
"ComfyCloud routes remain managed by the legacy checkout integration".into(),
));
}
let route = repository(&state)?.create_method_route(payload).await?;
Ok((
StatusCode::CREATED,
Json(json!({ "success": true, "route": route })),
))
}
async fn list_provider_routes(State(state): State<AppState>) -> AppResult<Json<Value>> {
let rows = crate::db::sql::query(
"SELECT r.public_id, r.name, r.payment_method, r.currency,
r.selection_strategy, r.enabled,
COUNT(m.id) AS member_count
FROM payment_method_routes r
LEFT JOIN payment_method_route_members m ON m.route_id = r.id
GROUP BY r.id ORDER BY r.id",
)
.fetch_all(&*state.db)
.await?;
let routes: Vec<Value> = rows
.into_iter()
.map(|row| {
json!({
"publicId": row.try_get::<String, _>("public_id").unwrap_or_default(),
"name": row.try_get::<String, _>("name").unwrap_or_default(),
"paymentMethod": row.try_get::<String, _>("payment_method").unwrap_or_default(),
"currency": row.try_get::<String, _>("currency").unwrap_or_default(),
"selectionStrategy": row.try_get::<String, _>("selection_strategy").unwrap_or_default(),
"enabled": row.try_get::<bool, _>("enabled").unwrap_or(false),
"memberCount": row.try_get::<i64, _>("member_count").unwrap_or(0),
})
})
.collect();
Ok(Json(json!({ "success": true, "routes": routes })))
}
async fn add_provider_route_member(
State(state): State<AppState>,
Path(public_id): Path<String>,
Json(payload): Json<AddRouteMember>,
) -> AppResult<(StatusCode, Json<Value>)> {
let member_id = repository(&state)?
.add_route_member(&public_id, payload)
.await?;
Ok((
StatusCode::CREATED,
Json(json!({ "success": true, "memberId": member_id })),
))
}
#[derive(Debug, Deserialize)]
struct PaymentMethodsQuery {
currency: Option<String>,
}
async fn list_payment_methods(
State(state): State<AppState>,
headers: HeaderMap,
Query(query): Query<PaymentMethodsQuery>,
) -> AppResult<Json<Value>> {
let _ = current_user_id(&headers, &state)?;
let currency = query.currency.as_deref().map(Currency::parse).transpose()?;
let rows = crate::db::sql::query(
"SELECT DISTINCT r.public_id, r.name, r.payment_method, r.currency
FROM payment_method_routes r
JOIN payment_method_route_members m ON m.route_id = r.id AND m.enabled = 1
JOIN payment_provider_instances i ON i.id = m.provider_instance_id AND i.enabled = 1
WHERE r.enabled = 1 AND (? IS NULL OR r.currency = ?)
ORDER BY r.currency, r.payment_method, r.id",
)
.bind(currency.map(|value| value.as_str()))
.bind(currency.map(|value| value.as_str()))
.fetch_all(&*state.db)
.await?;
let methods: Vec<Value> = rows
.into_iter()
.map(|row| {
json!({
"routeId": row.try_get::<String, _>("public_id").unwrap_or_default(),
"name": row.try_get::<String, _>("name").unwrap_or_default(),
"paymentMethod": row.try_get::<String, _>("payment_method").unwrap_or_default(),
"currency": row.try_get::<String, _>("currency").unwrap_or_default(),
})
})
.collect();
Ok(Json(json!({ "success": true, "methods": methods })))
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct CreateCheckoutPayload {
route_id: String,
}
fn public_payment_base(state: &AppState) -> AppResult<String> {
let value = state
.config
.payment
.public_base_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| {
AppError::ServiceUnavailable(
"payment public_base_url is required before creating checkout".into(),
)
})?;
let parsed = url::Url::parse(value)
.map_err(|_| AppError::ServiceUnavailable("payment public_base_url is invalid".into()))?;
let local = parsed
.host_str()
.is_some_and(|host| host == "localhost" || host == "127.0.0.1" || host == "::1");
if parsed.scheme() != "https" && !(parsed.scheme() == "http" && local) {
return Err(AppError::ServiceUnavailable(
"payment public_base_url must use HTTPS".into(),
));
}
Ok(value.trim_end_matches('/').to_string())
}
async fn create_package_checkout(
State(state): State<AppState>,
headers: HeaderMap,
Path(public_id): Path<String>,
Json(payload): Json<CreateCheckoutPayload>,
) -> AppResult<Json<Value>> {
let user_id = current_user_id(&headers, &state)?;
let order = crate::services::commerce::get_order(&state.db, user_id, &public_id).await?;
if order.status != "pending" {
return Err(AppError::BadRequest(format!(
"package order cannot create checkout from status {}",
order.status
)));
}
if let Some(row) = crate::db::sql::query(
"SELECT po.checkout_url, po.provider_session_id, b.status, b.capacity_state
FROM payment_orders po
JOIN payment_order_provider_bindings b ON b.payment_order_id = po.id
WHERE po.id = ? AND po.checkout_url IS NOT NULL",
)
.bind(order.payment_order_id)
.fetch_optional(&*state.db)
.await?
{
let binding_status: String = row.try_get("status")?;
let capacity_state: String = row.try_get("capacity_state")?;
if !checkout_reservation_is_active(&binding_status, &capacity_state) {
return Err(AppError::BadRequest(
"payment checkout cannot be retried after its capacity reservation was released; create a new order"
.into(),
));
}
return Ok(Json(json!({
"success": true,
"order": order,
"checkoutUrl": row.try_get::<String, _>("checkout_url")?,
"providerSessionId": row.try_get::<Option<String>, _>("provider_session_id")?,
})));
}
let base = public_payment_base(&state)?;
let repo = repository(&state)?;
let binding = repo
.bind_order(order.payment_order_id, payload.route_id.trim())
.await?;
let instance = match repo
.load_runtime_instance(&binding.instance_public_id)
.await
{
Ok(instance) => instance,
Err(error) => {
return Err(release_failed_checkout(&repo, order.payment_order_id, error).await)
}
};
let adapter = match adapter_for_instance(&instance, state.http.clone()) {
Ok(adapter) => adapter,
Err(error) => {
return Err(release_failed_checkout(&repo, order.payment_order_id, error).await)
}
};
let notify_url = format!(
"{base}/api/payment/callback/{}/{}",
instance.provider.as_str(),
instance.callback_token
);
let payment_method = match binding.payment_method.as_str() {
"alipay" => payment_providers::PaymentMethod::Alipay,
"wechatpay" => payment_providers::PaymentMethod::Wechatpay,
"card" => payment_providers::PaymentMethod::Card,
"comfy_cloud" => payment_providers::PaymentMethod::ComfyCloud,
_ => {
let error =
AppError::ServiceUnavailable("payment binding has an invalid method".into());
return Err(release_failed_checkout(&repo, order.payment_order_id, error).await);
}
};
let checkout = match adapter
.create_checkout(
&instance.credentials,
CheckoutRequest {
order_id: order.payment_order_public_id.clone(),
amount_minor: order.amount_minor,
currency: Currency::parse(&order.currency)?,
payment_method,
return_url: format!(
"{base}/user/recharge?payment=return&order={}",
order.public_id
),
notify_url,
description: order.product_name.clone(),
expires_at: order.expires_at.clone(),
},
)
.await
{
Ok(checkout) => checkout,
Err(error) => {
return Err(release_failed_checkout(&repo, order.payment_order_id, error).await)
}
};
let mut tx = state.db.begin().await?;
let binding_update = crate::db::sql::query(
"UPDATE payment_order_provider_bindings
SET provider_session_id = ?, status = 'processing',
updated_at = CURRENT_TIMESTAMP
WHERE payment_order_id = ?
AND status IN ('bound', 'processing')
AND capacity_state = 'reserved'",
)
.bind(&checkout.provider_session_id)
.bind(order.payment_order_id)
.execute(&mut *tx)
.await?;
let order_update = crate::db::sql::query(
"UPDATE payment_orders
SET provider_session_id = ?, checkout_url = ?,
updated_at = CURRENT_TIMESTAMP
WHERE id = ? AND status = 'pending'
AND expires_at IS NOT NULL
AND datetime(expires_at) > CURRENT_TIMESTAMP",
)
.bind(&checkout.provider_session_id)
.bind(&checkout.checkout_url)
.bind(order.payment_order_id)
.execute(&mut *tx)
.await?;
if binding_update.rows_affected() != 1 || order_update.rows_affected() != 1 {
tx.rollback().await?;
repo.release_expired_reservations().await?;
return Err(AppError::BadRequest(
"payment order expired before checkout creation completed; create a new order".into(),
));
}
tx.commit().await?;
Ok(Json(json!({
"success": true,
"order": order,
"binding": binding,
"providerSessionId": checkout.provider_session_id,
"checkoutUrl": checkout.checkout_url,
})))
}
fn checkout_reservation_is_active(status: &str, capacity_state: &str) -> bool {
capacity_state == "reserved" && matches!(status, "bound" | "processing")
}
async fn release_failed_checkout(
repo: &ProviderRepository,
payment_order_id: i64,
original_error: AppError,
) -> AppError {
match repo
.settle_binding(payment_order_id, BindingSettlement::Failed)
.await
{
Ok(_) => original_error,
Err(release_error) => {
tracing::error!(
payment_order_id,
error = %release_error,
original_error = %original_error,
"failed to release payment checkout capacity reservation"
);
release_error
}
}
}
async fn refresh_package_payment(
State(state): State<AppState>,
headers: HeaderMap,
Path(public_id): Path<String>,
) -> AppResult<Json<Value>> {
let user_id = current_user_id(&headers, &state)?;
let order = crate::services::commerce::get_order(&state.db, user_id, &public_id).await?;
if order.status == "expired" {
let released = repository(&state)?.release_expired_reservations().await?;
return Ok(Json(json!({
"success": true,
"order": order,
"releasedReservations": released,
})));
}
if matches!(
order.status.as_str(),
"fulfilled" | "partially_refunded" | "refunded"
) {
return Ok(Json(json!({ "success": true, "order": order })));
}
let repo = repository(&state)?;
let binding = repo.get_binding_by_order(order.payment_order_id).await?;
let row = crate::db::sql::query(
"SELECT provider_session_id, provider_payment_id
FROM payment_order_provider_bindings WHERE payment_order_id = ?",
)
.bind(order.payment_order_id)
.fetch_one(&*state.db)
.await?;
let instance = repo
.load_runtime_instance(&binding.instance_public_id)
.await?;
let adapter = adapter_for_instance(&instance, state.http.clone())?;
let result = adapter
.query_payment(
&instance.credentials,
PaymentQuery {
provider_session_id: row.try_get("provider_session_id")?,
provider_payment_id: row.try_get("provider_payment_id")?,
},
)
.await?;
if matches!(&result.status, ProviderPaymentStatus::Paid) {
verify_provider_amount(
binding.amount_minor,
&binding.currency,
result.amount_minor,
result.currency,
)?;
repo.assign_provider_payment_id(
order.payment_order_id,
result.provider_payment_id.as_deref(),
)
.await?;
repo.settle_binding(order.payment_order_id, BindingSettlement::Paid)
.await?;
crate::services::commerce::mark_order_paid(&state.db, &public_id).await?;
let fulfillment = crate::services::commerce::fulfill_order(&state.db, &public_id).await?;
return Ok(Json(json!({
"success": true,
"order": fulfillment.order,
"fulfillment": fulfillment,
})));
}
if let Some(settlement) = terminal_query_settlement(&result.status) {
let binding = repo
.settle_binding(order.payment_order_id, settlement)
.await?;
return Ok(Json(json!({
"success": true,
"order": order,
"binding": binding,
"providerStatus": result.status,
})));
}
Ok(Json(json!({
"success": true,
"order": order,
"providerStatus": result.status,
})))
}
fn terminal_query_settlement(status: &ProviderPaymentStatus) -> Option<BindingSettlement> {
match status {
ProviderPaymentStatus::Failed => Some(BindingSettlement::Failed),
ProviderPaymentStatus::Cancelled => Some(BindingSettlement::Cancelled),
_ => None,
}
}
fn verify_provider_amount(
expected_minor: i64,
expected_currency: &str,
actual_minor: Option<i64>,
actual_currency: Option<Currency>,
) -> AppResult<()> {
if actual_minor != Some(expected_minor)
|| actual_currency.map(|value| value.as_str()) != Some(expected_currency)
{
return Err(AppError::BadRequest(
"provider payment amount or currency does not match the order snapshot".into(),
));
}
Ok(())
}
async fn provider_callback(
State(state): State<AppState>,
Path((provider, callback_token)): Path<(String, String)>,
uri: Uri,
headers: HeaderMap,
body: Bytes,
) -> Response {
match process_provider_callback(&state, &provider, &callback_token, &uri, &headers, body).await
{
Ok(()) => "success".into_response(),
Err(error) => {
tracing::warn!(%error, provider, "payment provider callback rejected");
(app_error_status(&error), "fail").into_response()
}
}
}
fn app_error_status(error: &AppError) -> StatusCode {
match error {
AppError::NotFound(_) => StatusCode::NOT_FOUND,
AppError::Unauthorized(_) => StatusCode::UNAUTHORIZED,
AppError::Forbidden(_) => StatusCode::FORBIDDEN,
AppError::BadRequest(_) => StatusCode::BAD_REQUEST,
AppError::RateLimited
| AppError::RateLimitedWithRetry { .. }
| AppError::RequestRateLimited { .. } => StatusCode::TOO_MANY_REQUESTS,
AppError::BillingRequired { .. } => StatusCode::PAYMENT_REQUIRED,
AppError::ServiceUnavailable(_) => StatusCode::SERVICE_UNAVAILABLE,
AppError::GatewayTimeout(_) => StatusCode::GATEWAY_TIMEOUT,
AppError::UpstreamError(_) => StatusCode::BAD_GATEWAY,
AppError::Internal(_) | AppError::Database(_) => StatusCode::INTERNAL_SERVER_ERROR,
}
}
async fn process_provider_callback(
state: &AppState,
provider: &str,
callback_token: &str,
uri: &Uri,
headers: &HeaderMap,
body: Bytes,
) -> AppResult<()> {
let provider = ProviderKind::parse(provider)?;
let repo = repository(state)?;
let summary = repo
.find_callback_instance(provider, callback_token)
.await?;
let instance = repo.load_runtime_instance(&summary.public_id).await?;
let adapter = adapter_for_instance(&instance, state.http.clone())?;
let query = uri
.query()
.map(|query| {
url::form_urlencoded::parse(query.as_bytes())
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect()
})
.unwrap_or_default();
let callback = CallbackRequest {
headers: headers
.iter()
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.as_str().to_string(), value.to_string()))
})
.collect(),
query,
body: body.to_vec(),
};
let event = adapter
.verify_callback(&instance.credentials, callback)
.await?;
let receipt = repo
.ingest_event(
&instance.public_id,
&event.provider_event_id,
&event.event_type,
event.payload.clone(),
)
.await?;
if !receipt.inserted && receipt.status == "processed" {
return Ok(());
}
match event.event_type.as_str() {
"payment.refunded" | "refund.failed" | "refund.pending" => {
let refund_public_id = event.merchant_refund_id.as_deref().ok_or_else(|| {
AppError::BadRequest("provider refund callback has no merchant refund ID".into())
})?;
crate::services::payment_refunds::apply_provider_event(
&state.db,
refund_public_id,
&instance.public_id,
event.merchant_order_id.as_deref(),
&event.event_type,
event.provider_payment_id.as_deref(),
event.amount_minor,
event.currency.map(|currency| currency.as_str()),
event.payload.clone(),
)
.await?;
}
"payment.paid" | "payment.failed" | "payment.cancelled" | "payment.retryable_failed" => {
let order_public_id = event.merchant_order_id.as_deref().ok_or_else(|| {
AppError::BadRequest("provider payment callback has no merchant order ID".into())
})?;
let binding = crate::db::sql::query(
"SELECT b.payment_order_id, b.amount_minor, b.currency, b.status
FROM payment_order_provider_bindings b
JOIN payment_orders p ON p.id = b.payment_order_id
JOIN payment_provider_instances i ON i.id = b.provider_instance_id
WHERE p.public_id = ? AND i.public_id = ?",
)
.bind(order_public_id)
.bind(&instance.public_id)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("callback payment order binding not found".into()))?;
let payment_order_id: i64 = binding.try_get("payment_order_id")?;
if event.event_type == "payment.retryable_failed" {
crate::db::sql::query(
"UPDATE payment_provider_event_inbox
SET status = 'processed', processed_at = CURRENT_TIMESTAMP
WHERE public_id = ?",
)
.bind(receipt.public_id)
.execute(&*state.db)
.await?;
return Ok(());
}
if matches!(
event.event_type.as_str(),
"payment.failed" | "payment.cancelled"
) {
let status: String = binding.try_get("status")?;
if matches!(status.as_str(), "bound" | "processing") {
let settlement = if event.event_type == "payment.cancelled" {
BindingSettlement::Cancelled
} else {
BindingSettlement::Failed
};
repo.settle_binding(payment_order_id, settlement).await?;
}
crate::db::sql::query(
"UPDATE payment_provider_event_inbox
SET status = 'processed', processed_at = CURRENT_TIMESTAMP
WHERE public_id = ?",
)
.bind(receipt.public_id)
.execute(&*state.db)
.await?;
return Ok(());
}
let amount_minor: i64 = binding.try_get("amount_minor")?;
let currency: String = binding.try_get("currency")?;
verify_provider_amount(amount_minor, &currency, event.amount_minor, event.currency)?;
repo.assign_provider_payment_id(payment_order_id, event.provider_payment_id.as_deref())
.await?;
repo.settle_binding(payment_order_id, BindingSettlement::Paid)
.await?;
let commerce_order: Option<String> = crate::db::sql::query_scalar(
"SELECT public_id FROM commerce_orders WHERE payment_order_id = ?",
)
.bind(payment_order_id)
.fetch_optional(&*state.db)
.await?;
if let Some(commerce_order) = commerce_order {
crate::services::commerce::mark_order_paid(&state.db, &commerce_order).await?;
crate::services::commerce::fulfill_order(&state.db, &commerce_order).await?;
}
}
_ => {}
}
crate::db::sql::query(
"UPDATE payment_provider_event_inbox
SET status = 'processed', processed_at = CURRENT_TIMESTAMP
WHERE public_id = ?",
)
.bind(receipt.public_id)
.execute(&*state.db)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn refresh_releases_only_terminal_failed_or_cancelled_payments() {
assert_eq!(
terminal_query_settlement(&ProviderPaymentStatus::Failed),
Some(BindingSettlement::Failed)
);
assert_eq!(
terminal_query_settlement(&ProviderPaymentStatus::Cancelled),
Some(BindingSettlement::Cancelled)
);
assert_eq!(
terminal_query_settlement(&ProviderPaymentStatus::Pending),
None
);
assert_eq!(
terminal_query_settlement(&ProviderPaymentStatus::Paid),
None
);
}
#[test]
fn checkout_retry_requires_an_active_capacity_reservation() {
assert!(checkout_reservation_is_active("bound", "reserved"));
assert!(checkout_reservation_is_active("processing", "reserved"));
assert!(!checkout_reservation_is_active("failed", "released"));
assert!(!checkout_reservation_is_active("cancelled", "released"));
assert!(!checkout_reservation_is_active("processing", "released"));
}
}
+273
View File
@@ -0,0 +1,273 @@
use crate::{
auth::jwt::validate_jwt,
error::{AppError, AppResult},
services::payment_refunds,
AppState,
};
use axum::{
extract::{Path, State},
http::{HeaderMap, StatusCode},
routing::{get, post},
Json, Router,
};
use serde::Deserialize;
use serde_json::{json, Value};
pub(super) fn routes() -> Router<AppState> {
Router::new()
.route("/api/user/payment/refunds", get(list_user_refunds))
.route(
"/api/user/payment/package-orders/{public_id}/refunds",
post(create_user_refund),
)
}
pub(super) fn admin_routes() -> Router<AppState> {
Router::new()
.route("/api/admin/payment/refunds", get(list_admin_refunds))
.route(
"/api/admin/payment/refunds/{public_id}/approve",
post(approve_refund),
)
.route(
"/api/admin/payment/refunds/{public_id}/reject",
post(reject_refund),
)
.route(
"/api/admin/payment/refunds/{public_id}/retry",
post(retry_refund),
)
.route(
"/api/admin/payment/refunds/{public_id}/mark-succeeded",
post(mark_refund_succeeded),
)
.route(
"/api/admin/payment/refunds/{public_id}/release-failed",
post(release_failed_refund),
)
}
fn bearer_token<'a>(headers: &'a HeaderMap, message: &str) -> AppResult<&'a str> {
headers
.get("Authorization")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.strip_prefix("Bearer "))
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| AppError::Unauthorized(message.into()))
}
fn current_user_id(headers: &HeaderMap, state: &AppState) -> AppResult<i64> {
let claims = validate_jwt(
&state.config.auth.jwt_secret,
bearer_token(headers, "Missing or invalid Authorization header")?,
)
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
if claims.role != "user" {
return Err(AppError::Forbidden("User account is required".into()));
}
claims
.sub
.parse()
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))
}
fn reviewer(headers: &HeaderMap, state: &AppState) -> String {
let Ok(token) = bearer_token(headers, "Missing admin token") else {
return "admin".into();
};
if token == state.config.auth.auth_token {
return "admin:configured-token".into();
}
validate_jwt(&state.config.auth.jwt_secret, token)
.ok()
.map(|claims| format!("admin:{}", claims.sub))
.unwrap_or_else(|| "admin".into())
}
fn idempotency_key(headers: &HeaderMap) -> AppResult<&str> {
headers
.get("Idempotency-Key")
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| AppError::BadRequest("Idempotency-Key header is required".into()))
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct CreateRefundPayload {
amount_minor: Option<i64>,
reason: Option<String>,
}
async fn create_user_refund(
State(state): State<AppState>,
headers: HeaderMap,
Path(public_id): Path<String>,
Json(payload): Json<CreateRefundPayload>,
) -> AppResult<(StatusCode, Json<Value>)> {
let user_id = current_user_id(&headers, &state)?;
let refund = payment_refunds::create_user_request(
&state.db,
&state.config.payment,
user_id,
&public_id,
payload.amount_minor,
payload.reason.as_deref(),
idempotency_key(&headers)?,
)
.await?;
Ok((
StatusCode::CREATED,
Json(json!({ "success": true, "refund": refund })),
))
}
async fn list_user_refunds(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Json<Value>> {
let user_id = current_user_id(&headers, &state)?;
let refunds = payment_refunds::list_user(&state.db, user_id).await?;
Ok(Json(json!({ "success": true, "refunds": refunds })))
}
async fn list_admin_refunds(State(state): State<AppState>) -> AppResult<Json<Value>> {
let refunds = payment_refunds::list_admin(&state.db).await?;
Ok(Json(json!({ "success": true, "refunds": refunds })))
}
#[derive(Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ApproveRefundPayload {
#[serde(default)]
force: bool,
}
async fn approve_refund(
State(state): State<AppState>,
headers: HeaderMap,
Path(public_id): Path<String>,
Json(payload): Json<ApproveRefundPayload>,
) -> AppResult<Json<Value>> {
let refund = payment_refunds::approve_and_process(
&state.db,
state.http.clone(),
&state.config.payment,
&public_id,
&reviewer(&headers, &state),
payload.force,
)
.await?;
Ok(Json(json!({ "success": true, "refund": refund })))
}
#[derive(Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct RejectRefundPayload {
reason: Option<String>,
}
async fn reject_refund(
State(state): State<AppState>,
headers: HeaderMap,
Path(public_id): Path<String>,
Json(payload): Json<RejectRefundPayload>,
) -> AppResult<Json<Value>> {
let refund = payment_refunds::reject(
&state.db,
&public_id,
&reviewer(&headers, &state),
payload.reason.as_deref(),
)
.await?;
Ok(Json(json!({ "success": true, "refund": refund })))
}
async fn retry_refund(
State(state): State<AppState>,
headers: HeaderMap,
Path(public_id): Path<String>,
) -> AppResult<Json<Value>> {
let current = payment_refunds::get(&state.db, &public_id).await?;
if current.status != "processing" {
return Err(AppError::BadRequest(
"only processing refunds can be retried".into(),
));
}
crate::services::payment_jobs::revive_refund_job(&state.db, &public_id).await?;
let refund = payment_refunds::approve_and_process(
&state.db,
state.http.clone(),
&state.config.payment,
&public_id,
&reviewer(&headers, &state),
current.force_reversal,
)
.await?;
Ok(Json(json!({ "success": true, "refund": refund })))
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct MarkSucceededPayload {
confirmed_externally: bool,
provider_refund_id: String,
}
async fn mark_refund_succeeded(
State(state): State<AppState>,
headers: HeaderMap,
Path(public_id): Path<String>,
Json(payload): Json<MarkSucceededPayload>,
) -> AppResult<Json<Value>> {
if !payload.confirmed_externally {
return Err(AppError::BadRequest(
"confirmedExternally must be true after verifying the provider refund".into(),
));
}
let provider_refund_id = payload.provider_refund_id.trim();
if provider_refund_id.is_empty() || provider_refund_id.len() > 255 {
return Err(AppError::BadRequest(
"providerRefundId must contain 1 to 255 bytes".into(),
));
}
let refund = payment_refunds::mark_succeeded(
&state.db,
&public_id,
provider_refund_id,
&reviewer(&headers, &state),
)
.await?;
Ok(Json(json!({ "success": true, "refund": refund })))
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ReleaseFailedPayload {
confirmed_provider_failed: bool,
#[serde(default)]
reason: Option<String>,
}
async fn release_failed_refund(
State(state): State<AppState>,
headers: HeaderMap,
Path(public_id): Path<String>,
Json(payload): Json<ReleaseFailedPayload>,
) -> AppResult<Json<Value>> {
if !payload.confirmed_provider_failed {
return Err(AppError::BadRequest(
"confirmedProviderFailed must be true before restoring entitlements".into(),
));
}
let refund = payment_refunds::release_failed(
&state.db,
&public_id,
&reviewer(&headers, &state),
payload.reason.as_deref(),
)
.await?;
Ok(Json(json!({ "success": true, "refund": refund })))
}
+252 -54
View File
@@ -3,6 +3,9 @@ use crate::{
db::models::background_task::BackgroundTask,
db::models::downstream_api_key::DownstreamApiKey,
error::{AppError, AppResult},
services::workspace::{
record_base64_asset, record_remote_asset, StoredWorkspaceAsset, WorkspaceAssetInput,
},
AppState,
};
use axum::{
@@ -97,7 +100,7 @@ fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
}
async fn user_downstream_key(state: &AppState, user_id: i64) -> Result<DownstreamApiKey, String> {
sqlx::query_as::<_, DownstreamApiKey>(
crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT *
FROM downstream_api_keys
WHERE user_id = ?
@@ -105,7 +108,7 @@ async fn user_downstream_key(state: &AppState, user_id: i64) -> Result<Downstrea
AND (
expires_at IS NULL
OR expires_at = ''
OR datetime(expires_at) > datetime('now')
OR datetime(expires_at) > CURRENT_TIMESTAMP
)
ORDER BY updated_at DESC, created_at DESC, id DESC
LIMIT 1",
@@ -240,11 +243,11 @@ async fn models_handler(
) -> AppResult<Json<Value>> {
let _auth = resolve_picture_agent_auth(&headers, &state).await?;
let rows: Vec<String> = sqlx::query_scalar(
let rows: Vec<String> = crate::db::sql::query_scalar(
"SELECT model_pattern
FROM token_routes
WHERE COALESCE(enabled, 1) = 1
ORDER BY model_pattern COLLATE NOCASE",
ORDER BY LOWER(model_pattern), model_pattern",
)
.fetch_all(&*state.db)
.await?;
@@ -273,6 +276,26 @@ async fn models_handler(
mod tests {
use super::*;
fn picture_auth(role: &str, user_id: Option<i64>) -> PictureAgentAuth {
PictureAgentAuth {
user_id,
role: role.to_string(),
outbound_key: "test-key".to_string(),
key_id: None,
key_name: "test-key".to_string(),
}
}
fn stored_image(status: &str, has_local_file: bool) -> StoredWorkspaceAsset {
StoredWorkspaceAsset {
public_id: "wa_picture".to_string(),
file_name: Some("wa_picture.png".to_string()),
file_path: has_local_file.then(|| "/tmp/wa_picture.png".to_string()),
remote_url: Some("https://upstream.example/temporary.png".to_string()),
status: status.to_string(),
}
}
#[test]
fn picture_model_filter_keeps_image_models() {
assert!(is_picture_model_pattern("gpt-image-2"));
@@ -291,6 +314,68 @@ mod tests {
assert!(!is_picture_model_pattern("deepseek-chat"));
assert!(!is_picture_model_pattern("qwen-plus"));
}
#[test]
fn archived_user_image_keeps_display_data_and_adds_workspace_reference() {
let original = json!({
"b64_json": "temporary-base64",
"revised_prompt": "refined prompt",
});
let asset = stored_image("completed", true);
assert_eq!(
prefer_archived_image_item(&original, &picture_auth("user", Some(42)), Some(&asset),),
json!({
"b64_json": "temporary-base64",
"workspaceAssetId": "wa_picture",
"workspaceFileUrl": "/api/user/workspace/assets/wa_picture/file",
"revised_prompt": "refined prompt",
})
);
}
#[test]
fn archived_admin_image_uses_admin_workspace_url() {
let original = json!({
"url": "https://upstream.example/temporary.png",
"revised_prompt": "refined prompt",
});
let asset = stored_image("completed", true);
let result =
prefer_archived_image_item(&original, &picture_auth("admin", Some(7)), Some(&asset));
assert_eq!(result.get("url"), original.get("url"));
assert_eq!(
result.get("workspaceFileUrl").and_then(Value::as_str),
Some("/api/workspace/assets/wa_picture/file")
);
assert_eq!(
result.get("workspaceAssetId").and_then(Value::as_str),
Some("wa_picture")
);
}
#[test]
fn failed_or_incomplete_archive_preserves_original_image_item() {
let original = json!({
"url": "https://upstream.example/temporary.png",
"revised_prompt": "refined prompt",
});
let failed = stored_image("failed", false);
let missing_file = stored_image("completed", false);
let auth = picture_auth("user", Some(42));
assert_eq!(
prefer_archived_image_item(&original, &auth, Some(&failed)),
original
);
assert_eq!(
prefer_archived_image_item(&original, &auth, Some(&missing_file)),
original
);
assert_eq!(prefer_archived_image_item(&original, &auth, None), original);
}
}
fn prompt_title(prompt: &str, post_number: Option<i64>) -> String {
@@ -422,22 +507,136 @@ fn picture_generation_client_with_timeout(timeout_secs: u64) -> Client {
.unwrap_or_else(|_| Client::new())
}
enum PictureImageSource<'a> {
Remote(&'a str),
Base64(&'a str),
}
fn picture_image_source(item: &Value) -> Option<PictureImageSource<'_>> {
let value = item
.get("url")
.and_then(Value::as_str)
.or_else(|| item.as_str())
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(value) = value {
return if value.starts_with("data:") {
Some(PictureImageSource::Base64(value))
} else {
Some(PictureImageSource::Remote(value))
};
}
item.get("b64_json")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(PictureImageSource::Base64)
}
fn workspace_file_url(auth: &PictureAgentAuth, public_id: &str) -> String {
if auth.role == "user" {
format!("/api/user/workspace/assets/{public_id}/file")
} else {
format!("/api/workspace/assets/{public_id}/file")
}
}
fn prefer_archived_image_item(
original: &Value,
auth: &PictureAgentAuth,
archived: Option<&StoredWorkspaceAsset>,
) -> Value {
let Some(asset) = archived.filter(|asset| {
asset.status == "completed"
&& asset
.file_path
.as_deref()
.is_some_and(|path| !path.trim().is_empty())
}) else {
return original.clone();
};
let mut item = match original {
Value::Object(item) => item.clone(),
Value::String(url) => {
let mut item = serde_json::Map::new();
item.insert("url".to_string(), Value::String(url.clone()));
item
}
_ => return original.clone(),
};
item.insert(
"workspaceAssetId".to_string(),
Value::String(asset.public_id.clone()),
);
item.insert(
"workspaceFileUrl".to_string(),
Value::String(workspace_file_url(auth, &asset.public_id)),
);
Value::Object(item)
}
async fn archive_picture_agent_response(
state: &AppState,
auth: &PictureAgentAuth,
prompt: Option<String>,
model: Option<String>,
payload: &Value,
) {
crate::services::workspace::record_image_assets_from_response(
state,
auth.user_id,
"picture-agent",
prompt,
model,
payload,
)
.await;
) -> Vec<Value> {
let Some(items) = payload.get("data").and_then(Value::as_array) else {
return Vec::new();
};
let mut images = Vec::with_capacity(items.len());
for (index, item) in items.iter().enumerate() {
let Some(source) = picture_image_source(item) else {
images.push(item.clone());
continue;
};
let title = prompt
.as_deref()
.map(|value| value.chars().take(36).collect::<String>())
.filter(|value| !value.trim().is_empty())
.or_else(|| Some(format!("图片产出 #{}", index + 1)));
let input = WorkspaceAssetInput {
user_id: auth.user_id,
source: "picture-agent".to_string(),
kind: "image".to_string(),
title,
prompt: prompt.clone(),
model: model.clone(),
remote_url: match source {
PictureImageSource::Remote(url) => Some(url.to_string()),
PictureImageSource::Base64(_) => None,
},
metadata: json!({
"workspaceAsset": { "role": "output" },
"responseItem": item,
}),
};
let stored = match source {
PictureImageSource::Remote(_) => record_remote_asset(state, input).await,
PictureImageSource::Base64(value) => record_base64_asset(state, input, value).await,
};
match stored {
Ok(asset) => {
if asset.status != "completed" || asset.file_path.is_none() {
tracing::warn!(
public_id = %asset.public_id,
status = %asset.status,
"picture image archive did not produce a local file"
);
}
images.push(prefer_archived_image_item(item, auth, Some(&asset)));
}
Err(err) => {
tracing::warn!(%err, "failed to archive picture image; preserving upstream item");
images.push(item.clone());
}
}
}
images
}
async fn execute_generation_payload(
@@ -528,19 +727,21 @@ async fn execute_generation_payload_with_timeout(
}
async fn mark_stale_picture_jobs_failed(state: &AppState) -> AppResult<()> {
sqlx::query(
let stale_before =
(chrono::Utc::now() - chrono::Duration::minutes(STALE_RUNNING_JOB_MINUTES)).to_rfc3339();
crate::db::sql::query(
"UPDATE background_tasks
SET status = 'failed',
message = '生图任务已中断',
error = '后台任务长时间没有进展,可能是后端重启、上游超时或图片通道无响应。请重新提交。',
updated_at = datetime('now'),
finished_at = datetime('now')
updated_at = CURRENT_TIMESTAMP,
finished_at = CURRENT_TIMESTAMP
WHERE task_type = ?
AND status IN ('pending', 'running')
AND datetime(COALESCE(updated_at, created_at)) < datetime('now', ?)",
AND COALESCE(updated_at, created_at) < ?",
)
.bind(TASK_TYPE_PICTURE_GENERATION)
.bind(format!("-{STALE_RUNNING_JOB_MINUTES} minutes"))
.bind(stale_before)
.execute(&*state.db)
.await?;
Ok(())
@@ -588,13 +789,14 @@ fn is_terminal_job_status(status: &str) -> bool {
async fn load_picture_job(state: &AppState, id: &str, auth: &PictureAgentAuth) -> AppResult<Value> {
mark_stale_picture_jobs_failed(state).await?;
let task: BackgroundTask =
sqlx::query_as("SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1")
.bind(id)
.bind(TASK_TYPE_PICTURE_GENERATION)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("Picture job not found".into()))?;
let task: BackgroundTask = crate::db::sql::query_as(
"SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1",
)
.bind(id)
.bind(TASK_TYPE_PICTURE_GENERATION)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("Picture job not found".into()))?;
if !job_visible_to_auth(&task, auth) {
return Err(AppError::NotFound("Picture job not found".into()));
}
@@ -607,9 +809,9 @@ async fn update_job_progress(
progress: i64,
message: &str,
) -> AppResult<()> {
sqlx::query(
crate::db::sql::query(
"UPDATE background_tasks
SET progress = ?, message = ?, updated_at = datetime('now')
SET progress = ?, message = ?, updated_at = CURRENT_TIMESTAMP
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(progress)
@@ -622,7 +824,7 @@ async fn update_job_progress(
async fn is_job_cancelled(state: &AppState, id: &str) -> AppResult<bool> {
let status: Option<String> =
sqlx::query_scalar("SELECT status FROM background_tasks WHERE id = ? LIMIT 1")
crate::db::sql::query_scalar("SELECT status FROM background_tasks WHERE id = ? LIMIT 1")
.bind(id)
.fetch_optional(&*state.db)
.await?;
@@ -666,10 +868,10 @@ async fn create_generation_job(
"mode": "generate",
"request": body,
});
sqlx::query(
crate::db::sql::query(
"INSERT INTO background_tasks
(id, task_type, title, status, message, dedupe_key, result_json, progress, progress_total, started_at)
VALUES (?, ?, ?, 'running', '任务已提交,等待路由', ?, ?, 1, 4, datetime('now'))",
VALUES (?, ?, ?, 'running', '任务已提交,等待路由', ?, ?, 1, 4, CURRENT_TIMESTAMP)",
)
.bind(&task_id)
.bind(TASK_TYPE_PICTURE_GENERATION)
@@ -725,7 +927,7 @@ async fn create_generation_job(
"图片已返回,正在归档到工作空间",
)
.await?;
archive_picture_agent_response(
let urls = archive_picture_agent_response(
&state_for_task,
&auth_for_task,
prompt.clone(),
@@ -733,11 +935,6 @@ async fn create_generation_job(
&payload,
)
.await;
let urls = payload
.get("data")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
let final_result = json!({
"userId": auth_for_task.user_id,
"role": auth_for_task.role,
@@ -749,10 +946,10 @@ async fn create_generation_job(
"response": payload,
"images": urls,
});
sqlx::query(
crate::db::sql::query(
"UPDATE background_tasks
SET status = 'succeeded', message = '生图完成', result_json = ?, progress = 4,
updated_at = datetime('now'), finished_at = datetime('now')
updated_at = CURRENT_TIMESTAMP, finished_at = CURRENT_TIMESTAMP
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(final_result.to_string())
@@ -764,9 +961,9 @@ async fn create_generation_job(
.await;
if let Err(err) = result {
let _ = sqlx::query(
let _ = crate::db::sql::query(
"UPDATE background_tasks
SET status = 'failed', message = '生图失败', error = ?, updated_at = datetime('now'), finished_at = datetime('now')
SET status = 'failed', message = '生图失败', error = ?, updated_at = CURRENT_TIMESTAMP, finished_at = CURRENT_TIMESTAMP
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(err.to_string())
@@ -786,13 +983,14 @@ async fn cancel_job(
Path(id): Path<String>,
) -> AppResult<Json<Value>> {
let auth = resolve_picture_agent_auth(&headers, &state).await?;
let task: BackgroundTask =
sqlx::query_as("SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1")
.bind(&id)
.bind(TASK_TYPE_PICTURE_GENERATION)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("Picture job not found".into()))?;
let task: BackgroundTask = crate::db::sql::query_as(
"SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1",
)
.bind(&id)
.bind(TASK_TYPE_PICTURE_GENERATION)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("Picture job not found".into()))?;
if !job_visible_to_auth(&task, &auth) {
return Err(AppError::NotFound("Picture job not found".into()));
}
@@ -803,13 +1001,13 @@ async fn cancel_job(
})));
}
sqlx::query(
crate::db::sql::query(
"UPDATE background_tasks
SET status = 'cancelled',
message = '已取消',
error = NULL,
updated_at = datetime('now'),
finished_at = datetime('now')
updated_at = CURRENT_TIMESTAMP,
finished_at = CURRENT_TIMESTAMP
WHERE id = ? AND task_type = ? AND status IN ('pending', 'running')",
)
.bind(&id)
@@ -829,7 +1027,7 @@ async fn list_jobs(
let auth = resolve_picture_agent_auth(&headers, &state).await?;
mark_stale_picture_jobs_failed(&state).await?;
let limit = query.limit.unwrap_or(20).clamp(1, 100);
let tasks: Vec<BackgroundTask> = sqlx::query_as(
let tasks: Vec<BackgroundTask> = crate::db::sql::query_as(
"SELECT * FROM background_tasks
WHERE task_type = ?
ORDER BY created_at DESC
@@ -852,7 +1050,7 @@ async fn clear_jobs(State(state): State<AppState>, headers: HeaderMap) -> AppRes
let auth = resolve_picture_agent_auth(&headers, &state).await?;
mark_stale_picture_jobs_failed(&state).await?;
let tasks: Vec<BackgroundTask> = sqlx::query_as(
let tasks: Vec<BackgroundTask> = crate::db::sql::query_as(
"SELECT * FROM background_tasks
WHERE task_type = ?
ORDER BY created_at DESC",
@@ -866,7 +1064,7 @@ async fn clear_jobs(State(state): State<AppState>, headers: HeaderMap) -> AppRes
if !job_visible_to_auth(&task, &auth) || !is_terminal_job_status(&task.status) {
continue;
}
let affected = sqlx::query(
let affected = crate::db::sql::query(
"DELETE FROM background_tasks
WHERE id = ? AND task_type = ?",
)
+137
View File
@@ -0,0 +1,137 @@
use crate::{
auth::jwt::validate_jwt,
error::{AppError, AppResult},
services::profile_avatar,
AppState,
};
use axum::{
body::Body,
extract::{DefaultBodyLimit, Multipart, Path, State},
http::{header, HeaderMap, HeaderValue, StatusCode},
response::Response,
routing::{get, put},
Json, Router,
};
use serde_json::{json, Value};
const AVATAR_MULTIPART_BODY_LIMIT: usize = profile_avatar::MAX_AVATAR_UPLOAD_BYTES + 512 * 1024;
pub fn routes() -> Router<AppState> {
Router::new()
.route(
"/api/users/me/avatar",
put(update_avatar).delete(delete_avatar),
)
.route("/api/users/avatars/{public_id}", get(get_avatar))
.layer(DefaultBodyLimit::max(AVATAR_MULTIPART_BODY_LIMIT))
}
fn current_user_id(state: &AppState, headers: &HeaderMap) -> AppResult<i64> {
let token = headers
.get(header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.strip_prefix("Bearer "))
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
let claims = validate_jwt(&state.config.auth.jwt_secret, token)
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
claims
.sub
.parse::<i64>()
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))
}
async fn update_avatar(
State(state): State<AppState>,
headers: HeaderMap,
mut multipart: Multipart,
) -> AppResult<Json<Value>> {
let user_id = current_user_id(&state, &headers)?;
let mut upload = None;
while let Some(field) = multipart
.next_field()
.await
.map_err(|err| AppError::BadRequest(format!("头像上传请求无效: {err}")))?
{
if field.name() != Some("file") {
continue;
}
upload = Some(
field
.bytes()
.await
.map_err(|err| AppError::BadRequest(format!("头像文件读取失败: {err}")))?,
);
break;
}
let upload = upload.ok_or_else(|| AppError::BadRequest("请选择头像图片".into()))?;
let avatar_url = profile_avatar::replace_for_user(&state, user_id, &upload).await?;
Ok(Json(json!({ "success": true, "avatarUrl": avatar_url })))
}
async fn delete_avatar(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Json<Value>> {
let user_id = current_user_id(&state, &headers)?;
profile_avatar::delete_for_user(&state, user_id).await?;
Ok(Json(json!({ "success": true, "avatarUrl": Value::Null })))
}
async fn get_avatar(
State(state): State<AppState>,
Path(public_id): Path<String>,
) -> AppResult<Response<Body>> {
let avatar = profile_avatar::load_by_public_id(&state, &public_id)
.await?
.ok_or_else(|| AppError::NotFound("Avatar not found".into()))?;
avatar_response(avatar)
}
fn avatar_response(avatar: profile_avatar::StoredProfileAvatar) -> AppResult<Response<Body>> {
let mut response = Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, avatar.mime_type)
.header(header::CONTENT_LENGTH, avatar.bytes.len().to_string())
.header(header::CACHE_CONTROL, "public, max-age=31536000, immutable")
.header("X-Content-Type-Options", "nosniff")
.body(Body::from(avatar.bytes))
.map_err(|err| AppError::Internal(anyhow::anyhow!("build avatar response: {err}")))?;
response.headers_mut().insert(
header::ETAG,
HeaderValue::from_str(&format!("\"{}\"", avatar.public_id))
.map_err(|err| AppError::Internal(anyhow::anyhow!("build avatar etag: {err}")))?,
);
Ok(response)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn public_avatar_response_is_immutable_and_nosniff() {
let response = avatar_response(profile_avatar::StoredProfileAvatar {
public_id: "ava_test".into(),
bytes: vec![1, 2, 3],
mime_type: profile_avatar::AVATAR_MIME_TYPE.into(),
})
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.headers().get(header::CACHE_CONTROL).unwrap(),
"public, max-age=31536000, immutable"
);
assert_eq!(
response.headers().get("X-Content-Type-Options").unwrap(),
"nosniff"
);
assert_eq!(
response.headers().get(header::ETAG).unwrap(),
"\"ava_test\""
);
assert_eq!(
response.headers().get(header::CONTENT_TYPE).unwrap(),
profile_avatar::AVATAR_MIME_TYPE
);
}
}
+860 -316
View File
File diff suppressed because it is too large Load Diff
+7 -7
View File
@@ -37,7 +37,7 @@ async fn search_handler(
let per_category = (limit / 6).clamp(1, 10);
// 1. Search sites
let sites: Vec<Value> = sqlx::query_as::<_, (i64, String, String, String)>(
let sites: Vec<Value> = crate::db::sql::query_as::<_, (i64, String, String, String)>(
"SELECT id, name, url, platform FROM sites WHERE name LIKE ? OR url LIKE ? OR platform LIKE ? LIMIT ?"
)
.bind(&q).bind(&q).bind(&q).bind(per_category)
@@ -58,7 +58,7 @@ async fn search_handler(
.collect();
// 2. Search accounts (with site join)
let accounts: Vec<Value> = sqlx::query_as::<_, (i64, String, String, String, i64, String, String, Option<String>)>(
let accounts: Vec<Value> = crate::db::sql::query_as::<_, (i64, String, String, String, i64, String, String, Option<String>)>(
r#"SELECT a.id, a.username, a.status, a.extra_config, s.id, s.name, s.platform, a.access_token
FROM accounts a
INNER JOIN sites s ON s.id = a.site_id
@@ -80,7 +80,7 @@ async fn search_handler(
.collect();
// 3. Search account tokens (with account + site join)
let account_tokens: Vec<Value> = sqlx::query_as::<_, (i64, String, i64, String, i64, String, String)>(
let account_tokens: Vec<Value> = crate::db::sql::query_as::<_, (i64, String, i64, String, i64, String, String)>(
r#"SELECT t.id, t.name, a.id, a.username, s.id, s.name, s.platform
FROM account_tokens t
INNER JOIN accounts a ON a.id = t.account_id
@@ -103,7 +103,7 @@ async fn search_handler(
.collect();
// 4. Search checkin logs
let checkin_logs: Vec<Value> = sqlx::query_as::<_, (i64, String, i64, String)>(
let checkin_logs: Vec<Value> = crate::db::sql::query_as::<_, (i64, String, i64, String)>(
r#"SELECT c.id, c.message, a.id, a.username
FROM checkin_logs c
INNER JOIN accounts a ON a.id = c.account_id
@@ -121,7 +121,7 @@ async fn search_handler(
.collect();
// 5. Search proxy logs (by model)
let proxy_logs: Vec<Value> = sqlx::query_as::<_, (i64, String, String, i64, String, i64, String)>(
let proxy_logs: Vec<Value> = crate::db::sql::query_as::<_, (i64, String, String, i64, String, i64, String)>(
r#"SELECT p.id, p.model_requested, p.status, p.http_status, p.created_at,
COALESCE(a.id, 0), COALESCE(a.username, '')
FROM proxy_logs p
@@ -140,7 +140,7 @@ async fn search_handler(
.collect();
// 6. Search models (from model_availability + token_model_availability)
let token_model_rows: Vec<(String, i64, i64, i64)> = sqlx::query_as(
let token_model_rows: Vec<(String, i64, i64, i64)> = crate::db::sql::query_as(
r#"SELECT DISTINCT tma.model_name, t.id, a.id, s.id
FROM token_model_availability tma
INNER JOIN account_tokens t ON t.id = tma.token_id AND t.enabled = 1
@@ -154,7 +154,7 @@ async fn search_handler(
.fetch_all(&*state.db)
.await?;
let direct_model_rows: Vec<(String, i64, i64)> = sqlx::query_as(
let direct_model_rows: Vec<(String, i64, i64)> = crate::db::sql::query_as(
r#"SELECT DISTINCT ma.model_name, a.id, s.id
FROM model_availability ma
INNER JOIN accounts a ON a.id = ma.account_id AND a.status = 'active'
File diff suppressed because it is too large Load Diff
+27 -27
View File
@@ -91,7 +91,7 @@ async fn list_admin_announcements(
LIMIT ? OFFSET ?",
);
let mut query = sqlx::query_as::<_, AnnouncementRow>(&sql);
let mut query = crate::db::sql::query_as::<_, AnnouncementRow>(&sql);
if let Some(status) = status {
query = query.bind(status);
}
@@ -121,8 +121,8 @@ async fn list_user_announcements(
LEFT JOIN user_announcement_reads r
ON r.announcement_id = a.id AND r.user_id = ?
WHERE a.status = 'published'
AND (a.starts_at IS NULL OR a.starts_at <= datetime('now'))
AND (a.ends_at IS NULL OR a.ends_at >= datetime('now'))"
AND (a.starts_at IS NULL OR a.starts_at <= CURRENT_TIMESTAMP)
AND (a.ends_at IS NULL OR a.ends_at >= CURRENT_TIMESTAMP)"
.to_string();
if read_filter == Some("true") {
@@ -135,7 +135,7 @@ async fn list_user_announcements(
" ORDER BY COALESCE(a.published_at, a.created_at) DESC, a.id DESC LIMIT ? OFFSET ?",
);
let rows: Vec<UserAnnouncementRow> = sqlx::query_as(&sql)
let rows: Vec<UserAnnouncementRow> = crate::db::sql::query_as(&sql)
.bind(user_id)
.bind(limit)
.bind(offset)
@@ -154,10 +154,10 @@ async fn create_announcement(
) -> AppResult<Json<Value>> {
require_admin(&headers, &state)?;
let payload = normalize_payload(payload)?;
let result = sqlx::query(
let announcement_id: i64 = crate::db::sql::insert_id(
"INSERT INTO user_announcements
(title, content, level, status, audience, starts_at, ends_at, published_at)
VALUES (?, ?, ?, ?, 'all', ?, ?, CASE WHEN ? = 'published' THEN datetime('now') ELSE NULL END)",
VALUES (?, ?, ?, ?, 'all', ?, ?, CASE WHEN ? = 'published' THEN CURRENT_TIMESTAMP ELSE NULL END) RETURNING id",
)
.bind(&payload.title)
.bind(&payload.content)
@@ -166,10 +166,10 @@ async fn create_announcement(
.bind(&payload.starts_at)
.bind(&payload.ends_at)
.bind(&payload.status)
.execute(&*state.db)
.fetch_one(&*state.db)
.await?;
let row = fetch_admin_row(&state, result.last_insert_rowid()).await?;
let row = fetch_admin_row(&state, announcement_id).await?;
Ok(Json(row.into_admin_json()))
}
@@ -181,7 +181,7 @@ async fn update_announcement(
) -> AppResult<Json<Value>> {
require_admin(&headers, &state)?;
let payload = normalize_payload(payload)?;
let result = sqlx::query(
let result = crate::db::sql::query(
"UPDATE user_announcements
SET title = ?,
content = ?,
@@ -190,10 +190,10 @@ async fn update_announcement(
starts_at = ?,
ends_at = ?,
published_at = CASE
WHEN status != 'published' AND ? = 'published' THEN datetime('now')
WHEN status != 'published' AND ? = 'published' THEN CURRENT_TIMESTAMP
ELSE published_at
END,
updated_at = datetime('now')
updated_at = CURRENT_TIMESTAMP
WHERE id = ?",
)
.bind(&payload.title)
@@ -223,7 +223,7 @@ async fn delete_announcement(
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
require_admin(&headers, &state)?;
let result = sqlx::query("DELETE FROM user_announcements WHERE id = ?")
let result = crate::db::sql::query("DELETE FROM user_announcements WHERE id = ?")
.bind(id)
.execute(&*state.db)
.await?;
@@ -243,9 +243,9 @@ async fn mark_read_admin(
Path(id): Path<i64>,
) -> AppResult<Json<Value>> {
require_admin(&headers, &state)?;
sqlx::query(
crate::db::sql::query(
"INSERT INTO user_announcement_reads (announcement_id, user_id, read_at)
SELECT ?, id, datetime('now') FROM users WHERE role = 'user'
SELECT ?, id, CURRENT_TIMESTAMP FROM users WHERE role = 'user'
ON CONFLICT(announcement_id, user_id) DO UPDATE SET read_at = excluded.read_at",
)
.bind(id)
@@ -259,9 +259,9 @@ async fn mark_all_read_admin(
headers: HeaderMap,
) -> AppResult<Json<Value>> {
require_admin(&headers, &state)?;
sqlx::query(
crate::db::sql::query(
"INSERT INTO user_announcement_reads (announcement_id, user_id, read_at)
SELECT a.id, u.id, datetime('now')
SELECT a.id, u.id, CURRENT_TIMESTAMP
FROM user_announcements a
CROSS JOIN users u
WHERE u.role = 'user'
@@ -279,9 +279,9 @@ async fn mark_read_user(
) -> AppResult<Json<Value>> {
let user_id = current_user_id(&headers, &state)?;
ensure_visible_to_user(&state, id).await?;
sqlx::query(
crate::db::sql::query(
"INSERT INTO user_announcement_reads (announcement_id, user_id, read_at)
VALUES (?, ?, datetime('now'))
VALUES (?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(announcement_id, user_id) DO UPDATE SET read_at = excluded.read_at",
)
.bind(id)
@@ -296,13 +296,13 @@ async fn mark_all_read_user(
headers: HeaderMap,
) -> AppResult<Json<Value>> {
let user_id = current_user_id(&headers, &state)?;
sqlx::query(
crate::db::sql::query(
"INSERT INTO user_announcement_reads (announcement_id, user_id, read_at)
SELECT id, ?, datetime('now')
SELECT id, ?, CURRENT_TIMESTAMP
FROM user_announcements
WHERE status = 'published'
AND (starts_at IS NULL OR starts_at <= datetime('now'))
AND (ends_at IS NULL OR ends_at >= datetime('now'))
AND (starts_at IS NULL OR starts_at <= CURRENT_TIMESTAMP)
AND (ends_at IS NULL OR ends_at >= CURRENT_TIMESTAMP)
ON CONFLICT(announcement_id, user_id) DO UPDATE SET read_at = excluded.read_at",
)
.bind(user_id)
@@ -314,14 +314,14 @@ async fn mark_all_read_user(
async fn delete_all(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
require_admin(&headers, &state)?;
info!("Deleting all user announcements");
sqlx::query("DELETE FROM user_announcements")
crate::db::sql::query("DELETE FROM user_announcements")
.execute(&*state.db)
.await?;
Ok(Json(json!({ "success": true })))
}
async fn fetch_admin_row(state: &AppState, id: i64) -> AppResult<AnnouncementRow> {
let row = sqlx::query_as::<_, AnnouncementRow>(
let row = crate::db::sql::query_as::<_, AnnouncementRow>(
"SELECT
a.id, a.title, a.content, a.level, a.status, a.audience,
a.starts_at, a.ends_at, a.created_at, a.updated_at, a.published_at,
@@ -338,13 +338,13 @@ async fn fetch_admin_row(state: &AppState, id: i64) -> AppResult<AnnouncementRow
}
async fn ensure_visible_to_user(state: &AppState, id: i64) -> AppResult<()> {
let exists: Option<i64> = sqlx::query_scalar(
let exists: Option<i64> = crate::db::sql::query_scalar(
"SELECT id
FROM user_announcements
WHERE id = ?
AND status = 'published'
AND (starts_at IS NULL OR starts_at <= datetime('now'))
AND (ends_at IS NULL OR ends_at >= datetime('now'))",
AND (starts_at IS NULL OR starts_at <= CURRENT_TIMESTAMP)
AND (ends_at IS NULL OR ends_at >= CURRENT_TIMESTAMP)",
)
.bind(id)
.fetch_optional(&*state.db)
+244 -181
View File
@@ -1,14 +1,15 @@
// Stats routes — usage analytics queries
use crate::{error::AppResult, AppState};
use crate::{db::DbKind, error::AppResult, AppState};
use axum::{
extract::{Query, State},
routing::get,
Json, Router,
};
use chrono::{Datelike, Duration, FixedOffset, TimeZone, Utc};
use serde::Deserialize;
use serde_json::{json, Value};
use sqlx::{QueryBuilder, Row, Sqlite};
use sqlx::Row;
use std::collections::{BTreeMap, HashMap, HashSet};
use tracing::info;
@@ -91,6 +92,34 @@ fn parse_proxy_log_client_filter(value: Option<&String>) -> Option<ProxyLogClien
Some(ProxyLogClientFilter::Raw(text))
}
fn sqlite_utc_timestamp(value: chrono::DateTime<Utc>) -> String {
value.format("%Y-%m-%d %H:%M:%S").to_string()
}
fn dashboard_today_bounds_utc() -> (String, String) {
// The admin dashboard labels these fields as "today"; use the product's
// China-facing day boundary instead of SQLite's UTC (CURRENT_DATE::TEXT).
let offset = FixedOffset::east_opt(8 * 60 * 60).expect("valid UTC+8 offset");
let now_local = Utc::now().with_timezone(&offset);
let start_local = offset
.with_ymd_and_hms(
now_local.year(),
now_local.month(),
now_local.day(),
0,
0,
0,
)
.single()
.expect("fixed-offset midnight is unambiguous");
let end_local = start_local + chrono::Duration::days(1);
(
sqlite_utc_timestamp(start_local.with_timezone(&Utc)),
sqlite_utc_timestamp(end_local.with_timezone(&Utc)),
)
}
fn parse_proxy_log_filters(q: &std::collections::HashMap<String, String>) -> ProxyLogFilters {
let search = normalize_non_empty(q.get("search").or_else(|| q.get("q")));
@@ -110,11 +139,14 @@ fn parse_proxy_log_filters(q: &std::collections::HashMap<String, String>) -> Pro
}
}
fn push_proxy_log_condition_prefix(
builder: &mut QueryBuilder<'_, Sqlite>,
condition_count: &mut usize,
) {
builder.push(if *condition_count == 0 {
#[derive(Debug, Clone)]
enum ProxyLogBind {
Text(String),
I64(i64),
}
fn push_proxy_log_condition_prefix(sql: &mut String, condition_count: &mut usize) {
sql.push_str(if *condition_count == 0 {
" WHERE "
} else {
" AND "
@@ -122,107 +154,77 @@ fn push_proxy_log_condition_prefix(
*condition_count += 1;
}
fn push_proxy_log_raw_condition(
builder: &mut QueryBuilder<'_, Sqlite>,
condition_count: &mut usize,
condition: &str,
) {
push_proxy_log_condition_prefix(builder, condition_count);
builder.push(condition);
fn push_proxy_log_raw_condition(sql: &mut String, condition_count: &mut usize, condition: &str) {
push_proxy_log_condition_prefix(sql, condition_count);
sql.push_str(condition);
}
fn push_proxy_log_where(
builder: &mut QueryBuilder<'_, Sqlite>,
sql: &mut String,
binds: &mut Vec<ProxyLogBind>,
filters: &ProxyLogFilters,
) -> usize {
let mut condition_count = 0;
if let Some(status) = &filters.status {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("l.status = ");
builder.push_bind(status.clone());
push_proxy_log_condition_prefix(sql, &mut condition_count);
sql.push_str("l.status = ?");
binds.push(ProxyLogBind::Text(status.clone()));
}
if let Some(search_like) = &filters.search_like {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("(l.model_requested LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR l.model_actual LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR l.error_message LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR l.content_summary LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR dk.name LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR dk.group_name LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR dk.tags LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR a.username LIKE ");
builder.push_bind(search_like.clone());
builder.push(" OR s.name LIKE ");
builder.push_bind(search_like.clone());
builder.push(")");
push_proxy_log_condition_prefix(sql, &mut condition_count);
sql.push_str("(l.model_requested LIKE ? OR l.model_actual LIKE ? OR l.error_message LIKE ? OR l.content_summary LIKE ? OR dk.name LIKE ? OR dk.group_name LIKE ? OR dk.tags LIKE ? OR a.username LIKE ? OR s.name LIKE ?)");
binds.extend(std::iter::repeat_with(|| ProxyLogBind::Text(search_like.clone())).take(9));
}
if let Some(client) = &filters.client {
push_proxy_log_condition_prefix(builder, &mut condition_count);
push_proxy_log_condition_prefix(sql, &mut condition_count);
match client {
ProxyLogClientFilter::App(value) => {
builder.push("(l.client_app_id = ");
builder.push_bind(value.clone());
builder.push(" OR l.client_app_name = ");
builder.push_bind(value.clone());
builder.push(")");
sql.push_str("(l.client_app_id = ? OR l.client_app_name = ?)");
binds.extend([ProxyLogBind::Text(value.clone()), ProxyLogBind::Text(value.clone())]);
}
ProxyLogClientFilter::Family(value) => {
builder.push("LOWER(l.client_family) = ");
builder.push_bind(value.clone());
sql.push_str("LOWER(l.client_family) = ?");
binds.push(ProxyLogBind::Text(value.clone()));
}
ProxyLogClientFilter::Raw(value) => {
builder.push("(l.client_app_id = ");
builder.push_bind(value.clone());
builder.push(" OR l.client_family = ");
builder.push_bind(value.clone());
builder.push(" OR l.client_app_name = ");
builder.push_bind(value.clone());
builder.push(")");
sql.push_str("(l.client_app_id = ? OR l.client_family = ? OR l.client_app_name = ?)");
binds.extend(std::iter::repeat_with(|| ProxyLogBind::Text(value.clone())).take(3));
}
}
}
if let Some(site_id) = filters.site_id {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("a.site_id = ");
builder.push_bind(site_id);
push_proxy_log_condition_prefix(sql, &mut condition_count);
sql.push_str("a.site_id = ?");
binds.push(ProxyLogBind::I64(site_id));
}
if let Some(downstream_key_id) = filters.downstream_key_id {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("l.downstream_api_key_id = ");
builder.push_bind(downstream_key_id);
push_proxy_log_condition_prefix(sql, &mut condition_count);
sql.push_str("l.downstream_api_key_id = ?");
binds.push(ProxyLogBind::I64(downstream_key_id));
}
if let Some(from) = &filters.from {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("datetime(l.created_at) >= datetime(");
builder.push_bind(from.clone());
builder.push(")");
push_proxy_log_condition_prefix(sql, &mut condition_count);
sql.push_str("datetime(l.created_at) >= datetime(?)");
binds.push(ProxyLogBind::Text(from.clone()));
}
if let Some(to) = &filters.to {
push_proxy_log_condition_prefix(builder, &mut condition_count);
builder.push("datetime(l.created_at) <= datetime(");
builder.push_bind(to.clone());
builder.push(")");
push_proxy_log_condition_prefix(sql, &mut condition_count);
sql.push_str("datetime(l.created_at) <= datetime(?)");
binds.push(ProxyLogBind::Text(to.clone()));
}
condition_count
}
fn push_proxy_log_joined_from(builder: &mut QueryBuilder<'_, Sqlite>) {
builder.push(
fn push_proxy_log_joined_from(sql: &mut String) {
sql.push_str(
r#"
FROM proxy_logs l
LEFT JOIN accounts a ON a.id = l.account_id
@@ -269,21 +271,25 @@ async fn build_proxy_log_client_options(
let mut options = Vec::new();
let mut seen = HashSet::new();
let mut app_query =
QueryBuilder::<Sqlite>::new("SELECT DISTINCT l.client_app_id, l.client_app_name ");
push_proxy_log_joined_from(&mut app_query);
let mut condition_count = push_proxy_log_where(&mut app_query, &scoped);
let mut app_sql = "SELECT DISTINCT l.client_app_id, l.client_app_name ".to_string();
let mut app_binds = Vec::new();
push_proxy_log_joined_from(&mut app_sql);
let mut condition_count = push_proxy_log_where(&mut app_sql, &mut app_binds, &scoped);
push_proxy_log_raw_condition(
&mut app_query,
&mut app_sql,
&mut condition_count,
"((l.client_app_id IS NOT NULL AND TRIM(l.client_app_id) <> '') OR (l.client_app_name IS NOT NULL AND TRIM(l.client_app_name) <> ''))",
);
app_query.push(
" ORDER BY LOWER(COALESCE(NULLIF(l.client_app_name, ''), l.client_app_id)) ASC LIMIT 200",
);
app_sql.push_str(" ORDER BY l.client_app_name ASC, l.client_app_id ASC LIMIT 200");
let app_rows: Vec<(Option<String>, Option<String>)> =
app_query.build_query_as().fetch_all(&*state.db).await?;
let mut app_query = crate::db::sql::query_as(&app_sql);
for bind in app_binds {
app_query = match bind {
ProxyLogBind::Text(value) => app_query.bind(value),
ProxyLogBind::I64(value) => app_query.bind(value),
};
}
let app_rows: Vec<(Option<String>, Option<String>)> = app_query.fetch_all(&*state.db).await?;
for (client_app_id, client_app_name) in app_rows {
let app_id = normalize_owned_non_empty(client_app_id);
let app_name = normalize_owned_non_empty(client_app_name);
@@ -299,18 +305,25 @@ async fn build_proxy_log_client_options(
);
}
let mut family_query = QueryBuilder::<Sqlite>::new("SELECT DISTINCT l.client_family ");
push_proxy_log_joined_from(&mut family_query);
let mut condition_count = push_proxy_log_where(&mut family_query, &scoped);
let mut family_sql = "SELECT DISTINCT l.client_family ".to_string();
let mut family_binds = Vec::new();
push_proxy_log_joined_from(&mut family_sql);
let mut condition_count = push_proxy_log_where(&mut family_sql, &mut family_binds, &scoped);
push_proxy_log_raw_condition(
&mut family_query,
&mut family_sql,
&mut condition_count,
"l.client_family IS NOT NULL AND TRIM(l.client_family) <> ''",
);
family_query.push(" ORDER BY LOWER(l.client_family) ASC LIMIT 200");
family_sql.push_str(" ORDER BY l.client_family ASC LIMIT 200");
let family_rows: Vec<(Option<String>,)> =
family_query.build_query_as().fetch_all(&*state.db).await?;
let mut family_query = crate::db::sql::query_as(&family_sql);
for bind in family_binds {
family_query = match bind {
ProxyLogBind::Text(value) => family_query.bind(value),
ProxyLogBind::I64(value) => family_query.bind(value),
};
}
let family_rows: Vec<(Option<String>,)> = family_query.fetch_all(&*state.db).await?;
for (client_family,) in family_rows {
let Some(family) = normalize_owned_non_empty(client_family) else {
continue;
@@ -334,17 +347,21 @@ async fn build_proxy_log_site_options(
let mut scoped = filters.clone();
scoped.site_id = None;
let mut site_query = QueryBuilder::<Sqlite>::new(
"SELECT DISTINCT s.id, COALESCE(NULLIF(s.name, ''), '站点 #' || s.id) AS name, s.status ",
);
push_proxy_log_joined_from(&mut site_query);
let mut condition_count = push_proxy_log_where(&mut site_query, &scoped);
push_proxy_log_raw_condition(&mut site_query, &mut condition_count, "s.id IS NOT NULL");
site_query
.push(" ORDER BY LOWER(COALESCE(NULLIF(s.name, ''), '站点 #' || s.id)) ASC LIMIT 200");
let mut site_sql = "SELECT DISTINCT s.id, COALESCE(NULLIF(s.name, ''), '站点 #' || s.id) AS name, s.status ".to_string();
let mut site_binds = Vec::new();
push_proxy_log_joined_from(&mut site_sql);
let mut condition_count = push_proxy_log_where(&mut site_sql, &mut site_binds, &scoped);
push_proxy_log_raw_condition(&mut site_sql, &mut condition_count, "s.id IS NOT NULL");
site_sql.push_str(" ORDER BY name ASC LIMIT 200");
let rows: Vec<(i64, String, Option<String>)> =
site_query.build_query_as().fetch_all(&*state.db).await?;
let mut site_query = crate::db::sql::query_as(&site_sql);
for bind in site_binds {
site_query = match bind {
ProxyLogBind::Text(value) => site_query.bind(value),
ProxyLogBind::I64(value) => site_query.bind(value),
};
}
let rows: Vec<(i64, String, Option<String>)> = site_query.fetch_all(&*state.db).await?;
Ok(rows
.into_iter()
@@ -411,83 +428,117 @@ async fn dashboard(
}
async fn build_dashboard_summary(state: &AppState) -> AppResult<Value> {
let sites: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM sites")
let sites: i64 = crate::db::sql::query_scalar("SELECT COUNT(*) FROM sites")
.fetch_one(&*state.db)
.await?;
let accounts: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM accounts")
let accounts: i64 = crate::db::sql::query_scalar("SELECT COUNT(*) FROM accounts")
.fetch_one(&*state.db)
.await?;
let downstream_keys: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM downstream_api_keys")
.fetch_one(&*state.db)
.await?;
let requests_24h: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM proxy_logs WHERE created_at >= datetime('now', '-1 day')",
let downstream_keys: i64 =
crate::db::sql::query_scalar("SELECT COUNT(*) FROM downstream_api_keys")
.fetch_one(&*state.db)
.await?;
let requests_24h: i64 = crate::db::sql::query_scalar(
"SELECT COUNT(*) FROM proxy_logs WHERE datetime(created_at) >= datetime('now', '-1 day')",
)
.fetch_one(&*state.db)
.await?;
let success_24h: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM proxy_logs WHERE status = 'success' AND created_at >= datetime('now', '-1 day')",
let success_24h: i64 = crate::db::sql::query_scalar(
"SELECT COUNT(*) FROM proxy_logs
WHERE lower(COALESCE(status, '')) = 'success'
AND datetime(created_at) >= datetime('now', '-1 day')",
)
.fetch_one(&*state.db)
.await?;
let total_cost_24h: f64 = sqlx::query_scalar(
"SELECT COALESCE(SUM(estimated_cost), 0.0) FROM proxy_logs WHERE created_at >= datetime('now', '-1 day')",
let total_cost_24h: f64 = crate::db::sql::query_scalar(
"SELECT COALESCE(SUM(cost), 0.0)
FROM usage_logs
WHERE datetime(created_at) >= datetime('now', '-1 day')",
)
.fetch_one(&*state.db)
.await?;
let total_balance: f64 = crate::db::sql::query_scalar(
"SELECT COALESCE(SUM(
CASE
WHEN max_cost IS NULL THEN 0.0
ELSE MAX(max_cost - COALESCE(used_cost, 0.0), 0.0)
END
), 0.0)
FROM downstream_api_keys",
)
.fetch_one(&*state.db)
.await?;
let total_balance: f64 = sqlx::query_scalar("SELECT COALESCE(SUM(balance), 0.0) FROM accounts")
.fetch_one(&*state.db)
.await?;
let total_used: f64 =
sqlx::query_scalar("SELECT COALESCE(SUM(balance_used), 0.0) FROM accounts")
crate::db::sql::query_scalar("SELECT COALESCE(SUM(cost), 0.0) FROM usage_logs")
.fetch_one(&*state.db)
.await?;
let today_spend: f64 = sqlx::query_scalar(
"SELECT COALESCE(SUM(estimated_cost), 0.0)
FROM proxy_logs
WHERE date(created_at) = date('now')",
let (today_start_utc, today_end_utc) = dashboard_today_bounds_utc();
let today_spend: f64 = crate::db::sql::query_scalar(
"SELECT COALESCE(SUM(cost), 0.0)
FROM usage_logs
WHERE datetime(created_at) >= datetime(?) AND datetime(created_at) < datetime(?)",
)
.bind(&today_start_utc)
.bind(&today_end_utc)
.fetch_one(&*state.db)
.await?;
let today_reward: f64 = sqlx::query_scalar(
"SELECT COALESCE(SUM(CAST(reward AS REAL)), 0.0)
let reward_value = match crate::db::sql::runtime_kind() {
DbKind::Sqlite => "CAST(reward AS REAL)",
DbKind::Postgres => "CAST(reward AS DOUBLE PRECISION)",
DbKind::MySql => "CAST(reward AS DECIMAL(65, 30))",
};
let today_reward_sql = format!(
"SELECT COALESCE(SUM({reward_value}), 0.0)
FROM checkin_logs
WHERE status = 'success' AND date(created_at) = date('now')",
)
.fetch_one(&*state.db)
.await?;
WHERE status = 'success'
AND datetime(created_at) >= datetime(?)
AND datetime(created_at) < datetime(?)"
);
let today_reward: f64 = crate::db::sql::query_scalar(&today_reward_sql)
.bind(&today_start_utc)
.bind(&today_end_utc)
.fetch_one(&*state.db)
.await?;
let active_accounts: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM accounts WHERE status = 'active'")
crate::db::sql::query_scalar("SELECT COUNT(*) FROM accounts WHERE status = 'active'")
.fetch_one(&*state.db)
.await?;
let today_checkin_total: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM checkin_logs WHERE date(created_at) = date('now')",
)
.fetch_one(&*state.db)
.await?;
let today_checkin_success: i64 = sqlx::query_scalar(
let today_checkin_total: i64 = crate::db::sql::query_scalar(
"SELECT COUNT(*) FROM checkin_logs
WHERE status = 'success' AND date(created_at) = date('now')",
WHERE datetime(created_at) >= datetime(?) AND datetime(created_at) < datetime(?)",
)
.bind(&today_start_utc)
.bind(&today_end_utc)
.fetch_one(&*state.db)
.await?;
let total_tokens_24h: i64 = sqlx::query_scalar(
let today_checkin_success: i64 = crate::db::sql::query_scalar(
"SELECT COUNT(*) FROM checkin_logs
WHERE status = 'success'
AND datetime(created_at) >= datetime(?)
AND datetime(created_at) < datetime(?)",
)
.bind(&today_start_utc)
.bind(&today_end_utc)
.fetch_one(&*state.db)
.await?;
let total_tokens_24h: i64 = crate::db::sql::query_scalar(
"SELECT COALESCE(SUM(total_tokens), 0)
FROM proxy_logs
WHERE created_at >= datetime('now', '-1 day')",
WHERE datetime(created_at) >= datetime('now', '-1 day')",
)
.fetch_one(&*state.db)
.await?;
let perf_window_secs = 60_i64;
let perf_requests: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM proxy_logs WHERE created_at >= datetime('now', '-60 seconds')",
let perf_requests: i64 = crate::db::sql::query_scalar(
"SELECT COUNT(*) FROM proxy_logs
WHERE datetime(created_at) >= datetime('now', '-60 seconds')",
)
.fetch_one(&*state.db)
.await?;
let perf_tokens: i64 = sqlx::query_scalar(
let perf_tokens: i64 = crate::db::sql::query_scalar(
"SELECT COALESCE(SUM(total_tokens), 0)
FROM proxy_logs
WHERE created_at >= datetime('now', '-60 seconds')",
WHERE datetime(created_at) >= datetime('now', '-60 seconds')",
)
.fetch_one(&*state.db)
.await?;
@@ -536,11 +587,12 @@ async fn build_dashboard_summary(state: &AppState) -> AppResult<Value> {
async fn build_model_analysis(state: &AppState, days: i64) -> AppResult<Value> {
let days = days.max(1);
let rows = sqlx::query(
let since = (Utc::now() - Duration::days(days - 1)).to_rfc3339();
let rows = crate::db::sql::query(
r#"
SELECT
COALESCE(NULLIF(model_actual, ''), NULLIF(model_requested, ''), 'unknown') AS model,
date(created_at) AS day,
substr(created_at, 1, 10) AS day,
COUNT(*) AS calls,
COALESCE(SUM(CASE WHEN status = 'success' THEN 1 ELSE 0 END), 0) AS success_calls,
COALESCE(SUM(CASE WHEN status != 'success' OR status IS NULL THEN 1 ELSE 0 END), 0) AS failed_calls,
@@ -548,12 +600,12 @@ async fn build_model_analysis(state: &AppState, days: i64) -> AppResult<Value> {
COALESCE(SUM(total_tokens), 0) AS tokens,
AVG(latency_ms) AS avg_latency_ms
FROM proxy_logs
WHERE created_at >= datetime('now', ?)
WHERE created_at >= ?
GROUP BY model, day
ORDER BY day ASC
"#,
)
.bind(format!("-{} days", days - 1))
.bind(since)
.fetch_all(&*state.db)
.await?;
@@ -670,7 +722,7 @@ async fn build_site_availability(state: &AppState) -> AppResult<Vec<Value>> {
i64,
i64,
Option<f64>,
)> = sqlx::query_as(
)> = crate::db::sql::query_as(
r#"
SELECT
s.id,
@@ -743,8 +795,7 @@ async fn proxy_logs(
let items: Vec<Value> = if meta_only {
Vec::new()
} else {
let mut list_query = QueryBuilder::<Sqlite>::new(
r#"
let mut list_sql = r#"
SELECT
l.id, l.route_id, l.channel_id, l.account_id, l.downstream_api_key_id,
l.model_requested, l.model_actual, l.status, l.http_status,
@@ -761,22 +812,26 @@ async fn proxy_logs(
dk.name AS downstream_key_name,
dk.group_name AS downstream_key_group_name,
dk.tags AS downstream_key_tags
"#,
);
push_proxy_log_joined_from(&mut list_query);
push_proxy_log_where(&mut list_query, &filters);
list_query.push(
"#
.to_string();
let mut list_binds = Vec::new();
push_proxy_log_joined_from(&mut list_sql);
push_proxy_log_where(&mut list_sql, &mut list_binds, &filters);
list_sql.push_str(
r#"
ORDER BY l.created_at DESC
"#,
);
list_query.push(" LIMIT ");
list_query.push_bind(limit);
list_query.push(" OFFSET ");
list_query.push_bind(offset);
list_sql.push_str(&format!(" LIMIT {limit} OFFSET {offset}"));
let mut list_query = crate::db::sql::query_as::<_, ProxyLogRow>(&list_sql);
for bind in list_binds {
list_query = match bind {
ProxyLogBind::Text(value) => list_query.bind(value),
ProxyLogBind::I64(value) => list_query.bind(value),
};
}
list_query
.build_query_as::<ProxyLogRow>()
.fetch_all(&*state.db)
.await?
.into_iter()
@@ -784,20 +839,27 @@ async fn proxy_logs(
.collect()
};
let mut summary_query = QueryBuilder::<Sqlite>::new(
r#"
let mut summary_sql = r#"
SELECT
COUNT(l.id) AS total,
COALESCE(SUM(CASE WHEN l.status = 'success' THEN 1 ELSE 0 END), 0) AS success,
COALESCE(SUM(CASE WHEN l.status != 'success' OR l.status IS NULL THEN 1 ELSE 0 END), 0) AS failed,
COALESCE(SUM(l.estimated_cost), 0.0) AS cost,
COALESCE(SUM(l.total_tokens), 0) AS total_tokens
"#,
);
push_proxy_log_joined_from(&mut summary_query);
push_proxy_log_where(&mut summary_query, &filters);
CAST(COALESCE(SUM(l.total_tokens), 0) AS BIGINT) AS total_tokens
"#
.to_string();
let mut summary_binds = Vec::new();
push_proxy_log_joined_from(&mut summary_sql);
push_proxy_log_where(&mut summary_sql, &mut summary_binds, &filters);
let mut summary_query = crate::db::sql::query_as(&summary_sql);
for bind in summary_binds {
summary_query = match bind {
ProxyLogBind::Text(value) => summary_query.bind(value),
ProxyLogBind::I64(value) => summary_query.bind(value),
};
}
let (total, success, failed, cost, total_tokens): (i64, i64, i64, f64, i64) =
summary_query.build_query_as().fetch_one(&*state.db).await?;
summary_query.fetch_one(&*state.db).await?;
let (client_options, sites) = if include_meta {
(
build_proxy_log_client_options(&state, &filters).await?,
@@ -867,7 +929,7 @@ async fn proxy_log_detail(
State(state): State<AppState>,
axum::extract::Path(id): axum::extract::Path<i64>,
) -> AppResult<Json<Value>> {
let item: Option<ProxyLogRow> = sqlx::query_as(
let item: Option<ProxyLogRow> = crate::db::sql::query_as(
r#"
SELECT
l.id, l.route_id, l.channel_id, l.account_id, l.downstream_api_key_id,
@@ -946,7 +1008,7 @@ async fn proxy_debug_traces(
LIMIT ? OFFSET ?
"#
);
let mut list_query = sqlx::query_as::<_, ProxyDebugTraceRow>(&list_sql);
let mut list_query = crate::db::sql::query_as::<_, ProxyDebugTraceRow>(&list_sql);
for value in &bind_values {
list_query = list_query.bind(value);
}
@@ -960,7 +1022,7 @@ async fn proxy_debug_traces(
.collect();
let count_sql = format!("SELECT COUNT(*) FROM proxy_debug_traces{filter}");
let mut count_query = sqlx::query_scalar::<_, i64>(&count_sql);
let mut count_query = crate::db::sql::query_scalar::<_, i64>(&count_sql);
for value in &bind_values {
count_query = count_query.bind(value);
}
@@ -979,7 +1041,7 @@ async fn proxy_debug_trace_detail(
State(state): State<AppState>,
axum::extract::Path(id): axum::extract::Path<i64>,
) -> AppResult<Json<Value>> {
let trace: Option<ProxyDebugTraceDetailRow> = sqlx::query_as(
let trace: Option<ProxyDebugTraceDetailRow> = crate::db::sql::query_as(
r#"
SELECT
id, downstream_path, client_kind, session_id, trace_hint,
@@ -1004,7 +1066,7 @@ async fn proxy_debug_trace_detail(
));
};
let attempts: Vec<Value> = sqlx::query_as::<_, ProxyDebugAttemptRow>(
let attempts: Vec<Value> = crate::db::sql::query_as::<_, ProxyDebugAttemptRow>(
r#"
SELECT
id, trace_id, attempt_index, endpoint, request_path, target_url,
@@ -1128,7 +1190,7 @@ async fn model_by_site(
.to_string();
let rows = if let Some(site_id) = q.site_id {
sqlx::query_as::<_, (String, String, i64, f64, i64)>(
crate::db::sql::query_as::<_, (String, String, i64, f64, i64)>(
"SELECT model, local_day, total_calls, total_spend, total_tokens
FROM model_day_usage
WHERE site_id = ? AND local_day >= ?
@@ -1139,7 +1201,7 @@ async fn model_by_site(
.fetch_all(&*state.db)
.await?
} else {
sqlx::query_as::<_, (String, String, i64, f64, i64)>(
crate::db::sql::query_as::<_, (String, String, i64, f64, i64)>(
"SELECT model, local_day, total_calls, total_spend, total_tokens
FROM model_day_usage
WHERE local_day >= ?
@@ -1197,8 +1259,9 @@ async fn site_distribution(
.to_string();
// Sum per-site usage
let mut rows: Vec<(i64, String, Option<String>, i64, f64, i64, i64, f64)> = sqlx::query_as(
"SELECT
let mut rows: Vec<(i64, String, Option<String>, i64, f64, i64, i64, f64)> =
crate::db::sql::query_as(
"SELECT
s.id,
s.name,
s.platform,
@@ -1212,16 +1275,16 @@ async fn site_distribution(
LEFT JOIN model_day_usage m ON m.site_id = s.id AND m.local_day >= ?
GROUP BY s.id, s.name, s.platform
ORDER BY SUM(m.total_spend) DESC",
)
.bind(&since_day)
.fetch_all(&*state.db)
.await?;
)
.bind(&since_day)
.fetch_all(&*state.db)
.await?;
if rows
.iter()
.all(|(_, _, _, calls, spend, tokens, _, _)| *calls == 0 && *spend == 0.0 && *tokens == 0)
{
rows = sqlx::query_as(
rows = crate::db::sql::query_as(
"SELECT
s.id,
s.name,
@@ -1235,7 +1298,7 @@ async fn site_distribution(
LEFT JOIN accounts a ON a.site_id = s.id
LEFT JOIN proxy_logs l
ON l.account_id = a.id
AND date(l.created_at) >= ?
AND substr(l.created_at, 1, 10) >= ?
GROUP BY s.id, s.name, s.platform
ORDER BY COALESCE(SUM(l.estimated_cost), 0.0) DESC, COUNT(l.id) DESC",
)
@@ -1278,7 +1341,7 @@ async fn site_trend(
.format("%Y-%m-%d")
.to_string();
let mut rows: Vec<(String, String, i64, f64, i64)> = sqlx::query_as(
let mut rows: Vec<(String, String, i64, f64, i64)> = crate::db::sql::query_as(
"SELECT
m.local_day,
s.name,
@@ -1296,9 +1359,9 @@ async fn site_trend(
.await?;
if rows.is_empty() {
rows = sqlx::query_as(
rows = crate::db::sql::query_as(
"SELECT
date(l.created_at),
substr(l.created_at, 1, 10),
COALESCE(s.name, '未知站点'),
COUNT(l.id),
COALESCE(SUM(l.estimated_cost), 0.0),
@@ -1306,9 +1369,9 @@ async fn site_trend(
FROM proxy_logs l
LEFT JOIN accounts a ON a.id = l.account_id
LEFT JOIN sites s ON s.id = a.site_id
WHERE date(l.created_at) >= ?
GROUP BY date(l.created_at), s.id, s.name
ORDER BY date(l.created_at) ASC, s.name ASC",
WHERE substr(l.created_at, 1, 10) >= ?
GROUP BY substr(l.created_at, 1, 10), s.id, s.name
ORDER BY substr(l.created_at, 1, 10) ASC, s.name ASC",
)
.bind(&since_day)
.fetch_all(&*state.db)
+7 -6
View File
@@ -29,7 +29,7 @@ async fn list_tasks(
let limit = q.limit.unwrap_or(50).clamp(1, 500);
info!("Listing background tasks, limit={}", limit);
let tasks: Vec<BackgroundTask> = sqlx::query_as::<_, BackgroundTask>(
let tasks: Vec<BackgroundTask> = crate::db::sql::query_as::<_, BackgroundTask>(
"SELECT * FROM background_tasks ORDER BY created_at DESC LIMIT ?",
)
.bind(limit)
@@ -43,11 +43,12 @@ async fn list_tasks(
async fn get_task(State(state): State<AppState>, Path(id): Path<String>) -> AppResult<Json<Value>> {
info!("Getting background task: {}", id);
let task: Option<BackgroundTask> =
sqlx::query_as::<_, BackgroundTask>("SELECT * FROM background_tasks WHERE id = ?")
.bind(&id)
.fetch_optional(&*state.db)
.await?;
let task: Option<BackgroundTask> = crate::db::sql::query_as::<_, BackgroundTask>(
"SELECT * FROM background_tasks WHERE id = ?",
)
.bind(&id)
.fetch_optional(&*state.db)
.await?;
match task {
Some(t) => Ok(Json(json!({ "success": true, "task": t }))),
+428 -131
View File
File diff suppressed because it is too large Load Diff
+24 -21
View File
@@ -55,7 +55,7 @@ async fn list_user_keys(
let (user_id, _role) = extract_jwt_user(&headers, &state)?;
info!("Listing user API keys for user {}", user_id);
let keys: Vec<DownstreamApiKey> = sqlx::query_as::<_, DownstreamApiKey>(
let keys: Vec<DownstreamApiKey> = crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT * FROM downstream_api_keys WHERE user_id = ? ORDER BY created_at DESC",
)
.bind(user_id)
@@ -107,12 +107,12 @@ async fn create_user_key(
let excluded_site_ids = body.get("excludedSiteIds").map(|v| v.to_string());
let excluded_credential_refs = body.get("excludedCredentialRefs").map(|v| v.to_string());
let result = sqlx::query(
let key_id: i64 = crate::db::sql::insert_id(
r#"INSERT INTO downstream_api_keys
(name, key, description, tags, enabled, expires_at, max_cost, max_requests,
supported_models, allowed_route_ids, site_weight_multipliers,
excluded_site_ids, excluded_credential_refs, user_id)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"#,
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) RETURNING id"#,
)
.bind(name)
.bind(&key_str)
@@ -128,14 +128,15 @@ async fn create_user_key(
.bind(excluded_site_ids)
.bind(excluded_credential_refs)
.bind(user_id)
.execute(&*state.db)
.fetch_one(&*state.db)
.await?;
let new_key: DownstreamApiKey =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
.bind(result.last_insert_rowid())
.fetch_one(&*state.db)
.await?;
let new_key: DownstreamApiKey = crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT * FROM downstream_api_keys WHERE id = ?",
)
.bind(key_id)
.fetch_one(&*state.db)
.await?;
Ok((StatusCode::CREATED, Json(json!({ "key": new_key }))))
}
@@ -152,7 +153,7 @@ async fn update_user_key(
// Verify ownership
let existing: Option<DownstreamApiKey> =
sqlx::query_as("SELECT * FROM downstream_api_keys WHERE id = ? AND user_id = ?")
crate::db::sql::query_as("SELECT * FROM downstream_api_keys WHERE id = ? AND user_id = ?")
.bind(id)
.bind(user_id)
.fetch_optional(&*state.db)
@@ -229,7 +230,7 @@ async fn update_user_key(
bind_idx + 1,
bind_idx + 2
);
let mut q = sqlx::query(&sql);
let mut q = crate::db::sql::query(&sql);
for p in &params {
q = q.bind(p);
}
@@ -237,11 +238,12 @@ async fn update_user_key(
q.execute(&*state.db).await?;
}
let updated: DownstreamApiKey =
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
.bind(id)
.fetch_one(&*state.db)
.await?;
let updated: DownstreamApiKey = crate::db::sql::query_as::<_, DownstreamApiKey>(
"SELECT * FROM downstream_api_keys WHERE id = ?",
)
.bind(id)
.fetch_one(&*state.db)
.await?;
Ok(Json(json!({ "key": updated })))
}
@@ -255,11 +257,12 @@ async fn delete_user_key(
let (user_id, _role) = extract_jwt_user(&headers, &state)?;
info!("Deleting user API key {} for user {}", id, user_id);
let result = sqlx::query("DELETE FROM downstream_api_keys WHERE id = ? AND user_id = ?")
.bind(id)
.bind(user_id)
.execute(&*state.db)
.await?;
let result =
crate::db::sql::query("DELETE FROM downstream_api_keys WHERE id = ? AND user_id = ?")
.bind(id)
.bind(user_id)
.execute(&*state.db)
.await?;
if result.rows_affected() == 0 {
return Err(AppError::NotFound("Key not found".into()));
File diff suppressed because it is too large Load Diff
+13 -6
View File
@@ -657,14 +657,17 @@ mod tests {
services::{builtin_keys::BuiltinKeys, runtime_auth::RuntimeAuthSettings},
};
use axum::http::HeaderValue;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
async fn test_state() -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool");
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool");
AppState::new(
AppConfig {
@@ -678,6 +681,8 @@ mod tests {
url: "sqlite::memory:".to_string(),
max_connections: 1,
ssl: false,
ssl_mode: None,
ca_cert_path: None,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
@@ -696,9 +701,11 @@ mod tests {
template_proxy_url: None,
},
comfyui: ComfyUIConfig::default(),
agent: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
jimeng: Default::default(),
seedance: Default::default(),
},
pool,
BuiltinKeys {
+85 -2
View File
@@ -1,7 +1,7 @@
use super::{
cache_control_for_path, content_type_for_path, default_boosapi_workflow_graph,
embedded_workflow_ui_response, inject_bootstrap, local_workflow_template_asset,
native_workflow_base, simple_urlencode, text_to_video_workflow_graph,
native_workflow_base, native_workflow_request, simple_urlencode, text_to_video_workflow_graph,
workflow_template_not_found_response, EMBEDDED_WORKFLOW_UI, EMBEDDED_WORKFLOW_UI_CLOUD,
EMBEDDED_WORKFLOW_UI_EXTERNAL, OFFICIAL_WORKFLOW_TEMPLATES_BASE,
};
@@ -9,7 +9,7 @@ use crate::error::{AppError, AppResult};
use crate::AppState;
use axum::{
extract::{Path, State},
http::{header, StatusCode},
http::{header, Method, StatusCode},
response::{Html, IntoResponse, Response},
Json,
};
@@ -105,6 +105,89 @@ pub(super) async fn workflow_ui_runtime(State(state): State<AppState>) -> Json<V
}))
}
pub(super) async fn workflow_ui_health(State(state): State<AppState>) -> Json<Value> {
let config = &state.config.comfyui;
let mode = if config.uses_docker_backend() {
"docker"
} else if config.uses_cloud_backend() {
"cloud"
} else if config.uses_runninghub_backend() {
"runninghub"
} else if config.uses_native_backend() {
"external"
} else {
"boosapi"
};
if config.uses_cloud_backend() && workflow_cloud_api_key(config).is_none() {
return Json(json!({
"ready": false,
"mode": mode,
"message": "Comfy Cloud API Key 未配置,请在 [comfyui] 配置 api_key 或使用站点连接。",
}));
}
if !config.uses_native_backend() {
return Json(json!({
"ready": true,
"mode": mode,
}));
}
let probe_path = if config.uses_cloud_backend() {
"/user"
} else {
"/system_stats"
};
let result = native_workflow_request(&state, Method::GET, probe_path, None)
.send()
.await;
match result {
Ok(response) if response.status().is_success() => Json(json!({
"ready": true,
"mode": mode,
})),
Ok(response) => Json(json!({
"ready": false,
"mode": mode,
"status": response.status().as_u16(),
"message": workflow_health_failure_message(mode, response.status().as_u16()),
})),
Err(error) => {
tracing::warn!(%error, mode, "workflow backend health probe failed");
Json(json!({
"ready": false,
"mode": mode,
"message": workflow_health_unreachable_message(mode),
}))
}
}
}
fn workflow_health_failure_message(mode: &str, status: u16) -> String {
match (mode, status) {
("cloud", 401 | 403) => {
"Comfy Cloud 认证失败,请检查 comfyui.api_key 是否为有效的工作流 API Key。".to_string()
}
("cloud", _) => format!("Comfy Cloud API 返回 HTTP {status},请检查账号、套餐和后端配置。"),
("runninghub", 401 | 403) => "RunningHub 认证失败,请检查 comfyui.api_key。".to_string(),
("runninghub", _) => {
format!("RunningHub API 返回 HTTP {status},请检查后端地址和工作流服务。")
}
(_, 401 | 403) => "外部 ComfyUI 认证失败,请检查 API Key 或访问权限。".to_string(),
_ => format!("外部 ComfyUI 返回 HTTP {status},请检查 comfyui.base_url。"),
}
}
fn workflow_health_unreachable_message(mode: &str) -> &'static str {
match mode {
"cloud" => "Comfy Cloud API 不可达,请检查网络、代理和 comfyui.base_url。",
"runninghub" => "RunningHub API 不可达,请检查网络和 comfyui.base_url。",
_ => "外部 ComfyUI 不可达,请检查 comfyui.base_url 和服务状态。",
}
}
pub(super) fn workflow_ui_build_metadata(dir: &StdPath) -> Value {
let path = dir.join("boosapi-workflow-ui.json");
std::fs::read_to_string(path)
+4 -1
View File
@@ -57,6 +57,9 @@ pub(super) fn caller_is_admin(caller: &WorkflowCaller) -> bool {
matches!(caller, WorkflowCaller::Admin)
}
pub(super) async fn auth_caller(state: &AppState, headers: &HeaderMap) -> AppResult<WorkflowCaller> {
pub(super) async fn auth_caller(
state: &AppState,
headers: &HeaderMap,
) -> AppResult<WorkflowCaller> {
caller_from_headers(state, headers)
}
+692
View File
@@ -0,0 +1,692 @@
use super::{
auth::{auth_caller, caller_is_admin, caller_user_id},
native_workflow_request,
};
use crate::{
error::{AppError, AppResult},
services::comfy_cloud::ownership,
AppState,
};
use axum::{
body::Bytes,
http::{header, HeaderMap, HeaderValue, Method, StatusCode, Uri},
response::{IntoResponse, Response},
};
use serde_json::{json, Value};
use std::collections::{BTreeMap, HashSet};
const MAX_ASSET_PAGE_SIZE: usize = 500;
#[derive(Debug)]
struct UpstreamResponse {
status: StatusCode,
content_type: String,
cache_control: Option<HeaderValue>,
bytes: Bytes,
}
pub(super) async fn handle_cloud_asset_request(
state: AppState,
headers: HeaderMap,
path: String,
method: Method,
uri: Uri,
body: Bytes,
) -> AppResult<Response> {
let caller = auth_caller(&state, &headers).await?;
let target = match uri.query() {
Some(query) => format!("/{path}?{query}"),
None => format!("/{path}"),
};
if caller_is_admin(&caller) {
return forward_raw(&state, method, &target, &headers, body)
.await
.map(UpstreamResponse::into_response);
}
let user_id = caller_user_id(&caller).expect("non-admin workflow caller has user id");
let segments = path.split('/').collect::<Vec<_>>();
if segments.first() != Some(&"assets") {
return Err(AppError::NotFound("Cloud asset route not found".into()));
}
match (method.clone(), segments.as_slice()) {
(Method::GET, ["assets"]) => list_assets(&state, user_id, &uri).await,
(Method::POST, ["assets"]) => create_asset(&state, user_id, &headers, &target, body).await,
(Method::POST, ["assets", "download"]) => {
create_asset_download(&state, user_id, &headers, &target, body).await
}
(Method::GET, ["assets", "remote-metadata"]) => {
forward_raw(&state, Method::GET, &target, &headers, body)
.await
.map(UpstreamResponse::into_response)
}
(Method::GET, ["assets", asset_id]) => {
read_asset(&state, user_id, asset_id, &headers, &target).await
}
(Method::GET, ["assets", asset_id, "content"]) => {
read_asset_content(&state, user_id, asset_id, &headers, &target).await
}
(Method::PUT, ["assets", asset_id]) => {
update_asset(&state, user_id, asset_id, &headers, &target, body).await
}
(Method::DELETE, ["assets", asset_id]) => {
delete_asset(&state, user_id, asset_id, &headers, &target).await
}
(Method::POST, ["assets", asset_id, "tags"])
| (Method::DELETE, ["assets", asset_id, "tags"]) => {
update_asset_tags(&state, user_id, asset_id, method, &headers, &target, body).await
}
(Method::POST, ["assets", "export"]) => {
export_assets(&state, user_id, &headers, &target, body).await
}
_ => Err(AppError::NotFound("Cloud asset route not found".into())),
}
}
pub(super) async fn handle_cloud_job_request(
state: AppState,
headers: HeaderMap,
path: String,
method: Method,
uri: Uri,
body: Bytes,
) -> AppResult<Response> {
let caller = auth_caller(&state, &headers).await?;
let target = match uri.query() {
Some(query) => format!("/{path}?{query}"),
None => format!("/{path}"),
};
if caller_is_admin(&caller) {
return forward_raw(&state, method, &target, &headers, body)
.await
.map(UpstreamResponse::into_response);
}
let user_id = caller_user_id(&caller).expect("non-admin workflow caller has user id");
let segments = path.split('/').collect::<Vec<_>>();
match (method, segments.as_slice()) {
(Method::GET, ["jobs"]) => {
let response =
forward_raw(&state, Method::GET, &target, &headers, Bytes::new()).await?;
if !response.status.is_success() {
return Ok(response.into_response());
}
let mut value = parse_json(&response.bytes).ok_or_else(|| {
AppError::UpstreamError("Comfy Cloud returned invalid jobs JSON".into())
})?;
let owned = ownership::owned_job_ids(&state.db, user_id, 200)
.await?
.into_iter()
.collect::<HashSet<_>>();
let jobs = value
.get_mut("jobs")
.and_then(Value::as_array_mut)
.ok_or_else(|| {
AppError::UpstreamError("Comfy Cloud jobs response is invalid".into())
})?;
jobs.retain(|job| job_id_for_value(job).is_some_and(|id| owned.contains(id)));
let total = jobs.len();
value["total"] = json!(total);
value["hasMore"] = Value::Bool(false);
value["has_more"] = Value::Bool(false);
Ok(json_response(response.status, value))
}
(Method::GET, ["jobs", job_id]) => {
if !ownership::job_is_owned_by(&state.db, user_id, job_id).await? {
return Err(AppError::NotFound("Cloud job not found".into()));
}
sync_owned_assets(&state, user_id).await?;
forward_raw(&state, Method::GET, &target, &headers, Bytes::new())
.await
.map(UpstreamResponse::into_response)
}
_ => Err(AppError::NotFound("Cloud job route not found".into())),
}
}
pub(super) async fn sync_owned_assets(state: &AppState, user_id: i64) -> AppResult<()> {
let jobs = crate::db::sql::query_scalar::<_, String>(
"SELECT job_id FROM comfy_cloud_jobs
WHERE user_id = ?
AND (last_synced_at IS NULL OR last_synced_at < datetime('now', '-30 seconds'))
ORDER BY updated_at DESC, job_id DESC
LIMIT 50",
)
.bind(user_id)
.fetch_all(&*state.db)
.await?;
for job_id in jobs {
let query = url::form_urlencoded::Serializer::new(String::new())
.append_pair("job_ids", &job_id)
.append_pair("limit", "500")
.finish();
let value = match fetch_upstream_json(state, &format!("/assets?{query}")).await {
Ok(value) => value,
Err(err) => {
tracing::warn!(%err, user_id, job_id, "failed to sync owned Comfy Cloud assets");
continue;
}
};
if let Some(assets) = value.get("assets").and_then(Value::as_array) {
for asset in assets {
ownership::record_asset(&state.db, user_id, Some(&job_id), asset).await?;
}
}
crate::db::sql::query(
"UPDATE comfy_cloud_jobs SET last_synced_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP
WHERE job_id = ? AND user_id = ?",
)
.bind(&job_id)
.bind(user_id)
.execute(&*state.db)
.await?;
}
Ok(())
}
async fn list_assets(state: &AppState, user_id: i64, uri: &Uri) -> AppResult<Response> {
sync_owned_assets(state, user_id).await?;
let query = query_map(uri);
let mut assets = ownership::list_owned_assets(&state.db, user_id).await?;
assets.retain(|asset| !ownership::asset_is_shared_model(asset));
if should_include_models(&query) {
let models = fetch_upstream_json(state, &model_list_path(&query)).await?;
if let Some(model_assets) = models.get("assets").and_then(Value::as_array) {
assets.extend(
model_assets
.iter()
.filter(|asset| ownership::asset_is_shared_model(asset))
.cloned(),
);
}
}
assets.retain(|asset| asset_matches(asset, &query));
assets.sort_by(|left, right| {
asset_created_at(right)
.cmp(&asset_created_at(left))
.then_with(|| asset_id(right).cmp(&asset_id(left)))
});
let total = assets.len();
let offset = query_usize(&query, "offset").unwrap_or_default().min(total);
let limit = query_usize(&query, "limit")
.unwrap_or(MAX_ASSET_PAGE_SIZE)
.clamp(1, MAX_ASSET_PAGE_SIZE);
let end = offset.saturating_add(limit).min(total);
let items = assets[offset..end]
.iter()
.map(sanitize_asset)
.collect::<Vec<_>>();
Ok(axum::Json(json!({
"assets": items,
"total": total,
"has_more": end < total,
}))
.into_response())
}
async fn create_asset(
state: &AppState,
user_id: i64,
headers: &HeaderMap,
target: &str,
body: Bytes,
) -> AppResult<Response> {
validate_asset_create_body(&state.db, user_id, headers, &body).await?;
let response = forward_raw(state, Method::POST, target, headers, body).await?;
if response.status.is_success() {
if let Some(asset) = parse_json(&response.bytes) {
if !ownership::asset_is_shared_model(&asset) {
ownership::record_asset(&state.db, user_id, None, &asset).await?;
}
}
}
Ok(response.into_response())
}
async fn create_asset_download(
state: &AppState,
user_id: i64,
headers: &HeaderMap,
target: &str,
body: Bytes,
) -> AppResult<Response> {
validate_json_asset_write(&state.db, user_id, &body).await?;
let response = forward_raw(state, Method::POST, target, headers, body).await?;
if response.status.is_success() {
if let Some(value) = parse_json(&response.bytes) {
if let Some(asset) = value.get("asset") {
ownership::record_asset(&state.db, user_id, None, asset).await?;
}
for key in ["task_id", "job_id", "prompt_id"] {
if let Some(job_id) = value.get(key).and_then(Value::as_str) {
ownership::record_job(&state.db, user_id, job_id).await?;
}
}
}
}
Ok(response.into_response())
}
async fn read_asset(
state: &AppState,
user_id: i64,
asset_id: &str,
headers: &HeaderMap,
target: &str,
) -> AppResult<Response> {
ensure_asset_readable(state, user_id, asset_id).await?;
let response = forward_raw(state, Method::GET, target, headers, Bytes::new()).await?;
if response.status.is_success() {
if let Some(asset) = parse_json(&response.bytes) {
if !ownership::asset_is_shared_model(&asset) {
ownership::record_asset(&state.db, user_id, None, &asset).await?;
}
return Ok(json_response(response.status, sanitize_asset(&asset)));
}
}
Ok(response.into_response())
}
async fn read_asset_content(
state: &AppState,
user_id: i64,
asset_id: &str,
headers: &HeaderMap,
target: &str,
) -> AppResult<Response> {
ensure_asset_readable(state, user_id, asset_id).await?;
forward_raw(state, Method::GET, target, headers, Bytes::new())
.await
.map(UpstreamResponse::into_response)
}
async fn update_asset(
state: &AppState,
user_id: i64,
asset_id: &str,
headers: &HeaderMap,
target: &str,
body: Bytes,
) -> AppResult<Response> {
ensure_asset_owned(&state.db, user_id, asset_id).await?;
validate_json_asset_write(&state.db, user_id, &body).await?;
let response = forward_raw(state, Method::PUT, target, headers, body).await?;
if response.status.is_success() {
if let Some(asset) = parse_json(&response.bytes) {
ownership::record_asset(&state.db, user_id, None, &asset).await?;
return Ok(json_response(response.status, sanitize_asset(&asset)));
}
}
Ok(response.into_response())
}
async fn delete_asset(
state: &AppState,
user_id: i64,
asset_id: &str,
headers: &HeaderMap,
target: &str,
) -> AppResult<Response> {
ensure_asset_owned(&state.db, user_id, asset_id).await?;
let response = forward_raw(state, Method::DELETE, target, headers, Bytes::new()).await?;
if response.status.is_success() {
ownership::delete_asset(&state.db, user_id, asset_id).await?;
}
Ok(response.into_response())
}
async fn update_asset_tags(
state: &AppState,
user_id: i64,
asset_id: &str,
method: Method,
headers: &HeaderMap,
target: &str,
body: Bytes,
) -> AppResult<Response> {
ensure_asset_owned(&state.db, user_id, asset_id).await?;
reject_models_tag(&body)?;
forward_raw(state, method, target, headers, body)
.await
.map(UpstreamResponse::into_response)
}
async fn export_assets(
state: &AppState,
user_id: i64,
headers: &HeaderMap,
target: &str,
body: Bytes,
) -> AppResult<Response> {
let payload = parse_json(&body)
.ok_or_else(|| AppError::BadRequest("asset export payload must be JSON".into()))?;
let jobs = payload
.get("job_ids")
.and_then(Value::as_array)
.ok_or_else(|| AppError::BadRequest("asset export requires job_ids".into()))?;
for job_id in jobs.iter().filter_map(Value::as_str) {
if !ownership::job_is_owned_by(&state.db, user_id, job_id).await? {
return Err(AppError::NotFound("Cloud asset not found".into()));
}
}
forward_raw(state, Method::POST, target, headers, body)
.await
.map(UpstreamResponse::into_response)
}
async fn ensure_asset_owned(db: &sqlx::AnyPool, user_id: i64, asset_id: &str) -> AppResult<()> {
if ownership::asset_is_owned_by(db, user_id, asset_id).await? {
Ok(())
} else {
Err(AppError::NotFound("Cloud asset not found".into()))
}
}
async fn ensure_asset_readable(state: &AppState, user_id: i64, asset_id: &str) -> AppResult<()> {
if ownership::asset_is_owned_by(&state.db, user_id, asset_id).await? {
return Ok(());
}
let asset = fetch_upstream_json(state, &format!("/assets/{asset_id}")).await?;
if ownership::asset_is_shared_model(&asset) {
Ok(())
} else {
Err(AppError::NotFound("Cloud asset not found".into()))
}
}
async fn validate_asset_create_body(
db: &sqlx::AnyPool,
user_id: i64,
headers: &HeaderMap,
body: &[u8],
) -> AppResult<()> {
let is_json = headers
.get(header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.is_some_and(|value| value.to_ascii_lowercase().starts_with("application/json"));
if is_json {
return validate_json_asset_write(db, user_id, body).await;
}
// The Cloud frontend uses multipart uploads for image/media assets. The tags
// field is text within the body; reject a model tag before forwarding it.
reject_models_tag(body)
}
async fn validate_json_asset_write(db: &sqlx::AnyPool, user_id: i64, body: &[u8]) -> AppResult<()> {
let payload = parse_json(body)
.ok_or_else(|| AppError::BadRequest("asset payload must be valid JSON".into()))?;
if payload
.get("tags")
.and_then(Value::as_array)
.is_some_and(|tags| tags.iter().any(|tag| tag.as_str() == Some("models")))
{
return Err(AppError::Forbidden(
"普通用户不能创建或标记共享模型资产".into(),
));
}
if let Some(preview_id) = payload.get("preview_id").and_then(Value::as_str) {
ensure_asset_owned(db, user_id, preview_id).await?;
}
Ok(())
}
fn reject_models_tag(body: &[u8]) -> AppResult<()> {
if body
.windows("models".len())
.any(|window| window.eq_ignore_ascii_case(b"models"))
{
return Err(AppError::Forbidden(
"普通用户不能创建或标记共享模型资产".into(),
));
}
Ok(())
}
async fn fetch_upstream_json(state: &AppState, path: &str) -> AppResult<Value> {
let response = native_workflow_request(state, Method::GET, path, None)
.send()
.await
.map_err(|err| AppError::UpstreamError(format!("Comfy Cloud unreachable: {err}")))?;
let status = response.status();
let bytes = response
.bytes()
.await
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
if status == reqwest::StatusCode::NOT_FOUND {
return Err(AppError::NotFound("Cloud asset not found".into()));
}
if !status.is_success() {
return Err(AppError::UpstreamError(format!(
"Comfy Cloud asset request returned HTTP {status}"
)));
}
parse_json(&bytes)
.ok_or_else(|| AppError::UpstreamError("Comfy Cloud returned invalid asset JSON".into()))
}
async fn forward_raw(
state: &AppState,
method: Method,
path: &str,
headers: &HeaderMap,
body: Bytes,
) -> AppResult<UpstreamResponse> {
let mut request = native_workflow_request(state, method, path, None);
for (source, target) in [
(header::CONTENT_TYPE, reqwest::header::CONTENT_TYPE),
(header::ACCEPT, reqwest::header::ACCEPT),
] {
if let Some(value) = headers.get(&source).and_then(|value| value.to_str().ok()) {
request = request.header(target, value);
}
}
if !body.is_empty() {
request = request.body(body);
}
let response = request
.send()
.await
.map_err(|err| AppError::UpstreamError(format!("Comfy Cloud unreachable: {err}")))?;
let status =
StatusCode::from_u16(response.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or("application/json")
.to_string();
let cache_control = response
.headers()
.get(reqwest::header::CACHE_CONTROL)
.and_then(|value| HeaderValue::from_bytes(value.as_bytes()).ok());
let bytes = response
.bytes()
.await
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
Ok(UpstreamResponse {
status,
content_type,
cache_control,
bytes,
})
}
impl UpstreamResponse {
fn into_response(self) -> Response {
let mut builder = Response::builder()
.status(self.status)
.header(header::CONTENT_TYPE, self.content_type);
if let Some(cache_control) = self.cache_control {
builder = builder.header(header::CACHE_CONTROL, cache_control);
}
builder
.body(axum::body::Body::from(self.bytes))
.unwrap_or_else(|err| AppError::Internal(anyhow::anyhow!(err)).into_response())
}
}
fn json_response(status: StatusCode, value: Value) -> Response {
Response::builder()
.status(status)
.header(header::CONTENT_TYPE, "application/json")
.body(axum::body::Body::from(value.to_string()))
.unwrap_or_else(|err| AppError::Internal(anyhow::anyhow!(err)).into_response())
}
fn parse_json(bytes: &[u8]) -> Option<Value> {
serde_json::from_slice(bytes).ok()
}
fn query_map(uri: &Uri) -> BTreeMap<String, String> {
uri.query()
.map(|query| url::form_urlencoded::parse(query.as_bytes()))
.into_iter()
.flatten()
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect()
}
fn query_usize(query: &BTreeMap<String, String>, key: &str) -> Option<usize> {
query.get(key).and_then(|value| value.parse().ok())
}
fn split_tags(query: &BTreeMap<String, String>, key: &str) -> HashSet<String> {
query
.get(key)
.map(String::as_str)
.unwrap_or_default()
.split(',')
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.collect()
}
fn should_include_models(query: &BTreeMap<String, String>) -> bool {
let included = split_tags(query, "include_tags");
let excluded = split_tags(query, "exclude_tags");
!excluded.contains("models") && (included.is_empty() || included.contains("models"))
}
fn model_list_path(query: &BTreeMap<String, String>) -> String {
let mut serializer = url::form_urlencoded::Serializer::new(String::new());
serializer.append_pair("include_tags", "models");
serializer.append_pair("limit", "500");
for key in ["name_contains", "asset_hash", "order_by", "order"] {
if let Some(value) = query.get(key) {
serializer.append_pair(key, value);
}
}
format!("/assets?{}", serializer.finish())
}
fn asset_matches(asset: &Value, query: &BTreeMap<String, String>) -> bool {
let tags = asset
.get("tags")
.and_then(Value::as_array)
.map(|tags| {
tags.iter()
.filter_map(Value::as_str)
.map(str::to_string)
.collect::<HashSet<_>>()
})
.unwrap_or_default();
let required = split_tags(query, "include_tags");
if !required.iter().all(|tag| tags.contains(tag)) {
return false;
}
let excluded = split_tags(query, "exclude_tags");
if excluded.iter().any(|tag| tags.contains(tag)) {
return false;
}
if let Some(name) = query.get("name_contains") {
if !asset
.get("name")
.and_then(Value::as_str)
.is_some_and(|value| {
value
.to_ascii_lowercase()
.contains(&name.to_ascii_lowercase())
})
{
return false;
}
}
query
.get("asset_hash")
.is_none_or(|hash| asset.get("asset_hash").and_then(Value::as_str) == Some(hash.as_str()))
}
fn asset_created_at(asset: &Value) -> &str {
asset
.get("created_at")
.and_then(Value::as_str)
.unwrap_or_default()
}
fn asset_id(asset: &Value) -> &str {
asset.get("id").and_then(Value::as_str).unwrap_or_default()
}
fn job_id_for_value(job: &Value) -> Option<&str> {
for key in ["id", "job_id", "prompt_id", "task_id"] {
if let Some(value) = job.get(key).and_then(Value::as_str) {
return Some(value);
}
}
None
}
fn sanitize_asset(asset: &Value) -> Value {
let mut asset = asset.clone();
let content_id = asset
.get("preview_id")
.and_then(Value::as_str)
.or_else(|| asset.get("id").and_then(Value::as_str));
if let Some(content_id) = content_id {
let url = format!("/workflow-ui/api/assets/{content_id}/content");
asset["preview_url"] = Value::String(url.clone());
asset["thumbnail_url"] = Value::String(url);
} else if let Value::Object(object) = &mut asset {
object.remove("preview_url");
object.remove("thumbnail_url");
}
asset
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn normal_users_only_receive_matching_private_media() {
let query = BTreeMap::from([("include_tags".to_string(), "output".to_string())]);
assert!(asset_matches(&json!({ "tags": ["output"] }), &query));
assert!(!asset_matches(&json!({ "tags": ["input"] }), &query));
}
#[test]
fn models_are_only_added_when_the_query_allows_them() {
assert!(should_include_models(&BTreeMap::new()));
assert!(!should_include_models(&BTreeMap::from([(
"include_tags".to_string(),
"input".to_string()
)])));
assert!(!should_include_models(&BTreeMap::from([(
"exclude_tags".to_string(),
"models".to_string()
)])));
}
#[test]
fn sanitizing_preview_urls_keeps_reads_on_the_local_proxy() {
let asset = sanitize_asset(&json!({
"id": "11111111-1111-1111-1111-111111111111",
"preview_id": "22222222-2222-2222-2222-222222222222",
"preview_url": "https://cloud.example/private"
}));
assert_eq!(
asset["preview_url"],
"/workflow-ui/api/assets/22222222-2222-2222-2222-222222222222/content"
);
}
}
+26 -6
View File
@@ -147,15 +147,35 @@ pub(super) async fn userdata_get(
headers: HeaderMap,
Path(file): Path<String>,
) -> AppResult<Response> {
let caller = auth_caller(&state, &headers).await?;
let safe_file = safe_view_path(&file);
let optional_stylesheet = safe_file == "user.css";
let caller = match auth_caller(&state, &headers).await {
Ok(caller) => caller,
Err(_) if optional_stylesheet => {
return Ok((
StatusCode::OK,
[(header::CONTENT_TYPE, "text/css; charset=utf-8")],
Vec::<u8>::new(),
)
.into_response());
}
Err(err) => return Err(err),
};
let root = user_home(caller_user_id(&caller)).join("userdata");
let path = root.join(safe_view_path(&file));
let bytes = tokio::fs::read(&path)
.await
.map_err(|_| AppError::NotFound("User data file not found".into()))?;
let path = root.join(safe_file);
let bytes = match tokio::fs::read(&path).await {
Ok(bytes) => bytes,
Err(_) if optional_stylesheet => Vec::new(),
Err(_) => return Err(AppError::NotFound("User data file not found".into())),
};
let content_type = if optional_stylesheet {
"text/css; charset=utf-8"
} else {
"application/octet-stream"
};
Ok((
StatusCode::OK,
[(header::CONTENT_TYPE, "application/octet-stream")],
[(header::CONTENT_TYPE, content_type)],
bytes,
)
.into_response())
+46 -55
View File
@@ -1,14 +1,10 @@
use crate::{
error::{AppError, AppResult},
workflow::{
models::RunContext,
},
workflow::models::RunContext,
AppState,
};
use axum::{
extract::{
DefaultBodyLimit, Path, Query, State,
},
extract::{DefaultBodyLimit, Path, Query, State},
http::{header, HeaderMap, Method, StatusCode},
response::{IntoResponse, Response},
routing::{any, get, post},
@@ -24,8 +20,9 @@ use std::{
time::{Duration, Instant},
};
pub(super) static NATIVE_WORKFLOW_JSON_CACHE: Lazy<parking_lot::Mutex<HashMap<String, CachedJson>>> =
Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
pub(super) static NATIVE_WORKFLOW_JSON_CACHE: Lazy<
parking_lot::Mutex<HashMap<String, CachedJson>>,
> = Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
pub(super) const NATIVE_WORKFLOW_CACHE_MAX_AGE: Duration = Duration::from_secs(21600);
const WORKFLOW_UPLOAD_BODY_LIMIT_BYTES: usize = 256 * 1024 * 1024;
const OFFICIAL_WORKFLOW_TEMPLATES_BASE: &str =
@@ -120,7 +117,7 @@ use models::{
mod assets;
use assets::{
embedded_workflow_ui_for_config, official_workflow_template_asset, workflow_cloud_api_key,
workflow_ui_index, workflow_ui_runtime, workflow_ui_static,
workflow_ui_health, workflow_ui_index, workflow_ui_runtime, workflow_ui_static,
};
mod comfy_proxy;
use comfy_proxy::{
@@ -132,7 +129,7 @@ mod validate;
use validate::validate_workflow;
mod queue;
use queue::{
history, history_by_id, history_mutate, interrupt, internal_logs, internal_logs_raw, queue,
history, history_by_id, history_mutate, internal_logs, internal_logs_raw, interrupt, queue,
queue_mutate, queue_prompt, queue_status, view,
};
mod model_inject;
@@ -150,18 +147,17 @@ mod native_proxy;
use native_proxy::{
native_workflow_api_proxy, native_workflow_base, native_workflow_enabled, native_workflow_json,
native_workflow_json_cached, native_workflow_metadata_ttl, native_workflow_request,
native_workflow_response, native_workflow_status_ttl, native_workflow_video_response,
native_workflow_view_response, timestamp_ms, upload_image, upload_mask, workflow_cache_status,
native_workflow_status_ttl, native_workflow_video_response, native_workflow_view_response,
timestamp_ms, upload_image, upload_mask, workflow_cache_status,
};
mod cloud_assets;
mod run_convert;
use run_convert::{
node_output_for_comfy, run_outputs_for_comfy, workflow_graph_to_comfy_workflow,
workflow_run_to_history, workflow_run_to_history_task, workflow_run_to_queue_task,
};
mod ws;
use ws::{
publish_run_events, workflow_progress_callback_with_run_sender, workflow_ws,
};
use ws::{publish_run_events, workflow_progress_callback_with_run_sender, workflow_ws};
mod bootstrap;
use bootstrap::inject_bootstrap;
@@ -170,6 +166,7 @@ pub fn routes() -> Router<AppState> {
.route("/workflow-ui", get(workflow_ui_index))
.route("/workflow-ui/", get(workflow_ui_index))
.route("/api/workflow-ui/runtime", get(workflow_ui_runtime))
.route("/api/workflow-ui/health", get(workflow_ui_health))
.route("/workflow-ui/api/extensions", get(extensions))
.route(
"/workflow-ui/api/workflow_templates",
@@ -248,10 +245,20 @@ pub fn routes() -> Router<AppState> {
"/workflow-ui/templates/{*path}",
get(official_workflow_template_asset),
)
.route("/workflow-ui/user.css", get(workflow_ui_user_css))
.route("/workflow-ui/{*path}", any(workflow_ui_static))
.layer(DefaultBodyLimit::max(WORKFLOW_UPLOAD_BODY_LIMIT_BYTES))
}
async fn workflow_ui_user_css() -> Response {
(
StatusCode::OK,
[(header::CONTENT_TYPE, "text/css; charset=utf-8")],
"",
)
.into_response()
}
fn default_history_limit() -> i64 {
50
}
@@ -496,7 +503,10 @@ fn cache_control_for_path(path: &StdPath) -> &'static str {
}
}
pub(super) fn auth_caller_optional(state: &AppState, headers: &HeaderMap) -> Option<WorkflowCaller> {
pub(super) fn auth_caller_optional(
state: &AppState,
headers: &HeaderMap,
) -> Option<WorkflowCaller> {
caller_from_headers(state, headers).ok()
}
@@ -595,7 +605,10 @@ pub(super) fn workflow_validation_hints(errors: usize, warnings: usize) -> Vec<S
hints
}
pub(super) fn prompt_error_response(message: impl Into<String>, details: impl Into<String>) -> Response {
pub(super) fn prompt_error_response(
message: impl Into<String>,
details: impl Into<String>,
) -> Response {
let message = message.into();
let details = details.into();
(
@@ -742,7 +755,7 @@ async fn experiment_models_in_folder(
)));
}
let path = format!("/experiment/models/{}", simple_urlencode(&folder));
return native_workflow_json_cached(
let response = native_workflow_json_cached(
&state,
Method::GET,
&path,
@@ -750,6 +763,10 @@ async fn experiment_models_in_folder(
native_workflow_metadata_ttl(&state),
)
.await;
if folder.eq_ignore_ascii_case("video_models") && response.is_err() {
return Ok(Json(Value::Array(Vec::new())));
}
return response;
}
let route_models = workflow_route_models(&state).await.unwrap_or_default();
let api_model_items = workflow_folder_route_model_items(&folder, &route_models);
@@ -897,42 +914,16 @@ async fn view_video(
.into_response())
}
pub(super) async fn workflow_context(state: &AppState, caller: &WorkflowCaller) -> AppResult<RunContext> {
let bearer = match caller {
WorkflowCaller::Admin => state.builtin_keys.comfyui_agent.1.clone(),
WorkflowCaller::User { user_id } => {
let maybe_key =
sqlx::query_as::<_, crate::db::models::downstream_api_key::DownstreamApiKey>(
"SELECT * FROM downstream_api_keys
WHERE user_id = ?
AND COALESCE(enabled, 1) = 1
ORDER BY updated_at DESC, created_at DESC, id DESC
LIMIT 1",
)
.bind(user_id)
.fetch_optional(&*state.db)
.await?;
// Must run on the user's own key so consumption is billed to them.
// Falling back to the global builtin key would mix every user's
// cost into one shared account, breaking per-user accounting.
maybe_key.map(|key| key.key).ok_or_else(|| {
AppError::Forbidden(
"当前用户没有可用的 API Key,无法计费执行工作流,请先创建 API Key".into(),
)
})?
}
};
let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
"127.0.0.1".to_string()
} else {
state.config.server.host.clone()
};
Ok(RunContext {
user_id: caller_user_id(caller),
bearer,
base_url: format!("http://{}:{}/v1", host, state.config.server.port),
charge_points: !caller_is_admin(caller),
})
pub(super) async fn workflow_context(
state: &AppState,
caller: &WorkflowCaller,
) -> AppResult<RunContext> {
crate::workflow::routes::workflow_service_context(
state,
caller_user_id(caller),
!caller_is_admin(caller),
)
.await
}
pub(super) fn parse_workflow_cookie(raw: &str) -> Option<&str> {
@@ -945,9 +936,9 @@ pub(super) fn parse_workflow_cookie(raw: &str) -> Option<&str> {
#[cfg(test)]
mod tests {
use super::*;
use super::bootstrap::inject_enhanced_bootstrap;
use super::models::workflow_folder_models;
use super::*;
#[test]
fn video_model_folder_does_not_use_route_video_models() {
-1
View File
@@ -277,4 +277,3 @@ pub(super) fn normalize_model_folder_payload(value: &mut Value) {
}
}
}
+14 -5
View File
@@ -118,13 +118,13 @@ pub(super) fn is_chat_workflow_model(value: &str) -> bool {
}
pub(super) async fn workflow_route_models(state: &AppState) -> AppResult<Vec<String>> {
let rows = sqlx::query_scalar::<_, String>(
let rows = crate::db::sql::query_scalar::<_, String>(
"SELECT model_pattern
FROM token_routes
WHERE COALESCE(enabled, 1) = 1
AND model_pattern IS NOT NULL
AND model_pattern <> ''
ORDER BY model_pattern COLLATE NOCASE",
ORDER BY LOWER(model_pattern), model_pattern",
)
.fetch_all(&*state.db)
.await?;
@@ -212,7 +212,10 @@ pub(super) fn workflow_folder_has_route_models(folder: &str) -> bool {
)
}
pub(super) fn workflow_folder_route_model_items(folder: &str, route_models: &[String]) -> Vec<Value> {
pub(super) fn workflow_folder_route_model_items(
folder: &str,
route_models: &[String],
) -> Vec<Value> {
model_payload_items_from_names(workflow_folder_models(folder, route_models), 0)
}
@@ -242,7 +245,10 @@ pub(super) fn combo_choices_from_node_input(
.unwrap_or_default()
}
pub(super) fn combo_fields_from_node(object_info: &Value, node_type: &str) -> Vec<(String, Vec<String>)> {
pub(super) fn combo_fields_from_node(
object_info: &Value,
node_type: &str,
) -> Vec<(String, Vec<String>)> {
let mut fields = Vec::new();
let Some(node) = object_info.get(node_type) else {
return fields;
@@ -274,7 +280,10 @@ pub(super) fn combo_fields_from_node(object_info: &Value, node_type: &str) -> Ve
fields
}
pub(super) fn native_folder_models_from_object_info(folder: &str, object_info: &Value) -> Vec<String> {
pub(super) fn native_folder_models_from_object_info(
folder: &str,
object_info: &Value,
) -> Vec<String> {
let normalized = folder.trim().to_ascii_lowercase();
match normalized.as_str() {
"checkpoints" | "image_models" => {
+34 -4
View File
@@ -414,6 +414,22 @@ async fn native_workflow_file_response(
path: &str,
params: &HashMap<String, String>,
) -> AppResult<Response> {
if state.config.comfyui.uses_cloud_backend() && !matches!(caller, WorkflowCaller::Admin) {
let user_id = caller_user_id(caller).expect("non-admin workflow caller has user id");
let filename = params
.get("filename")
.map(String::as_str)
.unwrap_or("")
.trim();
if filename.is_empty()
|| !crate::services::comfy_cloud::ownership::asset_name_or_hash_is_owned_by(
&state.db, user_id, filename,
)
.await?
{
return Err(AppError::NotFound("Cloud asset not found".into()));
}
}
let query = build_query_string(params);
let response = native_workflow_request(state, Method::GET, &format!("{path}{query}"), None)
.send()
@@ -467,13 +483,13 @@ async fn archive_native_workflow_asset(
}
let user_id = caller_user_id(caller);
let source = "workflow-studio";
let already_exists = sqlx::query_scalar::<_, i64>(
let already_exists = crate::db::sql::query_scalar::<_, i64>(
"SELECT COUNT(*)
FROM workspace_assets
WHERE deleted_at IS NULL
AND (
(? = 1 AND user_id IS NULL)
OR (? = 0 AND user_id = ?)
(? AND user_id IS NULL)
OR (NOT ? AND user_id = ?)
)
AND source = ?
AND json_extract(metadata, '$.workflowUi.filename') = ?
@@ -525,10 +541,24 @@ pub(super) async fn native_workflow_api_proxy(
uri: Uri,
body: Bytes,
) -> AppResult<Response> {
let _ = auth_caller_optional(&state, &headers);
if !native_workflow_enabled(&state) {
return Err(AppError::NotFound("Workflow API route not found".into()));
}
if state.config.comfyui.uses_cloud_backend()
&& (path == "assets" || path.starts_with("assets/"))
{
return super::cloud_assets::handle_cloud_asset_request(
state, headers, path, method, uri, body,
)
.await;
}
if state.config.comfyui.uses_cloud_backend() && (path == "jobs" || path.starts_with("jobs/")) {
return super::cloud_assets::handle_cloud_job_request(
state, headers, path, method, uri, body,
)
.await;
}
let _ = auth_caller_optional(&state, &headers);
let target = match uri.query() {
Some(query) => format!("/{path}?{query}"),
None => format!("/{path}"),
+317 -17
View File
@@ -1,5 +1,6 @@
use crate::{
error::{AppError, AppResult},
services::comfy_cloud::ownership,
workflow::storage,
AppState,
};
@@ -11,14 +12,14 @@ use axum::{
};
use serde_json::{json, Value};
use sqlx::Row;
use std::{collections::HashMap, path::PathBuf, sync::mpsc};
use std::{collections::HashMap, path::PathBuf, sync::mpsc, time::Duration};
use super::auth::{auth_caller, caller_is_admin, caller_user_id};
use super::node_defs::normalize_prompt_payload;
use super::{
auth_caller_optional, guess_mime_from_path, inject_native_prompt_cloud_api_key,
native_workflow_json, native_workflow_json_cached, native_workflow_metadata_ttl,
native_workflow_response, native_workflow_view_response, prompt_contains_boosapi_nodes,
native_workflow_request, native_workflow_view_response, prompt_contains_boosapi_nodes,
prompt_error_response, publish_run_events, repair_dynamic_combo_nodes,
repair_resize_image_mask_nodes, repair_text_generate_ltx2_prompt_nodes, repair_wan2_api_nodes,
run_outputs_for_comfy, safe_view_path, simple_urlencode, user_home, workflow_context,
@@ -46,12 +47,17 @@ pub(super) async fn queue(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Json<Value>> {
if state.config.comfyui.uses_native_backend() {
let _ = auth_caller_optional(&state, &headers);
return native_workflow_json(&state, Method::GET, "/queue", None).await;
}
let caller =
auth_caller_optional(&state, &headers).unwrap_or(WorkflowCaller::User { user_id: 0 });
if state.config.comfyui.uses_native_backend() {
if !caller_is_admin(&caller) {
return Ok(Json(json!({
"Running": [],
"Pending": []
})));
}
return native_workflow_json(&state, Method::GET, "/queue", None).await;
}
let rows = storage::list_runs(
&state.db,
caller_user_id(&caller),
@@ -77,6 +83,13 @@ pub(super) async fn queue_mutate(
) -> AppResult<Json<Value>> {
let _ = auth_caller(&state, &headers).await?;
if state.config.comfyui.uses_native_backend() {
if state.config.comfyui.uses_cloud_backend()
&& !caller_is_admin(&auth_caller(&state, &headers).await?)
{
return Err(AppError::Forbidden(
"普通用户不能修改共享 Comfy Cloud 队列".into(),
));
}
let value = serde_json::to_value(&_payload).unwrap_or_else(|_| json!({}));
return native_workflow_json(&state, Method::POST, "/queue", Some(value)).await;
}
@@ -92,13 +105,15 @@ pub(super) async fn history(
headers: HeaderMap,
Query(query): Query<HistoryQuery>,
) -> AppResult<Json<Value>> {
let caller =
auth_caller_optional(&state, &headers).unwrap_or(WorkflowCaller::User { user_id: 0 });
if state.config.comfyui.uses_native_backend() {
let _ = auth_caller_optional(&state, &headers);
if !caller_is_admin(&caller) {
return Ok(Json(json!({ "History": [] })));
}
let path = format!("/history?max_items={}", query.max_items);
return native_workflow_json(&state, Method::GET, &path, None).await;
}
let caller =
auth_caller_optional(&state, &headers).unwrap_or(WorkflowCaller::User { user_id: 0 });
let rows = storage::list_runs(
&state.db,
caller_user_id(&caller),
@@ -146,13 +161,28 @@ pub(super) async fn history_by_id(
headers: HeaderMap,
Path(id): Path<String>,
) -> AppResult<Json<Value>> {
let caller =
auth_caller_optional(&state, &headers).unwrap_or(WorkflowCaller::User { user_id: 0 });
if state.config.comfyui.uses_native_backend() {
let _ = auth_caller_optional(&state, &headers);
if state.config.comfyui.uses_cloud_backend() {
let caller = auth_caller(&state, &headers).await?;
if !caller_is_admin(&caller) {
let user_id =
caller_user_id(&caller).expect("non-admin workflow caller has user id");
if !ownership::job_is_owned_by(&state.db, user_id, &id).await? {
return Err(AppError::NotFound("Workflow run not found".into()));
}
super::cloud_assets::sync_owned_assets(&state, user_id).await?;
}
let path = format!("/history/{}", simple_urlencode(&id));
return native_workflow_json(&state, Method::GET, &path, None).await;
}
if !caller_is_admin(&caller) {
return Err(AppError::NotFound("Workflow run not found".into()));
}
let path = format!("/history/{}", simple_urlencode(&id));
return native_workflow_json(&state, Method::GET, &path, None).await;
}
let caller =
auth_caller_optional(&state, &headers).unwrap_or(WorkflowCaller::User { user_id: 0 });
let run = storage::get_run(
&state.db,
&id,
@@ -172,8 +202,9 @@ pub(super) async fn queue_prompt(
if state.config.comfyui.uses_native_backend() {
if state.config.comfyui.uses_cloud_backend() {
repair_native_prompt_payload(&state, &mut payload).await;
return native_workflow_response(
return native_prompt_response_with_billing(
&state,
&caller,
Method::POST,
"/prompt",
Some(serde_json::to_value(payload).unwrap_or_else(|_| json!({}))),
@@ -184,8 +215,9 @@ pub(super) async fn queue_prompt(
return queue_boosapi_prompt(state, caller, payload).await;
}
repair_native_prompt_payload(&state, &mut payload).await;
return native_workflow_response(
return native_prompt_response_with_billing(
&state,
&caller,
Method::POST,
"/prompt",
Some(serde_json::to_value(payload).unwrap_or_else(|_| json!({}))),
@@ -195,6 +227,179 @@ pub(super) async fn queue_prompt(
queue_boosapi_prompt(state, caller, payload).await
}
async fn native_prompt_response_with_billing(
state: &AppState,
caller: &WorkflowCaller,
method: Method,
path: &str,
body: Option<Value>,
) -> AppResult<Response> {
let mut charge =
prepare_native_prompt_charge(&state.db, caller, state.config.comfyui.workflow_points_cost)
.await?;
let upstream = native_workflow_request(state, method, path, body)
.send()
.await
.map_err(|err| AppError::UpstreamError(format!("ComfyUI backend unreachable: {err}")));
let upstream = match upstream {
Ok(response) => response,
Err(err) => {
refund_native_prompt_charge(&state.db, charge.take()).await;
return Err(err);
}
};
let status =
StatusCode::from_u16(upstream.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
let content_type = upstream
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or("application/json")
.to_string();
let bytes = match upstream.bytes().await {
Ok(bytes) => bytes,
Err(err) => {
refund_native_prompt_charge(&state.db, charge.take()).await;
return Err(AppError::UpstreamError(err.to_string()));
}
};
if !status.is_success() {
refund_native_prompt_charge(&state.db, charge.take()).await;
} else if state.config.comfyui.uses_cloud_backend() {
if let Ok(value) = serde_json::from_slice::<Value>(&bytes) {
ownership::record_prompt_response(state, caller_user_id(caller), &value).await?;
}
}
if status.is_success() {
if let Some(charge) = charge {
let prompt_id = serde_json::from_slice::<Value>(&bytes)
.ok()
.and_then(|value| {
value
.get("prompt_id")
.and_then(Value::as_str)
.map(str::to_string)
});
if let Some(prompt_id) = prompt_id {
if let Err(err) = crate::services::comfy_billing::bind_hold_ref(
&state.db,
&charge.public_id,
&prompt_id,
)
.await
{
let _ = crate::services::comfy_billing::release(
&state.db,
&charge.public_id,
Some("prompt_binding_failed"),
)
.await;
return Err(err);
}
spawn_native_prompt_billing_monitor(state.clone(), charge.public_id, prompt_id);
} else {
let _ = crate::services::comfy_billing::release(
&state.db,
&charge.public_id,
Some("missing_prompt_id"),
)
.await;
}
}
}
Response::builder()
.status(status)
.header(axum::http::header::CONTENT_TYPE, content_type)
.body(axum::body::Body::from(bytes))
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))
}
async fn prepare_native_prompt_charge(
db: &sqlx::AnyPool,
caller: &WorkflowCaller,
cost: f64,
) -> AppResult<Option<crate::services::comfy_billing::ComfyHold>> {
if caller_is_admin(caller) {
return Ok(None);
}
let Some(user_id) = caller_user_id(caller) else {
return Err(AppError::Forbidden(
"工作流需要登录用户并绑定计费账户后才能执行".into(),
));
};
crate::services::comfy_billing::reserve_task(
db,
Some(user_id),
crate::services::comfy_billing::TASK_NATIVE_WORKFLOW,
&format!("prompt-submit-{}", uuid::Uuid::new_v4().simple()),
Some(cost),
)
.await
}
async fn refund_native_prompt_charge(
db: &sqlx::AnyPool,
charge: Option<crate::services::comfy_billing::ComfyHold>,
) {
if let Some(charge) = charge {
if let Err(err) = crate::services::comfy_billing::release(
db,
&charge.public_id,
Some("submission_failed"),
)
.await
{
tracing::warn!(
user_id = charge.user_id,
amount_micros = charge.amount_micros,
error = %err,
"failed to refund native workflow prompt points"
);
}
}
}
fn spawn_native_prompt_billing_monitor(state: AppState, hold_id: String, prompt_id: String) {
tokio::spawn(async move {
let deadline = tokio::time::Instant::now() + Duration::from_secs(24 * 60 * 60);
loop {
if tokio::time::Instant::now() >= deadline {
let _ = crate::services::comfy_billing::release(
&state.db,
&hold_id,
Some("workflow_status_timeout"),
)
.await;
break;
}
let path = format!("/history/{}", simple_urlencode(&prompt_id));
if let Ok(history) = native_workflow_json(&state, Method::GET, &path, None).await {
let entry = history.get(&prompt_id).unwrap_or(&history);
let status = entry.pointer("/status/status_str").and_then(Value::as_str);
let completed = entry.get("outputs").is_some_and(|value| {
value.as_object().is_some_and(|outputs| !outputs.is_empty())
});
if status == Some("error") {
let _ = crate::services::comfy_billing::release(
&state.db,
&hold_id,
Some("workflow_failed"),
)
.await;
break;
}
if status == Some("success") || completed {
let _ =
crate::services::comfy_billing::settle_success(&state.db, &hold_id).await;
break;
}
}
tokio::time::sleep(Duration::from_secs(3)).await;
}
});
}
pub(super) async fn repair_native_prompt_payload(state: &AppState, payload: &mut PromptPayload) {
inject_native_prompt_cloud_api_key(state, payload);
let object_info = native_workflow_json_cached(
@@ -319,8 +524,15 @@ pub(super) async fn interrupt(
State(state): State<AppState>,
headers: HeaderMap,
) -> AppResult<Json<Value>> {
let _ = auth_caller_optional(&state, &headers);
if state.config.comfyui.uses_native_backend() {
if state.config.comfyui.uses_cloud_backend()
&& !caller_is_admin(&auth_caller(&state, &headers).await?)
{
return Err(AppError::Forbidden(
"普通用户不能中断共享 Comfy Cloud 队列".into(),
));
}
let _ = auth_caller_optional(&state, &headers);
return native_workflow_json(&state, Method::POST, "/interrupt", Some(json!({}))).await;
}
Ok(Json(json!({ "success": true })))
@@ -367,15 +579,21 @@ pub(super) async fn view(
.into_response());
}
let row = sqlx::query(
let is_admin = caller_is_admin(&caller);
let scoped_user_id = caller_user_id(&caller);
let row = crate::db::sql::query(
"SELECT file_path, remote_url, mime_type
FROM workspace_assets
WHERE file_name = ? OR public_id = ?
WHERE (file_name = ? OR public_id = ?)
AND deleted_at IS NULL
AND (? = 1 OR user_id = ?)
ORDER BY id DESC
LIMIT 1",
)
.bind(&file)
.bind(&file)
.bind(if is_admin { 1_i64 } else { 0_i64 })
.bind(scoped_user_id)
.fetch_optional(&*state.db)
.await?;
let Some(row) = row else {
@@ -433,3 +651,85 @@ pub(super) async fn internal_logs_raw(
))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::services::billing::ledger;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
async fn mem_pool() -> sqlx::AnyPool {
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
ledger::ensure_ledger_schema(&pool).await.unwrap();
pool
}
#[tokio::test]
async fn native_prompt_charge_skips_admin() {
let pool = mem_pool().await;
let charge = prepare_native_prompt_charge(&pool, &WorkflowCaller::Admin, 10.0)
.await
.unwrap();
assert!(charge.is_none());
}
#[tokio::test]
async fn native_prompt_charge_rejects_user_without_points() {
let pool = mem_pool().await;
let err = prepare_native_prompt_charge(&pool, &WorkflowCaller::User { user_id: 7 }, 10.0)
.await
.unwrap_err();
assert!(matches!(
err,
AppError::BillingRequired {
billing_track: "comfy",
..
}
));
}
#[tokio::test]
async fn native_prompt_charge_debits_and_refunds_points() {
let pool = mem_pool().await;
crate::services::comfy_billing::grant(
&pool,
crate::services::comfy_billing::ComfyGrant {
user_id: 7,
amount_micros: 20_000_000,
grant_kind: "balance",
source_type: "test",
source_id: None,
payment_order_id: None,
starts_at: None,
expires_at: None,
idempotency_key: "native-prompt-test-balance",
operator: "test",
},
)
.await
.unwrap();
let charge =
prepare_native_prompt_charge(&pool, &WorkflowCaller::User { user_id: 7 }, 10.0)
.await
.unwrap()
.unwrap();
let reserved = crate::services::comfy_billing::balance(&pool, 7)
.await
.unwrap();
assert_eq!(reserved.balance_micros, 10_000_000);
assert_eq!(reserved.reserved_micros, 10_000_000);
refund_native_prompt_charge(&pool, Some(charge)).await;
let refunded = crate::services::comfy_billing::balance(&pool, 7)
.await
.unwrap();
assert_eq!(refunded.balance_micros, 20_000_000);
assert_eq!(refunded.reserved_micros, 0);
}
}
+6 -2
View File
@@ -23,7 +23,9 @@ pub(super) fn workflow_run_to_history(run: crate::workflow::models::WorkflowRunS
})
}
pub(super) fn workflow_run_to_history_task(run: crate::workflow::models::WorkflowRunSummary) -> Value {
pub(super) fn workflow_run_to_history_task(
run: crate::workflow::models::WorkflowRunSummary,
) -> Value {
let history = workflow_run_to_history(run);
json!({
"taskType": "History",
@@ -33,7 +35,9 @@ pub(super) fn workflow_run_to_history_task(run: crate::workflow::models::Workflo
})
}
pub(super) fn workflow_run_to_queue_task(run: crate::workflow::models::WorkflowRunSummary) -> Value {
pub(super) fn workflow_run_to_queue_task(
run: crate::workflow::models::WorkflowRunSummary,
) -> Value {
json!({
"taskType": "Running",
"prompt": workflow_run_to_prompt_tuple(&run),
-2
View File
@@ -301,5 +301,3 @@ async fn validate_workflow_input_file(
"localPath": path,
})
}
+6 -3
View File
@@ -18,12 +18,12 @@ use tokio::sync::broadcast;
use tokio_tungstenite::{connect_async, tungstenite};
use tungstenite::client::IntoClientRequest;
use super::auth::{auth_caller, caller_is_admin};
use super::workflow_node_label;
use super::{
build_query_string, native_workflow_base, native_workflow_enabled, node_output_for_comfy,
timestamp_ms,
};
use super::auth::auth_caller;
static WORKFLOW_WS_SESSIONS: Lazy<parking_lot::Mutex<HashMap<String, broadcast::Sender<Value>>>> =
Lazy::new(|| parking_lot::Mutex::new(HashMap::new()));
@@ -297,13 +297,16 @@ pub(super) async fn workflow_ws(
ws: WebSocketUpgrade,
Query(query): Query<HashMap<String, String>>,
) -> AppResult<Response> {
let _ = auth_caller(&state, &headers).await?;
let caller = auth_caller(&state, &headers).await?;
let sid = query
.get("clientId")
.or_else(|| query.get("client_id"))
.cloned()
.unwrap_or_else(|| "boosapi-workflow".to_string());
if native_workflow_enabled(&state) && !state.config.comfyui.uses_runninghub_backend() {
if native_workflow_enabled(&state)
&& !state.config.comfyui.uses_runninghub_backend()
&& !(state.config.comfyui.uses_cloud_backend() && !caller_is_admin(&caller))
{
return Ok(
ws.on_upgrade(move |socket| native_workflow_ws_bridge(state, socket, query, sid))
);
+135 -67
View File
@@ -146,7 +146,7 @@ fn normalize_query(
(page_size, offset, source, kind, q, user_id)
}
fn row_to_json(row: sqlx::sqlite::SqliteRow) -> Value {
fn row_to_json(row: sqlx::any::AnyRow) -> Value {
use sqlx::Row;
let metadata_raw: String = row.try_get("metadata").unwrap_or_else(|_| "{}".to_string());
let metadata = serde_json::from_str::<Value>(&metadata_raw).unwrap_or_else(|_| json!({}));
@@ -187,6 +187,21 @@ const WORKSPACE_OUTPUT_FILTER_SQL: &str = r#"
)
"#;
fn workspace_gallery_index(
scope_requires_user: bool,
has_source: bool,
has_kind: bool,
) -> &'static str {
match (scope_requires_user, has_source, has_kind) {
(true, true, _) => "workspace_assets_gallery_user_source_created_idx",
(true, false, true) => "workspace_assets_gallery_user_kind_created_idx",
(true, false, false) => "workspace_assets_gallery_user_created_idx",
(false, true, _) => "workspace_assets_gallery_source_created_idx",
(false, false, true) => "workspace_assets_gallery_kind_created_idx",
(false, false, false) => "workspace_assets_gallery_created_idx",
}
}
async fn list_assets(
state: AppState,
caller: WorkspaceCaller,
@@ -198,40 +213,57 @@ async fn list_assets(
WorkspaceCaller::User { user_id } => Some(user_id),
};
let is_admin = matches!(caller, WorkspaceCaller::Admin);
let list_sql = format!(
"SELECT *
FROM workspace_assets
WHERE deleted_at IS NULL
AND ((? = 1 AND (? IS NULL OR user_id = ?)) OR (? = 0 AND user_id = ?))
{WORKSPACE_OUTPUT_FILTER_SQL}
AND (? IS NULL OR source = ?)
AND (? IS NULL OR kind = ?)
AND (
? IS NULL
OR title LIKE ?
let scope_requires_user = !is_admin || scoped_user_id.is_some();
let scope_sql = if scope_requires_user {
"AND user_id = ?"
} else {
""
};
let source_sql = if source.is_some() {
"AND source = ?"
} else {
""
};
let kind_sql = if kind.is_some() { "AND kind = ?" } else { "" };
let search_sql = if q.is_some() {
"AND (
title LIKE ?
OR prompt LIKE ?
OR model LIKE ?
OR source LIKE ?
)
ORDER BY datetime(created_at) DESC, id DESC
)"
} else {
""
};
let gallery_index =
workspace_gallery_index(scope_requires_user, source.is_some(), kind.is_some());
let list_sql = format!(
"SELECT *
FROM workspace_assets INDEXED BY {gallery_index}
WHERE deleted_at IS NULL
{scope_sql}
{WORKSPACE_OUTPUT_FILTER_SQL}
{source_sql}
{kind_sql}
{search_sql}
ORDER BY created_at DESC, id DESC
LIMIT ? OFFSET ?"
);
let items = sqlx::query(&list_sql)
.bind(is_admin)
.bind(scoped_user_id)
.bind(scoped_user_id)
.bind(is_admin)
.bind(scoped_user_id)
.bind(&source)
.bind(&source)
.bind(&kind)
.bind(&kind)
.bind(&q)
.bind(&q)
.bind(&q)
.bind(&q)
.bind(&q)
let mut list_query = crate::db::sql::query(&list_sql);
if scope_requires_user {
list_query = list_query.bind(scoped_user_id);
}
if let Some(source) = &source {
list_query = list_query.bind(source);
}
if let Some(kind) = &kind {
list_query = list_query.bind(kind);
}
if let Some(q) = &q {
list_query = list_query.bind(q).bind(q).bind(q).bind(q);
}
let items = list_query
.bind(page_size)
.bind(offset)
.map(row_to_json)
@@ -240,37 +272,28 @@ async fn list_assets(
let count_sql = format!(
"SELECT COUNT(*)
FROM workspace_assets
FROM workspace_assets INDEXED BY {gallery_index}
WHERE deleted_at IS NULL
AND ((? = 1 AND (? IS NULL OR user_id = ?)) OR (? = 0 AND user_id = ?))
{scope_sql}
{WORKSPACE_OUTPUT_FILTER_SQL}
AND (? IS NULL OR source = ?)
AND (? IS NULL OR kind = ?)
AND (
? IS NULL
OR title LIKE ?
OR prompt LIKE ?
OR model LIKE ?
OR source LIKE ?
)"
{source_sql}
{kind_sql}
{search_sql}"
);
let total: i64 = sqlx::query_scalar(&count_sql)
.bind(is_admin)
.bind(scoped_user_id)
.bind(scoped_user_id)
.bind(is_admin)
.bind(scoped_user_id)
.bind(&source)
.bind(&source)
.bind(&kind)
.bind(&kind)
.bind(&q)
.bind(&q)
.bind(&q)
.bind(&q)
.bind(&q)
.fetch_one(&*state.db)
.await?;
let mut count_query = crate::db::sql::query_scalar(&count_sql);
if scope_requires_user {
count_query = count_query.bind(scoped_user_id);
}
if let Some(source) = &source {
count_query = count_query.bind(source);
}
if let Some(kind) = &kind {
count_query = count_query.bind(kind);
}
if let Some(q) = &q {
count_query = count_query.bind(q).bind(q).bind(q).bind(q);
}
let total: i64 = count_query.fetch_one(&*state.db).await?;
Ok(Json(json!({
"success": true,
@@ -285,13 +308,17 @@ async fn get_asset_row(
state: &AppState,
public_id: &str,
caller: WorkspaceCaller,
) -> AppResult<sqlx::sqlite::SqliteRow> {
) -> AppResult<sqlx::any::AnyRow> {
let scoped_user_id = match caller {
WorkspaceCaller::Admin => None,
WorkspaceCaller::User { user_id } => Some(user_id),
};
let is_admin = matches!(caller, WorkspaceCaller::Admin);
let row = sqlx::query(
let is_admin = if matches!(caller, WorkspaceCaller::Admin) {
1_i64
} else {
0_i64
};
let row = crate::db::sql::query(
"SELECT *
FROM workspace_assets
WHERE public_id = ?
@@ -331,10 +358,14 @@ async fn delete_asset(
WorkspaceCaller::Admin => None,
WorkspaceCaller::User { user_id } => Some(user_id),
};
let is_admin = matches!(caller, WorkspaceCaller::Admin);
let result = sqlx::query(
let is_admin = if matches!(caller, WorkspaceCaller::Admin) {
1_i64
} else {
0_i64
};
let result = crate::db::sql::query(
"UPDATE workspace_assets
SET deleted_at = datetime('now'), updated_at = datetime('now')
SET deleted_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP
WHERE public_id = ?
AND deleted_at IS NULL
AND ((? = 1 AND (? IS NULL OR user_id = ?)) OR (? = 0 AND user_id = ?))",
@@ -377,14 +408,18 @@ async fn delete_assets_batch(
WorkspaceCaller::Admin => None,
WorkspaceCaller::User { user_id } => Some(user_id),
};
let is_admin = matches!(caller, WorkspaceCaller::Admin);
let is_admin = if matches!(caller, WorkspaceCaller::Admin) {
1_i64
} else {
0_i64
};
let mut deleted_ids = Vec::new();
let mut failed_ids = Vec::new();
for public_id in public_ids {
let result = sqlx::query(
let result = crate::db::sql::query(
"UPDATE workspace_assets
SET deleted_at = datetime('now'), updated_at = datetime('now')
SET deleted_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP
WHERE public_id = ?
AND deleted_at IS NULL
AND ((? = 1 AND (? IS NULL OR user_id = ?)) OR (? = 0 AND user_id = ?))",
@@ -550,3 +585,36 @@ async fn batch_delete_user_assets(
let caller = resolve_caller(&headers, &state, false)?;
delete_assets_batch(state, caller, payload.public_ids).await
}
#[cfg(test)]
mod tests {
use super::workspace_gallery_index;
#[test]
fn gallery_index_matches_active_scope_and_filter_prefixes() {
assert_eq!(
workspace_gallery_index(false, false, false),
"workspace_assets_gallery_created_idx"
);
assert_eq!(
workspace_gallery_index(true, false, false),
"workspace_assets_gallery_user_created_idx"
);
assert_eq!(
workspace_gallery_index(false, true, false),
"workspace_assets_gallery_source_created_idx"
);
assert_eq!(
workspace_gallery_index(false, false, true),
"workspace_assets_gallery_kind_created_idx"
);
assert_eq!(
workspace_gallery_index(true, true, true),
"workspace_assets_gallery_user_source_created_idx"
);
assert_eq!(
workspace_gallery_index(true, false, true),
"workspace_assets_gallery_user_kind_created_idx"
);
}
}
+273
View File
@@ -0,0 +1,273 @@
use crate::{error::AppResult, AppState};
use futures::{stream, StreamExt};
use serde::Serialize;
use std::time::Instant;
use tokio::{task::JoinHandle, time::Duration};
pub(crate) const HEALTH_CHECK_MODEL: &str = "__account_health_check__";
const HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(15 * 60);
const HEALTH_CHECK_TIMEOUT: Duration = Duration::from_secs(20);
const HEALTH_CHECK_CONCURRENCY: usize = 4;
const HEALTH_CHECK_SETTING: &str = "runtime.model_availability_probe_enabled";
#[derive(Debug, Clone)]
struct ProbeTarget {
account_id: i64,
access_token: String,
site_id: i64,
site_url: String,
}
#[derive(Debug, Default, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountHealthProbeSummary {
pub checked: usize,
pub healthy: usize,
pub unhealthy: usize,
pub skipped: usize,
}
#[derive(Debug)]
struct ProbeOutcome {
healthy: bool,
}
pub fn spawn(state: AppState) -> JoinHandle<()> {
tokio::spawn(async move {
let mut ticker = tokio::time::interval_at(
tokio::time::Instant::now() + HEALTH_CHECK_INTERVAL,
HEALTH_CHECK_INTERVAL,
);
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
ticker.tick().await;
match periodic_probe_enabled(&state).await {
Ok(true) => match probe_accounts(&state, None).await {
Ok(summary) => tracing::info!(
checked = summary.checked,
healthy = summary.healthy,
unhealthy = summary.unhealthy,
skipped = summary.skipped,
"periodic account health detection completed"
),
Err(err) => tracing::warn!(%err, "periodic account health detection failed"),
},
Ok(false) => {}
Err(err) => tracing::warn!(%err, "failed to read account health detection setting"),
}
}
})
}
pub fn spawn_once(state: AppState, account_ids: Vec<i64>) {
tokio::spawn(async move {
for account_id in account_ids {
if let Err(err) = probe_accounts(&state, Some(account_id)).await {
tracing::warn!(account_id, %err, "automatic account health detection failed");
}
}
});
}
pub async fn probe_accounts(
state: &AppState,
account_id: Option<i64>,
) -> AppResult<AccountHealthProbeSummary> {
let targets = load_targets(state, account_id).await?;
let skipped = usize::from(account_id.is_some() && targets.is_empty());
let outcomes = stream::iter(targets.into_iter().map(|target| {
let state = state.clone();
async move { probe_one(&state, target).await }
}))
.buffer_unordered(HEALTH_CHECK_CONCURRENCY)
.collect::<Vec<_>>()
.await;
let mut summary = AccountHealthProbeSummary {
skipped,
..AccountHealthProbeSummary::default()
};
for outcome in outcomes {
summary.checked += 1;
match outcome {
Ok(ProbeOutcome { healthy: true }) => summary.healthy += 1,
Ok(ProbeOutcome { healthy: false }) | Err(_) => summary.unhealthy += 1,
}
}
Ok(summary)
}
async fn load_targets(state: &AppState, account_id: Option<i64>) -> AppResult<Vec<ProbeTarget>> {
let mut sql = String::from(
"SELECT a.id, a.access_token, s.id, s.url
FROM accounts a
JOIN sites s ON s.id = a.site_id
WHERE a.status = 'active'
AND s.status = 'active'
AND LOWER(s.platform) NOT IN ('jimeng', 'seedance', 'comfy_cloud')",
);
if account_id.is_some() {
sql.push_str(" AND a.id = ?");
}
sql.push_str(" ORDER BY a.id");
let mut query = crate::db::sql::query_as::<_, (i64, String, i64, String)>(&sql);
if let Some(account_id) = account_id {
query = query.bind(account_id);
}
let rows = query.fetch_all(&*state.db).await?;
Ok(rows
.into_iter()
.filter(|(_, token, _, _)| !token.trim().is_empty())
.map(
|(account_id, access_token, site_id, site_url)| ProbeTarget {
account_id,
access_token,
site_id,
site_url,
},
)
.collect())
}
async fn probe_one(state: &AppState, target: ProbeTarget) -> AppResult<ProbeOutcome> {
let base_url = resolve_api_base_url(state, &target).await?;
let models_url = crate::proxy::executors::openai_compatible_url(&base_url, "/models");
let started_at = Instant::now();
let response = tokio::time::timeout(
HEALTH_CHECK_TIMEOUT,
state
.http
.get(&models_url)
.bearer_auth(target.access_token.trim())
.header(reqwest::header::ACCEPT, "application/json")
.send(),
)
.await;
let latency_ms = started_at.elapsed().as_millis().min(i64::MAX as u128) as i64;
let (healthy, http_status, error_message) = match response {
Ok(Ok(response)) if response.status().is_success() => {
let status = response.status().as_u16();
match response.json::<serde_json::Value>().await {
Ok(body) if body.get("data").is_some_and(serde_json::Value::is_array) => {
(true, Some(i64::from(status)), None)
}
Ok(_) => (
false,
Some(i64::from(status)),
Some("自动检测失败:/models 响应缺少 data 数组".to_string()),
),
Err(err) => (
false,
Some(i64::from(status)),
Some(format!("自动检测失败:/models 响应不是有效 JSON:{err}")),
),
}
}
Ok(Ok(response)) => {
let status = response.status().as_u16();
let body = response.text().await.unwrap_or_default();
let detail = body.trim().chars().take(240).collect::<String>();
let message = if detail.is_empty() {
format!("自动检测失败:HTTP {status}")
} else {
format!("自动检测失败:HTTP {status} {detail}")
};
(false, Some(i64::from(status)), Some(message))
}
Ok(Err(err)) => (false, None, Some(format!("自动检测连接失败:{err}"))),
Err(_) => (false, None, Some("自动检测超时(20 秒)".to_string())),
};
record_probe_result(
state,
target.account_id,
healthy,
http_status,
latency_ms,
error_message.as_deref(),
)
.await?;
Ok(ProbeOutcome { healthy })
}
async fn resolve_api_base_url(state: &AppState, target: &ProbeTarget) -> AppResult<String> {
let endpoint: Option<String> = crate::db::sql::query_scalar(
"SELECT url
FROM site_api_endpoints
WHERE site_id = ?
AND COALESCE(enabled, 1) = 1
AND (cooldown_until IS NULL OR cooldown_until <= CURRENT_TIMESTAMP)
ORDER BY COALESCE(sort_order, 0), id
LIMIT 1",
)
.bind(target.site_id)
.fetch_optional(&*state.db)
.await?;
Ok(endpoint
.as_deref()
.unwrap_or(&target.site_url)
.trim()
.trim_end_matches('/')
.to_string())
}
async fn record_probe_result(
state: &AppState,
account_id: i64,
healthy: bool,
http_status: Option<i64>,
latency_ms: i64,
error_message: Option<&str>,
) -> AppResult<()> {
crate::db::sql::query(
"INSERT INTO proxy_logs (
account_id, model_requested, status, http_status, latency_ms,
error_message, retry_count, is_stream, client_family,
client_app_id, client_app_name, client_confidence, billing_details
) VALUES (?, ?, ?, ?, ?, ?, 0, 0, 'health-check',
'boosapi-health-check', 'BoosAPI 自动检测', 'high', ?)",
)
.bind(account_id)
.bind(HEALTH_CHECK_MODEL)
.bind(if healthy { "success" } else { "failed" })
.bind(http_status)
.bind(latency_ms)
.bind(error_message)
.bind(r#"{"source":"account-health","billable":false}"#)
.execute(&*state.db)
.await?;
Ok(())
}
async fn periodic_probe_enabled(state: &AppState) -> AppResult<bool> {
let value: Option<String> =
crate::db::sql::query_scalar("SELECT value FROM settings WHERE key = ?")
.bind(HEALTH_CHECK_SETTING)
.fetch_optional(&*state.db)
.await?;
Ok(value.as_deref().is_some_and(parse_bool_setting))
}
fn parse_bool_setting(value: &str) -> bool {
matches!(
value.trim().to_ascii_lowercase().as_str(),
"1" | "true" | "yes" | "on"
)
}
#[cfg(test)]
mod tests {
use super::parse_bool_setting;
#[test]
fn parses_runtime_probe_setting() {
assert!(parse_bool_setting("true"));
assert!(parse_bool_setting(" 1 "));
assert!(!parse_bool_setting("false"));
assert!(!parse_bool_setting("0"));
}
}
+312
View File
@@ -0,0 +1,312 @@
use base64::{engine::general_purpose, Engine as _};
use office_oxide::{Document, DocumentFormat};
use once_cell::sync::Lazy;
use serde::Serialize;
use sha2::{Digest, Sha256};
use sqlx::Row;
use std::{collections::HashMap, io::Cursor, sync::Arc};
use tokio::sync::{Mutex, Semaphore};
pub const MAX_ATTACHMENT_BYTES: usize = 64 * 1024 * 1024;
pub const MAX_ATTACHMENT_COUNT: usize = 8;
pub const MAX_EXTRACTED_CHARS: usize = 2_000_000;
pub const MAX_COMFY_FILE_CHARS: usize = 120_000;
pub const MAX_COMFY_TOTAL_CHARS: usize = 240_000;
pub const MAX_INLINE_IMAGE_BYTES: usize = 12 * 1024 * 1024;
static EXTRACTION_LIMIT: Lazy<Semaphore> = Lazy::new(|| Semaphore::new(2));
static EXTRACTION_CACHE: Lazy<Mutex<HashMap<String, Arc<ExtractionResult>>>> =
Lazy::new(|| Mutex::new(HashMap::new()));
#[derive(Debug, Clone)]
pub struct StoredAttachment {
pub asset_id: String,
pub name: String,
pub mime_type: String,
pub media_kind: String,
pub file_path: String,
pub size: usize,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ExtractionResult {
pub status: &'static str,
pub char_count: usize,
pub truncated: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub warning: Option<String>,
#[serde(skip)]
pub content: String,
}
impl ExtractionResult {
fn not_applicable() -> Self {
Self {
status: "not_applicable",
char_count: 0,
truncated: false,
warning: None,
content: String::new(),
}
}
fn failed(message: impl Into<String>) -> Self {
Self {
status: "failed",
char_count: 0,
truncated: false,
warning: Some(message.into()),
content: String::new(),
}
}
}
pub fn extension(name: &str) -> &str {
name.rsplit_once('.')
.map(|(_, ext)| ext.trim())
.unwrap_or("")
}
pub fn media_kind(name: &str, mime_type: &str) -> &'static str {
let mime = mime_type.to_ascii_lowercase();
if mime.starts_with("image/") {
return "image";
}
if mime.starts_with("video/") {
return "video";
}
if mime.starts_with("audio/") {
return "audio";
}
match extension(name).to_ascii_lowercase().as_str() {
"pdf" | "doc" | "docx" | "xls" | "xlsx" | "ppt" | "pptx" => "document",
"txt" | "md" | "markdown" | "csv" | "json" => "text",
_ if mime.starts_with("text/") || mime == "application/json" => "text",
_ => "file",
}
}
pub fn is_supported(name: &str, mime_type: &str) -> bool {
matches!(
media_kind(name, mime_type),
"image" | "video" | "audio" | "document" | "text"
)
}
fn truncate_chars(mut content: String, max: usize) -> (String, bool) {
if content.chars().count() <= max {
return (content, false);
}
content = content.chars().take(max).collect();
(content, true)
}
fn normalize_text(value: String) -> String {
value.replace("\r\n", "\n").replace('\r', "\n").trim().to_string()
}
fn extract_sync(name: &str, mime_type: &str, bytes: &[u8]) -> ExtractionResult {
let kind = media_kind(name, mime_type);
if !matches!(kind, "document" | "text") {
return ExtractionResult::not_applicable();
}
let ext = extension(name).to_ascii_lowercase();
let extracted = match ext.as_str() {
"pdf" => match pdf_extract::extract_text_from_mem_by_pages(bytes) {
Ok(pages) => pages
.into_iter()
.enumerate()
.map(|(index, page)| format!("## 第 {} 页\n\n{}", index + 1, page.trim()))
.collect::<Vec<_>>()
.join("\n\n"),
Err(error) => return ExtractionResult::failed(format!("PDF 解析失败:{error}")),
},
"doc" | "docx" | "xls" | "xlsx" | "ppt" | "pptx" => {
let Some(format) = DocumentFormat::from_extension(&ext) else {
return ExtractionResult::failed("无法识别 Office 文件格式");
};
match Document::from_reader(Cursor::new(bytes.to_vec()), format) {
Ok(document) => document.to_markdown(),
Err(error) => return ExtractionResult::failed(format!("Office 文档解析失败:{error}")),
}
}
"txt" | "md" | "markdown" | "csv" | "json" => {
match String::from_utf8(bytes.to_vec()) {
Ok(text) => text,
Err(_) => return ExtractionResult::failed("文本文件不是有效的 UTF-8 编码"),
}
}
_ if mime_type.starts_with("text/") || mime_type == "application/json" => {
match String::from_utf8(bytes.to_vec()) {
Ok(text) => text,
Err(_) => return ExtractionResult::failed("文本文件不是有效的 UTF-8 编码"),
}
}
_ => return ExtractionResult::failed("不支持读取该文件格式"),
};
let extracted = normalize_text(extracted);
if extracted.is_empty() {
let warning = if ext == "pdf" {
"PDF 中没有可提取文字;扫描版 PDF 暂不支持 OCR"
} else {
"文档中没有可提取文字"
};
return ExtractionResult::failed(warning);
}
let original_count = extracted.chars().count();
let (content, truncated) = truncate_chars(extracted, MAX_EXTRACTED_CHARS);
ExtractionResult {
status: "ready",
char_count: original_count.min(MAX_EXTRACTED_CHARS),
truncated,
warning: truncated.then(|| "文档内容过长,已截断至 200 万字符".to_string()),
content,
}
}
pub async fn extract_bytes(name: &str, mime_type: &str, bytes: &[u8]) -> ExtractionResult {
if !matches!(media_kind(name, mime_type), "document" | "text") {
return ExtractionResult::not_applicable();
}
let mut digest = Sha256::new();
digest.update(extension(name).to_ascii_lowercase().as_bytes());
digest.update(b"\0");
digest.update(mime_type.to_ascii_lowercase().as_bytes());
digest.update(b"\0");
digest.update(bytes);
let cache_key = hex::encode(digest.finalize());
if let Some(cached) = EXTRACTION_CACHE.lock().await.get(&cache_key).cloned() {
return (*cached).clone();
}
let Ok(_permit) = EXTRACTION_LIMIT.acquire().await else {
return ExtractionResult::failed("文档解析服务暂时不可用");
};
let name = name.to_string();
let mime_type = mime_type.to_string();
let bytes = bytes.to_vec();
let result = tokio::task::spawn_blocking(move || extract_sync(&name, &mime_type, &bytes))
.await
.unwrap_or_else(|error| ExtractionResult::failed(format!("文档解析任务失败:{error}")));
let mut cache = EXTRACTION_CACHE.lock().await;
if cache.len() >= 128 {
cache.clear();
}
cache.insert(cache_key, Arc::new(result.clone()));
result
}
pub async fn load_owned(
db: &sqlx::AnyPool,
user_id: Option<i64>,
asset_id: &str,
) -> Result<StoredAttachment, String> {
let row = crate::db::sql::query(
"SELECT public_id, title, file_name, mime_type, kind, file_path, size_bytes
FROM workspace_assets
WHERE public_id = ?
AND ((? IS NULL AND user_id IS NULL) OR user_id = ?)
AND deleted_at IS NULL AND status = 'completed'
LIMIT 1",
)
.bind(asset_id)
.bind(user_id)
.bind(user_id)
.fetch_optional(db)
.await
.map_err(|error| format!("读取附件失败:{error}"))?
.ok_or_else(|| format!("附件 '{asset_id}' 不存在或不属于当前用户"))?;
let file_path: Option<String> = row.try_get("file_path").ok();
Ok(StoredAttachment {
asset_id: row.try_get("public_id").unwrap_or_else(|_| asset_id.to_string()),
name: row
.try_get::<Option<String>, _>("title")
.ok()
.flatten()
.or_else(|| row.try_get::<Option<String>, _>("file_name").ok().flatten())
.unwrap_or_else(|| "attachment.bin".to_string()),
mime_type: row.try_get::<Option<String>, _>("mime_type").ok().flatten().unwrap_or_else(|| "application/octet-stream".to_string()),
media_kind: row.try_get("kind").unwrap_or_else(|_| "file".to_string()),
file_path: file_path.filter(|path| !path.trim().is_empty()).ok_or_else(|| "附件没有本地文件".to_string())?,
size: row.try_get::<Option<i64>, _>("size_bytes").ok().flatten().unwrap_or_default().max(0) as usize,
})
}
pub async fn extract_stored(attachment: &StoredAttachment) -> ExtractionResult {
let bytes = match tokio::fs::read(&attachment.file_path).await {
Ok(bytes) => bytes,
Err(error) => return ExtractionResult::failed(format!("读取附件内容失败:{error}")),
};
extract_bytes(&attachment.name, &attachment.mime_type, &bytes).await
}
pub async fn image_data_url(attachment: &StoredAttachment) -> Result<String, String> {
if attachment.size > MAX_INLINE_IMAGE_BYTES {
return Err("图片超过 12 MB,无法作为视觉内容发送".to_string());
}
let bytes = tokio::fs::read(&attachment.file_path)
.await
.map_err(|error| format!("读取图片失败:{error}"))?;
Ok(format!(
"data:{};base64,{}",
attachment.mime_type,
general_purpose::STANDARD.encode(bytes)
))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn classifies_supported_attachment_types() {
assert_eq!(media_kind("报告.PDF", "application/octet-stream"), "document");
assert_eq!(media_kind("数据.xls", "application/vnd.ms-excel"), "document");
assert_eq!(media_kind("notes.md", "text/markdown"), "text");
assert_eq!(media_kind("photo.webp", "image/webp"), "image");
}
#[tokio::test]
async fn extracts_utf8_text_and_rejects_empty_files() {
let ready = extract_bytes("中文.txt", "text/plain", "你好\n世界".as_bytes()).await;
assert_eq!(ready.status, "ready");
assert_eq!(ready.content, "你好\n世界");
let empty = extract_bytes("empty.txt", "text/plain", b"").await;
assert_eq!(empty.status, "failed");
}
#[tokio::test]
async fn extracts_generated_modern_office_documents() {
for (format, name, mime) in [
(
DocumentFormat::Docx,
"report.docx",
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
),
(
DocumentFormat::Xlsx,
"report.xlsx",
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
),
(
DocumentFormat::Pptx,
"report.pptx",
"application/vnd.openxmlformats-officedocument.presentationml.presentation",
),
] {
let mut writer = Cursor::new(Vec::new());
office_oxide::create::create_from_markdown_to_writer(
"# 季度报告\n\n| 项目 | 数值 |\n|---|---|\n| 收入 | 42 |",
format,
&mut writer,
)
.expect("generate office fixture");
let result = extract_bytes(name, mime, writer.get_ref()).await;
assert_eq!(result.status, "ready", "failed to extract {name}");
assert!(result.content.contains("季度报告"));
assert!(result.content.contains("42"));
}
}
}
+47 -3
View File
@@ -53,13 +53,22 @@ pub enum SettledState {
pub struct ApprovalEntry {
state: AtomicU8,
notify: Notify,
binding: Option<InteractionBinding>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InteractionBinding {
pub actor_key: String,
pub session_id: String,
pub run_id: String,
}
impl ApprovalEntry {
fn new() -> Self {
fn new(binding: Option<InteractionBinding>) -> Self {
Self {
state: AtomicU8::new(PENDING),
notify: Notify::new(),
binding,
}
}
@@ -102,7 +111,21 @@ impl Drop for PendingGuard {
/// [`PendingGuard`] that the caller must hold for the lifetime of the wait; when
/// the guard drops, the registry entry is removed.
pub fn register(request_id: &str) -> (Arc<ApprovalEntry>, PendingGuard) {
let entry = Arc::new(ApprovalEntry::new());
register_inner(request_id, None)
}
pub fn register_bound(
request_id: &str,
binding: InteractionBinding,
) -> (Arc<ApprovalEntry>, PendingGuard) {
register_inner(request_id, Some(binding))
}
fn register_inner(
request_id: &str,
binding: Option<InteractionBinding>,
) -> (Arc<ApprovalEntry>, PendingGuard) {
let entry = Arc::new(ApprovalEntry::new(binding));
REGISTRY
.lock()
.unwrap()
@@ -149,10 +172,31 @@ pub async fn wait_for_decision(
/// claimed the entry, so the frontend is never told its vote took effect when it
/// did not.
pub fn resolve(request_id: &str, outcome: ApprovalOutcome) -> bool {
resolve_inner(request_id, None, outcome)
}
pub fn resolve_bound(
request_id: &str,
binding: &InteractionBinding,
outcome: ApprovalOutcome,
) -> bool {
resolve_inner(request_id, Some(binding), outcome)
}
fn resolve_inner(
request_id: &str,
binding: Option<&InteractionBinding>,
outcome: ApprovalOutcome,
) -> bool {
let entry = REGISTRY.lock().unwrap().get(request_id).cloned();
let Some(entry) = entry else {
return false;
};
if let Some(binding) = binding {
if entry.binding.as_ref() != Some(binding) {
return false;
}
}
let target = match outcome {
ApprovalOutcome::Approve => APPROVED,
ApprovalOutcome::Deny => DENIED,
@@ -215,7 +259,7 @@ mod tests {
{
let (_entry, _guard) = register("r5");
} // guard dropped here
// Entry gone → resolve can't find it.
// Entry gone → resolve can't find it.
assert!(!resolve("r5", ApprovalOutcome::Approve));
}
+43 -9
View File
@@ -41,14 +41,23 @@ pub struct AskEntry {
state: AtomicU8,
answer: Mutex<Option<String>>,
notify: Notify,
binding: Option<InteractionBinding>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InteractionBinding {
pub actor_key: String,
pub session_id: String,
pub run_id: String,
}
impl AskEntry {
fn new() -> Self {
fn new(binding: Option<InteractionBinding>) -> Self {
Self {
state: AtomicU8::new(PENDING),
answer: Mutex::new(None),
notify: Notify::new(),
binding,
}
}
@@ -84,7 +93,21 @@ impl Drop for PendingGuard {
/// [`wait_for_answer`]) plus a [`PendingGuard`] the caller must hold for the
/// lifetime of the wait.
pub fn register(request_id: &str) -> (Arc<AskEntry>, PendingGuard) {
let entry = Arc::new(AskEntry::new());
register_inner(request_id, None)
}
pub fn register_bound(
request_id: &str,
binding: InteractionBinding,
) -> (Arc<AskEntry>, PendingGuard) {
register_inner(request_id, Some(binding))
}
fn register_inner(
request_id: &str,
binding: Option<InteractionBinding>,
) -> (Arc<AskEntry>, PendingGuard) {
let entry = Arc::new(AskEntry::new(binding));
REGISTRY
.lock()
.unwrap()
@@ -99,10 +122,7 @@ pub fn register(request_id: &str) -> (Arc<AskEntry>, PendingGuard) {
/// Wait until the entry is answered or `timeout` elapses. On timeout this claims
/// `Cancelled` (if still pending) so the result is always terminal.
pub async fn wait_for_answer(
entry: &AskEntry,
timeout: std::time::Duration,
) -> AskOutcome {
pub async fn wait_for_answer(entry: &AskEntry, timeout: std::time::Duration) -> AskOutcome {
// Arm the notified future before checking state to avoid a lost wakeup.
let notified = entry.notify.notified();
@@ -112,9 +132,10 @@ pub async fn wait_for_answer(
if tokio::time::timeout(timeout, notified).await.is_err() {
// Timer fired first — claim Cancelled if nobody answered meanwhile.
let _ = entry
.state
.compare_exchange(PENDING, CANCELLED, Ordering::SeqCst, Ordering::SeqCst);
let _ =
entry
.state
.compare_exchange(PENDING, CANCELLED, Ordering::SeqCst, Ordering::SeqCst);
}
entry.outcome().unwrap_or(AskOutcome::Cancelled)
@@ -126,10 +147,23 @@ pub async fn wait_for_answer(
/// question. Returns `false` for an unknown id (already cleaned up) or when the
/// timeout / another answer already claimed the entry.
pub fn answer(request_id: &str, text: String) -> bool {
answer_inner(request_id, None, text)
}
pub fn answer_bound(request_id: &str, binding: &InteractionBinding, text: String) -> bool {
answer_inner(request_id, Some(binding), text)
}
fn answer_inner(request_id: &str, binding: Option<&InteractionBinding>, text: String) -> bool {
let entry = REGISTRY.lock().unwrap().get(request_id).cloned();
let Some(entry) = entry else {
return false;
};
if let Some(binding) = binding {
if entry.binding.as_ref() != Some(binding) {
return false;
}
}
// Publish the answer *before* flipping state so a waiter that sees ANSWERED
// always reads the value.
*entry.answer.lock().unwrap() = Some(text);
+174 -73
View File
@@ -7,59 +7,61 @@
//! schema, registered via `custom_tools` so it overrides the built-in) and
//! confines the command with [bubblewrap](https://github.com/containers/bubblewrap):
//!
//! - **Default / AcceptEdits** (review modes): the command runs inside a
//! - **Every Web permission mode**: the command runs inside a
//! bwrap namespace with the system mounted read-only, the session sandbox
//! directory bind-mounted read-write at its real path, `/etc`, `/root`,
//! `/home` and every other host path *absent*, and the network unshared.
//! Filesystem reads/writes outside the session dir and all network egress are
//! therefore impossible at the kernel level, not merely discouraged.
//! - **FullAccess** (explicitly dangerous): the command runs unconfined, the
//! same behaviour the bare SDK had, but every command is audit-logged with a
//! loud warning.
//!
//! **Fail-closed:** in a review mode, if `bwrap` is unavailable the command is
//! refused rather than silently run unconfined — we never fabricate a boundary.
//! `full_access` only skips approval prompts; it never weakens this kernel
//! boundary. If the required runtime is unavailable, Bash is omitted from the
//! Web tool profile and direct calls fail closed.
use async_trait::async_trait;
use once_cell::sync::Lazy;
use serde_json::{json, Value};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::io::{AsyncRead, AsyncReadExt};
use tokio::process::Command;
use crate::services::agent_workbench::permission::PermissionMode;
use open_agent_sdk::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
const DEFAULT_TIMEOUT_MS: u64 = 120_000;
const MAX_TIMEOUT_MS: u64 = 600_000;
const MAX_TIMEOUT_MS: u64 = 120_000;
const MAX_OUTPUT_SIZE: usize = 100_000;
const CPU_LIMIT_SECS: u64 = 60;
const ADDRESS_SPACE_LIMIT_BYTES: u64 = 512 * 1024 * 1024;
const PROCESS_LIMIT: u64 = 64;
/// Read-only host paths exposed inside the sandbox so common tooling works.
/// Deliberately excludes `/etc`, `/root`, `/home`, `/var`, `/proc`-host, etc.
const RO_SYSTEM_PATHS: &[&str] = &["/usr", "/bin", "/sbin", "/lib", "/lib64", "/opt"];
const RO_SYSTEM_PATHS: &[&str] = &["/usr", "/bin", "/sbin", "/lib", "/lib64"];
/// A Bash tool that confines commands to the session sandbox via bubblewrap.
pub struct SandboxedBashTool {
/// Canonical session sandbox directory (bind-mounted rw inside bwrap).
session_dir: PathBuf,
/// Session permission mode — decides confined vs unconfined.
mode: PermissionMode,
/// Resolved `bwrap` binary path, if present on this host.
bwrap: Option<PathBuf>,
}
impl SandboxedBashTool {
pub fn new(session_dir: PathBuf, mode: PermissionMode) -> Self {
pub fn new(session_dir: PathBuf, _mode: PermissionMode) -> Self {
Self {
session_dir,
mode,
bwrap: find_bwrap(),
}
}
/// Whether this mode confines bash (everything except FullAccess).
/// Web Bash is always confined, including `full_access` sessions.
#[cfg(test)]
fn is_confined(&self) -> bool {
!matches!(self.mode, PermissionMode::FullAccess)
true
}
}
@@ -78,6 +80,47 @@ fn find_bwrap() -> Option<PathBuf> {
})
}
fn prlimit_available() -> bool {
Path::new("/usr/bin/prlimit").is_file()
}
static SANDBOX_AVAILABLE: Lazy<bool> = Lazy::new(probe_sandbox);
/// Run a real namespace probe. Merely finding the binary is insufficient on
/// hosts where unprivileged user namespaces are disabled.
pub fn sandbox_available() -> bool {
*SANDBOX_AVAILABLE
}
fn probe_sandbox() -> bool {
let Some(bwrap) = find_bwrap() else {
return false;
};
if !prlimit_available() {
return false;
}
let probe_dir = std::env::temp_dir().join(format!(
"boosapi-agent-bwrap-probe-{}",
uuid::Uuid::new_v4().simple()
));
if std::fs::create_dir_all(&probe_dir).is_err() {
return false;
}
let canonical = probe_dir.canonicalize().unwrap_or(probe_dir.clone());
let result = std::process::Command::new(bwrap)
.args(build_bwrap_args(&canonical, "true"))
.env_clear()
.env("PATH", "/usr/bin:/bin")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.map(|status| status.success())
.unwrap_or(false);
let _ = std::fs::remove_dir_all(probe_dir);
result
}
/// Build the bwrap argument vector that confines `command` to `session_dir`.
///
/// Pure and unit-testable. The session directory is bound read-write at its own
@@ -112,7 +155,15 @@ pub fn build_bwrap_args(session_dir: &Path, command: &str) -> Vec<String> {
args.push("--unshare-all".into());
args.push("--die-with-parent".into());
args.push("--new-session".into());
// Run the command under a shell inside the sandbox.
// Apply hard resource limits before running the shell. `prlimit` and its
// dynamic libraries are part of the read-only system mounts above.
args.push("/usr/bin/prlimit".into());
args.push(format!("--cpu={CPU_LIMIT_SECS}:{CPU_LIMIT_SECS}"));
args.push(format!(
"--as={ADDRESS_SPACE_LIMIT_BYTES}:{ADDRESS_SPACE_LIMIT_BYTES}"
));
args.push(format!("--nproc={PROCESS_LIMIT}:{PROCESS_LIMIT}"));
args.push("--".into());
args.push("/bin/bash".into());
args.push("-c".into());
args.push(command.into());
@@ -128,8 +179,8 @@ impl Tool for SandboxedBashTool {
fn description(&self) -> &str {
"Executes a bash command inside an OS-level sandbox confined to the \
session working directory (read-only system, no network) unless the \
session is in full-access mode. Use for shell operations."
session working directory (read-only system, no network). Use for \
shell operations."
}
fn input_schema(&self) -> ToolInputSchema {
@@ -142,7 +193,7 @@ impl Tool for SandboxedBashTool {
),
(
"timeout".to_string(),
json!({ "type": "number", "description": "Optional timeout in milliseconds (max 600000)" }),
json!({ "type": "number", "description": "Optional timeout in milliseconds (max 120000)" }),
),
(
"description".to_string(),
@@ -166,50 +217,35 @@ impl Tool for SandboxedBashTool {
.unwrap_or(DEFAULT_TIMEOUT_MS)
.min(MAX_TIMEOUT_MS);
// Build the runner per mode.
let (program, args): (String, Vec<String>) = if self.is_confined() {
let Some(bwrap) = &self.bwrap else {
// Fail closed: never downgrade a review-mode session to an
// unconfined shell just because the sandbox tool is missing.
return Ok(ToolResult::error(
"Bash is disabled: the OS sandbox (bwrap) is not available on \
this host, and this session runs in a review mode that requires \
confinement. Re-run in full-access mode (explicitly unsafe) or \
install bubblewrap."
.to_string(),
));
};
(
bwrap.to_string_lossy().to_string(),
build_bwrap_args(&self.session_dir, command),
)
} else {
// FullAccess: unconfined, but audit every command loudly.
tracing::warn!(
target: "boosapi::agent::bash",
mode = "full_access",
command = command,
"UNCONFINED agent bash command (full-access mode)"
);
(
"/bin/bash".to_string(),
vec!["-c".to_string(), command.to_string()],
)
let Some(bwrap) = &self.bwrap else {
return Ok(ToolResult::error(
"Bash is disabled: the required OS sandbox is not available".to_string(),
));
};
if !prlimit_available() {
return Ok(ToolResult::error(
"Bash is disabled: the required resource limiter is not available".to_string(),
));
}
let program = bwrap.to_string_lossy().to_string();
let args = build_bwrap_args(&self.session_dir, command);
// Unconfined mode honours the SDK's working dir; confined mode's cwd is
// set inside bwrap via --chdir, so the outer cwd is irrelevant.
let outer_cwd = if self.is_confined() {
self.session_dir.clone()
} else {
PathBuf::from(&context.working_dir)
};
let outer_cwd = self.session_dir.clone();
let fut = run_capture(&program, &args, &outer_cwd);
match tokio::time::timeout(std::time::Duration::from_millis(timeout_ms), fut).await {
Ok(Ok((stdout, stderr, code))) => Ok(format_result(stdout, stderr, code)),
Ok(Err(e)) => Ok(ToolResult::error(format!("Command failed: {e}"))),
Err(_) => Ok(ToolResult::error(format!(
tokio::pin!(fut);
match tokio::select! {
result = &mut fut => Some(result),
_ = tokio::time::sleep(std::time::Duration::from_millis(timeout_ms)) => None,
_ = context.abort_signal.cancelled() => {
return Ok(ToolResult::error("Command cancelled".to_string()));
}
} {
Some(Ok((stdout, stderr, code, truncated))) => {
Ok(format_result(stdout, stderr, code, truncated))
}
Some(Err(e)) => Ok(ToolResult::error(format!("Command failed: {e}"))),
None => Ok(ToolResult::error(format!(
"Command timed out after {timeout_ms}ms"
))),
}
@@ -220,23 +256,67 @@ async fn run_capture(
program: &str,
args: &[String],
cwd: &Path,
) -> Result<(String, String, i32), std::io::Error> {
let output = Command::new(program)
) -> Result<(String, String, i32, bool), std::io::Error> {
let mut command = Command::new(program);
command
.args(args)
.current_dir(cwd)
.env_clear()
.env("PATH", "/usr/bin:/bin")
.env("LANG", "C.UTF-8")
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.output()
.await?;
.kill_on_drop(true);
let mut child = command.spawn()?;
let stdout = child.stdout.take().ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::BrokenPipe, "missing command stdout")
})?;
let stderr = child.stderr.take().ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::BrokenPipe, "missing command stderr")
})?;
let budget = Arc::new(AtomicUsize::new(0));
let (status, stdout, stderr) = tokio::try_join!(
child.wait(),
read_bounded(stdout, budget.clone()),
read_bounded(stderr, budget.clone()),
)?;
let truncated = budget.load(Ordering::Relaxed) > MAX_OUTPUT_SIZE;
Ok((
String::from_utf8_lossy(&output.stdout).to_string(),
String::from_utf8_lossy(&output.stderr).to_string(),
output.status.code().unwrap_or(-1),
String::from_utf8_lossy(&stdout).to_string(),
String::from_utf8_lossy(&stderr).to_string(),
status.code().unwrap_or(-1),
truncated,
))
}
fn format_result(stdout: String, stderr: String, exit_code: i32) -> ToolResult {
async fn read_bounded(
mut reader: impl AsyncRead + Unpin,
budget: Arc<AtomicUsize>,
) -> Result<Vec<u8>, std::io::Error> {
let mut stored = Vec::new();
let mut chunk = [0_u8; 8192];
loop {
let read = reader.read(&mut chunk).await?;
if read == 0 {
break;
}
let prior = budget
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
Some(used.saturating_add(read).min(MAX_OUTPUT_SIZE + 1))
})
.unwrap_or(MAX_OUTPUT_SIZE + 1);
if prior < MAX_OUTPUT_SIZE {
let keep = read.min(MAX_OUTPUT_SIZE - prior);
stored.extend_from_slice(&chunk[..keep]);
}
// Continue draining after the cap so the child cannot block on a full
// pipe; bytes beyond the shared stdout+stderr budget are discarded.
}
Ok(stored)
}
fn format_result(stdout: String, stderr: String, exit_code: i32, truncated: bool) -> ToolResult {
let mut result = String::new();
if !stdout.is_empty() {
result.push_str(&stdout);
@@ -248,13 +328,19 @@ fn format_result(stdout: String, stderr: String, exit_code: i32) -> ToolResult {
result.push_str("STDERR:\n");
result.push_str(&stderr);
}
if result.len() > MAX_OUTPUT_SIZE {
result.truncate(MAX_OUTPUT_SIZE);
result.push_str("\n... (output truncated)");
let mut suffix = String::new();
if truncated {
suffix.push_str("\n... (output truncated)");
}
if exit_code != 0 {
result.push_str(&format!("\n\nExit code: {exit_code}"));
suffix.push_str(&format!("\n\nExit code: {exit_code}"));
}
if result.len().saturating_add(suffix.len()) > MAX_OUTPUT_SIZE {
let available = MAX_OUTPUT_SIZE.saturating_sub(suffix.len());
let boundary = open_agent_sdk::utils::text::floor_char_boundary(&result, available);
result.truncate(boundary);
}
result.push_str(&suffix);
if result.is_empty() {
result = "(no output)".to_string();
}
@@ -282,9 +368,10 @@ mod tests {
assert!(args.iter().any(|a| a == "--die-with-parent"));
// /etc is never bound.
assert!(!joined.contains("/etc"));
assert!(!joined.contains("/opt"));
// System paths are read-only.
assert!(joined.contains("--ro-bind-try /usr /usr"));
// The command is the last argument, run under bash -c.
// The command is the last argument, run under prlimit + bash -c.
assert_eq!(args.last().unwrap(), "echo hi");
let n = args.len();
assert_eq!(args[n - 3], "/bin/bash");
@@ -296,7 +383,14 @@ mod tests {
let sd = PathBuf::from("/tmp/x");
assert!(SandboxedBashTool::new(sd.clone(), PermissionMode::Default).is_confined());
assert!(SandboxedBashTool::new(sd.clone(), PermissionMode::AcceptEdits).is_confined());
assert!(!SandboxedBashTool::new(sd, PermissionMode::FullAccess).is_confined());
assert!(SandboxedBashTool::new(sd, PermissionMode::FullAccess).is_confined());
}
#[test]
fn utf8_truncation_never_splits_a_character() {
let result = format_result("界".repeat(40_000), String::new(), 0, true);
assert!(!result.is_error);
assert!(result.get_text().contains("output truncated"));
}
#[test]
@@ -304,4 +398,11 @@ mod tests {
let t = SandboxedBashTool::new(PathBuf::from("/tmp/x"), PermissionMode::Default);
assert_eq!(t.name(), "Bash");
}
#[test]
fn formatted_output_never_exceeds_wire_limit() {
let result = format_result("x".repeat(MAX_OUTPUT_SIZE), String::new(), 17, true);
assert!(result.get_text().len() <= MAX_OUTPUT_SIZE);
assert!(result.get_text().ends_with("Exit code: 17"));
}
}
+1
View File
@@ -17,5 +17,6 @@ pub mod approval;
pub mod ask;
pub mod bash_sandbox;
pub mod permission;
pub mod run_registry;
pub mod sandbox;
pub mod session_store;
+28 -5
View File
@@ -109,14 +109,28 @@ mod tests {
#[test]
fn classifies_execute_tools() {
for t in ["Bash", "bash", "shell", "run_command", "execute", "RunInTerminal"] {
for t in [
"Bash",
"bash",
"shell",
"run_command",
"execute",
"RunInTerminal",
] {
assert_eq!(classify_tool_risk(t), RiskLevel::Execute, "tool={t}");
}
}
#[test]
fn classifies_write_tools() {
for t in ["Write", "Edit", "create_file", "delete_file", "patch", "ApplyPatch"] {
for t in [
"Write",
"Edit",
"create_file",
"delete_file",
"patch",
"ApplyPatch",
] {
assert_eq!(classify_tool_risk(t), RiskLevel::Write, "tool={t}");
}
}
@@ -168,9 +182,18 @@ mod tests {
#[test]
fn mode_parsing_falls_back_to_default() {
assert_eq!(PermissionMode::from_wire("accept_edits"), PermissionMode::AcceptEdits);
assert_eq!(PermissionMode::from_wire("full_access"), PermissionMode::FullAccess);
assert_eq!(PermissionMode::from_wire("garbage"), PermissionMode::Default);
assert_eq!(
PermissionMode::from_wire("accept_edits"),
PermissionMode::AcceptEdits
);
assert_eq!(
PermissionMode::from_wire("full_access"),
PermissionMode::FullAccess
);
assert_eq!(
PermissionMode::from_wire("garbage"),
PermissionMode::Default
);
assert_eq!(PermissionMode::from_wire(""), PermissionMode::Default);
}
}
@@ -0,0 +1,78 @@
//! Process-local cancellation handles for durable Agent runs.
//!
//! The database owns run leases and authorization. This registry only bridges
//! an accepted cancel request to the SDK cancellation token in this process.
use once_cell::sync::Lazy;
use open_agent_sdk::CancellationToken;
use std::collections::HashMap;
use std::sync::Mutex;
#[derive(Clone)]
struct ActiveRun {
actor_key: String,
session_id: String,
cancel: CancellationToken,
}
static RUNS: Lazy<Mutex<HashMap<String, ActiveRun>>> = Lazy::new(|| Mutex::new(HashMap::new()));
pub struct RunGuard {
run_id: String,
}
impl Drop for RunGuard {
fn drop(&mut self) {
RUNS.lock().unwrap().remove(&self.run_id);
}
}
pub fn register(
run_id: &str,
session_id: &str,
actor_key: &str,
cancel: CancellationToken,
) -> Option<RunGuard> {
let mut runs = RUNS.lock().unwrap();
if runs.contains_key(run_id) {
return None;
}
runs.insert(
run_id.to_string(),
ActiveRun {
actor_key: actor_key.to_string(),
session_id: session_id.to_string(),
cancel,
},
);
Some(RunGuard {
run_id: run_id.to_string(),
})
}
pub fn cancel(run_id: &str, actor_key: &str, session_id: &str) -> bool {
let run = RUNS.lock().unwrap().get(run_id).cloned();
let Some(run) = run else {
return false;
};
if run.actor_key != actor_key || run.session_id != session_id {
return false;
}
run.cancel.cancel();
true
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cancellation_is_bound_to_actor_and_session() {
let token = CancellationToken::new();
let _guard = register("run", "session", "user:1", token.clone()).unwrap();
assert!(!cancel("run", "user:2", "session"));
assert!(!token.is_cancelled());
assert!(cancel("run", "user:1", "session"));
assert!(token.is_cancelled());
}
}
+101 -8
View File
@@ -29,12 +29,67 @@ pub fn ensure_session_dir(session_id: &str) -> std::io::Result<PathBuf> {
dir.canonicalize()
}
/// Create/resolve the directory identified by the database-owned workspace id.
/// New Web sessions never derive a filesystem path from caller-controlled text.
pub fn ensure_workspace_dir(workspace_id: &str) -> std::io::Result<PathBuf> {
let parsed = uuid::Uuid::parse_str(workspace_id).map_err(|_| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"invalid agent workspace id",
)
})?;
let dir = Path::new(SESSIONS_ROOT).join(parsed.hyphenated().to_string());
std::fs::create_dir_all(&dir)?;
dir.canonicalize()
}
/// Lazily move one unambiguous legacy session directory to its independent
/// workspace id. The caller must prove `legacy_name_is_unique`; sanitized-name
/// collisions are intentionally left untouched to prevent cross-session data
/// copying.
pub fn prepare_workspace_dir(
session_id: &str,
workspace_id: &str,
legacy_name_is_unique: bool,
) -> std::io::Result<PathBuf> {
let parsed = uuid::Uuid::parse_str(workspace_id).map_err(|_| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"invalid agent workspace id",
)
})?;
std::fs::create_dir_all(SESSIONS_ROOT)?;
let destination = Path::new(SESSIONS_ROOT).join(parsed.hyphenated().to_string());
let legacy = session_dir(session_id);
if !destination.exists() && legacy_name_is_unique && legacy != destination {
if let Ok(metadata) = std::fs::symlink_metadata(&legacy) {
if metadata.file_type().is_symlink() {
return Err(std::io::Error::new(
std::io::ErrorKind::PermissionDenied,
"legacy agent workspace is a symlink",
));
}
std::fs::rename(&legacy, &destination)?;
}
}
std::fs::create_dir_all(&destination)?;
destination.canonicalize()
}
/// Strip anything that could let a session id traverse outside SESSIONS_ROOT.
/// Keeps only `[A-Za-z0-9_-]`; everything else becomes `_`.
fn sanitize_session_id(session_id: &str) -> String {
pub(crate) fn sanitize_session_id(session_id: &str) -> String {
let cleaned: String = session_id
.chars()
.map(|c| if c.is_ascii_alphanumeric() || c == '-' || c == '_' { c } else { '_' })
.map(|c| {
if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
c
} else {
'_'
}
})
.collect();
if cleaned.is_empty() {
"default".to_string()
@@ -134,7 +189,9 @@ pub fn validate_path_within_cwd(path: &str, cwd: &Path) -> Result<PathBuf, Strin
}
ancestor = dir.parent();
}
Err(format!("cannot verify path '{path}' against the session sandbox"))
Err(format!(
"cannot verify path '{path}' against the session sandbox"
))
}
/// Resolve where a (possibly dangling) symlink points, as a real absolute path,
@@ -149,8 +206,7 @@ fn symlink_real_target(link: &Path) -> Result<PathBuf, String> {
} else {
link.parent().unwrap_or_else(|| Path::new("/")).join(raw)
};
real_location(&abs)
.ok_or_else(|| format!("cannot resolve symlink target '{}'", abs.display()))
real_location(&abs).ok_or_else(|| format!("cannot resolve symlink target '{}'", abs.display()))
}
/// Real absolute location of a path that may not fully exist: canonicalize the
@@ -193,9 +249,15 @@ pub fn check_tool_paths(tool_name: &str, input: &Value, cwd: &Path) -> Result<()
check_field(input, &["new_path", "destination", "dest"], cwd)?;
Ok(())
}
"glob" | "glob_search" | "grep" | "search" | "ripgrep" => {
// Optional explicit base path; the pattern itself is relative to cwd.
"glob" | "glob_search" => {
check_field(input, &["path", "dir", "directory"], cwd)?;
check_relative_pattern(input, "pattern")?;
Ok(())
}
"grep" | "search" | "ripgrep" => {
check_field(input, &["path", "dir", "directory"], cwd)?;
// `pattern` is a regex for Grep, while `glob` is path-like.
check_relative_pattern(input, "glob")?;
Ok(())
}
// Command execution and unknown tools: no path arg to validate here.
@@ -203,6 +265,26 @@ pub fn check_tool_paths(tool_name: &str, input: &Value, cwd: &Path) -> Result<()
}
}
fn check_relative_pattern(input: &Value, key: &str) -> Result<(), String> {
let Some(pattern) = input.get(key).and_then(Value::as_str) else {
return Ok(());
};
let trimmed = pattern.trim();
let path = Path::new(trimmed);
let escapes = path.is_absolute()
|| trimmed.starts_with('~')
|| trimmed
.split(['/', '\\'])
.any(|component| component == "..");
if escapes {
Err(format!(
"{key} must be a relative pattern within the session sandbox"
))
} else {
Ok(())
}
}
fn check_field(input: &Value, keys: &[&str], cwd: &Path) -> Result<(), String> {
for key in keys {
if let Some(p) = input.get(*key).and_then(|v| v.as_str()) {
@@ -268,6 +350,14 @@ mod tests {
assert!(check_tool_paths("bash", &input, &sb).is_ok());
}
#[test]
fn rejects_absolute_and_parent_glob_patterns() {
let sb = tmp_sandbox();
assert!(check_tool_paths("Glob", &json!({ "pattern": "/etc/**" }), &sb).is_err());
assert!(check_tool_paths("Glob", &json!({ "pattern": "../../**" }), &sb).is_err());
assert!(check_tool_paths("Glob", &json!({ "pattern": "src/**/*.rs" }), &sb).is_ok());
}
#[cfg(unix)]
#[test]
fn rejects_dangling_symlink_escape() {
@@ -277,7 +367,10 @@ mod tests {
let sb = tmp_sandbox();
std::os::unix::fs::symlink("/etc/cron.d/pwn_dangling", sb.join("evil")).unwrap();
let res = validate_path_within_cwd("evil", &sb);
assert!(res.is_err(), "dangling symlink escape must be rejected: {res:?}");
assert!(
res.is_err(),
"dangling symlink escape must be rejected: {res:?}"
);
}
#[cfg(unix)]
File diff suppressed because it is too large Load Diff
+9 -9
View File
@@ -1,25 +1,25 @@
use crate::{db::models::site_announcement::SiteAnnouncement, error::AppResult};
use sqlx::SqlitePool;
use sqlx::AnyPool;
use std::sync::Arc;
#[derive(Clone)]
pub struct AnnouncementService {
db: Arc<SqlitePool>,
db: Arc<AnyPool>,
}
impl AnnouncementService {
pub fn new(db: Arc<SqlitePool>) -> Self {
pub fn new(db: Arc<AnyPool>) -> Self {
Self { db }
}
pub async fn get_active_announcements(&self) -> AppResult<Vec<SiteAnnouncement>> {
self.expire_old_announcements().await?;
let rows = sqlx::query_as::<_, SiteAnnouncement>(
let rows = crate::db::sql::query_as::<_, SiteAnnouncement>(
"SELECT *
FROM site_announcements
WHERE dismissed_at IS NULL
AND (starts_at IS NULL OR starts_at <= datetime('now'))
AND (ends_at IS NULL OR ends_at > datetime('now'))
AND (starts_at IS NULL OR starts_at <= CURRENT_TIMESTAMP)
AND (ends_at IS NULL OR ends_at > CURRENT_TIMESTAMP)
ORDER BY first_seen_at DESC, id DESC",
)
.fetch_all(&*self.db)
@@ -28,12 +28,12 @@ impl AnnouncementService {
}
pub async fn expire_old_announcements(&self) -> AppResult<u64> {
let result = sqlx::query(
let result = crate::db::sql::query(
"UPDATE site_announcements
SET dismissed_at = datetime('now')
SET dismissed_at = CURRENT_TIMESTAMP
WHERE dismissed_at IS NULL
AND ends_at IS NOT NULL
AND ends_at <= datetime('now')",
AND ends_at <= CURRENT_TIMESTAMP",
)
.execute(&*self.db)
.await?;
File diff suppressed because it is too large Load Diff
+277 -59
View File
@@ -1,14 +1,14 @@
use crate::error::AppResult;
use serde_json::{json, Value};
use sqlx::{Row, SqlitePool, Transaction};
use sqlx::{AnyPool, Row, Transaction};
/// One balance-changing event. Written to `credit_ledger` in the same
/// transaction that updates the cached balance, so ledger and cache never diverge.
pub struct LedgerEntry<'a> {
pub user_id: i64,
pub track: &'a str, // "usd" | "points"
pub entry_type: &'a str, // "recharge" | "admin_adjust" | "consume" | "refund" | "grant"
pub amount: f64, // +credit / -debit
pub track: &'a str, // "usd" | "points"
pub entry_type: &'a str, // "recharge" | "admin_adjust" | "consume" | "refund" | "grant"
pub amount: f64, // +credit / -debit
pub balance_after: f64,
pub ref_type: Option<&'a str>,
pub ref_id: Option<&'a str>,
@@ -17,8 +17,11 @@ pub struct LedgerEntry<'a> {
pub note: Option<&'a str>,
}
pub async fn ensure_ledger_schema(db: &SqlitePool) -> AppResult<()> {
sqlx::query(
pub async fn ensure_ledger_schema(db: &AnyPool) -> AppResult<()> {
if crate::db::sql::runtime_kind() != crate::db::DbKind::Sqlite {
return Ok(());
}
crate::db::sql::query(
"CREATE TABLE IF NOT EXISTS credit_ledger (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
user_id INTEGER NOT NULL,
@@ -31,29 +34,29 @@ pub async fn ensure_ledger_schema(db: &SqlitePool) -> AppResult<()> {
model TEXT,
operator TEXT,
note TEXT,
created_at TEXT DEFAULT (datetime('now'))
created_at TEXT DEFAULT (CURRENT_TIMESTAMP)
)",
)
.execute(db)
.await?;
sqlx::query(
crate::db::sql::query(
"CREATE INDEX IF NOT EXISTS credit_ledger_user_created_idx
ON credit_ledger (user_id, created_at DESC)",
)
.execute(db)
.await?;
sqlx::query(
crate::db::sql::query(
"CREATE INDEX IF NOT EXISTS credit_ledger_track_type_idx
ON credit_ledger (track, entry_type)",
)
.execute(db)
.await?;
sqlx::query(
crate::db::sql::query(
"CREATE TABLE IF NOT EXISTS user_credits (
user_id INTEGER PRIMARY KEY NOT NULL,
points_balance REAL NOT NULL DEFAULT 0,
points_total_recharged REAL NOT NULL DEFAULT 0,
updated_at TEXT DEFAULT (datetime('now'))
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP)
)",
)
.execute(db)
@@ -64,10 +67,10 @@ pub async fn ensure_ledger_schema(db: &SqlitePool) -> AppResult<()> {
/// Append one ledger row inside an existing transaction. Caller is responsible
/// for also updating the cached balance in the same `tx`.
pub async fn record_ledger_entry(
tx: &mut Transaction<'_, sqlx::Sqlite>,
tx: &mut Transaction<'_, sqlx::Any>,
entry: LedgerEntry<'_>,
) -> AppResult<()> {
sqlx::query(
crate::db::sql::query(
"INSERT INTO credit_ledger
(user_id, track, entry_type, amount, balance_after,
ref_type, ref_id, model, operator, note)
@@ -89,13 +92,12 @@ pub async fn record_ledger_entry(
}
/// Current points balance for a user (0 if no row yet).
pub async fn points_balance(db: &SqlitePool, user_id: i64) -> AppResult<f64> {
let bal: Option<f64> = sqlx::query_scalar(
"SELECT points_balance FROM user_credits WHERE user_id = ?",
)
.bind(user_id)
.fetch_optional(db)
.await?;
pub async fn points_balance(db: &AnyPool, user_id: i64) -> AppResult<f64> {
let bal: Option<f64> =
crate::db::sql::query_scalar("SELECT points_balance FROM user_credits WHERE user_id = ?")
.bind(user_id)
.fetch_optional(db)
.await?;
Ok(bal.unwrap_or(0.0))
}
@@ -104,7 +106,7 @@ pub async fn points_balance(db: &SqlitePool, user_id: i64) -> AppResult<f64> {
/// in a single transaction. Returns the new balance.
#[allow(clippy::too_many_arguments)]
pub async fn apply_points_delta(
db: &SqlitePool,
db: &AnyPool,
user_id: i64,
delta: f64,
entry_type: &str,
@@ -115,38 +117,77 @@ pub async fn apply_points_delta(
) -> AppResult<f64> {
ensure_ledger_schema(db).await?;
let mut tx = db.begin().await?;
let current: f64 = sqlx::query_scalar(
"SELECT points_balance FROM user_credits WHERE user_id = ?",
let next = apply_points_delta_in_tx(
&mut tx, user_id, delta, entry_type, model, operator, ref_type, ref_id,
)
.bind(user_id)
.fetch_optional(&mut *tx)
.await?
.unwrap_or(0.0);
.await?;
tx.commit().await?;
Ok(next)
}
/// Apply a points delta inside a caller-owned transaction. Payment fulfillment
/// uses this form so the balance write and order state transition commit or
/// roll back together.
#[allow(clippy::too_many_arguments)]
pub async fn apply_points_delta_in_tx(
tx: &mut Transaction<'_, sqlx::Any>,
user_id: i64,
delta: f64,
entry_type: &str,
model: Option<&str>,
operator: Option<&str>,
ref_type: Option<&str>,
ref_id: Option<&str>,
) -> AppResult<f64> {
let current: f64 =
crate::db::sql::query_scalar("SELECT points_balance FROM user_credits WHERE user_id = ?")
.bind(user_id)
.fetch_optional(&mut **tx)
.await?
.unwrap_or(0.0);
let next = current + delta;
let total_recharged_inc = if delta > 0.0 && entry_type == "recharge" { delta } else { 0.0 };
sqlx::query(
let total_recharged_inc = if delta > 0.0 && entry_type == "recharge" {
delta
} else {
0.0
};
crate::db::sql::query(
"INSERT INTO user_credits (user_id, points_balance, points_total_recharged, updated_at)
VALUES (?, ?, ?, datetime('now'))
VALUES (?, ?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(user_id) DO UPDATE SET
points_balance = ?,
points_total_recharged = points_total_recharged + ?,
updated_at = datetime('now')",
updated_at = CURRENT_TIMESTAMP",
)
.bind(user_id).bind(next).bind(total_recharged_inc)
.bind(next).bind(total_recharged_inc)
.execute(&mut *tx).await?;
record_ledger_entry(&mut tx, LedgerEntry {
user_id, track: "points", entry_type,
amount: delta, balance_after: next,
ref_type, ref_id, model, operator, note: None,
}).await?;
tx.commit().await?;
.bind(user_id)
.bind(next)
.bind(total_recharged_inc)
.bind(next)
.bind(total_recharged_inc)
.execute(&mut **tx)
.await?;
record_ledger_entry(
tx,
LedgerEntry {
user_id,
track: "points",
entry_type,
amount: delta,
balance_after: next,
ref_type,
ref_id,
model,
operator,
note: None,
},
)
.await?;
Ok(next)
}
/// List a user's ledger rows, newest first, optional track/entry_type filter + paging.
pub async fn list_user_ledger(
db: &SqlitePool,
db: &AnyPool,
user_id: i64,
track: Option<&str>,
entry_type: Option<&str>,
@@ -156,7 +197,119 @@ pub async fn list_user_ledger(
ensure_ledger_schema(db).await?;
let limit = limit.clamp(1, 500);
let offset = offset.max(0);
let rows = sqlx::query(
let normalized_track = track.map(str::trim).unwrap_or_default();
if matches!(normalized_track, "points" | "comfy") {
crate::services::comfy_billing::ensure_schema(db).await?;
let rows = crate::db::sql::query(
"SELECT l.entry_type, l.amount_micros, l.balance_after_micros,
l.package_amount_micros, l.balance_amount_micros,
l.task_type, l.operator, l.ref_type, l.ref_id, l.note,
l.created_at, g.grant_kind
FROM comfy_entitlement_ledger l
LEFT JOIN comfy_credit_grants g ON g.id = l.grant_id
WHERE l.user_id = ? AND (? IS NULL OR l.entry_type = ?)
ORDER BY l.id DESC LIMIT ? OFFSET ?",
)
.bind(user_id)
.bind(entry_type)
.bind(entry_type)
.bind(limit)
.bind(offset)
.fetch_all(db)
.await?;
if !rows.is_empty() {
return Ok(rows
.iter()
.map(|row| {
let amount = row.try_get::<i64, _>("amount_micros").unwrap_or(0);
let balance_after = row.try_get::<i64, _>("balance_after_micros").unwrap_or(0);
json!({
"track": "points",
"billingTrack": "comfy",
"entryType": row.try_get::<String, _>("entry_type").unwrap_or_default(),
"amount": crate::services::comfy_billing::micros_to_points(amount),
"balanceAfter": crate::services::comfy_billing::micros_to_points(balance_after),
"grantKind": row.try_get::<Option<String>, _>("grant_kind").ok().flatten(),
"taskType": row.try_get::<Option<String>, _>("task_type").ok().flatten(),
"packageAmount": crate::services::comfy_billing::micros_to_points(
row.try_get::<i64, _>("package_amount_micros").unwrap_or(0)
),
"balanceAmount": crate::services::comfy_billing::micros_to_points(
row.try_get::<i64, _>("balance_amount_micros").unwrap_or(0)
),
"operator": row.try_get::<Option<String>, _>("operator").ok().flatten(),
"refType": row.try_get::<Option<String>, _>("ref_type").ok().flatten(),
"refId": row.try_get::<Option<String>, _>("ref_id").ok().flatten(),
"note": row.try_get::<Option<String>, _>("note").ok().flatten(),
"createdAt": row.try_get::<Option<String>, _>("created_at").ok().flatten(),
})
})
.collect());
}
} else if matches!(normalized_track, "usd" | "api") {
crate::services::billing::account::ensure_schema(db).await?;
let rows = crate::db::sql::query(
"SELECT l.entry_type, l.amount_micros, l.balance_after_micros,
l.package_amount_micros, l.balance_amount_micros,
l.model, l.operator, l.ref_type, l.ref_id, l.note,
l.created_at, g.grant_kind
FROM api_entitlement_ledger l
LEFT JOIN api_credit_grants g ON g.id = l.grant_id
WHERE l.user_id = ? AND (? IS NULL OR l.entry_type = ?)
ORDER BY l.id DESC LIMIT ? OFFSET ?",
)
.bind(user_id)
.bind(entry_type)
.bind(entry_type)
.bind(limit)
.bind(offset)
.fetch_all(db)
.await?;
if !rows.is_empty() {
return Ok(rows
.iter()
.map(|row| {
let package_amount =
row.try_get::<i64, _>("package_amount_micros").unwrap_or(0);
let balance_amount =
row.try_get::<i64, _>("balance_amount_micros").unwrap_or(0);
let grant_kind = if package_amount > 0 && balance_amount > 0 {
Some("mixed".to_string())
} else {
row.try_get::<Option<String>, _>("grant_kind")
.ok()
.flatten()
};
json!({
"track": "usd",
"billingTrack": "api",
"entryType": row.try_get::<String, _>("entry_type").unwrap_or_default(),
"amount": crate::services::billing::account::micros_to_units(
row.try_get::<i64, _>("amount_micros").unwrap_or(0)
),
"balanceAfter": crate::services::billing::account::micros_to_units(
row.try_get::<i64, _>("balance_after_micros").unwrap_or(0)
),
"grantKind": grant_kind,
"taskType": Value::Null,
"packageAmount": crate::services::billing::account::micros_to_units(
package_amount
),
"balanceAmount": crate::services::billing::account::micros_to_units(
balance_amount
),
"model": row.try_get::<Option<String>, _>("model").ok().flatten(),
"operator": row.try_get::<Option<String>, _>("operator").ok().flatten(),
"refType": row.try_get::<Option<String>, _>("ref_type").ok().flatten(),
"refId": row.try_get::<Option<String>, _>("ref_id").ok().flatten(),
"note": row.try_get::<Option<String>, _>("note").ok().flatten(),
"createdAt": row.try_get::<Option<String>, _>("created_at").ok().flatten(),
})
})
.collect());
}
}
let rows = crate::db::sql::query(
"SELECT track, entry_type, amount, balance_after, model, operator, ref_type, ref_id, note, created_at
FROM credit_ledger
WHERE user_id = ?
@@ -178,12 +331,18 @@ pub async fn list_user_ledger(
.map(|r| {
json!({
"track": r.try_get::<String, _>("track").unwrap_or_default(),
"billingTrack": if r.try_get::<String, _>("track").unwrap_or_default() == "points" { "comfy" } else { "api" },
"entryType": r.try_get::<String, _>("entry_type").unwrap_or_default(),
"amount": r.try_get::<f64, _>("amount").unwrap_or(0.0),
"balanceAfter": r.try_get::<f64, _>("balance_after").unwrap_or(0.0),
"model": r.try_get::<Option<String>, _>("model").ok().flatten(),
"operator": r.try_get::<Option<String>, _>("operator").ok().flatten(),
"refType": r.try_get::<Option<String>, _>("ref_type").ok().flatten(),
"refId": r.try_get::<Option<String>, _>("ref_id").ok().flatten(),
"grantKind": Value::Null,
"taskType": Value::Null,
"packageAmount": 0,
"balanceAmount": 0,
"note": r.try_get::<Option<String>, _>("note").ok().flatten(),
"createdAt": r.try_get::<Option<String>, _>("created_at").ok().flatten(),
})
@@ -192,9 +351,9 @@ pub async fn list_user_ledger(
}
/// Aggregate a user's ledger by (track, entry_type): sum amount + count.
pub async fn summarize_user_ledger(db: &SqlitePool, user_id: i64) -> AppResult<Vec<Value>> {
pub async fn summarize_user_ledger(db: &AnyPool, user_id: i64) -> AppResult<Vec<Value>> {
ensure_ledger_schema(db).await?;
let rows = sqlx::query(
let rows = crate::db::sql::query(
"SELECT track, entry_type, COALESCE(SUM(amount),0) AS total, COUNT(*) AS cnt
FROM credit_ledger WHERE user_id = ?
GROUP BY track, entry_type ORDER BY track, entry_type",
@@ -218,14 +377,17 @@ pub async fn summarize_user_ledger(db: &SqlitePool, user_id: i64) -> AppResult<V
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
async fn mem_pool() -> SqlitePool {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
async fn mem_pool() -> AnyPool {
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
ensure_ledger_schema(&pool).await.unwrap();
pool
}
@@ -253,7 +415,7 @@ mod tests {
.unwrap();
tx.commit().await.unwrap();
let row: (i64, String, String, f64, f64) = sqlx::query_as(
let row: (i64, String, String, f64, f64) = crate::db::sql::query_as(
"SELECT user_id, track, entry_type, amount, balance_after
FROM credit_ledger WHERE user_id = 7",
)
@@ -267,29 +429,85 @@ mod tests {
async fn apply_points_delta_upserts_balance_and_writes_ledger() {
let pool = mem_pool().await;
// 首次 +100(充值)
let bal = apply_points_delta(&pool, 9, 100.0, "recharge", None, Some("system"), Some("recharge_record"), None).await.unwrap();
let bal = apply_points_delta(
&pool,
9,
100.0,
"recharge",
None,
Some("system"),
Some("recharge_record"),
None,
)
.await
.unwrap();
assert_eq!(bal, 100.0);
// 再 -30(消费)
let bal2 = apply_points_delta(&pool, 9, -30.0, "consume", Some("wf/video"), Some("system"), Some("workflow"), None).await.unwrap();
let bal2 = apply_points_delta(
&pool,
9,
-30.0,
"consume",
Some("wf/video"),
Some("system"),
Some("workflow"),
None,
)
.await
.unwrap();
assert_eq!(bal2, 70.0);
// user_credits 缓存 = 70
let cached: f64 = sqlx::query_scalar("SELECT points_balance FROM user_credits WHERE user_id = 9").fetch_one(&pool).await.unwrap();
let cached: f64 = crate::db::sql::query_scalar(
"SELECT points_balance FROM user_credits WHERE user_id = 9",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(cached, 70.0);
// ledger 两行
let n: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM credit_ledger WHERE user_id=9 AND track='points'").fetch_one(&pool).await.unwrap();
let n: i64 = crate::db::sql::query_scalar(
"SELECT COUNT(*) FROM credit_ledger WHERE user_id=9 AND track='points'",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(n, 2);
}
#[tokio::test]
async fn list_and_summarize_user_ledger() {
let pool = mem_pool().await;
apply_points_delta(&pool, 9, 100.0, "recharge", None, Some("system"), Some("rr"), None).await.unwrap();
apply_points_delta(&pool, 9, -10.0, "consume", Some("wf:1"), Some("system"), Some("run"), None).await.unwrap();
apply_points_delta(
&pool,
9,
100.0,
"recharge",
None,
Some("system"),
Some("rr"),
None,
)
.await
.unwrap();
apply_points_delta(
&pool,
9,
-10.0,
"consume",
Some("wf:1"),
Some("system"),
Some("run"),
None,
)
.await
.unwrap();
// 明细:user 9 两行,track=points
let rows = list_user_ledger(&pool, 9, None, None, 50, 0).await.unwrap();
assert_eq!(rows.len(), 2);
// 聚合:points 轨 consume 合计 -10
let sum = summarize_user_ledger(&pool, 9).await.unwrap();
assert!(sum.iter().any(|v| v["track"]=="points" && v["entryType"]=="consume"));
assert!(sum
.iter()
.any(|v| v["track"] == "points" && v["entryType"] == "consume"));
}
}
+216 -157
View File
@@ -1,3 +1,4 @@
pub mod account;
pub mod ledger;
use crate::{
@@ -7,12 +8,12 @@ use crate::{
use chrono::Timelike;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sqlx::{FromRow, Row, SqlitePool};
use sqlx::{AnyPool, FromRow, Row};
use std::sync::Arc;
#[derive(Clone)]
pub struct UsageTracker {
db: Arc<SqlitePool>,
db: Arc<AnyPool>,
}
#[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq)]
@@ -43,7 +44,7 @@ pub struct ComfyTaskBillingEvent {
}
impl UsageTracker {
pub fn new(db: Arc<SqlitePool>) -> Self {
pub fn new(db: Arc<AnyPool>) -> Self {
Self { db }
}
@@ -57,6 +58,7 @@ impl UsageTracker {
) -> AppResult<()> {
Self::ensure_schema(&self.db).await?;
crate::services::billing::ledger::ensure_ledger_schema(&self.db).await?;
account::ensure_schema(&self.db).await?;
let prompt_tokens = prompt_tokens.max(0);
let completion_tokens = completion_tokens.max(0);
let total_tokens = prompt_tokens + completion_tokens;
@@ -64,10 +66,21 @@ impl UsageTracker {
let model = model.into();
let mut tx = self.db.begin().await?;
sqlx::query(
let owner_user_id: Option<i64> =
crate::db::sql::query_scalar("SELECT user_id FROM downstream_api_keys WHERE id = ?")
.bind(key_id)
.fetch_optional(&mut *tx)
.await?
.flatten();
if let Some(user_id) = owner_user_id {
// Bootstrap from the legacy per-key balance before incrementing
// used_cost, otherwise the first account debit is counted twice.
account::ensure_account_in_tx(&mut tx, user_id).await?;
}
let usage_log_id: i64 = crate::db::sql::insert_id(
"INSERT INTO usage_logs
(downstream_key_id, model, prompt_tokens, completion_tokens, total_tokens, cost)
VALUES (?, ?, ?, ?, ?, ?)",
VALUES (?, ?, ?, ?, ?, ?) RETURNING id",
)
.bind(key_id)
.bind(&model)
@@ -75,14 +88,14 @@ impl UsageTracker {
.bind(completion_tokens)
.bind(total_tokens)
.bind(cost)
.execute(&mut *tx)
.fetch_one(&mut *tx)
.await?;
sqlx::query(
crate::db::sql::query(
"UPDATE downstream_api_keys
SET used_requests = COALESCE(used_requests,0)+1,
used_cost = COALESCE(used_cost,0)+?,
last_used_at = datetime('now'), updated_at = datetime('now')
last_used_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP
WHERE id = ?",
)
.bind(cost)
@@ -90,16 +103,26 @@ impl UsageTracker {
.execute(&mut *tx)
.await?;
// Resolve owning user + new used_cost for the ledger snapshot.
let owner: Option<(Option<i64>, Option<f64>, Option<f64>)> = sqlx::query_as(
"SELECT user_id, max_cost, used_cost FROM downstream_api_keys WHERE id = ?",
)
.bind(key_id)
.fetch_optional(&mut *tx)
.await?;
if let Some((Some(user_id), max_cost, used_cost)) = owner {
let used = used_cost.unwrap_or(0.0);
let remaining = max_cost.map(|m| (m - used).max(0.0)).unwrap_or(used);
if let Some(user_id) = owner_user_id {
if cost > 0.0 {
account::debit_for_usage_in_tx(
&mut tx,
user_id,
account::units_to_micros(cost)?,
&model,
"usage_log",
&usage_log_id.to_string(),
&format!("usage:{usage_log_id}:api"),
)
.await?;
}
let account_remaining_micros: i64 = crate::db::sql::query_scalar(
"SELECT available_micros FROM api_billing_accounts WHERE user_id = ?",
)
.bind(user_id)
.fetch_one(&mut *tx)
.await?;
let remaining = account::micros_to_units(account_remaining_micros);
crate::services::billing::ledger::record_ledger_entry(
&mut tx,
crate::services::billing::ledger::LedgerEntry {
@@ -129,7 +152,7 @@ impl UsageTracker {
) -> AppResult<UsageSummary> {
Self::ensure_schema(&self.db).await?;
let summary = sqlx::query_as::<_, UsageSummary>(
let summary = crate::db::sql::query_as::<_, UsageSummary>(
"SELECT
? AS downstream_key_id,
COUNT(*) AS request_count,
@@ -157,7 +180,7 @@ impl UsageTracker {
pub async fn check_rate_limit(&self, key_id: i64) -> AppResult<()> {
enforce_key_request_limit(&self.db, key_id, 1).await?;
let row: Option<(Option<f64>, Option<f64>, Option<bool>)> = sqlx::query_as(
let row: Option<(Option<f64>, Option<f64>, Option<bool>)> = crate::db::sql::query_as(
"SELECT max_cost, used_cost, enabled
FROM downstream_api_keys
WHERE id = ?",
@@ -188,6 +211,7 @@ impl UsageTracker {
) -> AppResult<Value> {
Self::ensure_schema(&self.db).await?;
crate::services::billing::ledger::ensure_ledger_schema(&self.db).await?;
account::ensure_schema(&self.db).await?;
Self::ensure_proxy_log_task_columns(&self.db).await?;
let prompt_tokens = event.prompt_tokens.max(0);
let completion_tokens = event.completion_tokens.max(0);
@@ -197,10 +221,20 @@ impl UsageTracker {
let mut tx = self.db.begin().await?;
if let Some(key_id) = event.key_id {
sqlx::query(
let owner_user_id: Option<i64> = crate::db::sql::query_scalar(
"SELECT user_id FROM downstream_api_keys WHERE id = ?",
)
.bind(key_id)
.fetch_optional(&mut *tx)
.await?
.flatten();
if let Some(user_id) = owner_user_id {
account::ensure_account_in_tx(&mut tx, user_id).await?;
}
let usage_log_id: i64 = crate::db::sql::insert_id(
"INSERT INTO usage_logs
(downstream_key_id, model, prompt_tokens, completion_tokens, total_tokens, cost)
VALUES (?, ?, ?, ?, ?, ?)",
VALUES (?, ?, ?, ?, ?, ?) RETURNING id",
)
.bind(key_id)
.bind(&event.model)
@@ -208,15 +242,15 @@ impl UsageTracker {
.bind(completion_tokens)
.bind(total_tokens)
.bind(cost)
.execute(&mut *tx)
.fetch_one(&mut *tx)
.await?;
sqlx::query(
crate::db::sql::query(
"UPDATE downstream_api_keys
SET used_requests = COALESCE(used_requests, 0) + 1,
used_cost = COALESCE(used_cost, 0) + ?,
last_used_at = datetime('now'),
updated_at = datetime('now')
last_used_at = CURRENT_TIMESTAMP,
updated_at = CURRENT_TIMESTAMP
WHERE id = ?",
)
.bind(cost)
@@ -224,15 +258,26 @@ impl UsageTracker {
.execute(&mut *tx)
.await?;
let owner: Option<(Option<i64>, Option<f64>, Option<f64>)> = sqlx::query_as(
"SELECT user_id, max_cost, used_cost FROM downstream_api_keys WHERE id = ?",
)
.bind(key_id)
.fetch_optional(&mut *tx)
.await?;
if let Some((Some(user_id), max_cost, used_cost)) = owner {
let used = used_cost.unwrap_or(0.0);
let remaining = max_cost.map(|m| (m - used).max(0.0)).unwrap_or(used);
if let Some(user_id) = owner_user_id {
if cost > 0.0 {
account::debit_for_usage_in_tx(
&mut tx,
user_id,
account::units_to_micros(cost)?,
&event.model,
"usage_log",
&usage_log_id.to_string(),
&format!("usage:{usage_log_id}:api"),
)
.await?;
}
let account_remaining_micros: i64 = crate::db::sql::query_scalar(
"SELECT available_micros FROM api_billing_accounts WHERE user_id = ?",
)
.bind(user_id)
.fetch_one(&mut *tx)
.await?;
let remaining = account::micros_to_units(account_remaining_micros);
crate::services::billing::ledger::record_ledger_entry(
&mut tx,
crate::services::billing::ledger::LedgerEntry {
@@ -252,7 +297,7 @@ impl UsageTracker {
}
}
sqlx::query(
crate::db::sql::query(
"INSERT INTO proxy_logs (
downstream_api_key_id, model_requested, model_actual, status, http_status,
prompt_tokens, completion_tokens, total_tokens, estimated_cost,
@@ -290,7 +335,7 @@ impl UsageTracker {
Self::ensure_proxy_log_task_columns(&self.db).await?;
let billing_details = comfy_task_billing_details(&event, "failed", 0.0);
sqlx::query(
crate::db::sql::query(
"INSERT INTO proxy_logs (
downstream_api_key_id, model_requested, model_actual, status, http_status,
prompt_tokens, completion_tokens, total_tokens, estimated_cost, error_message,
@@ -325,39 +370,39 @@ impl UsageTracker {
Self::create_daily_usage_snapshot_for_db(&self.db).await
}
pub async fn create_daily_usage_snapshot_for_db(db: &SqlitePool) -> AppResult<()> {
pub async fn create_daily_usage_snapshot_for_db(db: &AnyPool) -> AppResult<()> {
Self::ensure_schema(db).await?;
sqlx::query(
crate::db::sql::query(
"INSERT INTO usage_daily_snapshots
(local_day, downstream_key_id, request_count, prompt_tokens,
completion_tokens, total_tokens, total_cost, updated_at)
SELECT
date(created_at, 'localtime') AS local_day,
substr(created_at, 1, 10) AS local_day,
downstream_key_id,
COUNT(*) AS request_count,
COALESCE(SUM(prompt_tokens), 0) AS prompt_tokens,
COALESCE(SUM(completion_tokens), 0) AS completion_tokens,
COALESCE(SUM(total_tokens), 0) AS total_tokens,
COALESCE(SUM(cost), 0.0) AS total_cost,
datetime('now') AS updated_at
CURRENT_TIMESTAMP AS updated_at
FROM usage_logs
GROUP BY date(created_at, 'localtime'), downstream_key_id
GROUP BY substr(created_at, 1, 10), downstream_key_id
ON CONFLICT(local_day, downstream_key_id) DO UPDATE SET
request_count = excluded.request_count,
prompt_tokens = excluded.prompt_tokens,
completion_tokens = excluded.completion_tokens,
total_tokens = excluded.total_tokens,
total_cost = excluded.total_cost,
updated_at = datetime('now')",
updated_at = CURRENT_TIMESTAMP",
)
.execute(db)
.await?;
Ok(())
}
pub async fn project_proxy_usage_for_db(db: &SqlitePool) -> AppResult<i64> {
pub async fn project_proxy_usage_for_db(db: &AnyPool) -> AppResult<i64> {
Self::ensure_schema(db).await?;
let last_id: i64 = sqlx::query_scalar(
let last_id: i64 = crate::db::sql::query_scalar(
"SELECT COALESCE(last_proxy_log_id, 0)
FROM analytics_projection_checkpoints
WHERE projector_key = 'usage-aggregates-v1'",
@@ -366,7 +411,7 @@ impl UsageTracker {
.await?
.unwrap_or(0);
let rows = sqlx::query(
let rows = crate::db::sql::query(
r#"
SELECT
l.id,
@@ -450,13 +495,13 @@ impl UsageTracker {
total_tokens as f64 / 500_000.0
};
sqlx::query(
crate::db::sql::query(
r#"
INSERT INTO site_day_usage
(local_day, site_id, total_calls, success_calls, failed_calls,
total_tokens, total_summary_spend, total_site_spend,
total_latency_ms, latency_count, updated_at)
VALUES (?, ?, 1, ?, ?, ?, ?, ?, ?, ?, datetime('now'))
VALUES (?, ?, 1, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(local_day, site_id) DO UPDATE SET
total_calls = total_calls + 1,
success_calls = success_calls + excluded.success_calls,
@@ -466,7 +511,7 @@ impl UsageTracker {
total_site_spend = total_site_spend + excluded.total_site_spend,
total_latency_ms = total_latency_ms + excluded.total_latency_ms,
latency_count = latency_count + excluded.latency_count,
updated_at = datetime('now')
updated_at = CURRENT_TIMESTAMP
"#,
)
.bind(&local_day)
@@ -481,13 +526,13 @@ impl UsageTracker {
.execute(&mut *tx)
.await?;
sqlx::query(
crate::db::sql::query(
r#"
INSERT INTO site_hour_usage
(bucket_start_utc, site_id, total_calls, success_calls, failed_calls,
total_tokens, total_summary_spend, total_site_spend,
total_latency_ms, latency_count, updated_at)
VALUES (?, ?, 1, ?, ?, ?, ?, ?, ?, ?, datetime('now'))
VALUES (?, ?, 1, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(bucket_start_utc, site_id) DO UPDATE SET
total_calls = total_calls + 1,
success_calls = success_calls + excluded.success_calls,
@@ -497,7 +542,7 @@ impl UsageTracker {
total_site_spend = total_site_spend + excluded.total_site_spend,
total_latency_ms = total_latency_ms + excluded.total_latency_ms,
latency_count = latency_count + excluded.latency_count,
updated_at = datetime('now')
updated_at = CURRENT_TIMESTAMP
"#,
)
.bind(&bucket_start_utc)
@@ -512,12 +557,12 @@ impl UsageTracker {
.execute(&mut *tx)
.await?;
sqlx::query(
crate::db::sql::query(
r#"
INSERT INTO model_day_usage
(local_day, site_id, model, total_calls, success_calls, failed_calls,
total_tokens, total_spend, total_latency_ms, latency_count, updated_at)
VALUES (?, ?, ?, 1, ?, ?, ?, ?, ?, ?, datetime('now'))
VALUES (?, ?, ?, 1, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(local_day, site_id, model) DO UPDATE SET
total_calls = total_calls + 1,
success_calls = success_calls + excluded.success_calls,
@@ -526,7 +571,7 @@ impl UsageTracker {
total_spend = total_spend + excluded.total_spend,
total_latency_ms = total_latency_ms + excluded.total_latency_ms,
latency_count = latency_count + excluded.latency_count,
updated_at = datetime('now')
updated_at = CURRENT_TIMESTAMP
"#,
)
.bind(&local_day)
@@ -542,19 +587,19 @@ impl UsageTracker {
.await?;
}
sqlx::query(
crate::db::sql::query(
r#"
INSERT INTO analytics_projection_checkpoints
(projector_key, time_zone, last_proxy_log_id, watermark_created_at,
last_projected_at, last_successful_at, last_error, updated_at)
VALUES ('usage-aggregates-v1', 'Local', ?, datetime('now'), datetime('now'), datetime('now'), NULL, datetime('now'))
VALUES ('usage-aggregates-v1', 'Local', ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, NULL, CURRENT_TIMESTAMP)
ON CONFLICT(projector_key) DO UPDATE SET
last_proxy_log_id = excluded.last_proxy_log_id,
watermark_created_at = excluded.watermark_created_at,
last_projected_at = excluded.last_projected_at,
last_successful_at = excluded.last_successful_at,
last_error = NULL,
updated_at = datetime('now')
updated_at = CURRENT_TIMESTAMP
"#,
)
.bind(max_id)
@@ -564,8 +609,12 @@ impl UsageTracker {
Ok(max_id - last_id)
}
async fn ensure_schema(db: &SqlitePool) -> AppResult<()> {
sqlx::query(
async fn ensure_schema(db: &AnyPool) -> AppResult<()> {
if crate::db::sql::runtime_kind() != crate::db::DbKind::Sqlite {
return Ok(());
}
crate::db::sql::query(
"CREATE TABLE IF NOT EXISTS usage_logs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
downstream_key_id INTEGER,
@@ -574,18 +623,18 @@ impl UsageTracker {
completion_tokens INTEGER DEFAULT 0,
total_tokens INTEGER DEFAULT 0,
cost REAL DEFAULT 0.0,
created_at TEXT DEFAULT (datetime('now'))
created_at TEXT DEFAULT (CURRENT_TIMESTAMP)
)",
)
.execute(db)
.await?;
sqlx::query(
crate::db::sql::query(
"CREATE INDEX IF NOT EXISTS usage_logs_downstream_key_created_idx
ON usage_logs (downstream_key_id, created_at)",
)
.execute(db)
.await?;
sqlx::query(
crate::db::sql::query(
"CREATE TABLE IF NOT EXISTS usage_daily_snapshots (
id INTEGER PRIMARY KEY AUTOINCREMENT,
local_day TEXT NOT NULL,
@@ -595,14 +644,14 @@ impl UsageTracker {
completion_tokens INTEGER NOT NULL DEFAULT 0,
total_tokens INTEGER NOT NULL DEFAULT 0,
total_cost REAL NOT NULL DEFAULT 0.0,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now')),
created_at TEXT DEFAULT (CURRENT_TIMESTAMP),
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP),
UNIQUE(local_day, downstream_key_id)
)",
)
.execute(db)
.await?;
sqlx::query(
crate::db::sql::query(
"CREATE TABLE IF NOT EXISTS analytics_projection_checkpoints (
projector_key TEXT PRIMARY KEY NOT NULL,
time_zone TEXT NOT NULL DEFAULT 'Local',
@@ -619,13 +668,13 @@ impl UsageTracker {
lease_owner TEXT,
lease_token TEXT,
lease_expires_at TEXT,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now'))
created_at TEXT DEFAULT (CURRENT_TIMESTAMP),
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP)
)",
)
.execute(db)
.await?;
sqlx::query(
crate::db::sql::query(
"CREATE TABLE IF NOT EXISTS model_day_usage (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
local_day TEXT NOT NULL,
@@ -638,14 +687,14 @@ impl UsageTracker {
total_spend REAL NOT NULL DEFAULT 0,
total_latency_ms INTEGER NOT NULL DEFAULT 0,
latency_count INTEGER NOT NULL DEFAULT 0,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now')),
created_at TEXT DEFAULT (CURRENT_TIMESTAMP),
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP),
UNIQUE(local_day, site_id, model)
)",
)
.execute(db)
.await?;
sqlx::query(
crate::db::sql::query(
"CREATE TABLE IF NOT EXISTS site_day_usage (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
local_day TEXT NOT NULL,
@@ -658,14 +707,14 @@ impl UsageTracker {
total_site_spend REAL NOT NULL DEFAULT 0,
total_latency_ms INTEGER NOT NULL DEFAULT 0,
latency_count INTEGER NOT NULL DEFAULT 0,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now')),
created_at TEXT DEFAULT (CURRENT_TIMESTAMP),
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP),
UNIQUE(local_day, site_id)
)",
)
.execute(db)
.await?;
sqlx::query(
crate::db::sql::query(
"CREATE TABLE IF NOT EXISTS site_hour_usage (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
bucket_start_utc TEXT NOT NULL,
@@ -678,8 +727,8 @@ impl UsageTracker {
total_site_spend REAL NOT NULL DEFAULT 0,
total_latency_ms INTEGER NOT NULL DEFAULT 0,
latency_count INTEGER NOT NULL DEFAULT 0,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now')),
created_at TEXT DEFAULT (CURRENT_TIMESTAMP),
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP),
UNIQUE(bucket_start_utc, site_id)
)",
)
@@ -688,7 +737,7 @@ impl UsageTracker {
Ok(())
}
async fn ensure_proxy_log_task_columns(db: &SqlitePool) -> AppResult<()> {
async fn ensure_proxy_log_task_columns(db: &AnyPool) -> AppResult<()> {
for (name, definition) in [
("client_family", "TEXT"),
("client_app_id", "TEXT"),
@@ -708,13 +757,13 @@ impl UsageTracker {
}
async fn ensure_table_column(
db: &SqlitePool,
db: &AnyPool,
table: &str,
name: &str,
definition: &str,
) -> AppResult<()> {
let info_sql = format!("PRAGMA table_info({table})");
let rows = sqlx::query(&info_sql).fetch_all(db).await?;
let rows = crate::db::sql::query(&info_sql).fetch_all(db).await?;
if rows.is_empty() {
return Ok(());
}
@@ -726,7 +775,7 @@ async fn ensure_table_column(
return Ok(());
}
let alter_sql = format!("ALTER TABLE {table} ADD COLUMN {name} {definition}");
sqlx::query(&alter_sql).execute(db).await?;
crate::db::sql::query(&alter_sql).execute(db).await?;
Ok(())
}
@@ -796,21 +845,21 @@ fn comfy_task_billing_details(
}
async fn touch_projection_checkpoint(
db: &SqlitePool,
db: &AnyPool,
last_proxy_log_id: i64,
error: Option<&str>,
) -> AppResult<()> {
sqlx::query(
crate::db::sql::query(
r#"
INSERT INTO analytics_projection_checkpoints
(projector_key, time_zone, last_proxy_log_id, last_projected_at,
last_successful_at, last_error, updated_at)
VALUES ('usage-aggregates-v1', 'Local', ?, datetime('now'), datetime('now'), ?, datetime('now'))
VALUES ('usage-aggregates-v1', 'Local', ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, ?, CURRENT_TIMESTAMP)
ON CONFLICT(projector_key) DO UPDATE SET
last_projected_at = datetime('now'),
last_successful_at = CASE WHEN ? IS NULL THEN datetime('now') ELSE last_successful_at END,
last_projected_at = CURRENT_TIMESTAMP,
last_successful_at = CASE WHEN ? IS NULL THEN CURRENT_TIMESTAMP ELSE last_successful_at END,
last_error = ?,
updated_at = datetime('now')
updated_at = CURRENT_TIMESTAMP
"#,
)
.bind(last_proxy_log_id)
@@ -884,16 +933,19 @@ fn fallback_token_cost(tokens: i64, platform: Option<&str>) -> f64 {
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
async fn pool() -> SqlitePool {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
async fn pool() -> AnyPool {
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE downstream_api_keys (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
@@ -905,13 +957,13 @@ mod tests {
used_requests INTEGER DEFAULT 0,
user_id INTEGER,
last_used_at TEXT,
updated_at TEXT DEFAULT (datetime('now'))
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP)
)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO downstream_api_keys
(id, name, key, enabled, max_cost, used_cost, max_requests, used_requests)
VALUES (1, 'test', 'sk-test', 1, 1.0, 0.0, 5, 0)",
@@ -923,13 +975,16 @@ mod tests {
pool
}
async fn billing_test_pool() -> SqlitePool {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
sqlx::query(
async fn billing_test_pool() -> AnyPool {
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
crate::db::sql::query(
"CREATE TABLE users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT NOT NULL,
@@ -944,7 +999,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE downstream_api_keys (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL,
@@ -956,20 +1011,20 @@ mod tests {
used_requests INTEGER DEFAULT 0,
user_id INTEGER,
last_used_at TEXT,
updated_at TEXT DEFAULT (datetime('now'))
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP)
)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO users (id, username, email, password_hash, role, status)
VALUES (9, 'creator', 'creator@example.test', 'hash', 'user', 'active')",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO downstream_api_keys (id, name, key, max_cost, used_cost, user_id)
VALUES (7, 'default', 'sk-user-test', 10.0, 2.0, 9)",
)
@@ -979,8 +1034,8 @@ mod tests {
pool
}
async fn create_proxy_logs(pool: &SqlitePool) {
sqlx::query(
async fn create_proxy_logs(pool: &AnyPool) {
crate::db::sql::query(
"CREATE TABLE proxy_logs (
id INTEGER PRIMARY KEY,
downstream_api_key_id INTEGER,
@@ -1003,7 +1058,7 @@ mod tests {
content_summary TEXT,
request_content TEXT,
response_content TEXT,
created_at TEXT DEFAULT (datetime('now'))
created_at TEXT DEFAULT (CURRENT_TIMESTAMP)
)",
)
.execute(pool)
@@ -1059,7 +1114,7 @@ mod tests {
.await
.unwrap();
let row: (String, String, f64, Option<String>) = sqlx::query_as(
let row: (String, String, f64, Option<String>) = crate::db::sql::query_as(
"SELECT track, entry_type, amount, model FROM credit_ledger
WHERE user_id = 9 ORDER BY id DESC LIMIT 1",
)
@@ -1077,7 +1132,7 @@ mod tests {
let tracker = UsageTracker::new(Arc::new(pool().await));
tracker.check_rate_limit(1).await.unwrap();
sqlx::query("UPDATE downstream_api_keys SET used_requests = 5 WHERE id = 1")
crate::db::sql::query("UPDATE downstream_api_keys SET used_requests = 5 WHERE id = 1")
.execute(&*tracker.db)
.await
.unwrap();
@@ -1096,7 +1151,7 @@ mod tests {
.unwrap();
tracker.create_daily_usage_snapshot().await.unwrap();
let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM usage_daily_snapshots")
let count: i64 = crate::db::sql::query_scalar("SELECT COUNT(*) FROM usage_daily_snapshots")
.fetch_one(&*tracker.db)
.await
.unwrap();
@@ -1111,16 +1166,13 @@ mod tests {
let mut event = test_comfy_event(0.35);
event.key_id = Some(7);
let billing = tracker
.record_comfy_task_success(event)
.await
.unwrap();
let billing = tracker.record_comfy_task_success(event).await.unwrap();
assert_eq!(billing["source"], "comfy-native-task");
assert_eq!(billing["charged"], true);
assert_eq!(billing["totalCost"], 0.35);
let usage: (i64, f64) = sqlx::query_as(
let usage: (i64, f64) = crate::db::sql::query_as(
"SELECT COUNT(*), COALESCE(SUM(cost), 0.0) FROM usage_logs WHERE downstream_key_id = 7",
)
.fetch_one(&*tracker.db)
@@ -1128,20 +1180,22 @@ mod tests {
.unwrap();
assert_eq!(usage, (1, 0.35));
let key_usage: (i64, f64) =
sqlx::query_as("SELECT used_requests, used_cost FROM downstream_api_keys WHERE id = 7")
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(key_usage, (1, 2.35));
let ledger: (String, String, f64, f64, Option<String>, Option<String>) = sqlx::query_as(
"SELECT track, entry_type, amount, balance_after, model, note
FROM credit_ledger WHERE user_id = 9 ORDER BY id DESC LIMIT 1",
let key_usage: (i64, f64) = crate::db::sql::query_as(
"SELECT used_requests, used_cost FROM downstream_api_keys WHERE id = 7",
)
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(key_usage, (1, 2.35));
let ledger: (String, String, f64, f64, Option<String>, Option<String>) =
crate::db::sql::query_as(
"SELECT track, entry_type, amount, balance_after, model, note
FROM credit_ledger WHERE user_id = 9 ORDER BY id DESC LIMIT 1",
)
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(ledger.0, "usd");
assert_eq!(ledger.1, "consume");
assert_eq!(ledger.2, -0.35);
@@ -1149,7 +1203,7 @@ mod tests {
assert_eq!(ledger.4.as_deref(), Some("seedance-comfy"));
assert_eq!(ledger.5.as_deref(), Some("comfy_task"));
let proxy_log: (String, f64, String, String) = sqlx::query_as(
let proxy_log: (String, f64, String, String) = crate::db::sql::query_as(
"SELECT status, estimated_cost, client_app_id, billing_details FROM proxy_logs WHERE downstream_api_key_id = 7",
)
.fetch_one(&*tracker.db)
@@ -1174,21 +1228,23 @@ mod tests {
assert_eq!(billing["charged"], false);
assert_eq!(billing["totalCost"], 0.0);
let usage_count: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM usage_logs WHERE downstream_key_id = 1")
.fetch_one(&*tracker.db)
.await
.unwrap();
let usage_count: i64 = crate::db::sql::query_scalar(
"SELECT COUNT(*) FROM usage_logs WHERE downstream_key_id = 1",
)
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(usage_count, 0);
let key_usage: (i64, f64) =
sqlx::query_as("SELECT used_requests, used_cost FROM downstream_api_keys WHERE id = 1")
.fetch_one(&*tracker.db)
.await
.unwrap();
let key_usage: (i64, f64) = crate::db::sql::query_as(
"SELECT used_requests, used_cost FROM downstream_api_keys WHERE id = 1",
)
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(key_usage, (0, 0.0));
let proxy_log: (String, f64, String) = sqlx::query_as(
let proxy_log: (String, f64, String) = crate::db::sql::query_as(
"SELECT status, estimated_cost, error_message FROM proxy_logs WHERE downstream_api_key_id = 1",
)
.fetch_one(&*tracker.db)
@@ -1203,7 +1259,7 @@ mod tests {
#[tokio::test]
async fn projects_proxy_logs_to_site_and_model_usage() {
let pool = pool().await;
sqlx::query(
crate::db::sql::query(
"CREATE TABLE sites (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
@@ -1213,7 +1269,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE accounts (
id INTEGER PRIMARY KEY,
site_id INTEGER NOT NULL
@@ -1222,7 +1278,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE proxy_logs (
id INTEGER PRIMARY KEY,
account_id INTEGER,
@@ -1238,15 +1294,17 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO sites (id, name, platform) VALUES (10, 'Site', 'new-api')")
crate::db::sql::query(
"INSERT INTO sites (id, name, platform) VALUES (10, 'Site', 'new-api')",
)
.execute(&pool)
.await
.unwrap();
crate::db::sql::query("INSERT INTO accounts (id, site_id) VALUES (20, 10)")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO accounts (id, site_id) VALUES (20, 10)")
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"INSERT INTO proxy_logs
(id, account_id, model_requested, model_actual, status, latency_ms, total_tokens, estimated_cost, created_at)
VALUES
@@ -1261,7 +1319,7 @@ mod tests {
.unwrap();
assert_eq!(projected, 1);
let model_usage: (i64, i64, f64, i64) = sqlx::query_as(
let model_usage: (i64, i64, f64, i64) = crate::db::sql::query_as(
"SELECT total_calls, total_tokens, total_spend, total_latency_ms
FROM model_day_usage
WHERE site_id = 10 AND model = 'gpt-test'",
@@ -1271,7 +1329,7 @@ mod tests {
.unwrap();
assert_eq!(model_usage, (1, 1000, 0.12, 120));
let site_usage: (i64, i64, f64) = sqlx::query_as(
let site_usage: (i64, i64, f64) = crate::db::sql::query_as(
"SELECT total_calls, success_calls, total_site_spend
FROM site_day_usage
WHERE site_id = 10 AND local_day = '2026-06-06'",
@@ -1281,11 +1339,12 @@ mod tests {
.unwrap();
assert_eq!(site_usage, (1, 1, 0.12));
let hour_bucket: String =
sqlx::query_scalar("SELECT bucket_start_utc FROM site_hour_usage WHERE site_id = 10")
.fetch_one(&pool)
.await
.unwrap();
let hour_bucket: String = crate::db::sql::query_scalar(
"SELECT bucket_start_utc FROM site_hour_usage WHERE site_id = 10",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(hour_bucket, "2026-06-06 08:00:00");
}
}
+31 -26
View File
@@ -13,7 +13,7 @@
//! consistent pair without rotating the secret.
use crate::error::{AppError, AppResult};
use sqlx::SqlitePool;
use sqlx::AnyPool;
use uuid::Uuid;
/// In-memory snapshot of system-owned downstream keys.
@@ -45,18 +45,18 @@ fn generate_comfyui_agent_secret() -> String {
/// | present | present (same) | no-op |
/// | present | present (diff) | rewrite settings to row secret |
/// | present | present, en=0 | re-enable + warn |
pub async fn ensure_comfyui_agent_key(db: &SqlitePool) -> AppResult<(i64, String)> {
pub async fn ensure_comfyui_agent_key(db: &AnyPool) -> AppResult<(i64, String)> {
let mut tx = db.begin().await.map_err(AppError::Database)?;
// 1. Read current state from both sides.
let settings_secret: Option<String> =
sqlx::query_scalar("SELECT value FROM settings WHERE key = ?")
crate::db::sql::query_scalar("SELECT value FROM settings WHERE key = ?")
.bind(COMFYUI_AGENT_SETTINGS_KEY)
.fetch_optional(&mut *tx)
.await
.map_err(AppError::Database)?;
let existing_row: Option<(i64, String, i64)> = sqlx::query_as(
let existing_row: Option<(i64, String, i64)> = crate::db::sql::query_as(
"SELECT id, key, enabled FROM downstream_api_keys \
WHERE name = ? AND user_id IS NULL \
LIMIT 1",
@@ -76,7 +76,8 @@ pub async fn ensure_comfyui_agent_key(db: &SqlitePool) -> AppResult<(i64, String
key_id = id,
"comfyui-agent builtin key was disabled; re-enabling",
);
sqlx::query("UPDATE downstream_api_keys SET enabled = 1 WHERE id = ?")
crate::db::sql::query("UPDATE downstream_api_keys SET enabled = ? WHERE id = ?")
.bind(true)
.bind(id)
.execute(&mut *tx)
.await
@@ -112,28 +113,29 @@ pub async fn ensure_comfyui_agent_key(db: &SqlitePool) -> AppResult<(i64, String
}
async fn insert_comfyui_agent_row(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
tx: &mut sqlx::Transaction<'_, sqlx::Any>,
secret: &str,
) -> AppResult<i64> {
let result = sqlx::query(
let id: i64 = crate::db::sql::insert_id(
"INSERT INTO downstream_api_keys (name, key, user_id, enabled, description) \
VALUES (?, ?, NULL, 1, ?)",
VALUES (?, ?, NULL, ?, ?) RETURNING id",
)
.bind(COMFYUI_AGENT_NAME)
.bind(secret)
.bind(true)
.bind(COMFYUI_AGENT_DESCRIPTION)
.execute(&mut **tx)
.fetch_one(&mut **tx)
.await
.map_err(AppError::Database)?;
Ok(result.last_insert_rowid())
Ok(id)
}
async fn upsert_settings(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
tx: &mut sqlx::Transaction<'_, sqlx::Any>,
key: &str,
value: &str,
) -> AppResult<()> {
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)")
crate::db::sql::query("INSERT INTO settings (key, value) VALUES (?, ?) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value")
.bind(key)
.bind(value)
.execute(&mut **tx)
@@ -145,14 +147,17 @@ async fn upsert_settings(
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
async fn fresh_pool() -> SqlitePool {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("memory sqlite");
async fn fresh_pool() -> AnyPool {
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("memory sqlite");
sqlx::migrate!("./migrations")
.run(&pool)
.await
@@ -160,16 +165,16 @@ mod tests {
pool
}
async fn read_settings(db: &SqlitePool, key: &str) -> Option<String> {
sqlx::query_scalar("SELECT value FROM settings WHERE key = ?")
async fn read_settings(db: &AnyPool, key: &str) -> Option<String> {
crate::db::sql::query_scalar("SELECT value FROM settings WHERE key = ?")
.bind(key)
.fetch_optional(db)
.await
.unwrap()
}
async fn read_row(db: &SqlitePool) -> Option<(i64, String, i64)> {
sqlx::query_as(
async fn read_row(db: &AnyPool) -> Option<(i64, String, i64)> {
crate::db::sql::query_as(
"SELECT id, key, enabled FROM downstream_api_keys \
WHERE name = 'comfyui-agent' AND user_id IS NULL",
)
@@ -199,7 +204,7 @@ mod tests {
async fn settings_present_row_missing_recovers_from_settings() {
let pool = fresh_pool().await;
let saved = "sk-comfyui-deadbeef".to_string();
sqlx::query("INSERT INTO settings (key, value) VALUES (?, ?)")
crate::db::sql::query("INSERT INTO settings (key, value) VALUES (?, ?)")
.bind(COMFYUI_AGENT_SETTINGS_KEY)
.bind(&saved)
.execute(&pool)
@@ -215,7 +220,7 @@ mod tests {
#[tokio::test]
async fn row_present_settings_missing_writes_settings() {
let pool = fresh_pool().await;
sqlx::query(
crate::db::sql::query(
"INSERT INTO downstream_api_keys (name, key, user_id, enabled, description) \
VALUES ('comfyui-agent', 'sk-comfyui-fromrow', NULL, 1, 'pre-existing')",
)
@@ -244,7 +249,7 @@ mod tests {
#[tokio::test]
async fn disabled_row_is_re_enabled() {
let pool = fresh_pool().await;
sqlx::query(
crate::db::sql::query(
"INSERT INTO downstream_api_keys (name, key, user_id, enabled, description) \
VALUES ('comfyui-agent', 'sk-comfyui-disabled', NULL, 0, 'pre-existing')",
)
File diff suppressed because it is too large Load Diff
+1
View File
@@ -1,3 +1,4 @@
//! Comfy Cloud 工作流执行适配器 service 层。
pub mod ownership;
pub mod site;
pub mod translate;
+269
View File
@@ -0,0 +1,269 @@
use crate::{error::AppResult, AppState};
use serde_json::{json, Value};
use sqlx::{AnyPool, Row};
fn string_field(asset: &Value, field: &str) -> Option<String> {
asset
.get(field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
}
fn tags_json(asset: &Value) -> String {
asset
.get("tags")
.and_then(Value::as_array)
.cloned()
.map(Value::Array)
.unwrap_or_else(|| json!([]))
.to_string()
}
pub fn asset_is_shared_model(asset: &Value) -> bool {
asset
.get("tags")
.and_then(Value::as_array)
.is_some_and(|tags| tags.iter().any(|tag| tag.as_str() == Some("models")))
}
pub async fn record_job(db: &AnyPool, user_id: i64, job_id: &str) -> AppResult<()> {
let job_id = job_id.trim();
if job_id.is_empty() {
return Ok(());
}
let sql = if crate::db::sql::runtime_kind() == crate::db::DbKind::MySql {
"INSERT INTO comfy_cloud_jobs (job_id, user_id, updated_at)
VALUES (?, ?, CURRENT_TIMESTAMP)
ON DUPLICATE KEY UPDATE updated_at = IF(user_id = VALUES(user_id), VALUES(updated_at), updated_at)"
} else {
"INSERT INTO comfy_cloud_jobs (job_id, user_id, updated_at)
VALUES (?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(job_id) DO UPDATE SET updated_at = CURRENT_TIMESTAMP
WHERE comfy_cloud_jobs.user_id = excluded.user_id"
};
crate::db::sql::query(sql)
.bind(job_id)
.bind(user_id)
.execute(db)
.await?;
Ok(())
}
pub async fn job_is_owned_by(db: &AnyPool, user_id: i64, job_id: &str) -> AppResult<bool> {
Ok(crate::db::sql::query_scalar::<_, i64>(
"SELECT 1 FROM comfy_cloud_jobs WHERE job_id = ? AND user_id = ? LIMIT 1",
)
.bind(job_id)
.bind(user_id)
.fetch_optional(db)
.await?
.is_some())
}
pub async fn owned_job_ids(db: &AnyPool, user_id: i64, limit: i64) -> AppResult<Vec<String>> {
crate::db::sql::query_scalar(
"SELECT job_id FROM comfy_cloud_jobs
WHERE user_id = ?
ORDER BY updated_at DESC, job_id DESC
LIMIT ?",
)
.bind(user_id)
.bind(limit.clamp(1, 200))
.fetch_all(db)
.await
.map_err(Into::into)
}
pub async fn record_asset(
db: &AnyPool,
user_id: i64,
job_id: Option<&str>,
asset: &Value,
) -> AppResult<()> {
let Some(asset_id) = string_field(asset, "id") else {
return Ok(());
};
let job_id = job_id
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| asset.get("job_id").and_then(Value::as_str))
.or_else(|| asset.get("prompt_id").and_then(Value::as_str));
let sql = if crate::db::sql::runtime_kind() == crate::db::DbKind::MySql {
"INSERT INTO comfy_cloud_assets
(asset_id, user_id, job_id, asset_hash, name, tags_json, snapshot_json, is_preview, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, 0, CURRENT_TIMESTAMP)
ON DUPLICATE KEY UPDATE
job_id = IF(user_id = VALUES(user_id), COALESCE(VALUES(job_id), job_id), job_id),
asset_hash = IF(user_id = VALUES(user_id), VALUES(asset_hash), asset_hash),
name = IF(user_id = VALUES(user_id), VALUES(name), name),
tags_json = IF(user_id = VALUES(user_id), VALUES(tags_json), tags_json),
snapshot_json = IF(user_id = VALUES(user_id), VALUES(snapshot_json), snapshot_json),
is_preview = IF(user_id = VALUES(user_id), 0, is_preview),
updated_at = IF(user_id = VALUES(user_id), VALUES(updated_at), updated_at)"
} else {
"INSERT INTO comfy_cloud_assets
(asset_id, user_id, job_id, asset_hash, name, tags_json, snapshot_json, is_preview, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, 0, CURRENT_TIMESTAMP)
ON CONFLICT(asset_id) DO UPDATE SET
user_id = excluded.user_id,
job_id = COALESCE(excluded.job_id, comfy_cloud_assets.job_id),
asset_hash = excluded.asset_hash,
name = excluded.name,
tags_json = excluded.tags_json,
snapshot_json = excluded.snapshot_json,
is_preview = 0,
updated_at = CURRENT_TIMESTAMP
WHERE comfy_cloud_assets.user_id = excluded.user_id"
};
crate::db::sql::query(sql)
.bind(&asset_id)
.bind(user_id)
.bind(job_id)
.bind(string_field(asset, "asset_hash"))
.bind(string_field(asset, "name"))
.bind(tags_json(asset))
.bind(asset.to_string())
.execute(db)
.await?;
if let Some(preview_id) = string_field(asset, "preview_id") {
let sql = if crate::db::sql::runtime_kind() == crate::db::DbKind::MySql {
"INSERT INTO comfy_cloud_assets (asset_id, user_id, is_preview, updated_at)
VALUES (?, ?, 1, CURRENT_TIMESTAMP)
ON DUPLICATE KEY UPDATE updated_at = IF(user_id = VALUES(user_id), VALUES(updated_at), updated_at)"
} else {
"INSERT INTO comfy_cloud_assets (asset_id, user_id, is_preview, updated_at)
VALUES (?, ?, 1, CURRENT_TIMESTAMP)
ON CONFLICT(asset_id) DO UPDATE SET updated_at = CURRENT_TIMESTAMP
WHERE comfy_cloud_assets.user_id = excluded.user_id"
};
crate::db::sql::query(sql)
.bind(preview_id)
.bind(user_id)
.execute(db)
.await?;
}
Ok(())
}
pub async fn asset_is_owned_by(db: &AnyPool, user_id: i64, asset_id: &str) -> AppResult<bool> {
Ok(crate::db::sql::query_scalar::<_, i64>(
"SELECT 1 FROM comfy_cloud_assets WHERE asset_id = ? AND user_id = ? LIMIT 1",
)
.bind(asset_id)
.bind(user_id)
.fetch_optional(db)
.await?
.is_some())
}
pub async fn asset_name_or_hash_is_owned_by(
db: &AnyPool,
user_id: i64,
value: &str,
) -> AppResult<bool> {
Ok(crate::db::sql::query_scalar::<_, i64>(
"SELECT 1 FROM comfy_cloud_assets
WHERE user_id = ? AND (name = ? OR asset_hash = ?) LIMIT 1",
)
.bind(user_id)
.bind(value)
.bind(value)
.fetch_optional(db)
.await?
.is_some())
}
pub async fn list_owned_assets(db: &AnyPool, user_id: i64) -> AppResult<Vec<Value>> {
let rows = crate::db::sql::query(
"SELECT snapshot_json FROM comfy_cloud_assets
WHERE user_id = ? AND is_preview = 0 AND snapshot_json IS NOT NULL
ORDER BY created_at DESC, asset_id DESC",
)
.bind(user_id)
.fetch_all(db)
.await?;
Ok(rows
.into_iter()
.filter_map(|row| row.try_get::<String, _>("snapshot_json").ok())
.filter_map(|raw| serde_json::from_str::<Value>(&raw).ok())
.collect())
}
pub async fn delete_asset(db: &AnyPool, user_id: i64, asset_id: &str) -> AppResult<()> {
crate::db::sql::query("DELETE FROM comfy_cloud_assets WHERE asset_id = ? AND user_id = ?")
.bind(asset_id)
.bind(user_id)
.execute(db)
.await?;
Ok(())
}
pub async fn record_prompt_response(
state: &AppState,
user_id: Option<i64>,
response: &Value,
) -> AppResult<()> {
let Some(user_id) = user_id else {
return Ok(());
};
for key in ["prompt_id", "job_id", "task_id"] {
if let Some(id) = response.get(key).and_then(Value::as_str) {
record_job(&state.db, user_id, id).await?;
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
async fn pool() -> AnyPool {
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
crate::db::sql::query("CREATE TABLE comfy_cloud_jobs (job_id TEXT PRIMARY KEY, user_id INTEGER NOT NULL, created_at TEXT, updated_at TEXT)")
.execute(&pool).await.unwrap();
crate::db::sql::query("CREATE TABLE comfy_cloud_assets (asset_id TEXT PRIMARY KEY, user_id INTEGER NOT NULL, job_id TEXT, asset_hash TEXT, name TEXT, tags_json TEXT NOT NULL DEFAULT '[]', snapshot_json TEXT, is_preview INTEGER NOT NULL DEFAULT 0, created_at TEXT DEFAULT (CURRENT_TIMESTAMP), updated_at TEXT)")
.execute(&pool).await.unwrap();
pool
}
#[tokio::test]
async fn records_assets_only_for_their_owner() {
let pool = pool().await;
let asset = json!({
"id": "asset-a",
"name": "first.png",
"asset_hash": "blake3:abc",
"tags": ["output"],
"preview_id": "preview-a"
});
record_asset(&pool, 7, Some("job-a"), &asset).await.unwrap();
assert!(asset_is_owned_by(&pool, 7, "asset-a").await.unwrap());
assert!(asset_is_owned_by(&pool, 7, "preview-a").await.unwrap());
assert!(!asset_is_owned_by(&pool, 8, "asset-a").await.unwrap());
assert!(asset_name_or_hash_is_owned_by(&pool, 7, "first.png")
.await
.unwrap());
assert_eq!(list_owned_assets(&pool, 7).await.unwrap(), vec![asset]);
}
#[test]
fn only_models_tag_marks_shared_assets() {
assert!(asset_is_shared_model(
&json!({ "tags": ["models", "checkpoints"] })
));
assert!(!asset_is_shared_model(&json!({ "tags": ["output"] })));
}
}
+18 -9
View File
@@ -1,10 +1,10 @@
use crate::error::{AppError, AppResult};
use sqlx::SqlitePool;
use sqlx::AnyPool;
/// 解析 Comfy Cloud 站点配置,返回 (base_url, api_key)。
/// 取 platform='comfy_cloud' 且 status='active' 的站点。
pub async fn resolve_comfy_cloud_site(db: &SqlitePool) -> AppResult<(String, String)> {
let row: Option<(String, Option<String>)> = sqlx::query_as(
pub async fn resolve_comfy_cloud_site(db: &AnyPool) -> AppResult<(String, String)> {
let row: Option<(String, Option<String>)> = crate::db::sql::query_as(
"SELECT url, api_key FROM sites
WHERE platform = 'comfy_cloud' AND status = 'active'
ORDER BY COALESCE(sort_order, 0), id LIMIT 1",
@@ -12,7 +12,9 @@ pub async fn resolve_comfy_cloud_site(db: &SqlitePool) -> AppResult<(String, Str
.fetch_optional(db)
.await?;
let (url, api_key) = row.ok_or_else(|| {
AppError::ServiceUnavailable("未配置 Comfy Cloud 站点(站点管理添加 comfy_cloud 平台)".into())
AppError::ServiceUnavailable(
"未配置 Comfy Cloud 站点(站点管理添加 comfy_cloud 平台)".into(),
)
})?;
let api_key = api_key
.map(|k| k.trim().to_string())
@@ -25,11 +27,18 @@ pub async fn resolve_comfy_cloud_site(db: &SqlitePool) -> AppResult<(String, Str
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
async fn pool() -> SqlitePool {
let pool = SqlitePoolOptions::new().max_connections(1).connect("sqlite::memory:").await.unwrap();
sqlx::query("CREATE TABLE sites (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT, url TEXT NOT NULL, platform TEXT NOT NULL, status TEXT NOT NULL, api_key TEXT, sort_order INTEGER)")
async fn pool() -> AnyPool {
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
crate::db::sql::query("CREATE TABLE sites (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT, url TEXT NOT NULL, platform TEXT NOT NULL, status TEXT NOT NULL, api_key TEXT, sort_order INTEGER)")
.execute(&pool).await.unwrap();
pool
}
@@ -37,7 +46,7 @@ mod tests {
#[tokio::test]
async fn resolves_active_comfy_cloud_site() {
let pool = pool().await;
sqlx::query("INSERT INTO sites (name,url,platform,status,api_key) VALUES ('cc','https://cloud.comfy.org/','comfy_cloud','active','key-123')")
crate::db::sql::query("INSERT INTO sites (name,url,platform,status,api_key) VALUES ('cc','https://cloud.comfy.org/','comfy_cloud','active','key-123')")
.execute(&pool).await.unwrap();
let (url, key) = resolve_comfy_cloud_site(&pool).await.unwrap();
assert_eq!(url, "https://cloud.comfy.org");
+5 -6
View File
@@ -6,11 +6,7 @@ use serde_json::{json, Value};
/// - status pending/in_progress → 返回空对象 {}(comfykit 据此继续轮询)
/// - status completed → { "{id}": { "status": {completed:true, status_str:"success"}, "outputs": {...} } }
/// - status failed/cancelled → { "{id}": { "status": {completed:false, status_str:"error"} } }
pub fn comfy_cloud_job_to_history(
prompt_id: &str,
status: &str,
outputs: Option<&Value>,
) -> Value {
pub fn comfy_cloud_job_to_history(prompt_id: &str, status: &str, outputs: Option<&Value>) -> Value {
match status {
"completed" => json!({
prompt_id: {
@@ -34,7 +30,10 @@ mod tests {
#[test]
fn in_progress_returns_empty() {
assert_eq!(comfy_cloud_job_to_history("p1", "in_progress", None), json!({}));
assert_eq!(
comfy_cloud_job_to_history("p1", "in_progress", None),
json!({})
);
assert_eq!(comfy_cloud_job_to_history("p1", "pending", None), json!({}));
}
File diff suppressed because it is too large Load Diff
+40 -24
View File
@@ -1,5 +1,5 @@
use crate::{error::AppResult, proxy::token_router::TokenRouter, services::billing::UsageTracker};
use sqlx::SqlitePool;
use sqlx::AnyPool;
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
@@ -11,7 +11,7 @@ use tokio::{
#[derive(Clone)]
pub struct CronScheduler {
db: Arc<SqlitePool>,
db: Arc<AnyPool>,
running: Arc<AtomicBool>,
}
@@ -21,7 +21,7 @@ pub struct CronSchedulerHandle {
}
impl CronScheduler {
pub fn new(db: Arc<SqlitePool>) -> Self {
pub fn new(db: Arc<AnyPool>) -> Self {
Self {
db,
running: Arc::new(AtomicBool::new(false)),
@@ -48,6 +48,11 @@ impl CronScheduler {
Duration::from_secs(120),
Self::recover_channels,
),
self.spawn_periodic(
"Comfy billing reconciliation",
Duration::from_secs(60),
Self::reconcile_comfy_billing,
),
];
CronSchedulerHandle {
@@ -58,7 +63,7 @@ impl CronScheduler {
fn spawn_periodic<F, Fut>(&self, name: &'static str, every: Duration, task: F) -> JoinHandle<()>
where
F: Fn(Arc<SqlitePool>) -> Fut + Send + Sync + Copy + 'static,
F: Fn(Arc<AnyPool>) -> Fut + Send + Sync + Copy + 'static,
Fut: std::future::Future<Output = AppResult<()>> + Send + 'static,
{
let db = self.db.clone();
@@ -77,15 +82,15 @@ impl CronScheduler {
})
}
async fn aggregate_usage(db: Arc<SqlitePool>) -> AppResult<()> {
async fn aggregate_usage(db: Arc<AnyPool>) -> AppResult<()> {
UsageTracker::create_daily_usage_snapshot_for_db(&db).await?;
let projected = UsageTracker::project_proxy_usage_for_db(&db).await?;
if projected > 0 {
tracing::debug!(projected, "usage aggregates projected");
}
sqlx::query(
crate::db::sql::query(
"UPDATE downstream_api_keys
SET updated_at = datetime('now')
SET updated_at = CURRENT_TIMESTAMP
WHERE used_requests > 0 OR used_cost > 0",
)
.execute(&*db)
@@ -93,10 +98,10 @@ impl CronScheduler {
Ok(())
}
async fn refresh_routes(db: Arc<SqlitePool>) -> AppResult<()> {
sqlx::query(
async fn refresh_routes(db: Arc<AnyPool>) -> AppResult<()> {
crate::db::sql::query(
"UPDATE token_routes
SET decision_refreshed_at = datetime('now'), updated_at = datetime('now')
SET decision_refreshed_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP
WHERE enabled = 1",
)
.execute(&*db)
@@ -104,9 +109,9 @@ impl CronScheduler {
Ok(())
}
async fn probe_models(db: Arc<SqlitePool>) -> AppResult<()> {
async fn probe_models(db: Arc<AnyPool>) -> AppResult<()> {
let router = TokenRouter::new(db.clone());
let models: Vec<String> = sqlx::query_scalar(
let models: Vec<String> = crate::db::sql::query_scalar(
"SELECT model_pattern FROM token_routes WHERE enabled = 1 ORDER BY id LIMIT 25",
)
.fetch_all(&*db)
@@ -119,17 +124,25 @@ impl CronScheduler {
Ok(())
}
async fn recover_channels(db: Arc<SqlitePool>) -> AppResult<()> {
sqlx::query(
async fn recover_channels(db: Arc<AnyPool>) -> AppResult<()> {
crate::db::sql::query(
"UPDATE route_channels
SET cooldown_until = NULL, cooldown_level = 0, consecutive_fail_count = 0
WHERE cooldown_until IS NOT NULL AND cooldown_until <= datetime('now')",
WHERE cooldown_until IS NOT NULL AND cooldown_until <= CURRENT_TIMESTAMP",
)
.execute(&*db)
.await?;
crate::proxy::token_router::router::invalidate_route_cache();
Ok(())
}
async fn reconcile_comfy_billing(db: Arc<AnyPool>) -> AppResult<()> {
let reconciled = crate::services::comfy_billing::reconcile_pending_holds(&db).await?;
if reconciled > 0 {
tracing::info!(reconciled, "Comfy billing holds reconciled");
}
Ok(())
}
}
impl CronSchedulerHandle {
@@ -153,17 +166,20 @@ impl Drop for CronSchedulerHandle {
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
#[tokio::test]
async fn scheduler_starts_and_stops_background_tasks() {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE downstream_api_keys (
id INTEGER PRIMARY KEY,
used_requests INTEGER DEFAULT 0,
@@ -174,7 +190,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE token_routes (
id INTEGER PRIMARY KEY,
model_pattern TEXT NOT NULL,
@@ -186,7 +202,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
crate::db::sql::query(
"CREATE TABLE route_channels (
id INTEGER PRIMARY KEY,
cooldown_until TEXT,
+409
View File
@@ -0,0 +1,409 @@
use std::{collections::HashMap, fmt, time::Duration};
use async_trait::async_trait;
use lettre::{
message::{header::ContentType, Attachment, Mailbox, MultiPart, SinglePart},
transport::smtp::authentication::Credentials,
Message, SmtpTransport, Transport,
};
use sqlx::AnyPool;
use crate::error::AppResult;
#[derive(Clone)]
pub struct SmtpMailConfig {
pub enabled: bool,
pub host: Option<String>,
pub port: u16,
pub secure: bool,
pub username: Option<String>,
pub password: Option<String>,
pub from: Option<String>,
}
impl fmt::Debug for SmtpMailConfig {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("SmtpMailConfig")
.field("enabled", &self.enabled)
.field("host", &self.host)
.field("port", &self.port)
.field("secure", &self.secure)
.field("username", &self.username)
.field("password", &self.password.as_ref().map(|_| "[REDACTED]"))
.field("from", &self.from)
.finish()
}
}
impl SmtpMailConfig {
pub fn is_configured(&self) -> bool {
self.port > 0
&& present(&self.host)
&& present(&self.username)
&& present(&self.password)
&& present(&self.from)
}
pub fn is_ready(&self) -> bool {
self.enabled && self.secure && self.port == 465 && self.is_configured()
}
pub(crate) fn from_settings(settings: &HashMap<String, String>) -> Self {
Self {
enabled: setting_bool(settings, "notify_smtp_enabled", true),
host: setting(settings, "notify_smtp_host"),
port: setting(settings, "notify_smtp_port")
.and_then(|value| value.parse().ok())
.unwrap_or(465),
secure: setting(settings, "notify_smtp_secure")
.map(|value| parse_bool(&value))
.unwrap_or(true),
username: setting(settings, "notify_smtp_user"),
password: setting(settings, "notify_smtp_pass"),
from: setting(settings, "notify_smtp_from"),
}
}
}
#[derive(Clone, PartialEq, Eq)]
pub struct InlineImage {
pub content_id: String,
pub content_type: String,
pub data: Vec<u8>,
}
impl fmt::Debug for InlineImage {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("InlineImage")
.field("content_id", &self.content_id)
.field("content_type", &self.content_type)
.field("data", &"[REDACTED]")
.finish()
}
}
#[derive(Clone, PartialEq, Eq)]
pub struct MailMessage {
pub to: String,
pub subject: String,
pub text: String,
pub html: Option<String>,
pub inline_images: Vec<InlineImage>,
}
impl fmt::Debug for MailMessage {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("MailMessage")
.field("to", &self.to)
.field("subject", &self.subject)
.field("text", &"[REDACTED]")
.field("html", &self.html.as_ref().map(|_| "[REDACTED]"))
.field("inline_images", &self.inline_images)
.finish()
}
}
#[async_trait]
pub trait MailSender: Send + Sync {
async fn send(&self, message: MailMessage) -> anyhow::Result<()>;
}
#[derive(Clone)]
pub struct SmtpMailSender {
config: SmtpMailConfig,
timeout: Duration,
}
impl SmtpMailSender {
pub fn new(config: SmtpMailConfig) -> Self {
Self::with_timeout(config, Duration::from_secs(30))
}
pub fn with_timeout(config: SmtpMailConfig, timeout: Duration) -> Self {
Self { config, timeout }
}
}
#[async_trait]
impl MailSender for SmtpMailSender {
async fn send(&self, message: MailMessage) -> anyhow::Result<()> {
if !self.config.enabled {
anyhow::bail!("smtp is disabled");
}
if !self.config.is_configured() {
anyhow::bail!("smtp is not configured");
}
let email = build_message(&self.config, message)?;
let config = self.config.clone();
let timeout = self.timeout;
let worker = tokio::task::spawn_blocking(move || send_blocking(config, timeout, email));
tokio::time::timeout(timeout, worker)
.await
.map_err(|_| anyhow::anyhow!("smtp send timed out"))?
.map_err(|error| anyhow::anyhow!("smtp worker failed: {error}"))??;
Ok(())
}
}
pub async fn load_smtp_mail_config(db: &AnyPool) -> AppResult<SmtpMailConfig> {
let rows: Vec<(String, Option<String>)> = crate::db::sql::query_as(
"SELECT key, value FROM settings WHERE key IN (
'notify_smtp_enabled',
'notify_smtp_host',
'notify_smtp_port',
'notify_smtp_secure',
'notify_smtp_user',
'notify_smtp_pass',
'notify_smtp_from'
)",
)
.fetch_all(db)
.await?;
let settings = rows
.into_iter()
.filter_map(|(key, value)| value.map(|value| (key, value)))
.collect();
Ok(SmtpMailConfig::from_settings(&settings))
}
fn build_message(config: &SmtpMailConfig, message: MailMessage) -> anyhow::Result<Message> {
let from = required(config.from.as_deref(), "smtp from")?;
let to = required(Some(message.to.trim()), "mail recipient")?;
if to.is_empty() {
anyhow::bail!("mail recipient is empty");
}
let builder = Message::builder()
.from(from.parse::<Mailbox>()?)
.to(to.parse::<Mailbox>()?)
.subject(message.subject);
match (message.html, message.inline_images.is_empty()) {
(Some(html), false) => {
let mut related = MultiPart::related().singlepart(SinglePart::html(html));
for image in message.inline_images {
let content_type = ContentType::parse(&image.content_type)?;
related = related.singlepart(
Attachment::new_inline(image.content_id).body(image.data, content_type),
);
}
Ok(builder.multipart(
MultiPart::alternative()
.singlepart(SinglePart::plain(message.text))
.multipart(related),
)?)
}
(Some(html), true) => {
Ok(builder.multipart(MultiPart::alternative_plain_html(message.text, html))?)
}
(None, _) => Ok(builder.body(message.text)?),
}
}
fn send_blocking(config: SmtpMailConfig, timeout: Duration, email: Message) -> anyhow::Result<()> {
let host = required(config.host.as_deref(), "smtp host")?;
let username = required(config.username.as_deref(), "smtp user")?;
let password = required(config.password.as_deref(), "smtp pass")?;
let builder = if config.secure {
// `relay` uses implicit TLS, including Feishu's SMTPS endpoint on port 465.
SmtpTransport::relay(host)?
} else {
SmtpTransport::builder_dangerous(host)
};
let mailer = builder
.port(config.port)
.credentials(Credentials::new(username.to_string(), password.to_string()))
.timeout(Some(timeout))
.build();
mailer.send(&email)?;
Ok(())
}
fn present(value: &Option<String>) -> bool {
value
.as_deref()
.is_some_and(|value| !value.trim().is_empty())
}
fn setting(settings: &HashMap<String, String>, key: &str) -> Option<String> {
settings
.get(key)
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
fn setting_bool(settings: &HashMap<String, String>, key: &str, default: bool) -> bool {
setting(settings, key)
.map(|value| parse_bool(&value))
.unwrap_or(default)
}
fn parse_bool(value: &str) -> bool {
matches!(
value.trim().to_ascii_lowercase().as_str(),
"1" | "true" | "yes" | "on" | "enabled"
)
}
fn required<'a>(value: Option<&'a str>, name: &str) -> anyhow::Result<&'a str> {
value.ok_or_else(|| anyhow::anyhow!("{name} is not configured"))
}
#[cfg(test)]
mod tests {
use super::*;
use parking_lot::Mutex;
use sqlx::any::{install_default_drivers, AnyPoolOptions};
#[derive(Default)]
struct RecordingSender {
messages: Mutex<Vec<MailMessage>>,
}
#[async_trait]
impl MailSender for RecordingSender {
async fn send(&self, message: MailMessage) -> anyhow::Result<()> {
self.messages.lock().push(message);
Ok(())
}
}
fn configured() -> SmtpMailConfig {
SmtpMailConfig {
enabled: true,
host: Some("smtp.feishu.cn".into()),
port: 465,
secure: true,
username: Some("no-reply@example.com".into()),
password: Some("secret".into()),
from: Some("BoosAPI <no-reply@example.com>".into()),
}
}
#[test]
fn readiness_requires_enabled_encrypted_complete_config() {
let config = configured();
assert!(config.is_configured());
assert!(config.is_ready());
let mut plaintext = config.clone();
plaintext.secure = false;
assert!(plaintext.is_configured());
assert!(!plaintext.is_ready());
let mut starttls = config.clone();
starttls.port = 587;
assert!(starttls.is_configured());
assert!(!starttls.is_ready());
let mut missing_password = config;
missing_password.password = None;
assert!(!missing_password.is_configured());
assert!(!missing_password.is_ready());
}
#[test]
fn builds_multipart_message_for_dynamic_recipient() {
let message = build_message(
&configured(),
MailMessage {
to: "new-user@example.net".into(),
subject: "邮箱验证码".into(),
text: "验证码:123456".into(),
html: Some("<p>验证码:<strong>123456</strong></p>".into()),
inline_images: Vec::new(),
},
)
.unwrap();
let formatted = String::from_utf8(message.formatted()).unwrap();
assert!(formatted.contains("To: new-user@example.net"));
assert!(formatted.contains("Content-Type: multipart/alternative"));
assert!(formatted.contains("Content-Type: text/plain"));
assert!(formatted.contains("Content-Type: text/html"));
}
#[tokio::test]
async fn mail_sender_trait_supports_injected_test_double() {
let sender = RecordingSender::default();
let message = MailMessage {
to: "dynamic@example.net".into(),
subject: "Subject".into(),
text: "sensitive body".into(),
html: Some("<p>sensitive body</p>".into()),
inline_images: Vec::new(),
};
sender.send(message.clone()).await.unwrap();
assert_eq!(sender.messages.lock().as_slice(), &[message]);
}
#[test]
fn builds_related_html_with_inline_image() {
let message = build_message(
&configured(),
MailMessage {
to: "new-user@example.net".into(),
subject: "品牌验证码".into(),
text: "验证码:123456".into(),
html: Some("<img src=\"cid:brand-logo\" alt=\"BossAI\">".into()),
inline_images: vec![InlineImage {
content_id: "brand-logo".into(),
content_type: "image/png".into(),
data: vec![137, 80, 78, 71],
}],
},
)
.unwrap();
let formatted = String::from_utf8(message.formatted()).unwrap();
assert!(formatted.contains("Content-Type: multipart/alternative"));
assert!(formatted.contains("Content-Type: multipart/related"));
assert!(formatted.contains("Content-Type: image/png"));
assert!(formatted.contains("Content-ID: <brand-logo>"));
assert!(formatted.contains("cid:brand-logo"));
}
#[tokio::test]
async fn loads_feishu_smtp_settings_without_exposing_password_in_debug() {
let pool = {
install_default_drivers();
AnyPoolOptions::new()
}
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
crate::db::sql::query("CREATE TABLE settings (key TEXT PRIMARY KEY NOT NULL, value TEXT)")
.execute(&pool)
.await
.unwrap();
for (key, value) in [
("notify_smtp_enabled", "true"),
("notify_smtp_host", "smtp.feishu.cn"),
("notify_smtp_port", "465"),
("notify_smtp_secure", "true"),
("notify_smtp_user", "no-reply@example.com"),
("notify_smtp_pass", "do-not-log-this"),
("notify_smtp_from", "no-reply@example.com"),
] {
crate::db::sql::query("INSERT INTO settings (key, value) VALUES (?, ?)")
.bind(key)
.bind(value)
.execute(&pool)
.await
.unwrap();
}
let config = load_smtp_mail_config(&pool).await.unwrap();
assert!(config.is_ready());
assert_eq!(config.host.as_deref(), Some("smtp.feishu.cn"));
assert_eq!(config.port, 465);
let debug = format!("{config:?}");
assert!(debug.contains("[REDACTED]"));
assert!(!debug.contains("do-not-log-this"));
}
}
+224 -29
View File
@@ -18,6 +18,15 @@ use types::{
};
use volc_visual_sdk::VisualClient;
/// Supported aspect ratios for Jimeng image generation. Keep these as strings
/// because the upstream API accepts symbolic ratios rather than floats.
pub const ALLOWED_IMAGE_RATIOS: &[&str] =
&["1:1", "4:3", "3:4", "16:9", "9:16", "3:2", "2:3", "21:9"];
const MIN_IMAGE_AREA: i64 = 1024 * 1024;
const MAX_IMAGE_AREA: i64 = 4096 * 4096;
const DEFAULT_MIN_IMAGE_RATIO: f64 = 1.0 / 3.0;
const DEFAULT_MAX_IMAGE_RATIO: f64 = 3.0;
/// Sentinel task id used by [`JimengProvider::verify`]. It is intentionally not
/// a real Volcengine task id: the gateway validates the V4 signature *before*
/// the business layer ever looks up the task, so a valid AK/SK is accepted
@@ -68,6 +77,50 @@ fn is_auth_error_code(code: &str) -> bool {
.any(|needle| lc.contains(needle))
}
fn dimensions_for_image_ratio(ratio: &str) -> Option<(i64, i64)> {
match ratio {
"1:1" => Some((2048, 2048)),
"4:3" => Some((2304, 1728)),
"3:4" => Some((1728, 2304)),
"16:9" => Some((2560, 1440)),
"9:16" => Some((1440, 2560)),
"3:2" => Some((2496, 1664)),
"2:3" => Some((1664, 2496)),
"21:9" => Some((3024, 1296)),
_ => None,
}
}
fn validate_image_dimensions(
width: i64,
height: i64,
min_ratio: Option<f64>,
max_ratio: Option<f64>,
) -> Result<(), JimengError> {
if width <= 0 || height <= 0 {
return Err(JimengError::Invalid(
"width and height must be greater than 0".into(),
));
}
let area = width
.checked_mul(height)
.ok_or_else(|| JimengError::Invalid("width * height is too large".into()))?;
if !(MIN_IMAGE_AREA..=MAX_IMAGE_AREA).contains(&area) {
return Err(JimengError::Invalid(format!(
"width * height must be between {MIN_IMAGE_AREA} and {MAX_IMAGE_AREA} pixels"
)));
}
let ratio = width as f64 / height as f64;
let min_ratio = min_ratio.unwrap_or(DEFAULT_MIN_IMAGE_RATIO);
let max_ratio = max_ratio.unwrap_or(DEFAULT_MAX_IMAGE_RATIO);
if ratio < min_ratio || ratio > max_ratio {
return Err(JimengError::Invalid(format!(
"width / height must be between {min_ratio} and {max_ratio}"
)));
}
Ok(())
}
/// Classify the raw response of a signature-probe call into a
/// [`JimengVerifyOutcome`]. Pure (no IO) so it is unit-testable.
///
@@ -302,7 +355,10 @@ impl JimengProvider {
.header("Content-Type", prepared.signed.content_type.as_str())
.header("Host", prepared.signed.host.as_str())
.header("X-Date", prepared.signed.x_date.as_str())
.header("X-Content-Sha256", prepared.signed.x_content_sha256.as_str())
.header(
"X-Content-Sha256",
prepared.signed.x_content_sha256.as_str(),
)
.header("Authorization", prepared.signed.authorization.as_str());
if let Some(token) = &prepared.signed.x_security_token {
req = req.header("X-Security-Token", token.as_str());
@@ -354,16 +410,42 @@ pub fn build_image_body(req: &ImageRequest) -> Result<Value, JimengError> {
"image_urls supports at most 10 images".into(),
));
}
if req.size.is_some_and(|size| size <= 0) {
return Err(JimengError::Invalid("size must be greater than 0".into()));
if req
.size
.is_some_and(|size| !(MIN_IMAGE_AREA..=MAX_IMAGE_AREA).contains(&size))
{
return Err(JimengError::Invalid(format!(
"size must be between {MIN_IMAGE_AREA} and {MAX_IMAGE_AREA} pixels"
)));
}
for (name, ratio) in [("min_ratio", req.min_ratio), ("max_ratio", req.max_ratio)] {
if ratio.is_some_and(|value| value <= 0.0) {
let ratio = req
.ratio
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(ratio) = ratio {
if !ALLOWED_IMAGE_RATIOS.contains(&ratio) {
return Err(JimengError::Invalid(format!(
"{name} must be greater than 0"
"ratio must be one of {ALLOWED_IMAGE_RATIOS:?} (got {ratio})"
)));
}
}
for (name, ratio) in [("min_ratio", req.min_ratio), ("max_ratio", req.max_ratio)] {
if ratio.is_some_and(|value| !value.is_finite() || !(1.0 / 16.0..16.0).contains(&value)) {
return Err(JimengError::Invalid(format!(
"{name} must be between 1/16 (inclusive) and 16 (exclusive)"
)));
}
}
if req
.min_ratio
.zip(req.max_ratio)
.is_some_and(|(min_ratio, max_ratio)| min_ratio > max_ratio)
{
return Err(JimengError::Invalid(
"min_ratio must not exceed max_ratio".into(),
));
}
let req_key = types::normalize_req_key(JimengCapability::Image, req.model.as_deref());
let mut body = json!({
"req_key": req_key,
@@ -373,10 +455,25 @@ pub fn build_image_body(req: &ImageRequest) -> Result<Value, JimengError> {
if !image_urls.is_empty() {
obj.insert("image_urls".into(), json!(image_urls));
}
if let Some(width) = req.width {
// `ratio` is a local UI helper, not an upstream API field. A known ratio
// always wins over stale client-provided dimensions so old clients cannot
// submit an area below Volcengine's 1K minimum.
let dimensions = if let Some(mapped) = ratio.and_then(dimensions_for_image_ratio) {
Some(mapped)
} else {
match (req.width, req.height) {
(Some(width), Some(height)) => Some((width, height)),
(None, None) => None,
_ => {
return Err(JimengError::Invalid(
"width and height must be provided together".into(),
));
}
}
};
if let Some((width, height)) = dimensions {
validate_image_dimensions(width, height, req.min_ratio, req.max_ratio)?;
obj.insert("width".into(), json!(width));
}
if let Some(height) = req.height {
obj.insert("height".into(), json!(height));
}
if let Some(size) = req.size {
@@ -403,7 +500,11 @@ pub fn build_image_body(req: &ImageRequest) -> Result<Value, JimengError> {
/// `image_urls` entry; `frames` must be one of [`ALLOWED_FRAMES`] and
/// `aspect_ratio`, when present, must be in [`ALLOWED_ASPECT_RATIOS`].
pub fn build_video_body(req: &VideoRequest) -> Result<Value, JimengError> {
let has_prompt = req.prompt.as_deref().map(str::trim).is_some_and(|p| !p.is_empty());
let has_prompt = req
.prompt
.as_deref()
.map(str::trim)
.is_some_and(|p| !p.is_empty());
let has_image = !req.image_urls.is_empty();
if !has_prompt && !has_image {
return Err(JimengError::Invalid(
@@ -455,6 +556,7 @@ mod tests {
model: None,
prompt: " ".into(),
image_urls: vec![],
ratio: None,
width: None,
height: None,
size: None,
@@ -463,7 +565,10 @@ mod tests {
max_ratio: None,
force_single: None,
};
assert!(matches!(build_image_body(&req), Err(JimengError::Invalid(_))));
assert!(matches!(
build_image_body(&req),
Err(JimengError::Invalid(_))
));
}
#[test]
@@ -472,9 +577,10 @@ mod tests {
model: None,
prompt: "a cat".into(),
image_urls: vec!["https://x/ref.png".into()],
width: Some(2048),
height: Some(2048),
size: Some(4096),
ratio: Some("9:16".into()),
width: Some(720),
height: Some(1280),
size: Some(2048 * 2048),
scale: Some(0.5),
min_ratio: Some(0.33),
max_ratio: Some(3.0),
@@ -484,9 +590,10 @@ mod tests {
assert_eq!(body["req_key"], "jimeng_t2i_v40");
assert_eq!(body["prompt"], "a cat");
assert_eq!(body["image_urls"][0], "https://x/ref.png");
assert_eq!(body["width"], 2048);
assert_eq!(body["height"], 2048);
assert_eq!(body["size"], 4096);
assert!(body.get("ratio").is_none());
assert_eq!(body["width"], 1440);
assert_eq!(body["height"], 2560);
assert_eq!(body["size"], 2048 * 2048);
assert_eq!(body["scale"], 0.5);
assert_eq!(body["min_ratio"], 0.33);
assert_eq!(body["max_ratio"], 3.0);
@@ -499,6 +606,7 @@ mod tests {
model: None,
prompt: "a cat".into(),
image_urls: vec![],
ratio: None,
width: None,
height: None,
size: None,
@@ -511,16 +619,18 @@ mod tests {
let obj = body.as_object().unwrap();
assert!(!obj.contains_key("width"));
assert!(!obj.contains_key("image_urls"));
assert!(!obj.contains_key("ratio"));
assert!(!obj.contains_key("size"));
assert!(!obj.contains_key("scale"));
}
#[test]
fn image_body_rejects_too_many_reference_images() {
fn image_body_maps_ratio_to_dimensions_when_width_height_absent() {
let req = ImageRequest {
model: None,
prompt: "a cat".into(),
image_urls: (0..11).map(|idx| format!("https://x/{idx}.png")).collect(),
prompt: "a vertical poster".into(),
image_urls: vec![],
ratio: Some("9:16".into()),
width: None,
height: None,
size: None,
@@ -529,7 +639,82 @@ mod tests {
max_ratio: None,
force_single: None,
};
assert!(matches!(build_image_body(&req), Err(JimengError::Invalid(_))));
let body = build_image_body(&req).unwrap();
assert!(body.get("ratio").is_none());
assert_eq!(body["width"], 1440);
assert_eq!(body["height"], 2560);
}
#[test]
fn image_body_rejects_incomplete_or_too_small_dimensions() {
let base = ImageRequest {
model: None,
prompt: "a vertical poster".into(),
image_urls: vec![],
ratio: None,
width: Some(720),
height: None,
size: None,
scale: None,
min_ratio: None,
max_ratio: None,
force_single: None,
};
assert!(matches!(
build_image_body(&base),
Err(JimengError::Invalid(_))
));
let too_small = ImageRequest {
height: Some(1280),
..base
};
assert!(matches!(
build_image_body(&too_small),
Err(JimengError::Invalid(_))
));
}
#[test]
fn image_body_rejects_unknown_ratio() {
let req = ImageRequest {
model: None,
prompt: "a cat".into(),
image_urls: vec![],
ratio: Some("7:5".into()),
width: None,
height: None,
size: None,
scale: None,
min_ratio: None,
max_ratio: None,
force_single: None,
};
assert!(matches!(
build_image_body(&req),
Err(JimengError::Invalid(_))
));
}
#[test]
fn image_body_rejects_too_many_reference_images() {
let req = ImageRequest {
model: None,
prompt: "a cat".into(),
image_urls: (0..11).map(|idx| format!("https://x/{idx}.png")).collect(),
ratio: None,
width: None,
height: None,
size: None,
scale: None,
min_ratio: None,
max_ratio: None,
force_single: None,
};
assert!(matches!(
build_image_body(&req),
Err(JimengError::Invalid(_))
));
}
#[test]
@@ -542,7 +727,10 @@ mod tests {
aspect_ratio: None,
seed: None,
};
assert!(matches!(build_video_body(&req), Err(JimengError::Invalid(_))));
assert!(matches!(
build_video_body(&req),
Err(JimengError::Invalid(_))
));
}
#[test]
@@ -590,7 +778,10 @@ mod tests {
aspect_ratio: None,
seed: None,
};
assert!(matches!(build_video_body(&req), Err(JimengError::Invalid(_))));
assert!(matches!(
build_video_body(&req),
Err(JimengError::Invalid(_))
));
}
#[test]
@@ -603,7 +794,10 @@ mod tests {
aspect_ratio: Some("21:9".into()),
seed: None,
};
assert!(matches!(build_video_body(&req), Err(JimengError::Invalid(_))));
assert!(matches!(
build_video_body(&req),
Err(JimengError::Invalid(_))
));
}
#[test]
@@ -629,8 +823,8 @@ mod tests {
fn video_body_accepts_all_allowed_aspect_ratios() {
for ratio in ALLOWED_ASPECT_RATIOS {
let req = VideoRequest {
model: None,
prompt: Some("x".into()),
model: None,
prompt: Some("x".into()),
image_urls: vec![],
frames: 121,
aspect_ratio: Some(ratio.to_string()),
@@ -671,8 +865,8 @@ mod tests {
fn video_body_accepts_both_frame_counts() {
for frames in ALLOWED_FRAMES {
let req = VideoRequest {
model: None,
prompt: Some("x".into()),
model: None,
prompt: Some("x".into()),
image_urls: vec![],
frames,
aspect_ratio: None,
@@ -754,7 +948,8 @@ mod tests {
#[test]
fn verify_classifies_invalid_access_key_as_auth_failed() {
let body = r#"{"ResponseMetadata":{"Error":{"Code":"InvalidAccessKeyId","Message":"bad ak"}}}"#;
let body =
r#"{"ResponseMetadata":{"Error":{"Code":"InvalidAccessKeyId","Message":"bad ak"}}}"#;
assert!(matches!(
classify_verify_response(403, body),
JimengVerifyOutcome::AuthFailed { .. }
+24 -9
View File
@@ -61,7 +61,10 @@ pub fn decide(capability: JimengCapability, resp: &GetResultResponse) -> PollDec
}
/// Extract the finished output from a `done` response, if present.
fn extract_output(capability: JimengCapability, data: Option<&GetResultData>) -> Option<JimengOutput> {
fn extract_output(
capability: JimengCapability,
data: Option<&GetResultData>,
) -> Option<JimengOutput> {
let data = data?;
match capability {
JimengCapability::Image => {
@@ -217,7 +220,10 @@ mod tests {
PollDecision::Fail(JimengError::Expired)
));
assert!(matches!(
decide(JimengCapability::Image, &resp(10000, Some("not_found"), None)),
decide(
JimengCapability::Image,
&resp(10000, Some("not_found"), None)
),
PollDecision::Fail(JimengError::NotFound)
));
}
@@ -382,7 +388,10 @@ mod tests {
#[test]
fn decide_not_found_is_terminal() {
assert!(matches!(
decide(JimengCapability::Video, &resp(10000, Some("not_found"), None)),
decide(
JimengCapability::Video,
&resp(10000, Some("not_found"), None)
),
PollDecision::Fail(JimengError::NotFound)
));
}
@@ -404,7 +413,10 @@ mod tests {
|| async {
let mut i = idx.borrow_mut();
let s = seq.borrow();
let item = s.get(*i).cloned().unwrap_or_else(|| s.last().cloned().unwrap());
let item = s
.get(*i)
.cloned()
.unwrap_or_else(|| s.last().cloned().unwrap());
*i += 1;
Ok(item)
},
@@ -457,14 +469,17 @@ mod tests {
// A transient code may be followed by a fatal one; the fatal must end
// the loop with its Api error rather than retrying forever.
let seq = vec![
resp(50429, None, None), // retryable QPS
resp(50411, None, None), // fatal input moderation
resp(50429, None, None), // retryable QPS
resp(50411, None, None), // fatal input moderation
];
let (out, _sleeps) =
run_sequence_capturing_sleeps(JimengCapability::Image, seq).await;
let (out, _sleeps) = run_sequence_capturing_sleeps(JimengCapability::Image, seq).await;
assert!(matches!(
out,
Err(JimengError::Api { code: 50411, retryable: false, .. })
Err(JimengError::Api {
code: 50411,
retryable: false,
..
})
));
}
}
+26 -6
View File
@@ -230,6 +230,8 @@ pub struct ImageRequest {
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub image_urls: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ratio: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub width: Option<i64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub height: Option<i64>,
@@ -378,7 +380,10 @@ mod tests {
assert_eq!(parsed.code, CODE_SUCCESS);
assert_eq!(parsed.task_status(), TaskStatus::Done);
let data = parsed.data.unwrap();
assert_eq!(data.image_urls.unwrap(), vec!["https://x/a.png".to_string()]);
assert_eq!(
data.image_urls.unwrap(),
vec!["https://x/a.png".to_string()]
);
}
#[test]
@@ -407,8 +412,14 @@ mod tests {
// gate them by capability — they are retryable for both. Pin that
// actual behavior so a future capability gate is a conscious change.
for code in [50516, 50517] {
assert_eq!(classify(code, JimengCapability::Video), CodeClass::Retryable);
assert_eq!(classify(code, JimengCapability::Image), CodeClass::Retryable);
assert_eq!(
classify(code, JimengCapability::Video),
CodeClass::Retryable
);
assert_eq!(
classify(code, JimengCapability::Image),
CodeClass::Retryable
);
}
}
@@ -417,7 +428,10 @@ mod tests {
// Input copyright pre-moderation (50518) is fatal; output copyright
// post-moderation (50519) is retryable — adjacent codes, opposite class.
assert_eq!(classify(50518, JimengCapability::Image), CodeClass::Fatal);
assert_eq!(classify(50519, JimengCapability::Image), CodeClass::Retryable);
assert_eq!(
classify(50519, JimengCapability::Image),
CodeClass::Retryable
);
}
#[test]
@@ -436,7 +450,10 @@ mod tests {
// assertion: 50500/50501 are Fatal for image, Retryable for video.
for code in [50500, 50501] {
assert_eq!(classify(code, JimengCapability::Image), CodeClass::Fatal);
assert_eq!(classify(code, JimengCapability::Video), CodeClass::Retryable);
assert_eq!(
classify(code, JimengCapability::Video),
CodeClass::Retryable
);
}
}
@@ -456,7 +473,10 @@ mod tests {
assert_eq!(parsed.task_status(), TaskStatus::Done);
let data = parsed.data.unwrap();
assert!(data.image_urls.is_none());
assert_eq!(data.binary_data_base64.unwrap(), vec!["aGVsbG8=".to_string()]);
assert_eq!(
data.binary_data_base64.unwrap(),
vec!["aGVsbG8=".to_string()]
);
}
#[test]
+19 -4
View File
@@ -1,19 +1,34 @@
pub mod announcement;
pub mod account_health;
pub mod agent_attachments;
pub mod agent_workbench;
pub mod announcement;
pub mod billing;
pub mod builtin_keys;
pub mod comfy_billing;
pub mod comfy_cloud;
pub mod comfyui_process;
pub mod commerce;
pub mod cron;
pub mod email;
pub mod jimeng;
pub mod notification;
pub mod oauth;
pub mod payment;
pub mod payment_jobs;
pub mod payment_providers;
pub mod payment_refunds;
pub mod pgrust;
pub mod pixelle_connection;
pub mod profile_avatar;
pub mod redis;
pub mod register_captcha;
pub mod registration_email;
pub mod registration_limiter;
pub mod route_rebuild;
pub mod runtime_auth;
pub mod runtime_logs;
pub mod seedance;
pub mod shared;
pub mod user_oauth;
pub mod user_oauth_flow;
pub mod workspace;
pub mod comfy_cloud;
pub mod pixelle_connection;
pub mod jimeng;

Some files were not shown because too many files have changed in this diff Show More