M4 收口：神识用户面与计量闭环——BYOK/配额三段式/usage 落库/Anthropic+Gemini 适配器/神识阁
- 神识 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 标记完成、能力矩阵神识行升已验证
This commit is contained in:
chenxin committed
1 parent
462d570751
commit
144415247e
51 files changed
+3533
-90
No files matched your search
@@ -0,0 +1,195 @@
|
||||
// BYOK / 配额三段式 / 用量归属的服务级场景测试。
|
||||
package server
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"hunjing/shenshi/internal/byok"
|
||||
"hunjing/shenshi/internal/quota"
|
||||
"hunjing/shenshi/internal/route"
|
||||
"hunjing/shenshi/internal/store"
|
||||
)
|
||||
|
||||
func testMasterKey() string {
|
||||
b := make([]byte, 32)
|
||||
for i := range b {
|
||||
b[i] = byte(i + 1)
|
||||
}
|
||||
return hex.EncodeToString(b)
|
||||
}
|
||||
|
||||
// fullServer 装配 store/byok/quota 的完整服务；eps 为注册表候选（按优先级升序尝试）。
|
||||
func fullServer(t *testing.T, budget int64, eps ...route.Endpoint) (*httptest.Server, store.Store, *byok.Vault) {
|
||||
t.Helper()
|
||||
st := store.NewMem()
|
||||
vault, err := byok.NewVault(testMasterKey(), st)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
registry := registryWith(t, time.Second, eps...)
|
||||
srv := New(Config{Store: st, Vault: vault, Quota: quota.New(st, budget)}, registry)
|
||||
ts := httptest.NewServer(srv)
|
||||
t.Cleanup(ts.Close)
|
||||
return ts, st, vault
|
||||
}
|
||||
|
||||
func postChat(ts *httptest.Server, body, userID string) (*http.Response, string) {
|
||||
req, _ := http.NewRequest(http.MethodPost, ts.URL+"/v1/chat/completions", strings.NewReader(body))
|
||||
req.Header.Set("content-type", "application/json")
|
||||
if userID != "" {
|
||||
req.Header.Set("X-Hunjing-User-Id", userID)
|
||||
}
|
||||
resp, err := ts.Client().Do(req)
|
||||
if err != nil {
|
||||
return nil, ""
|
||||
}
|
||||
var sb strings.Builder
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
for scanner.Scan() {
|
||||
sb.WriteString(scanner.Text())
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
return resp, sb.String()
|
||||
}
|
||||
|
||||
func byokChatBody(run, extra string) string {
|
||||
return `{"ir_version":1,"slot":"chat","messages":[{"role":"user","content":"hi"}],"run_id":"` + run + `"` + extra + `,"stream":true}`
|
||||
}
|
||||
|
||||
// 场景：BYOK 独占路由——请求携带 byok_key_id 时只走用户端点，不回落平台候选。
|
||||
func TestByokExclusiveRouting(t *testing.T) {
|
||||
platform := fakeOpenAI(t, 0, 0, "平台端点")
|
||||
defer platform.Close()
|
||||
ts, _, vault := fullServer(t, 0, oaEndpoint("platform", platform.URL, 10))
|
||||
|
||||
byokUpstream := fakeOpenAI(t, 0, 0, "用户端点")
|
||||
defer byokUpstream.Close()
|
||||
key, err := vault.Register("u1", "自带", "openai_compatible", byokUpstream.URL, "sk-mine", "my-model")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
body := byokChatBody("run_byok", `,"user_id":"u1","byok_key_id":"`+key.ID+`"`)
|
||||
resp, out := postChat(ts, body, "u1")
|
||||
if resp == nil || resp.StatusCode != 200 {
|
||||
t.Fatalf("BYOK 请求失败: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, "用户端点") || strings.Contains(out, "平台端点") {
|
||||
t.Fatalf("应独占 BYOK 端点: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
// 场景：BYOK 无效引用 → 整体失败，不回落平台端点。
|
||||
func TestByokInvalidReferenceRejected(t *testing.T) {
|
||||
platform := fakeOpenAI(t, 0, 0, "平台端点")
|
||||
defer platform.Close()
|
||||
ts, _, _ := fullServer(t, 0, oaEndpoint("platform", platform.URL, 10))
|
||||
|
||||
resp, out := postChat(ts, byokChatBody("run_bad", `,"user_id":"u1","byok_key_id":"byok_missing"`), "u1")
|
||||
if resp == nil || resp.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Fatalf("应 503: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, "HJ-2001") || strings.Contains(out, "平台端点") {
|
||||
t.Fatalf("应整体失败而非回落: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
// 场景：配额不足 → HJ-2004 预拒，且不产生用量记录。
|
||||
func TestQuotaRejectsWhenExhausted(t *testing.T) {
|
||||
upstream := fakeOpenAI(t, 0, 0, "内容")
|
||||
defer upstream.Close()
|
||||
ts, st, _ := fullServer(t, 10, oaEndpoint("up", upstream.URL, 10))
|
||||
|
||||
resp, out := postChat(ts, byokChatBody("run_q", `,"user_id":"u1","max_tokens":100`), "u1")
|
||||
if resp == nil || resp.StatusCode != http.StatusPaymentRequired {
|
||||
t.Fatalf("配额不足应 402: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, "HJ-2004") {
|
||||
t.Fatalf("应携带 HJ-2004: %s", out)
|
||||
}
|
||||
sample, _ := st.SampleUsage(10)
|
||||
if len(sample) != 0 {
|
||||
t.Fatalf("预拒请求不应产生用量记录: %+v", sample)
|
||||
}
|
||||
}
|
||||
|
||||
// 场景：成功结算释放差值；全败全额退款——余额不残留。
|
||||
func TestQuotaSettleAndRefundBalance(t *testing.T) {
|
||||
upstream := fakeOpenAI(t, 0, 0, "内容")
|
||||
bad := fakeOpenAI(t, 0, http.StatusForbidden, "")
|
||||
defer upstream.Close()
|
||||
defer bad.Close()
|
||||
ts, st, _ := fullServer(t, 1000, oaEndpoint("up", upstream.URL, 10))
|
||||
|
||||
// 成功：预扣 max_tokens=200，实际 3+4=7
|
||||
if resp, out := postChat(ts, byokChatBody("run_ok", `,"user_id":"u1","max_tokens":200`), "u1"); resp == nil || resp.StatusCode != 200 {
|
||||
t.Fatalf("成功请求失败: %s", out)
|
||||
}
|
||||
b, _ := st.Balance("u1")
|
||||
if b.Used != 7 || b.Reserved != 0 {
|
||||
t.Fatalf("结算后账目错误: %+v", b)
|
||||
}
|
||||
|
||||
// 已用 7，预算 1000：请求 max_tokens 994 应被拒（7+994 > 1000）
|
||||
if resp, out := postChat(ts, byokChatBody("run_over", `,"user_id":"u1","max_tokens":994`), "u1"); resp == nil || resp.StatusCode != http.StatusPaymentRequired {
|
||||
t.Fatalf("超预算应拒: %s", out)
|
||||
}
|
||||
|
||||
// 全候选失败 → 全额退款
|
||||
tsBad, stBad, _ := fullServer(t, 1000, oaEndpoint("bad", bad.URL, 10))
|
||||
if resp, out := postChat(tsBad, byokChatBody("run_af", `,"user_id":"u2","max_tokens":500`), "u2"); resp == nil || resp.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Fatalf("全败应 503: %s", out)
|
||||
}
|
||||
b2, _ := stBad.Balance("u2")
|
||||
if b2.Reserved != 0 || b2.Used != 0 {
|
||||
t.Fatalf("全败后应全额退款: %+v", b2)
|
||||
}
|
||||
}
|
||||
|
||||
// 场景：用量归属计量——/me/usage 按用户过滤并附配额总览。
|
||||
func TestUsageAttributionByUser(t *testing.T) {
|
||||
upstream := fakeOpenAI(t, 0, 0, "内容")
|
||||
defer upstream.Close()
|
||||
ts, _, _ := fullServer(t, 0, oaEndpoint("up", upstream.URL, 10))
|
||||
|
||||
for _, b := range []string{
|
||||
byokChatBody("run_u1", `,"user_id":"u1"`),
|
||||
byokChatBody("run_u2", `,"user_id":"u2"`),
|
||||
byokChatBody("run_u1b", `,"user_id":"u1"`),
|
||||
} {
|
||||
if resp, out := postChat(ts, b, ""); resp == nil || resp.StatusCode != 200 {
|
||||
t.Fatalf("请求失败: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
req, _ := http.NewRequest(http.MethodGet, ts.URL+"/me/usage", nil)
|
||||
req.Header.Set("X-Hunjing-User-Id", "u1")
|
||||
resp, err := ts.Client().Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var payload struct {
|
||||
Rows []struct {
|
||||
UserID string `json:"user_id"`
|
||||
Requests int `json:"requests"`
|
||||
} `json:"rows"`
|
||||
Quota *struct {
|
||||
Budget int64 `json:"budget"`
|
||||
Used int64 `json:"used"`
|
||||
} `json:"quota"`
|
||||
}
|
||||
_ = json.NewDecoder(resp.Body).Decode(&payload)
|
||||
if len(payload.Rows) != 1 || payload.Rows[0].UserID != "u1" || payload.Rows[0].Requests != 2 {
|
||||
t.Fatalf("u1 聚合错误: %+v", payload.Rows)
|
||||
}
|
||||
if payload.Quota == nil || payload.Quota.Used != 14 {
|
||||
t.Fatalf("配额总览错误: %+v", payload.Quota)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user