tsurumure пре 1 месец
родитељ
комит
edc82a27c7

+ 1 - 1
db/sys_user_role_relation.sql

@@ -15,7 +15,7 @@ CREATE TABLE `sys_user_role_relation` (
 ) ENGINE=INNODB AUTO_INCREMENT=1 DEFAULT CHARSET=utf8mb4 COMMENT='系统用户角色关系表';
 
 INSERT INTO sys_user_role_relation(user_id, role_id) VALUES
-    (1, 1), (2, 2), (3, 3), (4, 2), (5, 2), (6, 2),
+    (1, 1), (2, 1), (3, 3), (4, 2), (5, 2), (6, 2),
     (7, 2), (8, 2), (9, 2), (10, 2), (11, 2), (12, 2),
     (13, 1), (14, 1), (15, 1), (16, 1)
 ;

+ 3 - 1
src/main/java/com/backendsys/modules/sdk/comfyui/entity/ComfyuiResponse.java

@@ -4,7 +4,9 @@ import lombok.Data;
 
 @Data
 public class ComfyuiResponse {
-    private String client_id;       // 任务ID
+    private String url;             // 任务链接
+    private Integer url_port;        // 任务链接
+    private String client_id;       // 任务Client ID
     private String prompt_id;       // 任务ID
     private Integer number;         // 当前任务序号,可用于后续获取需要等待任务数的计算
     private Object node_errors;     // 错误信息

+ 20 - 2
src/main/java/com/backendsys/modules/sdk/comfyui/rabbitmq/ComfyuiRabbitListener.java

@@ -4,9 +4,12 @@ import cn.hutool.core.convert.Convert;
 import cn.hutool.core.util.StrUtil;
 import cn.hutool.json.JSONObject;
 import cn.hutool.json.JSONUtil;
+import com.backendsys.modules.sdk.comfyui.dao.ComfyuiTaskDao;
 import com.backendsys.modules.sdk.comfyui.entity.ComfyuiResponse;
+import com.backendsys.modules.sdk.comfyui.entity.ComfyuiTask;
 import com.backendsys.modules.sdk.comfyui.service.ComfyuiService;
 import com.backendsys.modules.sdk.comfyui.service.ComfyuiSocketService;
+import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
 import com.rabbitmq.client.Channel;
 import com.rabbitmq.client.GetResponse;
 import org.springframework.amqp.core.Message;
@@ -32,13 +35,16 @@ import java.util.concurrent.ConcurrentHashMap;
 @Lazy(false)
 public class ComfyuiRabbitListener {
 
-    @Autowired
-    private ComfyuiService comfyuiService;
     @Autowired
     private ComfyuiQueueConfig config;
     @Autowired
     private DeliveryTagHolder deliveryTagHolder;
 
+    @Autowired
+    private ComfyuiService comfyuiService;
+    @Autowired
+    private ComfyuiTaskDao comfyuiTaskDao;
+
 
     // [监听] 手动 ACK
     // (从 ComfyuiTaskService.tryPushNext 过来)
@@ -65,6 +71,18 @@ public class ComfyuiRabbitListener {
             ComfyuiResponse response = cfPromptResponseMono.block();
             System.out.println("[Comfyui][发起任务] response = " + response);
 
+            // -- [DB] 更新 任务ID、任务状态、任务URL、任务URL端口 --------------------------
+            ComfyuiTask entity = new ComfyuiTask();
+            entity.setExecute_prompt_id(response.getPrompt_id());
+            entity.setExecute_url(response.getUrl());
+            entity.setExecute_url_port(response.getUrl_port());
+            entity.setExecute_status(1);    // 1: 正在生成中
+
+            LambdaQueryWrapper<ComfyuiTask> wrapperTask = new LambdaQueryWrapper<>();
+            wrapperTask.eq(ComfyuiTask::getClient_id, client_id);
+            comfyuiTaskDao.update(entity, wrapperTask);
+            // ------------------------------------------------------------------------
+
             // 2. 暂存 deliveryTag
             System.out.println("put client_id = " + client_id);
             deliveryTagHolder.put(client_id, deliveryTag, channel);

+ 22 - 7
src/main/java/com/backendsys/modules/sdk/comfyui/service/impl/ComfyuiServiceImpl.java

@@ -6,12 +6,11 @@ import cn.hutool.json.JSONArray;
 import cn.hutool.json.JSONObject;
 import cn.hutool.json.JSONUtil;
 import com.backendsys.modules.common.Filter.WebClientFilter;
-import com.backendsys.modules.sdk.comfyui.entity.ComfyuiQueueItem;
-import com.backendsys.modules.sdk.comfyui.entity.ComfyuiRequest;
-import com.backendsys.modules.sdk.comfyui.entity.ComfyuiResponse;
-import com.backendsys.modules.sdk.comfyui.entity.ComfyuiQueue;
+import com.backendsys.modules.sdk.comfyui.dao.ComfyuiTaskDao;
+import com.backendsys.modules.sdk.comfyui.entity.*;
 import com.backendsys.modules.sdk.comfyui.service.ComfyuiService;
 import com.backendsys.modules.sdk.comfyui.service.ComfyuiSocketService;
+import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.http.HttpHeaders;
@@ -32,12 +31,13 @@ public class ComfyuiServiceImpl implements ComfyuiService {
     private String COMFYUI_TOKEN;
     @Value("${comfyui.ports}")
     private String COMFYUI_PORTS_STR;
-
-    @Autowired
-    private ComfyuiSocketService comfyuiSocketService;
     @Value("${comfyui.is-save}")
     private Boolean IS_SAVE;
 
+    @Autowired
+    private ComfyuiTaskDao comfyuiTaskDao;
+    @Autowired
+    private ComfyuiSocketService comfyuiSocketService;
 
     private WebClient webClient;
     public WebClient getWebClient() {
@@ -178,10 +178,25 @@ public class ComfyuiServiceImpl implements ComfyuiService {
                 .accept(MediaType.APPLICATION_JSON)
                 .bodyValue(bodyValue)
                 .exchangeToMono(response -> response.bodyToMono(ComfyuiResponse.class))
+                .map(r -> {
+                    r.setUrl(url);
+                    r.setUrl_port(port);
+                    return r;
+                })
                 .onErrorResume(e -> {
                     // 捕获所有异常(包括上面抛出的 RuntimeException)
                     ComfyuiResponse response = new ComfyuiResponse();
                     response.setNode_errors(e.getMessage());
+
+                    // -- [DB] 更新任务错误信息 --------------------------------------------------
+                    ComfyuiTask entity = new ComfyuiTask();
+                    entity.setReason(e.getMessage());
+                    entity.setExecute_status(3);
+                    LambdaQueryWrapper<ComfyuiTask> wrapperTask = new LambdaQueryWrapper<>();
+                    wrapperTask.eq(ComfyuiTask::getClient_id, client_id);
+                    comfyuiTaskDao.update(entity, wrapperTask);
+                    // ------------------------------------------------------------------------
+
                     return Mono.just(response);
                 });
 

+ 1 - 1
src/main/java/com/backendsys/modules/sdk/comfyui/service/impl/ComfyuiTaskServiceImpl.java

@@ -60,7 +60,7 @@ public class ComfyuiTaskServiceImpl implements ComfyuiTaskService {
         // 2. 抢令牌(原子)
         Boolean ok = redisTemplate.opsForValue().setIfAbsent(QUEUE_LOCK_KEY, "1", Duration.ofMinutes(10));
         if (!Boolean.TRUE.equals(ok)) {
-            return;          // 抢失败(并发极小概率)
+            return; // 抢失败(并发极小概率)
         }
 
         // 3. 该用户下一条未进主队列的任务