fix: use user api key for comfyui agent

src/routes/comfyui_agent.rs

@@ -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
+/// endpoints with the current caller's downstream key — every call then runs
-/// then runs through the regular accounting / routing pipeline.
+/// 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());
+                .unwrap_or_else(|| auth.outbound_key.clone());
-            (AgentMode::FallbackRemote, base, key)
+            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();
+        .map(str::trim)
-
+        .filter(|token| !token.is_empty())
-    if token == state.config.auth.auth_token {
-        return true;
 }
 
-    validate_jwt(&state.config.auth.jwt_secret, token)
+async fn user_downstream_key(state: &AppState, user_id: i64) -> Result<DownstreamApiKey, String> {
-        .map(|claims| claims.role == "admin" || claims.role == "user")
+    sqlx::query_as::<_, DownstreamApiKey>(
-        .unwrap_or(false)
+        "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) {
+    let auth = resolve_agent_auth(&headers, &state).await?;
-        return Err((
+    let (mode, upstream_base, _, key_name) = upstream_config(&state, &auth);
-            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",
-    };
     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) {
+    let _auth = resolve_agent_auth(&headers, &state).await?;
-        return Err((
-            axum::http::StatusCode::UNAUTHORIZED,
-            Json(json!({ "error": "Unauthorized" })),
-        ));
-    }
 
     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) {
+    let auth = resolve_agent_auth(&headers, &state).await?;
-        return Err((
-            axum::http::StatusCode::UNAUTHORIZED,
-            Json(json!({ "error": "Unauthorized" })),
-        ));
-    }
 
     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`
+    /// `/v1/chat/completions` endpoint with the current caller's downstream
-    /// downstream key as the bearer token.
+    /// 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;