Bläddra i källkod

【增强】MidjourneyWebSocket token无效会认证失败,增加 wss 通知回调,给业务方法处理

cherishsince 1 år sedan
förälder
incheckning
8bf5480bbf

+ 6 - 1
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/MidjourneyWebSocketStarter.java

@@ -56,6 +56,7 @@ public class MidjourneyWebSocketStarter implements WebSocketStarter {
 	 * 链接成功的 session
 	 */
 	private WebSocketSession webSocketSession = null;
+	private WssNotify wssNotify = null;
 
 	public MidjourneyWebSocketStarter(String wssServer,
 									  String resumeWss,
@@ -68,7 +69,8 @@ public class MidjourneyWebSocketStarter implements WebSocketStarter {
 	}
 
 	@Override
-	public void start() {
+	public void start(WssNotify wssNotify) {
+		this.wssNotify = wssNotify;
 		start(false);
 	}
 
@@ -178,6 +180,9 @@ public class MidjourneyWebSocketStarter implements WebSocketStarter {
 
 	private void notifyWssLock(int code, String reason) {
 		System.err.println("notifyWssLock: " + code + " - " + reason);
+		if (wssNotify != null) {
+			wssNotify.notify(code, reason);
+		}
 	}
 
 	/**

+ 13 - 0
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/WssNotify.java

@@ -0,0 +1,13 @@
+package cn.iocoder.yudao.framework.ai.midjourney.webSocket;
+
+/**
+ * 通知信息
+ *
+ * @author fansili
+ * @time 2024/4/29 14:21
+ * @since 1.0
+ */
+public interface WssNotify {
+
+    void notify(int code, String message);
+}