boos/BoosAPI
2
0
Fork 2
Code Issues Pull Requests Packages Projects Releases Wiki 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/routes/proxy/mod.rs
+323 -34
View File
@@ -13,6 +13,7 @@ use crate::{
},
proxy::executors::{self, UpstreamRoute},
proxy::token_router::TokenRouter,
services::billing::UsageTracker,
AppState,
};
use axum::{
@@ -50,6 +51,7 @@ pub enum DownstreamSource {
pub fn routes() -> Router<AppState> {
Router::new()
.route("/v1/chat/completions", post(chat_completions))
.route("/v1/responses", post(responses_handler))
.route("/v1/completions", post(completions_handler))
.route("/v1/embeddings", post(embeddings_handler))
.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/videos", post(videos_handler))
.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 ───────────────────────────────────────────────────────────
/// POST /v1/chat/completions
async fn chat_completions(
/// POST /v1/responses — OpenAI Responses API (non-stream + stream).
///
/// 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>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> 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
.get("model")
@@ -84,14 +110,50 @@ async fn chat_completions(
.unwrap_or(false);
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 =
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, 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())
}
}
@@ -104,7 +166,7 @@ async fn completions_handler(
) -> AppResult<Json<Value>> {
info!("POST /v1/completions");
check_auth(&state, &headers).await?;
let key_id = check_auth(&state, &headers).await?;
let model = body
.get("model")
@@ -113,7 +175,10 @@ async fn completions_handler(
.to_string();
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))
}
@@ -125,7 +190,7 @@ async fn embeddings_handler(
) -> AppResult<Json<Value>> {
info!("POST /v1/embeddings");
check_auth(&state, &headers).await?;
let key_id = check_auth(&state, &headers).await?;
let model = body
.get("model")
@@ -134,7 +199,10 @@ async fn embeddings_handler(
.to_string();
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))
}
@@ -146,7 +214,7 @@ async fn audio_speech(
) -> AppResult<Json<Value>> {
info!("POST /v1/audio/speech");
check_auth(&state, &headers).await?;
let key_id = check_auth(&state, &headers).await?;
let model = body
.get("model")
@@ -155,7 +223,9 @@ async fn audio_speech(
.to_string();
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))
}
@@ -167,7 +237,7 @@ async fn audio_transcriptions(
) -> AppResult<Json<Value>> {
info!("POST /v1/audio/transcriptions");
check_auth(&state, &headers).await?;
let key_id = check_auth(&state, &headers).await?;
let model = body
.get("model")
@@ -177,7 +247,9 @@ async fn audio_transcriptions(
let route = resolve_route(&state, &model).await?;
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))
}
@@ -189,7 +261,7 @@ async fn audio_translations(
) -> AppResult<Json<Value>> {
info!("POST /v1/audio/translations");
check_auth(&state, &headers).await?;
let key_id = check_auth(&state, &headers).await?;
let model = body
.get("model")
@@ -199,7 +271,9 @@ async fn audio_translations(
let route = resolve_route(&state, &model).await?;
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))
}
@@ -211,15 +285,17 @@ async fn videos_handler(
) -> AppResult<Json<Value>> {
info!("POST /v1/videos");
check_auth(&state, &headers).await?;
let key_id = check_auth(&state, &headers).await?;
let model = body
.get("model")
.and_then(|m| m.as_str())
.unwrap_or("default");
.unwrap_or("default")
.to_string();
let route = resolve_route(&state, model).await?;
let result = executors::execute_proxy_request(state, route, "/videos", body).await?;
let route = resolve_route(&state, &model).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))
}
@@ -231,7 +307,7 @@ async fn images_generations(
) -> AppResult<Json<Value>> {
info!("POST /v1/images/generations");
check_auth(&state, &headers).await?;
let key_id = check_auth(&state, &headers).await?;
let model = body
.get("model")
@@ -240,7 +316,8 @@ async fn images_generations(
.to_string();
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))
}
@@ -248,7 +325,8 @@ async fn images_generations(
async fn list_models(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
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(
"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 ───────────────────────────────────────────────────────
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
.get("Authorization")
.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)?;
match result {
DownstreamAuthResult::Managed { key, .. } => {
DownstreamAuthResult::Managed { key } => {
if !key.enabled.unwrap_or(true) {
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 ──────────────────────────────────────────────────
/// 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(
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,
@@ -404,15 +527,21 @@ async fn execute_chat_completions_with_conductor(
})
.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(
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,
@@ -446,7 +575,7 @@ async fn execute_chat_completions_stream_with_conductor(
})
.await;
match conductor_result_to_app_result(result)? {
let chunks = match conductor_result_to_app_result(result)? {
Value::Array(chunks) => chunks
.into_iter()
.map(|chunk| {
@@ -454,11 +583,36 @@ async fn execute_chat_completions_stream_with_conductor(
AppError::UpstreamError("Invalid streaming response from conductor".into())
})
})
.collect(),
_ => Err(AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
)),
}
.collect::<AppResult<Vec<String>>>()?,
_ => {
return Err(AppError::UpstreamError(
"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> {
@@ -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)]
mod tests {
use super::*;
@@ -651,6 +937,9 @@ mod tests {
comfyui: Default::default(),
},
pool,
crate::services::builtin_keys::BuiltinKeys {
comfyui_agent: (0, "sk-comfyui-test".to_string()),
},
)
}