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 编译不过。
This commit is contained in:
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(); ServerEndpointExporter exporter = new ServerEndpointExporter();
// 手动注册 WebSocket 端点 // 手动注册 WebSocket 端点。两条路径共用同一份会话表:
exporter.setAnnotatedEndpointClasses(WebSocketServer.class); // /chat 是老版本小程序在用的,/api/chat 用来穿过只放行 /api 前缀的网关。
exporter.setAnnotatedEndpointClasses(WebSocketServer.class, ApiChatWebSocketServer.class);
return exporter; return exporter;
@@ -1,106 +1,150 @@
package com.gxwebsoft.ai.config; package com.gxwebsoft.ai.config;
import org.slf4j.Logger; import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component; import org.springframework.stereotype.Component;
import javax.websocket.OnClose; import javax.websocket.OnClose;
import javax.websocket.OnOpen; import javax.websocket.OnMessage;
import javax.websocket.Session; import javax.websocket.OnOpen;
import javax.websocket.server.PathParam; import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint; import javax.websocket.server.PathParam;
import java.io.IOException; import javax.websocket.server.ServerEndpoint;
import java.util.concurrent.ConcurrentHashMap; import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
@ServerEndpoint(value = "/chat/{userId}")
@Component
public class WebSocketServer { @ServerEndpoint(value = "/chat/{userId}")
@Component
private static final Logger logger = LoggerFactory.getLogger(WebSocketServer.class); public class WebSocketServer {
/** private static final Logger logger = LoggerFactory.getLogger(WebSocketServer.class);
* concurrent包的线程安全Set,用来存放每个客户端对应的MyWebSocket对象。
*/ /**
private static ConcurrentHashMap<String, WebSocketServer> webSocketMap = new ConcurrentHashMap<>(); * 每个用户的在线会话:键是 userId,值是这条连接的 {@link Session}。
/** *
* 与某个客户端的连接会话,需要通过它来给客户端发送数据 * <p>存 {@link Session} 而不是端点实例,是因为同一用户可能从两条路径连进来:
*/ * 老版本小程序连 {@code /chat/{userId}},新版本连 {@code /api/chat/{userId}}
private Session session; * (见 {@link ApiChatWebSocketServer})。两个端点类共用这张表,推送才不会漏人。
/** */
* 接收userId private static final ConcurrentHashMap<String, Session> SESSIONS = new ConcurrentHashMap<>();
*/ /**
private String userId = ""; * 与某个客户端的连接会话,需要通过它来给客户端发送数据
*/
/** private Session session;
* 连接建立成功调用的方法 /**
* * 接收userId
* <p>这里才是值得打日志的位置:一次连接一条,低频且信息量大 */
* (排查「大纲/创意生成时端上收不到流式内容」时,第一件事就是确认端到底连没连上)。 private String userId = "";
* 推送路径 {@link #sendMessage(String, String)} 相反,单次请求会被调用上千次,必须静默。
*/ /**
@OnOpen * 连接建立成功的公共处理:登记会话、关掉空闲超时、回一条「连接成功」。
public void onOpen(Session session, @PathParam("userId") String userId) { *
this.session = session; * <p>这里才是值得打日志的位置:一次连接一条,低频且信息量大
this.userId = userId; * (排查「大纲/创意生成时端上收不到流式内容」时,第一件事就是确认端到底连没连上)。
if (webSocketMap.containsKey(userId)) { * 推送路径 {@link #send(String, String)} 相反,单次请求会被调用上千次,必须静默。
webSocketMap.remove(userId); */
webSocketMap.put(userId, this); static void register(Session session, String userId) {
//加入set中 SESSIONS.put(userId, session);
} else { try {
webSocketMap.put(userId, this); // 0 表示不超时。容器默认 20s 空闲就断,而 AI 流式一次要 50~80s,
} // 中间的静默期(上游在思考、还没吐字)正好落在默认值上,表现为「生成到一半没动静了」。
logger.info("websocket connected, userId={}", userId); session.setMaxIdleTimeout(0L);
} catch (RuntimeException e) {
try { // 个别容器/代理自带的 Session 实现不支持设置超时,不影响建连
sendMessage(userId, "连接成功"); logger.debug("websocket 空闲超时设置失败, userId={}", userId, e);
} catch (IOException e) { }
logger.warn("websocket 握手回应发送失败, userId={}", userId, e); logger.info("websocket connected, userId={}", userId);
}
} try {
send(userId, "连接成功");
/** } catch (IOException e) {
* 连接关闭调用的方法 logger.warn("websocket 握手回应发送失败, userId={}", userId, e);
*/ }
@OnClose }
public void onClose() {
if (webSocketMap.containsKey(userId)) { /**
webSocketMap.remove(userId); * 连接关闭的公共处理。按「userId + 会话」精确移除:同一用户重连后,
logger.info("websocket closed, userId={}", userId); * 旧连接的 close 回调不能把新连接从表里删掉。
} */
} static void unregister(String userId, Session session) {
if (SESSIONS.remove(userId, session)) {
logger.info("websocket closed, userId={}", userId);
/** }
* 实现服务器主动推送 }
*
* <p>两条 AI 流式链路(创意助手、活动策划大纲)**每个 SSE 事件**都会走到这里: /**
* 实测一次大纲生成 3062 个事件,即单次请求就会调 3000+ 次。 * 连接建立成功调用的方法
* 因此这里**一行日志都不能打**——原先的两行 {@code System.out.println} */
* ({@code userId} 与 {@code session1})在控制台刷出上万行「userId:833」, @OnOpen
* 既刷屏、又绕过日志级别、还会淹没真正的告警。 public void onOpen(Session session, @PathParam("userId") String userId) {
* this.session = session;
* <p>连接与断开另见 {@link #onOpen(Session, String)} / {@link #onClose()},那里是低频的。 this.userId = userId;
*/ register(session, userId);
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); */
} @OnClose
} public void onClose() {
} if (!userId.isEmpty()) {
unregister(userId, session);
}
/** }
* 实现服务器主动推送
*/ /**
public void sendAllMessage(String message) throws IOException { * 客户端消息入口,目前只用于心跳。
ConcurrentHashMap.KeySetView<String, WebSocketServer> userIds = webSocketMap.keySet(); *
for (String userId : userIds) { * <p>端上每 25s 发一个 {@code ping} 防止中间代理按空闲断开。必须存在这个
WebSocketServer webSocketServer = webSocketMap.get(userId); * {@code @OnMessage}:没有它,容器收到文本帧会按「不支持的消息」直接关连接。
webSocketServer.session.getBasicRemote().sendText(message); * 业务只做服务端推送,收到内容一律丢弃。
} */
} @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);
}
}
}
+3
View File
@@ -26,6 +26,9 @@ logging:
# 开放平台同步过程保留日志:一次同步动辄一两分钟, # 开放平台同步过程保留日志:一次同步动辄一两分钟,
# 若沿用 com.gxwebsoft 的 ERROR 级,进度与失败原因在生产完全不可见 # 若沿用 com.gxwebsoft 的 ERROR 级,进度与失败原因在生产完全不可见
com.gxwebsoft.gxmu.openplat: INFO com.gxwebsoft.gxmu.openplat: INFO
# WebSocket 只保留连接/断开(每次握手一条,低频)。推送路径每个 SSE 事件都会
# 走一遍 sendMessage,那里一行日志都不打,所以这里放开 INFO 不会刷屏
com.gxwebsoft.ai.config: INFO
# 框架配置 # 框架配置
config: config:
@@ -0,0 +1,100 @@
package com.gxwebsoft.ai.config;
import org.junit.jupiter.api.Test;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.function.Consumer;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
/**
* HTTP 流式通道(SSE)是生产上唯一能用的推送路径:网关在 {@code /api} 前丢掉
* WebSocket 的 {@code Upgrade} 头,WS 握手永远到不了 Tomcat。
*
* <p>因此这张登记表的几个行为是硬约束:没人接时静默、有人接时逐字转发、
* 重复登记后者覆盖前者、写失败要把死掉的 sink 摘掉
* (否则后续消息会一直往已经断开的长连接里写)。
*
* <p>每个用例用独立 userId,避免静态表互相串味,也就不需要清理。
*/
class ChatStreamRegistryTest {
private static String newUser() {
return "test-" + UUID.randomUUID();
}
@Test
void 没有登记通道时返回false() {
assertFalse(ChatStreamRegistry.emit(newUser(), "{}"));
}
@Test
void 登记后能收到原文() {
String user = newUser();
List<String> received = new ArrayList<>();
ChatStreamRegistry.bind(user, received::add);
assertTrue(ChatStreamRegistry.emit(user, "{\"answer\":\"你好\"}"));
assertEquals(Collections.singletonList("{\"answer\":\"你好\"}"), received);
}
@Test
void 重复登记时后者覆盖前者() {
String user = newUser();
List<String> first = new ArrayList<>();
List<String> second = new ArrayList<>();
ChatStreamRegistry.bind(user, first::add);
ChatStreamRegistry.bind(user, second::add);
ChatStreamRegistry.emit(user, "x");
assertTrue(first.isEmpty(), "被覆盖的旧连接不该再收到消息");
assertEquals(Collections.singletonList("x"), second);
}
@Test
void 解绑只摘自己的那个通道() {
String user = newUser();
List<String> first = new ArrayList<>();
List<String> second = new ArrayList<>();
Consumer<String> firstSink = first::add;
ChatStreamRegistry.bind(user, firstSink);
ChatStreamRegistry.bind(user, second::add);
// 旧连接晚到的 finally 块:不能把新连接顶掉
ChatStreamRegistry.unbind(user, firstSink);
assertTrue(ChatStreamRegistry.emit(user, "still-alive"));
assertEquals(Collections.singletonList("still-alive"), second);
}
@Test
void 写入失败后摘掉通道且不向上抛() {
String user = newUser();
ChatStreamRegistry.bind(user, payload -> {
throw new IllegalStateException("客户端已断开");
});
assertFalse(ChatStreamRegistry.emit(user, "boom"));
// 失败过的 sink 已被摘掉,下一次 emit 直接返回 false
assertFalse(ChatStreamRegistry.emit(user, "boom"));
}
@Test
void 没有WebSocket会话时推送落到流式通道() throws IOException {
String user = newUser();
List<String> received = new ArrayList<>();
ChatStreamRegistry.bind(user, received::add);
// 生产就是这样:SESSIONS 里没有这个用户,send 必须改投 SSE
WebSocketServer.send(user, "{\"answer\":\"fallback\"}");
assertEquals(Collections.singletonList("{\"answer\":\"fallback\"}"), received);
}
}
@@ -0,0 +1,42 @@
package com.gxwebsoft.ai.config;
import org.junit.jupiter.api.Test;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;
import java.lang.reflect.Proxy;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
/**
* 端点路径是**跨端契约**:前端 {@code utils/chat-socket.js} 的 {@code buildWebSocketUrl}
* 拼的就是这两个值,反代(nginx)也按它们配 location。
*
* <p>生产网关只放行 {@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));
}
}
@@ -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 末帧格式的回归测试。
*
* <p>出图这类「业务方法自己不推帧、只返回一个结果」的链路,结果靠末帧下发;
* 端上按 {@code type === 'final'} 识别。字段名一旦改了而端上没改,
* 表现为「流正常结束但页面一直转圈」,不容易在联调时发现,所以钉死格式。
*/
public class AiControllerSseFrameTest {
@Test
public void finalFrameMarksTypeAndKeepsResult() {
Map<String, Object> data = new LinkedHashMap<>();
data.put("photo_files", "[\"https://example.com/a.jpeg\"]");
String json = AiController.finalFrameJson(new ApiResult<>(0, "ok", data));
assertTrue(json.contains("\"type\":\"final\""), "末帧必须带 type=final,端上靠它区分业务帧");
assertTrue(json.contains("\"result\":{"), "业务返回值要原样放在 result 里");
assertTrue(json.contains("photo_files"), "返回体里的字段不能丢");
assertEquals(1, countOccurrences(json, "\"type\":"), "末帧只能有一个 type 字段");
}
@Test
public void finalFrameKeepsStringResult() {
String json = AiController.finalFrameJson("plain-text");
assertTrue(json.contains("\"result\":\"plain-text\""));
}
private static int countOccurrences(String text, String needle) {
int count = 0;
int index = text.indexOf(needle);
while (index != -1) {
count++;
index = text.indexOf(needle, index + needle.length());
}
return count;
}
}