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

feat: merge opus comfyui and login updates

Browse Source
This commit is contained in:
boos
2026-05-21 16:32:04 +02:00
parent c16479da7c
commit 9c84759498
105 changed files with 20160 additions and 1034 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/routes/comfyui_agent.rs
+331 -15
View File
@@ -1,4 +1,4 @@
use crate::AppState;
use crate::{auth::jwt::validate_jwt, AppState};
use axum::{
extract::State,
http::HeaderMap,
@@ -62,12 +62,14 @@ emotion: 平静
pub fn routes() -> Router<AppState> {
Router::new()
.route("/api/comfyui-agent/status", get(status_handler))
.route("/api/comfyui-agent/models", get(models_handler))
.route("/api/comfyui-agent/chat", post(chat_handler))
}
#[derive(Debug, Deserialize)]
struct AgentChatRequest {
messages: Vec<AgentMessage>,
model: Option<String>,
}
#[derive(Debug, Deserialize)]
@@ -76,18 +78,78 @@ struct AgentMessage {
content: String,
}
fn upstream_config(state: &AppState) -> (String, String) {
let base = std::env::var("BOOSAPI_UPSTREAM_BASE")
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AgentMode {
/// Route through this server's own `/v1/*` using the built-in `comfyui-agent` downstream key.
Local,
/// Legacy fallback: external upstream configured via `BOOSAPI_UPSTREAM_BASE` / `BOOSAPI_UPSTREAM_KEY`.
FallbackRemote,
}
impl AgentMode {
fn as_str(self) -> &'static str {
match self {
AgentMode::Local => "local",
AgentMode::FallbackRemote => "fallback_remote",
}
}
}
fn detect_agent_mode() -> AgentMode {
let base_override = std::env::var("BOOSAPI_UPSTREAM_BASE")
.ok()
.map(|value| value.trim().trim_end_matches('/').to_string())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| "https://api.boos.lat/v1".to_string());
let key = std::env::var("BOOSAPI_UPSTREAM_KEY")
.map(|v| v.trim().to_string())
.filter(|v| !v.is_empty());
let key_override = std::env::var("BOOSAPI_UPSTREAM_KEY")
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| state.config.auth.proxy_token.clone());
(base, key)
.map(|v| v.trim().to_string())
.filter(|v| !v.is_empty());
if base_override.is_some() || key_override.is_some() {
AgentMode::FallbackRemote
} else {
AgentMode::Local
}
}
fn local_base(state: &AppState) -> String {
let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
"127.0.0.1".to_string()
} else {
state.config.server.host.clone()
};
format!("http://{}:{}/v1", host, state.config.server.port)
}
/// Resolve `(mode, api_base, api_key)` 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.
///
/// 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) {
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::FallbackRemote => {
let base = std::env::var("BOOSAPI_UPSTREAM_BASE")
.ok()
.map(|value| value.trim().trim_end_matches('/').to_string())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| "https://api.boos.lat/v1".to_string());
let key = std::env::var("BOOSAPI_UPSTREAM_KEY")
.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)
}
}
}
fn authorize(headers: &HeaderMap, state: &AppState) -> bool {
@@ -101,7 +163,21 @@ fn authorize(headers: &HeaderMap, state: &AppState) -> bool {
.or_else(|| auth_header.strip_prefix("bearer "))
.unwrap_or_default();
token == state.config.auth.auth_token
if token == state.config.auth.auth_token {
return true;
}
validate_jwt(&state.config.auth.jwt_secret, token)
.map(|claims| claims.role == "admin" || claims.role == "user")
.unwrap_or(false)
}
fn is_exact_model_pattern(value: &str) -> bool {
let normalized = value.trim();
!normalized.is_empty()
&& !normalized.to_ascii_lowercase().starts_with("re:")
&& !normalized.contains('*')
&& !normalized.contains('?')
}
async fn status_handler(
@@ -114,15 +190,62 @@ async fn status_handler(
Json(json!({ "error": "Unauthorized" })),
));
}
let (upstream_base, _) = upstream_config(&state);
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,
"upstreamBase": upstream_base,
"comfyui": state.comfyui.status_response().await,
})))
}
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 rows: Vec<String> = sqlx::query_scalar(
"SELECT model_pattern
FROM token_routes
WHERE COALESCE(enabled, 1) = 1
ORDER BY model_pattern COLLATE NOCASE",
)
.fetch_all(&*state.db)
.await
.map_err(|err| {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": err.to_string() })),
)
})?;
let mut seen = std::collections::HashSet::new();
let models: Vec<String> = rows
.into_iter()
.map(|value| value.trim().to_string())
.filter(|value| is_exact_model_pattern(value))
.filter(|value| seen.insert(value.clone()))
.collect();
Ok(Json(json!({
"success": true,
"models": models,
"defaultModel": state.config.comfyui.agent_model,
})))
}
#[derive(Debug, Default)]
struct ParsedSections {
characters: Vec<CharacterSection>,
@@ -373,6 +496,13 @@ async fn chat_handler(
));
}
let requested_model = body
.model
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let messages: Vec<Value> = body
.messages
.into_iter()
@@ -383,8 +513,13 @@ async fn chat_handler(
})
})
.collect();
let (api_base, api_key) = upstream_config(&state);
let agent_model = state.config.comfyui.agent_model.clone();
let (mode, api_base, api_key) = upstream_config(&state);
info!(
mode = mode.as_str(),
upstream_base = %api_base,
"comfyui-agent: resolved upstream",
);
let agent_model = requested_model.unwrap_or_else(|| state.config.comfyui.agent_model.clone());
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
tokio::spawn(async move {
@@ -723,3 +858,184 @@ async fn chat_handler(
Ok(Sse::new(UnboundedReceiverStream::new(rx)))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig};
use crate::services::builtin_keys::BuiltinKeys;
use axum::http::HeaderValue;
use sqlx::sqlite::SqlitePoolOptions;
use tokio::sync::Mutex;
use wiremock::matchers::{header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
/// `BOOSAPI_UPSTREAM_*` is process-global. Serialize tests that read or
/// mutate it so cargo's parallel runner does not race their assertions.
/// Uses `tokio::sync::Mutex` so the guard can safely span `.await` points
/// (mock server setup, agent worker poll loop) without tripping
/// `clippy::await_holding_lock`.
static ENV_GUARD: Mutex<()> = Mutex::const_new(());
async fn make_state(host: &str, port: u16) -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
AppState::new(
AppConfig {
server: ServerConfig {
host: host.to_string(),
port,
log_level: "info".into(),
cors_origins: vec![],
},
db: DbConfig {
url: "sqlite::memory:".into(),
max_connections: 1,
ssl: false,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
},
auth: AuthConfig {
jwt_secret: "secret".into(),
jwt_expiry_hours: 1,
auth_token: "auth".into(),
proxy_token: "proxy".into(),
},
proxy: ProxyConfig {
upstream_timeout_secs: 30,
first_byte_timeout_secs: 30,
max_retries: 3,
proxy_log_retention_days: 7,
},
comfyui: Default::default(),
},
pool,
BuiltinKeys {
comfyui_agent: (1, "sk-comfyui-test".into()),
},
)
}
fn clear_upstream_env() {
std::env::remove_var("BOOSAPI_UPSTREAM_BASE");
std::env::remove_var("BOOSAPI_UPSTREAM_KEY");
}
fn auth_headers() -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert("Authorization", HeaderValue::from_static("Bearer auth"));
headers
}
async fn call_status(state: AppState) -> Value {
let result = status_handler(State(state), auth_headers()).await;
result.expect("status_handler returned error").0
}
/// 3.5 — `/api/comfyui-agent/status` reflects the env state in one shot:
/// no env → local; whitespace/empty env → still local (treated as unset);
/// real env → fallback_remote with trimmed `upstreamBase`.
#[tokio::test]
async fn status_reflects_env_state_across_three_scenarios() {
let _g = ENV_GUARD.lock().await;
clear_upstream_env();
// Scenario A: env unset → local; 0.0.0.0 host normalizes to 127.0.0.1.
{
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["upstreamBase"], "http://127.0.0.1:7860/v1");
}
// Scenario B: env set to whitespace/empty → treated as unset → local.
std::env::set_var("BOOSAPI_UPSTREAM_BASE", " ");
std::env::set_var("BOOSAPI_UPSTREAM_KEY", "");
{
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["upstreamBase"], "http://127.0.0.1:9000/v1");
}
// Scenario C: env set → fallback_remote, base trimmed of trailing slash.
std::env::set_var("BOOSAPI_UPSTREAM_BASE", "https://api.example.com/v1/");
std::env::remove_var("BOOSAPI_UPSTREAM_KEY");
{
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["upstreamBase"], "https://api.example.com/v1");
}
clear_upstream_env();
}
/// 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.
#[tokio::test]
async fn agent_local_mode_routes_chat_to_local_upstream() {
let _g = ENV_GUARD.lock().await;
clear_upstream_env();
let mock = MockServer::start().await;
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"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(sse),
)
.mount(&mock)
.await;
let addr = mock.address();
let host = addr.ip().to_string();
let port = addr.port();
let state = make_state(&host, port).await;
let result = chat_handler(
State(state),
auth_headers(),
Json(AgentChatRequest {
messages: vec![AgentMessage {
role: "user".to_string(),
content: "hello".to_string(),
}],
}),
)
.await;
assert!(result.is_ok(), "chat_handler should accept the request");
// The Sse stream is dropped here; the worker still runs the upstream
// POST before its streaming loop notices `tx.closed()`.
// Wait up to ~2s for the spawned worker's POST to land at wiremock.
let mut hit = false;
for _ in 0..40 {
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
let received = mock.received_requests().await.unwrap_or_default();
if received
.iter()
.any(|r| r.url.path() == "/v1/chat/completions")
{
hit = true;
break;
}
}
assert!(
hit,
"expected /v1/chat/completions to be hit by the local-mode agent"
);
}
}