aibo/BoosAPI
1
0
Fork 0
forked from boos/BoosAPI
Code Pull Requests Activity

feat(rust): wire downstream-key billing, add /v1/responses, ComfyUI shim

Browse Source 
Previously every executor took `_state: AppState` (unused) and `check_auth`
discarded the matched downstream key id, so `usage_logs` was orphan code.
This change plumbs the key id from auth → handlers → executors → billing,
adds the OpenAI Responses API endpoint, and exposes shim aliases under
`/api/comfyui-agent/proxy/openai/*` so embedded ComfyUI's OpenAI nodes
flow through the same auth/conductor/billing pipeline.

- check_auth now returns Option<i64> (managed key id) instead of unit
- record_billing helper writes usage_logs after each successful upstream call
- extract_usage handles both chat-completions and responses usage shapes
- conductor wrappers record billing post-Success; stream variant scans
  chunks in reverse for the final usage event
- POST /v1/responses (stream + non-stream) via ProxyConductor
- shim aliases mounted on the same Router so auth/failover/billing work for free
- ComfyUIProcess injects BOOSAPI_BUILTIN_KEY env + --comfy-api-base on spawn
- bootstrap built-in `comfyui-agent` downstream key at startup (idempotent)
- get_auth_header in vendored ComfyUI falls back to BOOSAPI_BUILTIN_KEY env

