feat(ai): 聊天 WebSocket 加 /api 端点与心跳，WS 不可用时退到 HTTP SSE

生产网关在 /api 前丢掉 Upgrade 头，wss://…/chat/{userId} 到不了后端，端上一直卡在「正在建立 WebSocket 连接…」。 - 新增 ApiChatWebSocketServer（/api/chat/{userId}），与老端点共用会话表；会话表改存 Session 而不是端点实例，两条路径都能收到推送 - 登记会话时关闭空闲超时（Tomcat 默认 20s，AI 首字慢时会被掐断），并补 @OnMessage 丢弃心跳（缺了它容器会因不支持的文本帧直接断连） - send() 找不到 WS 会话时改投 ChatStreamRegistry，把帧送到 HTTP SSE 连接，推送侧仍只有 sendMessage 一个入口，三千多次调用点不用改 - prod 放开 com.gxwebsoft.ai.config 到 INFO：只保留连接/断开，推送路径仍静默 - 补 ChatStreamRegistryTest / WebSocketEndpointPathTest / AiControllerSseFrameTest 顺带补上 31eb75c 漏加的 ChatStreamRegistry.java：AiController 已经 import 它，缺文件时 HEAD 编译不过。
2026-10-21 19:18:02 +08:00
parent 31eb75cdad
commit bf6813e71d
8 changed files with 469 additions and 108 deletions
@@ -0,0 +1,54 @@
+package com.gxwebsoft.ai.config;
+
+import javax.websocket.OnClose;
+import javax.websocket.OnMessage;
+import javax.websocket.OnOpen;
+import javax.websocket.Session;
+import javax.websocket.server.PathParam;
+import javax.websocket.server.ServerEndpoint;
+
+/**
+ * {@code /api} 前缀下的聊天端点，与 {@link WebSocketServer}（{@code /chat/{userId}}）等价。
+ *
+ * <p>为什么要有两条路径：
+ * <ul>
+ *   <li>生产网关只放行 {@code /api} 前缀（见 release/BUILD-INFO 第 7 条：
+ *       探测 {@code /chat/{userId}} 拿到 404），新客户端必须走 {@code /api/chat/{userId}}
+ *       才能穿过网关；</li>
+ *   <li>已经装在用户手机上的老版本小程序仍然连 {@code /chat/{userId}}，
+ *       反代侧补上 Upgrade 之后它们可以立刻恢复，不必等发版审核。</li>
+ * </ul>
+ *
+ * <p>两条路径共用 {@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) {
+        // 心跳与任何客户端文本都不参与业务
+    }
+}
@@ -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 连不上时的退路。
+ *
+ * <p>生产网关在 {@code /api} 前面就把 WebSocket 的 {@code Upgrade} 头丢了，
+ * 握手永远到不了 Tomcat（网关侧不能改），所以端上加了一条 HTTP 流式通道：
+ *
+ * <pre>
+ * POST /api/ai/chat/message/stream   ← 请求体与 WS 版完全一样
+ * data: {"data":{...}}               ← 每个 SSE 事件一行，内容与 WS 帧逐字节相同
+ * </pre>
+ *
+ * <p>推送侧仍然只认 {@code webSocketServer.sendMessage(userId, payload)} 一个入口
+ * （两条 AI 链路 3000+ 次调用的都是它），这里通过「同一用户登记一个 sink」把消息
+ * 改投给 HTTP 连接，业务代码一行不用改。
+ *
+ * <p>键用 userId 而不是连接 id：WS 那边也是按用户名建索引，语义保持一致，
+ * 同一时刻一个用户只跑一条 AI 流。
+ */
+public final class ChatStreamRegistry {
+
+    private static final Map<String, Consumer<String>> SINKS = new ConcurrentHashMap<>();
+
+    private ChatStreamRegistry() {
+    }
+
+    /** 登记某个用户的 HTTP 流通道；重复登记时后者覆盖前者（同一用户同时只该有一条流）。 */
+    public static void bind(String userId, Consumer<String> sink) {
+        if (userId == null || sink == null) {
+            return;
+        }
+        SINKS.put(userId, sink);
+    }
+
+    /** 解绑。只移除自己登记的那个 sink，避免把后来者的通道误删。 */
+    public static void unbind(String userId, Consumer<String> sink) {
+        if (userId == null || sink == null) {
+            return;
+        }
+        SINKS.remove(userId, sink);
+    }
+
+    /**
+     * 把一条消息投给该用户的 HTTP 流。
+     *
+     * @return 有没有人接（没人接返回 {@code false}，调用方可以继续走 WebSocket）
+     */
+    public static boolean emit(String userId, String payload) {
+        Consumer<String> sink = userId == null ? null : SINKS.get(userId);
+        if (sink == null) {
+            return false;
+        }
+        try {
+            sink.accept(payload);
+        } catch (RuntimeException e) {
+            // 端上断开、代理超时都会走到这里。推送是尽力而为：一次写失败不该
+            // 把生成流程整条打断（这正是当初 sendMessage 吞异常的理由）。
+            SINKS.remove(userId, sink);
+            return false;
+        }
+        return true;
+    }
+}
@@ -13,8 +13,9 @@ public class WebSocketConfig
    {
        ServerEndpointExporter exporter = new ServerEndpointExporter();

-        // 手动注册 WebSocket 端点
-        exporter.setAnnotatedEndpointClasses(WebSocketServer.class);
+        // 手动注册 WebSocket 端点。两条路径共用同一份会话表：
+        // /chat 是老版本小程序在用的，/api/chat 用来穿过只放行 /api 前缀的网关。
+        exporter.setAnnotatedEndpointClasses(WebSocketServer.class, ApiChatWebSocketServer.class);


        return exporter;
@@ -1,106 +1,150 @@
-package com.gxwebsoft.ai.config;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import org.springframework.stereotype.Component;
-
-import javax.websocket.OnClose;
-import javax.websocket.OnOpen;
-import javax.websocket.Session;
-import javax.websocket.server.PathParam;
-import javax.websocket.server.ServerEndpoint;
-import java.io.IOException;
-import java.util.concurrent.ConcurrentHashMap;
-
-
-@ServerEndpoint(value = "/chat/{userId}")
-@Component
-public class WebSocketServer {
-
-    private static final Logger logger = LoggerFactory.getLogger(WebSocketServer.class);
-
-    /**
-     * concurrent包的线程安全Set，用来存放每个客户端对应的MyWebSocket对象。
-     */
-    private static ConcurrentHashMap<String, WebSocketServer> webSocketMap = new ConcurrentHashMap<>();
-    /**
-     * 与某个客户端的连接会话，需要通过它来给客户端发送数据
-     */
-    private Session session;
-    /**
-     * 接收userId
-     */
-    private String userId = "";
-
-    /**
-     * 连接建立成功调用的方法
-     *
-     * <p>这里才是值得打日志的位置：一次连接一条，低频且信息量大
-     * （排查「大纲/创意生成时端上收不到流式内容」时，第一件事就是确认端到底连没连上）。
-     * 推送路径 {@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);
-        }
-    }
-
-
-    /**
-     * 实现服务器主动推送
-     *
-     * <p>两条 AI 流式链路（创意助手、活动策划大纲）**每个 SSE 事件**都会走到这里：
-     * 实测一次大纲生成 3062 个事件，即单次请求就会调 3000+ 次。
-     * 因此这里**一行日志都不能打**——原先的两行 {@code System.out.println}
-     * （{@code userId} 与 {@code session1}）在控制台刷出上万行「userId:833」，
-     * 既刷屏、又绕过日志级别、还会淹没真正的告警。
-     *
-     * <p>连接与断开另见 {@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<String, WebSocketServer> userIds = webSocketMap.keySet();
-        for (String userId : userIds) {
-            WebSocketServer webSocketServer = webSocketMap.get(userId);
-            webSocketServer.session.getBasicRemote().sendText(message);
-        }
-    }
-
-}
+package com.gxwebsoft.ai.config;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.stereotype.Component;
+
+import javax.websocket.OnClose;
+import javax.websocket.OnMessage;
+import javax.websocket.OnOpen;
+import javax.websocket.Session;
+import javax.websocket.server.PathParam;
+import javax.websocket.server.ServerEndpoint;
+import java.io.IOException;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+
+@ServerEndpoint(value = "/chat/{userId}")
+@Component
+public class WebSocketServer {
+
+    private static final Logger logger = LoggerFactory.getLogger(WebSocketServer.class);
+
+    /**
+     * 每个用户的在线会话：键是 userId，值是这条连接的 {@link Session}。
+     *
+     * <p>存 {@link Session} 而不是端点实例，是因为同一用户可能从两条路径连进来：
+     * 老版本小程序连 {@code /chat/{userId}}，新版本连 {@code /api/chat/{userId}}
+     * （见 {@link ApiChatWebSocketServer}）。两个端点类共用这张表，推送才不会漏人。
+     */
+    private static final ConcurrentHashMap<String, Session> SESSIONS = new ConcurrentHashMap<>();
+    /**
+     * 与某个客户端的连接会话，需要通过它来给客户端发送数据
+     */
+    private Session session;
+    /**
+     * 接收userId
+     */
+    private String userId = "";
+
+    /**
+     * 连接建立成功的公共处理：登记会话、关掉空闲超时、回一条「连接成功」。
+     *
+     * <p>这里才是值得打日志的位置：一次连接一条，低频且信息量大
+     * （排查「大纲/创意生成时端上收不到流式内容」时，第一件事就是确认端到底连没连上）。
+     * 推送路径 {@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);
+        }
+    }
+
+    /**
+     * 客户端消息入口，目前只用于心跳。
+     *
+     * <p>端上每 25s 发一个 {@code ping} 防止中间代理按空闲断开。必须存在这个
+     * {@code @OnMessage}：没有它，容器收到文本帧会按「不支持的消息」直接关连接。
+     * 业务只做服务端推送，收到内容一律丢弃。
+     */
+    @OnMessage
+    public void onMessage(String message, Session session) {
+        // 心跳与任何客户端文本都不参与业务
+    }
+
+    /**
+     * 实现服务器主动推送
+     *
+     * <p>两条 AI 流式链路（创意助手、活动策划大纲）**每个 SSE 事件**都会走到这里：
+     * 实测一次大纲生成 3062 个事件，即单次请求就会调 3000+ 次。
+     * 因此这里**一行日志都不能打**——原先的两行 {@code System.out.println}
+     * （{@code userId} 与 {@code session1}）在控制台刷出上万行「userId:833」，
+     * 既刷屏、又绕过日志级别、还会淹没真正的告警。
+     *
+     * <p>连接与断开另见 {@link #register(Session, String)} / {@link #unregister(String, Session)}，那里是低频的。
+     */
+    public void sendMessage(String userId, String message) throws IOException {
+        send(userId, message);
+    }
+
+    /**
+     * 向指定用户推送一条文本；用户不在线时静默丢弃。
+     *
+     * <p>优先走 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<String, Session> entry : SESSIONS.entrySet()) {
+            entry.getValue().getBasicRemote().sendText(message);
+        }
+    }
+
+}