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>

src/lib.rs

@@ -16,26 +16,34 @@ pub struct AppState {
     pub config: config::AppConfig,
     pub db: Arc<SqlitePool>,
     pub comfyui: Arc<services::comfyui_process::ComfyUIProcess>,
+    pub builtin_keys: Arc<services::builtin_keys::BuiltinKeys>,
 }
 
 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 == "::" {
             "127.0.0.1".to_string()
         } else {
             config.server.host.clone()
         };
         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(
             config.comfyui.clone(),
             boosapi_url,
             config.auth.proxy_token.clone(),
+            comfyui_agent_secret,
         );
 
         Self {
             config,
             db: Arc::new(db),
             comfyui,
+            builtin_keys: Arc::new(builtin_keys),
         }
     }
 }

src/main.rs

@@ -1,5 +1,6 @@
+use anyhow::Context;
 use axum::serve;
-use boosapi::{config::AppConfig, db, AppState};
+use boosapi::{config::AppConfig, db, services::builtin_keys, AppState};
 use clap::Parser;
 use std::ffi::OsString;
 use std::net::SocketAddr;
@@ -85,8 +86,20 @@ async fn main() -> anyhow::Result<()> {
         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
-    let state = AppState::new(config.clone(), pool);
+    let state = AppState::new(config.clone(), pool, builtin);
     if let Err(err) = state.comfyui.start().await {
         tracing::warn!(%err, "ComfyUI embedded process did not start");
     }

src/proxy/executors/mod.rs

@@ -70,9 +70,23 @@ pub async fn execute_chat_completions_stream(
     _state: AppState,
     route: UpstreamRoute,
     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>> {
     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
         .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 stream = resp.bytes_stream();
     use futures::StreamExt;
@@ -102,7 +115,6 @@ pub async fn execute_chat_completions_stream(
         match chunk_result {
             Ok(bytes) => {
                 let text = String::from_utf8_lossy(&bytes).to_string();
-                // Parse SSE lines from the chunk
                 for line in text.lines() {
                     if line.starts_with("data: ") {
                         chunks.push(line.to_string());

src/routes/proxy/mod.rs

@@ -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()),
+            },
         )
     }
 

src/services/builtin_keys.rs

@@ -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");
+    }
+}

src/services/comfyui_process.rs

@@ -35,10 +35,16 @@ pub struct ComfyUIProcess {
     started_at: Mutex<Option<Instant>>,
     boosapi_url: String,
     proxy_token: String,
+    builtin_key: String,
 }
 
 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 {
             ComfyUIStatus::Stopped
         } else {
@@ -52,6 +58,7 @@ impl ComfyUIProcess {
             started_at: Mutex::new(None),
             boosapi_url,
             proxy_token,
+            builtin_key,
         })
     }
 
@@ -174,6 +181,16 @@ impl ComfyUIProcess {
 
         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);
         command
             .arg("main.py")
@@ -181,6 +198,9 @@ impl ComfyUIProcess {
             .arg("127.0.0.1")
             .arg("--port")
             .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);
 
         for arg in &self.config.extra_args {

src/services/mod.rs

@@ -1,5 +1,6 @@
 pub mod announcement;
 pub mod billing;
+pub mod builtin_keys;
 pub mod comfyui_process;
 pub mod cron;
 pub mod notification;