220 lines
6.5 KiB
Go
220 lines
6.5 KiB
Go
// Package server 神识 HTTP 服务：数据面（IR 入站 → SSE 出站）、健康面与用量面。
|
||
//
|
||
// 故障切换红线：首帧到达前不提交 SSE 响应——失败可整体换道并返回 JSON 错误；
|
||
// 已写出字节后不再换道，如实以 error 终帧收尾（shenshi-gateway.md §4/§5）。
|
||
package server
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"net/http"
|
||
"time"
|
||
|
||
"hunjing/shenshi/internal/ir"
|
||
"hunjing/shenshi/internal/metering"
|
||
"hunjing/shenshi/internal/relay"
|
||
"hunjing/shenshi/internal/route"
|
||
"hunjing/shenshi/internal/transform"
|
||
)
|
||
|
||
// Config 服务配置。
|
||
type Config struct {
|
||
ServiceToken string
|
||
}
|
||
|
||
// Server 神识 HTTP 服务。
|
||
type Server struct {
|
||
cfg Config
|
||
routes *route.Registry
|
||
usage *metering.Log
|
||
mux *http.ServeMux
|
||
}
|
||
|
||
// New 构建服务并注册路由。
|
||
func New(cfg Config, routes *route.Registry) *Server {
|
||
s := &Server{cfg: cfg, routes: routes, usage: metering.New(10_000), mux: http.NewServeMux()}
|
||
s.mux.HandleFunc("GET /healthz", s.handleHealth)
|
||
s.mux.HandleFunc("GET /readyz", s.handleHealth)
|
||
s.mux.Handle("POST /v1/chat/completions", s.auth(http.HandlerFunc(s.handleChat)))
|
||
s.mux.Handle("GET /me/usage", s.auth(http.HandlerFunc(s.handleUsage)))
|
||
return s
|
||
}
|
||
|
||
// ServeHTTP 实现 http.Handler。
|
||
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { s.mux.ServeHTTP(w, r) }
|
||
|
||
// auth 服务令牌校验中间件；未配置令牌时放行（本地开发档）。
|
||
func (s *Server) auth(next http.Handler) http.Handler {
|
||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
if s.cfg.ServiceToken == "" {
|
||
next.ServeHTTP(w, r)
|
||
return
|
||
}
|
||
if r.Header.Get("Authorization") != "Bearer "+s.cfg.ServiceToken {
|
||
writeError(w, http.StatusUnauthorized, "HJ-2003", "服务令牌无效")
|
||
return
|
||
}
|
||
next.ServeHTTP(w, r)
|
||
})
|
||
}
|
||
|
||
// handleChat 数据面：IR → 候选迭代（首字超时/冷却标记/粘性）→ SSE。
|
||
func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) {
|
||
var req ir.ChatRequest
|
||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||
writeError(w, http.StatusBadRequest, "HJ-2002", "IR 请求解析失败: "+err.Error())
|
||
return
|
||
}
|
||
if err := req.Validate(); err != nil {
|
||
writeError(w, http.StatusBadRequest, "HJ-2002", err.Error())
|
||
return
|
||
}
|
||
candidates, err := s.routes.Candidates(req.Slot, req.RunID)
|
||
if err != nil {
|
||
writeError(w, http.StatusServiceUnavailable, "HJ-2001", err.Error())
|
||
return
|
||
}
|
||
|
||
var lastErr error
|
||
for _, endpoint := range candidates {
|
||
res := s.tryEndpoint(w, r, endpoint, &req)
|
||
s.usage.Record(res.Attempt)
|
||
if res.Committed {
|
||
// 已向客户端写出：无论成败都不能再换道
|
||
if res.Status == "ok" {
|
||
s.routes.Sticky.Remember(req.RunID, endpoint.Name)
|
||
}
|
||
return
|
||
}
|
||
// 未提交：分类冷却后换下一候选
|
||
if res.Upstream != nil {
|
||
if res.Upstream.StatusCode == 401 || res.Upstream.StatusCode == 403 {
|
||
s.routes.Cooldowns.Mark(route.CredentialCooldownKey(endpoint), true)
|
||
} else {
|
||
s.routes.Cooldowns.Mark(route.BindingCooldownKey(endpoint), false)
|
||
}
|
||
} else {
|
||
s.routes.Cooldowns.Mark(route.BindingCooldownKey(endpoint), false)
|
||
}
|
||
lastErr = res.Err
|
||
}
|
||
writeError(w, http.StatusServiceUnavailable, "HJ-2001", "全部候选端点失败："+errString(lastErr))
|
||
}
|
||
|
||
type attempt struct {
|
||
metering.Attempt
|
||
Committed bool // 是否已向客户端写出 SSE
|
||
Upstream *transform.UpstreamError // 上游分类（可空）
|
||
Err error // 原始错误
|
||
}
|
||
|
||
func errString(err error) string {
|
||
if err == nil {
|
||
return "未知错误"
|
||
}
|
||
return err.Error()
|
||
}
|
||
|
||
// tryEndpoint 单端点尝试。
|
||
func (s *Server) tryEndpoint(w http.ResponseWriter, r *http.Request, endpoint route.Endpoint, req *ir.ChatRequest) attempt {
|
||
start := time.Now()
|
||
ctx, cancel := context.WithTimeout(r.Context(), s.routes.FirstByteTimeout)
|
||
defer cancel()
|
||
|
||
deltas := make(chan ir.Delta, 16)
|
||
errCh := make(chan error, 1)
|
||
go func() {
|
||
defer close(deltas)
|
||
errCh <- endpoint.Adaptor.Stream(ctx, req, deltas)
|
||
}()
|
||
|
||
res := attempt{Attempt: metering.Attempt{
|
||
RunID: req.RunID, CharacterID: req.CharacterID,
|
||
Model: req.Model, Slot: req.Slot, Endpoint: endpoint.Name,
|
||
}}
|
||
|
||
// 阶段一：等首帧或失败——此阶段未提交响应，可整体回退
|
||
var first ir.Delta
|
||
select {
|
||
case d, ok := <-deltas:
|
||
if !ok {
|
||
res.Err = <-errCh
|
||
res.Upstream, _ = res.Err.(*transform.UpstreamError)
|
||
res.Status, res.TotalMS = failStatus(res.Err), time.Since(start).Milliseconds()
|
||
return res
|
||
}
|
||
first = d
|
||
case <-ctx.Done():
|
||
res.Status, res.Err = "first_byte_timeout", ctx.Err()
|
||
res.TotalMS = time.Since(start).Milliseconds()
|
||
return res
|
||
}
|
||
|
||
// 阶段二：首帧已到——提交 SSE，流式转发
|
||
relay.SSEHeaders(w)
|
||
w.WriteHeader(http.StatusOK)
|
||
res.Committed = true
|
||
res.FirstTokenMS = time.Since(start).Milliseconds()
|
||
_ = relay.WriteDelta(w, first)
|
||
|
||
for d := range deltas {
|
||
if d.Delta != "" || d.FinishReason != "" {
|
||
if err := relay.WriteDelta(w, d); err != nil {
|
||
break // 客户端断开
|
||
}
|
||
}
|
||
if d.FinishReason != "" {
|
||
res.Attempt.PromptTokens = d.UsagePromptTokens
|
||
res.Attempt.CompletionTokens = d.UsageCompletionTokens
|
||
}
|
||
}
|
||
err := <-errCh
|
||
res.TotalMS = time.Since(start).Milliseconds()
|
||
if err == nil {
|
||
_ = relay.WriteDone(w)
|
||
res.Status = "ok"
|
||
return res
|
||
}
|
||
// 流中失败：如实以 error 终帧收尾，不伪装完成
|
||
res.Err = err
|
||
res.Upstream, _ = err.(*transform.UpstreamError)
|
||
res.Status = "error"
|
||
res.ErrorCode = "stream_written"
|
||
_ = relay.WriteDelta(w, ir.Delta{FinishReason: "error"})
|
||
_ = relay.WriteDone(w)
|
||
return res
|
||
}
|
||
|
||
func failStatus(err error) string {
|
||
if up, ok := err.(*transform.UpstreamError); ok {
|
||
if up.StatusCode == 0 {
|
||
return "timeout_or_network"
|
||
}
|
||
return "upstream_error"
|
||
}
|
||
return "error"
|
||
}
|
||
|
||
// handleUsage 用量聚合（/me/usage；BYOK 归属维度随 M4b 接入）。
|
||
func (s *Server) handleUsage(w http.ResponseWriter, _ *http.Request) {
|
||
w.Header().Set("Content-Type", "application/json")
|
||
_ = json.NewEncoder(w).Encode(s.usage.Aggregate())
|
||
}
|
||
|
||
// handleHealth 健康面。
|
||
func (s *Server) handleHealth(w http.ResponseWriter, _ *http.Request) {
|
||
w.Header().Set("Content-Type", "application/json")
|
||
w.WriteHeader(http.StatusOK)
|
||
_, _ = w.Write([]byte(`{"status":"ok"}`))
|
||
}
|
||
|
||
// writeError 统一错误结构（API 契约 §5）。
|
||
func writeError(w http.ResponseWriter, status int, code, message string) {
|
||
w.Header().Set("Content-Type", "application/json")
|
||
w.WriteHeader(status)
|
||
body := map[string]any{
|
||
"error": map[string]string{"code": code, "message": message},
|
||
}
|
||
_ = json.NewEncoder(w).Encode(body)
|
||
}
|