diff --git a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/MidjourneyWebSocketStarter.java b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/MidjourneyWebSocketStarter.java
index e7ff41cc5..ddf7f3f7b 100644
--- a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/MidjourneyWebSocketStarter.java
+++ b/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);
+		}
 	}
 
 	/**
diff --git a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/WssNotify.java b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/webSocket/WssNotify.java
new file mode 100644
index 000000000..33c8a29e8
--- /dev/null
+++ b/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);
+}