为什么要有两条路径： + *
两条路径共用 {@link WebSocketServer} 里的会话表，所以推送逻辑一处即可。 + * 两个类分开而不是继承，是因为 JSR-356 对继承来的 {@code @OnOpen}/{@code @OnClose} + * 支持依赖容器实现，重复这十几行更保险。 + */ +@ServerEndpoint(value = "/api/chat/{userId}") +public class ApiChatWebSocketServer { + + /** + * 接收userId + */ + private String userId = ""; + + @OnOpen + public void onOpen(Session session, @PathParam("userId") String userId) { + this.userId = userId; + WebSocketServer.register(session, userId); + } + + @OnClose + public void onClose(Session session) { + if (!userId.isEmpty()) { + WebSocketServer.unregister(userId, session); + } + } + + /** + * 心跳入口，语义同 {@link WebSocketServer#onMessage(String, Session)}。 + */ + @OnMessage + public void onMessage(String message, Session session) { + // 心跳与任何客户端文本都不参与业务 + } +} diff --git a/src/main/java/com/gxwebsoft/ai/config/ChatStreamRegistry.java b/src/main/java/com/gxwebsoft/ai/config/ChatStreamRegistry.java new file mode 100644 index 0000000..7bb940a --- /dev/null +++ b/src/main/java/com/gxwebsoft/ai/config/ChatStreamRegistry.java @@ -0,0 +1,68 @@ +package com.gxwebsoft.ai.config; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Consumer; + +/** + * HTTP 流式（SSE）通道登记表：WebSocket 连不上时的退路。 + * + *
生产网关在 {@code /api} 前面就把 WebSocket 的 {@code Upgrade} 头丢了， + * 握手永远到不了 Tomcat（网关侧不能改），所以端上加了一条 HTTP 流式通道： + * + *
+ * POST /api/ai/chat/message/stream ← 请求体与 WS 版完全一样
+ * data: {"data":{...}} ← 每个 SSE 事件一行，内容与 WS 帧逐字节相同
+ *
+ *
+ * 推送侧仍然只认 {@code webSocketServer.sendMessage(userId, payload)} 一个入口 + * （两条 AI 链路 3000+ 次调用的都是它），这里通过「同一用户登记一个 sink」把消息 + * 改投给 HTTP 连接，业务代码一行不用改。 + * + *
键用 userId 而不是连接 id：WS 那边也是按用户名建索引，语义保持一致，
+ * 同一时刻一个用户只跑一条 AI 流。
+ */
+public final class ChatStreamRegistry {
+
+ private static final Map 这里才是值得打日志的位置：一次连接一条，低频且信息量大
- * （排查「大纲/创意生成时端上收不到流式内容」时，第一件事就是确认端到底连没连上）。
- * 推送路径 {@link #sendMessage(String, String)} 相反，单次请求会被调用上千次，必须静默。
- */
- @OnOpen
- public void onOpen(Session session, @PathParam("userId") String userId) {
- this.session = session;
- this.userId = userId;
- if (webSocketMap.containsKey(userId)) {
- webSocketMap.remove(userId);
- webSocketMap.put(userId, this);
- //加入set中
- } else {
- webSocketMap.put(userId, this);
- }
- logger.info("websocket connected, userId={}", userId);
-
- try {
- sendMessage(userId, "连接成功");
- } catch (IOException e) {
- logger.warn("websocket 握手回应发送失败, userId={}", userId, e);
- }
- }
-
- /**
- * 连接关闭调用的方法
- */
- @OnClose
- public void onClose() {
- if (webSocketMap.containsKey(userId)) {
- webSocketMap.remove(userId);
- logger.info("websocket closed, userId={}", userId);
- }
- }
-
-
- /**
- * 实现服务器主动推送
- *
- * 两条 AI 流式链路（创意助手、活动策划大纲）**每个 SSE 事件**都会走到这里：
- * 实测一次大纲生成 3062 个事件，即单次请求就会调 3000+ 次。
- * 因此这里**一行日志都不能打**——原先的两行 {@code System.out.println}
- * （{@code userId} 与 {@code session1}）在控制台刷出上万行「userId:833」，
- * 既刷屏、又绕过日志级别、还会淹没真正的告警。
- *
- * 连接与断开另见 {@link #onOpen(Session, String)} / {@link #onClose()}，那里是低频的。
- */
- public void sendMessage(String userId, String message) throws IOException {
- if (webSocketMap.containsKey(userId)) {
- Session session1 = webSocketMap.get(userId).session;
- if (session1 != null) {
- session1.getBasicRemote().sendText(message);
- }
- }
- }
-
-
- /**
- * 实现服务器主动推送
- */
- public void sendAllMessage(String message) throws IOException {
- ConcurrentHashMap.KeySetView 存 {@link Session} 而不是端点实例，是因为同一用户可能从两条路径连进来：
+ * 老版本小程序连 {@code /chat/{userId}}，新版本连 {@code /api/chat/{userId}}
+ * （见 {@link ApiChatWebSocketServer}）。两个端点类共用这张表，推送才不会漏人。
+ */
+ private static final ConcurrentHashMap 这里才是值得打日志的位置：一次连接一条，低频且信息量大
+ * （排查「大纲/创意生成时端上收不到流式内容」时，第一件事就是确认端到底连没连上）。
+ * 推送路径 {@link #send(String, String)} 相反，单次请求会被调用上千次，必须静默。
+ */
+ static void register(Session session, String userId) {
+ SESSIONS.put(userId, session);
+ try {
+ // 0 表示不超时。容器默认 20s 空闲就断，而 AI 流式一次要 50~80s，
+ // 中间的静默期（上游在思考、还没吐字）正好落在默认值上，表现为「生成到一半没动静了」。
+ session.setMaxIdleTimeout(0L);
+ } catch (RuntimeException e) {
+ // 个别容器/代理自带的 Session 实现不支持设置超时，不影响建连
+ logger.debug("websocket 空闲超时设置失败, userId={}", userId, e);
+ }
+ logger.info("websocket connected, userId={}", userId);
+
+ try {
+ send(userId, "连接成功");
+ } catch (IOException e) {
+ logger.warn("websocket 握手回应发送失败, userId={}", userId, e);
+ }
+ }
+
+ /**
+ * 连接关闭的公共处理。按「userId + 会话」精确移除：同一用户重连后，
+ * 旧连接的 close 回调不能把新连接从表里删掉。
+ */
+ static void unregister(String userId, Session session) {
+ if (SESSIONS.remove(userId, session)) {
+ logger.info("websocket closed, userId={}", userId);
+ }
+ }
+
+ /**
+ * 连接建立成功调用的方法
+ */
+ @OnOpen
+ public void onOpen(Session session, @PathParam("userId") String userId) {
+ this.session = session;
+ this.userId = userId;
+ register(session, userId);
+ }
+
+ /**
+ * 连接关闭调用的方法
+ */
+ @OnClose
+ public void onClose() {
+ if (!userId.isEmpty()) {
+ unregister(userId, session);
+ }
+ }
+
+ /**
+ * 客户端消息入口，目前只用于心跳。
+ *
+ * 端上每 25s 发一个 {@code ping} 防止中间代理按空闲断开。必须存在这个
+ * {@code @OnMessage}：没有它，容器收到文本帧会按「不支持的消息」直接关连接。
+ * 业务只做服务端推送，收到内容一律丢弃。
+ */
+ @OnMessage
+ public void onMessage(String message, Session session) {
+ // 心跳与任何客户端文本都不参与业务
+ }
+
+ /**
+ * 实现服务器主动推送
+ *
+ * 两条 AI 流式链路（创意助手、活动策划大纲）**每个 SSE 事件**都会走到这里：
+ * 实测一次大纲生成 3062 个事件，即单次请求就会调 3000+ 次。
+ * 因此这里**一行日志都不能打**——原先的两行 {@code System.out.println}
+ * （{@code userId} 与 {@code session1}）在控制台刷出上万行「userId:833」，
+ * 既刷屏、又绕过日志级别、还会淹没真正的告警。
+ *
+ * 连接与断开另见 {@link #register(Session, String)} / {@link #unregister(String, Session)}，那里是低频的。
+ */
+ public void sendMessage(String userId, String message) throws IOException {
+ send(userId, message);
+ }
+
+ /**
+ * 向指定用户推送一条文本；用户不在线时静默丢弃。
+ *
+ * 优先走 WebSocket；没有截活会话时退到 HTTP 流式通道
+ * （{@link ChatStreamRegistry}）。生产网关在 {@code /api} 前丢 Upgrade 头，
+ * WS 握手根本进不来，端上因此一直在用 HTTP 流那条路。
+ */
+ static void send(String userId, String message) throws IOException {
+ Session target = SESSIONS.get(userId);
+ if (target != null) {
+ target.getBasicRemote().sendText(message);
+ return;
+ }
+ ChatStreamRegistry.emit(userId, message);
+ }
+
+
+ /**
+ * 实现服务器主动推送
+ */
+ public void sendAllMessage(String message) throws IOException {
+ for (Map.Entry 因此这张登记表的几个行为是硬约束：没人接时静默、有人接时逐字转发、
+ * 重复登记后者覆盖前者、写失败要把死掉的 sink 摘掉
+ * （否则后续消息会一直往已经断开的长连接里写）。
+ *
+ * 每个用例用独立 userId，避免静态表互相串味，也就不需要清理。
+ */
+class ChatStreamRegistryTest {
+
+ private static String newUser() {
+ return "test-" + UUID.randomUUID();
+ }
+
+ @Test
+ void 没有登记通道时返回false() {
+ assertFalse(ChatStreamRegistry.emit(newUser(), "{}"));
+ }
+
+ @Test
+ void 登记后能收到原文() {
+ String user = newUser();
+ List 生产网关只放行 {@code /api} 前缀，所以 {@code /api/chat/{userId}} 必须存在；
+ * {@code /chat/{userId}} 留给还没更新的老版本小程序。改任一边都要同步改另外两处，
+ * 这条用例就是把「悄悄改掉」变成「测试失败」。
+ */
+class WebSocketEndpointPathTest {
+
+ @Test
+ void 根路径端点给老版本小程序用() {
+ ServerEndpoint endpoint = WebSocketServer.class.getAnnotation(ServerEndpoint.class);
+ assertEquals("/chat/{userId}", endpoint.value());
+ }
+
+ @Test
+ void api前缀端点是穿过网关的那条() {
+ ServerEndpoint endpoint = ApiChatWebSocketServer.class.getAnnotation(ServerEndpoint.class);
+ assertEquals("/api/chat/{userId}", endpoint.value());
+ }
+
+ @Test
+ void 心跳消息不参与业务也不抛异常() {
+ Session fake = (Session) Proxy.newProxyInstance(Session.class.getClassLoader(),
+ new Class[]{Session.class}, (proxy, method, args) -> null);
+
+ assertDoesNotThrow(() -> new ApiChatWebSocketServer().onMessage("ping", fake));
+ assertDoesNotThrow(() -> new WebSocketServer().onMessage("ping", fake));
+ }
+}
diff --git a/src/test/java/com/gxwebsoft/ai/controller/AiControllerSseFrameTest.java b/src/test/java/com/gxwebsoft/ai/controller/AiControllerSseFrameTest.java
new file mode 100644
index 0000000..c00dbef
--- /dev/null
+++ b/src/test/java/com/gxwebsoft/ai/controller/AiControllerSseFrameTest.java
@@ -0,0 +1,49 @@
+package com.gxwebsoft.ai.controller;
+
+import com.gxwebsoft.common.core.web.ApiResult;
+import org.junit.jupiter.api.Test;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * SSE 末帧格式的回归测试。
+ *
+ * 出图这类「业务方法自己不推帧、只返回一个结果」的链路，结果靠末帧下发；
+ * 端上按 {@code type === 'final'} 识别。字段名一旦改了而端上没改，
+ * 表现为「流正常结束但页面一直转圈」，不容易在联调时发现，所以钉死格式。
+ */
+public class AiControllerSseFrameTest {
+
+ @Test
+ public void finalFrameMarksTypeAndKeepsResult() {
+ Map