Co-Authored-By: Kiro <noreply@kiro.dev>
This commit is contained in:
boos
2026-05-20 06:20:39 +08:00
parent 22751effd5
commit 9c0d4ec69f
7 changed files with 644 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/services/builtin_keys.rs
+260
View File
@@ -0,0 +1,260 @@
//! Built-in downstream API keys.
//!
//! Some flows inside BoosAPI itself (e.g. ComfyUI Agent) call this server's
//! own `/v1/*` endpoints. To make those calls go through the same accounting
//! pipeline as a normal downstream consumer, we provision a system-level row
//! in `downstream_api_keys` (`name='comfyui-agent' AND user_id IS NULL`) and
//! mirror the plaintext secret into `settings` under
//! `builtin_key.comfyui_agent.secret`.
//!
//! `ensure_comfyui_agent_key` is the single entry point that boots this state
//! at startup and is intentionally idempotent: any combination of
//! "settings has it / row has it / both / neither" is converged to a
//! consistent pair without rotating the secret.
use crate::error::{AppError, AppResult};
use sqlx::SqlitePool;
use uuid::Uuid;
/// In-memory snapshot of system-owned downstream keys.
#[derive(Debug, Clone)]
pub struct BuiltinKeys {
/// `(downstream_api_keys.id, plaintext key)` for the ComfyUI Agent.
pub comfyui_agent: (i64, String),
}
const COMFYUI_AGENT_NAME: &str = "comfyui-agent";
const COMFYUI_AGENT_DESCRIPTION: &str = "Built-in key for ComfyUI Agent";
const COMFYUI_AGENT_SETTINGS_KEY: &str = "builtin_key.comfyui_agent.secret";
fn generate_comfyui_agent_secret() -> String {
let raw = Uuid::new_v4().simple().to_string();
format!("sk-comfyui-{}", raw)
}
/// Make sure the ComfyUI Agent built-in key exists in both `downstream_api_keys`
/// and `settings`. Returns the row id and plaintext secret.
///
/// This function is idempotent. The full state transition table:
///
/// | settings entry | downstream row | action |
/// |---------------|----------------|------------------------------------------|
/// | absent | absent | generate new secret, insert both |
/// | present | absent | reuse settings secret, insert row |
/// | absent | present | reuse row secret, write settings entry |
/// | 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)> {
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 = ?")
.bind(COMFYUI_AGENT_SETTINGS_KEY)
.fetch_optional(&mut *tx)
.await
.map_err(AppError::Database)?;
let existing_row: Option<(i64, String, i64)> = sqlx::query_as(
"SELECT id, key, enabled FROM downstream_api_keys \
WHERE name = ? AND user_id IS NULL \
LIMIT 1",
)
.bind(COMFYUI_AGENT_NAME)
.fetch_optional(&mut *tx)
.await
.map_err(AppError::Database)?;
// 2. Decide the canonical (id, secret) pair.
let (id, secret) = match (existing_row, settings_secret.as_deref()) {
// Both sides agree on a row. Make sure it's enabled and settings stays in sync.
(Some((id, row_key, enabled)), settings_value) => {
if enabled == 0 {
tracing::warn!(
target: "boosapi::builtin_keys",
key_id = id,
"comfyui-agent builtin key was disabled; re-enabling",
);
sqlx::query("UPDATE downstream_api_keys SET enabled = 1 WHERE id = ?")
.bind(id)
.execute(&mut *tx)
.await
.map_err(AppError::Database)?;
}
// settings absent or out-of-sync: write/overwrite with the row's secret.
// (Row wins because it's the source of truth used by the auth layer.)
if settings_value != Some(row_key.as_str()) {
upsert_settings(&mut tx, COMFYUI_AGENT_SETTINGS_KEY, &row_key).await?;
}
(id, row_key)
}
// Settings has a previous secret but the row got nuked — restore the row.
(None, Some(saved)) if !saved.trim().is_empty() => {
let id = insert_comfyui_agent_row(&mut tx, saved).await?;
(id, saved.to_string())
}
// Cold start (or settings has an empty/whitespace value): generate new.
(None, _) => {
let secret = generate_comfyui_agent_secret();
let id = insert_comfyui_agent_row(&mut tx, &secret).await?;
upsert_settings(&mut tx, COMFYUI_AGENT_SETTINGS_KEY, &secret).await?;
(id, secret)
}
};
tx.commit().await.map_err(AppError::Database)?;
Ok((id, secret))
}
async fn insert_comfyui_agent_row(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
secret: &str,
) -> AppResult<i64> {
let result = sqlx::query(
"INSERT INTO downstream_api_keys (name, key, user_id, enabled, description) \
VALUES (?, ?, NULL, 1, ?)",
)
.bind(COMFYUI_AGENT_NAME)
.bind(secret)
.bind(COMFYUI_AGENT_DESCRIPTION)
.execute(&mut **tx)
.await
.map_err(AppError::Database)?;
Ok(result.last_insert_rowid())
}
async fn upsert_settings(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
key: &str,
value: &str,
) -> AppResult<()> {
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)")
.bind(key)
.bind(value)
.execute(&mut **tx)
.await
.map_err(AppError::Database)?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
async fn fresh_pool() -> SqlitePool {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("memory sqlite");
sqlx::migrate!("./migrations")
.run(&pool)
.await
.expect("run migrations");
pool
}
async fn read_settings(db: &SqlitePool, key: &str) -> Option<String> {
sqlx::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(
"SELECT id, key, enabled FROM downstream_api_keys \
WHERE name = 'comfyui-agent' AND user_id IS NULL",
)
.fetch_optional(db)
.await
.unwrap()
}
#[tokio::test]
async fn cold_start_creates_both_sides() {
let pool = fresh_pool().await;
let (id, secret) = ensure_comfyui_agent_key(&pool).await.unwrap();
assert!(secret.starts_with("sk-comfyui-"));
assert_eq!(secret.len(), "sk-comfyui-".len() + 32);
let row = read_row(&pool).await.expect("row inserted");
assert_eq!(row.0, id);
assert_eq!(row.1, secret);
assert_eq!(row.2, 1);
assert_eq!(
read_settings(&pool, COMFYUI_AGENT_SETTINGS_KEY).await,
Some(secret)
);
}
#[tokio::test]
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 (?, ?)")
.bind(COMFYUI_AGENT_SETTINGS_KEY)
.bind(&saved)
.execute(&pool)
.await
.unwrap();
let (_, secret) = ensure_comfyui_agent_key(&pool).await.unwrap();
assert_eq!(secret, saved);
let row = read_row(&pool).await.expect("row restored");
assert_eq!(row.1, saved);
}
#[tokio::test]
async fn row_present_settings_missing_writes_settings() {
let pool = fresh_pool().await;
sqlx::query(
"INSERT INTO downstream_api_keys (name, key, user_id, enabled, description) \
VALUES ('comfyui-agent', 'sk-comfyui-fromrow', NULL, 1, 'pre-existing')",
)
.execute(&pool)
.await
.unwrap();
let (_, secret) = ensure_comfyui_agent_key(&pool).await.unwrap();
assert_eq!(secret, "sk-comfyui-fromrow");
assert_eq!(
read_settings(&pool, COMFYUI_AGENT_SETTINGS_KEY).await,
Some("sk-comfyui-fromrow".to_string()),
);
}
#[tokio::test]
async fn repeated_invocation_is_stable() {
let pool = fresh_pool().await;
let (id1, secret1) = ensure_comfyui_agent_key(&pool).await.unwrap();
let (id2, secret2) = ensure_comfyui_agent_key(&pool).await.unwrap();
assert_eq!(id1, id2);
assert_eq!(secret1, secret2);
}
#[tokio::test]
async fn disabled_row_is_re_enabled() {
let pool = fresh_pool().await;
sqlx::query(
"INSERT INTO downstream_api_keys (name, key, user_id, enabled, description) \
VALUES ('comfyui-agent', 'sk-comfyui-disabled', NULL, 0, 'pre-existing')",
)
.execute(&pool)
.await
.unwrap();
let (_, secret) = ensure_comfyui_agent_key(&pool).await.unwrap();
assert_eq!(secret, "sk-comfyui-disabled");
let row = read_row(&pool).await.expect("row");
assert_eq!(row.2, 1, "enabled should be restored to 1");
}
}
src/services/comfyui_process.rs
+21 -1
View File
@@ -35,10 +35,16 @@ pub struct ComfyUIProcess {
started_at: Mutex<Option<Instant>>,
boosapi_url: String,
proxy_token: String,
builtin_key: String,
}
impl ComfyUIProcess {
pub fn new(config: ComfyUIConfig, boosapi_url: String, proxy_token: String) -> Arc<Self> {
pub fn new(
config: ComfyUIConfig,
boosapi_url: String,
proxy_token: String,
builtin_key: String,
) -> Arc<Self> {
let initial_status = if config.enabled {
ComfyUIStatus::Stopped
} else {
@@ -52,6 +58,7 @@ impl ComfyUIProcess {
started_at: Mutex::new(None),
boosapi_url,
proxy_token,
builtin_key,
})
}
@@ -174,6 +181,16 @@ impl ComfyUIProcess {
write_boosapi_config(install_dir, &self.model_backend_url(), &self.proxy_token)?;
// Point ComfyUI's OpenAI nodes at our shim and inject the
// built-in `comfyui-agent` downstream key as the Bearer token via
// env var. The `--comfy-api-base` CLI flag controls
// `comfy.cli_args.args.comfy_api_base`, which is what the
// OpenAI nodes prepend to paths like `/proxy/openai/...`.
let shim_api_base = format!(
"{}/api/comfyui-agent",
self.boosapi_url.trim_end_matches('/')
);
let mut command = Command::new(&self.config.python_bin);
command
.arg("main.py")
@@ -181,6 +198,9 @@ impl ComfyUIProcess {
.arg("127.0.0.1")
.arg("--port")
.arg(self.config.port.to_string())
.arg("--comfy-api-base")
.arg(&shim_api_base)
.env("BOOSAPI_BUILTIN_KEY", &self.builtin_key)
.current_dir(install_dir);
for arg in &self.config.extra_args {
src/services/mod.rs
+1
View File
@@ -1,5 +1,6 @@
pub mod announcement;
pub mod billing;
pub mod builtin_keys;
pub mod comfyui_process;
pub mod cron;
pub mod notification;