Integrate ComfyUI agents and video workflows

src/proxy/executors/mod.rs

@@ -28,6 +28,126 @@ fn http_client() -> Client {
         .expect("Failed to build HTTP client")
 }
 
+fn parse_upstream_json(
+    text: &str,
+    upstream_url: &str,
+    status: reqwest::StatusCode,
+) -> AppResult<Value> {
+    serde_json::from_str(text).map_err(|e| {
+        let snippet = text.chars().take(500).collect::<String>();
+        AppError::UpstreamError(format!(
+            "Invalid JSON from upstream {upstream_url} ({status}): {e}; body: {snippet:?}"
+        ))
+    })
+}
+
+fn parse_chat_completion_response(
+    text: &str,
+    upstream_url: &str,
+    status: reqwest::StatusCode,
+    fallback_model: &str,
+) -> AppResult<Value> {
+    match serde_json::from_str(text) {
+        Ok(json) => Ok(json),
+        Err(json_error) if looks_like_sse(text) => aggregate_chat_sse(text, fallback_model)
+            .ok_or_else(|| {
+                AppError::UpstreamError(format!(
+                    "Invalid SSE chat response from upstream {upstream_url} ({status}); JSON parse error: {json_error}; body: {:?}",
+                    text.chars().take(500).collect::<String>()
+                ))
+            }),
+        Err(json_error) => {
+            let snippet = text.chars().take(500).collect::<String>();
+            Err(AppError::UpstreamError(format!(
+                "Invalid JSON from upstream {upstream_url} ({status}): {json_error}; body: {snippet:?}"
+            )))
+        }
+    }
+}
+
+fn looks_like_sse(text: &str) -> bool {
+    text.trim_start().starts_with("data:")
+}
+
+fn aggregate_chat_sse(text: &str, fallback_model: &str) -> Option<Value> {
+    let mut id: Option<String> = None;
+    let mut created: Option<i64> = None;
+    let mut model: Option<String> = None;
+    let mut role: Option<String> = None;
+    let mut content = String::new();
+    let mut finish_reason: Option<Value> = None;
+
+    for line in text.lines() {
+        let Some(payload) = line.trim().strip_prefix("data:") else {
+            continue;
+        };
+        let payload = payload.trim();
+        if payload.is_empty() || payload == "[DONE]" {
+            continue;
+        }
+        let Ok(chunk) = serde_json::from_str::<Value>(payload) else {
+            continue;
+        };
+        if id.is_none() {
+            id = chunk.get("id").and_then(Value::as_str).map(str::to_string);
+        }
+        if created.is_none() {
+            created = chunk.get("created").and_then(Value::as_i64);
+        }
+        if model.is_none() {
+            model = chunk
+                .get("model")
+                .and_then(Value::as_str)
+                .map(str::to_string);
+        }
+        let Some(choice) = chunk
+            .get("choices")
+            .and_then(Value::as_array)
+            .and_then(|choices| choices.first())
+        else {
+            continue;
+        };
+        if let Some(delta) = choice.get("delta") {
+            if role.is_none() {
+                role = delta
+                    .get("role")
+                    .and_then(Value::as_str)
+                    .map(str::to_string);
+            }
+            if let Some(piece) = delta.get("content").and_then(Value::as_str) {
+                content.push_str(piece);
+            }
+        }
+        if finish_reason.is_none()
+            && !choice
+                .get("finish_reason")
+                .unwrap_or(&Value::Null)
+                .is_null()
+        {
+            finish_reason = choice.get("finish_reason").cloned();
+        }
+    }
+
+    if id.is_none() && content.is_empty() {
+        return None;
+    }
+
+    Some(serde_json::json!({
+        "id": id.unwrap_or_else(|| "chatcmpl-sse-aggregated".to_string()),
+        "object": "chat.completion",
+        "created": created.unwrap_or(0),
+        "model": model.unwrap_or_else(|| fallback_model.to_string()),
+        "choices": [{
+            "index": 0,
+            "message": {
+                "role": role.unwrap_or_else(|| "assistant".to_string()),
+                "content": content,
+            },
+            "finish_reason": finish_reason.unwrap_or_else(|| Value::String("stop".to_string())),
+        }],
+    }))
+}
+
 /// Execute a non-streaming chat completions request.
 pub async fn execute_chat_completions(
     _state: AppState,
@@ -58,8 +178,7 @@ pub async fn execute_chat_completions(
         )));
     }
 
-    let json: Value = serde_json::from_str(&text)
-        .map_err(|e| AppError::UpstreamError(format!("Invalid JSON from upstream: {e}")))?;
+    let json = parse_chat_completion_response(&text, &upstream_url, status, &route.model)?;
 
     Ok(json)
 }
@@ -165,12 +284,126 @@ pub async fn execute_proxy_request(
         )));
     }
 
-    let json: Value = serde_json::from_str(&text)
-        .map_err(|e| AppError::UpstreamError(format!("Invalid JSON from upstream: {e}")))?;
+    let json = parse_upstream_json(&text, &upstream_url, status)?;
 
     Ok(json)
 }
 
