mirror of
https://gitee.com/wanwujie/sub2api
synced 2026-04-19 14:24:45 +08:00
refactor: decouple claude max cache simulation from RecordUsage
Extract setupClaudeMaxStreamingHook and applyClaudeMaxNonStreamingRewrite facade functions to helpers file. RecordUsage now uses detect-only (no mutation), client response rewriting handled at Forward layer.
This commit is contained in:
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
@@ -51,6 +52,18 @@ func claudeMaxGroupFromGinContext(c *gin.Context) *Group {
|
||||
return apiKey.Group
|
||||
}
|
||||
|
||||
func parsedRequestFromGinContext(c *gin.Context) *ParsedRequest {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
raw, exists := c.Get("parsed_request")
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
parsed, _ := raw.(*ParsedRequest)
|
||||
return parsed
|
||||
}
|
||||
|
||||
func applyClaudeMaxSimulationToUsage(ctx context.Context, usage *ClaudeUsage, model string, accountID int64) claudeMaxCacheBillingOutcome {
|
||||
var out claudeMaxCacheBillingOutcome
|
||||
if usage == nil {
|
||||
@@ -145,3 +158,39 @@ func usageIntFromAny(v any) int {
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// setupClaudeMaxStreamingHook 为 Antigravity 流式路径设置 SSE usage 改写 hook。
|
||||
func setupClaudeMaxStreamingHook(c *gin.Context, processor *antigravity.StreamingProcessor, originalModel string, accountID int64) {
|
||||
group := claudeMaxGroupFromGinContext(c)
|
||||
parsed := parsedRequestFromGinContext(c)
|
||||
if !shouldApplyClaudeMaxBillingRulesForUsage(group, originalModel, parsed) {
|
||||
return
|
||||
}
|
||||
processor.SetUsageMapHook(func(usageMap map[string]any) {
|
||||
svcUsage := claudeUsageFromJSONMap(usageMap)
|
||||
outcome := applyClaudeMaxCacheBillingPolicyToUsage(&svcUsage, parsed, group, originalModel, accountID)
|
||||
if outcome.Simulated {
|
||||
rewriteClaudeUsageJSONMap(usageMap, svcUsage)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// applyClaudeMaxNonStreamingRewrite 为 Antigravity 非流式路径改写响应体中的 usage。
|
||||
func applyClaudeMaxNonStreamingRewrite(c *gin.Context, claudeResp []byte, agUsage *antigravity.ClaudeUsage, originalModel string, accountID int64) []byte {
|
||||
group := claudeMaxGroupFromGinContext(c)
|
||||
parsed := parsedRequestFromGinContext(c)
|
||||
if !shouldApplyClaudeMaxBillingRulesForUsage(group, originalModel, parsed) {
|
||||
return claudeResp
|
||||
}
|
||||
svcUsage := &ClaudeUsage{
|
||||
InputTokens: agUsage.InputTokens,
|
||||
OutputTokens: agUsage.OutputTokens,
|
||||
CacheCreationInputTokens: agUsage.CacheCreationInputTokens,
|
||||
CacheReadInputTokens: agUsage.CacheReadInputTokens,
|
||||
}
|
||||
outcome := applyClaudeMaxCacheBillingPolicyToUsage(svcUsage, parsed, group, originalModel, accountID)
|
||||
if outcome.Simulated {
|
||||
return rewriteClaudeUsageJSONBytes(claudeResp, *svcUsage)
|
||||
}
|
||||
return claudeResp
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user