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 admin;
|
||||||
mod auth;
|
mod auth;
|
||||||
mod checkin;
|
mod checkin;
|
||||||
|
mod comfy_cloud_adapter;
|
||||||
mod comfyui;
|
mod comfyui;
|
||||||
mod comfyui_agent;
|
mod comfyui_agent;
|
||||||
mod digital_human;
|
mod digital_human;
|
||||||
|
|||||||
Reference in New Issue
Block a user