From 0f87322f49d4ae28ae13da9a786565b56d50bee3 Mon Sep 17 00:00:00 2001 From: martin <406615373@qq.com> Date: Sun, 26 Jul 2026 11:03:22 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=20P1=20=E5=8A=9F?= =?UTF-8?q?=E8=83=BD=E6=AD=A3=E7=A1=AE=E6=80=A7=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/handler/relay/response_capture.go | 17 ++++++ .../handler/relay/response_capture_test.go | 48 +++++++++++++++ internal/logic/billing/pricing.go | 6 +- internal/logic/common/cache.go | 44 ++++++++++---- internal/logic/common/cache_test.go | 21 +++++++ internal/logic/common/verify_code.go | 13 ++++- internal/logic/open/open_api.go | 30 ++++++---- internal/logic/relay/provider.go | 58 ++++++++++++++++--- internal/logic/relay/provider_test.go | 27 +++++++++ relay/handler/chat_handler.go | 12 ++++ relay/handler/claude_handler.go | 3 + relay/handler/error_writer_test.go | 34 +++++++++++ relay/handler/gemini_handler.go | 4 ++ 13 files changed, 282 insertions(+), 35 deletions(-) create mode 100644 internal/handler/relay/response_capture_test.go create mode 100644 internal/logic/common/cache_test.go diff --git a/internal/handler/relay/response_capture.go b/internal/handler/relay/response_capture.go index d8249db..09e3585 100644 --- a/internal/handler/relay/response_capture.go +++ b/internal/handler/relay/response_capture.go @@ -33,6 +33,10 @@ func NewResponseCaptureWriter(w http.ResponseWriter) *ResponseCaptureWriter { // WriteHeader 捕获状态码并委托给底层 writer func (w *ResponseCaptureWriter) WriteHeader(code int) { w.mu.Lock() + if w.status != 0 { + w.mu.Unlock() + return + } w.status = code w.mu.Unlock() w.ResponseWriter.WriteHeader(code) @@ -40,6 +44,12 @@ func (w *ResponseCaptureWriter) WriteHeader(code int) { // Write 写入客户端并同时捕获到 buffer（head + tail 环形缓冲区） func (w *ResponseCaptureWriter) Write(b []byte) (int, error) { + w.mu.Lock() + if w.status == 0 { + w.status = http.StatusOK + } + w.mu.Unlock() + n, err := w.ResponseWriter.Write(b) w.bytesWritten.Add(int64(n)) w.mu.Lock() @@ -119,6 +129,13 @@ func (w *ResponseCaptureWriter) BytesWritten() int64 { return w.bytesWritten.Load() } +// ResponseCommitted reports whether headers or body bytes have already been sent. +func (w *ResponseCaptureWriter) ResponseCommitted() bool { + w.mu.Lock() + defer w.mu.Unlock() + return w.status != 0 || w.bytesWritten.Load() > 0 +} + // ResponseHeaders 返回响应头快照（调用时从底层 ResponseWriter 提取） func (w *ResponseCaptureWriter) ResponseHeaders() map[string]string { headers := make(map[string]string) diff --git a/internal/handler/relay/response_capture_test.go b/internal/handler/relay/response_capture_test.go new file mode 100644 index 0000000..34f822c --- /dev/null +++ b/internal/handler/relay/response_capture_test.go @@ -0,0 +1,48 @@ +package relay + +import ( + "net/http" + "net/http/httptest" + "testing" +) + +func TestResponseCaptureWriterCommitted(t *testing.T) { + t.Run("fresh", func(t *testing.T) { + capture := NewResponseCaptureWriter(httptest.NewRecorder()) + if capture.ResponseCommitted() { + t.Fatal("fresh writer should not be committed") + } + }) + + t.Run("header", func(t *testing.T) { + capture := NewResponseCaptureWriter(httptest.NewRecorder()) + capture.WriteHeader(http.StatusAccepted) + if !capture.ResponseCommitted() { + t.Fatal("writer should be committed after WriteHeader") + } + if got := capture.StatusCode(); got != http.StatusAccepted { + t.Fatalf("status = %d, want %d", got, http.StatusAccepted) + } + }) + + t.Run("body", func(t *testing.T) { + capture := NewResponseCaptureWriter(httptest.NewRecorder()) + _, _ = capture.Write([]byte("partial")) + if !capture.ResponseCommitted() { + t.Fatal("writer should be committed after Write") + } + capture.WriteHeader(http.StatusInternalServerError) + if got := capture.StatusCode(); got != http.StatusOK { + t.Fatalf("status = %d, want implicit status %d", got, http.StatusOK) + } + }) +} + +func TestResponseCaptureWriterKeepsFirstStatus(t *testing.T) { + capture := NewResponseCaptureWriter(httptest.NewRecorder()) + capture.WriteHeader(http.StatusAccepted) + capture.WriteHeader(http.StatusInternalServerError) + if got := capture.StatusCode(); got != http.StatusAccepted { + t.Fatalf("status = %d, want first status %d", got, http.StatusAccepted) + } +} diff --git a/internal/logic/billing/pricing.go b/internal/logic/billing/pricing.go index 5d12a29..7d15bb5 100644 --- a/internal/logic/billing/pricing.go +++ b/internal/logic/billing/pricing.go @@ -409,7 +409,7 @@ type CostBreakdown struct { func EstimatePreDeductAmount(ctx context.Context, tenantID int64, modelName string, inputTokens, requestedMaxTokens int, isStream bool) (float64, error) { pricing, err := GetModelPrice(ctx, tenantID, modelName) if err != nil { - return 0.01, nil + return 0, gerror.Wrapf(err, "estimate pre-deduct: get model price") } // 按次计费：直接用单价 @@ -430,7 +430,7 @@ func EstimatePreDeductAmount(ctx context.Context, tenantID int64, modelName stri Fields("max_output_tokens"). Scan(&model) if err != nil { - return 0.01, nil + return 0, gerror.Wrapf(err, "estimate pre-deduct: query model output limit") } maxOutput := 4096 @@ -448,7 +448,7 @@ func EstimatePreDeductAmount(ctx context.Context, tenantID int64, modelName stri breakdown, err := CalculateCost(ctx, tenantID, modelName, inputTokens, estimatedOutput) if err != nil { - return 0.01, nil + return 0, gerror.Wrapf(err, "estimate pre-deduct: calculate cost") } if breakdown.TotalCost > 1.0 { diff --git a/internal/logic/common/cache.go b/internal/logic/common/cache.go index df19188..4576e0b 100644 --- a/internal/logic/common/cache.go +++ b/internal/logic/common/cache.go @@ -3,12 +3,15 @@ package common import ( "context" "encoding/json" + "math" + "math/rand/v2" "time" "github.com/gogf/gf/v2/database/gredis" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/os/gcache" "github.com/gogf/gf/v2/util/gconv" + "golang.org/x/sync/singleflight" ) // Cache provides a two-level caching mechanism: @@ -20,6 +23,19 @@ import ( type Cache struct { prefix string ttl time.Duration + group singleflight.Group +} + +func jitterTTL(ttl time.Duration) time.Duration { + if ttl <= 0 { + return ttl + } + spread := ttl / 10 + if spread <= 0 { + return ttl + } + delta := time.Duration(rand.Int64N(int64(spread)*2+1)) - spread + return ttl + delta } // NewCache creates a new Cache instance with the given prefix and default TTL. @@ -43,13 +59,14 @@ func (c *Cache) Set(ctx context.Context, key string, value any, ttl ...time.Dura if len(ttl) > 0 && ttl[0] > 0 { expire = ttl[0] } + expire = jitterTTL(expire) fullKey := c.fullKey(key) // L1: memory cache (native Go struct storage) gcache.Set(ctx, fullKey, value, expire) // L2: Redis cache (JSON serialized) - ttlSeconds := int64(expire.Seconds()) + ttlSeconds := int64(math.Ceil(expire.Seconds())) jsonBytes, err := json.Marshal(value) if err != nil { g.Log().Warningf(ctx, "[Cache] JSON marshal failed key=%s: %v", fullKey, err) @@ -83,11 +100,11 @@ func (c *Cache) Get(ctx context.Context, key string) (any, bool) { // Unmarshal JSON to get the original Go value var raw any if unmarshalErr := json.Unmarshal([]byte(jsonStr), &raw); unmarshalErr == nil { - gcache.Set(ctx, fullKey, raw, c.ttl) + gcache.Set(ctx, fullKey, raw, jitterTTL(c.ttl)) return raw, true } // Fallback: treat as plain string (backward compat) - gcache.Set(ctx, fullKey, jsonStr, c.ttl) + gcache.Set(ctx, fullKey, jsonStr, jitterTTL(c.ttl)) return jsonStr, true } } @@ -119,7 +136,7 @@ func (c *Cache) GetJSON(ctx context.Context, key string, target any) bool { return false } // Backfill L1 - gcache.Set(ctx, fullKey, target, c.ttl) + gcache.Set(ctx, fullKey, target, jitterTTL(c.ttl)) return true } } @@ -133,13 +150,18 @@ func (c *Cache) GetOrSet(ctx context.Context, key string, fn func(ctx context.Co return val, nil } - val, err := fn(ctx) - if err != nil { - return nil, err - } - - c.Set(ctx, key, val) - return val, nil + val, err, _ := c.group.Do(c.fullKey(key), func() (any, error) { + if cached, ok := c.Get(ctx, key); ok { + return cached, nil + } + loaded, loadErr := fn(ctx) + if loadErr != nil { + return nil, loadErr + } + c.Set(ctx, key, loaded) + return loaded, nil + }) + return val, err } // Delete removes a value from both L1 and L2. diff --git a/internal/logic/common/cache_test.go b/internal/logic/common/cache_test.go new file mode 100644 index 0000000..9e13419 --- /dev/null +++ b/internal/logic/common/cache_test.go @@ -0,0 +1,21 @@ +package common + +import ( + "testing" + "time" +) + +func TestJitterTTLRange(t *testing.T) { + base := 10 * time.Minute + minTTL := base - base/10 + maxTTL := base + base/10 + for range 1000 { + got := jitterTTL(base) + if got < minTTL || got > maxTTL { + t.Fatalf("jitterTTL(%v) = %v, want range [%v, %v]", base, got, minTTL, maxTTL) + } + } + if got := jitterTTL(0); got != 0 { + t.Fatalf("jitterTTL(0) = %v, want 0", got) + } +} diff --git a/internal/logic/common/verify_code.go b/internal/logic/common/verify_code.go index e946df7..8836e16 100644 --- a/internal/logic/common/verify_code.go +++ b/internal/logic/common/verify_code.go @@ -139,10 +139,12 @@ func VerifyCode(ctx context.Context, email, code, purpose string) error { UsedAt *time.Time `json:"used_at"` } - err := dao.SysEmailVerifyCodes.Ctx(ctx). + baseQuery := dao.SysEmailVerifyCodes.Ctx(ctx). Where("email", email). Where("purpose", purpose). - Where("used_at IS NULL"). + Where("used_at IS NULL") + err := baseQuery. + Where("expires_at > NOW()"). OrderDesc("created_at"). Limit(1). Scan(&record) @@ -151,6 +153,13 @@ func VerifyCode(ctx context.Context, email, code, purpose string) error { } if record == nil { + expiredCount, countErr := baseQuery.Where("expires_at <= NOW()").Count() + if countErr != nil { + return gerror.Wrapf(countErr, "query expired verify code") + } + if expiredCount > 0 { + return NewBusinessError(consts.CodeEmailVerifyExpired, consts.MsgEmailVerifyExpired) + } incrVerifyFail(ctx, failKey) return NewBusinessError(consts.CodeVerifyCodeInvalid, "验证码错误") } diff --git a/internal/logic/open/open_api.go b/internal/logic/open/open_api.go index 6741e82..7734e5a 100644 --- a/internal/logic/open/open_api.go +++ b/internal/logic/open/open_api.go @@ -551,22 +551,25 @@ func (s *sOpen) OpenUsageQuery(ctx context.Context, req *v1.OpenUsageQueryReq) ( pageSize = 20 } - m := dao.BilUsageLogs.Ctx(ctx).Where("tenant_id", tenantID) - m = m.Where("created_at >= ?", req.StartDate+" 00:00:00") - m = m.Where("created_at <= ?", req.EndDate+" 23:59:59") + m := dao.BilUsageLogs.Ctx(ctx). + As("u"). + Where("u.tenant_id", tenantID). + Where("u.created_at >= ?", req.StartDate+" 00:00:00"). + Where("u.created_at <= ?", req.EndDate+" 23:59:59") selectFields := "" groupBy := "" switch req.GroupBy { case "model": - selectFields = "model_name, COUNT(*) as request_count, SUM(input_tokens) as prompt_tokens, SUM(output_tokens) as completion_tokens, SUM(input_tokens+output_tokens) as total_tokens, SUM(actual_cost) as cost" - groupBy = "model_name" + selectFields = "u.model_name, COUNT(*) as request_count, SUM(u.input_tokens) as prompt_tokens, SUM(u.output_tokens) as completion_tokens, SUM(u.input_tokens+u.output_tokens) as total_tokens, SUM(u.actual_cost) as cost" + groupBy = "u.model_name" case "key": - selectFields = "key_prefix as key_name, COUNT(*) as request_count, SUM(input_tokens) as prompt_tokens, SUM(output_tokens) as completion_tokens, SUM(input_tokens+output_tokens) as total_tokens, SUM(actual_cost) as cost" - groupBy = "key_prefix" + m = m.LeftJoin("api_keys", "ak", "u.api_key_id = ak.id AND u.tenant_id = ak.tenant_id") + selectFields = "COALESCE(ak.key_prefix, '') as key_name, COUNT(*) as request_count, SUM(u.input_tokens) as prompt_tokens, SUM(u.output_tokens) as completion_tokens, SUM(u.input_tokens+u.output_tokens) as total_tokens, SUM(u.actual_cost) as cost" + groupBy = "COALESCE(ak.key_prefix, '')" default: - selectFields = "DATE(created_at) as date, COUNT(*) as request_count, SUM(input_tokens) as prompt_tokens, SUM(output_tokens) as completion_tokens, SUM(input_tokens+output_tokens) as total_tokens, SUM(actual_cost) as cost" - groupBy = "DATE(created_at)" + selectFields = "DATE(u.created_at) as date, COUNT(*) as request_count, SUM(u.input_tokens) as prompt_tokens, SUM(u.output_tokens) as completion_tokens, SUM(u.input_tokens+u.output_tokens) as total_tokens, SUM(u.actual_cost) as cost" + groupBy = "DATE(u.created_at)" } type usageRow struct { @@ -580,8 +583,13 @@ func (s *sOpen) OpenUsageQuery(ctx context.Context, req *v1.OpenUsageQueryReq) ( Cost float64 `json:"cost"` } + total, err := m.Group(groupBy).Count() + if err != nil { + return nil, err + } + var rows []usageRow - err := m.Fields(selectFields).Group(groupBy).OrderDesc(groupBy).Page(page, pageSize).Scan(&rows) + err = m.Fields(selectFields).Group(groupBy).OrderDesc(groupBy).Page(page, pageSize).Scan(&rows) if err != nil { return nil, err } @@ -606,7 +614,7 @@ func (s *sOpen) OpenUsageQuery(ctx context.Context, req *v1.OpenUsageQueryReq) ( items = append(items, item) } - return &v1.OpenUsageQueryRes{List: items, Total: len(items), Page: page, PageSize: pageSize}, nil + return &v1.OpenUsageQueryRes{List: items, Total: total, Page: page, PageSize: pageSize}, nil } // ============================================================ diff --git a/internal/logic/relay/provider.go b/internal/logic/relay/provider.go index 0babb6a..7ecd1f7 100644 --- a/internal/logic/relay/provider.go +++ b/internal/logic/relay/provider.go @@ -7,6 +7,7 @@ import ( "sort" "strconv" "strings" + "sync" "time" "unicode/utf8" @@ -25,6 +26,37 @@ import ( // modelCache 模型信息缓存（TTL 600s） var modelCache = lcommon.NewCache("model", 600*time.Second) +const channelKeyLastUsedInterval = time.Minute + +var channelKeyUsage = newChannelKeyUsageTracker() + +type channelKeyUsageTracker struct { + mu sync.Mutex + touched map[int64]time.Time +} + +func newChannelKeyUsageTracker() *channelKeyUsageTracker { + return &channelKeyUsageTracker{touched: make(map[int64]time.Time)} +} + +func (t *channelKeyUsageTracker) claim(keyID int64, now time.Time) bool { + t.mu.Lock() + defer t.mu.Unlock() + if last, ok := t.touched[keyID]; ok && now.Sub(last) < channelKeyLastUsedInterval { + return false + } + t.touched[keyID] = now + return true +} + +func (t *channelKeyUsageTracker) release(keyID int64, claimedAt time.Time) { + t.mu.Lock() + defer t.mu.Unlock() + if last, ok := t.touched[keyID]; ok && last.Equal(claimedAt) { + delete(t.touched, keyID) + } +} + // tenantGroupModelCache 租户分组模型集合缓存（通过 lcommon.TenantGroupModelCache 共享） // 实际缓存实例在 lcommon 包中，admin 和 relay 共用 @@ -1050,19 +1082,12 @@ func getChannelKey(ctx context.Context, channelID int64) (string, error) { return "", common.ErrChannelUnavailable } - // 更新最后使用时间（用于监控） - if _, err := dao.ChnChannelKeys.Ctx(ctx). - Where("id", key.ID). - Data(do.ChnChannelKeys{LastUsedAt: gtime.Now()}). - Update(); err != nil { - g.Log().Debugf(ctx, "[DataProvider] update channel key LastUsedAt failed: keyID=%d, err=%v", key.ID, err) - } - encKey := GetEncryptionKey() decrypted, err := uc.DecryptString(encKey, key.EncryptedKey) if err != nil { return "", err } + touchChannelKeyLastUsed(ctx, key.ID) // OAuth 按需刷新 if key.KeyType == "oauth" && loauth.IsOAuthKeyData(decrypted) { @@ -1083,6 +1108,23 @@ func getChannelKey(ctx context.Context, channelID int64) (string, error) { return decrypted, nil } +func touchChannelKeyLastUsed(ctx context.Context, keyID int64) { + now := time.Now() + if !channelKeyUsage.claim(keyID, now) { + return + } + + _, err := dao.ChnChannelKeys.Ctx(ctx). + Where("id", keyID). + Where("last_used_at IS NULL OR last_used_at < ?", now.Add(-channelKeyLastUsedInterval)). + Data(do.ChnChannelKeys{LastUsedAt: gtime.New(now)}). + Update() + if err != nil { + channelKeyUsage.release(keyID, now) + g.Log().Debugf(ctx, "[DataProvider] update channel key LastUsedAt failed: keyID=%d, err=%v", keyID, err) + } +} + // refreshOAuthKey 刷新 OAuth 令牌并更新数据库 func refreshOAuthKey(ctx context.Context, keyID int64, oauthData *loauth.OAuthKeyData, encKey []byte) (string, error) { var newToken *loauth.OAuthKeyData diff --git a/internal/logic/relay/provider_test.go b/internal/logic/relay/provider_test.go index a7bf528..4bfc872 100644 --- a/internal/logic/relay/provider_test.go +++ b/internal/logic/relay/provider_test.go @@ -4,11 +4,38 @@ import ( "fmt" "strings" "testing" + "time" "unicode/utf8" "github.com/qianfree/team-api/relay/common" ) +func TestChannelKeyUsageTracker(t *testing.T) { + tracker := newChannelKeyUsageTracker() + now := time.Now() + if !tracker.claim(1, now) { + t.Fatal("first claim should update") + } + if tracker.claim(1, now.Add(channelKeyLastUsedInterval-time.Second)) { + t.Fatal("claim inside interval should be throttled") + } + if !tracker.claim(1, now.Add(channelKeyLastUsedInterval)) { + t.Fatal("claim at interval boundary should update") + } +} + +func TestChannelKeyUsageTrackerRelease(t *testing.T) { + tracker := newChannelKeyUsageTracker() + now := time.Now() + if !tracker.claim(1, now) { + t.Fatal("first claim should update") + } + tracker.release(1, now) + if !tracker.claim(1, now.Add(time.Second)) { + t.Fatal("released claim should be immediately retryable") + } +} + func TestSafeUTF8Truncate(t *testing.T) { tests := []struct { name string diff --git a/relay/handler/chat_handler.go b/relay/handler/chat_handler.go index a842399..d5130e9 100644 --- a/relay/handler/chat_handler.go +++ b/relay/handler/chat_handler.go @@ -144,6 +144,9 @@ func WriteRelayError(w http.ResponseWriter, err error) { g.Log().Infof(context.Background(), "[RelayError] Client disconnected during stream") return } + if responseCommitted(w) { + return + } // adaptor 已直接写入响应体（如 Gemini 原生格式透传），跳过二次写入 var prewritten *constant.RelayError @@ -197,3 +200,12 @@ func WriteRelayError(w http.ResponseWriter, err error) { }) _, _ = w.Write(errBody) } + +type responseCommitter interface { + ResponseCommitted() bool +} + +func responseCommitted(w http.ResponseWriter) bool { + committer, ok := w.(responseCommitter) + return ok && committer.ResponseCommitted() +} diff --git a/relay/handler/claude_handler.go b/relay/handler/claude_handler.go index a089a59..3ed41ff 100644 --- a/relay/handler/claude_handler.go +++ b/relay/handler/claude_handler.go @@ -24,6 +24,9 @@ func WriteClaudeRelayError(w http.ResponseWriter, err error) { g.Log().Infof(context.Background(), "[ClaudeRelayError] Client disconnected during stream") return } + if responseCommitted(w) { + return + } // adaptor 已直接写入响应体（如 Gemini 原生格式透传），跳过二次写入 var prewritten *constant.RelayError diff --git a/relay/handler/error_writer_test.go b/relay/handler/error_writer_test.go index 5aab9b9..e588a17 100644 --- a/relay/handler/error_writer_test.go +++ b/relay/handler/error_writer_test.go @@ -1,6 +1,7 @@ package handler import ( + "errors" "net/http" "net/http/httptest" "testing" @@ -8,6 +9,14 @@ import ( "github.com/qianfree/team-api/relay/constant" ) +type committedTestWriter struct { + *httptest.ResponseRecorder +} + +func (w *committedTestWriter) ResponseCommitted() bool { + return true +} + // TestErrorWriters_SkipResponseWritten 验证当 adaptor 已置 ResponseWritten=true 时， // 三个错误写入器都跳过二次写入（不写 header、不写 body）。 // 这是 Gemini/OpenAI「双重写入」修复的核心机制兜底。 @@ -65,3 +74,28 @@ func TestErrorWriters_WriteNormalError(t *testing.T) { }) } } + +func TestErrorWriters_SkipCommittedResponse(t *testing.T) { + cases := []struct { + name string + fn func(http.ResponseWriter, error) + }{ + {"WriteRelayError", WriteRelayError}, + {"WriteClaudeRelayError", WriteClaudeRelayError}, + {"WriteGeminiRelayError", WriteGeminiRelayError}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + _, _ = recorder.WriteString("data: partial\n\n") + writer := &committedTestWriter{ResponseRecorder: recorder} + before := recorder.Body.String() + + c.fn(writer, errors.New("upstream stream failed")) + + if got := recorder.Body.String(); got != before { + t.Fatalf("committed body changed: got %q, want %q", got, before) + } + }) + } +} diff --git a/relay/handler/gemini_handler.go b/relay/handler/gemini_handler.go index 1168a92..a3fa637 100644 --- a/relay/handler/gemini_handler.go +++ b/relay/handler/gemini_handler.go @@ -61,6 +61,10 @@ func HandleGeminiModelDetail(ctx context.Context, tenantID int64, modelName stri // WriteGeminiRelayError 写入 Gemini 格式的错误响应 func WriteGeminiRelayError(w http.ResponseWriter, err error) { + if responseCommitted(w) { + return + } + // adaptor 已直接写入响应体（如 Gemini 原生格式透传），跳过二次写入 var prewritten *constant.RelayError if errors.As(err, &prewritten) && prewritten.ResponseWritten {