Nvidia/BossAI
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: use user api key for comfyui agent

Browse Source
This commit is contained in:
boos
2026-05-21 16:55:32 +02:00
parent d0912e44b4
commit 0207ba748f
1 changed files with 133 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/routes/comfyui_agent.rs
+134 -49
View File
@@ -1,4 +1,4 @@
use crate::{auth::jwt::validate_jwt, AppState};
use crate::{auth::jwt::validate_jwt, db::models::downstream_api_key::DownstreamApiKey, AppState};
use axum::{
extract::State,
http::HeaderMap,
@@ -80,12 +80,20 @@ struct AgentMessage {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AgentMode {
/// Route through this server's own `/v1/*` using the built-in `comfyui-agent` downstream key.
/// Route through this server's own `/v1/*` using the caller's downstream key.
Local,
/// Legacy fallback: external upstream configured via `BOOSAPI_UPSTREAM_BASE` / `BOOSAPI_UPSTREAM_KEY`.
FallbackRemote,
}
#[derive(Debug, Clone)]
struct AgentAuthContext {
user_id: Option<i64>,
role: String,
outbound_key: String,
key_name: String,
}
impl AgentMode {
fn as_str(self) -> &'static str {
match self {
@@ -120,21 +128,29 @@ fn local_base(state: &AppState) -> String {
format!("http://{}:{}/v1", host, state.config.server.port)
}
/// Resolve `(mode, api_base, api_key)` for the agent's outbound calls.
/// Resolve `(mode, api_base, api_key, key_name)` for the agent's outbound calls.
///
/// Defaults to `AgentMode::Local`, which targets this server's own `/v1`
/// endpoints with the built-in `comfyui-agent` downstream key — every call
/// then runs through the regular accounting / routing pipeline.
/// endpoints with the current caller's downstream key — every call then runs
/// through the regular accounting / routing pipeline and is billed/logged
/// against that user key.
///
/// Setting `BOOSAPI_UPSTREAM_BASE` or `BOOSAPI_UPSTREAM_KEY` switches to
/// `FallbackRemote` for operators that still want to point straight at an
/// external aggregator.
fn upstream_config(state: &AppState) -> (AgentMode, String, String) {
fn upstream_config(
state: &AppState,
auth: &AgentAuthContext,
) -> (AgentMode, String, String, String) {
match detect_agent_mode() {
AgentMode::Local => {
let base = local_base(state);
let key = state.builtin_keys.comfyui_agent.1.clone();
(AgentMode::Local, base, key)
(
AgentMode::Local,
base,
auth.outbound_key.clone(),
auth.key_name.clone(),
)
}
AgentMode::FallbackRemote => {
let base = std::env::var("BOOSAPI_UPSTREAM_BASE")
@@ -146,30 +162,113 @@ fn upstream_config(state: &AppState) -> (AgentMode, String, String) {
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| state.config.auth.proxy_token.clone());
(AgentMode::FallbackRemote, base, key)
.unwrap_or_else(|| auth.outbound_key.clone());
let key_name = if std::env::var("BOOSAPI_UPSTREAM_KEY")
.ok()
.map(|value| !value.trim().is_empty())
.unwrap_or(false)
{
"external".to_string()
} else {
auth.key_name.clone()
};
(AgentMode::FallbackRemote, base, key, key_name)
}
}
}
fn authorize(headers: &HeaderMap, state: &AppState) -> bool {
fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
let auth_header = headers
.get("authorization")
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
let token = auth_header
auth_header
.strip_prefix("Bearer ")
.or_else(|| auth_header.strip_prefix("bearer "))
.unwrap_or_default();
if token == state.config.auth.auth_token {
return true;
.map(str::trim)
.filter(|token| !token.is_empty())
}
validate_jwt(&state.config.auth.jwt_secret, token)
.map(|claims| claims.role == "admin" || claims.role == "user")
.unwrap_or(false)
async fn user_downstream_key(state: &AppState, user_id: i64) -> Result<DownstreamApiKey, String> {
sqlx::query_as::<_, DownstreamApiKey>(
"SELECT *
FROM downstream_api_keys
WHERE user_id = ?
AND COALESCE(enabled, 1) = 1
AND (
expires_at IS NULL
OR expires_at = ''
OR datetime(expires_at) > datetime('now')
)
ORDER BY updated_at DESC, created_at DESC, id DESC
LIMIT 1",
)
.bind(user_id)
.fetch_optional(&*state.db)
.await
.map_err(|err| err.to_string())?
.ok_or_else(|| "当前用户没有可用的 API Key，请先在用户中心创建或启用 API Key".to_string())
}
async fn resolve_agent_auth(
headers: &HeaderMap,
state: &AppState,
) -> Result<AgentAuthContext, (axum::http::StatusCode, Json<Value>)> {
let token = extract_bearer_token(headers).ok_or_else(|| {
(
axum::http::StatusCode::UNAUTHORIZED,
Json(json!({ "error": "Unauthorized" })),
)
})?;
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
return Ok(AgentAuthContext {
user_id: None,
role: "admin".to_string(),
outbound_key: state.config.auth.proxy_token.clone(),
key_name: "global-proxy".to_string(),
});
}
let claims = validate_jwt(&state.config.auth.jwt_secret, token).map_err(|_| {
(
axum::http::StatusCode::UNAUTHORIZED,
Json(json!({ "error": "Unauthorized" })),
)
})?;
if claims.role != "admin" && claims.role != "user" {
return Err((
axum::http::StatusCode::FORBIDDEN,
Json(json!({ "error": "Forbidden" })),
));
}
let user_id = claims.sub.parse::<i64>().map_err(|_| {
(
axum::http::StatusCode::UNAUTHORIZED,
Json(json!({ "error": "Invalid user ID in token" })),
)
})?;
match user_downstream_key(state, user_id).await {
Ok(key) => Ok(AgentAuthContext {
user_id: Some(user_id),
role: claims.role,
outbound_key: key.key,
key_name: key.name,
}),
Err(_) if claims.role == "admin" => Ok(AgentAuthContext {
user_id: Some(user_id),
role: claims.role,
outbound_key: state.config.auth.proxy_token.clone(),
key_name: "global-proxy".to_string(),
}),
Err(message) => Err((
axum::http::StatusCode::FORBIDDEN,
Json(json!({ "error": message })),
)),
}
}
fn is_exact_model_pattern(value: &str) -> bool {
@@ -184,22 +283,15 @@ async fn status_handler(
State(state): State<AppState>,
headers: HeaderMap,
) -> Result<Json<Value>, (axum::http::StatusCode, Json<Value>)> {
if !authorize(&headers, &state) {
return Err((
axum::http::StatusCode::UNAUTHORIZED,
Json(json!({ "error": "Unauthorized" })),
));
}
let (mode, upstream_base, _) = upstream_config(&state);
let key_name = match mode {
AgentMode::Local => "comfyui-agent",
AgentMode::FallbackRemote => "external",
};
let auth = resolve_agent_auth(&headers, &state).await?;
let (mode, upstream_base, _, key_name) = upstream_config(&state, &auth);
Ok(Json(json!({
"success": true,
"model": state.config.comfyui.agent_model,
"mode": mode.as_str(),
"keyName": key_name,
"userId": auth.user_id,
"role": auth.role,
"upstreamBase": upstream_base,
"comfyui": state.comfyui.status_response().await,
})))
@@ -209,12 +301,7 @@ async fn models_handler(
State(state): State<AppState>,
headers: HeaderMap,
) -> Result<Json<Value>, (axum::http::StatusCode, Json<Value>)> {
if !authorize(&headers, &state) {
return Err((
axum::http::StatusCode::UNAUTHORIZED,
Json(json!({ "error": "Unauthorized" })),
));
}
let _auth = resolve_agent_auth(&headers, &state).await?;
let rows: Vec<String> = sqlx::query_scalar(
"SELECT model_pattern
@@ -477,12 +564,7 @@ async fn chat_handler(
Json(body): Json<AgentChatRequest>,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, (axum::http::StatusCode, Json<Value>)>
{
if !authorize(&headers, &state) {
return Err((
axum::http::StatusCode::UNAUTHORIZED,
Json(json!({ "error": "Unauthorized" })),
));
}
let auth = resolve_agent_auth(&headers, &state).await?;
if body.messages.is_empty()
|| body
@@ -513,10 +595,12 @@ async fn chat_handler(
})
})
.collect();
let (mode, api_base, api_key) = upstream_config(&state);
let (mode, api_base, api_key, key_name) = upstream_config(&state, &auth);
info!(
mode = mode.as_str(),
upstream_base = %api_base,
key_name = %key_name,
user_id = ?auth.user_id,
"comfyui-agent: resolved upstream",
);
let agent_model = requested_model.unwrap_or_else(|| state.config.comfyui.agent_model.clone());
@@ -949,7 +1033,7 @@ mod tests {
let state = make_state("0.0.0.0", 7860).await;
let body = call_status(state).await;
assert_eq!(body["mode"], "local");
assert_eq!(body["keyName"], "comfyui-agent");
assert_eq!(body["keyName"], "global-proxy");
assert_eq!(body["upstreamBase"], "http://127.0.0.1:7860/v1");
}
@@ -960,7 +1044,7 @@ mod tests {
let state = make_state("127.0.0.1", 9000).await;
let body = call_status(state).await;
assert_eq!(body["mode"], "local");
assert_eq!(body["keyName"], "comfyui-agent");
assert_eq!(body["keyName"], "global-proxy");
assert_eq!(body["upstreamBase"], "http://127.0.0.1:9000/v1");
}
@@ -971,7 +1055,7 @@ mod tests {
let state = make_state("0.0.0.0", 7860).await;
let body = call_status(state).await;
assert_eq!(body["mode"], "fallback_remote");
assert_eq!(body["keyName"], "external");
assert_eq!(body["keyName"], "global-proxy");
assert_eq!(body["upstreamBase"], "https://api.example.com/v1");
}
@@ -979,8 +1063,8 @@ mod tests {
}
/// 3.6 — In Local mode the agent's chat call must hit the local
/// `/v1/chat/completions` endpoint with the built-in `comfyui-agent`
/// downstream key as the bearer token.
/// `/v1/chat/completions` endpoint with the current caller's downstream
/// key as the bearer token.
#[tokio::test]
async fn agent_local_mode_routes_chat_to_local_upstream() {
let _g = ENV_GUARD.lock().await;
@@ -990,7 +1074,7 @@ mod tests {
let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n\n";
Mock::given(method("POST"))
.and(path("/v1/chat/completions"))
.and(header("authorization", "Bearer sk-comfyui-test"))
.and(header("authorization", "Bearer proxy"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
@@ -1013,6 +1097,7 @@ mod tests {
role: "user".to_string(),
content: "hello".to_string(),
}],
model: None,
}),
)
.await;