+/// Generic GET proxy request to an upstream endpoint path.
+pub async fn execute_proxy_get(route: UpstreamRoute, path_suffix: &str) -> AppResult<Value> {
+    let client = http_client();
+    let upstream_url = format!(
+        "{}/{}",
+        route.base_url.trim_end_matches('/'),
+        path_suffix.trim_start_matches('/')
+    );
+
+    let resp = client
+        .get(&upstream_url)
+        .header("Authorization", format!("Bearer {}", route.api_key))
+        .send()
+        .await
+        .map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
+
+    let status = resp.status();
+    let text = resp
+        .text()
+        .await
+        .map_err(|e| AppError::UpstreamError(format!("Failed to read upstream response: {e}")))?;
+
+    if !status.is_success() {
+        return Err(AppError::UpstreamError(format!(
+            "Upstream returned {status}: {text}"
+        )));
+    }
+
+    parse_upstream_json(&text, &upstream_url, status)
+}
+
+/// Execute an OpenAI-shaped video task request, with Volcengine Ark native
+/// task endpoints as a first-class compatibility target.
+pub async fn execute_video_generation(
+    state: AppState,
+    route: UpstreamRoute,
+    body: Value,
+) -> AppResult<Value> {
+    if is_volcengine_ark_base(&route.base_url) {
+        let body = volcengine_video_create_body(body, &route.model)?;
+        return execute_proxy_request(state, route, "/contents/generations/tasks", body).await;
+    }
+
+    execute_proxy_request(state, route, "/videos", body).await
+}
+
+fn is_volcengine_ark_base(base_url: &str) -> bool {
+    base_url.contains("ark.cn-beijing.volces.com") || base_url.contains("volces.com/api/v3")
+}
+
+fn volcengine_video_create_body(mut body: Value, fallback_model: &str) -> AppResult<Value> {
+    if body.get("content").is_some() {
+        return Ok(body);
+    }
+
+    let prompt = body
+        .get("prompt")
+        .and_then(Value::as_str)
+        .map(str::trim)
+        .filter(|value| !value.is_empty())
+        .ok_or_else(|| {
+            AppError::BadRequest("prompt is required for Volcengine video tasks".into())
+        })?
+        .to_string();
+    let model = body
+        .get("model")
+        .and_then(Value::as_str)
+        .map(str::trim)
+        .filter(|value| !value.is_empty())
+        .unwrap_or(fallback_model)
+        .to_string();
+
+    let mut text = prompt;
+    for (field, flag) in [
+        ("ratio", "--ratio"),
+        ("fps", "--fps"),
+        ("duration", "--dur"),
+        ("dur", "--dur"),
+        ("resolution", "--resolution"),
+    ] {
+        if let Some(value) = body.get(field).and_then(scalar_to_string) {
+            let marker = format!("{flag} ");
+            if !text.contains(&marker) {
+                text.push(' ');
+                text.push_str(flag);
+                text.push(' ');
+                text.push_str(&value);
+            }
+        }
+    }
+
+    if let Some(object) = body.as_object_mut() {
+        object.remove("prompt");
+        object.insert("model".to_string(), Value::String(model));
+        object.insert(
+            "content".to_string(),
+            Value::Array(vec![serde_json::json!({
+                "type": "text",
+                "text": text,
+            })]),
+        );
+    }
+
+    Ok(body)
+}
+
+fn scalar_to_string(value: &Value) -> Option<String> {
+    match value {
+        Value::String(s) => Some(s.trim().to_string()).filter(|s| !s.is_empty()),
+        Value::Number(n) => Some(n.to_string()),
+        Value::Bool(v) => Some(v.to_string()),
+        _ => None,
+    }
+}
+
 /// Execute an image generation request.
 pub async fn execute_image_generation(
     _state: AppState,
@@ -204,8 +437,7 @@ pub async fn execute_image_generation(
         )));
     }
 
-    let json: Value = serde_json::from_str(&text)
-        .map_err(|e| AppError::UpstreamError(format!("Invalid JSON from upstream: {e}")))?;
+    let json = parse_upstream_json(&text, &upstream_url, status)?;
 
     Ok(json)
 }

src/proxy/token_router/router.rs

@@ -232,15 +232,22 @@ impl TokenRouter {
         Ok(())
     }
 
-    pub async fn record_failure(&self, channel_id: i64, _info: FailureInfo) -> AppResult<()> {
+    pub async fn record_failure(&self, channel_id: i64, info: FailureInfo) -> AppResult<()> {
+        let cooldown_seconds = cooldown_seconds_for_failure(&info);
         sqlx::query(
             "UPDATE route_channels
              SET fail_count = COALESCE(fail_count, 0) + 1,
                  consecutive_fail_count = consecutive_fail_count + 1,
                  cooldown_level = cooldown_level + 1,
+                 cooldown_until = CASE
+                    WHEN ? > 0 THEN datetime('now', '+' || ? || ' seconds')
+                    ELSE cooldown_until
+                 END,
                  last_fail_at = datetime('now')
              WHERE id = ?",
         )
+        .bind(cooldown_seconds)
+        .bind(cooldown_seconds)
         .bind(channel_id)
         .execute(&*self.db)
         .await?;
@@ -248,6 +255,22 @@ impl TokenRouter {
     }
 }
 
+fn cooldown_seconds_for_failure(info: &FailureInfo) -> i64 {
+    match info.action.as_str() {
+        "retry_same_channel" => 10,
+        "refresh_auth" => 60,
+        "failover" => {
+            if matches!(info.status, Some(429 | 500..=599) | None) {
+                300
+            } else {
+                60
+            }
+        }
+        "terminal" => 900,
+        _ => 60,
+    }
+}
+
 trait IfEmptyThen {
     fn if_empty_then(self, fallback: String) -> String;
 }