feat: enhance API and session management with Nacos and Redis integration

- Add Nacos registry for service registration and deregistration.
- Implement Redis registry for session management with heartbeat and session claiming.
- Improve completion service with session handling and request validation.
- Enhance WebSocket handling for completion requests with JSON-RPC support.
- Add tests for new registry implementations and completion manager functionalities.
- Refactor existing code for better readability and maintainability.
This commit is contained in:
2026-02-15 17:46:34 +08:00
parent 57afb90bc0
commit 3284ce07c7
22 changed files with 1863 additions and 87 deletions

View File

@@ -15,15 +15,18 @@ type CompletionService interface {
Complete(ctx context.Context, req completion.Request) (completion.Response, error)
}
// SessionStatsProvider 暴露会话统计信息,供就绪探针输出。
type SessionStatsProvider interface {
ActiveSessions() map[string]int
}
// RouteOptions 控制 HTTP/WS 接口的超时与请求体上限。
type RouteOptions struct {
RequestTimeout time.Duration
MaxBodyBytes int64
RequestTimeout time.Duration // 单次补全调用超时时间。
MaxBodyBytes int64 // 请求体最大字节数HTTP/WS 共用)。
}
// RegisterRoutes 注册健康检查、HTTP 补全接口和 WebSocket 补全接口。
func RegisterRoutes(router *gin.Engine, service CompletionService, options ...RouteOptions) {
opts := RouteOptions{
RequestTimeout: 10 * time.Second,
@@ -60,6 +63,7 @@ func RegisterRoutes(router *gin.Engine, service CompletionService, options ...Ro
registerWSRoutes(router, service, opts)
handleCompletion := func(c *gin.Context) {
// 为单次请求限制 body 大小,避免异常大包占满内存。
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, opts.MaxBodyBytes)
var req completion.Request
if err := c.ShouldBindJSON(&req); err != nil {
@@ -68,10 +72,12 @@ func RegisterRoutes(router *gin.Engine, service CompletionService, options ...Ro
}
routeLang := c.Param("language")
// 若 body 未显式给出 language则使用路由参数。
if req.Language == "" {
req.Language = routeLang
}
// 对下游补全调用增加超时保护,防止请求长时间悬挂。
ctx, cancel := context.WithTimeout(c.Request.Context(), opts.RequestTimeout)
defer cancel()

View File

@@ -28,6 +28,7 @@ func (f *fakeCompletionService) Complete(_ context.Context, _ completion.Request
return f.resp, nil
}
// 验证 HTTP 补全接口的成功路径。
func TestRegisterRoutesCompletionSuccess(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
@@ -63,6 +64,7 @@ func TestRegisterRoutesCompletionSuccess(t *testing.T) {
}
}
// 验证非法 JSON 会返回 400。
func TestRegisterRoutesCompletionBadJSON(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
@@ -78,6 +80,7 @@ func TestRegisterRoutesCompletionBadJSON(t *testing.T) {
}
}
// 验证业务校验错误会映射为 400。
func TestRegisterRoutesCompletionValidationError(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
@@ -100,6 +103,7 @@ func TestRegisterRoutesCompletionValidationError(t *testing.T) {
}
}
// 验证未知内部错误会映射为 500。
func TestRegisterRoutesCompletionServerError(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
@@ -121,6 +125,7 @@ func TestRegisterRoutesCompletionServerError(t *testing.T) {
}
}
// 验证 WebSocket 补全协议的基础成功流程。
func TestRegisterRoutesCompletionWebSocketSuccess(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()

View File

@@ -21,6 +21,7 @@ var wsUpgrader = websocket.Upgrader{
},
}
// wsCompletionRequest 是普通 WS 消息的补全请求格式。
type wsCompletionRequest struct {
ID string `json:"id"`
Language string `json:"language,omitempty"`
@@ -31,6 +32,7 @@ type wsCompletionRequest struct {
Character int `json:"character"`
}
// wsCompletionResponse 是普通 WS 消息的补全响应格式。
type wsCompletionResponse struct {
ID string `json:"id"`
Items []completion.Item `json:"items,omitempty"`
@@ -40,6 +42,7 @@ type wsCompletionResponse struct {
Error string `json:"error,omitempty"`
}
// wsRPCRequest/wsRPCResponse 用于兼容 JSON-RPC 2.0 客户端。
type wsRPCRequest struct {
JSONRPC string `json:"jsonrpc"`
ID json.RawMessage `json:"id"`
@@ -54,6 +57,7 @@ type wsRPCResponse struct {
Error any `json:"error,omitempty"`
}
// registerWSRoutes 注册 WebSocket 补全入口(含可选语言路由)。
func registerWSRoutes(router *gin.Engine, service CompletionService, opts RouteOptions) {
handler := func(c *gin.Context, defaultLanguage string) {
conn, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil)
@@ -67,6 +71,7 @@ func registerWSRoutes(router *gin.Engine, service CompletionService, opts RouteO
var writeMu sync.Mutex
for {
// 单连接串行读取消息,写操作通过 writeMu 保证并发安全。
_, payload, err := conn.ReadMessage()
if err != nil {
break
@@ -83,6 +88,7 @@ func registerWSRoutes(router *gin.Engine, service CompletionService, opts RouteO
})
}
// handleWSMessage 先尝试按 JSON-RPC 处理;失败后回退到普通 JSON 协议。
func handleWSMessage(
conn *websocket.Conn,
writeMu *sync.Mutex,
@@ -110,6 +116,7 @@ func handleWSMessage(
processWSCompletion(conn, writeMu, service, req, opts)
}
// tryHandleRPCMessage 处理 JSON-RPC 2.0 请求,返回 true 表示消息已消费。
func tryHandleRPCMessage(
conn *websocket.Conn,
writeMu *sync.Mutex,
@@ -127,6 +134,7 @@ func tryHandleRPCMessage(
}
if rpcReq.Method != "completion/complete" && rpcReq.Method != "completion.complete" {
// 非补全方法按 JSON-RPC 规范返回 method not found。
sendWSRPCResponse(conn, writeMu, wsRPCResponse{
JSONRPC: "2.0",
ID: rpcReq.ID,
@@ -140,6 +148,7 @@ func tryHandleRPCMessage(
var req wsCompletionRequest
if err := json.Unmarshal(rpcReq.Params, &req); err != nil {
// 参数反序列化失败按 invalid params 处理。
sendWSRPCResponse(conn, writeMu, wsRPCResponse{
JSONRPC: "2.0",
ID: rpcReq.ID,
@@ -154,6 +163,7 @@ func tryHandleRPCMessage(
req.Language = defaultLanguage
}
if req.ID == "" {
// 兼容未在 params 提供业务 ID 的客户端。
req.ID = string(rpcReq.ID)
}
@@ -188,6 +198,7 @@ func tryHandleRPCMessage(
return true
}
// processWSCompletion 处理普通 WS 协议下的补全请求。
func processWSCompletion(
conn *websocket.Conn,
writeMu *sync.Mutex,
@@ -220,6 +231,7 @@ func processWSCompletion(
default:
var ownedErr *completion.ErrSessionOwnedByOtherInstance
if errors.As(err, &ownedErr) {
// 会话在其他实例上时返回路由提示,客户端可重连对应节点。
msg = err.Error()
routeTo = ownedErr.OwnerEndpoint
ownerID = ownedErr.OwnerID
@@ -241,12 +253,14 @@ func processWSCompletion(
})
}
// sendWSResponse 统一串行写回普通 WS 响应。
func sendWSResponse(conn *websocket.Conn, writeMu *sync.Mutex, resp wsCompletionResponse) {
writeMu.Lock()
defer writeMu.Unlock()
_ = conn.WriteJSON(resp)
}
// sendWSRPCResponse 统一串行写回 JSON-RPC 响应。
func sendWSRPCResponse(conn *websocket.Conn, writeMu *sync.Mutex, resp wsRPCResponse) {
writeMu.Lock()
defer writeMu.Unlock()