feat(comfy-cloud): 适配端点骨架 + POST /prompt 转发
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,101 @@
|
||||
use crate::error::{AppError, AppResult};
|
||||
use crate::services::comfy_cloud::site::resolve_comfy_cloud_site;
|
||||
use crate::services::comfy_cloud::translate::comfy_cloud_job_to_history;
|
||||
use crate::AppState;
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::{Path, RawQuery, State},
|
||||
http::{header, StatusCode},
|
||||
response::Response,
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/comfy-cloud/prompt", post(submit_prompt))
|
||||
.route("/comfy-cloud/history/{prompt_id}", get(get_history))
|
||||
.route("/comfy-cloud/view", get(get_view))
|
||||
}
|
||||
|
||||
fn client() -> reqwest::Client {
|
||||
reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()
|
||||
.unwrap_or_else(|_| reqwest::Client::new())
|
||||
}
|
||||
|
||||
async fn submit_prompt(
|
||||
State(state): State<AppState>,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let (base_url, api_key) = resolve_comfy_cloud_site(&state.db).await?;
|
||||
let resp = client()
|
||||
.post(format!("{base_url}/api/prompt"))
|
||||
.header("X-API-Key", api_key)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||
let status = resp.status();
|
||||
let text = resp.text().await.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(format!("Comfy Cloud prompt [{status}]: {text}")));
|
||||
}
|
||||
let value: Value = serde_json::from_str(&text).unwrap_or_else(|_| serde_json::json!({ "body": text }));
|
||||
Ok(Json(value))
|
||||
}
|
||||
|
||||
async fn get_history(
|
||||
State(state): State<AppState>,
|
||||
Path(prompt_id): Path<String>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let _ = (&state, &prompt_id);
|
||||
Ok(Json(serde_json::json!({})))
|
||||
}
|
||||
|
||||
async fn get_view(
|
||||
State(state): State<AppState>,
|
||||
RawQuery(query): RawQuery,
|
||||
) -> AppResult<Response> {
|
||||
let _ = (&state, &query);
|
||||
Ok(Response::builder().status(StatusCode::NOT_IMPLEMENTED).body(Body::empty())
|
||||
.map_err(|e| AppError::Internal(anyhow::anyhow!(e)))?)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use wiremock::matchers::{header as match_header, method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
use sqlx::sqlite::SqlitePoolOptions;
|
||||
|
||||
async fn db_with_site(base_url: &str) -> sqlx::SqlitePool {
|
||||
let pool = SqlitePoolOptions::new().max_connections(1).connect("sqlite::memory:").await.unwrap();
|
||||
sqlx::query("CREATE TABLE sites (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT, url TEXT NOT NULL, platform TEXT NOT NULL, status TEXT NOT NULL, api_key TEXT, sort_order INTEGER)")
|
||||
.execute(&pool).await.unwrap();
|
||||
sqlx::query("INSERT INTO sites (name,url,platform,status,api_key) VALUES ('cc',?,'comfy_cloud','active','key-123')")
|
||||
.bind(base_url).execute(&pool).await.unwrap();
|
||||
pool
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn submit_forwards_with_api_key() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/api/prompt"))
|
||||
.and(match_header("X-API-Key", "key-123"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "prompt_id": "p1" })))
|
||||
.mount(&server).await;
|
||||
|
||||
let pool = db_with_site(&server.uri()).await;
|
||||
let (base_url, api_key) = resolve_comfy_cloud_site(&pool).await.unwrap();
|
||||
let resp = client().post(format!("{base_url}/api/prompt"))
|
||||
.header("X-API-Key", api_key)
|
||||
.json(&serde_json::json!({ "prompt": {} }))
|
||||
.send().await.unwrap();
|
||||
let v: Value = resp.json().await.unwrap();
|
||||
assert_eq!(v["prompt_id"], "p1");
|
||||
}
|
||||
}
|
||||
@@ -24,6 +24,7 @@ const _: &[u8] = include_bytes!(concat!(
|
||||
mod admin;
|
||||
mod auth;
|
||||
mod checkin;
|
||||
mod comfy_cloud_adapter;
|
||||
mod comfyui;
|
||||
mod comfyui_agent;
|
||||
mod digital_human;
|
||||
|
||||
Reference in New Issue
Block a user