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

cherishsince
1 parent ab0a49a1
Showing 2 changed files with 19 additions and 1 deletions
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/MidjourneyWebSocketStarter.java
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/WssNotify.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);
+		}
 	}
 	/**
+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);
+}