chenxin
144415247e
- 神识 Go：store 双实现（PG 自动建表+60s 就绪重试 / 内存回退）承载 usage_logs·byok_keys·quota_balances；BYOK 密钥库（AES-GCM 加密、API 脱敏、撤回即失效、主密钥未配置 fail-closed）；配额三段式（预扣 max_tokens 钳位 1M→结算实际 token→全败幂等退款，HJ-2004）；计量归属（IR user_id→usage_logs，/me/usage 用户维度聚合+配额总览，/admin/usage/sample 对账抽样）；Anthropic/Gemini 出站适配器（SSE 解析、system 提升、用量映射）；stub 端点补名修复计量空 endpoint - IR 契约：user_id/byok_key_id 双端同步（Go types.go + Py types.py） - 灵台：runs.owner_id（迁移 005，server_default 兼容存量）经 kernel 全链透传为 IR user_id；shenshi_client 扩展用户面方法（X-Hunjing-User-Id 代理，Echo 桩带内存实现）；/me/byok-keys 登记·列表·撤回与 /me/usage 端点（4 用例） - 前端神识阁（/models 重建）：偏好（生成参数表单）/自带密钥（登记加密保存、脱敏列表、撤回二次确认）/用量（聚合表+配额进度条）三区；me API 库+mock+5 用例；侧栏恢复入口 - 配置与部署：.env.example 增 SHENSHI_DATABASE_URL/SHENSHI_BYOK_MASTER_KEY/SHENSHI_USER_BUDGET_TOKENS；compose 神识容器接自有库（init-shenshi.sql 预建 hunjing_shenshi，与灵台不共库共表） - 证据：Go 34 用例（BYOK 独占/配额拒绝/结算退款/归属聚合）、后端 532、前端 210、BFF 4 全绿；跨进程 E2E 四场景（灵台 HttpShenshiClient→真 Go 神识：登记脱敏/BYOK 独占不可达整体失败/用量归属聚合/撤回即失效）通过；rewrite-plan M4 标记完成、能力矩阵神识行升已验证
408 lines
13 KiB
Go
408 lines
13 KiB
Go
// Package server 神识 HTTP 服务：数据面（IR 入站 → SSE 出站）、健康面、用量面与 BYOK 面。
|
||
//
|
||
// 故障切换红线：首帧到达前不提交 SSE 响应——失败可整体换道并返回 JSON 错误；
|
||
// 已写出字节后不再换道，如实以 error 终帧收尾（shenshi-gateway.md §4/§5）。
|
||
//
|
||
// 配额三段式：预扣（max_tokens 上限，钳位）→ 结算（实际 token）→ 退款（失败幂等）。
|
||
// BYOK：请求携带 byok_key_id 时独占候选（不回落平台端点，防止用户预期自带密钥却被平台计费）。
|
||
package server
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"net/http"
|
||
"time"
|
||
|
||
"hunjing/shenshi/internal/byok"
|
||
"hunjing/shenshi/internal/ir"
|
||
"hunjing/shenshi/internal/metering"
|
||
"hunjing/shenshi/internal/quota"
|
||
"hunjing/shenshi/internal/relay"
|
||
"hunjing/shenshi/internal/route"
|
||
"hunjing/shenshi/internal/store"
|
||
"hunjing/shenshi/internal/transform"
|
||
)
|
||
|
||
// Config 服务配置。
|
||
type Config struct {
|
||
ServiceToken string
|
||
Store store.Store // nil 时计量退化为进程内环形
|
||
Vault *byok.Vault // nil 时 BYOK 面整体 503（fail-closed）
|
||
Quota *quota.Limiter
|
||
}
|
||
|
||
// 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, cfg.Store),
|
||
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)))
|
||
s.mux.Handle("GET /admin/usage/sample", s.auth(http.HandlerFunc(s.handleUsageSample)))
|
||
s.mux.Handle("GET /me/byok-keys", s.auth(http.HandlerFunc(s.handleByokList)))
|
||
s.mux.Handle("POST /me/byok-keys", s.auth(http.HandlerFunc(s.handleByokRegister)))
|
||
s.mux.Handle("DELETE /me/byok-keys/{id}", s.auth(http.HandlerFunc(s.handleByokRevoke)))
|
||
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)
|
||
})
|
||
}
|
||
|
||
// defaultReserve 预扣默认上限（max_tokens 未指定时按 4K 输出预扣）。
|
||
const defaultReserve = 4096
|
||
|
||
// handleChat 数据面：配额预扣 → 候选迭代（首字超时/冷却标记/粘性/BYOK 独占）→ 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
|
||
}
|
||
|
||
// 三段式之一：预扣（按输出上限；钳位集中在 quota 包）
|
||
reserveTokens := int64(defaultReserve)
|
||
if req.MaxTokens != nil {
|
||
reserveTokens = int64(*req.MaxTokens)
|
||
}
|
||
reserved, err := s.quotaReserve(req.UserID, reserveTokens)
|
||
if err != nil {
|
||
writeError(w, http.StatusPaymentRequired, "HJ-2004", err.Error())
|
||
return
|
||
}
|
||
|
||
candidates, err := s.candidates(&req)
|
||
if err != nil {
|
||
_ = s.quotaRefund(req.UserID, reserved)
|
||
writeError(w, http.StatusServiceUnavailable, "HJ-2001", err.Error())
|
||
return
|
||
}
|
||
|
||
var lastErr error
|
||
for i, endpoint := range candidates {
|
||
res := s.tryEndpoint(w, r, endpoint, &req, i)
|
||
s.usage.Record(res.Attempt)
|
||
if res.Committed {
|
||
// 已向客户端写出：无论成败都不能再换道
|
||
if res.Status == "ok" {
|
||
s.routes.Sticky.Remember(req.RunID, endpoint.Name)
|
||
s.quotaSettle(req.UserID, reserved, res.Attempt.PromptTokens+res.Attempt.CompletionTokens)
|
||
} else {
|
||
// 流中失败：已发生的 token 照实结算，其余退款
|
||
s.quotaSettle(req.UserID, reserved, res.Attempt.PromptTokens+res.Attempt.CompletionTokens)
|
||
}
|
||
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
|
||
}
|
||
// 三段式之三：全部候选失败 → 全额退款
|
||
_ = s.quotaRefund(req.UserID, reserved)
|
||
writeError(w, http.StatusServiceUnavailable, "HJ-2001", "全部候选端点失败："+errString(lastErr))
|
||
}
|
||
|
||
// candidates 候选序列：BYOK 独占（不回落平台端点）；否则注册表候选。
|
||
func (s *Server) candidates(req *ir.ChatRequest) ([]route.Endpoint, error) {
|
||
if req.ByokKeyID == "" {
|
||
return s.routes.Candidates(req.Slot, req.RunID)
|
||
}
|
||
if s.cfg.Vault == nil {
|
||
return nil, fmt.Errorf("HJ-2003 BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）")
|
||
}
|
||
key, plain, err := s.cfg.Vault.Resolve(req.UserID, req.ByokKeyID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
spec := route.EndpointSpec{Type: key.Type, Name: "byok:" + key.ID, BaseURL: key.BaseURL, APIKey: plain, Model: key.Model}
|
||
adaptor, ok := route.BuildAdaptor(spec)
|
||
if !ok {
|
||
return nil, fmt.Errorf("HJ-2003 BYOK 密钥类型不受支持: %s", key.Type)
|
||
}
|
||
return []route.Endpoint{{
|
||
Name: spec.Name, Type: key.Type, Adaptor: adaptor,
|
||
Priority: -100, CredentialKey: spec.Name, UpstreamModel: key.Model,
|
||
}}, nil
|
||
}
|
||
|
||
func (s *Server) quotaReserve(userID string, tokens int64) (int64, error) {
|
||
if s.cfg.Quota == nil {
|
||
return 0, nil
|
||
}
|
||
return s.cfg.Quota.Reserve(userID, tokens)
|
||
}
|
||
|
||
func (s *Server) quotaSettle(userID string, reserved int64, actual int) error {
|
||
if s.cfg.Quota == nil {
|
||
return nil
|
||
}
|
||
return s.cfg.Quota.Settle(userID, reserved, int64(actual))
|
||
}
|
||
|
||
func (s *Server) quotaRefund(userID string, reserved int64) error {
|
||
if s.cfg.Quota == nil {
|
||
return nil
|
||
}
|
||
return s.cfg.Quota.Refund(userID, reserved)
|
||
}
|
||
|
||
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, seq int) 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{
|
||
ID: fmt.Sprintf("%s#%s#%d", req.RunID, endpoint.Name, seq),
|
||
RunID: req.RunID, CharacterID: req.CharacterID, UserID: req.UserID,
|
||
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"
|
||
}
|
||
|
||
// meUser 从代理头取用户标识（灵台鉴权后注入，服务令牌即信任边界）。
|
||
func meUser(r *http.Request) string { return r.Header.Get("X-Hunjing-User-Id") }
|
||
|
||
// handleUsage 用户用量聚合（/me/usage；附配额余额）。
|
||
func (s *Server) handleUsage(w http.ResponseWriter, r *http.Request) {
|
||
userID := meUser(r)
|
||
rows, err := s.usage.Aggregate(userID)
|
||
if err != nil {
|
||
writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error())
|
||
return
|
||
}
|
||
resp := map[string]any{"rows": rows}
|
||
if s.cfg.Quota != nil {
|
||
budget, used, reserved, err := s.cfg.Quota.Budget(userID)
|
||
if err == nil {
|
||
resp["quota"] = map[string]int64{"budget": budget, "used": used, "reserved": reserved}
|
||
}
|
||
}
|
||
w.Header().Set("Content-Type", "application/json")
|
||
_ = json.NewEncoder(w).Encode(resp)
|
||
}
|
||
|
||
// handleUsageSample 对账抽样（/admin/usage/sample?limit=N；最近 N 条原始记录）。
|
||
func (s *Server) handleUsageSample(w http.ResponseWriter, r *http.Request) {
|
||
limit := 50
|
||
if v := r.URL.Query().Get("limit"); v != "" {
|
||
if _, err := fmt.Sscanf(v, "%d", &limit); err != nil || limit <= 0 || limit > 1000 {
|
||
writeError(w, http.StatusBadRequest, "HJ-2002", "limit 须为 1..1000")
|
||
return
|
||
}
|
||
}
|
||
rows, err := s.usage.Sample(limit)
|
||
if err != nil {
|
||
writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error())
|
||
return
|
||
}
|
||
w.Header().Set("Content-Type", "application/json")
|
||
_ = json.NewEncoder(w).Encode(map[string]any{"rows": rows})
|
||
}
|
||
|
||
func (s *Server) handleByokList(w http.ResponseWriter, r *http.Request) {
|
||
if s.cfg.Vault == nil {
|
||
writeError(w, http.StatusServiceUnavailable, "HJ-2003", "BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）")
|
||
return
|
||
}
|
||
keys, err := s.cfg.Vault.List(meUser(r))
|
||
if err != nil {
|
||
writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error())
|
||
return
|
||
}
|
||
w.Header().Set("Content-Type", "application/json")
|
||
_ = json.NewEncoder(w).Encode(map[string]any{"keys": keys})
|
||
}
|
||
|
||
type byokRegisterRequest struct {
|
||
Name string `json:"name"`
|
||
Type string `json:"type"` // openai_compatible / anthropic / gemini
|
||
BaseURL string `json:"base_url"`
|
||
APIKey string `json:"api_key"`
|
||
Model string `json:"model"`
|
||
}
|
||
|
||
func (s *Server) handleByokRegister(w http.ResponseWriter, r *http.Request) {
|
||
if s.cfg.Vault == nil {
|
||
writeError(w, http.StatusServiceUnavailable, "HJ-2003", "BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）")
|
||
return
|
||
}
|
||
var body byokRegisterRequest
|
||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||
writeError(w, http.StatusBadRequest, "HJ-2002", "请求解析失败: "+err.Error())
|
||
return
|
||
}
|
||
switch body.Type {
|
||
case "openai_compatible", "anthropic", "gemini":
|
||
default:
|
||
writeError(w, http.StatusBadRequest, "HJ-2002", "type 须为 openai_compatible / anthropic / gemini")
|
||
return
|
||
}
|
||
if body.APIKey == "" || body.BaseURL == "" || body.Model == "" {
|
||
writeError(w, http.StatusBadRequest, "HJ-2002", "api_key / base_url / model 不能为空")
|
||
return
|
||
}
|
||
key, err := s.cfg.Vault.Register(meUser(r), body.Name, body.Type, body.BaseURL, body.APIKey, body.Model)
|
||
if err != nil {
|
||
writeError(w, http.StatusBadRequest, "HJ-2003", err.Error())
|
||
return
|
||
}
|
||
w.Header().Set("Content-Type", "application/json")
|
||
w.WriteHeader(http.StatusCreated)
|
||
_ = json.NewEncoder(w).Encode(key)
|
||
}
|
||
|
||
func (s *Server) handleByokRevoke(w http.ResponseWriter, r *http.Request) {
|
||
if s.cfg.Vault == nil {
|
||
writeError(w, http.StatusServiceUnavailable, "HJ-2003", "BYOK 功能未启用（SHENSHI_BYOK_MASTER_KEY 未配置）")
|
||
return
|
||
}
|
||
if err := s.cfg.Vault.Revoke(meUser(r), r.PathValue("id")); err != nil {
|
||
if errors.Is(err, store.ErrByokNotFound) {
|
||
writeError(w, http.StatusNotFound, "HJ-2003", err.Error())
|
||
return
|
||
}
|
||
writeError(w, http.StatusInternalServerError, "HJ-2001", err.Error())
|
||
return
|
||
}
|
||
w.WriteHeader(http.StatusNoContent)
|
||
}
|
||
|
||
// 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)
|
||
}
|