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:
boos
2026-05-20 06:20:39 +08:00
parent 22751effd5
commit 9c0d4ec69f
7 changed files with 644 additions and 41 deletions
+9 -1
View File
@@ -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
View File
@@ -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");
} }
+15 -3
View File
@@ -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
View File
@@ -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()),
},
) )
} }
+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");
}
}
+21 -1
View File
@@ -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
View File
@@ -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;