feat(rust): wire downstream-key billing, add /v1/responses, ComfyUI shim
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:
+9
-1
@@ -16,26 +16,34 @@ pub struct AppState {
|
|||||||
pub config: config::AppConfig,
|
pub config: config::AppConfig,
|
||||||
pub db: Arc<SqlitePool>,
|
pub db: Arc<SqlitePool>,
|
||||||
pub comfyui: Arc<services::comfyui_process::ComfyUIProcess>,
|
pub comfyui: Arc<services::comfyui_process::ComfyUIProcess>,
|
||||||
|
pub builtin_keys: Arc<services::builtin_keys::BuiltinKeys>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl AppState {
|
impl AppState {
|
||||||
pub fn new(config: config::AppConfig, db: SqlitePool) -> Self {
|
pub fn new(
|
||||||
|
config: config::AppConfig,
|
||||||
|
db: SqlitePool,
|
||||||
|
builtin_keys: services::builtin_keys::BuiltinKeys,
|
||||||
|
) -> Self {
|
||||||
let boosapi_host = if config.server.host == "0.0.0.0" || config.server.host == "::" {
|
let boosapi_host = if config.server.host == "0.0.0.0" || config.server.host == "::" {
|
||||||
"127.0.0.1".to_string()
|
"127.0.0.1".to_string()
|
||||||
} else {
|
} else {
|
||||||
config.server.host.clone()
|
config.server.host.clone()
|
||||||
};
|
};
|
||||||
let boosapi_url = format!("http://{}:{}", boosapi_host, config.server.port);
|
let boosapi_url = format!("http://{}:{}", boosapi_host, config.server.port);
|
||||||
|
let comfyui_agent_secret = builtin_keys.comfyui_agent.1.clone();
|
||||||
let comfyui = services::comfyui_process::ComfyUIProcess::new(
|
let comfyui = services::comfyui_process::ComfyUIProcess::new(
|
||||||
config.comfyui.clone(),
|
config.comfyui.clone(),
|
||||||
boosapi_url,
|
boosapi_url,
|
||||||
config.auth.proxy_token.clone(),
|
config.auth.proxy_token.clone(),
|
||||||
|
comfyui_agent_secret,
|
||||||
);
|
);
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
config,
|
config,
|
||||||
db: Arc::new(db),
|
db: Arc::new(db),
|
||||||
comfyui,
|
comfyui,
|
||||||
|
builtin_keys: Arc::new(builtin_keys),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+15
-2
@@ -1,5 +1,6 @@
|
|||||||
|
use anyhow::Context;
|
||||||
use axum::serve;
|
use axum::serve;
|
||||||
use boosapi::{config::AppConfig, db, AppState};
|
use boosapi::{config::AppConfig, db, services::builtin_keys, AppState};
|
||||||
use clap::Parser;
|
use clap::Parser;
|
||||||
use std::ffi::OsString;
|
use std::ffi::OsString;
|
||||||
use std::net::SocketAddr;
|
use std::net::SocketAddr;
|
||||||
@@ -85,8 +86,20 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Bootstrap built-in system keys (idempotent).
|
||||||
|
let comfyui_agent_key = builtin_keys::ensure_comfyui_agent_key(&pool)
|
||||||
|
.await
|
||||||
|
.context("failed to bootstrap built-in ComfyUI Agent key")?;
|
||||||
|
tracing::info!(
|
||||||
|
key_id = comfyui_agent_key.0,
|
||||||
|
"Built-in ComfyUI Agent downstream key ready"
|
||||||
|
);
|
||||||
|
let builtin = builtin_keys::BuiltinKeys {
|
||||||
|
comfyui_agent: comfyui_agent_key,
|
||||||
|
};
|
||||||
|
|
||||||
// Build app state
|
// Build app state
|
||||||
let state = AppState::new(config.clone(), pool);
|
let state = AppState::new(config.clone(), pool, builtin);
|
||||||
if let Err(err) = state.comfyui.start().await {
|
if let Err(err) = state.comfyui.start().await {
|
||||||
tracing::warn!(%err, "ComfyUI embedded process did not start");
|
tracing::warn!(%err, "ComfyUI embedded process did not start");
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -70,9 +70,23 @@ pub async fn execute_chat_completions_stream(
|
|||||||
_state: AppState,
|
_state: AppState,
|
||||||
route: UpstreamRoute,
|
route: UpstreamRoute,
|
||||||
body: Value,
|
body: Value,
|
||||||
|
) -> AppResult<Vec<String>> {
|
||||||
|
execute_streaming_request(route, "/chat/completions", body).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute a streaming request against any upstream path (e.g. /responses).
|
||||||
|
/// Returns SSE `data: ...` lines verbatim.
|
||||||
|
pub async fn execute_streaming_request(
|
||||||
|
route: UpstreamRoute,
|
||||||
|
path_suffix: &str,
|
||||||
|
body: Value,
|
||||||
) -> AppResult<Vec<String>> {
|
) -> AppResult<Vec<String>> {
|
||||||
let client = http_client();
|
let client = http_client();
|
||||||
let upstream_url = format!("{}/chat/completions", route.base_url.trim_end_matches('/'));
|
let upstream_url = format!(
|
||||||
|
"{}/{}",
|
||||||
|
route.base_url.trim_end_matches('/'),
|
||||||
|
path_suffix.trim_start_matches('/')
|
||||||
|
);
|
||||||
|
|
||||||
let resp = client
|
let resp = client
|
||||||
.post(&upstream_url)
|
.post(&upstream_url)
|
||||||
@@ -94,7 +108,6 @@ pub async fn execute_chat_completions_stream(
|
|||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Read the streaming response as chunks
|
|
||||||
let mut chunks = Vec::new();
|
let mut chunks = Vec::new();
|
||||||
let mut stream = resp.bytes_stream();
|
let mut stream = resp.bytes_stream();
|
||||||
use futures::StreamExt;
|
use futures::StreamExt;
|
||||||
@@ -102,7 +115,6 @@ pub async fn execute_chat_completions_stream(
|
|||||||
match chunk_result {
|
match chunk_result {
|
||||||
Ok(bytes) => {
|
Ok(bytes) => {
|
||||||
let text = String::from_utf8_lossy(&bytes).to_string();
|
let text = String::from_utf8_lossy(&bytes).to_string();
|
||||||
// Parse SSE lines from the chunk
|
|
||||||
for line in text.lines() {
|
for line in text.lines() {
|
||||||
if line.starts_with("data: ") {
|
if line.starts_with("data: ") {
|
||||||
chunks.push(line.to_string());
|
chunks.push(line.to_string());
|
||||||
|
|||||||
+321
-32
@@ -13,6 +13,7 @@ use crate::{
|
|||||||
},
|
},
|
||||||
proxy::executors::{self, UpstreamRoute},
|
proxy::executors::{self, UpstreamRoute},
|
||||||
proxy::token_router::TokenRouter,
|
proxy::token_router::TokenRouter,
|
||||||
|
services::billing::UsageTracker,
|
||||||
AppState,
|
AppState,
|
||||||
};
|
};
|
||||||
use axum::{
|
use axum::{
|
||||||
@@ -50,6 +51,7 @@ pub enum DownstreamSource {
|
|||||||
pub fn routes() -> Router<AppState> {
|
pub fn routes() -> Router<AppState> {
|
||||||
Router::new()
|
Router::new()
|
||||||
.route("/v1/chat/completions", post(chat_completions))
|
.route("/v1/chat/completions", post(chat_completions))
|
||||||
|
.route("/v1/responses", post(responses_handler))
|
||||||
.route("/v1/completions", post(completions_handler))
|
.route("/v1/completions", post(completions_handler))
|
||||||
.route("/v1/embeddings", post(embeddings_handler))
|
.route("/v1/embeddings", post(embeddings_handler))
|
||||||
.route("/v1/audio/transcriptions", post(audio_transcriptions))
|
.route("/v1/audio/transcriptions", post(audio_transcriptions))
|
||||||
@@ -58,19 +60,43 @@ pub fn routes() -> Router<AppState> {
|
|||||||
.route("/v1/images/generations", post(images_generations))
|
.route("/v1/images/generations", post(images_generations))
|
||||||
.route("/v1/videos", post(videos_handler))
|
.route("/v1/videos", post(videos_handler))
|
||||||
.route("/v1/models", get(list_models))
|
.route("/v1/models", get(list_models))
|
||||||
|
// ── ComfyUI shim aliases ──────────────────────────────────────
|
||||||
|
// ComfyUI's OpenAI nodes hit /proxy/openai/... paths (see
|
||||||
|
// comfy_api_nodes/nodes_openai.py). Mounting them here as aliases
|
||||||
|
// means auth, conductor failover, and usage_logs all work for
|
||||||
|
// free — the built-in `comfyui-agent` downstream key authenticates
|
||||||
|
// via the same authorize_downstream_token path.
|
||||||
|
.route(
|
||||||
|
"/api/comfyui-agent/proxy/openai/images/generations",
|
||||||
|
post(images_generations),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/api/comfyui-agent/proxy/openai/v1/responses",
|
||||||
|
post(responses_handler),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/api/comfyui-agent/proxy/openai/v1/chat/completions",
|
||||||
|
post(chat_completions),
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Handlers ───────────────────────────────────────────────────────────
|
// ── Handlers ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
/// POST /v1/chat/completions
|
/// POST /v1/responses — OpenAI Responses API (non-stream + stream).
|
||||||
async fn chat_completions(
|
///
|
||||||
|
/// Behaves like /v1/chat/completions: routes through ProxyConductor for
|
||||||
|
/// failover, supports SSE streaming, and records usage_logs against the
|
||||||
|
/// matched downstream key. Capability gating is left to the upstream —
|
||||||
|
/// providers without Responses support will return 4xx and the error
|
||||||
|
/// surfaces to the caller through AppError::UpstreamError.
|
||||||
|
async fn responses_handler(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
Json(body): Json<Value>,
|
Json(body): Json<Value>,
|
||||||
) -> Result<Response, AppError> {
|
) -> Result<Response, AppError> {
|
||||||
info!("POST /v1/chat/completions");
|
info!("POST /v1/responses");
|
||||||
|
|
||||||
check_auth(&state, &headers).await?;
|
let key_id = check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
let model = body
|
let model = body
|
||||||
.get("model")
|
.get("model")
|
||||||
@@ -84,14 +110,50 @@ async fn chat_completions(
|
|||||||
.unwrap_or(false);
|
.unwrap_or(false);
|
||||||
|
|
||||||
if stream {
|
if stream {
|
||||||
let chunks = execute_chat_completions_stream_with_conductor(state, model, body).await?;
|
let chunks =
|
||||||
|
execute_responses_stream_with_conductor(state, key_id, model, body).await?;
|
||||||
let sse =
|
let sse =
|
||||||
Sse::new(stream::iter(chunks).map(|chunk| {
|
Sse::new(stream::iter(chunks).map(|chunk| {
|
||||||
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
|
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
|
||||||
}));
|
}));
|
||||||
Ok(sse.into_response())
|
Ok(sse.into_response())
|
||||||
} else {
|
} else {
|
||||||
let result = execute_chat_completions_with_conductor(state, model, body).await?;
|
let result = execute_responses_with_conductor(state, key_id, model, body).await?;
|
||||||
|
Ok(Json(result).into_response())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /v1/chat/completions
|
||||||
|
async fn chat_completions(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> Result<Response, AppError> {
|
||||||
|
info!("POST /v1/chat/completions");
|
||||||
|
|
||||||
|
let key_id = check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
|
let model = body
|
||||||
|
.get("model")
|
||||||
|
.and_then(|m| m.as_str())
|
||||||
|
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
let stream = body
|
||||||
|
.get("stream")
|
||||||
|
.and_then(|s| s.as_bool())
|
||||||
|
.unwrap_or(false);
|
||||||
|
|
||||||
|
if stream {
|
||||||
|
let chunks =
|
||||||
|
execute_chat_completions_stream_with_conductor(state, key_id, model, body).await?;
|
||||||
|
let sse =
|
||||||
|
Sse::new(stream::iter(chunks).map(|chunk| {
|
||||||
|
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
|
||||||
|
}));
|
||||||
|
Ok(sse.into_response())
|
||||||
|
} else {
|
||||||
|
let result = execute_chat_completions_with_conductor(state, key_id, model, body).await?;
|
||||||
Ok(Json(result).into_response())
|
Ok(Json(result).into_response())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -104,7 +166,7 @@ async fn completions_handler(
|
|||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
info!("POST /v1/completions");
|
info!("POST /v1/completions");
|
||||||
|
|
||||||
check_auth(&state, &headers).await?;
|
let key_id = check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
let model = body
|
let model = body
|
||||||
.get("model")
|
.get("model")
|
||||||
@@ -113,7 +175,10 @@ async fn completions_handler(
|
|||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
let route = resolve_route(&state, &model).await?;
|
let route = resolve_route(&state, &model).await?;
|
||||||
let result = executors::execute_proxy_request(state, route, "/completions", body).await?;
|
let result =
|
||||||
|
executors::execute_proxy_request(state.clone(), route, "/completions", body).await?;
|
||||||
|
let (prompt, completion) = extract_usage(&result);
|
||||||
|
record_billing(&state, key_id, &model, prompt, completion, 0.0).await;
|
||||||
Ok(Json(result))
|
Ok(Json(result))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -125,7 +190,7 @@ async fn embeddings_handler(
|
|||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
info!("POST /v1/embeddings");
|
info!("POST /v1/embeddings");
|
||||||
|
|
||||||
check_auth(&state, &headers).await?;
|
let key_id = check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
let model = body
|
let model = body
|
||||||
.get("model")
|
.get("model")
|
||||||
@@ -134,7 +199,10 @@ async fn embeddings_handler(
|
|||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
let route = resolve_route(&state, &model).await?;
|
let route = resolve_route(&state, &model).await?;
|
||||||
let result = executors::execute_proxy_request(state, route, "/embeddings", body).await?;
|
let result =
|
||||||
|
executors::execute_proxy_request(state.clone(), route, "/embeddings", body).await?;
|
||||||
|
let (prompt, completion) = extract_usage(&result);
|
||||||
|
record_billing(&state, key_id, &model, prompt, completion, 0.0).await;
|
||||||
Ok(Json(result))
|
Ok(Json(result))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -146,7 +214,7 @@ async fn audio_speech(
|
|||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
info!("POST /v1/audio/speech");
|
info!("POST /v1/audio/speech");
|
||||||
|
|
||||||
check_auth(&state, &headers).await?;
|
let key_id = check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
let model = body
|
let model = body
|
||||||
.get("model")
|
.get("model")
|
||||||
@@ -155,7 +223,9 @@ async fn audio_speech(
|
|||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
let route = resolve_route(&state, &model).await?;
|
let route = resolve_route(&state, &model).await?;
|
||||||
let result = executors::execute_proxy_request(state, route, "/audio/speech", body).await?;
|
let result =
|
||||||
|
executors::execute_proxy_request(state.clone(), route, "/audio/speech", body).await?;
|
||||||
|
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
|
||||||
Ok(Json(result))
|
Ok(Json(result))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -167,7 +237,7 @@ async fn audio_transcriptions(
|
|||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
info!("POST /v1/audio/transcriptions");
|
info!("POST /v1/audio/transcriptions");
|
||||||
|
|
||||||
check_auth(&state, &headers).await?;
|
let key_id = check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
let model = body
|
let model = body
|
||||||
.get("model")
|
.get("model")
|
||||||
@@ -177,7 +247,9 @@ async fn audio_transcriptions(
|
|||||||
|
|
||||||
let route = resolve_route(&state, &model).await?;
|
let route = resolve_route(&state, &model).await?;
|
||||||
let result =
|
let result =
|
||||||
executors::execute_proxy_request(state, route, "/audio/transcriptions", body).await?;
|
executors::execute_proxy_request(state.clone(), route, "/audio/transcriptions", body)
|
||||||
|
.await?;
|
||||||
|
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
|
||||||
Ok(Json(result))
|
Ok(Json(result))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -189,7 +261,7 @@ async fn audio_translations(
|
|||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
info!("POST /v1/audio/translations");
|
info!("POST /v1/audio/translations");
|
||||||
|
|
||||||
check_auth(&state, &headers).await?;
|
let key_id = check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
let model = body
|
let model = body
|
||||||
.get("model")
|
.get("model")
|
||||||
@@ -199,7 +271,9 @@ async fn audio_translations(
|
|||||||
|
|
||||||
let route = resolve_route(&state, &model).await?;
|
let route = resolve_route(&state, &model).await?;
|
||||||
let result =
|
let result =
|
||||||
executors::execute_proxy_request(state, route, "/audio/translations", body).await?;
|
executors::execute_proxy_request(state.clone(), route, "/audio/translations", body)
|
||||||
|
.await?;
|
||||||
|
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
|
||||||
Ok(Json(result))
|
Ok(Json(result))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -211,15 +285,17 @@ async fn videos_handler(
|
|||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
info!("POST /v1/videos");
|
info!("POST /v1/videos");
|
||||||
|
|
||||||
check_auth(&state, &headers).await?;
|
let key_id = check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
let model = body
|
let model = body
|
||||||
.get("model")
|
.get("model")
|
||||||
.and_then(|m| m.as_str())
|
.and_then(|m| m.as_str())
|
||||||
.unwrap_or("default");
|
.unwrap_or("default")
|
||||||
|
.to_string();
|
||||||
|
|
||||||
let route = resolve_route(&state, model).await?;
|
let route = resolve_route(&state, &model).await?;
|
||||||
let result = executors::execute_proxy_request(state, route, "/videos", body).await?;
|
let result = executors::execute_proxy_request(state.clone(), route, "/videos", body).await?;
|
||||||
|
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
|
||||||
Ok(Json(result))
|
Ok(Json(result))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -231,7 +307,7 @@ async fn images_generations(
|
|||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
info!("POST /v1/images/generations");
|
info!("POST /v1/images/generations");
|
||||||
|
|
||||||
check_auth(&state, &headers).await?;
|
let key_id = check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
let model = body
|
let model = body
|
||||||
.get("model")
|
.get("model")
|
||||||
@@ -240,7 +316,8 @@ async fn images_generations(
|
|||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
let route = resolve_route(&state, &model).await?;
|
let route = resolve_route(&state, &model).await?;
|
||||||
let result = executors::execute_image_generation(state, route, body).await?;
|
let result = executors::execute_image_generation(state.clone(), route, body).await?;
|
||||||
|
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
|
||||||
Ok(Json(result))
|
Ok(Json(result))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -248,7 +325,8 @@ async fn images_generations(
|
|||||||
async fn list_models(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
async fn list_models(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||||
info!("GET /v1/models");
|
info!("GET /v1/models");
|
||||||
|
|
||||||
check_auth(&state, &headers).await?;
|
// /v1/models is a directory listing — no upstream call, no billing.
|
||||||
|
let _ = check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
let rows: Vec<(String, String)> = sqlx::query_as(
|
let rows: Vec<(String, String)> = sqlx::query_as(
|
||||||
"SELECT tr.model_pattern, COALESCE(tr.model_mapping, tr.model_pattern) FROM token_routes tr WHERE tr.enabled = 1 ORDER BY tr.id",
|
"SELECT tr.model_pattern, COALESCE(tr.model_mapping, tr.model_pattern) FROM token_routes tr WHERE tr.enabled = 1 ORDER BY tr.id",
|
||||||
@@ -276,7 +354,10 @@ async fn list_models(State(state): State<AppState>, headers: HeaderMap) -> AppRe
|
|||||||
|
|
||||||
// ── Auth helper ───────────────────────────────────────────────────────
|
// ── Auth helper ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppError> {
|
/// Authenticate and return the managed downstream key id (if matched).
|
||||||
|
/// Global proxy token returns Ok(None) — billing is skipped for those calls
|
||||||
|
/// since they don't belong to a tracked downstream key.
|
||||||
|
async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<Option<i64>, AppError> {
|
||||||
let auth = headers
|
let auth = headers
|
||||||
.get("Authorization")
|
.get("Authorization")
|
||||||
.and_then(|v| v.to_str().ok())
|
.and_then(|v| v.to_str().ok())
|
||||||
@@ -289,16 +370,55 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
|
|||||||
.map_err(AppError::Unauthorized)?;
|
.map_err(AppError::Unauthorized)?;
|
||||||
|
|
||||||
match result {
|
match result {
|
||||||
DownstreamAuthResult::Managed { key, .. } => {
|
DownstreamAuthResult::Managed { key } => {
|
||||||
if !key.enabled.unwrap_or(true) {
|
if !key.enabled.unwrap_or(true) {
|
||||||
return Err(AppError::Forbidden("API key is disabled".into()));
|
return Err(AppError::Forbidden("API key is disabled".into()));
|
||||||
}
|
}
|
||||||
Ok(())
|
Ok(Some(key.id))
|
||||||
}
|
}
|
||||||
DownstreamAuthResult::Global { .. } => Ok(()),
|
DownstreamAuthResult::Global { .. } => Ok(None),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Best-effort usage logging. Only fires when the request was authenticated
|
||||||
|
/// via a managed downstream key (key_id is Some). Failures are logged but
|
||||||
|
/// never bubbled — billing must not break a successful proxy response.
|
||||||
|
async fn record_billing(
|
||||||
|
state: &AppState,
|
||||||
|
key_id: Option<i64>,
|
||||||
|
model: &str,
|
||||||
|
prompt_tokens: i64,
|
||||||
|
completion_tokens: i64,
|
||||||
|
cost: f64,
|
||||||
|
) {
|
||||||
|
let Some(key_id) = key_id else { return };
|
||||||
|
let tracker = UsageTracker::new(state.db.clone());
|
||||||
|
if let Err(err) = tracker
|
||||||
|
.record_usage(key_id, model, prompt_tokens, completion_tokens, cost)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
tracing::warn!(%err, key_id, model, "failed to record usage_logs entry");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pull `usage` block out of an OpenAI-shaped response. Returns
|
||||||
|
/// (prompt_tokens, completion_tokens). Missing fields default to 0.
|
||||||
|
fn extract_usage(response: &Value) -> (i64, i64) {
|
||||||
|
let usage = response.get("usage");
|
||||||
|
let prompt = usage
|
||||||
|
.and_then(|u| u.get("prompt_tokens").or_else(|| u.get("input_tokens")))
|
||||||
|
.and_then(|v| v.as_i64())
|
||||||
|
.unwrap_or(0);
|
||||||
|
let completion = usage
|
||||||
|
.and_then(|u| {
|
||||||
|
u.get("completion_tokens")
|
||||||
|
.or_else(|| u.get("output_tokens"))
|
||||||
|
})
|
||||||
|
.and_then(|v| v.as_i64())
|
||||||
|
.unwrap_or(0);
|
||||||
|
(prompt, completion)
|
||||||
|
}
|
||||||
|
|
||||||
// ── Route Resolution ──────────────────────────────────────────────────
|
// ── Route Resolution ──────────────────────────────────────────────────
|
||||||
|
|
||||||
/// Resolve upstream endpoint + API key for a given model.
|
/// Resolve upstream endpoint + API key for a given model.
|
||||||
@@ -367,10 +487,13 @@ fn build_conductor(state: &AppState) -> ProxyConductor {
|
|||||||
|
|
||||||
async fn execute_chat_completions_with_conductor(
|
async fn execute_chat_completions_with_conductor(
|
||||||
state: AppState,
|
state: AppState,
|
||||||
|
key_id: Option<i64>,
|
||||||
model: String,
|
model: String,
|
||||||
body: Value,
|
body: Value,
|
||||||
) -> AppResult<Value> {
|
) -> AppResult<Value> {
|
||||||
let conductor = build_conductor(&state);
|
let conductor = build_conductor(&state);
|
||||||
|
let billing_state = state.clone();
|
||||||
|
let billing_model = model.clone();
|
||||||
let result = conductor
|
let result = conductor
|
||||||
.execute(ExecuteInput {
|
.execute(ExecuteInput {
|
||||||
requested_model: model,
|
requested_model: model,
|
||||||
@@ -404,15 +527,21 @@ async fn execute_chat_completions_with_conductor(
|
|||||||
})
|
})
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
conductor_result_to_app_result(result)
|
let response = conductor_result_to_app_result(result)?;
|
||||||
|
let (prompt, completion) = extract_usage(&response);
|
||||||
|
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
|
||||||
|
Ok(response)
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn execute_chat_completions_stream_with_conductor(
|
async fn execute_chat_completions_stream_with_conductor(
|
||||||
state: AppState,
|
state: AppState,
|
||||||
|
key_id: Option<i64>,
|
||||||
model: String,
|
model: String,
|
||||||
body: Value,
|
body: Value,
|
||||||
) -> AppResult<Vec<String>> {
|
) -> AppResult<Vec<String>> {
|
||||||
let conductor = build_conductor(&state);
|
let conductor = build_conductor(&state);
|
||||||
|
let billing_state = state.clone();
|
||||||
|
let billing_model = model.clone();
|
||||||
let result = conductor
|
let result = conductor
|
||||||
.execute(ExecuteInput {
|
.execute(ExecuteInput {
|
||||||
requested_model: model,
|
requested_model: model,
|
||||||
@@ -446,7 +575,7 @@ async fn execute_chat_completions_stream_with_conductor(
|
|||||||
})
|
})
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
match conductor_result_to_app_result(result)? {
|
let chunks = match conductor_result_to_app_result(result)? {
|
||||||
Value::Array(chunks) => chunks
|
Value::Array(chunks) => chunks
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|chunk| {
|
.map(|chunk| {
|
||||||
@@ -454,11 +583,36 @@ async fn execute_chat_completions_stream_with_conductor(
|
|||||||
AppError::UpstreamError("Invalid streaming response from conductor".into())
|
AppError::UpstreamError("Invalid streaming response from conductor".into())
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
.collect(),
|
.collect::<AppResult<Vec<String>>>()?,
|
||||||
_ => Err(AppError::UpstreamError(
|
_ => {
|
||||||
|
return Err(AppError::UpstreamError(
|
||||||
"Invalid streaming response from conductor".into(),
|
"Invalid streaming response from conductor".into(),
|
||||||
)),
|
))
|
||||||
}
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Streaming responses don't include a usage block in the aggregate.
|
||||||
|
// OpenAI emits usage on the final `data: {...}` chunk only when
|
||||||
|
// `stream_options.include_usage = true` was set; parse it best-effort.
|
||||||
|
let (prompt, completion) = chunks
|
||||||
|
.iter()
|
||||||
|
.rev()
|
||||||
|
.find_map(|chunk| {
|
||||||
|
let payload = chunk.strip_prefix("data: ")?;
|
||||||
|
if payload.trim() == "[DONE]" {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let value: Value = serde_json::from_str(payload).ok()?;
|
||||||
|
let (p, c) = extract_usage(&value);
|
||||||
|
if p == 0 && c == 0 {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some((p, c))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.unwrap_or((0, 0));
|
||||||
|
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
|
||||||
|
Ok(chunks)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn conductor_result_to_app_result(result: ExecuteResult) -> AppResult<Value> {
|
fn conductor_result_to_app_result(result: ExecuteResult) -> AppResult<Value> {
|
||||||
@@ -478,6 +632,138 @@ fn conductor_result_to_app_result(result: ExecuteResult) -> AppResult<Value> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn execute_responses_with_conductor(
|
||||||
|
state: AppState,
|
||||||
|
key_id: Option<i64>,
|
||||||
|
model: String,
|
||||||
|
body: Value,
|
||||||
|
) -> AppResult<Value> {
|
||||||
|
let conductor = build_conductor(&state);
|
||||||
|
let billing_state = state.clone();
|
||||||
|
let billing_model = model.clone();
|
||||||
|
let result = conductor
|
||||||
|
.execute(ExecuteInput {
|
||||||
|
requested_model: model,
|
||||||
|
downstream_policy: None,
|
||||||
|
attempt_fn: Box::new(move |ctx| {
|
||||||
|
let state = state.clone();
|
||||||
|
let body = body_with_selected_model(body.clone(), &ctx.selected);
|
||||||
|
async move {
|
||||||
|
match executors::execute_proxy_request(
|
||||||
|
state,
|
||||||
|
selected_channel_to_route(ctx.selected),
|
||||||
|
"/responses",
|
||||||
|
body,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(response) => AttemptResult::Success {
|
||||||
|
response,
|
||||||
|
latency_ms: None,
|
||||||
|
cost: None,
|
||||||
|
},
|
||||||
|
Err(err) => AttemptResult::Failure {
|
||||||
|
action: crate::proxy::conductor::AttemptFailureAction::Failover,
|
||||||
|
status: None,
|
||||||
|
raw_error_text: Some(err.to_string()),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}),
|
||||||
|
on_terminal_failure: None,
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let response = conductor_result_to_app_result(result)?;
|
||||||
|
let (prompt, completion) = extract_usage(&response);
|
||||||
|
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
|
||||||
|
Ok(response)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute_responses_stream_with_conductor(
|
||||||
|
state: AppState,
|
||||||
|
key_id: Option<i64>,
|
||||||
|
model: String,
|
||||||
|
body: Value,
|
||||||
|
) -> AppResult<Vec<String>> {
|
||||||
|
let conductor = build_conductor(&state);
|
||||||
|
let billing_state = state.clone();
|
||||||
|
let billing_model = model.clone();
|
||||||
|
let result = conductor
|
||||||
|
.execute(ExecuteInput {
|
||||||
|
requested_model: model,
|
||||||
|
downstream_policy: None,
|
||||||
|
attempt_fn: Box::new(move |ctx| {
|
||||||
|
let body = body_with_selected_model(body.clone(), &ctx.selected);
|
||||||
|
async move {
|
||||||
|
match executors::execute_streaming_request(
|
||||||
|
selected_channel_to_route(ctx.selected),
|
||||||
|
"/responses",
|
||||||
|
body,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(chunks) => AttemptResult::Success {
|
||||||
|
response: Value::Array(chunks.into_iter().map(Value::String).collect()),
|
||||||
|
latency_ms: None,
|
||||||
|
cost: None,
|
||||||
|
},
|
||||||
|
Err(err) => AttemptResult::Failure {
|
||||||
|
action: crate::proxy::conductor::AttemptFailureAction::Failover,
|
||||||
|
status: None,
|
||||||
|
raw_error_text: Some(err.to_string()),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}),
|
||||||
|
on_terminal_failure: None,
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let chunks = match conductor_result_to_app_result(result)? {
|
||||||
|
Value::Array(chunks) => chunks
|
||||||
|
.into_iter()
|
||||||
|
.map(|chunk| {
|
||||||
|
chunk.as_str().map(str::to_string).ok_or_else(|| {
|
||||||
|
AppError::UpstreamError("Invalid streaming response from conductor".into())
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect::<AppResult<Vec<String>>>()?,
|
||||||
|
_ => {
|
||||||
|
return Err(AppError::UpstreamError(
|
||||||
|
"Invalid streaming response from conductor".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Responses API streams `response.completed` events that contain a
|
||||||
|
// `usage` object with input_tokens / output_tokens. extract_usage already
|
||||||
|
// handles both naming conventions.
|
||||||
|
let (prompt, completion) = chunks
|
||||||
|
.iter()
|
||||||
|
.rev()
|
||||||
|
.find_map(|chunk| {
|
||||||
|
let payload = chunk.strip_prefix("data: ")?;
|
||||||
|
if payload.trim() == "[DONE]" {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let value: Value = serde_json::from_str(payload).ok()?;
|
||||||
|
// Responses API nests usage under `response.usage`.
|
||||||
|
let usage_holder = value.get("response").unwrap_or(&value);
|
||||||
|
let (p, c) = extract_usage(usage_holder);
|
||||||
|
if p == 0 && c == 0 {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some((p, c))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.unwrap_or((0, 0));
|
||||||
|
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
|
||||||
|
Ok(chunks)
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -651,6 +937,9 @@ mod tests {
|
|||||||
comfyui: Default::default(),
|
comfyui: Default::default(),
|
||||||
},
|
},
|
||||||
pool,
|
pool,
|
||||||
|
crate::services::builtin_keys::BuiltinKeys {
|
||||||
|
comfyui_agent: (0, "sk-comfyui-test".to_string()),
|
||||||
|
},
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -35,10 +35,16 @@ pub struct ComfyUIProcess {
|
|||||||
started_at: Mutex<Option<Instant>>,
|
started_at: Mutex<Option<Instant>>,
|
||||||
boosapi_url: String,
|
boosapi_url: String,
|
||||||
proxy_token: String,
|
proxy_token: String,
|
||||||
|
builtin_key: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ComfyUIProcess {
|
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 {
|
let initial_status = if config.enabled {
|
||||||
ComfyUIStatus::Stopped
|
ComfyUIStatus::Stopped
|
||||||
} else {
|
} else {
|
||||||
@@ -52,6 +58,7 @@ impl ComfyUIProcess {
|
|||||||
started_at: Mutex::new(None),
|
started_at: Mutex::new(None),
|
||||||
boosapi_url,
|
boosapi_url,
|
||||||
proxy_token,
|
proxy_token,
|
||||||
|
builtin_key,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -174,6 +181,16 @@ impl ComfyUIProcess {
|
|||||||
|
|
||||||
write_boosapi_config(install_dir, &self.model_backend_url(), &self.proxy_token)?;
|
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);
|
let mut command = Command::new(&self.config.python_bin);
|
||||||
command
|
command
|
||||||
.arg("main.py")
|
.arg("main.py")
|
||||||
@@ -181,6 +198,9 @@ impl ComfyUIProcess {
|
|||||||
.arg("127.0.0.1")
|
.arg("127.0.0.1")
|
||||||
.arg("--port")
|
.arg("--port")
|
||||||
.arg(self.config.port.to_string())
|
.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);
|
.current_dir(install_dir);
|
||||||
|
|
||||||
for arg in &self.config.extra_args {
|
for arg in &self.config.extra_args {
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
pub mod announcement;
|
pub mod announcement;
|
||||||
pub mod billing;
|
pub mod billing;
|
||||||
|
pub mod builtin_keys;
|
||||||
pub mod comfyui_process;
|
pub mod comfyui_process;
|
||||||
pub mod cron;
|
pub mod cron;
|
||||||
pub mod notification;
|
pub mod notification;
|
||||||
|
|||||||
Reference in New Issue
Block a user